Native JNI S3 PyTorch #2274
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
name: Native JNI S3 PyTorch | |
on: | |
workflow_dispatch: | |
inputs: | |
pt_version: | |
description: 'pytorch version' | |
required: false | |
cuda: | |
description: 'CUDA version' | |
required: true | |
default: 'cu128' | |
schedule: | |
- cron: '0 5 * * *' | |
permissions: | |
id-token: write | |
contents: read | |
jobs: | |
build-pytorch-jni-linux: | |
if: github.repository == 'deepjavalibrary/djl' | |
runs-on: ubuntu-latest | |
container: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu20.04 | |
steps: | |
- uses: actions/checkout@v4 | |
- name: Set up JDK 17 | |
uses: actions/setup-java@v4 | |
with: | |
distribution: 'corretto' | |
java-version: 17 | |
- uses: actions/cache@v3 | |
with: | |
path: ~/.gradle/caches | |
key: ${{ runner.os }}-gradle-${{ hashFiles('*/build.gradle.kts', 'engines/**/build.gradle.kts', 'extensions/**/build.gradle.kts') }} | |
restore-keys: | | |
${{ runner.os }}-gradle- | |
- name: Install Environment | |
run: | | |
apt-get update | |
DEBIAN_FRONTEND=noninteractive apt-get install -y locales curl unzip software-properties-common | |
apt-get install -y python3-pip python3-distutils | |
pip3 install awscli cmake | |
ln -s /usr/local/bin/cmake /usr/bin/cmake | |
- name: Release JNI prep | |
shell: bash | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
export PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
echo $PYTORCH_VERSION | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION | |
./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch" | |
./gradlew :engines:pytorch:pytorch-native:cleanJNI | |
# precxx11 is not supported since PyTorch 2.7.1 | |
if [[ $PYTORCH_VERSION != "2.7.1" ]]; then \ | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION | |
./gradlew :engines:pytorch:pytorch-jni:clean | |
PYTORCH_PRECXX11=true ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:clean :integration:test "-Dai.djl.default_engine=PyTorch" | |
./gradlew :engines:pytorch:pytorch-native:cleanJNI | |
fi | |
export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0" | |
CUDA_VERSION=${{ github.event.inputs.cuda }} | |
export CUDA_VERSION=${CUDA_VERSION:-cu128} | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=$CUDA_VERSION -Ppt_version=$PYTORCH_VERSION | |
./gradlew :engines:pytorch:pytorch-native:cleanJNI | |
# precxx11 is not supported since PyTorch 2.7.1 | |
if [[ $PYTORCH_VERSION != "2.7.1" ]]; then \ | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Pcuda=$CUDA_VERSION -Ppt_version=$PYTORCH_VERSION; \ | |
fi | |
- name: Configure AWS Credentials | |
uses: aws-actions/configure-aws-credentials@v4 | |
with: | |
role-to-assume: arn:aws:iam::425969335547:role/djl-ci-publish-role | |
aws-region: us-east-2 | |
- name: Copy files to S3 with the AWS CLI | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib | |
aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*" | |
build-pytorch-jni-windows: | |
if: github.repository == 'deepjavalibrary/djl' | |
runs-on: windows-latest | |
steps: | |
- uses: actions/checkout@v4 | |
- name: Set up JDK 17 | |
uses: actions/setup-java@v4 | |
with: | |
distribution: 'corretto' | |
java-version: 17 | |
- uses: actions/cache@v4 | |
with: | |
path: ~/.gradle/caches | |
key: ${{ runner.os }}-gradle-${{ hashFiles('*/build.gradle.kts', 'engines/**/build.gradle.kts', 'extensions/**/build.gradle.kts') }} | |
restore-keys: | | |
${{ runner.os }}-gradle- | |
- name: Release CPU JNI | |
shell: cmd | |
run: | | |
call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64 | |
gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=${{ github.event.inputs.pt_version }} | |
- name: Install CUDA 11.7 | |
shell: cmd | |
run: | | |
call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64 | |
curl.exe -L https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_516.94_windows.exe -o cuda.exe | |
curl.exe -L https://developer.download.nvidia.com/compute/redist/cudnn/v8.5.0/local_installers/11.7/cudnn-windows-x86_64-8.5.0.96_cuda11-archive.zip -o cudnn.zip | |
cuda.exe -s | |
unzip.exe cudnn.zip | |
cp.exe -a cudnn*/include cudnn*/lib cudnn*/bin "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.7/" | |
rm.exe -Rf cuda.exe cuda.exe cudnn.zip cudnn* | |
- name: Release cuda11 JNI | |
shell: cmd | |
run: | | |
rm.exe -Rf engines/pytorch/pytorch-native/libtorch | |
rm.exe -Rf engines/pytorch/pytorch-native/build | |
call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64 | |
set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v11.7" | |
set "PATH=%CUDA_PATH%\bin;%CUDA_PATH%\libnvvp;%PATH%" | |
set CUDA_VERSION=${{ github.event.inputs.cuda }} | |
if "%CUDA_VERSION%" == "" set CUDA_VERSION=cu128 | |
gradlew --rerun-tasks :engines:pytorch:pytorch-native:compileJNI -Pcuda=%CUDA_VERSION% -Ppt_version=${{ github.event.inputs.pt_version }} | |
- name: Configure Deployment AWS Credentials | |
uses: aws-actions/configure-aws-credentials@v4 | |
with: | |
role-to-assume: arn:aws:iam::425969335547:role/djl-ci-publish-role | |
aws-region: us-east-2 | |
- name: Copy files to S3 with the AWS CLI | |
shell: bash | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib | |
aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*" | |
build-pytorch-jni-arm64-macos: | |
if: github.repository == 'deepjavalibrary/djl' | |
runs-on: macos-latest-xlarge | |
steps: | |
- uses: actions/checkout@v4 | |
- name: Set up JDK 17 | |
uses: actions/setup-java@v4 | |
with: | |
java-version: 17 | |
distribution: corretto | |
architecture: aarch64 | |
- uses: actions/cache@v4 | |
with: | |
path: ~/.gradle/caches | |
key: ${{ runner.os }}-gradle-${{ hashFiles('*/build.gradle.kts', 'engines/**/build.gradle.kts', 'extensions/**/build.gradle.kts') }} | |
restore-keys: | | |
${{ runner.os }}-gradle- | |
- name: Release JNI prep | |
shell: bash | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
export PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
echo $PYTORCH_VERSION | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION | |
./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch" | |
- name: Configure Deployment AWS Credentials | |
uses: aws-actions/configure-aws-credentials@v4 | |
with: | |
role-to-assume: arn:aws:iam::425969335547:role/djl-ci-publish-role | |
aws-region: us-east-2 | |
- name: Copy files to S3 with the AWS CLI | |
shell: bash | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib | |
aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*" | |
build-pytorch-jni-aarch64: | |
if: github.repository == 'deepjavalibrary/djl' | |
runs-on: ubuntu-24.04-arm | |
container: | |
image: amazonlinux:2023 | |
env: | |
JAVA_HOME: /usr/lib/jvm/java-17-amazon-corretto.aarch64 | |
steps: | |
- name: Install compiler environment | |
run: | | |
dnf -y group install "Development Tools" | |
dnf -y install patch git cmake3 python3-devel python3-pip java-17-amazon-corretto-devel jq | |
pip3 install awscli --upgrade | |
- uses: actions/checkout@v4 | |
- name: Release JNI prep | |
run: | | |
export PATH=$PATH:$JAVA_HOME/bin | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
export PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
echo $PYTORCH_VERSION | |
if [[ $PYTORCH_VERSION == "2.7.1" ]]; then \ | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION; \ | |
else \ | |
./gradlew :engines:pytorch:pytorch-native:compileJNI -Pprecxx11 -Ppt_version=$PYTORCH_VERSION; \ | |
fi | |
./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch" | |
- name: Configure AWS Credentials | |
run: | | |
oidc_token=$(curl -H "Authorization: bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" \ | |
"$ACTIONS_ID_TOKEN_REQUEST_URL&audience=sts.amazonaws.com" | jq -r ".value") | |
echo "::add-mask::$oidc_token" | |
read -r AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN <<<"$(aws sts assume-role-with-web-identity \ | |
--region "us-east-2" \ | |
--role-arn "arn:aws:iam::425969335547:role/djl-ci-publish-role" \ | |
--role-session-name "native-jni-s3-pytorch-jni-aarch64" \ | |
--web-identity-token "$oidc_token" \ | |
--query "[Credentials.AccessKeyId, Credentials.SecretAccessKey, Credentials.SessionToken]" \ | |
--output text)" | |
echo "::add-mask::$AWS_ACCESS_KEY_ID" | |
echo "::add-mask::$AWS_SECRET_ACCESS_KEY" | |
echo "::add-mask::$AWS_SESSION_TOKEN" | |
echo "AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID" >> "$GITHUB_ENV" | |
echo "AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY" >> "$GITHUB_ENV" | |
echo "AWS_SESSION_TOKEN=$AWS_SESSION_TOKEN" >> "$GITHUB_ENV" | |
- name: Copy files to S3 with the AWS CLI | |
run: | | |
PYTORCH_VERSION=${{ github.event.inputs.pt_version }} | |
PYTORCH_VERSION=${PYTORCH_VERSION:-$(awk -F '=' '/pytorch/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)} | |
aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib | |
aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*" |