Run NCCL tests on the JAX-specific base container.

Author: olupton

.github/workflows/_ci.yaml

@@ -47,6 +47,13 @@ jobs:
       MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
     secrets: inherit
 
+  test-nccl:
+    needs: build-base
+    uses: ./.github/workflows/_test_nccl.yaml
+    with:
+      CONTAINER: ${{ needs.build-base.outputs.DOCKER_TAG }}
+    secrets: inherit
+
   build-jax:
     needs: build-base
     uses: ./.github/workflows/_build.yaml

.github/workflows/_test_nccl.yaml

@@ -3,12 +3,9 @@ name: ~run NCCL tests
 on:
   workflow_call:
     inputs:
-      # Note that cuda-dl-base installs the NCCL tests, while the vanilla nvidia/cuda
-      # images do not; when JAX-Toolbox moves to using cuda-dl-base this workflow ought
-      # to be modified to test one of the JAX-Toolbox containers.
       CONTAINER:
         type: string
-        description: CUDA image to use as base, e.g. nvcr.io/nvidia/cuda-dl-base:24.11-cuda12.6-devel-ubuntu24.04
+        description: Container to test; assumed to already contain {all_gather,all_reduce,broadcast,reduce_scatter}_perf_mpi in $PATH
         required: true
 
 permissions:

.github/workflows/nccl-k8s.yaml

@@ -1,25 +1,12 @@
 name: NCCL on Kubernetes
 on:
-  schedule:
-    - cron: '30 8 * * *'
-  pull_request:
-    types:
-      - opened
-      - reopened
-      - ready_for_review
-      - synchronize
-    paths-ignore:
-      - '**.md'
   workflow_dispatch:
     inputs:
-      # Note that cuda-dl-base installs the NCCL tests, while the vanilla nvidia/cuda
-      # images do not; when JAX-Toolbox moves to using cuda-dl-base this workflow ought
-      # to be modified to test one of the JAX-Toolbox containers.
       CONTAINER:
         type: string
-        description: Container to test, this is assumed to already contain the NCCL tests e.g. cuda-dl-base or derived
-        default: ''
-        required: false
+        default: ghcr.io/nvidia/jax:base
+        description: Container to test, this is assumed to already contain the NCCL tests
+        required: true
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -34,5 +21,5 @@ jobs:
   nccl-tests:
     uses: ./.github/workflows/_test_nccl.yaml
     with:
-      CONTAINER: ${{ inputs.CONTAINER || 'nvcr.io/nvidia/cuda-dl-base:25.02-cuda12.8-devel-ubuntu24.04' }}
+      CONTAINER: ${{ inputs.CONTAINER }}
     secrets: inherit