Move the remaining GPU workflows onto the per-task GPU queue

.github/scripts/gpu_task_queue.py

@@ -13,9 +13,14 @@
 Usage:
     gpu_task_queue.py --install git unittests:8 unittests:4 unittests:2 unittests:1
     gpu_task_queue.py --install editable new_examples:8 new_examples:2
+    gpu_task_queue.py --task 8 "GEMM All-Scatter" "bash .github/scripts/run_perf_benchmark.sh ..."
 
-Each task is <test_dir>:<num_ranks>. "new_examples" runs run_new_examples.sh;
-anything else runs run_tests.sh on tests/<test_dir>.
+A positional task is <test_dir>:<num_ranks>. "new_examples" runs
+run_new_examples.sh; anything else runs run_tests.sh on tests/<test_dir>.
+
+--task RANKS NAME COMMAND runs COMMAND with bash, for jobs that are not a test
+directory. Every task runs in a private copy of the checkout with GPU_DEVICES
+set to its GPUs, and must pass them on (container_exec.sh --gpus).
 
 Scheduling: largest task first, smaller tasks backfill whatever GPUs are free.
 If the largest pending task has waited RESERVE_AFTER seconds, it posts a
@@ -58,14 +63,14 @@ def log(msg):
 
 
 class Task:
-    def __init__(self, spec, install):
-        test_dir, _, ranks = spec.rpartition(":")
-        if not test_dir or not ranks.isdigit() or not 1 <= int(ranks) <= MAX_GPUS:
-            raise SystemExit("bad task '{}': expected <test_dir>:<1-{}>".format(spec, MAX_GPUS))
+    def __init__(self, ranks, name, test_dir=None, install=None, shell=None):
+        if not ranks.isdigit() or not 1 <= int(ranks) <= MAX_GPUS:
+            raise SystemExit("bad task '{}': ranks must be 1-{}, got '{}'".format(name, MAX_GPUS, ranks))
         self.test_dir = test_dir
         self.ranks = int(ranks)
         self.install = install
-        self.name = "{} ({} ranks, {})".format(test_dir, self.ranks, install)
+        self.shell = shell
+        self.name = name
         self.proc = None
         self.mask = 0
         self.gpus = ""
@@ -77,7 +82,17 @@ def __init__(self, spec, install):
         self.ended_at = None
         self.rc = None
 
+    @classmethod
+    def from_spec(cls, spec, install):
+        """A positional <test_dir>:<ranks> task."""
+        test_dir, _, ranks = spec.rpartition(":")
+        if not test_dir or not ranks.isdigit() or not 1 <= int(ranks) <= MAX_GPUS:
+            raise SystemExit("bad task '{}': expected <test_dir>:<1-{}>".format(spec, MAX_GPUS))
+        return cls(ranks, "{} ({} ranks, {})".format(test_dir, ranks, install), test_dir=test_dir, install=install)
+
     def command(self):
+        if self.shell is not None:
+            return ["bash", "-c", self.shell]
         scripts = os.path.join(self.workdir, ".github", "scripts")
         if self.test_dir == "new_examples":
             return ["bash", os.path.join(scripts, "run_new_examples.sh"), str(self.ranks), self.install]
@@ -135,7 +150,6 @@ def __init__(self, tasks):
         self.pending = sorted(tasks, key=lambda t: -t.ranks)
         self.running = []
         self.done = []
-        self.head_waiting_since = None
         self.reserved_for = None
 
     def held_mask(self):
@@ -183,14 +197,11 @@ def try_start(self):
             alloc.write_bitmap(bitmap)
             self.record_held()
 
-            # Anti-starvation: track how long the largest pending task has waited.
-            # The head only changes when it launches, since nothing is added later.
-            if head is None or head in to_launch:
-                self.head_waiting_since = None
-            else:
-                if self.head_waiting_since is None:
-                    self.head_waiting_since = time.time()
-                waited = time.time() - self.head_waiting_since
+            # Anti-starvation: the largest pending task reserves the node once it
+            # has been queued for RESERVE_AFTER. Counted from when the job queued
+            # it, so a task that already waited behind a sibling does not restart.
+            if head is not None and head not in to_launch:
+                waited = time.time() - head.queued_at
                 if self.reserved_for is head:
                     alloc.reserve(head.ranks)  # refresh
                 elif self.reserved_for is None and waited >= RESERVE_AFTER and (not foreign or head.ranks > res[1]):
@@ -330,11 +341,25 @@ def summary(self):
 
 def main():
     parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    parser.add_argument("--install", required=True, choices=["git", "editable", "install"])
-    parser.add_argument("tasks", nargs="+", metavar="TEST_DIR:RANKS")
+    parser.add_argument("--install", choices=["git", "editable", "install"], help="required for TEST_DIR:RANKS tasks")
+    parser.add_argument(
+        "--task",
+        nargs=3,
+        action="append",
+        default=[],
+        metavar=("RANKS", "NAME", "COMMAND"),
+        help="run COMMAND with bash on RANKS GPUs; repeatable",
+    )
+    parser.add_argument("tasks", nargs="*", metavar="TEST_DIR:RANKS")
     args = parser.parse_args()
-
-    queue = Queue([Task(spec, args.install) for spec in args.tasks])
+    if not args.tasks and not args.task:
+        parser.error("no tasks given")
+    if args.tasks and not args.install:
+        parser.error("--install is required for TEST_DIR:RANKS tasks")
+
+    tasks = [Task(ranks, name, shell=command) for ranks, name, command in args.task]
+    tasks += [Task.from_spec(spec, args.install) for spec in args.tasks]
+    queue = Queue(tasks)
 
     def on_signal(signum, _frame):
         raise SystemExit(128 + signum)

.github/scripts/release_gpus.sh

@@ -5,7 +5,7 @@
 # Release GPUs for CI workflows - to be called as a workflow step with if: always()
 # Usage: release_gpus.sh
 #
-# Reads GPU allocation details from environment variables set by acquire_gpus.sh
+# Returns whatever GPUs gpu_task_queue.py still holds if it was killed mid-run.
 
 set -e
 

.github/scripts/run_external_validation.sh

@@ -0,0 +1,44 @@
+#!/bin/bash
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Install Iris from this commit and run an external validation test on 2 ranks.
+# Usage: run_external_validation.sh <triton|gluon>
+#
+# The tests live in gists, outside the repo, so they exercise Iris the way an
+# external user would: a pip install from git rather than the checkout.
+
+set -e
+
+case "$1" in
+    triton)
+        TEST_URL="https://gist.githubusercontent.com/mawad-amd/6375dc078e39e256828f379e03310ec7/raw/0827d023eaf8e9755b17cbe8ab06f2ce258e746a/test_iris_distributed.py"
+        ;;
+    gluon)
+        TEST_URL="https://gist.githubusercontent.com/mawad-amd/2666dde8ebe2755eb0c4f2108709fcd5/raw/c5544943e2832c75252160bd9084600bf01a6b06/test_iris_gluon_distributed.py"
+        ;;
+    *)
+        echo "[ERROR] Usage: $0 <triton|gluon>"
+        exit 1
+        ;;
+esac
+
+if [ -z "$GPU_DEVICES" ]; then
+    echo "[ERROR] GPU_DEVICES is not set; run this through gpu_task_queue.py"
+    exit 1
+fi
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO=${GITHUB_REPOSITORY:-"ROCm/iris"}
+SHA=${GITHUB_SHA:-"HEAD"}
+TEST_FILE=$(basename "$TEST_URL")
+
+"$SCRIPT_DIR/container_exec.sh" --gpus "$GPU_DEVICES" "
+    set -e
+    cd /iris_workspace
+    pip install git+https://github.com/${REPO}.git@${SHA}
+    wget -O ${TEST_FILE} ${TEST_URL}
+    torchrun --rdzv-backend=c10d --rdzv-endpoint=localhost:0 --nnodes=1 --nproc_per_node=2 ${TEST_FILE}
+"
+
+echo "✅ External $1 validation test passed!"

.github/scripts/run_new_examples.sh

@@ -14,7 +14,7 @@ INSTALL_METHOD=${2:-"editable"}
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-# GPU_DEVICES should be set by the workflow-level acquire_gpus.sh step
+# GPU_DEVICES is set by gpu_task_queue.py
 GPU_ARG=""
 if [ -n "$GPU_DEVICES" ]; then
     GPU_ARG="--gpus $GPU_DEVICES"

.github/scripts/run_tests.sh

@@ -11,6 +11,9 @@
 #     - "git": pip install git+https://github.com/${{ github.repository }}.git@${{ github.sha }}
 #     - "editable": pip install -e .
 #     - "install": pip install .
+#
+# CONTAINER_IMAGE, if set, selects the container image (e.g. the nightly Triton
+# image) instead of the one container_build.sh produced.
 
 set -e
 
@@ -43,18 +46,21 @@ fi
 
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-# GPU_DEVICES should be provided by workflow-level acquire_gpus.sh step
-# or via command-line argument for backward compatibility
+# GPU_DEVICES is set by gpu_task_queue.py, or passed as an argument
 if [ -z "$GPU_DEVICES" ]; then
     echo "[RUN-TESTS] WARNING: No GPUs allocated. GPU_DEVICES not set."
     echo "[RUN-TESTS] Tests may fail if they require GPUs."
 fi
 
-# Build GPU argument
+# Build GPU and image arguments
 GPU_ARG=""
 if [ -n "$GPU_DEVICES" ]; then
     GPU_ARG="--gpus $GPU_DEVICES"
 fi
+IMAGE_ARG=()
+if [ -n "$CONTAINER_IMAGE" ]; then
+    IMAGE_ARG=(--image "$CONTAINER_IMAGE")
+fi
 
 # Build install command based on method
 INSTALL_CMD=""
@@ -72,9 +78,10 @@ fi
 # Run tests in container
 EXIT_CODE=0
 # shellcheck disable=SC2086
-"$SCRIPT_DIR/container_exec.sh" $GPU_ARG "
+"$SCRIPT_DIR/container_exec.sh" "${IMAGE_ARG[@]}" $GPU_ARG "
     set -e
 
+    echo \"Triton version: \$(pip show triton 2>/dev/null | grep Version || echo unknown)\"
     echo \"Installing iris using method: $INSTALL_METHOD\"
     $INSTALL_CMD
 

.github/workflows/iris-ccl-benchmark.yml

@@ -83,26 +83,23 @@ jobs:
         run: |
           bash .github/scripts/container_build.sh
 
-      - name: Acquire GPUs
+      - name: Run the sweep
         run: |
-          # Enough GPUs for the largest rank count; smaller ones run inside it.
-          # Validate first: this runs before run_ccl_benchmark.sh checks inputs.
+          # Validate first: num_ranks sizes the GPU request below, before
+          # run_ccl_benchmark.sh checks its inputs.
           if ! [[ "$NUM_RANKS" =~ ^[1-8](,[1-8])*$ ]]; then
             echo "[ERROR] num_ranks must be comma-separated integers 1-8, got '$NUM_RANKS'"
             exit 1
           fi
+          # Enough GPUs for the largest rank count; smaller ones run inside it.
           MAX_RANKS=$(echo "$NUM_RANKS" | tr ',' '\n' | sort -n | tail -1)
-          bash .github/scripts/acquire_gpus.sh "$MAX_RANKS"
-
-      - name: Run the sweep
-        run: |
-          set -e
-          bash .github/scripts/run_ccl_benchmark.sh \
-            "$NUM_RANKS" \
-            "$DTYPES" \
-            "$M_VALUES" \
-            "$N_VALUES" \
-            "$N_REPEAT"
+          # Single-quoted: the task's shell expands the inputs from the
+          # environment, so they are never spliced into the command text. The
+          # task runs in a private copy of the checkout that is deleted when it
+          # ends, so copy the results out for the steps below, even on failure.
+          python3 .github/scripts/gpu_task_queue.py --task "$MAX_RANKS" "CCL sweep" \
+            'bash .github/scripts/run_ccl_benchmark.sh "$NUM_RANKS" "$DTYPES" "$M_VALUES" "$N_VALUES" "$N_REPEAT"; rc=$?
+             cp -r ccl_results "$GITHUB_WORKSPACE/" 2>/dev/null; exit $rc'
 
       - name: Publish the comparison table
         if: always()