Merge pull request #62 from nebius/SCHED920

asteny · web-flow · commit ae481ae98a71 · 2026-02-10T11:42:37.000+01:00
Build slurm_training_diag images for different cuda versions on Github runners
diff --git a/.github/workflows/slurm_training_diag.yml b/.github/workflows/slurm_training_diag.yml
@@ -13,48 +13,150 @@ jobs:
         name: Compute shared DAYTIME
         run: echo "daytime=$(date -u +%Y%m%d%H%M%S)" >> "$GITHUB_OUTPUT"
 
-  build_and_push:
-    runs-on: [self-hosted, X64]
+  build_single_arch:
+    name: Build slurm=${{ matrix.slurm_version }} ${{ matrix.variant }} arch=${{ matrix.arch }}
     needs: prepare
+    runs-on: ${{ matrix.runner }}
 
     strategy:
       fail-fast: false
       matrix:
+        slurm_version: ["25.05.6", "25.11.2"]
+        variant: ["cuda12", "cuda13"]
+        arch: ["amd64", "arm64"]
         include:
-          - name: slurm_training_diag-cuda12
+          # runner/platform mapping for arch
+          - arch: amd64
+            platform: linux/amd64
+            runner: ubuntu-24.04
+          - arch: arm64
+            platform: linux/arm64
+            runner: ubuntu-24.04-arm
+
+          # variant mapping
+          - variant: cuda12
             cuda_version: "12.9.0"
             ubuntu_version: "ubuntu24.04"
             nccl_tests_version: "2.16.4"
-            slurm_version: 25.11.2
-          - name: slurm_training_diag-cuda13
+          - variant: cuda13
             cuda_version: "13.0.2"
             ubuntu_version: "ubuntu24.04"
             nccl_tests_version: "2.17.6"
-            slurm_version: 25.11.2
 
-    name: Build and push image for ${{ matrix.name }}
     env:
       DAYTIME: ${{ needs.prepare.outputs.daytime }}
+      # https://console.eu.nebius.com/project-e00managed-schedulers/registry/registry-e00hrt9na9xsn2px9f
+      IMAGE_BASE: cr.eu-north1.nebius.cloud/ml-containers/slurm_training_diag
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v6
+        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
 
-      - name: Add docker-credential-nebius to PATH
-        run: echo "/home/ml_containers/.nebius/bin" >> $GITHUB_PATH
+      - name: Install Nebius CLI
+        shell: bash
+        run: |
+          set -euo pipefail
+          curl -sSL https://storage.eu-north1.nebius.cloud/cli/install.sh | bash
+          echo "$HOME/.nebius/bin" >> "$GITHUB_PATH"
+
+      - name: Restore Nebius config
+        shell: bash
+        run: |
+          set -euo pipefail
+          mkdir -p "$HOME/.nebius"
+          echo "${{ secrets.NEBIUS_CONFIG_YAML_B64 }}" | base64 -d > "$HOME/.nebius/config.yaml"
+          chmod 600 "$HOME/.nebius/config.yaml"
 
-      - name: Get image version name
-        run: echo "cr.eu-north1.nebius.cloud/ml-containers/slurm_training_diag:slurm${{ matrix.slurm_version }}-cuda${{ matrix.cuda_version }}-${{ matrix.ubuntu_version }}-${{ env.DAYTIME }}"
+      - name: Configure Nebius docker credential-helper for auth
+        shell: bash
+        run: nebius registry configure-helper
 
-      - name: Build and push images
-        # https://console.eu.nebius.com/project-e00managed-schedulers/registry/registry-e00hrt9na9xsn2px9f
+      - name: Build and push single-arch image
+        shell: bash
         run: |
+          set -euo pipefail
+
+          IMAGE="${IMAGE_BASE}:slurm${{ matrix.slurm_version }}-cuda${{ matrix.cuda_version }}-${{ matrix.ubuntu_version }}-${DAYTIME}"
+          ARCH_TAG="${IMAGE}-${{ matrix.arch }}"
+
+          echo "Pushing: ${ARCH_TAG} (${{ matrix.platform }})"
+          echo "  SLURM_VERSION=${{ matrix.slurm_version }}"
+          echo "  CUDA_VERSION=${{ matrix.cuda_version }}"
+          echo "  NCCL_TESTS_VERSION=${{ matrix.nccl_tests_version }}"
+
           docker buildx build \
-            --platform linux/amd64,linux/arm64 \
-            --build-arg CUDA_VERSION=${{ matrix.cuda_version }} \
-            --build-arg NCCL_TESTS_VERSION=${{ matrix.nccl_tests_version }} \
-            -t cr.eu-north1.nebius.cloud/ml-containers/slurm_training_diag:slurm${{ matrix.slurm_version }}-cuda${{ matrix.cuda_version }}-${{ matrix.ubuntu_version }}-${{ env.DAYTIME }} \
+            --platform "${{ matrix.platform }}" \
+            --build-arg SLURM_VERSION="${{ matrix.slurm_version }}" \
+            --build-arg CUDA_VERSION="${{ matrix.cuda_version }}" \
+            --build-arg NCCL_TESTS_VERSION="${{ matrix.nccl_tests_version }}" \
+            -t "${ARCH_TAG}" \
             --target slurm_training_diag \
             --push \
             --progress=plain \
             .
+
+  create_manifest:
+    name: Create manifest slurm=${{ matrix.slurm_version }} ${{ matrix.variant }}
+    needs: [prepare, build_single_arch]
+    runs-on: ubuntu-24.04
+
+    strategy:
+      fail-fast: false
+      matrix:
+        slurm_version: ["25.05.6", "25.11.2"]
+        variant: ["cuda12", "cuda13"]
+        include:
+          - variant: cuda12
+            cuda_version: "12.9.0"
+            ubuntu_version: "ubuntu24.04"
+          - variant: cuda13
+            cuda_version: "13.0.2"
+            ubuntu_version: "ubuntu24.04"
+
+    env:
+      DAYTIME: ${{ needs.prepare.outputs.daytime }}
+      IMAGE_BASE: cr.eu-north1.nebius.cloud/ml-containers/slurm_training_diag
+
+    steps:
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
+
+      - name: Install Nebius CLI
+        shell: bash
+        run: |
+          set -euo pipefail
+          curl -sSL https://storage.eu-north1.nebius.cloud/cli/install.sh | bash
+          echo "$HOME/.nebius/bin" >> "$GITHUB_PATH"
+
+      - name: Restore Nebius config
+        shell: bash
+        run: |
+          set -euo pipefail
+          mkdir -p "$HOME/.nebius"
+          echo "${{ secrets.NEBIUS_CONFIG_YAML_B64 }}" | base64 -d > "$HOME/.nebius/config.yaml"
+          chmod 600 "$HOME/.nebius/config.yaml"
+
+      - name: Configure Nebius docker credential-helper for auth
+        shell: bash
+        run: nebius registry configure-helper
+
+      - name: Create and push manifest list
+        shell: bash
+        run: |
+          set -euo pipefail
+
+          IMAGE="${IMAGE_BASE}:slurm${{ matrix.slurm_version }}-cuda${{ matrix.cuda_version }}-${{ matrix.ubuntu_version }}-${DAYTIME}"
+          AMD="${IMAGE}-amd64"
+          ARM="${IMAGE}-arm64"
+
+          echo "Creating multi-arch manifest: ${IMAGE}"
+          echo "  - ${AMD}"
+          echo "  - ${ARM}"
+
+          docker buildx imagetools create \
+            -t "${IMAGE}" \
+            "${AMD}" \
+            "${ARM}"
diff --git a/Dockerfile b/Dockerfile
@@ -260,9 +260,12 @@ FROM training_diag AS slurm_training_diag
 # CUDA Image for training and diagnostics with a slurm client
 ######
 
+ARG SLURM_VERSION
+ENV SLURM_VERSION=$SLURM_VERSION
+
 # Install slurm client and divert files
 COPY ansible/slurm-client.yml /opt/ansible/slurm-client.yml
 COPY ansible/roles/slurm-client /opt/ansible/roles/slurm-client
 COPY ansible/roles/slurm-divert /opt/ansible/roles/slurm-divert
 RUN cd /opt/ansible && \
-    ansible-playbook -i inventory/ -c local slurm-client.yml
+    ansible-playbook -i inventory/ -c local slurm-client.yml -e "slurm_version=${SLURM_VERSION}"
