Updating docker cuda and github branch (#3233)

* cuda update and branch fix * adding print * testing cuda * testing branch --------- Co-authored-by: Ankith Gunapal <[email protected]>
pytorch · Jul 9, 2024 · 1125bb1 · 1125bb1
1 parent a1c8eb2
commit 1125bb1
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/.github/workflows/docker-ci.yaml b/.github/workflows/docker-ci.yaml
@@ -24,12 +24,22 @@ jobs:
         working-directory: docker
         run: ./test_build_image_tagging.sh ${{ matrix.python-version }}
 
+      - name: Determine the branch name
+        id: branch-name
+        run: |
+          if [[ "${GITHUB_REF}" == refs/heads/* ]]; then
+            GITHUB_BRANCH=${GITHUB_REF#refs/heads/}
+          elif [[ "${GITHUB_REF}" == refs/pull/*/merge ]]; then
+            GITHUB_BRANCH=${GITHUB_HEAD_REF}
+          fi
+          echo "GITHUB_BRANCH=${GITHUB_BRANCH}" >> $GITHUB_OUTPUT
+
       - name: Build Image for container test
         id: image_build
         working-directory: docker
         run: |
           IMAGE_TAG=test-image-${{ matrix.python-version }}
-          ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b ${GITHUB_HEAD_REF} -s
+          ./build_image.sh -py "${{ matrix.python-version }}" -t "${IMAGE_TAG}" -b ${{ steps.branch-name.outputs.GITHUB_BRANCH }} -s
           echo "IMAGE_TAG=${IMAGE_TAG}" >> $GITHUB_OUTPUT
 
       - name: Container Healthcheck

diff --git a/docker/build_image.sh b/docker/build_image.sh
@@ -54,8 +54,8 @@ do
         -g|--gpu)
           MACHINE=gpu
           DOCKER_TAG="pytorch/torchserve:latest-gpu"
-          BASE_IMAGE="nvidia/cuda:11.8.0-base-ubuntu20.04"
-          CUDA_VERSION="cu117"
+          BASE_IMAGE="nvidia/cuda:12.1.1-base-ubuntu20.04"
+          CUDA_VERSION="cu121"
           shift
           ;;
         -bi|--baseimage)