ci: Bump runner

.github/scripts/check_for_ngc_images.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+
+# Configuration
+BASE_IMAGE="nvcr.io/nvidia/pytorch"
+TAG_SUFFIX="-py3"
+MONTHS_TO_CHECK=7 # Check current month and previous 6 months (total 7)
+
+# Initialize an array to store existing tags
+EXISTING_TAGS=()
+
+echo "Checking for existence of the last ${MONTHS_TO_CHECK} NGC PyTorch images: ${BASE_IMAGE}:YY.MM${TAG_SUFFIX}"
+echo "---------------------------------------------------------------------"
+
+# Loop through the last N months
+for i in $(seq 0 $((MONTHS_TO_CHECK - 1))); do
+    # Calculate Year and Month for the tag
+    CURRENT_YEAR=$(date +%Y)
+    CURRENT_MONTH=$(date +%m)
+
+    # Calculate target month and year
+    TARGET_DATE=$(date -d "$CURRENT_YEAR-$CURRENT_MONTH-01 -$i months" +%y.%m)
+
+    # Construct the full image tag and the tag-only string
+    IMAGE_TAG="${TARGET_DATE}${TAG_SUFFIX}"
+    FULL_IMAGE="${BASE_IMAGE}:${IMAGE_TAG}"
+
+    echo "Checking: ${FULL_IMAGE}"
+
+    # Use 'docker manifest inspect' to check for image existence without pulling.
+    if docker manifest inspect "${FULL_IMAGE}" > /dev/null 2>&1; then
+        echo "✅ EXISTS: Found."
+        # Add the tag-only string to the array
+        EXISTING_TAGS+=("nvcr.io/nvidia/pytorch:${IMAGE_TAG}")
+    else
+        echo "❌ MISSING: Not found."
+    fi
+done
+
+echo "---------------------------------------------------------------------"
+
+## JSON Output Generation
+# This uses the collected array to build a JSON string.
+
+# 1. Convert the shell array to a newline-separated string.
+TAGS_NL_SEP=$(printf "%s\n" "${EXISTING_TAGS[@]}")
+
+# 2. Use jq to read the newline-separated list and format it into a JSON array.
+# . | split("\n") | .[:-1] reads the input, splits it by newline, and removes the trailing empty element.
+if command -v jq &> /dev/null; then
+    JSON_STRING=$(echo -e "${TAGS_NL_SEP}" | jq -R -s 'split("\n") | .[:-1]')
+
+    echo "Generated JSON String of Existing Tags:"
+    echo "${JSON_STRING}"
+
+    # Optional: Save the JSON string to a variable for further use
+    # echo "JSON_STRING is now available in the shell if you source this script."
+else
+    echo "WARNING: 'jq' is not installed. Cannot format output as JSON."
+    echo "Found Tags: ${EXISTING_TAGS[*]}"
+fi
+
+echo "---"
+echo "Check complete."
+
+echo "${JSON_STRING}" > ngc_images.json

.github/workflows/publish.yaml

@@ -42,7 +42,7 @@ jobs:
       matrix:
         # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
         # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
-        os: [ubuntu-20.04]
+        os: [ubuntu-22.04]
         python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
         torch-version: ["2.1.2", "2.2.2", "2.3.1", "2.4.0", "2.5.1", "2.6.0.dev20241001"]
         cuda-version: ["11.8.0", "12.3.2"]
@@ -76,20 +76,31 @@ jobs:
       release-version: ${{ needs.setup_release.outputs.release-version }}
       upload-to-release: true
 
+  check_for_ngc_images:
+    runs-on: ubuntu-latest
+    outputs:
+      images: ${{ steps.check_for_ngc_images.outputs.IMAGES }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Check for NGC PyTorch images
+        id: check_for_ngc_images
+        run: |
+          bash ./.github/scripts/check_for_ngc_images.sh
+          echo "IMAGES=$(cat ngc_images.json| jq -cr)" | tee -a $GITHUB_OUTPUT
+
   build_ngc_wheels:
     name: Build Wheel for NGC PyTorch
-    needs: setup_release
+    needs: [setup_release, check_for_ngc_images]
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-20.04]
-        container-image:
-          - nvcr.io/nvidia/pytorch:25.05-py3
-          - nvcr.io/nvidia/pytorch:25.06-py3
-          - nvcr.io/nvidia/pytorch:25.08-py3
+        os: [ubuntu-22.04]
+        container-image: ${{ fromJson(needs.check_for_ngc_images.outputs.images) }}
     uses: ./.github/workflows/_build_in_container.yml
     with:
-      runs-on: ${{ matrix.runs-on }}
+      runs-on: ${{ matrix.os }}
       container-image: ${{ matrix.container-image }}
       release-version: ${{ needs.setup_release.outputs.release-version }}
       upload-to-release: true