[FASTTEXT] Adds linux-aarch64 support (#3584)

* [FASTTEXT] arm support for linux * Add ci --------- Co-authored-by: Xin Yang <[email protected]>
deepjavalibrary · Jan 24, 2025 · d381c4d · d381c4d
1 parent c458401
commit d381c4d
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 2 deletions.
diff --git a/.github/workflows/native_s3_fasttext.yml b/.github/workflows/native_s3_fasttext.yml
@@ -54,6 +54,7 @@ jobs:
         run: |
           FASTTEXT_VERSION="$(awk -F '=' '/fasttext/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
           aws s3 sync extensions/fasttext/jnilib s3://djl-ai/publish/fasttext-${FASTTEXT_VERSION}/jnilib/
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/fasttext-${FASTTEXT_VERSION}/jnilib/*"
 
   build-fasttext-jni-arm64-osx:
     if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
@@ -85,3 +86,80 @@ jobs:
         run: |
           FASTTEXT_VERSION="$(awk -F '=' '/fasttext/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
           aws s3 sync extensions/fasttext/jnilib s3://djl-ai/publish/fasttext-${FASTTEXT_VERSION}/jnilib/
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/fasttext-${FASTTEXT_VERSION}/jnilib/*"
+
+  build-fasttext-jni-aarch64:
+    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
+    runs-on: [ self-hosted, aarch64 ]
+    timeout-minutes: 30
+    needs: create-aarch64-runner
+    container:
+      image: amazonlinux:2
+      env:
+        JAVA_HOME: /usr/lib/jvm/java-17-amazon-corretto.aarch64
+    steps:
+      - name: Install Environment
+        run: |
+          yum -y groupinstall "Development Tools"
+          yum -y install patch perl-IPC-Cmd cmake3 curl jq
+          yum -y install java-17-amazon-corretto-devel
+          ln -s /usr/bin/cmake3 /usr/bin/cmake
+          pip3 install awscli --upgrade
+      - uses: taiki-e/checkout-action@v1
+      - name: Release JNI prep
+        run: |
+          ./gradlew :extensions:fasttext:compileJNI
+          ./gradlew -Pjni :extensions:fasttext:test
+      - 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 "build-tokenizer-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: |
+          FASTTEXT_VERSION="$(awk -F '=' '/fasttext/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
+          aws s3 sync extensions/fasttext/jnilib s3://djl-ai/publish/fasttext-${FASTTEXT_VERSION}/jnilib/
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/fasttext-${FASTTEXT_VERSION}/jnilib/*"
+
+  create-aarch64-runner:
+    if: github.repository == 'deepjavalibrary/djl'
+    runs-on: [ self-hosted, scheduler ]
+    steps:
+      - name: Create new Graviton instance
+        id: create_aarch64
+        run: |
+          cd /home/ubuntu/djl_benchmark_script/scripts
+          token=$( curl -X POST -H "Authorization: token ${{ secrets.ACTION_RUNNER_PERSONAL_TOKEN }}" \
+          https://api.github.com/repos/deepjavalibrary/djl/actions/runners/registration-token \
+          --fail \
+          | jq '.token' | tr -d '"' )
+          ./start_instance.sh action_graviton $token djl
+    outputs:
+      aarch64_instance_id: ${{ steps.create_aarch64.outputs.action_graviton_instance_id }}
+
+  stop-runners:
+    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
+    runs-on: [ self-hosted, scheduler ]
+    needs: [ create-aarch64-runner, build-fasttext-jni-aarch64 ]
+    steps:
+      - name: Stop all instances
+        run: |
+          cd /home/ubuntu/djl_benchmark_script/scripts
+          instance_id=${{ needs.create-aarch64-runner.outputs.aarch64_instance_id }}
+          ./stop_instance.sh $instance_id
diff --git a/extensions/fasttext/build.gradle.kts b/extensions/fasttext/build.gradle.kts
@@ -33,6 +33,7 @@ tasks {
                 "https://publish.djl.ai/fasttext-${libs.versions.fasttext.get()}/jnilib/${libs.versions.djl.get()}"
             val files = mapOf(
                 "linux-x86_64" to "libjni_fasttext.so",
+                "linux-aarch64" to "libjni_fasttext.so",
                 "osx-aarch64" to "libjni_fasttext.dylib"
             )
             for ((key, value) in files) {

diff --git a/extensions/fasttext/build.sh b/extensions/fasttext/build.sh
@@ -45,6 +45,11 @@ if [[ $PLATFORM == 'darwin' ]]; then
     cp -f build/libjni_fasttext.dylib build/jnilib/osx-x86_64/
   fi
 elif [[ $PLATFORM == 'linux' ]]; then
-  mkdir -p build/jnilib/linux-x86_64
-  cp -f build/libjni_fasttext.so build/jnilib/linux-x86_64/
+  if [[ $(uname -m) == 'aarch64' ]]; then
+    mkdir -p build/jnilib/linux-aarch64
+    cp -f build/libjni_fasttext.so build/jnilib/linux-aarch64/
+  else
+    mkdir -p build/jnilib/linux-x86_64
+    cp -f build/libjni_fasttext.so build/jnilib/linux-x86_64/
+  fi
 fi