Port PR CI workflow from rocm-main (#312)

Author: charleshofer

.github/workflows/rocm-ci.yml

@@ -0,0 +1,72 @@
+name: ROCm GPU CI
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the rocm-main branch
+  push:
+    branches:
+      - rocm-main
+      - 'rocm-jaxlib-v*'
+  pull_request:
+    branches:
+      - rocm-main
+      - 'rocm-jaxlib-v*'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-jax-in-docker:
+    runs-on:  mi-250
+    strategy:
+      matrix:
+        python: ["3.10", "3.11", "3.12"]
+        rocm: [ "6.3.3"]
+    env:
+      BASE_IMAGE: "ubuntu:22.04"
+      TEST_IMAGE: ubuntu-jax-${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
+      PYTHON_VERSION: ${{ matrix.python }}
+      ROCM_VERSION: ${{ matrix.rocm }}
+      WORKSPACE_DIR: workdir_${{ github.run_id }}_${{ github.run_number }}_${{ github.run_attempt }}
+    steps:
+      - name: Clean up old runs
+        run: |
+          ls
+          # Make sure that we own all of the files so that we have permissions to delete them
+          docker run -v "./:/jax" ubuntu /bin/bash -c "chown -R $UID /jax/workdir_* || true"
+          # Remove any old work directories from this machine
+          rm -rf workdir_*
+          ls
+      - name: Print system info
+        run: |
+          whoami
+          printenv
+          df -h
+          rocm-smi
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+        with:
+          path: ${{ env.WORKSPACE_DIR }}
+      - name: Build JAX
+        run: |
+          pushd $WORKSPACE_DIR
+          python3 build/rocm/ci_build \
+            --rocm-version $ROCM_VERSION \
+            --base-docker $BASE_IMAGE \
+            --python-versions $PYTHON_VERSION \
+            --compiler=clang \
+            dist_docker \
+            --image-tag $TEST_IMAGE
+      - name: Archive jax wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
+          path: ./dist/*.whl
+      - name: Run tests
+        env:
+          GPU_COUNT: "8"
+          GFX: "gfx90a"
+        run: |
+          cd $WORKSPACE_DIR
+          python3 build/rocm/ci_build test $TEST_IMAGE --test-cmd "pytest tests/core_test.py"
+

build/rocm/ci_build

@@ -193,7 +193,7 @@ def dist_docker(
     subprocess.check_call(cmd)
 
 
-def test(image_name):
+def test(image_name, test_cmd):
     """Run unit tests like CI would inside a JAX image."""
 
     gpu_args = [
@@ -211,10 +211,12 @@ def test(image_name):
     cmd = [
         "docker",
         "run",
-        "-it",
         "--rm",
     ]
 
+    if os.isatty(sys.stdout.fileno()):
+        cmd.append("-it")
+
     # NOTE(mrodden): we need jax source dir for the unit test code only,
     # JAX and jaxlib are already installed from wheels
     mounts = [
@@ -225,7 +227,7 @@ def test(image_name):
     cmd.extend(mounts)
     cmd.extend(gpu_args)
 
-    container_cmd = "cd /jax && ./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh"
+    container_cmd = "cd /jax && " + test_cmd 
     cmd.append(image_name)
     cmd.extend(
         [
@@ -288,6 +290,7 @@ def parse_args():
 
     testp = subp.add_parser("test")
     testp.add_argument("image_name")
+    testp.add_argument("--test-cmd", default="./build/rocm/build_rocm.sh && ./build/rocm/run_single_gpu.py -c && ./build/rocm/run_multi_gpu.sh")
 
     ddp = subp.add_parser("dist_docker")
     ddp.add_argument("--dockerfile", default="build/rocm/Dockerfile.ms")
@@ -311,7 +314,7 @@ def main():
         )
 
     elif args.action == "test":
-        test(args.image_name)
+        test(args.image_name, args.test_cmd)
 
     elif args.action == "dist_docker":
         dist_wheels(