sGPU Test Scheduling: Global Work Queue

[VeeraRajasekhar]

sGPU Test Scheduling — Global Work Queue

Design document for .github/scripts/run_queue_sgpu.sh and the scheduler package under .github/scripts/scheduler/, which replace the static one-suite-per-GPU sGPU runner with a single LPT-ordered, self-calibrating work queue.

Contents

1. Analysis — old parallel run vs unordered queue vs ordered queue
2. Problem
3. Design — the three enabling mechanisms
4. Flow chart
5. Phase
6. The weight table
7. Usage

---

1. Analysis

Three ways to spread the same body of sGPU work across 4 GPUs. Measured on gfx950 at TEST_LEVEL=3: 10497s of total work, 80 work items.

Scheduler	Unit of work	Order	Wall clock	Utilisation	Speed-up
Old — static parallel	one suite per GPU	fixed, by config line	6935s · 1h55m35s	38%	1.00x
Unordered queue	one item, pulled on completion	arbitrary	~2890s mean	91%	2.40x
Ordered queue (LPT)	one item, pulled on completion	longest first, from the weight table	2643s · 44m03s	99.3%	2.62x

Refer: https://github.com/ROCm/TransformerEngine/actions/runs/31151497276

---

2. Problem

run_parallel_sgpu.sh assigns one suite per GPU, statically, by config-file line order (ci_sgpu_jobs.conf line i → GPU i). The suites are very unevenly sized, so most GPUs finish early and idle.

Measured on gfx950 at TEST_LEVEL=3:

Suite	Work	Share
torch	6935s	66%
jax	1987s	19%
core	1251s	12%
examples	324s	3%
Total	10497s

Wall clock equals the largest suite (6935s), so 4 GPUs deliver 10497 / (6935 × 4) = ~38% utilisation. Two structural causes:

1. The scheduling unit is too coarse. A suite is indivisible, so no assignment policy can improve on a 66/19/12/3 split across 4 bins.
2. Assignment is static. When jax finishes, its GPU cannot take torch work.

The imbalance is self-worsening: it reflects whichever suite grew last, and every new test file makes it more permanent.

Goals

• Cut sGPU wall clock without changing which tests run.
• Work unchanged at 4 GPUs (today) and 8 GPUs (future) with no reconfiguration.
• Preserve the per-suite log / exit-code / JUnit-XML contract so rocm-ci.yml needs a one-line change.
• Keep the suite scripts the single source of truth for what runs.

---

3. Design

One global queue; N workers, each pinned to one GPU, each pulling the next item when it goes idle. A GPU stops working only when the queue is empty.

        +-- LPT-ordered queue: 80 items, all four suites interleaved --+
GPU0  --| ####  ##  #  #####  #  ##  ###  #  ####  #  ##  #            |
GPU1  --| ##  ######  #  ###  ####  #  ####  #                         |  2643s
GPU2  --| #####  ##  ######  #  ###  ##                                |
GPU3  --| ############  #  ##                                          |
        +---------------------------------------------------------------+

The queue itself is easy. The hard parts are (a) obtaining a correct item list without duplicating suite logic, and (b) dispatching an item without losing its configuration. Three mechanisms address them. All three are env-var hooks inside ci/_utils.sh, so the suite scripts pick them up without knowing a scheduler exists — and each is explained in the phase that uses it rather than in the abstract:

Mechanism	Hook	What it solves	Explained in
List mode	TE_CI_LIST_ITEMS=1	a suite enumerates its own work items, so the scheduler never has to reimplement its gating	Phase 1
Setup hoisting	TE_CI_SETUP_ONLY=1, TE_CI_SKIP_SETUP=1	pip and CK-JIT setup run once per suite, not once per item	Phase 3
Dispatch by re-entry	TEST_FILTER=<tag>	an item runs in its exact original configuration, because the suite script itself replays it	Phase 4

---

4. Flow chart

Rounded boxes are steps, cylinders are files. Each file says what is in it.

flowchart TD
    A[("<b>the config</b><br/><i>which suites exist, and<br/>whether each one can be split</i>")] --> P0

    subgraph SETUP [" "]
        P0["<b>Phase 0</b> — read the config<br/><i>load the suites; stop now if any<br/>label is missing or repeated</i>"]
        S["<b>Run setup</b><br/><i>which GPUs, which chip,<br/>where the logs go</i>"]
        P0 --> S
    end

    S --> P1

    P1["<b>Phase 1</b> — expand<br/><i>ask each suite to list its own<br/>work items instead of running them</i>"]
    P1 -->|"what this machine will run"| Q0[("<b>the queue, unsorted</b><br/><i>one line per item:<br/>suite, command, tag</i>")]
    P1 -->|"what exists at this level"| IT[("<b>the item census</b><br/><i>every item that still exists,<br/>even ones skipped here</i>")]

    W0[("<b>the weight table</b><br/><i>how many seconds each item<br/>took the last few runs</i>")]

    subgraph PH2 ["<b>Phase 2</b> — weight and order"]
        P2A["<b>load the cached weights</b><br/><i>look up a duration for every queued item;<br/>anything never seen before counts as huge</i>"]
        P2B["<b>order the queue — longest first (LPT)</b><br/><i>big items go out early, so only short ones<br/>are left to fill the end of the run</i>"]
        P2A --> P2B
    end

    W0 -.Update the cached weights.-> P2A
    Q0 --> P2B
    P2B --> Q1[("<b>the run plan</b><br/><i>the same items, now in<br/>the order they get handed out</i>")]

    Q1 --> P3
    P3["<b>Phase 3</b> — set up once<br/><i>pip install and CK-JIT build,<br/>one time per suite, not per item</i>"]

    P3 --> P4
    P4["<b>Phase 4</b> — run the queue<br/><i>one worker per GPU; a worker takes<br/>new work only when it is free</i>"]
    P4 --> W1["worker on gpu0"]
    P4 --> W2["worker on gpu1"]
    P4 --> W3["worker on gpuN"]
    W1 & W2 & W3 -->|"a lock, so two workers<br/>never take the same item"| TK{{"take the next item<br/>off the plan"}}
    TK --> RUN["<b>re-enter the suite's own script</b><br/>ci/pytorch.sh, ci/jax.sh, …<br/>TEST_FILTER=tag<br/>TE_CI_SKIP_SETUP=1<br/>HIP_VISIBLE_DEVICES=gpu<br/><i>the script replays just that one<br/>call line, in its original config</i>"]
    RUN --> LOGS[("<b>per-item logs</b><br/><i>what each item printed,<br/>and whether it passed</i>")]
    RUN --> TIM[("<b>timings</b><br/><i>per item: which GPU, how long,<br/>pass/fail, was it cut off</i>")]
    RUN --> XML[("<b>JUnit XML</b><br/><i>the test results the CI<br/>report is built from</i>")]

    LOGS --> P5
    P5["<b>Phase 5</b> — per-suite verdict<br/><i>did this suite pass? the worst<br/>item decides</i>"]
    P5 --> SL[("<b>per-suite result</b><br/><i>one exit code, plus an index<br/>of that suite's items</i>")]

    TIM --> P6
    IT -.what to keep.-> P6
    P6["<b>Phase 6</b> — learn<br/><i>blend today's durations into the<br/>weights; drop items that are gone</i>"]
    P6 --> W0

    TIM --> P7
    W0 -.read back.-> P7
    P7["<b>Phase 7</b> — report<br/><i>how well the GPUs were packed,<br/>estimate vs actual per item</i>"]
    P7 --> RM[("<b>schedule report</b><br/><i>posted to the CI job summary</i>")]

    P5 --> P8["<b>Phase 8</b> — failure summary<br/><i>print the failed items,<br/>so the terminal ends useful</i>"]
    P8 --> EX["exit with the worst<br/>exit code seen"]

Loading

Same flow in a terminal, with the real file names — the layout of each is in §6:

conf ──▶ [0] parse ──▶ setup: GPUs, arch, log tree
                              │
                              ▼
                        [1] expand ──┬──▶ queue.tsv.raw ─┐
                                     └──▶ items.tsv ───┐ │
                                                       │ ▼
         weight table ──────────────▶ [2] order (LPT) ─│─┴─▶ queue.tsv
                                            │          │
                                     [3] setup once per suite
                                            │          │
                                     [4] run: N workers pull under flock
                                            ├──▶ items/*.log + .rc ──┐
                                            └──▶ timings.tsv ──┬─────┤
                                                               │     ▼
                                                               │  [5] verdict:
                                                               │  suites/*.log.rc
                                                               │     │
                        [6] learn ◀────── items.tsv ◀──────────┤     │
                              │                                │     │
                              └──▶ weight table (next run)     │     │
                                                               ▼     │
                                                    [7] report/schedule.md
                                                                     │
                                                    [8] failed-item list
                                                                     │
                                                              exit OVERALL_RC

Reading the dispatch step (Phase 4). A worker never invokes pytest itself. It runs the suite's own entry script again — ci/pytorch.sh, ci/jax.sh, whatever the config named — with three variables set, and that script replays exactly one of its own call lines:

Variable	Meaning
TEST_FILTER=<tag>	run only the call line that owns this tag; skip the other ~79. The item therefore keeps its original NVTE_* prefix, pytest args and backend flags, because the line that runs it is unchanged
TE_CI_SKIP_SETUP=1	do not redo pip / CK-JIT setup — Phase 3 already did it once for the whole run
HIP_VISIBLE_DEVICES=<gpu>	pin this item to the worker's own GPU, so concurrent items never share a device

Full detail in Phase 4.

---

5. Phase

Phase 1 — expand every suite into work items

Each suite is list (expandable) or opaque (a single item — core and examples have no per-item filter mechanism).

A list suite enumerates itself. The suite scripts encode substantial policy: TEST_LEVEL gating, the fused-attn backend matrix, runtime capability probes (check_supported mxfp8, check_supported flash_attn), and per-item env prefixes. Reimplementing that here would guarantee drift, so the scheduler asks the script instead. TE_CI_LIST_ITEMS=1 makes pytest_run print the item's tag and return rather than run it — placed after the existing level, filter and uniqueness checks, so every gate still applies (ci/_utils.sh:416):

if [ -n "$TE_CI_LIST_ITEMS" ]; then
    echo "TE_CI_ITEM $_test_name_tag"
    return
fi

TE_CI_LIST_ITEMS=1 ci/pytorch.sh therefore walks the real schedule and prints exactly what it would have executed, and the tags it prints are the same ones Phase 4 dispatches by.

It is listed twice, because "will run" and "exists" are different questions:

pass	environment	answers	goes to
1	TE_CI_LIST_ITEMS=1	what this host will run, after capability probes	queue.tsv.raw
2	+ TE_CI_SKIP_CHECK_SUPPORTED=1	what exists at this TEST_LEVEL	items.tsv

pass2 − pass1 is what this host skipped — no flash-attn, say — and those tests do still exist. Only a tag in neither list is gone for good, and that is what Phase 6 prunes on (§7.3).

Phase 2 — weight and order

build_weights.py order joins the raw queue against the weight table and sorts descending by weight, emitting weight ⇥ label ⇥ cmd ⇥ tag ⇥ rest.

Longest-processing-time-first is the standard makespan heuristic: dispatch the long items first and what is left to fill the tail is short, so no GPU is still starting a big item once the others have gone idle.

An item the table has never seen takes DEFAULT_WEIGHT=999999 and therefore sorts first. unknown-first costs at most one item's runtime of slightly-wrong ordering; unknown-last makes the new item the tail, idling N−1 GPUs behind it.

The phase prints the full dispatch plan before anything runs, ending with the lower bound total_work / N — which is what the wall clock is later compared against.

Phase 4 — run the queue

One background worker per GPU. Workers share a single index file guarded by flock:

take_next() {
    local i
    { flock 9; i=$(cat "$IDX_FILE"); echo $((i + 1)) > "$IDX_FILE"; } 9<>"$LOCK_FILE"
    echo "$i"
}

Each worker loops: take an index, sed -n "${i}p" the queue, dispatch, record.

Dispatch is a re-entry of the suite script, not a pytest invocation. The obvious alternative — serialise each item (path, -k expression, NVTE_* vars) and have the scheduler call pytest itself — is fragile. Items such as

NVTE_USE_ATOMIC_AMAX=1 NVTE_USE_CAST_TRANSPOSE_TRITON=1 \
    run_default_fa_lbl "amax+triton" 3 test_numerics.py

carry state that is easy to drop, and dropping it means the test silently runs in the wrong configuration — a failure mode that produces green runs with no coverage. So the scheduler reconstructs nothing. It runs the suite's own command again with three variables set:

HIP_VISIBLE_DEVICES=$gpu TE_CI_SKIP_SETUP=1 TEST_FILTER="$tag" \
    JUNITXML_PREFIX="$junit_dir" "$cmd" ${rest:-} > "$itemlog" 2>&1

Variable	Value	Effect on the re-entered script
TEST_FILTER	the item's tag, e.g. test_numerics.ck.amax+triton	check_test_filter returns early on every call line but the one owning that tag, which then runs with its original env prefix, level and arguments — unmodified
TE_CI_SKIP_SETUP	1	check_setup_needed reports nothing to do, so pip and CK-JIT are not repeated; Phase 3 did them once
HIP_VISIBLE_DEVICES	this worker's GPU index	the item sees one GPU, so N items running at once never share a device

"$cmd" is the suite's own entry point from the config (ci/pytorch.sh, ci/jax.sh, …), so a dispatched item is the same invocation the suite would have made unattended. The tag is the only thing that crosses the scheduler boundary, and check_test_tag_unique (ci/_utils.sh:408) guarantees it selects exactly one call line. That round trip — a tag printed by Phase 1 selecting the identical item in Phase 4 — is the load-bearing assumption of the design, so it was verified directly: 123 tags round-tripped, 0 mismatches.

An opaque item has an empty tag and takes the else branch: no TEST_FILTER, no TE_CI_SKIP_SETUP, so the script runs end to end exactly as it does today.

Pull-on-completion is what makes the system robust to bad weights. A poor ordering costs a little tail latency; it can never cost correctness or balance. The bootstrap weight table was once wrong by 16x on one file and the run still hit 99% utilisation.

Per item the worker writes four things:

• items/<label>.<tag>.log — the test output;
• items/<label>.<tag>.log.rc — the exit code, which is what makes selective per-item reporting possible later;
• one row of timings.tsv, appended as the item ends, so a killed run still records everything that had finished;
• one console line: [HH:MM:SS] gpu1 t+1204s 260s rc=0 est=245s torch test_numerics.ck.amax+triton.

It also decides the incomplete flag. A te_ci_result_sink .partial sidecar that outlived the process means pytest never reached its end-of-session write, so the duration is where the item was cut off, not what it costs. The check happens here because this is the only moment it is unambiguous — see §7.2.

No scheduler-imposed deadline is applied: the suite scripts' own PYTEST_TIMEOUT and the workflow's timeout-minutes remain the only limits, exactly as outside the queue.

6. The weight table Storage and update rule

Storage is the GHA cache alone, keyed te-sgpu-weights-<arch>-l<level>-<run_id> with a prefix restore-keys fallback — the standard append-only idiom, since cache entries are immutable and a fresh key per run is how the table gets updated at all. There is deliberately no committed seed table: it would go stale silently and show up in every CI diff, whereas one unordered run per cache eviction is self-healing (that run writes the table the next one reads).

The update is an asymmetric EWMA, per key:

measured > old:  new = old + 0.5 * (measured - old)     # ALPHA_UP
measured < old:  new = old + 0.1 * (measured - old)     # ALPHA_DOWN

Rises in ~2 runs, decays over ~10. The asymmetry follows from LPT's error profile: over-estimating an item only starts it earlier than necessary, while under-estimating puts a long item late and creates the tail. So react fast to growth, treat a single fast run as probably noise.

---

7. Usage

Inside the dev container, the whole sGPU set at level 1:

TEST_LEVEL=1 .github/scripts/run_queue_sgpu.sh

That is the complete command. The arch comes from rocminfo, the config defaults to ci_sgpu_queue.conf, logs go to test-results/logs/, and the weight table to ci-weights/test_weights.<arch>.l1.txt.

The queue uses every GPU it can see; restrict it as you would any ROCm program:

HIP_VISIBLE_DEVICES=0,1 TEST_LEVEL=1 .github/scripts/run_queue_sgpu.sh

The whole surface is --log-dir DIR plus an optional list of config files, taken positionally:

TEST_LEVEL=1 .github/scripts/run_queue_sgpu.sh my_subset.conf

[Copilot]

Pull request overview

Introduces a self-calibrating, LPT-ordered global queue for distributing sGPU tests dynamically across available GPUs.

Changes:

• Adds test enumeration, filtered dispatch, and setup-hoisting hooks.
• Adds queue execution, weight learning, and scheduling reports.
• Integrates cached weights and queue artifacts into ROCm CI.

Reviewed changes

Copilot reviewed 9 out of 10 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
.gitignore	Ignores queue outputs and weights.
.github/workflows/rocm-ci.yml	Integrates queue execution and caching.
.github/scripts/run_queue_sgpu.sh	Implements the global worker queue.
.github/scripts/ci_sgpu_queue.conf	Defines queued suites and modes.
.github/scripts/scheduler/build_weights.py	Orders items and updates weights.
.github/scripts/scheduler/queue_files.py	Parses timing records.
.github/scripts/scheduler/schedule_report.py	Generates scheduling reports.
ci/_utils.sh	Adds scheduler hooks and tag validation.
ci/pytorch.sh	Supports PyTorch item listing and dispatch.
ci/jax.sh	Supports JAX setup hoisting and unique tags.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.