override image name

.github/workflows/rocm-ci.yaml

@@ -11,7 +11,7 @@ jobs:
       - uses: actions/checkout@v4
       - name: Build Docker 
         env: 
-          BUILD_TAG: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
+          DOCKER_IMG_NAME: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
           # XLA_CLONE_DIR: 
         run: |
           ./build/rocm/ci_build.sh --rocm_version 6.1.3 \
@@ -23,12 +23,12 @@ jobs:
           path: ${{ github.workspace }}/wheelhouse/*.whl
 
 
-#       - name: Try rocm-smi
-#         run: |
-#           docker run \
-#             --group-add video \
-#             $RENDER_DEVICES \
-#             --cap-add=SYS_PTRACE \
-#             --security-opt seccomp=unconfined \
-#             --shm-size=64G \
-#             rocm/rocm-terminal rocm-smi
+      - name: Detect GPUs
+        run: |
+          docker run \
+            --group-add video \
+            $RENDER_DEVICES \
+            --cap-add=SYS_PTRACE \
+            --security-opt seccomp=unconfined \
+            --shm-size=64G \
+            rocm/rocm-terminal rocm-smi

build/rocm/ci_build.sh

@@ -117,8 +117,10 @@ WORKSPACE="${WORKSPACE:-$(upsearch WORKSPACE)}"
 BUILD_TAG="${BUILD_TAG:-jax}"
 
 # Determine the docker image name and BUILD_TAG.
-DOCKER_IMG_NAME="${BUILD_TAG}.${CONTAINER_TYPE}"
+DOCKER_IMG_NAME_DEFAULT="${BUILD_TAG}.${CONTAINER_TYPE}"
 
+# Let the env override the image name
+DOCKER_IMG_NAME="${DOCKER_IMG_NAME:-$DOCKER_IMG_NAME_DEFAULT}"
 # Under Jenkins matrix build, the build tag may contain characters such as
 # commas (,) and equal signs (=), which are not valid inside docker image names.
 DOCKER_IMG_NAME=$(echo "${DOCKER_IMG_NAME}" | sed -e 's/=/_/g' -e 's/,/-/g')