name: Native S3 Sentencepiece

on:
  workflow_dispatch:

permissions:
  id-token: write
  contents: read

jobs:
  build-sentencepiece-jni-linux:
    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
    runs-on: ubuntu-latest
    container:
      image: amazonlinux:2
      env:
        JAVA_HOME: /usr/lib/jvm/java-17-amazon-corretto
    steps:
      - name: Install Environment
        run: |
          yum -y update
          yum -y groupinstall "Development Tools"
          yum -y install patch cmake3 jq
          yum -y install java-17-amazon-corretto-devel
          ln -sf /usr/bin/cmake3 /usr/bin/cmake
          pip3 install awscli --upgrade
      - uses: taiki-e/checkout-action@v1
      - name: Release JNI prep
        run: |
          ./gradlew :extensions:sentencepiece:compileJNI
          ./gradlew -Pjni :extensions:sentencepiece: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-fasttext-jni-linux" \
            --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: |
          SENTENCEPIECE_VERSION="$(awk -F '=' '/sentencepiece/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
          aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/*"

  build-sentencepiece-jni-windows:
    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
    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 JNI prep
        run: |
          ./gradlew :extensions:sentencepiece:compileJNI
          ./gradlew -Pjni :extensions:sentencepiece:test
      - 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
        shell: bash
        run: |
          SENTENCEPIECE_VERSION="$(awk -F '=' '/sentencepiece/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
          aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/*"

  build-sentencepiece-jni-arm64-osx:
    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
    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-rs/toolchain@v1
        with:
          toolchain: stable
      - 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
        run: |
          ./gradlew :extensions:sentencepiece:compileJNI
          ./gradlew -Pjni :extensions:sentencepiece:test
      - 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: |
          SENTENCEPIECE_VERSION="$(awk -F '=' '/sentencepiece/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
          aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/*"

  build-sentencepiece-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:sentencepiece:compileJNI
          ./gradlew -Pjni :extensions:sentencepiece: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: |
          SENTENCEPIECE_VERSION="$(awk -F '=' '/sentencepiece/ {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)"
          aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_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-sentencepiece-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
