CI: Migrate MoRI to ROCm 10.1

.claude/skills/deploy-mori/SKILL.md

@@ -52,7 +52,7 @@ done
 ## Step 1: Start the Docker container
 
 Use the container name from the user's args if provided; ask if not.
-Default image: `rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0`
+Default image: `rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0`
 (use unless the user specifies another) — same base the CI workflows build on.
 
 ```bash
@@ -427,35 +427,44 @@ mst status -v
 
 ## Step 4: Install MORI
 
-### 4.0 — ROCm 7.14 dev toolchain (rocm7.14 base images only)
+### 4.0 — ROCm dev toolchain (rocm10.x / rocm7.14 base images only)
 
 **Skip on pre-7.14 images** — they already ship a full `/opt/rocm` and must not
 get a second ROCm stacked on top.
 
-The rocm7.14 pytorch image has **no `/opt/rocm` at all**: ROCm comes from the pip
-`rocm-sdk-core` wheel, which ships runtime `.so` files and headers but no CMake
-package configs. Without this step the build dies at configure time with
+The rocm10.x and rocm7.14 pytorch images have **no `/opt/rocm` at all**: ROCm
+comes from the pip `rocm-sdk-core` wheel, which ships runtime `.so` files and
+headers but no CMake package configs. Without this step the build dies at
+configure time with
 `does not contain the HIP runtime CMake package ... hip-lang-config.cmake`.
 
+Install the release in the image's tag — mixing them makes every `import mori`
+warn that it was built against another ROCm. For ROCm 10.1 (an Ubuntu 22.04
+image takes `ubuntu2204` in the repo URL):
+
 ```bash
 sudo docker exec $CONTAINER_NAME bash -c '
 set -e
 [ -d /opt/rocm/bin ] && { echo "/opt/rocm already present — skipping"; exit 0; }
 apt-get install -y --no-install-recommends wget gnupg ca-certificates
 mkdir -p --mode=0755 /etc/apt/keyrings
-wget -qO - https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg \
+wget -qO - https://stable.repo.amd.com/rocm/gpg/packages.gpg \
   | gpg --dearmor > /etc/apt/keyrings/amdrocm.gpg
-echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main" \
+echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://stable.repo.amd.com/rocm/core/packages/ubuntu2404 stable main" \
   > /etc/apt/sources.list.d/amdrocm.list
 apt-get update
-apt-get install -y --no-install-recommends amdrocm-core-dev7.14-gfx942 amdrocm-core-dev7.14-gfx950
-hipconfig --version
+apt-get install -y --no-install-recommends amdrocm-core-dev10.1-gfx942 amdrocm-core-dev10.1-gfx950
+cat /opt/rocm/core/.info/version
 '
 ```
 
+A rocm7.14 image takes `amdrocm-core-dev7.14-<arch>` from the
+`https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404` repo instead, signed
+by `https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg`.
+
 No environment setup is needed for the rest of the session. CMake resolves ROCm
-as `-DROCM_PATH` → `$ROCM_PATH` → `$HIP_PATH` → `/opt/rocm` (where both the
-7.2.4 and 7.14 stacks live) and pins the compiler to an absolute
+as `-DROCM_PATH` → `$ROCM_PATH` → `$HIP_PATH` → `/opt/rocm` (where the 7.2.4,
+7.14 and 10.x stacks all live) and pins the compiler to an absolute
 `<rocm>/llvm/bin/amdclang++`; mori's libraries carry
 `RUNPATH "$ORIGIN:<rocm>/lib"`, and the JIT builds absolute `<rocm>/bin/...`
 paths off the same default. Set `ROCM_PATH` only for ROCm somewhere else (e.g. a

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_ci_${{ github.run_id }}
   CT: docker
 
@@ -83,6 +83,12 @@ jobs:
           $CT exec $CONTAINER \
             git config --global --add safe.directory $GITHUB_WORKSPACE
 
+      - name: Report ROCm version
+        run: |
+          $CT exec $CONTAINER bash -c "
+            ls -d /opt/rocm/core-* && /opt/rocm/bin/hipconfig --version && echo
+          "
+
       - name: Build and install mori
         run: $CT exec $CONTAINER pip install .
 
@@ -454,19 +460,17 @@ jobs:
                 chown -R $(id -u):$(id -g) $GITHUB_WORKSPACE 2>/dev/null || true
           fi
 
-  # Gate for "does mori still build and run on the next ROCm major". Kept to the
-  # single EP intranode test on one runner: the rest of the intranode suite is
-  # ROCm-version agnostic, and the ROCm 10 base image is a separate ~30GB pull.
-  intranode-test-rocm10:
-    name: mori intranode test (ROCm 10, MI355X_AINIC)
+  # Gate for "does mori still build and run on ROCm 7.14"; every other job runs
+  # on 10.1. Kept to the EP intranode test on one runner: the rest of the
+  # intranode suite is ROCm-version agnostic.
+  intranode-test-rocm714:
+    name: mori intranode test (ROCm 7.14, MI355X_AINIC)
     needs: build
     runs-on: [self-hosted, MI355X-AINIC-PIT-G17]
     env:
-      # py3.11 rather than the 7.14 lane's py3.12 — this is the tag the ROCm 10
-      # build was validated on.
-      BASE_IMAGE: rocm/pytorch:rocm10.0_ubuntu24.04_py3.11_pytorch_release_2.13.0
-      IMAGE: rocm/mori:ci-rocm10
-      CONTAINER: mori_ci_rocm10_${{ github.run_id }}
+      BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+      IMAGE: rocm/mori:ci-rocm714
+      CONTAINER: mori_ci_rocm714_${{ github.run_id }}
       MORI_RDMA_DEVICES: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
       MORI_RDMA_SL: 3
       MORI_RDMA_TC: 104
@@ -491,14 +495,14 @@ jobs:
           # main; otherwise it would reuse the merge commit of the first attempt.
           ref: ${{ github.event_name == 'pull_request' && github.run_attempt > 1 && format('refs/pull/{0}/merge', github.event.pull_request.number) || '' }}
 
-      # ROCm 10 ships in a different apt repo than 7.x, hence the ROCM_APT_* args.
+      # ROCm 7.14 ships in a different apt repo than 10.x, hence the ROCM_APT_* args.
       - name: Build CI image
         run: |
           $CT build --network=host \
             --build-arg BASE_IMAGE=$BASE_IMAGE \
-            --build-arg ROCM_APT_VERSION=10.0 \
-            --build-arg ROCM_APT_GPG=https://stable.repo.amd.com/rocm/gpg/packages.gpg \
-            --build-arg ROCM_APT_LIST="https://stable.repo.amd.com/rocm/core/packages/ubuntu2404 stable main" \
+            --build-arg ROCM_APT_VERSION=7.14 \
+            --build-arg ROCM_APT_GPG=https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg \
+            --build-arg ROCM_APT_LIST="https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main" \
             -t $IMAGE -f docker/Dockerfile.dev .
 
       - name: Start container
@@ -522,10 +526,12 @@ jobs:
             ls -d /opt/rocm/core-* && /opt/rocm/bin/hipconfig --version && echo
           "
 
+      # The same build as intranode-test's, so a ROCm 10-only API anywhere in it
+      # fails to compile here.
       - name: Install mori
         run: |
           $CT exec $CONTAINER bash -c \
-            "cd $GITHUB_WORKSPACE && pip install ."
+            "cd $GITHUB_WORKSPACE && BUILD_BENCHMARK=ON BUILD_CCO_SDMA=ON BUILD_EXAMPLES=ON pip install ."
 
       - name: Verify installation
         run: $CT exec $CONTAINER bash -c "cd $GITHUB_WORKSPACE && python tools/verify_install.py"
@@ -556,20 +562,19 @@ jobs:
           - platform: MI355X_AINIC
             runner: [self-hosted, MI355X-AINIC-PIT-G17]
             rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
-            gpu_arch: gfx950
             socket_ifname: eno0
           # Runner 325_bnxt_108 is decommissioned.
           # - platform: MI325X_BNXT
           #   runner: [self-hosted, 325_bnxt_108]
           #   rdma_devices: bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re7,bnxt_re8
-          #   gpu_arch: gfx942
           #   socket_ifname: enp49s0f1np1
     env:
       # jax.distributed rendezvouses over localhost, but mori.jax.shmem_init_attr
       # goes through MORI's own bootstrap, which needs the same eno0 pin as
       # intranode-test.
       MORI_SOCKET_IFNAME: ${{ matrix.socket_ifname }}
       NCCL_SOCKET_IFNAME: ${{ matrix.socket_ifname }}
+      JAX_BASE_IMAGE: rocm/jax:rocm10.1.0_ubuntu24.04_py3.12_jax_release_0.10.2
     timeout-minutes: 90
     steps:
       - name: Fix workspace permissions
@@ -588,15 +593,18 @@ jobs:
           # main; otherwise it would reuse the merge commit of the first attempt.
           ref: ${{ github.event_name == 'pull_request' && github.run_attempt > 1 && format('refs/pull/{0}/merge', github.event.pull_request.number) || '' }}
 
+      # The JAX image has no ROCm dev stack either, so it takes the same
+      # Dockerfile as the torch images.
       - name: Build JAX CI image
         run: |
-          BUILD_ARGS="--network=host --build-arg GPU_TARGET=${{ matrix.gpu_arch }}"
+          BUILD_ARGS="--network=host"
           if [ "$CT" = "podman" ]; then
             BUILD_ARGS="$BUILD_ARGS --isolation=chroot"
           fi
           $CT build $BUILD_ARGS \
+            --build-arg BASE_IMAGE=$JAX_BASE_IMAGE \
             -t rocm/mori:jax-ci \
-            -f docker/Dockerfile.jax091-rocm711 .
+            -f docker/Dockerfile.dev .
 
       - name: Start container
         run: |
@@ -616,10 +624,11 @@ jobs:
           # Resolve the XLA FFI headers from the system jax *outside* pip's
           # isolated build env (build isolation hides system site-packages, so
           # CMake can't `import jax` during configure). Pass it via env var.
+          # The JAX image ships no pytest.
           $CT exec ${CONTAINER}_jax bash -c \
             "cd $GITHUB_WORKSPACE && \
              export XLA_FFI_INCLUDE_DIR=\$(python3 -c 'from jax import ffi; print(ffi.include_dir())') && \
-             BUILD_UMBP=OFF BUILD_XLA_FFI_OPS=ON pip install --break-system-packages ."
+             BUILD_UMBP=OFF BUILD_XLA_FFI_OPS=ON pip install . && pip install pytest"
 
       - name: MORI-EP JAX test
         run: |

.github/workflows/ci_cco.yml

@@ -12,7 +12,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci-cco
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_cco_ci_${{ github.run_id }}
   CT: podman  # MI355X-AINIC runner is podman rootless (no docker daemon)
 

.github/workflows/codeql.yml

@@ -32,7 +32,7 @@ concurrency:
 
 env:
   IMAGE: rocm/mori:ci
-  BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+  BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
   CONTAINER: mori_codeql_${{ github.run_id }}
   CT: docker
   # falls back to cpp,python on push events (where inputs.languages is empty)

.github/workflows/nightly.yml

@@ -31,16 +31,13 @@ jobs:
           # Built in docker/Dockerfile.manylinux, so a wheel needs nothing newer
           # than glibc 2.28 instead of the glibc/libstdc++ of the Ubuntu image that
           # ships its Python. Each lane builds against the ROCm its test image
-          # below runs: 7.2.4 for py3.10, 7.14 for py3.12.
+          # below runs: 10.1 for both, the Dockerfile's default (its header lists
+          # the args for other ROCm releases).
           - python: "3.10"
-            rocm: "7.2.4"
-            rocm_build_args: >-
-              --build-arg ROCM_VERSION=7.2.4
-              --build-arg ROCM_REPO=https://repo.radeon.com/rocm/rhel8/7.2.4/main
-              --build-arg ROCM_GPG=https://repo.radeon.com/rocm/rocm.gpg.key
-              --build-arg "ROCM_PACKAGES=rocm-hip-runtime-devel rocm-smi-lib"
+            rocm: "10.1"
+            rocm_build_args: ""
           - python: "3.12"
-            rocm: "7.14"
+            rocm: "10.1"
             rocm_build_args: ""
     steps:
       # Must precede Checkout: later steps run as root inside the CI container
@@ -189,7 +186,7 @@ jobs:
           - platform: MI355X_AINIC
             python: "3.10"
             runner: [self-hosted, MI355X-AINIC-PIT-G17]
-            base_image: rocm/pytorch:rocm7.2.4_ubuntu22.04_py3.10_pytorch_release_2.8.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu22.04_py3.10_pytorch_release_2.12.0
             rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
             rdma_sl: 3
             rdma_tc: 104
@@ -199,7 +196,7 @@ jobs:
           - platform: MI355X_AINIC
             python: "3.12"
             runner: [self-hosted, MI355X-AINIC-PIT-G17]
-            base_image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
             rdma_devices: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7
             rdma_sl: 3
             rdma_tc: 104
@@ -213,7 +210,7 @@ jobs:
           - platform: MI300X_MLX
             python: "3.12"
             runner: [self-hosted, 300_mlx_p19_29]
-            base_image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
             rdma_devices: mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7
             socket_ifname: ens14np0
             timeout_minutes: 90
@@ -436,7 +433,7 @@ jobs:
           - platform: MI355X_AINIC
             python: "3.10"
             runner: [self-hosted, MI355X-AINIC]
-            base_image: rocm/pytorch:rocm7.2.4_ubuntu22.04_py3.10_pytorch_release_2.8.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu22.04_py3.10_pytorch_release_2.12.0
             node1_host: 10.2.80.20
             node2_host: 10.2.80.10
             node2_port: 22
@@ -450,7 +447,7 @@ jobs:
           - platform: MI355X_AINIC
             python: "3.12"
             runner: [self-hosted, MI355X-AINIC]
-            base_image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
             node1_host: 10.2.80.20
             node2_host: 10.2.80.10
             node2_port: 22
@@ -471,7 +468,7 @@ jobs:
           - platform: MI300X_MLX
             python: "3.12"
             runner: [self-hosted, 300_mlx_p19_29]
-            base_image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
             node1_host: 10.194.132.110
             node2_host: 10.194.132.71
             node2_port: 22
@@ -935,9 +932,9 @@ jobs:
       matrix:
         include:
           - python: "3.10"
-            base_image: rocm/pytorch:rocm7.2.4_ubuntu22.04_py3.10_pytorch_release_2.8.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu22.04_py3.10_pytorch_release_2.12.0
           - python: "3.12"
-            base_image: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+            base_image: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
     env:
       IMAGE: rocm/mori:ci-py${{ matrix.python }}
     steps:
@@ -1017,7 +1014,7 @@ jobs:
       MORI_RDMA_SL: 3
       MORI_RDMA_TC: 104
       IMAGE: rocm/mori:ci-flydsl-${{ matrix.flydsl }}
-      BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+      BASE_IMAGE: rocm/pytorch:rocm10.1.0_ubuntu24.04_py3.12_pytorch_release_2.12.0
       CONTAINER: mori_flydsl_compat_${{ github.run_id }}_${{ strategy.job-index }}
       CT: podman  # MI355X-AINIC runner is podman rootless (no docker daemon)
     timeout-minutes: 90