[ci] update docker build of temp images to properly build release images (#2739)

siddvenk · web-flow · commit 6a76ee125ac9 · 2025-02-18T14:04:16.000-08:00
diff --git a/.github/workflows/docker-nightly-publish.yml b/.github/workflows/docker-nightly-publish.yml
@@ -114,55 +114,48 @@ jobs:
         with:
           path: ~/.gradle/caches
           key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
-      - name: Extract DJL and DJL Serving versions from TOML
+      - name: Set env vars for container build and tag
         id: get-versions
         run: |
           DJL_VERSION=$(awk -F '=' '/djl / {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)
           SERVING_VERSION=$(awk -F '=' '/serving / {gsub(/ ?"/, "", $2); print $2}' gradle/libs.versions.toml)
           echo "DJL_VERSION=$DJL_VERSION" >> $GITHUB_ENV
-          echo "SERVING_VERSION=$SERVING_VERSION" >> $GITHUB_ENV
           echo "DJL_VERSION=$DJL_VERSION" >> $GITHUB_OUTPUT
-      - name: Build release candidate docker image
-        if: ${{ inputs.mode == 'release' }}
-        working-directory: serving/docker
+          echo "SERVING_VERSION=$SERVING_VERSION" >> $GITHUB_ENV
+          mode=${{ inputs.mode }}
+          if [[ "$mode" != "release" ]]; then
+            export NIGHTLY="-nightly"
+            echo "NIGHTLY=$NIGHTLY" >> $GITHUB_ENV
+            ./gradlew --refresh-dependencies :serving:dockerDeb -Psnapshot
+            export BUILD_VERSION_ARG_DJL="${DJL_VERSION}-SNAPSHOT"
+            export BUILD_VERSION_ARG_SERVING="${SERVING_VERSION}-SNAPSHOT"
+          else
+            rm -rf serving/docker/distributions
+            export BUILD_VERSION_ARG_DJL="${DJL_VERSION}-SNAPSHOT"
+            export BUILD_VERSION_ARG_SERVING="${SERVING_VERSION}-SNAPSHOT" 
+          fi
+          echo "BUILD_VERSION_ARG_DJL=$BUILD_VERSION_ARG_DJL" >> $GITHUB_ENV
+          echo "BUILD_VERSION_ARG_SERVING=$BUILD_VERSION_ARG_SERVING" >> $GITHUB_ENV
+      - name: Build docker image
         run: |
-          export BASE_RELEASE_VERSION="${{ env.SERVING_VERSION }}"
           export RELEASE_VERSION="${{ env.SERVING_VERSION }}-"
-          docker compose build --no-cache \
-          --build-arg djl_version=${{ env.DJL_VERSION }} \
-          --build-arg djl_serving_version=${{ env.SERVING_VERSION }} \
-          ${{ matrix.arch }}
-      - name: Build temp docker image
-        if: ${{ inputs.mode == '' || inputs.mode == 'temp' || inputs.mode == 'nightly' }}
-        run: |
-          ./gradlew --refresh-dependencies :serving:dockerDeb -Psnapshot
           cd serving/docker
-          export NIGHTLY="-nightly"
-          echo "NIGHTLY=$NIGHTLY" >> $GITHUB_ENV
           docker compose build --no-cache \
-          --build-arg djl_version=${{ env.DJL_VERSION }}-SNAPSHOT \
-          --build-arg djl_serving_version=${{ env.SERVING_VERSION }}-SNAPSHOT \
+          --build-arg djl_version=${{ env.BUILD_VERSION_ARG_DJL }} \
+          --build-arg djl_serving_version=${{ env.BUILD_VERSION_ARG_SERVING }} \
           ${{ matrix.arch }}
       - name: Tag and push temp image to ECR repo
         working-directory: serving/docker
         run: |
           ECR_REGION=$(echo "${{ env.AWS_ECR_REPO }}" | awk -F. '{print $4}')
           aws ecr get-login-password --region $ECR_REGION | docker login --username AWS --password-stdin ${{env.AWS_ECR_REPO}}
-          mode=${{ inputs.mode }}
-          if [[ "${{ inputs.mode }}" == "release" ]]; then
-            mode=${{ env.DJL_VERSION }}
-          fi
-          if [[ -z "${{ inputs.mode }}" ]]; then
-            mode="nightly"
-          fi
-          tempRunIdTag="${{ env.AWS_ECR_REPO }}:${{ matrix.arch }}-$mode-${GITHUB_RUN_ID}"
-          tempCommitTag="${{ env.AWS_ECR_REPO }}:${{ matrix.arch }}-$mode-${GITHUB_SHA}"
+          baseDockerTag="${{ env.AWS_ECR_REPO }}:${{ env.SERVING_VERSION}}-${{ matrix.arch }}${{ env.NIGHTLY }}"
+          tempRunIdTag="${baseDockerTag}-${GITHUB_RUN_ID}"
+          tempCommitTag="${baseDockerTag}-${GITHUB_SHA}"
 
-          docker tag ${{ env.DOCKER_HUB_REPO }}:${{ matrix.arch }}${{ env.NIGHTLY }} $tempRunIdTag
-          docker tag ${{ env.DOCKER_HUB_REPO }}:${{ matrix.arch }}${{ env.NIGHTLY }} $tempCommitTag
-          if [[ "$mode" == "nightly" ]]; then
-            docker tag ${{ env.DOCKER_HUB_REPO }}:${{ matrix.arch }}${{ env.NIGHTLY }} ${{ env.AWS_ECR_REPO }}:${{ matrix.arch }}-nightly
-          fi
+          docker tag ${{ env.DOCKER_HUB_REPO }}:${{ env.SERVING_VERSION}}-${{ matrix.arch }}${{ env.NIGHTLY }} $baseDockerTag 
+          docker tag ${{ env.DOCKER_HUB_REPO }}:${{ env.SERVING_VERSION}}-${{ matrix.arch }}${{ env.NIGHTLY }} $tempRunIdTag
+          docker tag ${{ env.DOCKER_HUB_REPO }}:${{ env.SERVING_VERSION}}-${{ matrix.arch }}${{ env.NIGHTLY }} $tempCommitTag
           time docker push --all-tags  ${{ env.AWS_ECR_REPO }}
 
   stop-runners:
diff --git a/.github/workflows/docker_publish.yml b/.github/workflows/docker_publish.yml
@@ -48,26 +48,50 @@ env:
   ECR_REPO_REGION: "us-east-1"
 
 jobs:
-  create-aarch64-runner:
+  create-runners:
     runs-on: [ self-hosted, scheduler ]
     steps:
-      - name: Create new Graviton instance
-        id: create_aarch64
+      - name: Create new CPU instance
+        id: create_cpu_1
         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-serving/actions/runners/registration-token \
           --fail \
           | jq '.token' | tr -d '"' )
-          ./start_instance.sh action_graviton $token djl-serving
+          ./start_instance.sh action_cpu $token djl-serving
+      - name: Create new CPU instance
+        id: create_cpu_2
+        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-serving/actions/runners/registration-token \
+          --fail \
+          | jq '.token' | tr -d '"' )
+          ./start_instance.sh action_cpu $token djl-serving
+      - name: Create new CPU instance
+        id: create_cpu_3
+        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-serving/actions/runners/registration-token \
+          --fail \
+          | jq '.token' | tr -d '"' )
+          ./start_instance.sh action_cpu $token djl-serving
     outputs:
-      aarch64_instance_id: ${{ steps.create_aarch64.outputs.action_graviton_instance_id }}
-
+      cpu_instance_id1: ${{ steps.create_cpu_1.outputs.action_cpu_instance_id }}
+      cpu_instance_id2: ${{ steps.create_cpu_2.outputs.action_cpu_instance_id }}
+      cpu_instance_id3: ${{ steps.create_cpu_3.outputs.action_cpu_instance_id }}
 
   docker-sync:
-    runs-on:  ${{ contains( '[tensorrt-llm]', matrix.arch) && needs.create-aarch64-runner.outputs.aarch64_instance_id || 'ubuntu-latest' }}
+    runs-on:
+      - self-hosted
+      - cpu
+      - RUN_ID-${{ github.run_id }}
+      - RUN_NUMBER-${{ github.run_number }}
+      - SHA-${{ github.sha }}
     timeout-minutes: 60
-    needs: [create-aarch64-runner]
+    needs: [create-runners]
     strategy:
       fail-fast: false
       matrix:
@@ -112,23 +136,21 @@ jobs:
         working-directory: serving/docker
         run: |
           ./scripts/push_image_from_ECR.sh $DJL_VERSION $AWS_STAGING_ECR_REPO ${{ inputs.mode }} ${{ matrix.arch }} ${{ inputs.commit_sha }}
-      - name: Retag image for release latest
-        if: ${{ inputs.mode == 'release' }}
-        working-directory: serving/docker
-        run: |
-          docker tag deepjavalibrary/djl-serving:${DJL_VERSION} deepjavalibrary/djl-serving:latest
-          docker push deepjavalibrary/djl-serving:latest
       - name: Clean docker env
         working-directory: serving/docker
         run: |
           yes | docker system prune -a --volumes
-  stop-aarch64-runner:
+  stop-runners:
     if: always()
     runs-on: [ self-hosted, scheduler ]
-    needs: [docker-sync, create-aarch64-runner]
+    needs: [docker-sync, create-runners]
     steps:
       - name: Stop all instances
         run: |
           cd /home/ubuntu/djl_benchmark_script/scripts
-          instance_id=${{ needs.create-aarch64-runner.outputs.aarch64_instance_id }}
+          instance_id=${{ needs.create-runners.outputs.cpu_instance_id1 }}
+          ./stop_instance.sh $instance_id
+          instance_id=${{ needs.create-runners.outputs.cpu_instance_id2 }}
+          ./stop_instance.sh $instance_id
+          instance_id=${{ needs.create-runners.outputs.cpu_instance_id3 }}
           ./stop_instance.sh $instance_id
diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml
@@ -5,25 +5,35 @@ on:
     inputs:
       djl-version:
         description: 'The released version of DJL.'
-        required: false
+        required: true
         default: ''
       tag-suffix:
         description: 'Run tests on the specific tags suffix i.e. arch-{suffix}'
         required: false
         type: string
         default: 'nightly'
+      image-repo:
+        description: 'The repository to fetch images from'
+        required: false
+        type: string
+        default: '185921645874.dkr.ecr.us-east-1.amazonaws.com/djl-ci-temp'
   workflow_call:
     inputs:
       djl-version:
         description: 'The released version of DJL.'
         required: false
         type: string
-        default: 'nightly'
+        default: ''
       tag-suffix:
         description: 'Run tests on the specific tags suffix i.e. arch-{suffix}'
         required: false
         type: string
-        default: ''
+        default: 'nightly'
+      image-repo:
+        description: 'The repository to fetch images from'
+        required: false
+        type: string
+        default: '185921645874.dkr.ecr.us-east-1.amazonaws.com/djl-ci-temp'
     outputs:
       failure_cpu:
         value: ${{ jobs.test.outputs.failure_cpu || '0' }}
@@ -42,9 +52,6 @@ permissions:
   id-token: write
   contents: read
 
-env:
-  AWS_ECR_REPO: "185921645874.dkr.ecr.us-east-1.amazonaws.com/djl-ci-temp"
-
 jobs:
   create-runners:
     runs-on: [self-hosted, scheduler]
@@ -244,7 +251,6 @@ jobs:
           chmod +x awscurl
           mkdir outputs
       - name: Configure AWS Credentials
-        if: matrix.test.instance == 'ubuntu-latest'
         uses: aws-actions/configure-aws-credentials@v4
         with:
           role-to-assume: arn:aws:iam::185921645874:role/github-actions-djl-serving
@@ -253,11 +259,13 @@ jobs:
         working-directory: tests/integration
         env:
           TEST_DJL_VERSION: ${{ inputs.djl-version }}
-          OVERRIDE_IMAGE_TAG_SUFFIX: ${{ inputs.tag-suffix }}
-          IMAGE_REPO: ${{ env.AWS_ECR_REPO }}
+          IMAGE_TAG_SUFFIX: ${{ inputs.tag-suffix }}
+          IMAGE_REPO: ${{ inputs.image-repo }}
         run: |
-          ECR_REGION=$(echo "${{ env.AWS_ECR_REPO }}" | awk -F. '{print $4}')
-          aws ecr get-login-password --region $ECR_REGION | docker login --username AWS --password-stdin ${{env.AWS_ECR_REPO}}
+          ECR_REGION=$(echo "$IMAGE_REPO" | awk -F. '{print $4}')
+          if [[ -n "$ECR_REGION" ]]; then
+            aws ecr get-login-password --region $ECR_REGION | docker login --username AWS --password-stdin "$IMAGE_REPO"
+          fi
           python -m pytest -s -k ${{ matrix.test.test }} tests.py
       - name: Cleanup
         working-directory: tests/integration
@@ -314,16 +322,16 @@ jobs:
           aws-region: us-east-1
       - name: Download models and dockers
         run: |
-          if [ "${{ github.event.inputs.djl-version }}" == "temp"  ]; then
-            DOCKER_IMAGE_URI="${{ env.AWS_ECR_REPO }}:pytorch-inf2-temp-${GITHUB_SHA}"
-          elif [ -n "${{ inputs.tag-suffix }}" ]; then
-            DOCKER_IMAGE_URI="${{ env.AWS_ECR_REPO }}:pytorch-inf2-${{ inputs.tag-suffix }}"
+          if [ -n "${{ inputs.tag-suffix }}" ]; then
+            DOCKER_IMAGE_URI="${{ inputs.image-repo }}:${{ inputs.djl-version }}-pytorch-inf2-${{ inputs.tag-suffix }}"
           else
-            DOCKER_IMAGE_URI="${{ env.AWS_ECR_REPO }}:pytorch-inf2-nightly"
+            DOCKER_IMAGE_URI="${{ inputs.image-repo }}:${{ inputs.djl-version }}-pytorch-inf2"
           fi
           echo "DOCKER_IMAGE_URI=$DOCKER_IMAGE_URI" >>$GITHUB_ENV
-          ECR_REGION=$(echo "${{ env.AWS_ECR_REPO }}" | awk -F. '{print $4}')
-          aws ecr get-login-password --region $ECR_REGION | docker login --username AWS --password-stdin ${{env.AWS_ECR_REPO}}
+          ECR_REGION=$(echo "${{ inputs.image-repo }}" | awk -F. '{print $4}')
+          if [[ -n "$ECR_REGION" ]]; then 
+            aws ecr get-login-password --region $ECR_REGION | docker login --username AWS --password-stdin ${{ inputs.image-repo }}
+          fi
           echo $DOCKER_IMAGE_URI
           docker pull $DOCKER_IMAGE_URI
       - name: Run djl_python unit/integration tests on container
diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -32,29 +32,13 @@ jobs:
         secrets: inherit
         with:
             mode: ${{ inputs.mode || 'nightly' }}
-    get_image_tag_suffix:
-        outputs:
-            test_image_tag_suffix: ${{ steps.get_image_tag_suffix.test_image_tag_suffix }}
-        needs: [build]
-        runs-on: ubuntu-latest
-        steps:
-            - name: get_image_tag_suffix
-              id: get_image_tag_suffix
-              run: |
-                if [[ "${{ inputs.mode || 'nightly' }}" == "nightly" ]]; then
-                    echo "test_image_tag_suffix=nightly" >> $GITHUB_OUTPUT
-                elif [[ "${{ inputs.mode }}" == "release" ]]; then
-                    echo "test_image_tag_suffix=${{ needs.build.outputs.djl_version }}-${GITHUB_RUN_ID}" >> $GITHUB_OUTPUT
-                else
-                    echo "Invalid mode specified"
-                    exit 1
-                fi
     integration-test:
-        needs: [get_image_tag_suffix]
+        needs: [build]
         uses: ./.github/workflows/integration.yml
         secrets: inherit
         with:
-            tag-suffix: ${{ needs.get_image_tag_suffix.outputs.test_image_tag_suffix }}
+            djl-version: ${{ needs.build.outputs.djl_version }}
+            tag-suffix: ${{ inputs.mode == 'nightly' && format('{0}-{1}', 'nightly', github.sha) || github.sha }}
     determine_images_to_publish:
         if: always()
         needs: [ integration-test ]
diff --git a/serving/docker/docker-compose.yml b/serving/docker/docker-compose.yml
@@ -5,7 +5,7 @@ services:
       context: .
       target: base
       dockerfile: Dockerfile
-    image: "deepjavalibrary/djl-serving:${BASE_RELEASE_VERSION:-cpu}${NIGHTLY}"
+    image: "deepjavalibrary/djl-serving:${RELEASE_VERSION}cpu${NIGHTLY}"
   cpu-full:
     build:
       context: .
diff --git a/serving/docker/partition/partition.py b/serving/docker/partition/partition.py
@@ -199,11 +199,9 @@ def run_partition(self) -> str:
             self.cleanup()
             return ''.join(partition_stdout)
         else:
-            logging.error(
-                f"Partitioning was not successful.")
+            logging.error(f"Partitioning was not successful.")
             raise Exception(
-                f"Partitioning exited with return code: {proc.returncode}."
-            )
+                f"Partitioning exited with return code: {proc.returncode}.")
 
     def run_quantization(self):
         quant_method = self.properties['option.quantize']
diff --git a/serving/docker/scripts/push_image_from_ECR.sh b/serving/docker/scripts/push_image_from_ECR.sh
@@ -24,17 +24,22 @@ from_repo=$AWS_TMP_ECR_REPO
 
 set -x
 
-if [[ "$mode" == "release" ]]; then
-  if [[ "$image" == "cpu" ]]; then
-    tag=$version
-  else
-    tag="$version-$image"
-  fi
-fi
+base_tag="$version-$image"
 
 if [[ "$mode" == "nightly" ]]; then
-  tag="$image-nightly"
+  base_tag="$base_tag-nightly"
+fi
+docker pull "$from_repo:$base_tag-$commit_sha"
+docker tag "$from_repo:$base_tag-$commit_sha" "$to_repo:$base_tag"
+docker push "$to_repo:$base_tag"
+if [[ "$image" == "cpu" ]]; then
+  if [[ "$mode" == "release" ]]; then
+    docker tag "$from_repo:$base_tag-$commit_sha" "$to_repo:$version"
+    docker push "$to_repo:$version"
+    docker tag "$from_repo:$base_tag-$commit_sha" "$to_repo:latest"
+    docker push "$to_repo:latest"
+  elif [[ "$mode" == "nightly" ]]; then
+    docker tag "$from_repo:$base_tag-$commit_sha" "$to_repo:$version-nightly"
+    docker push "$to_repo:$version-nightly"
+  fi
 fi
-docker pull $from_repo:$image-$mode-$commit_sha
-docker tag $from_repo:$image-$mode-$commit_sha $to_repo:$tag
-docker push $to_repo:$tag
diff --git a/tests/integration/tests.py b/tests/integration/tests.py
