Smoke test docker

Author: alekstheod

.github/workflows/ci.yml

@@ -97,6 +97,7 @@ jobs:
     strategy:
       matrix:
         rocm-version: ["7.1.1"]
+        python-version: ["3.11", "3.12", "3.13", "3.14"]
     steps:
       - uses: actions/checkout@v4
 
@@ -114,3 +115,24 @@ jobs:
 
       - name: Build JAX docker image
         run: bash jax_rocm_plugin/build/rocm/build_docker.sh ${{ matrix.rocm-version }}
+
+      - name: Smoke test docker image
+        env:
+          ROCM_VERSION: ${{ matrix.rocm-version }}
+        run: |
+          ROCM_VERSION_TAG="${ROCM_VERSION//./}"
+          JAX_TAG="jax-ubu24.rocm${ROCM_VERSION_TAG}"
+          eval "$(bash jax_rocm_plugin/build/rocm/get_commits.sh)"
+          docker run --rm \
+            --device=/dev/kfd \
+            --device=/dev/dri \
+            --group-add video \
+            --cap-add=SYS_PTRACE \
+            --security-opt seccomp=unconfined \
+            --shm-size 16G \
+            "${JAX_TAG}" \
+            bash -c "
+              curl -sL https://github.com/ROCm/jax/archive/${JAX_COMMIT}.tar.gz | tar xz -C /tmp && \
+              cd /tmp/jax-${JAX_COMMIT} && \
+              python${{ matrix.python-version }} -m pytest -x -v tests/nn_test.py tests/random_test.py
+            "

jax_rocm_plugin/build/rocm/build_docker.sh

@@ -10,10 +10,8 @@ REGISTRY="ghcr.io/rocm"
 BASE_TAG="jax-base-ubu24.rocm${ROCM_VERSION_TAG}"
 JAX_TAG="jax-ubu24.rocm${ROCM_VERSION_TAG}"
 
-# Extract commit hashes from workspace files
-XLA_COMMIT=$(grep -oP 'XLA_COMMIT = "\K[0-9a-f]+' jax_rocm_plugin/third_party/xla/workspace.bzl)
-JAX_COMMIT=$(grep -oP 'COMMIT = "\K[0-9a-f]+' jax_rocm_plugin/third_party/jax/workspace.bzl)
-ROCM_JAX_COMMIT=$(git rev-parse HEAD)
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+eval "$(bash "${SCRIPT_DIR}/get_commits.sh")"
 
 # Build base image if not available in the registry
 if ! docker manifest inspect "${REGISTRY}/${BASE_TAG}:latest" >/dev/null 2>&1; then

jax_rocm_plugin/build/rocm/get_commits.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+# Extracts commit hashes from workspace files and prints them as
+# shell-evaluable assignments. Usage:
+#   eval "$(bash jax_rocm_plugin/build/rocm/get_commits.sh)"
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+THIRD_PARTY="${SCRIPT_DIR}/../../third_party"
+
+XLA_COMMIT=$(grep -oP 'XLA_COMMIT = "\K[0-9a-f]+' "${THIRD_PARTY}/xla/workspace.bzl")
+JAX_COMMIT=$(grep -oP 'COMMIT = "\K[0-9a-f]+' "${THIRD_PARTY}/jax/workspace.bzl")
+ROCM_JAX_COMMIT=$(git rev-parse HEAD)
+
+echo "XLA_COMMIT=${XLA_COMMIT}"
+echo "JAX_COMMIT=${JAX_COMMIT}"
+echo "ROCM_JAX_COMMIT=${ROCM_JAX_COMMIT}"