diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83b4edfb55..68dbfb8d2d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,36 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
hipBLASLt, and the fusion backend. `deploy/` is also absent -- every file in
it targets the retired repository.
+- **`scripts/partition_mode_sweep.py` measures which compute-partition mode a
+ workload wants.** Sets each mode on one card in turn, runs the same benchmark
+ on every partition that mode creates, sums the throughput, and restores the
+ card's entry mode on the way out — including after a failure or a Ctrl-C.
+ Modes whose partitions provably cannot hold the configured streams are skipped
+ with the arithmetic shown rather than run into an out-of-memory failure.
+ The fan-out is the substance of it. A benchmark that loads one partition and
+ ignores the rest measures a fraction of the card, which reports `CPX` as eight
+ times worse than it is; every figure here is the sum over a mode's partitions
+ with all of them loaded together, and a mode is reported only when every one of
+ its partitions returned a measurement. Partitions are selected by matching CU
+ count within the swept card's PCI bus, never by device index: `amd-smi` orders
+ by PCI address while HSA/HIP enumerates whole cards first, so on an 8-card
+ MI355X node with card 0 in `CPX` the two tools disagree about which devices the
+ partitions are — 0-7 against 7-14.
+ This is where the privileged `amd-smi set` lives, and the only place it does.
+ A card-wide mutation that evicts every GPU context is reasonable between
+ benchmarks in a script an operator ran on purpose, and unreasonable inside an
+ optimization loop that also runs agent-authored code, so `optimize` continues
+ to only read the mode. Together the two halves are a boundary: the sweep
+ chooses the shape, the session asserts it.
+ Because that set evicts work, the check standing in front of it fails closed:
+ an `amd-smi` process listing in a shape the parser does not model is a refusal,
+ not an empty one, since the only wrong answer that destroys anything is reading
+ a busy node as free. It is scoped to the card being swept, so a neighbour's
+ benchmark on a shared node no longer forces `--allow-busy` and with it the loss
+ of the guard on the target card. Every exit from a started sweep runs the
+ restore and the report, including on an error the script does not model — which
+ exits `4`, keeps the modes already measured, and still yields `3` if the card
+ could not be put back.
- **The card's compute-partition shape is now recorded, checked, and published.**
An MI300-series card can be split into independent partitions (`SPX`, `DPX`,
`QPX`, `CPX`), and splitting one trades per-request latency for aggregate
diff --git a/README.md b/README.md
index 18ffceab54..968b61fd24 100755
--- a/README.md
+++ b/README.md
@@ -86,6 +86,12 @@ feedback on how to improve Hyperloom by completing the
- Main agent instructions: [`src/hyperloom/inference_optimizer/SKILL.md`](src/hyperloom/inference_optimizer/SKILL.md)
- CLI entry point: `python -m hyperloom.inference_optimizer.cli optimize`
- Operator tools: `python -m hyperloom.inference_optimizer.tools.*`
+- Compute-partition sweep: `python3 scripts/partition_mode_sweep.py` — sets each
+ AMD partition mode (`SPX`/`DPX`/`QPX`/`CPX`) on one card in turn, runs the same
+ benchmark on every partition that mode creates, sums the throughput and restores
+ the entry mode. Answers which shape a workload wants before a session commits to
+ one; `optimize` itself only ever reads the mode. Needs privilege for the set, so
+ it is a script rather than part of the loop.
- Platform tuning audit: `python3 scripts/platform_audit.py` — checks the host CPU
tuning that silently changes benchmark results. Judges Core Performance Boost and
the cpufreq governor against [AMD's BIOS & Workload Tuning Guide for EPYC 9004][58011];
diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md
index cf5f4f7674..c677e93ad9 100644
--- a/docs/reference/environment-variables.md
+++ b/docs/reference/environment-variables.md
@@ -501,6 +501,52 @@ drive, and passing the flags with one warns. Its shape is still recorded in the
report and the fingerprint — that is provenance, not a hand-off — and the report
states plainly that the figure cannot be read as an aggregate.
+### Choosing the mode: `scripts/partition_mode_sweep.py`
+
+`optimize` treats the mode as fixed and asserts it. Deciding *which* mode to be
+in is a separate job, done before the session, by
+`python3 scripts/partition_mode_sweep.py`. It sets each mode on one card in
+turn, runs the same benchmark on every partition that mode creates, sums the
+result, and restores the card's entry mode on the way out.
+
+```bash
+# what it would do, nothing set
+python3 scripts/partition_mode_sweep.py --benchmark-config bench.yaml --dry-run
+
+# sweep every mode the card reports, skipping any that cannot hold the workload
+python3 scripts/partition_mode_sweep.py \
+ --benchmark-config bench.yaml --output-dir /shared/sweep \
+ --per-stream-gib 20.7 --sudo
+```
+
+The fan-out is the point rather than a detail: a benchmark that loads one
+partition and ignores the rest measures a fraction of the card, which makes
+`CPX` look eight times worse than it is. The sweep therefore launches every
+partition at once, and reports a mode only when all of its partitions returned
+a measurement — a mode with six of eight reporting is unmeasured, not slow.
+
+It publishes the same `HYPERLOOM_PARTITION_*` variables as a session, so a
+benchmark entrypoint written against the table above works unchanged under
+either. It pins each process with `ROCR_VISIBLE_DEVICES` and removes any
+inherited `HIP_VISIBLE_DEVICES`, because two masks apply in sequence and the
+second indexes into the first.
+
+The privileged `amd-smi set` lives here and nowhere else. An operator-run script
+between benchmarks is a reasonable place for a card-wide mutation that evicts
+every GPU context; an optimization loop that also runs agent-authored code is
+not. Before setting anything it refuses if a process holds a context on the card
+being swept — and only that card, since no other card is repartitioned. A
+neighbour's benchmark on a shared node is not a reason to stop. If `amd-smi`
+reports its process list in a shape the script cannot read, that is also a
+refusal rather than an assumption that the card is idle: `--allow-busy` is the
+way past both, and `--dry-run` never asks.
+
+Exit codes: `0` swept, `1` nothing measurable, `2` refused before anything
+changed, `3` swept but the card could not be restored to its entry mode, `4`
+stopped on an error it does not model. Every path out of a started sweep goes
+through the restore and the report, so a mode that fails unexpectedly costs its
+own result and nothing else.
+
---
## Multi-node / prefill-decode (PD)
diff --git a/scripts/partition_mode_sweep.py b/scripts/partition_mode_sweep.py
new file mode 100644
index 0000000000..645ed84eff
--- /dev/null
+++ b/scripts/partition_mode_sweep.py
@@ -0,0 +1,1248 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: MIT
+
+"""Sweep the AMD compute-partition modes and report which one the workload wants.
+
+Sets each requested mode on one card in turn, runs the same fixed benchmark on
+every partition that mode creates, and sums the result. The comparison it prints
+is the one an operator needs before committing a multi-hour optimization
+session: under which shape does this workload go fastest, and what does that
+cost in per-request latency.
+
+**The fan-out is the whole point.** A partition gets a fraction of the card, so
+a benchmark that loads one partition and ignores the rest measures a fraction of
+the card. Measured on this node, one MI355X in ``CPX`` presents eight 32-CU
+partitions -- a single-partition run reports roughly an eighth of the card's
+throughput, making ``CPX`` look catastrophic when in aggregate it may well win.
+Every figure here is therefore the sum over a mode's partitions with all of them
+loaded at once, and a mode whose partitions cannot all be loaded is reported as
+unmeasured rather than as slow.
+
+**Why the privileged set lives here and not in the optimizer.** Changing the
+mode is a card-wide privileged operation that evicts every process holding a GPU
+context. That is reasonable between benchmarks in an operator-run script and
+unreasonable inside an optimization loop that also runs agent-authored code, so
+``hyperloom`` itself only ever *reads* the mode -- see
+``hyperloom.common.gpu_partition`` -- and refuses at launch any session whose
+streams will not fit one partition. This script is the other half of that
+split: the boundary that establishes the shape the optimizer then treats as
+fixed for the whole session.
+
+**Device indices are not portable between tools, so partitions are matched on CU
+count.** Measured on an 8-card MI355X node with card 0 in ``CPX``: ``amd-smi``
+orders by PCI address and calls the eight partitions devices 0-7, while HSA/HIP
+enumerates whole cards first and calls them devices 7-14. A device list computed
+with one tool and handed to the other is wrong, and wrong invisibly -- the
+benchmark runs to completion, on the wrong silicon. Partitions are selected here
+the way :func:`~hyperloom.common.gpu_partition.partition_device_predicate`
+documents: by matching the expected CU count, in the HIP index space the
+benchmark itself will use, and narrowed to the card being swept so that an
+identical CU count on a neighbouring card cannot be mistaken for a partition.
+
+Only the target card is repartitioned; the rest of the node is left alone. That
+keeps total silicon constant across the sweep, which is what makes the modes
+comparable to each other.
+
+Usage::
+
+ # print the plan, touch nothing
+ python3 scripts/partition_mode_sweep.py --benchmark-config bench.yaml --dry-run
+
+ # sweep the modes the card reports it supports
+ python3 scripts/partition_mode_sweep.py \\
+ --benchmark-config /path/to/benchmark.yaml \\
+ --modes SPX,DPX,QPX,CPX \\
+ --output-dir /shared/partition-sweep
+
+ # arbitrary workload; {device} and {output_dir} are substituted per partition
+ python3 scripts/partition_mode_sweep.py \\
+ --benchmark-command 'my_bench --gpu {device} --out {output_dir}' \\
+ --per-stream-gib 20.7
+
+Exit codes:
+
+ 0 the sweep completed and a winner was reported
+ 1 the sweep ran but no mode produced a valid measurement
+ 2 the request was refused before anything was changed
+ 3 the sweep finished but the card could not be restored to its entry mode
+ 4 the sweep stopped on an error it does not model, after reporting whatever
+ it had already measured and restoring the card
+
+Every path out of a started sweep goes through the restore and the report,
+including an unexpected exception, so ``3`` stays reachable and the modes
+measured before a failure are never lost to it.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shlex
+import signal
+import subprocess # nosec B404 - fixed argv, never a shell.
+import sys
+import time
+import traceback
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Sequence
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+_SRC = _REPO_ROOT / "src"
+if str(_SRC) not in sys.path:
+ sys.path.insert(0, str(_SRC))
+
+from hyperloom.common.coerce import to_float # noqa: E402
+from hyperloom.common.gpu_partition import ( # noqa: E402
+ MODE_PARTITION_COUNTS,
+ PARTITION_COUNT_ENV,
+ PARTITION_CU_ENV,
+ PARTITION_MODE_ENV,
+ PARTITION_STREAMS_ENV,
+ PARTITION_TOTAL_STREAMS_ENV,
+ PartitionError,
+ PartitionLayout,
+ fits_in_partition,
+ layout_for,
+ parse_mode,
+ partition_device_predicate,
+ read_device_gib,
+)
+
+#: Read timeout for an ``amd-smi`` query. Generous because a card mid-transition
+#: answers slowly, and a spurious timeout here reads as a hardware fault.
+_READ_TIMEOUT_S = 30.0
+
+#: Timeout for the set itself.
+_SET_TIMEOUT_S = 120.0
+
+#: How long to keep retrying a set the card refuses because it still holds
+#: processes. The condition clears on its own once the previous benchmark's
+#: contexts are torn down, so the correct response is to wait rather than fail:
+#: this script's own predecessor run is the most likely holder.
+_DRAIN_TIMEOUT_S = 120.0
+_DRAIN_POLL_S = 2.0
+
+#: Pause between a successful set and trusting the enumeration. The set returns
+#: before the new devices appear; measured here at one to two seconds, so this
+#: is that with margin. Without it the first device scan after a set sees the
+#: old topology and the run is attributed to the wrong mode.
+_SETTLE_S = 5.0
+
+#: Substrings marking a refusal as "busy, try again" rather than "wrong".
+#: Matched narrowly so that permanent failures -- unknown mode, missing binary,
+#: denied permission -- fail fast instead of retrying for two minutes.
+#: ``AMDSMI_STATUS_BUSY`` is the authoritative one; the human-readable half of
+#: the message is not stable enough to match on alone.
+_BUSY_MARKERS = ("amdsmi_status_busy", "device busy", "resident process", "try again")
+
+#: What ``amd-smi process`` puts in ``process_list`` when a GPU is idle: a bare
+#: string where a caller would reasonably expect a list of dicts. Parsed
+#: explicitly because treating it as a process would make every sweep refuse to
+#: start on a perfectly free node.
+_NO_PROCESS_MARKER = "no running processes"
+
+#: Throughput fields summed across partitions, in preference order for the
+#: headline figure. Named to match ``benchmark_result.extract_benchmark_measurement``.
+_THROUGHPUT_FIELDS = ("output_throughput", "total_token_throughput", "request_throughput")
+
+#: Per-request latency fields, in the same preference order the optimizer's
+#: latency budget uses.
+_LATENCY_FIELDS = ("e2el_mean_ms", "mean_e2el_ms", "e2el_ms")
+
+
+class SweepError(RuntimeError):
+ """Raised when the sweep cannot proceed and nothing has been left changed."""
+
+
+@dataclass(frozen=True)
+class HsaAgent:
+ """One GPU agent as HSA/HIP enumerates it.
+
+ Attributes:
+ index: Position among GPU agents, which is the index
+ ``ROCR_VISIBLE_DEVICES`` and HIP both use.
+ cu: Compute units the agent reports.
+ bus: PCI bus byte, identifying the physical card. Partitions of one card
+ share it and differ only in PCI function.
+ device: PCI device number.
+ function: PCI function number.
+ """
+
+ index: int
+ cu: int
+ bus: int
+ device: int
+ function: int
+
+ @property
+ def bdf(self) -> str:
+ """Human-readable ``bus:device.function``."""
+ return f"{self.bus:02x}:{self.device:02x}.{self.function}"
+
+
+@dataclass
+class PartitionRun:
+ """One benchmark process on one partition."""
+
+ partition_index: int
+ hsa_device: int
+ output_dir: Path
+ returncode: int | None = None
+ measurement: dict[str, Any] = field(default_factory=dict)
+ error: str = ""
+
+ @property
+ def ok(self) -> bool:
+ return self.returncode == 0 and not self.error and bool(self.measurement)
+
+
+@dataclass
+class ModeResult:
+ """What one mode produced, or why it produced nothing."""
+
+ mode: str
+ layout: PartitionLayout | None = None
+ runs: list[PartitionRun] = field(default_factory=list)
+ skipped: str = ""
+ error: str = ""
+
+ @property
+ def measured(self) -> bool:
+ """Whether every partition returned a usable measurement.
+
+ Deliberately all-or-nothing. A mode with six of eight partitions
+ reporting is not a slower mode, it is an unmeasured one, and summing the
+ six would understate it by a quarter while looking like a result.
+ """
+ return bool(self.runs) and all(run.ok for run in self.runs)
+
+ def total(self, fields: Sequence[str] = _THROUGHPUT_FIELDS) -> float | None:
+ """Sum the first available throughput field across partitions."""
+ for name in fields:
+ values = [to_float(run.measurement.get(name)) for run in self.runs]
+ if values and all(v is not None for v in values):
+ return sum(v for v in values if v is not None)
+ return None
+
+ def worst_latency_ms(self) -> float | None:
+ """The slowest partition's mean latency.
+
+ The worst partition rather than the average of them, because a request
+ landing on the slow one is not consoled by the mean.
+ """
+ for name in _LATENCY_FIELDS:
+ values = [to_float(run.measurement.get(name)) for run in self.runs]
+ present = [v for v in values if v is not None]
+ if present and len(present) == len(values):
+ return max(present)
+ return None
+
+
+# ------------------------------------------------------------------ amd-smi
+
+
+def _sudo_prefix(use_sudo: bool) -> list[str]:
+ """Return the privilege prefix for a mutating call.
+
+ ``-n`` matters: an unattended sweep that stops at an interactive password
+ prompt hangs until someone notices, with nothing in the output saying why.
+ """
+ return ["sudo", "-n"] if use_sudo else []
+
+
+def _amd_smi_json(args: Sequence[str], timeout_s: float = _READ_TIMEOUT_S, *, sudo: bool = False) -> object:
+ """Run an ``amd-smi`` subcommand with ``--json`` and parse its output.
+
+ Args:
+ args: Subcommand and flags.
+ timeout_s: Per-call timeout.
+ sudo: Route through ``sudo -n``. Needed for the profile query, which
+ degrades every field to ``"N/A"`` rather than failing when it lacks
+ privilege.
+
+ Raises:
+ SweepError: If ``amd-smi`` is missing, fails, times out, or returns
+ output that is not JSON.
+ """
+ cmd = [*_sudo_prefix(sudo), "amd-smi", *args, "--json"]
+ try:
+ proc = subprocess.run( # nosec B603 B607 - fixed argv, no shell.
+ cmd,
+ capture_output=True,
+ text=True,
+ timeout=timeout_s,
+ check=False,
+ )
+ except FileNotFoundError as exc:
+ raise SweepError("amd-smi not found; this sweep needs it on PATH") from exc
+ except subprocess.TimeoutExpired as exc:
+ raise SweepError(f"amd-smi {' '.join(args)} timed out after {timeout_s:.0f}s") from exc
+ if proc.returncode != 0:
+ raise SweepError(f"amd-smi {' '.join(args)} failed ({proc.returncode}): {proc.stderr.strip()}")
+ try:
+ return json.loads(proc.stdout)
+ except ValueError as exc:
+ raise SweepError(f"amd-smi {' '.join(args)} returned unparseable JSON") from exc
+
+
+def _is_busy(detail: str) -> bool:
+ """Whether a failed set was refused because the card was still occupied."""
+ low = (detail or "").lower()
+ return any(marker in low for marker in _BUSY_MARKERS)
+
+
+def current_modes(payload: object) -> dict[int, str]:
+ """Extract GPU id to mode from an ``amd-smi partition --json`` payload."""
+ rows: list[dict] = []
+ if isinstance(payload, dict):
+ raw = payload.get("current_partition")
+ if isinstance(raw, list):
+ rows = [r for r in raw if isinstance(r, dict)]
+ elif isinstance(payload, list):
+ rows = [r for r in payload if isinstance(r, dict)]
+ modes: dict[int, str] = {}
+ for row in rows:
+ try:
+ gpu_id = int(row.get("gpu_id"))
+ except (TypeError, ValueError):
+ continue
+ mode = str(row.get("accelerator_type") or "").strip().upper().rstrip("*")
+ if mode and mode != "N/A":
+ modes[gpu_id] = mode
+ return modes
+
+
+def read_mode(gpu_id: int) -> str:
+ """Read one card's live compute-partition mode."""
+ modes = current_modes(_amd_smi_json(["partition"]))
+ if gpu_id not in modes:
+ raise SweepError(f"amd-smi reported no compute-partition state for GPU {gpu_id}")
+ return modes[gpu_id]
+
+
+def supported_modes(payload: object) -> tuple[str, ...]:
+ """Extract the modes a card reports it can enter, in profile order.
+
+ An empty result means "the card did not say", not "the card supports
+ nothing": the profile query returns ``"N/A"`` for everything when run
+ without privilege, and those two cases must not collapse into one answer.
+ """
+ rows: list[dict] = []
+ if isinstance(payload, dict):
+ raw = payload.get("partition_profiles")
+ if isinstance(raw, list):
+ rows = [r for r in raw if isinstance(r, dict)]
+ elif isinstance(payload, list):
+ rows = [r for r in payload if isinstance(r, dict)]
+ seen: list[str] = []
+ for row in rows:
+ # The table is sparse: a profile's first row names it, and the rows
+ # after it continue the same profile with its other resources under
+ # blank identity fields. Only named rows describe a profile.
+ mode = str(row.get("accelerator_type") or "").strip().upper().rstrip("*")
+ if mode and mode != "N/A" and mode in MODE_PARTITION_COUNTS and mode not in seen:
+ seen.append(mode)
+ return tuple(seen)
+
+
+def _live_processes(listing: object, *, gpu_id: int) -> int:
+ """Count the entries of one card's ``process_list`` that are real processes.
+
+ The idle sentinel turns up at either depth -- as the whole ``process_list``,
+ or as the lone entry's ``process_info`` -- and a process turns up either
+ wrapped in ``process_info`` or as the entry itself. All four were observed.
+ Anything else raises, because a shape this function does not recognise is
+ indistinguishable from an idle card once it has been counted as zero.
+ """
+ if isinstance(listing, str):
+ entries: list[object] = [listing]
+ elif isinstance(listing, list):
+ entries = list(listing)
+ else:
+ raise SweepError(
+ f"amd-smi process gave GPU {gpu_id} a process_list of type "
+ f"{type(listing).__name__}; expected a list or the idle string"
+ )
+ live = 0
+ for entry in entries:
+ info = entry.get("process_info", entry) if isinstance(entry, dict) else entry
+ if isinstance(info, str):
+ # Only the known idle sentinel means idle. An unrecognised string
+ # counts as a process: over-counting refuses a sweep, under-counting
+ # evicts somebody's work.
+ if _NO_PROCESS_MARKER not in info.lower():
+ live += 1
+ elif isinstance(info, dict):
+ if str(info.get("name") or "").strip():
+ live += 1
+ else:
+ raise SweepError(
+ f"amd-smi process gave GPU {gpu_id} a process entry of type "
+ f"{type(info).__name__}; expected a mapping or a string"
+ )
+ return live
+
+
+def resident_processes(payload: object) -> dict[int, int]:
+ """Count real processes holding a context on each GPU.
+
+ ``amd-smi`` reports an idle GPU as a ``process_list`` holding the *string*
+ ``"No running processes detected"`` rather than an empty list, so a naive
+ length check finds one process on every idle card and this sweep would
+ refuse to start on a free node.
+
+ Every departure from the documented shape raises instead of being skipped.
+ This count is the only thing between a payload this parser does not
+ understand and an ``amd-smi set`` that evicts whatever is running, and a
+ parser that answers ``{}`` for a payload it cannot read reports a busy node
+ as a free one -- the single wrong answer here that destroys work. Refusing
+ costs an operator one ``--allow-busy``; guessing costs somebody a job.
+
+ Raises:
+ SweepError: If the payload is not a list of per-GPU rows, or a row is
+ missing ``gpu`` or ``process_list``, or either field has a type
+ this parser does not model.
+ """
+ if not isinstance(payload, list):
+ raise SweepError(
+ f"amd-smi process returned {type(payload).__name__}, not the expected list of "
+ f"per-GPU rows, so which cards are in use cannot be read from it"
+ )
+ counts: dict[int, int] = {}
+ for row in payload:
+ if not isinstance(row, dict):
+ raise SweepError(f"amd-smi process returned a {type(row).__name__} where a GPU row was expected")
+ missing = [key for key in ("gpu", "process_list") if key not in row]
+ if missing:
+ raise SweepError(f"amd-smi process returned a GPU row without {' or '.join(missing)}")
+ try:
+ gpu_id = int(row["gpu"])
+ except (TypeError, ValueError) as exc:
+ raise SweepError(f"amd-smi process reported a GPU id of {row['gpu']!r}, which is not a number") from exc
+ counts[gpu_id] = _live_processes(row["process_list"], gpu_id=gpu_id)
+ return counts
+
+
+def card_bus(payload: object, gpu_id: int) -> int:
+ """Read the PCI bus byte of one card from an ``amd-smi list`` payload.
+
+ The bus identifies the physical card and does not change when it is
+ repartitioned, so it is captured once and used afterwards to tell this
+ card's partitions from an identically-shaped neighbour.
+ """
+ rows = payload if isinstance(payload, list) else []
+ for row in rows:
+ if not isinstance(row, dict):
+ continue
+ try:
+ if int(row.get("gpu")) != gpu_id:
+ continue
+ except (TypeError, ValueError):
+ continue
+ bdf = str(row.get("bdf") or "")
+ match = re.search(r"(?:[0-9a-fA-F]{4}:)?([0-9a-fA-F]{2}):([0-9a-fA-F]{2})\.(\d)", bdf)
+ if match:
+ return int(match.group(1), 16)
+ raise SweepError(f"amd-smi list reported no PCI address for GPU {gpu_id}")
+
+
+def set_mode(
+ gpu_id: int,
+ mode: str,
+ *,
+ sudo: bool,
+ drain_timeout_s: float = _DRAIN_TIMEOUT_S,
+ settle_s: float = _SETTLE_S,
+) -> str:
+ """Set one card's compute-partition mode and verify it took effect.
+
+ The verification is the point. ``amd-smi set`` reports success for a change
+ that has only been staged, and it exits zero on some permission failures, so
+ a caller trusting the exit code goes on to benchmark the old topology while
+ labelling the results with the new mode -- a wrong number with a reassuring
+ log line above it.
+
+ Args:
+ gpu_id: Card to reconfigure.
+ mode: Target mode.
+ sudo: Whether to route the set through ``sudo -n``.
+ drain_timeout_s: How long to keep retrying while the card reports
+ resident processes. Zero fails on the first refusal.
+ settle_s: Pause before reading back, so the new devices have appeared.
+
+ Returns:
+ The mode read back from the card, equal to ``mode`` on success.
+
+ Raises:
+ SweepError: If the mode is unknown, the set fails, or the read-back
+ disagrees with what was asked for.
+ """
+ canonical = parse_mode(mode)
+ if canonical not in MODE_PARTITION_COUNTS:
+ raise SweepError(f"unknown compute-partition mode {mode!r}")
+ if read_mode(gpu_id) == canonical:
+ return canonical
+
+ cmd = [*_sudo_prefix(sudo), "amd-smi", "set", "-g", str(gpu_id), "--compute-partition", canonical]
+ deadline = time.monotonic() + max(0.0, drain_timeout_s)
+ waited = False
+ while True:
+ try:
+ proc = subprocess.run( # nosec B603 B607 - fixed argv, no shell.
+ cmd,
+ capture_output=True,
+ text=True,
+ timeout=_SET_TIMEOUT_S,
+ # Some builds ask for confirmation; answer rather than block.
+ input="Y\n",
+ check=False,
+ )
+ except FileNotFoundError as exc:
+ raise SweepError("amd-smi not found; this sweep needs it on PATH") from exc
+ except subprocess.TimeoutExpired as exc:
+ raise SweepError(f"setting {canonical} on GPU {gpu_id} timed out") from exc
+ if proc.returncode == 0:
+ break
+ detail = (proc.stderr or proc.stdout).strip()
+ if _is_busy(detail) and time.monotonic() < deadline:
+ if not waited:
+ print(f" GPU {gpu_id} busy; waiting for it to drain before setting {canonical}")
+ waited = True
+ time.sleep(_DRAIN_POLL_S)
+ continue
+ hint = (
+ f" The card still held processes after {drain_timeout_s:.0f}s."
+ if _is_busy(detail)
+ else " A card with resident processes refuses to repartition; stop them first."
+ )
+ raise SweepError(f"setting {canonical} on GPU {gpu_id} failed ({proc.returncode}): {detail}.{hint}")
+
+ time.sleep(max(0.0, settle_s))
+ observed = read_mode(gpu_id)
+ if observed != canonical:
+ raise SweepError(
+ f"GPU {gpu_id} reports {observed} after being set to {canonical}. The command "
+ f"returned success but the mode did not change, so any measurement now would be "
+ f"attributed to the wrong mode. Most often this is a permission failure that "
+ f"amd-smi did not report as one -- try --sudo."
+ )
+ return observed
+
+
+# ------------------------------------------------------- HIP enumeration
+
+
+def parse_hsa_agents(text: str) -> tuple[HsaAgent, ...]:
+ """Parse ``rocminfo`` output into GPU agents in HIP index order.
+
+ HIP indices are what the benchmark will use, and they are not ``amd-smi``
+ indices: on this node under ``CPX``, ``amd-smi`` calls card 0's partitions
+ devices 0-7 while HSA calls them 7-14, because HSA enumerates whole cards
+ first. Reading the order from HSA is the only way to hand the benchmark a
+ device it agrees with.
+
+ Within an agent block ``BDFID`` appears *before* ``Compute Unit``, so a
+ line-at-a-time parser that prints on ``BDFID`` attributes every agent the
+ previous one's CU count. Fields are therefore collected per block and only
+ interpreted once the block ends.
+ """
+ agents: list[HsaAgent] = []
+ block: dict[str, str] = {}
+ index = 0
+
+ def _flush(block: dict[str, str]) -> None:
+ nonlocal index
+ if block.get("Device Type") != "GPU":
+ return
+ try:
+ cu = int(block.get("Compute Unit", ""))
+ bdfid = int(block.get("BDFID", ""))
+ except ValueError:
+ return
+ agents.append(
+ HsaAgent(
+ index=index,
+ cu=cu,
+ bus=(bdfid >> 8) & 0xFF,
+ device=(bdfid >> 3) & 0x1F,
+ function=bdfid & 0x7,
+ )
+ )
+ index += 1
+
+ for line in text.splitlines():
+ if re.match(r"^Agent \d+", line):
+ _flush(block)
+ block = {}
+ continue
+ match = re.match(r"\s+(Device Type|Compute Unit|BDFID):\s*(.+?)\s*$", line)
+ if match and match.group(1) not in block:
+ block[match.group(1)] = match.group(2)
+ _flush(block)
+ return tuple(agents)
+
+
+def read_hsa_agents() -> tuple[HsaAgent, ...]:
+ """Run ``rocminfo`` and parse its GPU agents."""
+ try:
+ proc = subprocess.run( # nosec B603 B607 - fixed argv, no shell.
+ ["rocminfo"],
+ capture_output=True,
+ text=True,
+ timeout=_READ_TIMEOUT_S,
+ check=False,
+ )
+ except FileNotFoundError as exc:
+ raise SweepError("rocminfo not found; it is how HIP device order is read") from exc
+ except subprocess.TimeoutExpired as exc:
+ raise SweepError("rocminfo timed out") from exc
+ if proc.returncode != 0:
+ raise SweepError(f"rocminfo failed ({proc.returncode}): {proc.stderr.strip()}")
+ return parse_hsa_agents(proc.stdout)
+
+
+def partition_cu_on_bus(agents: Sequence[HsaAgent], bus: int) -> int:
+ """The CU count the swept card's devices report, once they agree.
+
+ Every GPU device on the swept card's bus is one of its partitions, and a
+ mode's partitions are identical, so a disagreement means the enumeration was
+ read while the card was still transitioning. Nothing measured against a
+ half-applied topology is worth keeping, so that is an error rather than a
+ figure to pick from.
+ """
+ counts = {a.cu for a in agents if a.bus == bus}
+ if not counts:
+ raise SweepError(f"HIP reports no GPU on bus {bus:02x}")
+ if len(counts) > 1:
+ raise SweepError(
+ f"devices on bus {bus:02x} report different CU counts ({sorted(counts)}), which a "
+ f"settled card does not do -- the enumeration was read mid-transition"
+ )
+ return counts.pop()
+
+
+def card_total_gib(device_gib: float | None, device_mode: str) -> float | None:
+ """Scale one device's HBM back up to the whole card's.
+
+ The reading is per device, so under a split mode it is one partition's share
+ and has to be multiplied by the partition count to describe the card. Doing
+ that once at entry means each mode's per-partition memory is a division of
+ the same known total, rather than a fresh probe whose device-index mapping
+ changes with every set.
+ """
+ if not device_gib or device_gib <= 0:
+ return None
+ return device_gib * MODE_PARTITION_COUNTS.get(parse_mode(device_mode), 1)
+
+
+def select_partition_devices(
+ agents: Sequence[HsaAgent],
+ layout: PartitionLayout,
+ *,
+ bus: int,
+) -> tuple[int, ...]:
+ """Return the HIP indices of the swept card's partitions.
+
+ Selection is by CU count, never by index, and is narrowed to one PCI bus.
+ Both halves matter: the CU match is what distinguishes a partition from a
+ whole card, and the bus match is what stops seven untouched 256-CU
+ neighbours being mistaken for partitions when the mode under test is
+ ``SPX``, whose "partition" is a whole card.
+
+ Args:
+ agents: GPU agents in HIP order.
+ layout: The mode's expected shape.
+ bus: PCI bus of the card being swept.
+
+ Returns:
+ HIP indices, ascending, one per partition.
+
+ Raises:
+ SweepError: If the number found is not the number the mode implies,
+ which is what a set that did not really take looks like from here.
+ """
+ is_partition = partition_device_predicate(layout.cu_per_partition)
+ found = tuple(a.index for a in agents if a.bus == bus and is_partition(a.cu))
+ if len(found) != layout.partitions:
+ on_bus = [f"{a.bdf}={a.cu}CU" for a in agents if a.bus == bus]
+ raise SweepError(
+ f"{layout.mode} implies {layout.partitions} partitions of "
+ f"{layout.cu_per_partition} CU on bus {bus:02x}, but HIP reports "
+ f"{len(found)}: {', '.join(on_bus) or 'nothing on that bus'}. "
+ f"Benchmarking now would measure the wrong silicon."
+ )
+ return found
+
+
+# ------------------------------------------------------------- the fan-out
+
+
+def build_partition_command(
+ template: Sequence[str],
+ *,
+ device: int,
+ output_dir: Path,
+ layout: PartitionLayout,
+ partition_index: int,
+) -> list[str]:
+ """Substitute per-partition values into a benchmark command template.
+
+ Substitution is per already-split token, so a path containing a space cannot
+ turn into two arguments and no shell is involved at any point.
+ """
+ values = {
+ "device": str(device),
+ "output_dir": str(output_dir),
+ "mode": layout.mode,
+ "partitions": str(layout.partitions),
+ "partition_index": str(partition_index),
+ "cu": str(layout.cu_per_partition),
+ }
+ out: list[str] = []
+ for token in template:
+ for key, value in values.items():
+ token = token.replace("{" + key + "}", value)
+ out.append(token)
+ return out
+
+
+def magpie_command(benchmark_config: Path, python_exe: str = sys.executable) -> list[str]:
+ """The in-tree benchmark invocation, with placeholders for the fan-out."""
+ return [
+ python_exe,
+ "-m",
+ "Magpie",
+ "-v",
+ "benchmark",
+ "--benchmark-config",
+ str(benchmark_config),
+ "--output-dir",
+ "{output_dir}",
+ "--run-mode",
+ "local",
+ ]
+
+
+def partition_env(
+ base: dict[str, str],
+ layout: PartitionLayout,
+ *,
+ device: int,
+ streams_per_partition: int,
+) -> dict[str, str]:
+ """Environment for one partition's benchmark process.
+
+ Pins the process to its partition and publishes the session shape in the
+ same variables the optimizer publishes, so a benchmark entrypoint written
+ against that contract behaves identically whether it was launched by a
+ session or by this sweep.
+
+ ``HIP_VISIBLE_DEVICES`` and ``CUDA_VISIBLE_DEVICES`` are removed rather than
+ set. Leaving an inherited one alongside ``ROCR_VISIBLE_DEVICES`` means two
+ masks apply in sequence, and the second is interpreted as an index into the
+ first -- so a stale ``HIP_VISIBLE_DEVICES=0`` silently redirects every
+ partition's work onto whichever device the first mask selected.
+ """
+ env = dict(base)
+ env["ROCR_VISIBLE_DEVICES"] = str(device)
+ env.pop("HIP_VISIBLE_DEVICES", None)
+ env.pop("CUDA_VISIBLE_DEVICES", None)
+ env[PARTITION_MODE_ENV] = layout.mode
+ env[PARTITION_COUNT_ENV] = str(layout.partitions)
+ env[PARTITION_CU_ENV] = str(layout.cu_per_partition)
+ env[PARTITION_STREAMS_ENV] = str(streams_per_partition)
+ env[PARTITION_TOTAL_STREAMS_ENV] = str(layout.partitions * streams_per_partition)
+ return env
+
+
+def read_measurement(output_dir: Path) -> dict[str, Any]:
+ """Parse one partition's benchmark report using the in-tree extractor.
+
+ Imported here rather than at module scope: the extractor pulls in the
+ orchestrator, and the pure logic in this file -- command construction,
+ device selection, mode gating -- is worth testing without that weight.
+ """
+ reports = sorted(output_dir.rglob("benchmark_report.json"))
+ if not reports:
+ return {}
+ try:
+ payload = json.loads(reports[-1].read_text(encoding="utf-8"))
+ except (OSError, ValueError):
+ return {}
+
+ import hyperloom.orchestrator.actions.executors._grid_runner # noqa: F401
+ from hyperloom.orchestrator.actions.executors.benchmark_result import extract_benchmark_measurement
+
+ return extract_benchmark_measurement(payload if isinstance(payload, dict) else None, workspace=output_dir)
+
+
+def run_mode(
+ layout: PartitionLayout,
+ devices: Sequence[int],
+ *,
+ template: Sequence[str],
+ output_root: Path,
+ streams_per_partition: int,
+ timeout_s: float,
+ base_env: dict[str, str] | None = None,
+) -> list[PartitionRun]:
+ """Run the benchmark on every partition at once and collect each result.
+
+ All partitions are launched before any is waited on, which is the only
+ arrangement that measures what partitioning is for. Running them in sequence
+ would measure one partition at a time on an otherwise idle card and report a
+ fraction of the mode's throughput.
+ """
+ base = dict(os.environ if base_env is None else base_env)
+ runs: list[PartitionRun] = []
+ procs: list[tuple[PartitionRun, subprocess.Popen[str] | None, Any]] = []
+
+ for position, device in enumerate(devices):
+ out_dir = output_root / layout.mode.lower() / f"partition{position}"
+ out_dir.mkdir(parents=True, exist_ok=True)
+ run = PartitionRun(partition_index=position, hsa_device=device, output_dir=out_dir)
+ runs.append(run)
+ cmd = build_partition_command(
+ template,
+ device=device,
+ output_dir=out_dir,
+ layout=layout,
+ partition_index=position,
+ )
+ env = partition_env(base, layout, device=device, streams_per_partition=streams_per_partition)
+ log = None
+ try:
+ log = (out_dir / "benchmark.log").open("w", encoding="utf-8")
+ proc = subprocess.Popen( # nosec B603 - argv built from a template, no shell.
+ cmd,
+ env=env,
+ stdout=log,
+ stderr=subprocess.STDOUT,
+ text=True,
+ )
+ except (OSError, ValueError) as exc:
+ if log is not None:
+ log.close()
+ run.error = f"could not launch: {exc}"
+ procs.append((run, None, None))
+ continue
+ procs.append((run, proc, log))
+
+ deadline = time.monotonic() + max(1.0, timeout_s)
+ for run, proc, log in procs:
+ if proc is None:
+ continue
+ remaining = max(1.0, deadline - time.monotonic())
+ try:
+ run.returncode = proc.wait(timeout=remaining)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ proc.wait(timeout=30)
+ run.error = f"exceeded the {timeout_s:.0f}s budget"
+ run.returncode = -1
+ finally:
+ if log is not None:
+ log.close()
+ if run.error:
+ continue
+ if run.returncode != 0:
+ run.error = f"exited {run.returncode}; see {run.output_dir / 'benchmark.log'}"
+ continue
+ run.measurement = read_measurement(run.output_dir)
+ if not run.measurement:
+ run.error = f"no benchmark_report.json under {run.output_dir}"
+ return runs
+
+
+# ------------------------------------------------------------- reporting
+
+
+def render(results: Sequence[ModeResult], *, entry_mode: str) -> list[str]:
+ """Render the comparison an operator reads to choose a mode."""
+ lines = ["", "=" * 78, "Compute-partition sweep", "=" * 78, ""]
+ measured = [r for r in results if r.measured and r.total() is not None]
+ baseline = next((r for r in measured if r.mode == entry_mode), None)
+ best = max(measured, key=lambda r: r.total() or 0.0, default=None)
+
+ lines.append(f"{'mode':<6} {'parts':>5} {'CU':>5} {'throughput':>13} {'vs entry':>9} {'worst lat':>11}")
+ lines.append("-" * 78)
+ for result in results:
+ if not result.measured or result.total() is None:
+ why = result.skipped or result.error or "not measured"
+ lines.append(f"{result.mode:<6} {'-':>5} {'-':>5} {'-':>13} {'-':>9} {'-':>11} {why}")
+ continue
+ layout = result.layout
+ total = result.total() or 0.0
+ rel = ""
+ if baseline is not None and (baseline.total() or 0.0) > 0:
+ rel = f"{total / (baseline.total() or 1.0):.2f}x"
+ latency = result.worst_latency_ms()
+ lines.append(
+ f"{result.mode:<6} {layout.partitions if layout else '-':>5} "
+ f"{layout.cu_per_partition if layout else '-':>5} "
+ f"{total:>13.1f} {rel:>9} {(f'{latency:.1f} ms' if latency else '-'):>11}"
+ )
+ lines.append("")
+
+ if best is None:
+ lines.append("No mode produced a measurement, so there is nothing to choose between.")
+ lines.append("")
+ return lines
+
+ lines.append(f"Fastest measured mode: {best.mode} at {best.total():.1f} aggregate throughput.")
+ if baseline is not None and best.mode != baseline.mode:
+ gain = (best.total() or 0.0) / (baseline.total() or 1.0)
+ best_lat = best.worst_latency_ms()
+ base_lat = baseline.worst_latency_ms()
+ lines.append(
+ f"That is {gain:.2f}x the {baseline.mode} the card was in when this started."
+ + (
+ f" Worst-partition latency moves from {base_lat:.1f} ms to {best_lat:.1f} ms."
+ if best_lat and base_lat
+ else ""
+ )
+ )
+ lines.append("")
+ lines.append(
+ "Throughput is the sum over a mode's partitions, all loaded together; a mode is\n"
+ "reported only when every one of its partitions returned a measurement. The\n"
+ "latency column is the worst partition's mean, not the average of them.\n"
+ "\n"
+ "This chooses a mode. It does not tune one: run the optimizer in the winning\n"
+ "mode to do that, and pass --compute-partition-mode so the session refuses to\n"
+ "start if the card is not actually in it."
+ )
+ lines.append("")
+ return lines
+
+
+def summary_json(results: Sequence[ModeResult], *, entry_mode: str, gpu_id: int) -> dict[str, Any]:
+ """Machine-readable form of the same comparison."""
+ return {
+ "gpu": gpu_id,
+ "entry_mode": entry_mode,
+ "modes": [
+ {
+ "mode": r.mode,
+ "measured": r.measured,
+ "skipped": r.skipped,
+ "error": r.error,
+ "partitions": r.layout.partitions if r.layout else None,
+ "cu_per_partition": r.layout.cu_per_partition if r.layout else None,
+ "gib_per_partition": r.layout.gib_per_partition if r.layout else None,
+ "aggregate_throughput": r.total(),
+ "worst_partition_latency_ms": r.worst_latency_ms(),
+ "partition_runs": [
+ {
+ "partition_index": run.partition_index,
+ "hsa_device": run.hsa_device,
+ "returncode": run.returncode,
+ "error": run.error,
+ "output_throughput": to_float(run.measurement.get("output_throughput")),
+ "e2el_mean_ms": to_float(run.measurement.get("e2el_mean_ms")),
+ }
+ for run in r.runs
+ ],
+ }
+ for r in results
+ ],
+ }
+
+
+# ------------------------------------------------------------------- main
+
+
+def resolve_modes(requested: str | None, available: Sequence[str]) -> tuple[str, ...]:
+ """Parse the requested mode list, defaulting to what the card reports.
+
+ Raises:
+ SweepError: If a requested mode is not a mode, or the card says it
+ cannot enter it. A card that reported no profiles at all is not
+ second-guessed -- the request stands and the set will judge it.
+ """
+ if not (requested or "").strip():
+ if available:
+ return tuple(available)
+ raise SweepError(
+ "the card reported no partition profiles, so there is no mode list to "
+ "default to; name the modes with --modes (and try --sudo, since the "
+ "profile query needs privilege)"
+ )
+ modes: list[str] = []
+ for raw in str(requested).replace(",", " ").split():
+ # parse_mode refuses an unknown spelling with PartitionError. Converted
+ # here so that every way of mistyping --modes leaves by the same door as
+ # the other usage errors, rather than as a traceback.
+ try:
+ mode = parse_mode(raw)
+ except PartitionError as exc:
+ raise SweepError(f"unknown compute-partition mode {raw!r}") from exc
+ if mode not in MODE_PARTITION_COUNTS:
+ raise SweepError(f"unknown compute-partition mode {raw!r}")
+ if mode not in modes:
+ modes.append(mode)
+ if available:
+ rejected = [m for m in modes if m not in available]
+ if rejected:
+ raise SweepError(
+ f"the card does not report support for {', '.join(rejected)}; it lists {', '.join(available)}"
+ )
+ return tuple(modes)
+
+
+def _restore_entry_mode(gpu_id: int, entry_mode: str, *, sudo: bool) -> bool:
+ """Put the card back in the mode it was found in. True if it could not be.
+
+ Raises nothing, because it is called from a ``finally``: an exception here
+ would replace whatever sent the sweep into the restore and would take the
+ report down with it, losing both the diagnosis and the modes already
+ measured.
+
+ A read-back that fails is treated as "mode unknown" and the set is attempted
+ regardless. The alternative is skipping the restore because the check that
+ would have proved it necessary is the thing that broke, which leaves a card
+ in a shape nobody asked for.
+ """
+ try:
+ if read_mode(gpu_id) == entry_mode:
+ return False
+ except Exception as exc:
+ print(f"note: could not read GPU {gpu_id}'s mode ({exc}); attempting the restore anyway")
+ try:
+ print(f"\nrestoring {entry_mode} on GPU {gpu_id}")
+ set_mode(gpu_id, entry_mode, sudo=sudo)
+ except Exception as exc:
+ print(
+ f"ERROR: could not restore {entry_mode} on GPU {gpu_id}: {exc}\n"
+ f"The card is NOT in the mode it started in. Anything that runs on it now "
+ f"will be measured under a shape nobody asked for.",
+ file=sys.stderr,
+ )
+ return True
+ return False
+
+
+def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
+ ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+ ap.add_argument("--gpu", type=int, default=0, help="Card to repartition and sweep (default 0).")
+ ap.add_argument(
+ "--modes",
+ default="",
+ help="Comma-separated modes to try, in order. Default: every mode the card reports.",
+ )
+ source = ap.add_mutually_exclusive_group()
+ source.add_argument("--benchmark-config", type=Path, help="Run the in-tree Magpie benchmark with this config.")
+ source.add_argument(
+ "--benchmark-command",
+ help="Command to run per partition. {device} {output_dir} {mode} {partitions} "
+ "{partition_index} {cu} are substituted.",
+ )
+ ap.add_argument(
+ "--output-dir",
+ type=Path,
+ default=Path("partition-sweep"),
+ help="Where per-partition benchmark output and the summary go.",
+ )
+ ap.add_argument(
+ "--streams-per-partition",
+ type=int,
+ default=2,
+ help="Concurrent streams the entrypoint should place on each partition (default 2).",
+ )
+ ap.add_argument(
+ "--per-stream-gib",
+ type=float,
+ default=0.0,
+ help="Per-stream HBM footprint. When given, modes whose partitions provably cannot hold "
+ "the streams are skipped instead of being run into an out-of-memory failure.",
+ )
+ ap.add_argument(
+ "--timeout",
+ type=float,
+ default=3600.0,
+ help="Per-mode budget in seconds for the whole fan-out (default 3600).",
+ )
+ ap.add_argument("--sudo", action="store_true", help="Route privileged amd-smi calls through 'sudo -n'.")
+ ap.add_argument("--dry-run", action="store_true", help="Print the plan and exit without changing anything.")
+ ap.add_argument(
+ "--allow-busy",
+ action="store_true",
+ help="Proceed even though processes hold a context on the swept card; the set evicts them. "
+ "Skips the check entirely, so it is also the way past an amd-smi process payload this "
+ "script cannot parse.",
+ )
+ return ap.parse_args(argv)
+
+
+def main(argv: Sequence[str] | None = None) -> int: # noqa: C901
+ args = _parse_args(argv)
+
+ if args.benchmark_command:
+ template: list[str] = shlex.split(args.benchmark_command)
+ elif args.benchmark_config:
+ template = magpie_command(args.benchmark_config)
+ else:
+ print("ERROR: one of --benchmark-config or --benchmark-command is required.", file=sys.stderr)
+ return 2
+ if args.streams_per_partition < 1:
+ print("ERROR: --streams-per-partition must be at least 1.", file=sys.stderr)
+ return 2
+
+ try:
+ entry_mode = read_mode(args.gpu)
+ bus = card_bus(_amd_smi_json(["list"]), args.gpu)
+ card_gib = card_total_gib(read_device_gib(args.gpu), entry_mode)
+ try:
+ available = supported_modes(_amd_smi_json(["partition", "-a", "-g", str(args.gpu)], sudo=args.sudo))
+ except SweepError as exc:
+ print(f"note: could not read partition profiles ({exc}); not restricting the mode list")
+ available = ()
+ modes = resolve_modes(args.modes, available)
+ # Scoped to the swept card because that is the only card a set touches.
+ # A neighbour's benchmark is not a reason to refuse, and refusing on one
+ # left --allow-busy as the only way forward -- which drops the guard on
+ # the target card too, the one card it exists to protect. Skipped
+ # entirely when no set will follow, so a payload it cannot read never
+ # blocks a caller it was not protecting.
+ if not (args.allow_busy or args.dry_run):
+ busy = resident_processes(_amd_smi_json(["process"]))
+ if args.gpu not in busy:
+ raise SweepError(
+ f"amd-smi process listed no GPU {args.gpu}, so whether the card is in use "
+ f"is unknown. Repartitioning evicts every context on it, so this refuses "
+ f"rather than assume it is idle. Pass --allow-busy to sweep anyway."
+ )
+ if busy[args.gpu]:
+ raise SweepError(
+ f"{busy[args.gpu]} process(es) still hold a context on GPU {args.gpu}. "
+ f"Repartitioning evicts them, so this refuses rather than killing "
+ f"someone's work. Stop them, or pass --allow-busy if they are yours."
+ )
+ except (SweepError, PartitionError) as exc:
+ print(f"ERROR: {exc}", file=sys.stderr)
+ return 2
+ except Exception:
+ # Nothing has been set at this point, so there is no card to restore --
+ # but the exit code still has to mean something. 2 is for a refusal this
+ # script decided on, so an unmodelled failure gets its own code rather
+ # than borrowing that one.
+ traceback.print_exc()
+ print("ERROR: unexpected error while reading the card; nothing was changed.", file=sys.stderr)
+ return 4
+
+ print(f"GPU {args.gpu} on PCI bus {bus:02x} is in {entry_mode}; sweeping {', '.join(modes)}.")
+ if args.dry_run:
+ print("\n--dry-run: nothing will be set. Planned per mode:")
+ for mode in modes:
+ partitions = MODE_PARTITION_COUNTS[mode]
+ print(f" {mode:<4} set GPU {args.gpu}, then {partitions} concurrent benchmark(s), one per partition")
+ print(f"\nCommand template: {' '.join(template)}")
+ print(f"Would restore {entry_mode} on exit.")
+ return 0
+
+ output_root = args.output_dir.expanduser().resolve()
+ output_root.mkdir(parents=True, exist_ok=True)
+ results: list[ModeResult] = []
+ unexpected: str | None = None
+ interrupted = {"flag": False}
+
+ def _on_signal(signum: int, _frame: Any) -> None:
+ interrupted["flag"] = True
+ print(f"\nsignal {signum} received; finishing the current mode, then restoring {entry_mode}.")
+
+ for sig in (signal.SIGINT, signal.SIGTERM):
+ signal.signal(sig, _on_signal)
+
+ try:
+ for mode in modes:
+ if interrupted["flag"]:
+ results.append(ModeResult(mode=mode, skipped="interrupted before this mode ran"))
+ continue
+ result = ModeResult(mode=mode)
+ results.append(result)
+ print(f"\n--- {mode} ---")
+ try:
+ set_mode(args.gpu, mode, sudo=args.sudo)
+ agents = read_hsa_agents()
+ cu = partition_cu_on_bus(agents, bus)
+ gib = None if card_gib is None else card_gib / MODE_PARTITION_COUNTS[mode]
+ layout = layout_for(mode, cu_per_partition=cu, gib_per_partition=gib)
+ result.layout = layout
+ if args.per_stream_gib > 0 and not fits_in_partition(
+ args.per_stream_gib, layout, args.streams_per_partition
+ ):
+ result.skipped = (
+ f"{args.streams_per_partition} x {args.per_stream_gib:.1f} GiB will not fit a "
+ f"{layout.gib_per_partition:.1f} GiB partition"
+ )
+ print(f" skipped: {result.skipped}")
+ continue
+ devices = select_partition_devices(agents, layout, bus=bus)
+ print(
+ f" {layout.partitions} partition(s) of {layout.cu_per_partition} CU at HIP devices {list(devices)}"
+ )
+ result.runs = run_mode(
+ layout,
+ devices,
+ template=template,
+ output_root=output_root,
+ streams_per_partition=args.streams_per_partition,
+ timeout_s=args.timeout,
+ )
+ ok = sum(1 for r in result.runs if r.ok)
+ print(f" {ok}/{len(result.runs)} partition(s) reported")
+ for run in result.runs:
+ if run.error:
+ print(f" partition {run.partition_index} (HIP {run.hsa_device}): {run.error}")
+ except (SweepError, PartitionError) as exc:
+ result.error = str(exc)
+ print(f" failed: {exc}")
+ except Exception as exc:
+ # Everything this script anticipates arrives as a SweepError or a
+ # PartitionError. Anything else is a bug here or an amd-smi
+ # behaviour not modelled, which means the assumptions driving
+ # privileged sets no longer hold -- so stop sweeping. Stopping by
+ # breaking rather than propagating is the point: the card still
+ # gets restored, and the modes already measured still get
+ # reported. Letting it escape lost the table, the summary file,
+ # and the exit code that says the card was left wrong.
+ unexpected = f"{type(exc).__name__}: {exc}"
+ result.error = f"unexpected error: {unexpected}"
+ print(f" aborted: {result.error}", file=sys.stderr)
+ traceback.print_exc()
+ break
+ finally:
+ restore_failed = _restore_entry_mode(args.gpu, entry_mode, sudo=args.sudo)
+
+ summary = output_root / "sweep_summary.json"
+ try:
+ print("\n".join(render(results, entry_mode=entry_mode)))
+ summary.write_text(
+ json.dumps(summary_json(results, entry_mode=entry_mode, gpu_id=args.gpu), indent=2, sort_keys=True) + "\n",
+ encoding="utf-8",
+ )
+ print(f"wrote {summary}")
+ except Exception as exc:
+ # A full disk or an unrenderable result must not cost the caller the exit
+ # code, which is the one thing it cannot reconstruct for itself -- least
+ # of all the code saying the card was left in the wrong mode.
+ traceback.print_exc()
+ print(f"ERROR: could not write the report to {summary}: {exc}", file=sys.stderr)
+ unexpected = unexpected or f"{type(exc).__name__}: {exc}"
+
+ # A card left in the wrong shape outranks everything else: it mislabels
+ # whatever runs on the node next, not just this sweep.
+ if restore_failed:
+ return 3
+ if unexpected:
+ print(f"ERROR: sweep stopped on an unexpected error ({unexpected}).", file=sys.stderr)
+ return 4
+ return 0 if any(r.measured for r in results) else 1
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/scripts/tests/test_partition_mode_sweep.py b/scripts/tests/test_partition_mode_sweep.py
new file mode 100644
index 0000000000..1a209485eb
--- /dev/null
+++ b/scripts/tests/test_partition_mode_sweep.py
@@ -0,0 +1,897 @@
+# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
+# SPDX-License-Identifier: MIT
+
+"""Hardware-free tests for ``scripts/partition_mode_sweep.py``.
+
+The script sits outside the coverage denominator because ``scripts/`` is not
+shipped as a package, but the logic tested here decides which silicon a
+benchmark runs on -- and getting that wrong produces a plausible number
+attributed to the wrong mode. So the parsing, device selection and aggregation
+are pinned here. Nothing below sets a partition mode, shells out, or needs a
+GPU.
+
+Payload shapes are the ones observed on an 8-card MI355X node rather than
+invented: ``amd-smi`` reporting an idle card's ``process_list`` as a bare
+string, ``rocminfo`` putting ``BDFID`` before ``Compute Unit`` in each agent
+block, and HSA enumerating whole cards ahead of partitions.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import pytest
+
+_SCRIPT = Path(__file__).resolve().parents[1] / "partition_mode_sweep.py"
+
+
+def _load():
+ spec = importlib.util.spec_from_file_location("partition_mode_sweep", _SCRIPT)
+ assert spec and spec.loader
+ mod = importlib.util.module_from_spec(spec)
+ # Registered before exec: the module's dataclasses resolve their annotations
+ # through sys.modules, and a loader that skips this raises AttributeError on
+ # the first field with a default_factory.
+ sys.modules[spec.name] = mod
+ spec.loader.exec_module(mod)
+ return mod
+
+
+pms = _load()
+
+
+#: One MI355X (card 0, PCI bus 0x09) in CPX beside seven untouched cards, as
+#: ``rocminfo`` actually printed it. Two properties matter and both are load
+#: bearing: ``BDFID`` precedes ``Compute Unit``, and the whole cards come first.
+ROCMINFO_CPX = """
+Agent 1
+ Name: AMD EPYC 9575F 64-Core Processor
+ Device Type: CPU
+ Compute Unit: 128
+Agent 2
+ Name: AMD EPYC 9575F 64-Core Processor
+ Device Type: CPU
+ Compute Unit: 128
+Agent 3
+ Name: gfx950
+ Device Type: GPU
+ BDFID: 31744
+ Compute Unit: 256
+Agent 4
+ Name: gfx950
+ Device Type: GPU
+ BDFID: 26880
+ Compute Unit: 256
+Agent 5
+ Name: gfx950
+ Device Type: GPU
+ BDFID: 2304
+ Compute Unit: 32
+Agent 6
+ Name: gfx950
+ Device Type: GPU
+ BDFID: 2305
+ Compute Unit: 32
+"""
+
+
+# ------------------------------------------------------- rocminfo parsing
+
+
+class TestParseHsaAgents:
+ def test_cu_is_not_shifted_by_the_field_order(self):
+ """BDFID precedes Compute Unit, so a line-at-a-time parser is off by one.
+
+ Printing on ``BDFID`` with the last-seen CU attributes each agent the
+ previous one's count -- which, on the node this came from, silently
+ turned the first CPX partition into a 256-CU card and would have pointed
+ the benchmark at whole silicon while labelling it CPX.
+ """
+ agents = pms.parse_hsa_agents(ROCMINFO_CPX)
+ by_bdf = {a.bdf: a.cu for a in agents}
+ assert by_bdf == {"7c:00.0": 256, "69:00.0": 256, "09:00.0": 32, "09:00.1": 32}
+
+ def test_cpu_agents_are_not_gpus(self):
+ agents = pms.parse_hsa_agents(ROCMINFO_CPX)
+ assert len(agents) == 4
+ assert all(a.cu in (32, 256) for a in agents)
+
+ def test_indices_are_positions_among_gpus_only(self):
+ """The index has to be the one ROCR_VISIBLE_DEVICES uses, so CPUs cannot count."""
+ agents = pms.parse_hsa_agents(ROCMINFO_CPX)
+ assert [a.index for a in agents] == [0, 1, 2, 3]
+
+ def test_bdfid_decodes_to_bus_device_function(self):
+ agents = pms.parse_hsa_agents(ROCMINFO_CPX)
+ first = next(a for a in agents if a.bdf == "09:00.0")
+ assert (first.bus, first.device, first.function) == (0x09, 0x00, 0)
+
+ def test_partitions_of_one_card_share_a_bus_and_differ_by_function(self):
+ agents = pms.parse_hsa_agents(ROCMINFO_CPX)
+ on_card = [a for a in agents if a.bus == 0x09]
+ assert {a.function for a in on_card} == {0, 1}
+
+ def test_an_agent_missing_its_fields_is_dropped_not_guessed(self):
+ assert pms.parse_hsa_agents("Agent 1\n Device Type: GPU\n") == ()
+
+ def test_empty_input_is_no_agents(self):
+ assert pms.parse_hsa_agents("") == ()
+
+
+# --------------------------------------------------------- device selection
+
+
+def _agents(*specs: tuple[int, int, int]) -> tuple:
+ """Build agents from ``(cu, bus, function)`` triples in enumeration order."""
+ return tuple(pms.HsaAgent(index=i, cu=cu, bus=bus, device=0, function=fn) for i, (cu, bus, fn) in enumerate(specs))
+
+
+class TestSelectPartitionDevices:
+ def test_whole_cards_enumerate_first_so_indices_are_not_zero_based(self):
+ """The measured trap: under DPX the partitions are HIP devices 7 and 8.
+
+ Seven untouched cards take indices 0-6, so a driver that assumed the
+ partitions of the card it just split start at 0 would benchmark a
+ neighbour and report the number as DPX.
+ """
+ agents = _agents(*[(256, 0x19 + i, 0) for i in range(7)], (128, 0x09, 0), (128, 0x09, 1))
+ layout = pms.layout_for("DPX", cu_per_partition=128)
+ assert pms.select_partition_devices(agents, layout, bus=0x09) == (7, 8)
+
+ def test_spx_does_not_collect_the_untouched_neighbours(self):
+ """SPX's "partition" is a whole card, so CU alone matches all eight."""
+ agents = _agents(*[(256, 0x19 + i, 0) for i in range(7)], (256, 0x09, 0))
+ layout = pms.layout_for("SPX", cu_per_partition=256)
+ assert pms.select_partition_devices(agents, layout, bus=0x09) == (7,)
+
+ def test_a_set_that_did_not_take_is_an_error_not_a_short_list(self):
+ """Finding fewer partitions than the mode implies means the topology is stale."""
+ agents = _agents((256, 0x09, 0))
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ with pytest.raises(pms.SweepError, match="CPX implies 8 partitions"):
+ pms.select_partition_devices(agents, layout, bus=0x09)
+
+ def test_the_error_names_what_was_actually_on_the_bus(self):
+ agents = _agents((256, 0x09, 0))
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ with pytest.raises(pms.SweepError, match=r"09:00\.0=256CU"):
+ pms.select_partition_devices(agents, layout, bus=0x09)
+
+ def test_nothing_on_the_bus_says_so(self):
+ layout = pms.layout_for("SPX", cu_per_partition=256)
+ with pytest.raises(pms.SweepError, match="nothing on that bus"):
+ pms.select_partition_devices(_agents((256, 0x19, 0)), layout, bus=0x09)
+
+ def test_selection_uses_the_shared_predicate(self):
+ """The rule lives in gpu_partition; this script must not re-invent it."""
+ agents = _agents((256, 0x09, 0), (32, 0x09, 1))
+ assert pms.partition_device_predicate(32)(32) is True
+ assert pms.partition_device_predicate(32)(256) is False
+ layout = pms.layout_for("DPX", cu_per_partition=32)
+ with pytest.raises(pms.SweepError):
+ # Two devices on the bus, but only one is 32 CU: DPX wants two.
+ pms.select_partition_devices(agents, layout, bus=0x09)
+
+
+class TestPartitionCuOnBus:
+ def test_a_settled_card_reports_one_width(self):
+ assert pms.partition_cu_on_bus(_agents((32, 0x09, 0), (32, 0x09, 1)), 0x09) == 32
+
+ def test_neighbours_do_not_contribute(self):
+ assert pms.partition_cu_on_bus(_agents((32, 0x09, 0), (256, 0x19, 0)), 0x09) == 32
+
+ def test_mixed_widths_mean_the_read_was_mid_transition(self):
+ with pytest.raises(pms.SweepError, match="mid-transition"):
+ pms.partition_cu_on_bus(_agents((32, 0x09, 0), (256, 0x09, 1)), 0x09)
+
+ def test_an_empty_bus_is_an_error(self):
+ with pytest.raises(pms.SweepError, match="no GPU on bus"):
+ pms.partition_cu_on_bus(_agents((256, 0x19, 0)), 0x09)
+
+
+# ------------------------------------------------------------ amd-smi payloads
+
+
+class TestResidentProcesses:
+ def test_an_idle_card_reports_a_string_not_an_empty_list(self):
+ """The observed payload. Counting it as a process refuses every free node."""
+ payload = [{"gpu": 0, "process_list": [{"process_info": "No running processes detected"}]}]
+ assert pms.resident_processes(payload) == {0: 0}
+
+ def test_a_real_process_is_counted(self):
+ payload = [{"gpu": 0, "process_list": [{"process_info": {"name": "python", "pid": 42}}]}]
+ assert pms.resident_processes(payload) == {0: 1}
+
+ def test_a_nameless_entry_is_not_a_process(self):
+ payload = [{"gpu": 0, "process_list": [{"process_info": {"name": " "}}]}]
+ assert pms.resident_processes(payload) == {0: 0}
+
+ def test_an_unexpected_string_is_counted_rather_than_ignored(self):
+ """Fail closed: an unrecognised sentinel might really be a process."""
+ payload = [{"gpu": 0, "process_list": [{"process_info": "something new"}]}]
+ assert pms.resident_processes(payload) == {0: 1}
+
+ def test_cards_are_reported_separately(self):
+ payload = [
+ {"gpu": 0, "process_list": [{"process_info": "No running processes detected"}]},
+ {"gpu": 1, "process_list": [{"process_info": {"name": "vllm"}}]},
+ ]
+ assert pms.resident_processes(payload) == {0: 0, 1: 1}
+
+ def test_the_idle_sentinel_is_also_recognised_unwrapped(self):
+ """Observed at both depths: as the list, and as the entry's process_info."""
+ assert pms.resident_processes([{"gpu": 0, "process_list": "No running processes detected"}]) == {0: 0}
+
+ def test_a_process_is_counted_without_the_process_info_wrapper(self):
+ assert pms.resident_processes([{"gpu": 0, "process_list": [{"name": "python", "pid": 42}]}]) == {0: 1}
+
+ def test_an_empty_list_is_an_idle_card(self):
+ assert pms.resident_processes([{"gpu": 0, "process_list": []}]) == {0: 0}
+
+
+class TestResidentProcessesRefusesDrift:
+ """Schema drift must raise, never read as "nobody is using the node".
+
+ This count is the only thing between an ``amd-smi`` payload this parser does
+ not understand and a partition set that evicts whatever is running. Every
+ case below used to return ``{}`` or silently skip the row, which the sweep
+ read as an idle node and acted on.
+ """
+
+ @pytest.mark.parametrize(
+ "payload",
+ [
+ {"process": [{"gpu": 0, "process_list": []}]}, # a wrapper key appears
+ "No running processes detected", # the whole payload degrades to a string
+ None, # the query answers nothing
+ 42,
+ ],
+ )
+ def test_a_payload_that_is_not_a_list_of_rows_is_refused(self, payload):
+ with pytest.raises(pms.SweepError, match="amd-smi process returned"):
+ pms.resident_processes(payload)
+
+ def test_a_row_that_is_not_a_mapping_is_refused(self):
+ with pytest.raises(pms.SweepError, match="where a GPU row was expected"):
+ pms.resident_processes([["gpu", 0]])
+
+ @pytest.mark.parametrize("row", [{"process_list": []}, {"gpu": 0}, {}])
+ def test_a_row_missing_either_field_is_refused(self, row):
+ with pytest.raises(pms.SweepError, match="without gpu|without process_list|without gpu or process_list"):
+ pms.resident_processes([row])
+
+ def test_a_renamed_gpu_id_is_refused_rather_than_skipped(self):
+ """A row filed under a name this parser cannot read is not an absent row."""
+ with pytest.raises(pms.SweepError, match="not a number"):
+ pms.resident_processes([{"gpu": "card0", "process_list": []}])
+
+ def test_a_process_list_of_an_unmodelled_type_is_refused(self):
+ with pytest.raises(pms.SweepError, match="process_list of type dict"):
+ pms.resident_processes([{"gpu": 0, "process_list": {"pid": 42}}])
+
+ def test_a_process_entry_of_an_unmodelled_type_is_refused(self):
+ with pytest.raises(pms.SweepError, match="process entry of type NoneType"):
+ pms.resident_processes([{"gpu": 0, "process_list": [None]}])
+
+
+class TestCurrentModes:
+ def test_reads_the_observed_payload(self):
+ payload = {
+ "current_partition": [
+ {"gpu_id": 0, "memory": "NPS1", "accelerator_type": "SPX"},
+ {"gpu_id": 1, "memory": "NPS1", "accelerator_type": "CPX"},
+ ]
+ }
+ assert pms.current_modes(payload) == {0: "SPX", 1: "CPX"}
+
+ def test_the_current_marker_asterisk_is_stripped(self):
+ payload = {"current_partition": [{"gpu_id": 0, "accelerator_type": "SPX*"}]}
+ assert pms.current_modes(payload) == {0: "SPX"}
+
+ def test_not_available_is_absent_rather_than_a_mode(self):
+ payload = {"current_partition": [{"gpu_id": 0, "accelerator_type": "N/A"}]}
+ assert pms.current_modes(payload) == {}
+
+
+class TestSupportedModes:
+ def test_only_named_rows_describe_a_profile(self):
+ """The profile table continues each profile in rows with blank identity."""
+ payload = {
+ "partition_profiles": [
+ {"profile_index": 0, "accelerator_type": "SPX*", "num_partitions": 1},
+ {"profile_index": None, "accelerator_type": "", "resource_type": "DMA"},
+ {"profile_index": 3, "accelerator_type": "CPX", "num_partitions": 8},
+ ]
+ }
+ assert pms.supported_modes(payload) == ("SPX", "CPX")
+
+ def test_an_unprivileged_query_reports_nothing_not_everything(self):
+ payload = {"partition_profiles": [{"accelerator_type": "N/A"}]}
+ assert pms.supported_modes(payload) == ()
+
+ def test_duplicates_collapse(self):
+ payload = {
+ "partition_profiles": [
+ {"accelerator_type": "DPX"},
+ {"accelerator_type": "DPX"},
+ ]
+ }
+ assert pms.supported_modes(payload) == ("DPX",)
+
+
+class TestCardBus:
+ def test_reads_the_bus_from_a_domain_qualified_address(self):
+ assert pms.card_bus([{"gpu": 0, "bdf": "0000:09:00.0"}], 0) == 0x09
+
+ def test_reads_the_bus_without_a_domain(self):
+ assert pms.card_bus([{"gpu": 1, "bdf": "7c:00.0"}], 1) == 0x7C
+
+ def test_a_missing_card_is_an_error(self):
+ with pytest.raises(pms.SweepError, match="no PCI address"):
+ pms.card_bus([{"gpu": 0, "bdf": "0000:09:00.0"}], 3)
+
+
+# -------------------------------------------------------------- the fan-out
+
+
+class TestBuildPartitionCommand:
+ def test_substitutes_the_per_partition_values(self):
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ cmd = pms.build_partition_command(
+ ["bench", "--gpu", "{device}", "--out", "{output_dir}", "--mode", "{mode}"],
+ device=9,
+ output_dir=Path("/tmp/out"),
+ layout=layout,
+ partition_index=2,
+ )
+ assert cmd == ["bench", "--gpu", "9", "--out", "/tmp/out", "--mode", "CPX"]
+
+ def test_a_path_with_a_space_stays_one_argument(self):
+ """Substitution is per already-split token, so no shell can re-split it."""
+ layout = pms.layout_for("SPX", cu_per_partition=256)
+ cmd = pms.build_partition_command(
+ ["bench", "{output_dir}"],
+ device=0,
+ output_dir=Path("/tmp/my runs"),
+ layout=layout,
+ partition_index=0,
+ )
+ assert cmd == ["bench", "/tmp/my runs"]
+
+ def test_shape_placeholders_are_available(self):
+ layout = pms.layout_for("QPX", cu_per_partition=64)
+ cmd = pms.build_partition_command(
+ ["b", "{partitions}", "{partition_index}", "{cu}"],
+ device=0,
+ output_dir=Path("/x"),
+ layout=layout,
+ partition_index=3,
+ )
+ assert cmd == ["b", "4", "3", "64"]
+
+ def test_a_template_without_placeholders_is_left_alone(self):
+ layout = pms.layout_for("SPX", cu_per_partition=256)
+ assert pms.build_partition_command(
+ ["b", "-v"], device=0, output_dir=Path("/x"), layout=layout, partition_index=0
+ ) == [
+ "b",
+ "-v",
+ ]
+
+
+class TestPartitionEnv:
+ def test_pins_the_process_to_its_partition(self):
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ env = pms.partition_env({}, layout, device=11, streams_per_partition=2)
+ assert env["ROCR_VISIBLE_DEVICES"] == "11"
+
+ def test_an_inherited_hip_mask_is_removed_not_kept(self):
+ """Two masks apply in sequence, the second indexing into the first.
+
+ A stale HIP_VISIBLE_DEVICES=0 alongside ROCR would send every
+ partition's work to one device, and the sweep would report the mode as
+ uniformly slow with every process apparently succeeding.
+ """
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ env = pms.partition_env(
+ {"HIP_VISIBLE_DEVICES": "0", "CUDA_VISIBLE_DEVICES": "0"},
+ layout,
+ device=11,
+ streams_per_partition=2,
+ )
+ assert "HIP_VISIBLE_DEVICES" not in env
+ assert "CUDA_VISIBLE_DEVICES" not in env
+
+ def test_publishes_the_same_contract_the_optimizer_publishes(self):
+ """An entrypoint written against a session's env must work here unchanged."""
+ layout = pms.layout_for("CPX", cu_per_partition=32)
+ env = pms.partition_env({}, layout, device=7, streams_per_partition=2)
+ assert env[pms.PARTITION_MODE_ENV] == "CPX"
+ assert env[pms.PARTITION_COUNT_ENV] == "8"
+ assert env[pms.PARTITION_CU_ENV] == "32"
+ assert env[pms.PARTITION_STREAMS_ENV] == "2"
+ assert env[pms.PARTITION_TOTAL_STREAMS_ENV] == "16"
+
+ def test_the_caller_environment_is_not_mutated(self):
+ base = {"PATH": "/usr/bin"}
+ layout = pms.layout_for("SPX", cu_per_partition=256)
+ pms.partition_env(base, layout, device=0, streams_per_partition=1)
+ assert base == {"PATH": "/usr/bin"}
+
+
+# ------------------------------------------------------------- aggregation
+
+
+def _run(index: int, *, throughput: float | None = 100.0, latency: float | None = 10.0, ok: bool = True):
+ measurement: dict = {}
+ if throughput is not None:
+ measurement["output_throughput"] = throughput
+ if latency is not None:
+ measurement["e2el_mean_ms"] = latency
+ return pms.PartitionRun(
+ partition_index=index,
+ hsa_device=7 + index,
+ output_dir=Path("/x"),
+ returncode=0 if ok else 1,
+ measurement=measurement if ok else {},
+ error="" if ok else "exited 1",
+ )
+
+
+class TestModeResult:
+ def test_throughput_is_the_sum_over_partitions(self):
+ """The whole point: a mode's number is every partition, loaded together."""
+ result = pms.ModeResult(mode="CPX", runs=[_run(i, throughput=320.0) for i in range(8)])
+ assert result.total() == pytest.approx(2560.0)
+
+ def test_a_partially_reporting_mode_is_unmeasured_not_slow(self):
+ """Summing six of eight understates by a quarter while looking like a result."""
+ runs = [_run(i) for i in range(6)] + [_run(6, ok=False), _run(7, ok=False)]
+ result = pms.ModeResult(mode="CPX", runs=runs)
+ assert result.measured is False
+
+ def test_every_partition_reporting_is_measured(self):
+ result = pms.ModeResult(mode="DPX", runs=[_run(0), _run(1)])
+ assert result.measured is True
+
+ def test_a_mode_that_never_ran_is_not_measured(self):
+ assert pms.ModeResult(mode="CPX", skipped="will not fit").measured is False
+
+ def test_a_field_missing_on_one_partition_blocks_the_sum(self):
+ """A partial sum is a wrong number, not an approximate one."""
+ result = pms.ModeResult(mode="DPX", runs=[_run(0), _run(1, throughput=None)])
+ assert result.total() is None
+
+ def test_the_throughput_field_falls_back_in_preference_order(self):
+ runs = [
+ pms.PartitionRun(0, 7, Path("/x"), returncode=0, measurement={"total_token_throughput": 5.0}),
+ pms.PartitionRun(1, 8, Path("/x"), returncode=0, measurement={"total_token_throughput": 7.0}),
+ ]
+ assert pms.ModeResult(mode="DPX", runs=runs).total() == pytest.approx(12.0)
+
+ def test_latency_is_the_worst_partition_not_the_mean_of_them(self):
+ """A request landing on the slow partition is not consoled by the average."""
+ runs = [_run(0, latency=10.0), _run(1, latency=90.0)]
+ assert pms.ModeResult(mode="DPX", runs=runs).worst_latency_ms() == pytest.approx(90.0)
+
+ def test_latency_accepts_the_alternate_spellings(self):
+ runs = [pms.PartitionRun(0, 7, Path("/x"), returncode=0, measurement={"mean_e2el_ms": 12.0})]
+ assert pms.ModeResult(mode="SPX", runs=runs).worst_latency_ms() == pytest.approx(12.0)
+
+ def test_latency_missing_on_one_partition_is_not_reported(self):
+ runs = [_run(0, latency=10.0), _run(1, latency=None)]
+ assert pms.ModeResult(mode="DPX", runs=runs).worst_latency_ms() is None
+
+
+class TestToFloat:
+ """The canonical coercion from hyperloom.common.coerce, not a local copy.
+
+ Pinned here because the aggregation depends on its rejections: a stray bool
+ becoming 1.0, or an inf reaching the sum, produces a throughput comparison
+ that is wrong rather than absent.
+ """
+
+ def test_a_bool_is_not_a_measurement(self):
+ assert pms.to_float(True) is None
+
+ @pytest.mark.parametrize("bad", ["", "abc", None, float("inf"), float("nan")])
+ def test_unusable_values_are_none(self, bad):
+ assert pms.to_float(bad) is None
+
+ def test_a_numeric_string_is_accepted(self):
+ assert pms.to_float("12.5") == pytest.approx(12.5)
+
+
+class TestCardTotalGib:
+ def test_a_split_mode_reading_is_scaled_back_to_the_card(self):
+ """36 GiB on a CPX partition is a 288 GiB MI355X."""
+ assert pms.card_total_gib(36.0, "CPX") == pytest.approx(288.0)
+
+ def test_an_unpartitioned_reading_is_already_the_card(self):
+ assert pms.card_total_gib(288.0, "SPX") == pytest.approx(288.0)
+
+ @pytest.mark.parametrize("bad", [None, 0.0, -1.0])
+ def test_an_unknown_reading_stays_unknown(self, bad):
+ assert pms.card_total_gib(bad, "SPX") is None
+
+
+class TestFeasibilityGate:
+ def test_the_mi355x_case_that_motivated_the_gate(self):
+ """20.7 GiB x 2 streams does not fit a 36 GiB CPX partition."""
+ layout = pms.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0)
+ assert pms.fits_in_partition(20.7, layout, 2) is False
+
+ def test_one_stream_of_the_same_workload_does_fit(self):
+ """Which is why gating on the single-stream figure is the trap."""
+ layout = pms.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0)
+ assert pms.fits_in_partition(20.7, layout, 1) is True
+
+ def test_the_whole_card_holds_it_comfortably(self):
+ layout = pms.layout_for("SPX", cu_per_partition=256, gib_per_partition=288.0)
+ assert pms.fits_in_partition(20.7, layout, 2) is True
+
+
+# ----------------------------------------------------------- mode resolution
+
+
+class TestResolveModes:
+ def test_defaults_to_what_the_card_reports(self):
+ assert pms.resolve_modes("", ("SPX", "DPX", "QPX", "CPX")) == ("SPX", "DPX", "QPX", "CPX")
+
+ def test_an_explicit_list_keeps_its_order(self):
+ assert pms.resolve_modes("CPX,SPX", ("SPX", "DPX", "QPX", "CPX")) == ("CPX", "SPX")
+
+ def test_whitespace_and_commas_both_separate(self):
+ assert pms.resolve_modes("spx dpx", ("SPX", "DPX")) == ("SPX", "DPX")
+
+ def test_duplicates_collapse(self):
+ assert pms.resolve_modes("SPX,SPX", ("SPX",)) == ("SPX",)
+
+ def test_a_misspelled_mode_is_refused(self):
+ with pytest.raises(pms.SweepError, match="unknown compute-partition mode"):
+ pms.resolve_modes("SPXX", ("SPX",))
+
+ def test_a_mode_the_card_disclaims_is_refused_before_anything_is_set(self):
+ with pytest.raises(pms.SweepError, match="does not report support for CPX"):
+ pms.resolve_modes("CPX", ("SPX", "DPX"))
+
+ def test_an_unqueryable_card_is_not_second_guessed(self):
+ """No profile table means "unknown", so the request stands and the set judges it."""
+ assert pms.resolve_modes("CPX", ()) == ("CPX",)
+
+ def test_no_request_and_no_profiles_says_which_flag_to_pass(self):
+ with pytest.raises(pms.SweepError, match="--modes"):
+ pms.resolve_modes("", ())
+
+
+# --------------------------------------------------------------- reporting
+
+
+class TestRender:
+ def test_a_skipped_mode_is_shown_with_its_reason(self):
+ results = [
+ pms.ModeResult(mode="SPX", layout=pms.layout_for("SPX", cu_per_partition=256), runs=[_run(0)]),
+ pms.ModeResult(mode="CPX", skipped="2 x 20.7 GiB will not fit a 36.0 GiB partition"),
+ ]
+ text = "\n".join(pms.render(results, entry_mode="SPX"))
+ assert "will not fit" in text
+ assert "CPX" in text
+
+ def test_the_winner_is_named(self):
+ results = [
+ pms.ModeResult(
+ mode="SPX", layout=pms.layout_for("SPX", cu_per_partition=256), runs=[_run(0, throughput=100.0)]
+ ),
+ pms.ModeResult(
+ mode="DPX",
+ layout=pms.layout_for("DPX", cu_per_partition=128),
+ runs=[_run(0, throughput=90.0), _run(1, throughput=90.0)],
+ ),
+ ]
+ text = "\n".join(pms.render(results, entry_mode="SPX"))
+ assert "Fastest measured mode: DPX" in text
+ assert "1.80x" in text
+
+ def test_nothing_measured_says_so_rather_than_naming_a_winner(self):
+ results = [pms.ModeResult(mode="CPX", error="setting CPX failed")]
+ text = "\n".join(pms.render(results, entry_mode="SPX"))
+ assert "No mode produced a measurement" in text
+
+ def test_the_report_says_the_number_is_a_sum(self):
+ """The claim has to travel with the figure or it gets read as a single run."""
+ results = [pms.ModeResult(mode="SPX", layout=pms.layout_for("SPX", cu_per_partition=256), runs=[_run(0)])]
+ text = "\n".join(pms.render(results, entry_mode="SPX"))
+ assert "sum over a mode's partitions" in text
+
+ def test_it_points_at_the_assertion_flag_for_the_session_that_follows(self):
+ results = [pms.ModeResult(mode="SPX", layout=pms.layout_for("SPX", cu_per_partition=256), runs=[_run(0)])]
+ text = "\n".join(pms.render(results, entry_mode="SPX"))
+ assert "--compute-partition-mode" in text
+
+
+class TestSummaryJson:
+ def test_records_the_shape_and_the_per_partition_devices(self):
+ result = pms.ModeResult(
+ mode="DPX",
+ layout=pms.layout_for("DPX", cu_per_partition=128, gib_per_partition=144.0),
+ runs=[_run(0), _run(1)],
+ )
+ payload = pms.summary_json([result], entry_mode="SPX", gpu_id=0)
+ assert payload["entry_mode"] == "SPX"
+ mode = payload["modes"][0]
+ assert mode["partitions"] == 2
+ assert mode["cu_per_partition"] == 128
+ assert [r["hsa_device"] for r in mode["partition_runs"]] == [7, 8]
+
+ def test_a_skipped_mode_records_why_rather_than_a_null_row(self):
+ payload = pms.summary_json([pms.ModeResult(mode="CPX", skipped="will not fit")], entry_mode="SPX", gpu_id=0)
+ assert payload["modes"][0]["skipped"] == "will not fit"
+ assert payload["modes"][0]["measured"] is False
+
+
+class TestMagpieCommand:
+ def test_the_output_dir_is_left_as_a_placeholder_for_the_fan_out(self):
+ cmd = pms.magpie_command(Path("/cfg/bench.yaml"), python_exe="/usr/bin/python3")
+ assert "{output_dir}" in cmd
+ assert "/cfg/bench.yaml" in cmd
+
+ def test_it_is_the_local_run_mode(self):
+ assert "local" in pms.magpie_command(Path("/cfg/bench.yaml"))
+
+
+class TestBusyDetection:
+ @pytest.mark.parametrize(
+ "detail",
+ ["AMDSMI_STATUS_BUSY", "Device busy, try again", "resident process present"],
+ )
+ def test_a_transient_refusal_is_retried(self, detail):
+ assert pms._is_busy(detail) is True
+
+ @pytest.mark.parametrize("detail", ["permission denied", "unknown partition mode", ""])
+ def test_a_permanent_failure_fails_fast(self, detail):
+ """Retrying these for two minutes turns a clear error into a hang."""
+ assert pms._is_busy(detail) is False
+
+
+class TestSudoPrefix:
+ def test_privilege_is_opt_in(self):
+ assert pms._sudo_prefix(False) == []
+
+ def test_the_prefix_never_prompts(self):
+ """An unattended sweep that stops at a password prompt hangs silently."""
+ assert pms._sudo_prefix(True) == ["sudo", "-n"]
+
+
+# ------------------------------------------------------------ control flow
+#
+# The tests above are pure functions. What follows drives main() end to end
+# against a fake node, because the decisions being pinned -- which card's
+# processes block a set, and what happens on the way out of a failure -- live in
+# its control flow and cannot be reached any other way. Nothing here touches a
+# GPU: every call that would is replaced.
+
+
+IDLE = "No running processes detected"
+
+
+class FakeNode:
+ """One MI355X-shaped card on PCI bus 0x09, plus an untouched neighbour."""
+
+ def __init__(self, tmp_path):
+ self.mode = "SPX"
+ self.sets: list[str] = []
+ self.processes: object = [{"gpu": 0, "process_list": IDLE}, {"gpu": 1, "process_list": IDLE}]
+ self.output_dir = tmp_path / "out"
+ #: Raised by the next set of this mode, once. Simulates a mid-sweep fault.
+ self.fail_set_on: dict[str, BaseException] = {}
+ #: Raised instead of enumerating agents for this mode.
+ self.fail_agents_on: dict[str, BaseException] = {}
+
+ def summary(self) -> dict:
+ return json.loads((self.output_dir / "sweep_summary.json").read_text())
+
+
+@pytest.fixture
+def node(monkeypatch, tmp_path):
+ fake = FakeNode(tmp_path)
+
+ def _amd_smi_json(args, timeout_s=None, *, sudo=False):
+ if args[0] == "list":
+ return [{"gpu": 0, "bdf": "0000:09:00.0"}, {"gpu": 1, "bdf": "0000:0a:00.0"}]
+ if args[0] == "partition":
+ return [] # no profile table, so the mode list is taken as given
+ if args[0] == "process":
+ return fake.processes
+ raise AssertionError(f"unexpected amd-smi call: {args}")
+
+ def _set_mode(gpu_id, mode, *, sudo=False):
+ boom = fake.fail_set_on.pop(mode, None)
+ if boom is not None:
+ raise boom
+ fake.sets.append(mode)
+ fake.mode = mode
+ return mode
+
+ def _read_hsa_agents():
+ boom = fake.fail_agents_on.pop(fake.mode, None)
+ if boom is not None:
+ raise boom
+ partitions = pms.MODE_PARTITION_COUNTS[fake.mode]
+ cu = 256 // partitions
+ return tuple(pms.HsaAgent(index=i, cu=cu, bus=0x09, device=0, function=i) for i in range(partitions))
+
+ def _run_mode(layout, devices, **kwargs):
+ return [
+ pms.PartitionRun(i, d, tmp_path, returncode=0, measurement={"output_throughput": 100.0})
+ for i, d in enumerate(devices)
+ ]
+
+ monkeypatch.setattr(pms, "_amd_smi_json", _amd_smi_json)
+ monkeypatch.setattr(pms, "read_mode", lambda gpu_id: fake.mode)
+ monkeypatch.setattr(pms, "set_mode", _set_mode)
+ monkeypatch.setattr(pms, "read_hsa_agents", _read_hsa_agents)
+ monkeypatch.setattr(pms, "read_device_gib", lambda gpu_id: 288.0)
+ monkeypatch.setattr(pms, "run_mode", _run_mode)
+ # main() installs handlers and never removes them; leaving pytest's own
+ # SIGINT replaced for the rest of the session is not this test's business.
+ monkeypatch.setattr(pms.signal, "signal", lambda *a, **k: None)
+ return fake
+
+
+def _sweep(node, *extra, modes="SPX,DPX", gpu=0):
+ return pms.main(
+ [
+ "--benchmark-command",
+ "bench --gpu {device}",
+ "--modes",
+ modes,
+ "--gpu",
+ str(gpu),
+ "--output-dir",
+ str(node.output_dir),
+ *extra,
+ ]
+ )
+
+
+class TestBusyCheckScope:
+ """Only the swept card is repartitioned, so only its contexts are at risk.
+
+ The check was node-wide while the set is per-card, so any busy card on a
+ shared node blocked sweeping an idle one -- and the only way past it,
+ ``--allow-busy``, also gave up the protection on the target card itself.
+ """
+
+ def test_a_busy_neighbour_does_not_block_an_idle_target(self, node):
+ """The set never touches card 1, so card 1's tenant is not this sweep's business."""
+ node.processes = [
+ {"gpu": 0, "process_list": IDLE},
+ {"gpu": 1, "process_list": [{"process_info": {"name": "vllm"}}]},
+ ]
+ assert _sweep(node, gpu=0) == 0
+ assert node.sets # the sweep actually ran
+
+ def test_a_process_on_the_target_card_still_refuses(self, node):
+ node.processes = [{"gpu": 0, "process_list": [{"process_info": {"name": "vllm"}}]}]
+ assert _sweep(node, gpu=0) == 2
+ assert node.sets == []
+
+ def test_the_refusal_names_the_card_and_the_count(self, node, capsys):
+ node.processes = [{"gpu": 0, "process_list": [{"process_info": {"name": "vllm"}}]}]
+ _sweep(node, gpu=0)
+ assert "1 process(es) still hold a context on GPU 0" in capsys.readouterr().err
+
+ def test_a_target_missing_from_the_listing_is_refused_not_assumed_idle(self, node):
+ """An absent row is "unknown", and unknown does not license an eviction."""
+ node.processes = [{"gpu": 1, "process_list": IDLE}]
+ assert _sweep(node, gpu=0) == 2
+ assert node.sets == []
+
+ def test_allow_busy_proceeds_on_the_target_card(self, node):
+ node.processes = [{"gpu": 0, "process_list": [{"process_info": {"name": "vllm"}}]}]
+ assert _sweep(node, "--allow-busy", gpu=0) == 0
+
+ def test_an_unreadable_payload_refuses_by_default(self, node):
+ node.processes = {"process": [{"gpu": 0, "process_list": []}]}
+ assert _sweep(node, gpu=0) == 2
+ assert node.sets == []
+
+ def test_allow_busy_is_the_way_past_an_unreadable_payload(self, node):
+ """The check cannot change what happens next, so it is not consulted."""
+ node.processes = {"process": [{"gpu": 0, "process_list": []}]}
+ assert _sweep(node, "--allow-busy", gpu=0) == 0
+
+ def test_a_dry_run_does_not_consult_it_at_all(self, node):
+ """--dry-run sets nothing, so nothing it might evict is at stake."""
+ node.processes = {"unparseable": True}
+ assert _sweep(node, "--dry-run", gpu=0) == 0
+ assert node.sets == []
+
+
+class TestExitCodeContract:
+ def test_a_clean_sweep_reports_and_restores(self, node):
+ assert _sweep(node) == 0
+ assert node.mode == "SPX"
+ assert node.summary()["entry_mode"] == "SPX"
+
+ def test_an_expected_failure_moves_on_to_the_next_mode(self, node):
+ node.fail_agents_on["DPX"] = pms.SweepError("rocminfo timed out")
+ assert _sweep(node) == 0
+ assert node.summary()["modes"][0]["mode"] == "SPX"
+
+ def test_an_unexpected_error_still_renders_and_writes_the_summary(self, node):
+ """The finding: anything not a SweepError escaped main() past the report."""
+ node.fail_agents_on["DPX"] = KeyError("compute_partition")
+ assert _sweep(node) == 4
+ assert node.summary()["modes"][0]["mode"] == "SPX"
+
+ def test_an_unexpected_error_does_not_lose_the_modes_already_measured(self, node, capsys):
+ """SPX measured before DPX broke; escaping main() threw that away."""
+ node.fail_agents_on["DPX"] = KeyError("compute_partition")
+ _sweep(node)
+ out = capsys.readouterr().out
+ assert "Fastest measured mode: SPX" in out
+ assert node.summary()["modes"][0]["aggregate_throughput"] == pytest.approx(100.0)
+
+ def test_an_unexpected_error_still_restores_the_card(self, node):
+ node.fail_agents_on["DPX"] = KeyError("compute_partition")
+ _sweep(node)
+ assert node.mode == "SPX"
+
+ def test_an_unexpected_error_stops_the_sweep_rather_than_setting_more_modes(self, node):
+ """An assumption this script does not understand is broken; stop mutating."""
+ node.fail_agents_on["DPX"] = KeyError("compute_partition")
+ _sweep(node, modes="SPX,DPX,QPX")
+ assert "QPX" not in node.sets
+
+ def test_the_unexpected_error_is_named_not_swallowed(self, node, capsys):
+ node.fail_agents_on["DPX"] = KeyError("compute_partition")
+ _sweep(node)
+ assert "KeyError" in capsys.readouterr().err
+
+ def test_a_failed_restore_after_a_clean_sweep_exits_3(self, node):
+ # Entry mode is SPX and only DPX is swept, so the one set of SPX is the
+ # restore itself.
+ node.fail_set_on["SPX"] = pms.SweepError("permission denied")
+ assert _sweep(node, modes="DPX") == 3
+
+ def test_a_failed_restore_outranks_an_unexpected_error(self, node):
+ """A card left in the wrong shape mislabels whatever runs next, so it wins."""
+ node.fail_agents_on["QPX"] = KeyError("compute_partition")
+ node.fail_set_on["SPX"] = pms.SweepError("permission denied")
+ assert _sweep(node, modes="DPX,QPX") == 3
+
+ def test_an_unexpected_error_in_the_restore_is_still_a_restore_failure(self, node):
+ """Not a SweepError, so it used to propagate out of the finally and exit 1."""
+ node.fail_set_on["SPX"] = RuntimeError("amd-smi vanished")
+ assert _sweep(node, modes="DPX") == 3
+
+ def test_the_reported_reproduction_a_key_error_from_set_mode(self, node):
+ """As reviewed: KeyError out of set_mode escaped main() and exited 1."""
+ node.fail_set_on["DPX"] = KeyError("compute_partition")
+ assert _sweep(node) == 4
+ assert node.mode == "SPX"
+ assert node.summary()["modes"][0]["measured"] is True
+
+ def test_an_unwritable_report_does_not_cost_the_restore_failure_its_code(self, node, monkeypatch):
+ """The exit code is the one thing a caller cannot reconstruct itself."""
+ node.fail_set_on["SPX"] = pms.SweepError("permission denied")
+ monkeypatch.setattr(pms, "render", lambda results, *, entry_mode: (_ for _ in ()).throw(OSError("no space")))
+ assert _sweep(node, modes="DPX") == 3
+
+ def test_an_unwritable_report_is_not_reported_as_a_clean_sweep(self, node, monkeypatch):
+ monkeypatch.setattr(pms, "render", lambda results, *, entry_mode: (_ for _ in ()).throw(OSError("no space")))
+ assert _sweep(node) == 4
+
+ def test_a_sweep_that_measures_nothing_exits_1(self, node, monkeypatch):
+ monkeypatch.setattr(pms, "run_mode", lambda layout, devices, **kw: [])
+ assert _sweep(node) == 1
+
+ def test_an_unexpected_error_before_anything_is_set_does_not_report_a_refusal(self, node, monkeypatch):
+ """Exit 2 means this script decided to refuse; a bug is not a decision."""
+ monkeypatch.setattr(pms, "read_device_gib", lambda gpu_id: (_ for _ in ()).throw(KeyError("vram")))
+ assert _sweep(node) == 4
+ assert node.sets == []