[experimental] rocSHMEM as an allocation provider

.dockerignore

@@ -0,0 +1,15 @@
+# The Docker build context is the repo root (see .github/scripts/container_build.sh)
+# so that the Dockerfile can COPY in .github/scripts/install_rocshmem.sh. Nothing
+# else in the tree is needed at image build time, and .git alone is tens of MB.
+.git
+.github/workflows
+**/__pycache__
+**/*.pyc
+.pytest_cache
+*.egg-info
+build
+dist
+docs
+examples
+tests
+.claude

.github/scripts/container_build.sh

@@ -43,40 +43,71 @@ if [ "$CONTAINER_RUNTIME" = "apptainer" ]; then
     DEF_CHECKSUM=$(sha256sum "$DEF_FILE" | awk '{print $1}')
 
     # Create persistent Apptainer directory with checksum subdirectory
-    mkdir -p "${HOME}/iris-apptainer-images/${DEF_CHECKSUM}"
+    CACHE_DIR="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}"
+    mkdir -p "$CACHE_DIR"
 
-    # Define paths
-    IMAGE_PATH="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}/iris-dev.sif"
-    CHECKSUM_FILE="${HOME}/iris-apptainer-images/${DEF_CHECKSUM}/iris-dev.sif.checksum"
+    # Define paths. $HOME is shared across the runners, so every job in a run
+    # reads and writes this one directory.
+    IMAGE_PATH="$CACHE_DIR/iris-dev.sif"
+    CHECKSUM_FILE="$CACHE_DIR/iris-dev.sif.checksum"
+    LOCK_FILE="$CACHE_DIR/.build.lock"
 
-    # Check if image exists and has a valid checksum
-    REBUILD_NEEDED=true
-    if [ -f "$IMAGE_PATH" ] && [ -f "$CHECKSUM_FILE" ]; then
-        OLD_CHECKSUM=$(head -n1 "$CHECKSUM_FILE" 2>/dev/null)
+    image_is_current() {
+        [ -f "$IMAGE_PATH" ] && [ -f "$CHECKSUM_FILE" ] || return 1
+        local old
+        old=$(head -n1 "$CHECKSUM_FILE" 2>/dev/null)
         # Validate checksum format (64 hex characters for SHA256)
-        if [[ "$OLD_CHECKSUM" =~ ^[a-f0-9]{64}$ ]] && [ "$OLD_CHECKSUM" = "$DEF_CHECKSUM" ]; then
-            echo "[INFO] Def file unchanged (checksum: $DEF_CHECKSUM)"
-            echo "[INFO] Skipping rebuild, using existing image at $IMAGE_PATH"
-            REBUILD_NEEDED=false
-        else
-            echo "[INFO] Def file changed (old: ${OLD_CHECKSUM:-<invalid>}, new: $DEF_CHECKSUM)"
-            echo "[INFO] Rebuilding Apptainer image..."
-        fi
-    else
-        echo "[INFO] Image or checksum not found, building new Apptainer image..."
-    fi
+        [[ "$old" =~ ^[a-f0-9]{64}$ ]] && [ "$old" = "$DEF_CHECKSUM" ]
+    }
 
-    # Build the image if needed
-    if [ "$REBUILD_NEEDED" = true ]; then
-        if apptainer build --force "$IMAGE_PATH" "$DEF_FILE"; then
+    # Build to a private path and rename into place. Renaming is atomic and
+    # leaves the inode alone, so a job already executing the old image keeps
+    # running against it; `apptainer build --force` straight to IMAGE_PATH would
+    # truncate the file out from under it.
+    build_image() {
+        local tmp
+        tmp=$(mktemp -u "$CACHE_DIR/.iris-dev.XXXXXXXX.sif")
+        if apptainer build --force "$tmp" "$DEF_FILE"; then
+            mv -f "$tmp" "$IMAGE_PATH"
             # Store the checksum only if build succeeded
             echo "$DEF_CHECKSUM" > "$CHECKSUM_FILE"
             echo "[INFO] Built image: $IMAGE_PATH"
             echo "[INFO] Checksum saved: $DEF_CHECKSUM"
         else
+            rm -f "$tmp"
             echo "[ERROR] Apptainer build failed"
             exit 1
         fi
+    }
+
+    if image_is_current; then
+        echo "[INFO] Def file unchanged (checksum: $DEF_CHECKSUM)"
+        echo "[INFO] Skipping rebuild, using existing image at $IMAGE_PATH"
+    else
+        echo "[INFO] No current image for this def file"
+        # Serialize builders. Without this, jobs that start together all see no
+        # image and all build concurrently into the same path -- observed, with
+        # two runners building at once. The re-check inside the lock is the
+        # point: whoever waits usually finds the image already built and skips a
+        # redundant half-hour build.
+        if command -v flock > /dev/null 2>&1; then
+            exec 9> "$LOCK_FILE"
+            if ! flock -w 5400 9; then
+                echo "[ERROR] Timed out waiting for the image build lock"
+                exit 1
+            fi
+            if image_is_current; then
+                echo "[INFO] Another job built it while we waited; using $IMAGE_PATH"
+            else
+                echo "[INFO] Building new Apptainer image..."
+                build_image
+            fi
+            exec 9>&-
+        else
+            echo "[WARN] flock not available; building without a lock"
+            echo "[INFO] Building new Apptainer image..."
+            build_image
+        fi
     fi
 
 elif [ "$CONTAINER_RUNTIME" = "docker" ]; then
@@ -89,8 +120,11 @@ elif [ "$CONTAINER_RUNTIME" = "docker" ]; then
         echo "[INFO] Using existing Docker image: $IMAGE_NAME"
     else
         echo "[INFO] Docker image $IMAGE_NAME not found, building..."
-        DOCKER_DIR="$(dirname "$(realpath "$0")")/../../docker"
-        if docker build -t "$IMAGE_NAME" "$DOCKER_DIR"; then
+        REPO_ROOT="$(dirname "$(realpath "$0")")/../.."
+        # Build from the repo root, not docker/, so the Dockerfile can COPY in
+        # .github/scripts/install_rocshmem.sh -- the same installer the Apptainer
+        # def file pulls in via %files. A docker/-only context cannot see it.
+        if docker build -t "$IMAGE_NAME" -f "$REPO_ROOT/docker/Dockerfile" "$REPO_ROOT"; then
             echo "[INFO] Built Docker image: $IMAGE_NAME"
         else
             echo "[ERROR] Docker build failed"

.github/scripts/install_rocshmem.sh

@@ -0,0 +1,105 @@
+#!/bin/bash
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Install rocSHMEM and its Python bindings, for tests/unittests/test_rocshmem_provider.py.
+#
+# Without this the provider tests skip rather than run: rocshmem4py has no
+# prebuilt wheel anywhere (not PyPI, not the ROCm wheel indexes), and its
+# bindings do not build rocSHMEM themselves -- python/rocshmem does
+# find_package(rocshmem 3.5.0 CONFIG REQUIRED) with no FetchContent. So rocSHMEM
+# has to be built first, and then the bindings against it.
+#
+# Scope is deliberately IPC-only: that is what the provider uses, and it is all
+# a single-node CI runner can exercise. Upstream already defaults USE_IPC=ON and
+# USE_GDA=OFF, so no conduit flags are passed -- which also keeps MPI and the
+# RDMA provider libraries out of the picture entirely.
+set -euo pipefail
+
+ROCSHMEM_PREFIX="${ROCSHMEM_PREFIX:-/opt/rocshmem}"
+# Both arches on purpose. The runner label says mi325 (gfx942) but the hardware
+# reports gfx950, and rocSHMEM's device code must match the GPU it runs on: build
+# for the wrong one and hipModuleGetGlobal fails on rocSHMEM's device globals
+# ("Cannot create GlobalVar Obj for symbol: _ZN8rocshmem14logd_constantsE") and
+# init aborts. Semicolon-separated cmake list.
+ROCSHMEM_GPU_TARGETS="${ROCSHMEM_GPU_TARGETS:-gfx942;gfx950}"
+ROCSHMEM_REPO="${ROCSHMEM_REPO:-https://github.com/ROCm/rocm-systems.git}"
+ROCSHMEM_REF="${ROCSHMEM_REF:-develop}"
+ROCM_PATH="${ROCM_PATH:-/opt/rocm}"
+SRC="$(mktemp -d)"
+
+# TODO: drop this source build once the CI base images reach ROCm 7.14+, whose
+# artifacts ship rocSHMEM's static library and headers -- only the bindings below
+# would still be needed. The bases are ROCm 7.2.1 (apptainer) and 7.1 (docker)
+# today. rocshmem4py stays a source build either way until its TheRock packaging
+# lands.
+echo "==> rocSHMEM ${ROCSHMEM_REF} -> ${ROCSHMEM_PREFIX} (GPU_TARGETS=${ROCSHMEM_GPU_TARGETS})"
+
+# rocm-systems is a large monorepo and we need two directories out of it. Sparse
+# checkout keeps this from dominating image build time and size.
+git clone --depth 1 --branch "${ROCSHMEM_REF}" --filter=blob:none --sparse \
+    "${ROCSHMEM_REPO}" "${SRC}"
+git -C "${SRC}" sparse-checkout set projects/rocshmem python/rocshmem
+
+[ -f "${SRC}/projects/rocshmem/CMakeLists.txt" ] || {
+    echo "ERROR: projects/rocshmem missing after sparse checkout" >&2; exit 1; }
+
+# rocSHMEM's cmake/setup_project.cmake does a REQUIRED find_file for
+# .info/version under ROCM_PATH. Images that lack that file fail to configure
+# with "Could not find rocm_version_file", so use the documented escape hatch and
+# read the version from rocm_version.h, which is authoritative. hipconfig
+# --version is not used: it reports a build number that parses as the patch level.
+EXPLICIT_ROCM_VERSION="${EXPLICIT_ROCM_VERSION:-}"
+if [ -z "${EXPLICIT_ROCM_VERSION}" ] && [ ! -f "${ROCM_PATH}/.info/version" ]; then
+    _vh="$(find "${ROCM_PATH}" -name rocm_version.h 2>/dev/null | head -1)"
+    if [ -n "${_vh}" ]; then
+        EXPLICIT_ROCM_VERSION="$(awk '
+            /ROCM_VERSION_MAJOR/ {maj=$3} /ROCM_VERSION_MINOR/ {min=$3}
+            /ROCM_VERSION_PATCH/ {pat=$3}
+            END {if (maj != "") printf "%s.%s.%s", maj, min, pat}' "${_vh}")"
+        echo "==> ROCm ${EXPLICIT_ROCM_VERSION} detected from ${_vh}"
+    fi
+fi
+
+cmake -S "${SRC}/projects/rocshmem" -B "${SRC}/build" -G Ninja \
+    ${EXPLICIT_ROCM_VERSION:+-DEXPLICIT_ROCM_VERSION="${EXPLICIT_ROCM_VERSION}"} \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_PREFIX}" \
+    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
+    -DGPU_TARGETS="${ROCSHMEM_GPU_TARGETS}"
+cmake --build "${SRC}/build" --parallel "$(nproc)"
+cmake --install "${SRC}/build"
+
+# USE_IPC must be ON or rocshmem_ptr returns NULL for every peer and the provider
+# refuses to build a table. It is the upstream default, so this asserts rather
+# than sets it -- a silent flip would otherwise surface much later as skipped tests.
+if ! grep -qi "define ROCSHMEM_USE_IPC\|USE_IPC" \
+     "${ROCSHMEM_PREFIX}"/include/rocshmem/*.hpp 2>/dev/null; then
+    echo "==> note: could not confirm USE_IPC from headers; provider will report at run time"
+fi
+
+# A pip install from source, same as the one-liner in iris/experimental/README.md
+# but pointed at the checkout above instead of a git+ URL. That is deliberate: a
+# git+ URL makes pip clone the monorepo again, independently, at whatever HEAD
+# develop happens to be at -- so the bindings could be built from a different
+# revision than the core installed above. find_package would not catch it, since
+# it only compares versions, and the bindings statically link the core. One
+# checkout for both makes the skew impossible, and saves a second clone.
+#
+# CMAKE_PREFIX_PATH is the documented way to point the bindings at an install;
+# setup.py forwards it to CMake as a cache variable so a rocSHMEM shipped under
+# /opt/rocm cannot shadow it. ROCSHMEM_HOME is no longer required.
+echo "==> building rocshmem4py against ${ROCSHMEM_PREFIX}"
+CMAKE_PREFIX_PATH="${ROCSHMEM_PREFIX}" ROCM_PATH="${ROCM_PATH}" \
+    pip3 install --no-cache-dir "${SRC}/python/rocshmem"
+
+python3 -c "
+import rocshmem4py, importlib.metadata as md
+print('    rocshmem4py', md.version('rocshmem4py'), '->', rocshmem4py.__file__)
+for n in ('rocshmem_my_pe', 'rocshmem_n_pes', 'rocshmem_ptr'):
+    assert hasattr(rocshmem4py, n), f'missing {n}'
+print('    provider API present')
+"
+
+rm -rf "${SRC}"
+echo "==> rocSHMEM install complete"

apptainer/iris.def

@@ -4,8 +4,25 @@
 Bootstrap: docker
 From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0
 
+# The rocSHMEM installer is copied in rather than inlined, so it stays one
+# implementation shared with docker/Dockerfile. Caveat: container_build.sh caches
+# the image on the checksum of THIS def file only, so editing
+# install_rocshmem.sh alone reuses a stale image -- touch this file too.
+%files
+    .github/scripts/install_rocshmem.sh /opt/install_rocshmem.sh
+
 %post
     /bin/bash -c "
+    # Fail the build on the first error. Without this a step can fail silently
+    # and still produce an image that Apptainer reports as built -- which
+    # container_build.sh then caches by def-file checksum and every later job
+    # reuses. That happened: a transient 'fetch-pack: unexpected disconnect'
+    # during the Triton clone left /opt/triton absent, so the pinned checkout and
+    # its editable install never ran, 'import triton' silently fell back to the
+    # older pytorch-triton-rocm in site-packages, and every Triton test failed
+    # far downstream with an unrelated-looking error.
+    set -e
+
     # Set environment variables
     export TRITON_PATH=/opt/triton
     export ROCM_PATH=/opt/rocm
@@ -28,15 +45,27 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0
     pip3 install --upgrade pip && \
     pip3 install wheel jupyter
 
-    # Clone and install Triton
+    # Clone and install Triton. Retried: this is a large clone and a transient
+    # disconnect here used to poison the cached image rather than fail the build.
     cd /opt
-    git clone https://github.com/triton-lang/triton.git \$TRITON_PATH
+    for attempt in 1 2 3; do
+        rm -rf \$TRITON_PATH
+        git clone https://github.com/triton-lang/triton.git \$TRITON_PATH && break
+        echo \"triton clone failed (attempt \$attempt)\" >&2
+        [ \$attempt -lt 3 ] || exit 1
+        sleep 10
+    done
     cd \$TRITON_PATH
     git checkout f7c1d69401e9f09050451f30776562954b05e850
     pip3 install -e .
 
     # Make the venv writable by all
     chmod -R 777 /opt/venv
+
+    # rocSHMEM + rocshmem4py, so the provider tests run instead of skipping.
+    # See install_rocshmem.sh for the rationale, including the two arches.
+    ROCSHMEM_GPU_TARGETS='gfx942;gfx950' ROCSHMEM_PREFIX=/opt/rocshmem \
+        bash /opt/install_rocshmem.sh
     "
 
 %environment
@@ -52,6 +81,7 @@ From: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.14_pytorch_2.10.0
     export OMPI_ALLOW_RUN_AS_ROOT=1
     # Set required RCCL environment variable for ROCm
     export HSA_NO_SCRATCH_RECLAIM=1
+    export ROCSHMEM_PREFIX=/opt/rocshmem
 
 %runscript
     echo "Welcome to the ROCm-aware Apptainer image!"

docker/Dockerfile

@@ -47,6 +47,13 @@ RUN git checkout bcbcabdd0cff6539c7168299075992b2a23ff38e
 RUN pip3 install -e .
 ENV PYTHONPATH=$TRITON_PATH
 
+# rocSHMEM + rocshmem4py, so tests/unittests/test_rocshmem_provider.py runs
+# instead of skipping. See install_rocshmem.sh for why it is a source build and
+# why it builds for two arches.
+ENV ROCSHMEM_PREFIX=/opt/rocshmem
+COPY .github/scripts/install_rocshmem.sh /tmp/install_rocshmem.sh
+RUN ROCSHMEM_GPU_TARGETS='gfx942;gfx950' bash /tmp/install_rocshmem.sh && rm /tmp/install_rocshmem.sh
+
 # Set up workspace
 WORKDIR /workspace