From 89fa548e6edbf2f718035aaae1a79408027efd7e Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Mon, 24 Aug 2026 00:16:17 +0000 Subject: [PATCH 01/10] feat(partition): AMD compute-partition modes as an optimizer lever Adds SPX/DPX/QPX/CPX as a first-class lever so a throughput-sensitive session can trade single-stream latency for aggregate throughput, and adds the constraint that keeps that trade honest. Off by default: with no ``--compute-partition-modes`` the hardware is never touched and the KEEP ladder is unchanged. Measured on one MI355X over HunyuanWorld-Mirror's 30 example scenes, all five configurations run through the production ``run_scriptable`` path: SPX x1 @2 12.24 fwd/s 163 ms (anchor) DPX x2 @2 14.14 fwd/s 283 ms +15.5% QPX x4 @2 14.64 fwd/s 543 ms +19.6% CPX x8 @1 14.66 fwd/s 540 ms +19.8% CPX x8 @2 OOM, predicted from the per-stream footprint Throughput saturates long before the hardware does: partitioning past DPX buys 3.7% for 1.9x the latency. That shape is why the lever ships with a budget rather than alone -- unconstrained, the search picks CPX and makes every request three times slower to get there. * ``common/gpu_partition.py`` is the hardware boundary, and enforces the two invariants that make a partition measurement trustworthy. A set is not a set until it reads back: ``amd-smi set`` reports success for a change that was only staged, and exits 0 on a permission failure, so every mutation re-reads and compares. And partitions are not identified by device index: HIP enumerates whole cards first, so on one card of eight split into DPX the partitions are devices 7 and 8, and selecting by index measures a full card while labelling it a partition. Callers get a CU-count predicate instead. * A card refuses to repartition while a process is resident, and teardown is not synchronous -- ``docker rm -f`` returns before the runtime has released the device. The set waits that refusal out, matching on ``AMDSMI_STATUS_BUSY``. Treating it as fatal cost a sweep two configurations: the restore failed, left a shared card in QPX, and the next run inherited that as the state to restore to. * ``_latency_budget.py`` gates KEEP on ``--max-latency-ms`` between the throughput and accuracy checks. Absolute, because an SLA is; fail-closed on an unmeasured latency, because a throughput win with no latency evidence is exactly the case the budget exists to stop. * ``_partition_lever.py`` carries the mode as ordinary variant env, so fingerprinting, dedup and the journal need no changes, and holds the hardware only around the benchmark via a context manager that restores on both paths. Device enumeration stays with the benchmark script, which is the only layer that can see the partitions. Co-Authored-By: Claude Opus 5 Co-authored-by: Cursor (cherry picked from commit 2900bdbe3440c20e0137d80848d36b73c2fb0d7f) --- src/hyperloom/common/gpu_partition.py | 462 ++++++++++++++++++ .../common/tests/test_gpu_partition.py | 311 ++++++++++++ .../inference_optimizer/cli/__init__.py | 89 ++++ .../inference_optimizer/cli/bootstrap.py | 9 + .../inference_optimizer/cli/parser.py | 44 ++ .../tests/test_partition_lever.py | 193 ++++++++ .../actions/executors/_latency_budget.py | 143 ++++++ .../actions/executors/_partition_lever.py | 194 ++++++++ .../actions/executors/bypass_scriptable.py | 74 ++- .../orchestrator/actions/executors/explore.py | 24 + .../orchestrator/state/shared_state.py | 11 + 11 files changed, 1533 insertions(+), 21 deletions(-) create mode 100644 src/hyperloom/common/gpu_partition.py create mode 100644 src/hyperloom/common/tests/test_gpu_partition.py create mode 100644 src/hyperloom/inference_optimizer/tests/test_partition_lever.py create mode 100644 src/hyperloom/orchestrator/actions/executors/_latency_budget.py create mode 100644 src/hyperloom/orchestrator/actions/executors/_partition_lever.py diff --git a/src/hyperloom/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py new file mode 100644 index 0000000000..9af9a11246 --- /dev/null +++ b/src/hyperloom/common/gpu_partition.py @@ -0,0 +1,462 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""AMD compute-partition modes (SPX/DPX/QPX/CPX) as a hardware-level lever. + +Every other lever in the optimizer is an environment variable or a server flag: +it travels in a ``GridVariant``, applies by materializing a benchmark YAML, and +reverts by not setting it again. A compute-partition mode is none of those +things. It is privileged, it is global to the card, it evicts every process +resident on that card, and it renumbers the devices underneath a running +process. So it gets its own module rather than a row in an existing table. + +What the modes buy, and what they cost: partitioning a card only ever gives a +single stream *fewer* CUs, so it cannot improve single-stream latency and will +always make it worse. It pays only in aggregate, when there are at least as many +concurrent streams as partitions -- which is why ``streams_per_partition`` +belongs next to the mode in any configuration that sets one. Measured on one +MI355X with a 1.26B-parameter vision model, CPX at two streams per partition +carried ~20% more aggregate throughput than the best SPX configuration while +per-request latency went from 183 ms to 1211 ms. That trade is a decision for +whoever owns the SLA, never a default. + +Two invariants this module exists to enforce: + +* **A set is not a set until it reads back.** ``amd-smi set`` reports success + for a memory-partition change that has only been staged pending a driver + reload, so trusting the exit code silently measures the old topology as if it + were the new one. Every mutation here re-reads the mode and compares. +* **Partitions are not identified by device index.** HIP enumerates whole GPUs + before partitions, so with one card of eight split into DPX the partitions are + devices 7 and 8, not 0 and 1. Selecting by index measures a full card and + reports it as a partition -- a wrong number with no error attached. Callers + get :func:`partition_device_predicate` for a CU-count test instead. +""" + +from __future__ import annotations + +import json +import logging +import os +import subprocess +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Iterator, Sequence + +from .gpu_identity import AMD_GPU_DISPATCH_IDENTITIES + +log = logging.getLogger(__name__) + +#: Compute-partition mode -> partitions per card. Both gfx942 and gfx950 boards +#: carry eight XCDs, so the ladder is the same width on each; the CU count per +#: partition differs and is derived from the board's own total. +MODE_PARTITION_COUNTS: dict[str, int] = { + "SPX": 1, + "DPX": 2, + "QPX": 4, + "CPX": 8, +} + +#: Mode a session is restored to. SPX is the only mode every board supports and +#: the only one compatible with NPS1, so it is the safe terminal state. +DEFAULT_MODE = "SPX" + +_SET_TIMEOUT_S = 120.0 +_READ_TIMEOUT_S = 30.0 + +#: How long a set waits out a card that still holds processes, and how often it +#: retries. The card genuinely refuses to repartition while a process is +#: resident, but teardown is not synchronous: ``docker rm -f`` returns before +#: the runtime has released the device, so a repartition issued immediately +#: afterwards races it. Treating that refusal as fatal cost a sweep two +#: configurations -- the restore after one run failed and left a shared card in +#: QPX, and the next run then could not set its own mode. The condition clears +#: on its own in seconds, so the correct response is to wait for it. +_DRAIN_TIMEOUT_S = 120.0 +_DRAIN_POLL_S = 2.0 + +#: Substrings that mark a refusal as "busy, try again" rather than "wrong". +#: Matching narrowly keeps permanent failures -- an unknown mode, a missing +#: binary, a denied permission -- fast instead of retrying them for two minutes. +#: ``AMDSMI_STATUS_BUSY`` is the authoritative one and the reason the others are +#: only a safety net: the library reports the refusal as status code 30, and the +#: human-readable half of the message ("Device busy") is not stable enough to +#: match on alone. +_BUSY_MARKERS = ("amdsmi_status_busy", "device busy", "resident process", "try again") + +#: Opt-in elevation for the set path. Reading the partition state is +#: unprivileged; changing it is not, so a session running as an ordinary user +#: needs a way to say how. Set to ``1`` to route the set through +#: ``sudo -n``, which requires a NOPASSWD sudoers entry for ``amd-smi``. +#: +#: Off by default and never inferred. Escalating privilege because a command +#: failed is not a fallback, it is a decision, and it belongs to the operator. +PARTITION_SUDO_ENV = "HYPERLOOM_PARTITION_SUDO" + + +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) + + +class PartitionError(RuntimeError): + """Raised when a partition mode cannot be read, set, or verified.""" + + +@dataclass(frozen=True) +class PartitionLayout: + """What one compute-partition mode does to one card. + + Attributes: + mode: Canonical mode name (``SPX``/``DPX``/``QPX``/``CPX``). + partitions: Independent devices the card presents in this mode. + cu_per_partition: Compute units each partition gets. + gib_per_partition: HBM each partition gets, or ``None`` when the caller + did not supply the card's capacity. + """ + + mode: str + partitions: int + cu_per_partition: int + gib_per_partition: float | None = None + + def describe(self) -> str: + """Return a one-line summary for logs and reports.""" + mem = f", {self.gib_per_partition:.0f} GiB" if self.gib_per_partition else "" + return f"{self.mode} ({self.partitions} x {self.cu_per_partition} CU{mem})" + + +def parse_modes(raw: str | Sequence[str] | None) -> tuple[str, ...]: + """Parse an operator-supplied mode list into canonical order-preserving modes. + + Args: + raw: Comma-separated string (``"spx,dpx"``) or an already-split + sequence. Empty or ``None`` means the lever is off. + + Returns: + Canonical upper-case mode names, deduplicated, in the order given. + + Raises: + PartitionError: If any entry is not a known mode. Refused at parse time + rather than at apply time, because the apply site is a privileged + hardware mutation partway through a session. + """ + if not raw: + return () + items = raw.split(",") if isinstance(raw, str) else list(raw) + modes: list[str] = [] + for item in items: + name = str(item).strip().upper() + if not name: + continue + if name not in MODE_PARTITION_COUNTS: + raise PartitionError( + f"unknown compute-partition mode {name!r}; expected one of {', '.join(MODE_PARTITION_COUNTS)}" + ) + if name not in modes: + modes.append(name) + return tuple(modes) + + +def layout_for(gpu_type: str | None, mode: str, hbm_gib: float | None = None) -> PartitionLayout: + """Describe what ``mode`` does to a ``gpu_type`` card. + + CU counts come from :data:`AMD_GPU_DISPATCH_IDENTITIES` rather than a second + table, so a board added there is described here without a further edit. + + Args: + gpu_type: Board name (``mi300x``, ``mi355x``, ...). + mode: Compute-partition mode. + hbm_gib: The card's total HBM, when known. Supplied by the caller from a + live device rather than tabled here, since capacity varies across + boards that share an ISA. + + Returns: + The resulting layout. + + Raises: + PartitionError: If the mode is unknown or the board is unrecognised. + """ + canonical = str(mode or "").strip().upper() + partitions = MODE_PARTITION_COUNTS.get(canonical) + if partitions is None: + raise PartitionError(f"unknown compute-partition mode {mode!r}") + identity = AMD_GPU_DISPATCH_IDENTITIES.get(str(gpu_type or "").strip().lower()) + if identity is None: + raise PartitionError(f"unknown gpu_type {gpu_type!r}; cannot size partitions without the board's CU count") + cu_total = identity[1] + return PartitionLayout( + mode=canonical, + partitions=partitions, + cu_per_partition=cu_total // partitions, + gib_per_partition=(hbm_gib / partitions) if hbm_gib else None, + ) + + +def partition_device_predicate(cu_per_partition: int): + """Return a predicate selecting partition devices by CU count. + + The index-based alternative is what makes a partitioning measurement quietly + wrong: HIP enumerates whole cards first, so under DPX on one card of eight + ``device 0`` is a full 256-CU GPU and the partitions are devices 7 and 8. + Callers apply this to ``torch.cuda.get_device_properties(i).multi_processor_count`` + (kept out of this module so it stays importable without torch). + + Args: + cu_per_partition: CU count a real partition must report. + + Returns: + A callable taking a device's CU count and returning whether it is a + partition of the requested shape. + """ + + def _is_partition(device_cu: int) -> bool: + return int(device_cu) == int(cu_per_partition) + + return _is_partition + + +def fits_in_partition( + required_gib: float, + layout: PartitionLayout, + streams_per_partition: int = 1, +) -> bool: + """Report whether the streams sharing one partition fit in its memory. + + This is the real criterion behind "partitioning helps smaller models": a + partition gets a fraction of the card's HBM, and every stream on it holds + its own copy of the weights plus activations. A model needing the whole card + cannot be partitioned at all, however throughput-bound it is. + + ``streams_per_partition`` is not a refinement -- it is the question. A mode + is only worth setting at two streams per partition, and one stream fitting + says nothing about two: a 20 GiB footprint fits a 36 GiB CPX partition alone + and exhausts it in pairs. Gating on the single-stream figure is how a + configuration gets declared feasible and then dies at the second worker. + + Args: + required_gib: Peak footprint of one stream, weights included. + layout: The layout under consideration. + streams_per_partition: Concurrent streams intended per partition. + + Returns: + ``True`` when they fit, or when capacity is unknown -- an unknown is + reported by the caller that has the number, not guessed at here. + """ + if not layout.gib_per_partition or required_gib <= 0: + return True + return required_gib * max(1, int(streams_per_partition)) <= layout.gib_per_partition + + +def _amd_smi_json(args: Sequence[str], timeout_s: float) -> object: + """Run an ``amd-smi`` subcommand with ``--json`` and parse its output.""" + cmd = ["amd-smi", *args, "--json"] + try: + proc = subprocess.run( # noqa: S603 — fixed argv, no shell + cmd, + capture_output=True, + text=True, + timeout=timeout_s, + check=False, + ) + except FileNotFoundError as exc: + raise PartitionError("amd-smi not found; compute partitioning needs it on PATH") from exc + except subprocess.TimeoutExpired as exc: + raise PartitionError(f"amd-smi {' '.join(args)} timed out after {timeout_s}s") from exc + if proc.returncode != 0: + raise PartitionError(f"amd-smi {' '.join(args)} failed ({proc.returncode}): {proc.stderr.strip()}") + try: + return json.loads(proc.stdout) + except ValueError as exc: + raise PartitionError(f"amd-smi {' '.join(args)} returned unparseable JSON") from exc + + +def _set_prefix() -> list[str]: + """Return the command prefix for the privileged set (``sudo -n`` or none). + + ``-n`` matters: an automated optimization loop that stops at an interactive + password prompt hangs until its budget expires, with no indication why. + """ + if os.environ.get(PARTITION_SUDO_ENV, "").strip().lower() in ("1", "true", "yes", "on"): + return ["sudo", "-n"] + return [] + + +def read_partition_modes() -> dict[int, str]: + """Read the live compute-partition mode of every GPU. + + Returns: + Mapping of GPU id to mode name. + + Raises: + PartitionError: If ``amd-smi`` is absent, fails, or reports no cards. + """ + payload = _amd_smi_json(["partition"], _READ_TIMEOUT_S) + rows: list[dict] = [] + if isinstance(payload, dict): + raw_rows = payload.get("current_partition") + if isinstance(raw_rows, list): + rows = [r for r in raw_rows 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() + if mode: + modes[gpu_id] = mode + if not modes: + raise PartitionError("amd-smi partition reported no compute-partition state") + return modes + + +def read_partition_mode(gpu_id: int) -> str: + """Read one GPU's compute-partition mode. + + Raises: + PartitionError: If that GPU is not present in the report. + """ + modes = read_partition_modes() + if gpu_id not in modes: + raise PartitionError(f"amd-smi partition reported no state for GPU {gpu_id}") + return modes[gpu_id] + + +def set_partition_mode(gpu_id: int, mode: str, drain_timeout_s: float = _DRAIN_TIMEOUT_S) -> str: + """Set one GPU's compute-partition mode and verify it took effect. + + The verification is the point of this function. ``amd-smi set`` reports + success for a change that has only been staged, so a caller trusting the + exit code proceeds to benchmark the old topology while labelling the results + with the new mode -- a wrong number with a reassuring log line above it. + + A card still holding processes refuses to repartition, which is a race + rather than a wall: see :data:`_DRAIN_TIMEOUT_S`. + + Args: + gpu_id: GPU to reconfigure. + mode: Target mode. + drain_timeout_s: How long to keep retrying while the card reports + resident processes. Zero fails on the first refusal. + + Returns: + The mode read back from the hardware, which equals ``mode`` on success. + + Raises: + PartitionError: If the mode is unknown, the set fails, or the read-back + disagrees with what was requested. + """ + canonical = str(mode or "").strip().upper() + if canonical not in MODE_PARTITION_COUNTS: + raise PartitionError(f"unknown compute-partition mode {mode!r}") + + current = read_partition_mode(gpu_id) + if current == canonical: + return current + + cmd = [*_set_prefix(), "amd-smi", "set", "-g", str(gpu_id), "--compute-partition", canonical] + deadline = time.monotonic() + max(0.0, drain_timeout_s) + waited_for_drain = False + while True: + try: + proc = subprocess.run( # noqa: S603 — fixed argv, no shell + cmd, + capture_output=True, + text=True, + timeout=_SET_TIMEOUT_S, + # Some builds prompt for confirmation; answer it rather than block. + input="Y\n", + check=False, + ) + except FileNotFoundError as exc: + raise PartitionError("amd-smi not found; compute partitioning needs it on PATH") from exc + except subprocess.TimeoutExpired as exc: + raise PartitionError(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_for_drain: + log.info("GPU %d busy; waiting for it to drain before setting %s", gpu_id, canonical) + waited_for_drain = 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 PartitionError(f"setting {canonical} on GPU {gpu_id} failed ({proc.returncode}): {detail}.{hint}") + + observed = read_partition_mode(gpu_id) + if observed != canonical: + raise PartitionError( + f"GPU {gpu_id} reports {observed} after being set to {canonical}. The " + f"command returned success but the mode did not change, so any " + f"measurement now would be attributed to the wrong mode. Two known " + f"causes: amd-smi exits 0 on a permission failure (set " + f"{PARTITION_SUDO_ENV}=1 for a NOPASSWD sudo path), and a memory " + f"partition change needs an amdgpu reload to take effect." + ) + log.info("GPU %d compute partition: %s -> %s", gpu_id, current, canonical) + return observed + + +@contextmanager +def partitioned(gpu_id: int, mode: str, restore_to: str | None = None) -> Iterator[str]: + """Hold a GPU in ``mode`` for the duration of the block, then restore it. + + Restoration runs on the way out of both the happy and the failing path, + because the alternative is leaving a shared card in a mode the next tenant + did not ask for. A failure to restore is logged rather than raised, so it + cannot mask the exception that caused the exit. + + Args: + gpu_id: GPU to reconfigure. + mode: Mode to hold during the block. + restore_to: Mode to return to; defaults to whatever was observed on + entry, falling back to :data:`DEFAULT_MODE`. + + Yields: + The mode the hardware confirmed. + """ + try: + entry_mode = read_partition_mode(gpu_id) + except PartitionError: + entry_mode = DEFAULT_MODE + target_restore = str(restore_to or entry_mode or DEFAULT_MODE).strip().upper() + observed = set_partition_mode(gpu_id, mode) + try: + yield observed + finally: + if target_restore != observed: + try: + set_partition_mode(gpu_id, target_restore) + except PartitionError as exc: + log.error( + "GPU %d left in %s: restore to %s failed: %s", + gpu_id, + observed, + target_restore, + exc, + ) + + +__all__ = [ + "DEFAULT_MODE", + "MODE_PARTITION_COUNTS", + "PARTITION_SUDO_ENV", + "PartitionError", + "PartitionLayout", + "fits_in_partition", + "layout_for", + "parse_modes", + "partition_device_predicate", + "partitioned", + "read_partition_mode", + "read_partition_modes", + "set_partition_mode", +] diff --git a/src/hyperloom/common/tests/test_gpu_partition.py b/src/hyperloom/common/tests/test_gpu_partition.py new file mode 100644 index 0000000000..73a237353d --- /dev/null +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Tests for the compute-partition lever. + +``amd-smi`` is faked throughout, so these run on a CPU host and on a card that +must not be repartitioned by a test suite. The cases that matter are the two +silent-wrong-answer paths: a set that reports success without taking effect, and +a feasibility check that passes one stream where two will run. +""" + +from __future__ import annotations + +import json +import subprocess + +import pytest + +from hyperloom.common import gpu_partition +from hyperloom.common.gpu_partition import ( + DEFAULT_MODE, + PartitionError, + fits_in_partition, + layout_for, + parse_modes, + partition_device_predicate, + partitioned, + read_partition_mode, + read_partition_modes, + set_partition_mode, +) + + +class _FakeSmi: + """An ``amd-smi`` whose partition state a test can drive. + + ``stage_only`` reproduces the mode change that is accepted and then does not + take effect, which is the failure the read-back exists to catch. + """ + + def __init__(self, modes: dict[int, str], *, stage_only: bool = False, set_rc: int = 0): + self.modes = dict(modes) + self.stage_only = stage_only + self.set_rc = set_rc + self.set_calls: list[tuple[int, str]] = [] + self.argv: list[list[str]] = [] + + def __call__(self, cmd, **kwargs): + if "set" in cmd and "--compute-partition" in cmd: + self.argv.append(list(cmd)) + gpu_id = int(cmd[cmd.index("-g") + 1]) + mode = cmd[cmd.index("--compute-partition") + 1] + self.set_calls.append((gpu_id, mode)) + if self.set_rc == 0 and not self.stage_only: + self.modes[gpu_id] = mode + return subprocess.CompletedProcess(cmd, self.set_rc, "", "busy" if self.set_rc else "") + rows = [ + {"gpu_id": gid, "memory": "NPS1", "accelerator_type": mode, "partition_id": "0"} + for gid, mode in sorted(self.modes.items()) + ] + payload = json.dumps({"current_partition": rows}) + return subprocess.CompletedProcess(cmd, 0, payload, "") + + +@pytest.fixture +def smi(monkeypatch): + """Install a fake amd-smi with all eight cards in SPX.""" + fake = _FakeSmi({i: "SPX" for i in range(8)}) + monkeypatch.setattr(subprocess, "run", fake) + return fake + + +def test_parse_modes_canonicalizes_dedupes_and_keeps_order(): + assert parse_modes("spx, dpx ,cpx,dpx") == ("SPX", "DPX", "CPX") + assert parse_modes(["CPX", "spx"]) == ("CPX", "SPX") + + +def test_parse_modes_off_by_default(): + assert parse_modes(None) == () + assert parse_modes("") == () + assert parse_modes(" , ") == () + + +def test_parse_modes_refuses_unknown_mode_at_parse_time(): + with pytest.raises(PartitionError, match="unknown compute-partition mode 'NPS2'"): + parse_modes("spx,nps2") + + +@pytest.mark.parametrize( + ("mode", "partitions", "cu"), + [("SPX", 1, 256), ("DPX", 2, 128), ("QPX", 4, 64), ("CPX", 8, 32)], +) +def test_layout_matches_mi355x_hardware(mode, partitions, cu): + layout = layout_for("mi355x", mode, hbm_gib=288.0) + assert (layout.partitions, layout.cu_per_partition) == (partitions, cu) + assert layout.gib_per_partition == pytest.approx(288.0 / partitions) + + +def test_layout_derives_cu_from_the_board_not_a_second_table(): + # mi300x carries 304 CU, so its CPX partitions are 38 CU, not mi355x's 32. + assert layout_for("mi300x", "CPX").cu_per_partition == 38 + + +def test_layout_refuses_unknown_board_and_mode(): + with pytest.raises(PartitionError, match="unknown gpu_type"): + layout_for("mi999x", "SPX") + with pytest.raises(PartitionError, match="unknown compute-partition mode"): + layout_for("mi355x", "OPX") + + +def test_fits_in_partition_gates_on_the_streams_that_will_run(): + cpx = layout_for("mi355x", "CPX", hbm_gib=288.0) + # 20 GiB fits a 36 GiB partition alone and exhausts it in pairs. Gating on + # the single-stream figure is what lets a doomed config be declared feasible. + assert fits_in_partition(20.0, cpx, streams_per_partition=1) + assert not fits_in_partition(20.0, cpx, streams_per_partition=2) + # The 6-view footprint that did run at two streams per CPX partition. + assert fits_in_partition(8.3, cpx, streams_per_partition=2) + + +def test_fits_in_partition_is_permissive_when_capacity_is_unknown(): + assert fits_in_partition(999.0, layout_for("mi355x", "CPX"), streams_per_partition=8) + + +def test_partition_device_predicate_selects_by_cu_not_index(): + is_dpx = partition_device_predicate(layout_for("mi355x", "DPX").cu_per_partition) + # Under DPX on one card of eight, devices 0-6 are whole 256-CU cards and the + # partitions are the two 128-CU devices enumerated after them. + enumerated = [256] * 7 + [128, 128] + assert [i for i, cu in enumerate(enumerated) if is_dpx(cu)] == [7, 8] + + +def test_read_partition_modes_parses_amd_smi_json(smi): + assert read_partition_modes() == {i: "SPX" for i in range(8)} + assert read_partition_mode(3) == "SPX" + + +def test_read_partition_mode_refuses_absent_gpu(smi): + with pytest.raises(PartitionError, match="no state for GPU 99"): + read_partition_mode(99) + + +def test_set_partition_mode_verifies_the_read_back(smi): + assert set_partition_mode(0, "cpx") == "CPX" + assert smi.set_calls == [(0, "CPX")] + + +def test_set_partition_mode_refuses_a_change_that_did_not_take_effect(monkeypatch): + # The NPS2 failure shape: exit code 0, mode unchanged. Trusting the exit + # code here attributes a whole measurement to a mode that never applied. + fake = _FakeSmi({0: "SPX"}, stage_only=True) + monkeypatch.setattr(subprocess, "run", fake) + with pytest.raises(PartitionError, match="reports SPX after being set to CPX"): + set_partition_mode(0, "CPX") + + +#: Verbatim stderr from a card that refused to repartition while a container was +#: still shutting down. The retry matches on the status code because the prose +#: half of this message is not stable across amd-smi builds. +BUSY_STDERR = ( + "amdsmi.amdsmi_exception.AmdSmiLibraryException: Error code:\n" + "\t30 | AMDSMI_STATUS_BUSY - Device busy\n\n" + "The above exception was the direct cause of the following exception:\n\n" + "ValueError: Unable to set accelerator partition to CPX on GPU ID: 0 BDF:0000:09:00.0." +) + + +def _busy_until(succeed_on: int, fake: _FakeSmi): + """An amd-smi that refuses ``succeed_on - 1`` sets, then behaves.""" + attempts = {"n": 0} + + def run(cmd, **kwargs): + if "set" in cmd and "--compute-partition" in cmd: + attempts["n"] += 1 + if attempts["n"] < succeed_on: + return subprocess.CompletedProcess(cmd, 1, "", BUSY_STDERR) + return fake(cmd, **kwargs) + + return run, attempts + + +def test_set_partition_mode_waits_for_a_busy_card_to_drain(monkeypatch): + # docker rm -f returns before the runtime has released the device, so the + # repartition that follows it races teardown. The refusal clears on its own. + fake = _FakeSmi({0: "SPX"}) + run, attempts = _busy_until(3, fake) + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setattr(gpu_partition.time, "sleep", lambda _s: None) + + assert set_partition_mode(0, "CPX") == "CPX" + assert attempts["n"] == 3 + + +def test_set_partition_mode_gives_up_on_a_card_that_never_drains(monkeypatch): + fake = _FakeSmi({0: "SPX"}) + run, attempts = _busy_until(10**6, fake) + monkeypatch.setattr(subprocess, "run", run) + with pytest.raises(PartitionError, match="still held processes"): + set_partition_mode(0, "CPX", drain_timeout_s=0.0) + assert attempts["n"] == 1 + + +def test_set_partition_mode_does_not_retry_a_permanent_failure(monkeypatch): + # Retrying an unknown flag or a denied permission for two minutes turns a + # clear error into a hang, so only a busy card is waited out. + fake = _FakeSmi({0: "SPX"}) + attempts = {"n": 0} + + def denied(cmd, **kwargs): + if "set" in cmd and "--compute-partition" in cmd: + attempts["n"] += 1 + return subprocess.CompletedProcess(cmd, 1, "", "permission denied") + return fake(cmd, **kwargs) + + monkeypatch.setattr(subprocess, "run", denied) + with pytest.raises(PartitionError, match="permission denied"): + set_partition_mode(0, "CPX", drain_timeout_s=60.0) + assert attempts["n"] == 1 + + +def test_restore_also_waits_out_a_busy_card(monkeypatch): + # The restore is the path that matters most: a failure there leaves a shared + # card in a mode its next tenant did not ask for. + fake = _FakeSmi({0: "SPX"}) + attempts = {"n": 0} + + def busy_on_restore(cmd, **kwargs): + if "set" in cmd and "--compute-partition" in cmd: + attempts["n"] += 1 + # Refuse the first restore attempt only. + if attempts["n"] == 2: + return subprocess.CompletedProcess(cmd, 1, "", BUSY_STDERR) + return fake(cmd, **kwargs) + + monkeypatch.setattr(subprocess, "run", busy_on_restore) + monkeypatch.setattr(gpu_partition.time, "sleep", lambda _s: None) + with partitioned(0, "QPX"): + pass + assert fake.modes[0] == "SPX" + + +def test_set_partition_mode_is_a_no_op_when_already_there(smi): + assert set_partition_mode(0, "SPX") == "SPX" + assert smi.set_calls == [] + + +def test_set_partition_mode_refuses_unknown_mode_without_touching_hardware(smi): + with pytest.raises(PartitionError): + set_partition_mode(0, "NPS2") + assert smi.set_calls == [] + + +def test_partitioned_restores_the_entry_mode(smi): + with partitioned(0, "CPX") as mode: + assert mode == "CPX" + assert smi.modes[0] == "CPX" + assert smi.modes[0] == "SPX" + + +def test_partitioned_restores_after_a_failure_inside_the_block(smi): + with pytest.raises(ZeroDivisionError): + with partitioned(0, "QPX"): + assert smi.modes[0] == "QPX" + raise ZeroDivisionError + assert smi.modes[0] == "SPX" + + +def test_partitioned_restore_failure_does_not_mask_the_real_error(monkeypatch, caplog): + fake = _FakeSmi({0: "SPX"}) + + def flaky(cmd, **kwargs): + # Let the entry set through, then refuse every later set. + if cmd[:2] == ["amd-smi", "set"] and fake.modes.get(0) != "SPX": + # A permanent refusal, so the restore fails immediately instead of + # waiting out the drain timeout. + return subprocess.CompletedProcess(cmd, 1, "", "permission denied") + return fake(cmd, **kwargs) + + monkeypatch.setattr(subprocess, "run", flaky) + with pytest.raises(ValueError, match="workload blew up"): + with partitioned(0, "DPX"): + raise ValueError("workload blew up") + assert "left in DPX" in caplog.text + + +def test_partitioned_honours_an_explicit_restore_target(smi): + with partitioned(0, "CPX", restore_to=DEFAULT_MODE): + pass + assert smi.modes[0] == DEFAULT_MODE + + +def test_set_is_unprivileged_by_default(smi, monkeypatch): + monkeypatch.delenv("HYPERLOOM_PARTITION_SUDO", raising=False) + set_partition_mode(0, "CPX") + assert smi.argv[0][:1] == ["amd-smi"] + + +def test_set_routes_through_sudo_when_opted_in(smi, monkeypatch): + monkeypatch.setenv("HYPERLOOM_PARTITION_SUDO", "1") + set_partition_mode(0, "CPX") + # -n so an automated loop fails fast instead of waiting on a password prompt. + assert smi.argv[0][:3] == ["sudo", "-n", "amd-smi"] + + +def test_permission_failure_that_exits_zero_is_still_caught(monkeypatch): + # Observed on ROCm 7.x: amd-smi prints AmdSmiPermissionDeniedException and + # exits 0. Only the read-back distinguishes this from a real change. + fake = _FakeSmi({0: "SPX"}, stage_only=True) + monkeypatch.setattr(subprocess, "run", fake) + with pytest.raises(PartitionError, match="returned success but the mode did not change"): + set_partition_mode(0, "CPX") diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 775d439b38..54f222cc93 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -65,6 +65,7 @@ resolve_model_display_name, ) from hyperloom.orchestrator.actions.executors._aiter_jit import clean_stale_aiter_locks +from hyperloom.orchestrator.actions.executors._latency_budget import LATENCY_BUDGET_ENV from hyperloom.orchestrator.actions.executors._workload_envs import ( agentx_enabled as _agentx_enabled, ) @@ -1631,6 +1632,79 @@ def _export_operator_launch_shape( os.environ.pop("INFERENCE_OPTIMIZER_EXTRA_ENV", None) +#: Internal handoff for the compute-partition lever. Operators pass +#: ``--compute-partition-modes`` / ``--streams-per-partition`` / +#: ``--max-latency-ms``; these carry the resolved values to the executors. +PARTITION_MODES_ENV = "HYPERLOOM_COMPUTE_PARTITION_MODES" +STREAMS_PER_PARTITION_ENV = "HYPERLOOM_STREAMS_PER_PARTITION" + + +def _export_partition_lever( + *, + modes_raw: str | None, + streams_per_partition: int, + max_latency_ms: float | None, +) -> tuple[str, ...]: + """Validate and project the compute-partition lever into env. + + Validation happens here, at launch, rather than at the apply site: applying + a mode is a privileged mutation of a shared card partway through a session, + and discovering a typo there means unwinding a live topology instead of + printing a usage error. + + Empty inputs clear the variables, so a second session in the same shell + cannot inherit a partition lever the operator did not ask for this time. + + Args: + modes_raw: The raw ``--compute-partition-modes`` value. + streams_per_partition: The resolved ``--streams-per-partition``. + max_latency_ms: The resolved ``--max-latency-ms``, if any. + + Returns: + The canonical modes, empty when the lever is off. + """ + from hyperloom.common.gpu_partition import PartitionError, parse_modes + + budget = float(max_latency_ms or 0.0) + if budget > 0: + os.environ[LATENCY_BUDGET_ENV] = repr(budget) + else: + os.environ.pop(LATENCY_BUDGET_ENV, None) + + try: + modes = parse_modes(modes_raw) + except PartitionError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(2) + + if not modes: + os.environ.pop(PARTITION_MODES_ENV, None) + os.environ.pop(STREAMS_PER_PARTITION_ENV, None) + return () + + if streams_per_partition < 1: + print( + f"ERROR: --streams-per-partition must be >= 1, got {streams_per_partition}", + file=sys.stderr, + ) + sys.exit(2) + if budget <= 0: + # Not fatal: maximizing offline throughput regardless of per-request + # latency is a legitimate goal. But it is not usually what someone + # means, and throughput is the only gate without a budget, so the + # search will end up choosing the narrowest partition offered. + print( + "WARN: --compute-partition-modes without --max-latency-ms. Narrower " + "partitions raise aggregate throughput by making each stream " + "slower, and throughput is the only KEEP gate, so the search is " + "free to trade away per-request latency without bound.", + file=sys.stderr, + ) + os.environ[PARTITION_MODES_ENV] = ",".join(modes) + os.environ[STREAMS_PER_PARTITION_ENV] = str(int(streams_per_partition)) + return modes + + # Terminal stop_reasons that represent a clean, successful optimizer run (exit 0). # Anything else (baseline / preflight failures, crashes, enablement stalls) exits # non-zero so CI surfaces genuine problems. @@ -1715,6 +1789,11 @@ async def _run_optimize(args: argparse.Namespace) -> int: server_args=str(getattr(args, "server_args", "") or "").strip(), extra_env=parse_operator_extra_env(args), ) + _export_partition_lever( + modes_raw=getattr(args, "compute_partition_modes", None), + streams_per_partition=int(getattr(args, "streams_per_partition", 2) or 2), + max_latency_ms=getattr(args, "max_latency_ms", None), + ) # Project resolved workload knobs into env for the fresh-launch path only. # A resume must NOT export here: ``args.tp``/etc. are still unresolved # (``None`` -> 1) because the persisted SharedState is loaded later; the @@ -1934,6 +2013,16 @@ async def _run_optimize(args: argparse.Namespace) -> int: server_args=_resume_server_args, extra_env=_resume_extra_env, ) + # The lever is part of the measurement contract: a resume that dropped + # it would compare candidates measured under partitioning against a + # baseline that no longer is. + _export_partition_lever( + modes_raw=getattr(args, "compute_partition_modes", None) or state.compute_partition_modes, + streams_per_partition=int( + getattr(args, "streams_per_partition", 0) or state.streams_per_partition or 2 + ), + max_latency_ms=getattr(args, "max_latency_ms", None) or state.latency_budget_ms, + ) state.operator_server_args = _resume_server_args state.operator_extra_env = _resume_extra_env if _resume_server_args: diff --git a/src/hyperloom/inference_optimizer/cli/bootstrap.py b/src/hyperloom/inference_optimizer/cli/bootstrap.py index 6aa5d5400e..5b878629f6 100644 --- a/src/hyperloom/inference_optimizer/cli/bootstrap.py +++ b/src/hyperloom/inference_optimizer/cli/bootstrap.py @@ -355,6 +355,15 @@ def _resolve_framework_version(args_in: Any) -> str: bypass_scripts_dir=os.environ.get("HYPERLOOM_BYPASS_SCRIPTS_DIR", "").strip(), framework_repo_path=os.environ.get("FRAMEWORK_REPO_PATH", "").strip(), benchmark_backend=os.environ.get("HYPERLOOM_BENCHMARK_BACKEND", "").strip().lower(), + # Compute-partition lever, read back from the env ``_export_partition_lever`` + # published rather than re-derived from argv: that env is the validated, + # clamped, canonical form, and seeding from the raw flag would let the + # manifest disagree with what the executors were actually handed. + compute_partition_modes=[ + m for m in os.environ.get("HYPERLOOM_COMPUTE_PARTITION_MODES", "").split(",") if m + ], + streams_per_partition=max(1, int(os.environ.get("HYPERLOOM_STREAMS_PER_PARTITION", "") or 2)), + latency_budget_ms=max(0.0, float(os.environ.get("HYPERLOOM_MAX_LATENCY_MS", "") or 0.0)), nodes=max(1, int(getattr(args, "nodes", 1) or 1)), robustness_options=_build_robustness_options(args), warm_replay_enabled=not bool(getattr(args, "no_warm_replay", False)), diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index 6de63ab826..883b260a40 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -287,6 +287,35 @@ def _build_parser() -> argparse.ArgumentParser: "mi308x and mi325x currently run with mi300x runner scripts because " "Magpie does not yet ship MI308X/MI325X-specific SGLang/vLLM scripts.", ) + opt.add_argument( + "--compute-partition-modes", + type=str, + default=None, + metavar="MODES", + help="Comma-separated compute-partition modes to evaluate, e.g. " + "'spx,dpx,cpx'. Off by default. Partitioning a card only ever gives " + "one stream fewer CUs, so it cannot improve single-stream latency " + "and pays only in aggregate throughput at concurrency -- see " + "--streams-per-partition. Requires a privileged amd-smi on the " + "host: the mode is a property of the card, not of the process, so a " + "session confined to an unprivileged container cannot set it. The " + "session restores the mode it found on the way out. Pair with " + "--max-latency-ms; without a budget the search picks the narrowest " + "partition on offer, which is the slowest one per request.", + ) + opt.add_argument( + "--streams-per-partition", + type=int, + default=2, + metavar="N", + help="Concurrent streams to place on each partition when " + "--compute-partition-modes is in use. Defaults to 2, which is where " + "every mode measured on MI355X peaked: one stream leaves each " + "partition idle through the fixed per-pass cost, a second fills it, " + "a third only adds queueing. Raise it only with evidence, and note " + "that every stream on a partition holds its own copy of the " + "weights, so this multiplies the memory a mode has to fit.", + ) opt.add_argument( "--framework", choices=list(framework_registry.names()), @@ -588,6 +617,21 @@ def _build_parser() -> argparse.ArgumentParser: grp.add_argument( "--target-baseline-dir", type=str, default=None, help="Stop when current best matches the baseline in DIR" ) + opt.add_argument( + "--max-latency-ms", + type=float, + default=None, + metavar="MS", + help="Refuse any candidate whose mean end-to-end latency exceeds MS. A " + "constraint, not a target, so it combines with --target-gain / " + "--target-tput rather than replacing one. Off by default, which " + "leaves throughput the only gate. Set it whenever the search can " + "trade latency for throughput -- most of all with " + "--compute-partition-modes, where the highest-throughput " + "configuration is by construction the slowest one per request. The " + "gate fails closed: a candidate that reported no latency is " + "refused, because an unmeasured constraint is not a satisfied one.", + ) opt.add_argument( "--resume-from", type=str, diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py new file mode 100644 index 0000000000..169388500d --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""The compute-partition lever, end to end around one scriptable benchmark. + +``amd-smi`` is faked so these run anywhere, but the benchmark script is a real +bash process: the point of most of these cases is what the script does or does +not see in its environment, and whether it ran at all. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest + +from hyperloom.orchestrator.actions.executors import _partition_lever as pl +from hyperloom.orchestrator.actions.executors import bypass_scriptable as bs +from hyperloom.orchestrator.actions.executors._latency_budget import ( + REASON_OVER_BUDGET, + REASON_UNMEASURED, + latency_keep_block, + resolve_latency_budget_ms, +) + + +class _FakeSmi: + """An ``amd-smi`` whose partition state the test drives.""" + + def __init__(self, mode: str = "SPX", *, settable: bool = True): + self.modes = {i: mode for i in range(8)} + self.settable = settable + self.history: list[str] = [] + + def __call__(self, cmd, **kwargs): + if cmd[:2] == ["amd-smi", "set"]: + gpu_id = int(cmd[cmd.index("-g") + 1]) + mode = cmd[cmd.index("--compute-partition") + 1] + if not self.settable: + return subprocess.CompletedProcess(cmd, 1, "", "permission denied") + self.modes[gpu_id] = mode + self.history.append(mode) + return subprocess.CompletedProcess(cmd, 0, "", "") + rows = [ + {"gpu_id": gid, "memory": "NPS1", "accelerator_type": m} + for gid, m in sorted(self.modes.items()) + ] + return subprocess.CompletedProcess(cmd, 0, json.dumps({"current_partition": rows}), "") + + +def _scripts_dir(tmp_path: Path) -> Path: + """A benchmark entrypoint that dumps its partition env and records that it ran.""" + scripts = tmp_path / "scripts" + scripts.mkdir(exist_ok=True) + (scripts / "custom_mi355x.sh").write_text( + "#!/bin/bash\n" + 'echo "ran" > "$RESULT_DIR/ran.marker"\n' + 'env | grep -E "^HYPERLOOM_PARTITION" | sort > "$RESULT_DIR/partition.env" || true\n', + encoding="utf-8", + ) + return scripts + + +def _run(tmp_path: Path, monkeypatch, envs: dict | None = None): + monkeypatch.setenv("HYPERLOOM_BYPASS_SCRIPTS_DIR", str(_scripts_dir(tmp_path))) + monkeypatch.delenv("MAGPIE_PATH", raising=False) + workspace = tmp_path / "ws" + workspace.mkdir(exist_ok=True) + rc, error = bs.run_scriptable( + framework="custom", + runner_type="mi355x", + inferencex_root=str(tmp_path / "InferenceX"), + bench={"model": "/models/hwmirror", "envs": dict(envs or {})}, + workspace=workspace, + timeout_s=60.0, + ) + partition_env = workspace / "partition.env" + seen = dict( + line.split("=", 1) + for line in (partition_env.read_text(encoding="utf-8").splitlines() if partition_env.is_file() else []) + if "=" in line + ) + return rc, error, (workspace / "ran.marker").is_file(), seen + + +@pytest.fixture(autouse=True) +def _clean_lever_env(monkeypatch): + """No ambient lever: these cases each state their own.""" + for name in ( + pl.PARTITION_MODE_ENV, + pl.PARTITION_MODES_ENV, + pl.STREAMS_PER_PARTITION_ENV, + pl.PARTITION_GPU_ENV, + ): + monkeypatch.delenv(name, raising=False) + + +def test_lever_off_leaves_the_run_untouched(tmp_path, monkeypatch): + fake = _FakeSmi() + monkeypatch.setattr(subprocess, "run", fake) + rc, error, ran, seen = _run(tmp_path, monkeypatch) + assert (rc, error, ran) == (0, None, True) + # Nothing published, and above all nothing set: a session that did not ask + # for partitioning must not touch a shared card. + assert seen == {} + assert fake.history == [] + + +def test_a_requested_mode_is_established_published_and_restored(tmp_path, monkeypatch): + fake = _FakeSmi() + monkeypatch.setattr(subprocess, "run", fake) + rc, error, ran, seen = _run(tmp_path, monkeypatch, {pl.PARTITION_MODE_ENV: "cpx"}) + assert (rc, error, ran) == (0, None, True) + assert seen[pl.RUNTIME_MODE_ENV] == "CPX" + assert seen[pl.RUNTIME_COUNT_ENV] == "8" + assert seen[pl.RUNTIME_CU_ENV] == "32" + # Two per partition by default, so sixteen streams in total. + assert seen[pl.RUNTIME_STREAMS_ENV] == "2" + assert seen[pl.RUNTIME_TOTAL_STREAMS_ENV] == "16" + # Held for the run, then handed back in the mode it was found in. + assert fake.history == ["CPX", "SPX"] + assert fake.modes[0] == "SPX" + + +def test_streams_per_partition_flows_through_to_the_benchmark(tmp_path, monkeypatch): + monkeypatch.setattr(subprocess, "run", _FakeSmi()) + monkeypatch.setenv(pl.STREAMS_PER_PARTITION_ENV, "3") + _, _, _, seen = _run(tmp_path, monkeypatch, {pl.PARTITION_MODE_ENV: "DPX"}) + assert seen[pl.RUNTIME_STREAMS_ENV] == "3" + assert seen[pl.RUNTIME_TOTAL_STREAMS_ENV] == "6" + + +def test_an_unsettable_mode_aborts_before_the_benchmark_runs(tmp_path, monkeypatch): + # No privilege to repartition. The alternative to failing here is measuring + # the topology that happens to be present and labelling it CPX. + monkeypatch.setattr(subprocess, "run", _FakeSmi(settable=False)) + rc, error, ran, _ = _run(tmp_path, monkeypatch, {pl.PARTITION_MODE_ENV: "CPX"}) + assert rc == 2 + assert "compute-partition lever" in error + assert not ran + + +def test_an_unknown_mode_aborts_before_the_benchmark_runs(tmp_path, monkeypatch): + monkeypatch.setattr(subprocess, "run", _FakeSmi()) + rc, error, ran, _ = _run(tmp_path, monkeypatch, {pl.PARTITION_MODE_ENV: "NPS2"}) + assert rc == 2 + assert "unknown compute-partition mode" in error + assert not ran + + +def test_a_variant_mode_overrides_the_ambient_one(monkeypatch): + monkeypatch.setenv(pl.PARTITION_MODE_ENV, "SPX") + assert pl.requested_mode({pl.PARTITION_MODE_ENV: "cpx"}) == "CPX" + assert pl.requested_mode({}) == "SPX" + + +def test_partition_plan_is_empty_when_no_mode_is_requested(): + assert pl.plan_partition_run({}, gpu_type="mi355x") == ("", {}) + + +def test_latency_budget_is_off_by_default(): + # The gate must not change behaviour for the sessions that never set it. + assert latency_keep_block(1211.0, budget_ms=0.0) == (False, "") + assert resolve_latency_budget_ms({}, None) == 0.0 + + +def test_latency_budget_blocks_the_partition_trade_it_exists_for(): + # CPX-16 measured 13.07 fwd/s at 1211 ms against SPX-2's 10.90 at 183 ms: + # a +20% throughput gain the old gate would have kept. + blocked, reason = latency_keep_block(1211.0, budget_ms=300.0) + assert blocked and REASON_OVER_BUDGET in reason and "4.04x" in reason + assert latency_keep_block(183.0, budget_ms=300.0) == (False, "") + + +def test_latency_budget_fails_closed_on_an_unmeasured_candidate(): + blocked, reason = latency_keep_block(None, budget_ms=300.0) + assert blocked and REASON_UNMEASURED in reason + # Nonsense readings are treated as no reading, not as a pass. + assert latency_keep_block(0.0, budget_ms=300.0)[0] + assert latency_keep_block(float("nan"), budget_ms=300.0)[0] + + +def test_latency_budget_precedence_is_most_specific_first(monkeypatch): + monkeypatch.setenv("HYPERLOOM_MAX_LATENCY_MS", "250") + + class _State: + latency_budget_ms = 400.0 + + assert resolve_latency_budget_ms({"latency_budget_ms": 500}, _State()) == 500.0 + assert resolve_latency_budget_ms({}, _State()) == 400.0 + assert resolve_latency_budget_ms({}, None) == 250.0 diff --git a/src/hyperloom/orchestrator/actions/executors/_latency_budget.py b/src/hyperloom/orchestrator/actions/executors/_latency_budget.py new file mode 100644 index 0000000000..2d3affee29 --- /dev/null +++ b/src/hyperloom/orchestrator/actions/executors/_latency_budget.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Latency budget — the constraint that makes a throughput-only KEEP safe. + +The optimizer maximizes ``output_throughput`` and nothing else. Latency is +measured, reported and fed to the prompts, but no latency number has ever +blocked a KEEP. For every lever that was in the search space when that was +decided the omission is survivable, because a serving-config flag that doubles +throughput rarely destroys per-request latency. + +Compute partitioning is not like that. Splitting a card raises aggregate +throughput precisely *by* making each stream slower, so a throughput-only gate +does not merely tolerate a latency regression -- it selects for the largest one +available. On one MI355X the ladder tops out at CPX with two streams per +partition: ~20% more throughput than the best SPX configuration, with +per-request latency going from 183 ms to 1211 ms. A gate reading only throughput +calls that a win and reports +20%. + +So a session that opts into partition modes should also state what latency it +can live with. The budget is off by default and absolute rather than relative: an +SLA is a fixed number the workload owner already knows, and a percentage cap +against a baseline would silently ratchet as the baseline improves. + +The gate fails closed. An operator who names a budget has asserted the +constraint matters, and a candidate whose latency was never measured cannot be +shown to satisfy it -- so it is refused rather than admitted on the assumption +that unmeasured means acceptable. The reason string says which of the two +happened, because the remedy differs: one needs a different candidate, the other +needs the benchmark to report latency at all. +""" + +from __future__ import annotations + +import logging +import math +import os +from typing import Any + +log = logging.getLogger(__name__) + +#: Env override for the session latency budget, in milliseconds. Non-positive or +#: unparseable disables the gate, matching the CLI default. +LATENCY_BUDGET_ENV = "HYPERLOOM_MAX_LATENCY_MS" + +#: Journal/ledger reasons. Distinct so a report can tell "too slow" apart from +#: "never timed", which are different bugs with different owners. +REASON_OVER_BUDGET = "latency_budget_exceeded" +REASON_UNMEASURED = "latency_unmeasured_under_budget" + + +def _finite_positive(value: Any) -> float | None: + """Return ``value`` as a finite positive float, else ``None``.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + val = float(value) + if not math.isfinite(val) or val <= 0: + return None + return val + + +def resolve_latency_budget_ms( + params: dict[str, Any] | None = None, + shared_state: Any = None, +) -> float: + """Resolve the session's latency budget in milliseconds. + + Precedence is most-specific-first: an explicit task parameter, then the + session state the CLI seeded, then the environment. Anything non-positive or + unparseable means no budget, which is the default and leaves KEEP behaviour + exactly as it was. + + Args: + params: Task params, which may carry ``latency_budget_ms``. + shared_state: Live SharedState, which may carry ``latency_budget_ms``. + + Returns: + The budget in ms, or ``0.0`` when the gate is off. + """ + for candidate in ( + (params or {}).get("latency_budget_ms"), + getattr(shared_state, "latency_budget_ms", None), + os.environ.get(LATENCY_BUDGET_ENV), + ): + if candidate in (None, ""): + continue + try: + val = float(candidate) + except (TypeError, ValueError): + log.warning("ignoring unparseable latency budget %r", candidate) + continue + if val > 0: + return val + return 0.0 + + +def latency_keep_block( + observed_ms: Any, + *, + budget_ms: float, +) -> tuple[bool, str]: + """Decide whether the latency budget blocks a KEEP. + + Args: + observed_ms: The candidate's mean end-to-end latency, or ``None`` when + the benchmark reported none. + budget_ms: The session budget; ``<= 0`` disables the gate. + + Returns: + ``(blocked, reason)``. ``reason`` is empty when nothing blocks. + """ + budget = _finite_positive(budget_ms) + if budget is None: + return False, "" + observed = _finite_positive(observed_ms) + if observed is None: + return True, ( + f"{REASON_UNMEASURED}: a {budget:.0f} ms budget is set but this " + f"candidate reported no end-to-end latency, so the constraint " + f"cannot be shown to hold. Have the benchmark report mean_e2el_ms." + ) + if observed > budget: + return True, ( + f"{REASON_OVER_BUDGET}: {observed:.0f} ms exceeds the " + f"{budget:.0f} ms budget ({observed / budget:.2f}x)" + ) + return False, "" + + +def describe_latency_budget(budget_ms: float) -> str: + """Render the budget for a log line or report header.""" + budget = _finite_positive(budget_ms) + return f"latency budget {budget:.0f} ms" if budget else "no latency budget" + + +__all__ = [ + "LATENCY_BUDGET_ENV", + "REASON_OVER_BUDGET", + "REASON_UNMEASURED", + "describe_latency_budget", + "latency_keep_block", + "resolve_latency_budget_ms", +] diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py new file mode 100644 index 0000000000..382131272e --- /dev/null +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Applying a compute-partition mode around one benchmark run. + +The mode travels as an ordinary variant env, ``HYPERLOOM_PARTITION_MODE``. That +is deliberate: a candidate identified by its envs is already fingerprinted, +deduplicated in the explore ledger, graded by the KEEP gate and recorded in the +journal, so expressing the mode that way inherits all of it instead of growing a +parallel search axis that each of those would have to learn about. + +What cannot be inherited is the application. Every other env is delivered *to* +the benchmark process; this one has to change the card before that process +starts, and change it back afterwards. Hence this module, called from the +scriptable runner rather than from the env materializer. + +Two things it deliberately refuses to do: + +* **Guess.** If the mode cannot be set -- no privilege, no ``amd-smi``, a card + with resident processes -- the run fails. Falling back to the current topology + would produce a perfectly good measurement labelled with the wrong mode, which + is worse than no measurement, because it is indistinguishable from a real one. +* **Enumerate devices for the benchmark.** The runner publishes the shape it + established (how many partitions, how many CUs each) and the benchmark selects + its own devices by CU count. Index-based selection is the classic way to + measure a whole card and report it as a partition, and the process that owns + the GPU context is the one positioned to check. +""" + +from __future__ import annotations + +import logging +import os +from contextlib import contextmanager, nullcontext +from typing import Any, Iterator + +from hyperloom.common.gpu_partition import ( + PartitionError, + PartitionLayout, + layout_for, + partitioned, +) + +log = logging.getLogger(__name__) + +#: Per-variant selector. Present => this run is measured under that mode. +PARTITION_MODE_ENV = "HYPERLOOM_PARTITION_MODE" + +#: Session-level lever published by the CLI. +PARTITION_MODES_ENV = "HYPERLOOM_COMPUTE_PARTITION_MODES" +STREAMS_PER_PARTITION_ENV = "HYPERLOOM_STREAMS_PER_PARTITION" + +#: GPU whose partition mode is managed. Single-card today: the measurement is +#: aggregate throughput on one physical GPU, and repartitioning a whole node to +#: measure that is a much larger blast radius for no extra signal. +PARTITION_GPU_ENV = "HYPERLOOM_PARTITION_GPU" + +#: Published to the benchmark so its entrypoint can fan out and, crucially, +#: verify what it is running on. +RUNTIME_MODE_ENV = "HYPERLOOM_PARTITION_MODE" +RUNTIME_COUNT_ENV = "HYPERLOOM_PARTITION_COUNT" +RUNTIME_CU_ENV = "HYPERLOOM_PARTITION_CU" +RUNTIME_STREAMS_ENV = "HYPERLOOM_PARTITION_STREAMS_PER_PARTITION" +RUNTIME_TOTAL_STREAMS_ENV = "HYPERLOOM_PARTITION_TOTAL_STREAMS" + + +def partition_gpu_id() -> int: + """Return the GPU whose mode this session manages (default 0).""" + try: + return max(0, int(os.environ.get(PARTITION_GPU_ENV, "0").strip() or 0)) + except ValueError: + return 0 + + +def streams_per_partition() -> int: + """Return the configured streams per partition (default 2). + + Two is where every mode measured on MI355X peaked. It is a default rather + than a constant because it is a property of the workload's fixed per-pass + cost, not of the hardware. + """ + try: + return max(1, int(os.environ.get(STREAMS_PER_PARTITION_ENV, "2").strip() or 2)) + except ValueError: + return 2 + + +def requested_mode(envs: dict[str, Any] | None) -> str: + """Return the partition mode this run asks for, or ``""`` when none. + + The variant's own env wins over the process environment, so a grid can hold + one variant at SPX while the session lever lists several modes. + """ + for source in (envs or {}, os.environ): + value = str((source or {}).get(PARTITION_MODE_ENV) or "").strip().upper() + if value: + return value + return "" + + +def runtime_env(layout: PartitionLayout, streams: int) -> dict[str, str]: + """Build the env describing an established partition shape. + + Args: + layout: The layout the hardware confirmed. + streams: Streams to place on each partition. + + Returns: + Env mapping for the benchmark subprocess. + """ + return { + RUNTIME_MODE_ENV: layout.mode, + RUNTIME_COUNT_ENV: str(layout.partitions), + RUNTIME_CU_ENV: str(layout.cu_per_partition), + RUNTIME_STREAMS_ENV: str(streams), + RUNTIME_TOTAL_STREAMS_ENV: str(layout.partitions * streams), + } + + +def plan_partition_run( + envs: dict[str, Any] | None, + *, + gpu_type: str | None, +) -> tuple[str, dict[str, str]]: + """Resolve the mode for this run and the env describing it. + + Args: + envs: The materialized benchmark envs, which may select a mode. + gpu_type: Board name, needed to size the partitions. + + Returns: + ``(mode, env)``. ``("", {})`` when the lever is not engaged, which is + the default and leaves the run exactly as it was. + + Raises: + PartitionError: If a mode was requested but cannot be described -- an + unknown mode or an unrecognised board. Raised rather than ignored: + the request was explicit, so silently not honouring it would + mislabel the measurement. + """ + mode = requested_mode(envs) + if not mode: + return "", {} + layout = layout_for(gpu_type, mode) + streams = streams_per_partition() + return mode, runtime_env(layout, streams) + + +@contextmanager +def hold_partition_mode(mode: str, *, gpu_type: str | None) -> Iterator[PartitionLayout | None]: + """Hold the managed GPU in ``mode`` for the duration of the block. + + Yields ``None`` when no mode was requested, so callers can wrap + unconditionally. + + Raises: + PartitionError: If the mode cannot be established or verified. + """ + if not mode: + yield None + return + layout = layout_for(gpu_type, mode) + gpu_id = partition_gpu_id() + log.info("benchmark: holding GPU %d at %s", gpu_id, layout.describe()) + with partitioned(gpu_id, mode): + yield layout + + +def maybe_hold_partition_mode(mode: str, *, gpu_type: str | None): + """Return a context manager for ``mode``, or a no-op when the lever is off.""" + if not mode: + return nullcontext(None) + return hold_partition_mode(mode, gpu_type=gpu_type) + + +__all__ = [ + "PARTITION_GPU_ENV", + "PARTITION_MODES_ENV", + "PARTITION_MODE_ENV", + "PartitionError", + "RUNTIME_COUNT_ENV", + "RUNTIME_CU_ENV", + "RUNTIME_MODE_ENV", + "RUNTIME_STREAMS_ENV", + "RUNTIME_TOTAL_STREAMS_ENV", + "STREAMS_PER_PARTITION_ENV", + "hold_partition_mode", + "maybe_hold_partition_mode", + "partition_gpu_id", + "plan_partition_run", + "requested_mode", + "runtime_env", + "streams_per_partition", +] diff --git a/src/hyperloom/orchestrator/actions/executors/bypass_scriptable.py b/src/hyperloom/orchestrator/actions/executors/bypass_scriptable.py index 987bb72c07..264892c53b 100644 --- a/src/hyperloom/orchestrator/actions/executors/bypass_scriptable.py +++ b/src/hyperloom/orchestrator/actions/executors/bypass_scriptable.py @@ -29,6 +29,12 @@ from hyperloom.common.env_safety import build_benchmark_env from hyperloom.inference_optimizer.session.paths import asset_root +from ._partition_lever import ( + PartitionError, + maybe_hold_partition_mode, + plan_partition_run, +) + def _scriptable_script_name(framework: str, runner_type: str) -> str: """Return the scriptable entrypoint name (e.g. xdit_mi300x.sh).""" @@ -116,6 +122,7 @@ def build_scriptable_env( *, profile: bool = False, profile_dir: str | None = None, + partition_env: dict[str, str] | None = None, ) -> dict[str, str]: """Build the env for a scriptable benchmark script. @@ -125,6 +132,10 @@ def build_scriptable_env( workspace: Per-run workspace directory. profile: Whether the torch profiler is enabled for this run. profile_dir: Directory the profiler traces should be written to. + partition_env: Description of the compute-partition shape established + for this run, when the lever is engaged. Run-scoped rather than + overridable: it reports what the hardware was actually set to, so a + YAML env claiming otherwise must not win. Returns: The environment mapping for the scriptable subprocess. @@ -145,6 +156,8 @@ def build_scriptable_env( if profile_dir: run_scoped["VLLM_TORCH_PROFILER_DIR"] = profile_dir run_scoped["SGLANG_TORCH_PROFILER_DIR"] = profile_dir + if partition_env: + run_scoped.update(partition_env) return build_benchmark_env(defaults, bench.get("envs"), run_scoped) @@ -187,31 +200,50 @@ def run_scriptable( # carry the diagnostic instead of a blank abort_reason.json. _write_logs(workspace, "", f"{error}\ntried:\n{tried}\n") return 2, error - env = build_scriptable_env(bench, runner_type, workspace, profile=profile, profile_dir=profile_dir) + # A requested partition mode is established before the benchmark starts and + # restored after it ends, so the measurement and its label agree. A mode + # that cannot be set aborts the run rather than measuring the topology that + # happens to be there, which would be indistinguishable from a real result. + try: + mode, partition_env = plan_partition_run(bench.get("envs"), gpu_type=runner_type) + except PartitionError as exc: + return 2, f"compute-partition lever: {exc}" + env = build_scriptable_env( + bench, + runner_type, + workspace, + profile=profile, + profile_dir=profile_dir, + partition_env=partition_env, + ) cmd = ["bash", str(script)] # Streamed straight to disk instead of captured in memory: a runner killed # from outside (lease reap / OOM) must still leave a forensic trail. - with ExitStack() as stack: - stdout_sink = stack.enter_context(_open_log_sink(workspace, "scriptable_stdout.log")) - stderr_sink = stack.enter_context(_open_log_sink(workspace, "scriptable_stderr.log")) - proc = subprocess.Popen( # noqa: S603 — cmd is this module's own bash entrypoint - cmd, - env=env, - stdout=stdout_sink, - stderr=stderr_sink, - ) - try: - proc.wait(timeout=timeout_s) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - _write_logs( - workspace, - "", - f"scriptable benchmark timed out after {timeout_s}s", - append=True, + try: + with maybe_hold_partition_mode(mode, gpu_type=runner_type), ExitStack() as stack: + stdout_sink = stack.enter_context(_open_log_sink(workspace, "scriptable_stdout.log")) + stderr_sink = stack.enter_context(_open_log_sink(workspace, "scriptable_stderr.log")) + proc = subprocess.Popen( # noqa: S603 — cmd is this module's own bash entrypoint + cmd, + env=env, + stdout=stdout_sink, + stderr=stderr_sink, ) - return 124, None + try: + proc.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + _write_logs( + workspace, + "", + f"scriptable benchmark timed out after {timeout_s}s", + append=True, + ) + return 124, None + except PartitionError as exc: + _write_logs(workspace, "", f"compute-partition lever: {exc}\n", append=True) + return 2, f"compute-partition lever: {exc}" return proc.returncode, None diff --git a/src/hyperloom/orchestrator/actions/executors/explore.py b/src/hyperloom/orchestrator/actions/executors/explore.py index 5643dd31cb..fe7abf4ef2 100644 --- a/src/hyperloom/orchestrator/actions/executors/explore.py +++ b/src/hyperloom/orchestrator/actions/executors/explore.py @@ -80,6 +80,11 @@ session_grid_bounds, ) from ._grid_server_args import compose_server_args, server_args_env_name +from ._latency_budget import ( + describe_latency_budget, + latency_keep_block, + resolve_latency_budget_ms, +) from ._ray_serving import maybe_serving_lease # DEFAULT_STACK_STABLE_PCT: post-KEEP confirmation floor; override via @@ -799,6 +804,9 @@ async def __call__(self, ctx) -> dict[str, Any]: self.stack_stable_threshold_pct, ) ) + latency_budget_ms = resolve_latency_budget_ms(params, ss) + if latency_budget_ms > 0: + log.info("explore: enforcing %s on KEEP", describe_latency_budget(latency_budget_ms)) enable_stack_rebench = bool( params.get( "enable_stack_rebench", @@ -1514,6 +1522,13 @@ def _stopped_by_the_run(result: Any, *, variant: GridVariant, idx: int, round_la # clear keep_threshold (and the accuracy gate) earn a warm # stack-rebench round. gain = gain_pct(r.output_throughput, running_base_tput) + # A throughput gain bought with latency is still a gain to + # the gate above, so a session that named a latency budget + # has it enforced here rather than only reported. + latency_blocked, latency_reason = latency_keep_block( + r.e2el_mean_ms, + budget_ms=latency_budget_ms, + ) outcome = "FAILED" reason: str = "" if r.status != "succeeded" or gain is None: @@ -1521,6 +1536,15 @@ def _stopped_by_the_run(result: Any, *, variant: GridVariant, idx: int, round_la elif gain < keep_threshold_pct: outcome = "REVERT" reason = "gain_below_threshold" + elif latency_blocked: + outcome = "REVERT" + reason = latency_reason + log.warning( + "explore: variant %s REVERT (+%.2f%% throughput, %s)", + gv.name, + gain, + latency_reason, + ) else: # Accuracy gate. For serving it runs only for high-risk # variants. For scriptable frameworks the image-quality diff --git a/src/hyperloom/orchestrator/state/shared_state.py b/src/hyperloom/orchestrator/state/shared_state.py index c791ac01bd..04c6f01174 100644 --- a/src/hyperloom/orchestrator/state/shared_state.py +++ b/src/hyperloom/orchestrator/state/shared_state.py @@ -666,6 +666,17 @@ class SharedState(_RenderMixin, _ExploreStateMixin): framework_repo_path: str = "" # ``HYPERLOOM_BENCHMARK_BACKEND`` at seed time (``bypass`` for custom). benchmark_backend: str = "" + # ``--compute-partition-modes``: canonical SPX/DPX/QPX/CPX modes the session + # may set on the card. Empty (the default) leaves the hardware alone. + compute_partition_modes: list[str] = field(default_factory=list) + # ``--streams-per-partition``: concurrent streams placed on each partition. + # Part of the measurement contract, not a tuning detail -- a mode measured + # at one stream per partition is a different experiment from the same mode + # at two, and only the latter is where partitioning pays. + streams_per_partition: int = 2 + # ``--max-latency-ms``: constraint, not target. 0 disables, which restores + # the throughput-only KEEP the optimizer had before. + latency_budget_ms: float = 0.0 # ``--nodes``, feeding the robustness defaults and the IR-8 check. NOT the # cluster hand-off, which is resolved from argv before this state loads. nodes: int = 1 From f824723540e71c2974a911e80352782664cb5a17 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Mon, 24 Aug 2026 00:39:23 +0000 Subject: [PATCH 02/10] refactor(partition): resume the lever on main's restore/apply/persist path The rebase brought in ``ea45bf687``, which gave operator-supplied session knobs a three-step shape: restore from archive when the resume omits them, apply (validating) so a re-passed flag wins, then persist the live env back onto state for the resume after. The lever had hand-rolled its own half of that in one expression and got two things wrong that the convention gets right. * The env tier was missing. Resume resolved CLI flag > archive, so a ``HYPERLOOM_MAX_LATENCY_MS`` exported for this resume was overwritten by the older archived value. ``_restore_partition_lever_from_state`` now applies flag > env > archive, writing onto ``args`` so ``_export_partition_lever`` stays the only writer of that env and the flag still gets validated rather than trusted. * Nothing was persisted. A resume that re-passed ``--max-latency-ms`` published the new budget to the executors while the manifest kept the old one, so the session was measured under one contract and recorded under another -- and the next resume restored the stale number. ``_persist_partition_lever`` mirrors it back. * ``--streams-per-partition`` defaulted to 2, which made "not passed" indistinguishable from "passed 2" and would overwrite a persisted 4 on every resume. Now defaults to None, with the 2 applied where the flag is resolved. A parser-level test locks this, because re-adding the 2 reads as harmless -- 2 *is* the documented default -- and silently changes the experiment. Also drops a duplicated pair of env-name constants: ``cli/__init__`` had its own copy of ``PARTITION_MODES_ENV`` / ``STREAMS_PER_PARTITION_ENV`` alongside ``_partition_lever``'s. The names now come from the module that owns them, and ``read_session_lever`` is the one parser of those three variables, shared by the manifest seed and the resume persist. Conflicts in the rebase were all additive and taken whole from both sides, except ``run_scriptable``, where main's candidate-list diagnostic for a missing entrypoint replaced the one-line error the lever's branch still had above its partition planning. Co-Authored-By: Claude Opus 5 Co-authored-by: Cursor (cherry picked from commit 259c6d32e86f3fb7c0b0b819439dcdf208dd3289) --- .../inference_optimizer/cli/__init__.py | 74 ++++++++--- .../inference_optimizer/cli/bootstrap.py | 13 +- .../inference_optimizer/cli/parser.py | 5 +- .../tests/test_partition_lever.py | 125 ++++++++++++++++++ .../actions/executors/_partition_lever.py | 23 ++++ 5 files changed, 219 insertions(+), 21 deletions(-) diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 54f222cc93..370444db7b 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -66,6 +66,11 @@ ) from hyperloom.orchestrator.actions.executors._aiter_jit import clean_stale_aiter_locks from hyperloom.orchestrator.actions.executors._latency_budget import LATENCY_BUDGET_ENV +from hyperloom.orchestrator.actions.executors._partition_lever import ( + PARTITION_MODES_ENV, + STREAMS_PER_PARTITION_ENV, + read_session_lever, +) from hyperloom.orchestrator.actions.executors._workload_envs import ( agentx_enabled as _agentx_enabled, ) @@ -1632,13 +1637,6 @@ def _export_operator_launch_shape( os.environ.pop("INFERENCE_OPTIMIZER_EXTRA_ENV", None) -#: Internal handoff for the compute-partition lever. Operators pass -#: ``--compute-partition-modes`` / ``--streams-per-partition`` / -#: ``--max-latency-ms``; these carry the resolved values to the executors. -PARTITION_MODES_ENV = "HYPERLOOM_COMPUTE_PARTITION_MODES" -STREAMS_PER_PARTITION_ENV = "HYPERLOOM_STREAMS_PER_PARTITION" - - def _export_partition_lever( *, modes_raw: str | None, @@ -1705,6 +1703,49 @@ def _export_partition_lever( return modes +def _restore_partition_lever_from_state(args: Any, state: SharedState) -> None: + """Fill the partition lever from env or archive when this resume omitted it. + + Priority is CLI flag > already-exported env > archived ``SharedState``, the + same chain :func:`_restore_operator_supplied_paths_from_state` applies to + the custom-workload paths. Resolved values are written back onto ``args`` so + :func:`_export_partition_lever` remains the only writer of the env it owns, + and so the flag still gets validated rather than trusted. + + The archive tier is what lets a resume reproduce the session's measurement + contract: partition mode and streams-per-partition are part of *how* a + number was obtained, so a resume that silently dropped them would compare + candidates measured under partitioning against a baseline that was not. + + Args: + args: Parsed CLI namespace, updated in place. + state: Resumed session state. + """ + if not str(getattr(args, "compute_partition_modes", None) or "").strip(): + env_modes = os.environ.get(PARTITION_MODES_ENV, "").strip() + args.compute_partition_modes = env_modes or ",".join(state.compute_partition_modes or []) + if not (getattr(args, "streams_per_partition", None) or 0): + env_streams = os.environ.get(STREAMS_PER_PARTITION_ENV, "").strip() + args.streams_per_partition = int(env_streams or state.streams_per_partition or 2) + if not (getattr(args, "max_latency_ms", None) or 0.0): + env_budget = os.environ.get(LATENCY_BUDGET_ENV, "").strip() + args.max_latency_ms = float(env_budget or state.latency_budget_ms or 0.0) + + +def _persist_partition_lever(state: SharedState) -> None: + """Mirror the live lever env back onto ``state`` for the next resume. + + Without this a resume that re-passed ``--max-latency-ms`` would publish the + new budget to the executors while the manifest kept the old one, so the + session would be measured under one contract and recorded under another -- + and the resume after it would restore the stale number. + """ + modes, streams, budget = read_session_lever() + state.compute_partition_modes = list(modes) + state.streams_per_partition = streams + state.latency_budget_ms = budget + + # Terminal stop_reasons that represent a clean, successful optimizer run (exit 0). # Anything else (baseline / preflight failures, crashes, enablement stalls) exits # non-zero so CI surfaces genuine problems. @@ -2013,16 +2054,19 @@ async def _run_optimize(args: argparse.Namespace) -> int: server_args=_resume_server_args, extra_env=_resume_extra_env, ) - # The lever is part of the measurement contract: a resume that dropped - # it would compare candidates measured under partitioning against a - # baseline that no longer is. + # The lever is part of the measurement contract, so it resumes on the + # same restore / apply / persist path as the custom-workload paths below. + _restore_partition_lever_from_state(args, state) _export_partition_lever( - modes_raw=getattr(args, "compute_partition_modes", None) or state.compute_partition_modes, - streams_per_partition=int( - getattr(args, "streams_per_partition", 0) or state.streams_per_partition or 2 - ), - max_latency_ms=getattr(args, "max_latency_ms", None) or state.latency_budget_ms, + modes_raw=getattr(args, "compute_partition_modes", None), + streams_per_partition=int(getattr(args, "streams_per_partition", None) or 2), + max_latency_ms=getattr(args, "max_latency_ms", None), ) + _persist_partition_lever(state) + if state.compute_partition_modes: + print(f" re-exported partition modes: {','.join(state.compute_partition_modes)}") + if state.latency_budget_ms: + print(f" re-exported max_latency_ms : {state.latency_budget_ms:g}") state.operator_server_args = _resume_server_args state.operator_extra_env = _resume_extra_env if _resume_server_args: diff --git a/src/hyperloom/inference_optimizer/cli/bootstrap.py b/src/hyperloom/inference_optimizer/cli/bootstrap.py index 5b878629f6..6fd4887dd9 100644 --- a/src/hyperloom/inference_optimizer/cli/bootstrap.py +++ b/src/hyperloom/inference_optimizer/cli/bootstrap.py @@ -22,6 +22,7 @@ from hyperloom.common.coerce import to_unix from hyperloom.common.env import forge_explicitly_enabled from hyperloom.common.timeutil import now_iso +from hyperloom.orchestrator.actions.executors._partition_lever import read_session_lever from hyperloom.orchestrator.actions.executors._workload_envs import ( agentx_enabled as _agentx_enabled, ) @@ -308,6 +309,10 @@ def _resolve_framework_version(args_in: Any) -> str: # Canonical model identity (prefers the quantize prelude's pinned source name). _model_identity = resolve_model_display_name(args) + + # Compute-partition lever as the launcher published it, read through the one + # parser that owns those variables. + _partition_lever_seed = read_session_lever() state = SharedState( session_id=session_id, claw_session_id=(os.environ.get("CLAW_SESSION_ID") or "").strip(), @@ -359,11 +364,9 @@ def _resolve_framework_version(args_in: Any) -> str: # published rather than re-derived from argv: that env is the validated, # clamped, canonical form, and seeding from the raw flag would let the # manifest disagree with what the executors were actually handed. - compute_partition_modes=[ - m for m in os.environ.get("HYPERLOOM_COMPUTE_PARTITION_MODES", "").split(",") if m - ], - streams_per_partition=max(1, int(os.environ.get("HYPERLOOM_STREAMS_PER_PARTITION", "") or 2)), - latency_budget_ms=max(0.0, float(os.environ.get("HYPERLOOM_MAX_LATENCY_MS", "") or 0.0)), + compute_partition_modes=list(_partition_lever_seed[0]), + streams_per_partition=_partition_lever_seed[1], + latency_budget_ms=_partition_lever_seed[2], nodes=max(1, int(getattr(args, "nodes", 1) or 1)), robustness_options=_build_robustness_options(args), warm_replay_enabled=not bool(getattr(args, "no_warm_replay", False)), diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index 883b260a40..77b30ee4f8 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -306,7 +306,10 @@ def _build_parser() -> argparse.ArgumentParser: opt.add_argument( "--streams-per-partition", type=int, - default=2, + # None, not 2, so a resume can tell "not passed" from "passed 2" and + # let the persisted value stand. The 2 is applied where the flag is + # resolved. + default=None, metavar="N", help="Concurrent streams to place on each partition when " "--compute-partition-modes is in use. Defaults to 2, which is where " diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py index 169388500d..007068ebdf 100644 --- a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -10,12 +10,15 @@ from __future__ import annotations +import argparse import json import subprocess from pathlib import Path +from types import SimpleNamespace import pytest +from hyperloom.inference_optimizer import cli from hyperloom.orchestrator.actions.executors import _partition_lever as pl from hyperloom.orchestrator.actions.executors import bypass_scriptable as bs from hyperloom.orchestrator.actions.executors._latency_budget import ( @@ -191,3 +194,125 @@ class _State: assert resolve_latency_budget_ms({"latency_budget_ms": 500}, _State()) == 500.0 assert resolve_latency_budget_ms({}, _State()) == 400.0 assert resolve_latency_budget_ms({}, None) == 250.0 + + +class TestResumeContract: + """The lever must survive a resume, because it is part of how a number was got. + + A resume that quietly dropped the mode or the streams count would compare + candidates measured under partitioning against a baseline that was not, and + the comparison would look ordinary. These lock the precedence chain (CLI + flag > exported env > archived state) and the write-back that keeps the + manifest agreeing with what the executors were handed. + """ + + @staticmethod + def _args(**over): + ns = argparse.Namespace( + compute_partition_modes=None, streams_per_partition=None, max_latency_ms=None + ) + for key, value in over.items(): + setattr(ns, key, value) + return ns + + @staticmethod + def _state(modes=(), streams=2, budget=0.0): + state = SimpleNamespace( + compute_partition_modes=list(modes), + streams_per_partition=streams, + latency_budget_ms=budget, + ) + return state + + def test_archive_supplies_the_lever_when_the_resume_omits_it(self, monkeypatch): + monkeypatch.delenv(pl.PARTITION_MODES_ENV, raising=False) + monkeypatch.delenv(pl.STREAMS_PER_PARTITION_ENV, raising=False) + monkeypatch.delenv("HYPERLOOM_MAX_LATENCY_MS", raising=False) + args = self._args() + cli._restore_partition_lever_from_state(args, self._state(("DPX", "CPX"), 3, 400.0)) + assert args.compute_partition_modes == "DPX,CPX" + assert args.streams_per_partition == 3 + assert args.max_latency_ms == 400.0 + + def test_an_exported_env_outranks_the_archive(self, monkeypatch): + # Someone who exports the budget for this resume means it, and the + # archived value is older information. + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "QPX") + monkeypatch.setenv("HYPERLOOM_MAX_LATENCY_MS", "250") + args = self._args() + cli._restore_partition_lever_from_state(args, self._state(("DPX",), 2, 400.0)) + assert args.compute_partition_modes == "QPX" + assert args.max_latency_ms == 250.0 + + def test_a_re_passed_flag_outranks_both(self, monkeypatch): + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "QPX") + monkeypatch.setenv("HYPERLOOM_MAX_LATENCY_MS", "250") + args = self._args(compute_partition_modes="cpx", max_latency_ms=99.0) + cli._restore_partition_lever_from_state(args, self._state(("DPX",), 2, 400.0)) + assert args.compute_partition_modes == "cpx" + assert args.max_latency_ms == 99.0 + + def test_streams_two_is_distinguishable_from_streams_unset(self, monkeypatch): + # The reason --streams-per-partition defaults to None: with a default of + # 2 a resume cannot tell "not passed" from "passed 2", and would + # overwrite a persisted 4 every time. + monkeypatch.delenv(pl.STREAMS_PER_PARTITION_ENV, raising=False) + unset = self._args() + cli._restore_partition_lever_from_state(unset, self._state(("DPX",), 4, 0.0)) + assert unset.streams_per_partition == 4 + + explicit = self._args(streams_per_partition=2) + cli._restore_partition_lever_from_state(explicit, self._state(("DPX",), 4, 0.0)) + assert explicit.streams_per_partition == 2 + + def test_persist_writes_the_live_contract_back_onto_state(self, monkeypatch): + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "DPX,CPX") + monkeypatch.setenv(pl.STREAMS_PER_PARTITION_ENV, "3") + monkeypatch.setenv("HYPERLOOM_MAX_LATENCY_MS", "275.5") + state = self._state() + cli._persist_partition_lever(state) + assert state.compute_partition_modes == ["DPX", "CPX"] + assert state.streams_per_partition == 3 + assert state.latency_budget_ms == 275.5 + + def test_persist_records_the_lever_being_off(self, monkeypatch): + monkeypatch.delenv(pl.PARTITION_MODES_ENV, raising=False) + monkeypatch.delenv(pl.STREAMS_PER_PARTITION_ENV, raising=False) + monkeypatch.delenv("HYPERLOOM_MAX_LATENCY_MS", raising=False) + state = self._state(("DPX",), 3, 400.0) + cli._persist_partition_lever(state) + assert state.compute_partition_modes == [] + assert state.latency_budget_ms == 0.0 + + +def test_streams_per_partition_parses_to_none_when_not_passed(): + """Locks the parser half of the resume contract. + + ``_restore_partition_lever_from_state`` can only tell "not passed" from + "passed 2" if the flag's default stays None. Giving it a default of 2 -- + which reads as harmless, since 2 is the documented default -- would make + every resume overwrite a persisted streams count with 2 and quietly change + the experiment. + """ + from hyperloom.inference_optimizer.cli.parser import _build_parser + + args = _build_parser().parse_args(["optimize", "--model", "/tmp/m"]) + assert args.streams_per_partition is None + assert args.compute_partition_modes is None + assert args.max_latency_ms is None + + passed = _build_parser().parse_args( + ["optimize", "--model", "/tmp/m", "--streams-per-partition", "2"] + ) + assert passed.streams_per_partition == 2 + + +def test_read_session_lever_tolerates_a_malformed_budget(monkeypatch): + # The env is machine-written, but a hand-edited resume script is not, and a + # crash in the seed path would lose the session rather than the value. + monkeypatch.setenv("HYPERLOOM_MAX_LATENCY_MS", "not-a-number") + monkeypatch.setenv(pl.PARTITION_MODES_ENV, " DPX , , CPX ") + modes, streams, budget = pl.read_session_lever() + assert modes == ("DPX", "CPX") + assert budget == 0.0 + assert streams >= 1 diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py index 382131272e..eeb20db759 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -85,6 +85,29 @@ def streams_per_partition() -> int: return 2 +def read_session_lever() -> tuple[tuple[str, ...], int, float]: + """Read the session lever back out of the env the CLI published. + + One reader for the several places that need it -- seeding the manifest, + persisting across a resume, reporting the session's shape -- because that + env is the canonical post-validation form of these three flags. Parsing it + independently at each site is that many chances for the manifest to record + a contract the executors were never handed. + + Returns: + ``(modes, streams_per_partition, latency_budget_ms)``. Empty modes and a + zero budget each mean off. + """ + from ._latency_budget import LATENCY_BUDGET_ENV + + modes = tuple(m.strip() for m in os.environ.get(PARTITION_MODES_ENV, "").split(",") if m.strip()) + try: + budget = max(0.0, float(os.environ.get(LATENCY_BUDGET_ENV, "").strip() or 0.0)) + except ValueError: + budget = 0.0 + return modes, streams_per_partition(), budget + + def requested_mode(envs: dict[str, Any] | None) -> str: """Return the partition mode this run asks for, or ``""`` when none. From f4c59763e6a4229b98b4f0e47989e7f0951ca9f9 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Mon, 24 Aug 2026 19:06:49 +0000 Subject: [PATCH 03/10] feat(partition): scope the lever to what the card reports, not to a table The mode list was validated only against the four known names, and the partition ladder came from a static table. Neither was checked against the board. ``amd-smi partition -a`` states both, so it is now asked. * ``read_partition_profiles`` reads the card's own profiles: mode, index, ``num_partitions``, XCC instances per partition, and the memory-partition caps each compute mode can pair with. The report is sparse -- a profile's first row names it and carries its XCC count, the rows after it continue the same profile with its DECODER/DMA/JPEG resources under blank identity fields -- so only the named rows are taken. A parser treating every row as a profile would invent four per mode. * The scope is checked at launch. A mode this board does not offer now exits 2 with the card's actual list, instead of reaching ``set_partition_mode`` and failing there: partway through a session, on a privileged mutation of shared hardware. That is the same argument ``parse_modes`` already made for refusing a typo at parse time; the board's real capabilities were simply not available to it before. * The query needs the same elevation as the set -- unprivileged, amd-smi fills every field with "N/A" -- so it degrades to *no answer*, never to *supports nothing*, and an unelevated session is told its request went unvalidated rather than being blocked. "Not validated" and "validated as fine" do not look alike in the log. * ``partition_count_conflicts`` turns ``MODE_PARTITION_COUNTS`` from an assumption into a checked one. That table drives every CU calculation, and partition devices are then found by matching the CU count exactly, so a board whose ladder is not 1/2/4/8 would not disagree loudly -- the benchmark would just find no device of the expected width. * ``layout_for`` refuses a CU count that does not divide evenly instead of flooring it, for the same reason: a floored width matches no device, and the eventual error names the wrong cause. All four boards in the identity table divide evenly today; this is what catches the one that does not. Memory partitioning is still not a lever. The per-profile NPS caps are captured because the card reports them and the pairing is a real constraint -- on MI355X, SPX is NPS1-only while the split modes accept NPS2 -- but nothing acts on them yet, and switching NPS needs a driver reload. Also corrects the DEFAULT_MODE comment, which had that relationship backwards. Co-Authored-By: Claude Opus 5 Co-authored-by: Cursor (cherry picked from commit f2b58a2e146c305edce30828efb0b66e7060bb09) --- src/hyperloom/common/gpu_partition.py | 194 +++++++++++++++++- .../common/tests/test_gpu_partition.py | 184 +++++++++++++++++ .../inference_optimizer/cli/__init__.py | 52 ++++- .../inference_optimizer/cli/parser.py | 6 +- 4 files changed, 429 insertions(+), 7 deletions(-) diff --git a/src/hyperloom/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py index 9af9a11246..dca11baeb2 100644 --- a/src/hyperloom/common/gpu_partition.py +++ b/src/hyperloom/common/gpu_partition.py @@ -31,6 +31,14 @@ devices 7 and 8, not 0 and 1. Selecting by index measures a full card and reports it as a partition -- a wrong number with no error attached. Callers get :func:`partition_device_predicate` for a CU-count test instead. +* **The card decides what it supports, not this table.** A mode's name being one + of the four known ones says nothing about whether this board offers it, and + :data:`MODE_PARTITION_COUNTS` is an assumption about the ladder's width. + ``amd-smi partition -a`` states both, so :func:`read_partition_profiles` asks + and :func:`partition_count_conflicts` checks the assumption against the + answer. The query needs the same privilege as the set and degrades to *no + answer* -- never to *supports nothing* -- so an unelevated session stays + usable and is told its request went unvalidated. """ from __future__ import annotations @@ -58,8 +66,10 @@ "CPX": 8, } -#: Mode a session is restored to. SPX is the only mode every board supports and -#: the only one compatible with NPS1, so it is the safe terminal state. +#: Mode a session is restored to. SPX is the only mode every board supports, and +#: on MI355X it is the only profile whose memory-partition caps are NPS1 alone +#: (the split modes accept NPS1 or NPS2), so it is the safe terminal state under +#: the NPS1 this optimizer assumes. DEFAULT_MODE = "SPX" _SET_TIMEOUT_S = 120.0 @@ -128,6 +138,156 @@ def describe(self) -> str: return f"{self.mode} ({self.partitions} x {self.cu_per_partition} CU{mem})" +@dataclass(frozen=True) +class PartitionProfile: + """One compute-partition profile the card reports it can enter. + + This is the card's own answer, not a derivation. ``partitions`` and + ``xcc_per_partition`` come from ``num_partitions`` and the profile's ``XCC`` + resource count, so a board whose ladder is not the usual 1/2/4/8 over eight + XCDs describes itself correctly without an edit here. + + Attributes: + mode: Canonical mode name, with amd-smi's "current" asterisk stripped. + index: The card's own profile index. + partitions: Devices the card presents in this profile. + xcc_per_partition: XCC (compute die) instances each partition gets. + memory_modes: NPS modes this profile can be combined with. Captured + because the card reports it and the pairing is a real constraint -- + on MI355X, SPX is NPS1-only while the split modes accept NPS2 -- + but nothing here acts on it yet: memory partitioning is not a lever, + and switching NPS needs a driver reload. + """ + + mode: str + index: int + partitions: int + xcc_per_partition: int + memory_modes: tuple[str, ...] = () + + +def read_partition_profiles(gpu_id: int) -> tuple[PartitionProfile, ...]: + """Read the compute-partition profiles a card reports it supports. + + Needs the same privilege as the set: ``amd-smi partition -a`` fills every + field with ``"N/A"`` when run unprivileged, so an unelevated caller gets an + empty tuple rather than a wrong answer. That is why this returns "unknown" + instead of raising -- a session that cannot query capabilities is the normal + case, not an error, and the caller decides whether to proceed unvalidated. + + Args: + gpu_id: GPU to interrogate. + + Returns: + The profiles the card reports, or ``()`` when it reported none. + """ + try: + payload = _amd_smi_json( + ["partition", "-a", "-g", str(gpu_id)], + _READ_TIMEOUT_S, + privileged=True, + ) + except PartitionError as exc: + log.debug("could not read partition profiles for GPU %d: %s", gpu_id, exc) + return () + + 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)] + + profiles: list[PartitionProfile] = [] + for row in rows: + # The report is sparse: a profile's first row names it and carries its + # XCC count, and the rows after it continue the same profile with its + # other resources (DECODER/DMA/JPEG) under blank identity fields. Only + # the named rows describe a profile, so the blanks are skipped rather + # than carried forward. + mode = str(row.get("accelerator_type") or "").strip().upper().rstrip("*") + if not mode or mode == "N/A" or mode not in MODE_PARTITION_COUNTS: + continue + try: + index = int(row.get("profile_index")) + partitions = int(row.get("num_partitions")) + except (TypeError, ValueError): + continue + if partitions < 1: + continue + xcc = 0 + if str(row.get("resource_type") or "").strip().upper() == "XCC": + try: + xcc = int(row.get("resource_instances")) + except (TypeError, ValueError): + xcc = 0 + caps = tuple( + part.strip().upper() + for part in str(row.get("memory_partition_caps") or "").split(",") + if part.strip() and part.strip().upper() != "N/A" + ) + profiles.append( + PartitionProfile( + mode=mode, + index=index, + partitions=partitions, + xcc_per_partition=xcc, + memory_modes=caps, + ) + ) + return tuple(profiles) + + +def supported_modes(gpu_id: int) -> tuple[str, ...]: + """Return the compute-partition modes a card reports, or ``()`` if unknown.""" + return tuple(p.mode for p in read_partition_profiles(gpu_id)) + + +def unsupported_modes(modes: Sequence[str], gpu_id: int = 0) -> tuple[str, ...]: + """Return which of ``modes`` the card says it cannot enter. + + Empty when the card reported no profile table, because "the query failed" + and "the card supports everything asked" must not collapse into the same + answer. Callers distinguish the two with :func:`supported_modes`. + + Args: + modes: Canonical modes the session wants to evaluate. + gpu_id: GPU whose capabilities decide. + + Returns: + The requested modes the card does not list, in the order given. + """ + available = supported_modes(gpu_id) + if not available: + return () + return tuple(m for m in modes if str(m).strip().upper() not in available) + + +def partition_count_conflicts(gpu_id: int = 0) -> tuple[str, ...]: + """Return modes where the card contradicts :data:`MODE_PARTITION_COUNTS`. + + The table drives every CU calculation in this module, and partition devices + are then found by matching that CU count exactly. If a board's real ladder + differs, nothing downstream disagrees loudly -- the benchmark simply finds + no device of the expected width. Now that the card states its own + ``num_partitions``, the assumption is checkable, so it gets checked. + + Args: + gpu_id: GPU whose profiles to compare. + + Returns: + Descriptions of each disagreement, empty when the card agrees or could + not be queried. + """ + conflicts: list[str] = [] + for profile in read_partition_profiles(gpu_id): + expected = MODE_PARTITION_COUNTS.get(profile.mode) + if expected is not None and expected != profile.partitions: + conflicts.append(f"{profile.mode}: card reports {profile.partitions} partitions, table says {expected}") + return tuple(conflicts) + + def parse_modes(raw: str | Sequence[str] | None) -> tuple[str, ...]: """Parse an operator-supplied mode list into canonical order-preserving modes. @@ -187,6 +347,18 @@ def layout_for(gpu_type: str | None, mode: str, hbm_gib: float | None = None) -> if identity is None: raise PartitionError(f"unknown gpu_type {gpu_type!r}; cannot size partitions without the board's CU count") cu_total = identity[1] + if cu_total % partitions: + # Flooring here would be silent and then fatal much later: partition + # devices are selected by matching this exact CU count, so a floored + # value matches nothing and the benchmark reports "mode did not take + # effect" -- true, but about the wrong cause. Every board in the + # identity table divides evenly today; this is what catches the one + # that does not. + raise PartitionError( + f"{gpu_type} has {cu_total} CU, which does not divide into {partitions} " + f"{canonical} partitions; the per-partition CU count would be wrong and " + f"device selection matches on it exactly" + ) return PartitionLayout( mode=canonical, partitions=partitions, @@ -250,9 +422,17 @@ def fits_in_partition( return required_gib * max(1, int(streams_per_partition)) <= layout.gib_per_partition -def _amd_smi_json(args: Sequence[str], timeout_s: float) -> object: - """Run an ``amd-smi`` subcommand with ``--json`` and parse its output.""" - cmd = ["amd-smi", *args, "--json"] +def _amd_smi_json(args: Sequence[str], timeout_s: float, privileged: bool = False) -> object: + """Run an ``amd-smi`` subcommand with ``--json`` and parse its output. + + Args: + args: Subcommand and its flags. + timeout_s: Per-call timeout. + privileged: Route through the opt-in sudo prefix. Needed for the + accelerator-profile query, which silently degrades every field to + ``"N/A"`` rather than failing when it lacks privilege. + """ + cmd = [*(_set_prefix() if privileged else []), "amd-smi", *args, "--json"] try: proc = subprocess.run( # noqa: S603 — fixed argv, no shell cmd, @@ -451,6 +631,7 @@ def partitioned(gpu_id: int, mode: str, restore_to: str | None = None) -> Iterat "PARTITION_SUDO_ENV", "PartitionError", "PartitionLayout", + "PartitionProfile", "fits_in_partition", "layout_for", "parse_modes", @@ -458,5 +639,8 @@ def partitioned(gpu_id: int, mode: str, restore_to: str | None = None) -> Iterat "partitioned", "read_partition_mode", "read_partition_modes", + "read_partition_profiles", "set_partition_mode", + "supported_modes", + "unsupported_modes", ] diff --git a/src/hyperloom/common/tests/test_gpu_partition.py b/src/hyperloom/common/tests/test_gpu_partition.py index 73a237353d..c064daa2f1 100644 --- a/src/hyperloom/common/tests/test_gpu_partition.py +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -309,3 +309,187 @@ def test_permission_failure_that_exits_zero_is_still_caught(monkeypatch): monkeypatch.setattr(subprocess, "run", fake) with pytest.raises(PartitionError, match="returned success but the mode did not change"): set_partition_mode(0, "CPX") + + +#: Verbatim shape of ``amd-smi partition -a --json`` on an MI355X, trimmed to +#: two profiles. The sparse continuation rows are the point: a profile's first +#: row names it and carries its XCC count, and the rows after it describe the +#: same profile's other resources under blank identity fields. A parser that +#: treats every row as a profile invents four per mode. +PROFILE_PAYLOAD = { + "partition_profiles": [ + { + "gpu_id": 0, + "profile_index": 0, + "memory_partition_caps": "NPS1", + # amd-smi marks the live profile with a trailing asterisk. + "accelerator_type": "SPX*", + "num_partitions": 1, + "resource_index": 0, + "resource_type": "XCC", + "resource_instances": 8, + }, + { + "gpu_id": "", + "profile_index": "", + "memory_partition_caps": "", + "accelerator_type": "", + "num_partitions": "", + "resource_index": 1, + "resource_type": "DECODER", + "resource_instances": 4, + }, + { + "gpu_id": "", + "profile_index": 3, + "memory_partition_caps": "NPS1,NPS2", + "accelerator_type": "CPX", + "num_partitions": 8, + "resource_index": 4, + "resource_type": "XCC", + "resource_instances": 1, + }, + { + "gpu_id": "", + "profile_index": "", + "memory_partition_caps": "", + "accelerator_type": "", + "num_partitions": "", + "resource_index": 5, + "resource_type": "JPEG", + "resource_instances": 20, + }, + ] +} + +#: What the same command returns without privilege: the shape is there and every +#: value is gone. This is the reason the query reports "unknown" rather than +#: raising -- and the reason it must not read as "supports nothing". +PROFILE_PAYLOAD_UNPRIVILEGED = { + "partition_profiles": [ + { + "gpu_id": 0, + "profile_index": "N/A", + "memory_partition_caps": "N/A", + "accelerator_type": "N/A", + "num_partitions": "N/A", + "resource_type": "N/A", + "resource_instances": "N/A", + } + ] +} + + +def _profiles(monkeypatch, payload): + """Install a fake amd-smi answering the profile query with ``payload``.""" + + def run(cmd, **kwargs): + return subprocess.CompletedProcess(cmd, 0, json.dumps(payload), "") + + monkeypatch.setattr(subprocess, "run", run) + + +class TestCapabilityQuery: + def test_reads_the_profiles_the_card_reports(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD) + profiles = gpu_partition.read_partition_profiles(0) + assert [p.mode for p in profiles] == ["SPX", "CPX"] + assert [p.partitions for p in profiles] == [1, 8] + assert [p.xcc_per_partition for p in profiles] == [8, 1] + # The pairing constraint is captured even though nothing acts on it yet: + # SPX is NPS1-only here while CPX accepts NPS2. + assert profiles[0].memory_modes == ("NPS1",) + assert profiles[1].memory_modes == ("NPS1", "NPS2") + + def test_strips_the_current_profile_marker(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD) + assert gpu_partition.supported_modes(0) == ("SPX", "CPX") + + def test_unprivileged_query_is_unknown_not_empty_support(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD_UNPRIVILEGED) + assert gpu_partition.supported_modes(0) == () + # The distinction that matters: an unanswerable query must not reject + # every mode, or an unprivileged session cannot ask for anything. + assert gpu_partition.unsupported_modes(["DPX", "CPX"]) == () + + def test_a_missing_amd_smi_is_unknown_rather_than_fatal(self, monkeypatch): + def missing(cmd, **kwargs): + raise FileNotFoundError("amd-smi") + + monkeypatch.setattr(subprocess, "run", missing) + assert gpu_partition.read_partition_profiles(0) == () + assert gpu_partition.supported_modes(0) == () + + def test_names_the_modes_the_card_does_not_offer(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD) + assert gpu_partition.unsupported_modes(["SPX", "DPX", "QPX"]) == ("DPX", "QPX") + assert gpu_partition.unsupported_modes(["spx", "cpx"]) == () + + def test_profile_query_is_routed_through_sudo_when_opted_in(self, monkeypatch): + seen: list[list[str]] = [] + + def run(cmd, **kwargs): + seen.append(list(cmd)) + return subprocess.CompletedProcess(cmd, 0, json.dumps(PROFILE_PAYLOAD), "") + + monkeypatch.setattr(subprocess, "run", run) + monkeypatch.setenv("HYPERLOOM_PARTITION_SUDO", "1") + gpu_partition.read_partition_profiles(0) + assert seen[0][:2] == ["sudo", "-n"] + assert "amd-smi" in seen[0] + + seen.clear() + monkeypatch.delenv("HYPERLOOM_PARTITION_SUDO") + gpu_partition.read_partition_profiles(0) + assert seen[0][0] == "amd-smi" + + +class TestStaticTableIsVerifiedNotTrusted: + def test_no_conflict_when_the_card_agrees(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD) + assert gpu_partition.partition_count_conflicts(0) == () + + def test_conflict_is_reported_when_the_ladder_differs(self, monkeypatch): + payload = { + "partition_profiles": [ + { + "gpu_id": 0, + "profile_index": 3, + "accelerator_type": "CPX", + # A board with six XCDs would land here. + "num_partitions": 6, + "memory_partition_caps": "NPS1", + "resource_type": "XCC", + "resource_instances": 1, + } + ] + } + _profiles(monkeypatch, payload) + conflicts = gpu_partition.partition_count_conflicts(0) + assert len(conflicts) == 1 + assert "card reports 6" in conflicts[0] and "table says 8" in conflicts[0] + + def test_unqueryable_card_reports_no_conflict(self, monkeypatch): + _profiles(monkeypatch, PROFILE_PAYLOAD_UNPRIVILEGED) + assert gpu_partition.partition_count_conflicts(0) == () + + +def test_layout_refuses_a_cu_count_that_does_not_divide(monkeypatch): + # Flooring would be silent, and then fatal much later and for an apparently + # unrelated reason: device selection matches the per-partition CU count + # exactly, so a floored value matches no device at all. + monkeypatch.setitem( + gpu_partition.AMD_GPU_DISPATCH_IDENTITIES, "oddboard", ("gfx950", 300, "x") + ) + with pytest.raises(PartitionError, match="does not divide"): + layout_for("oddboard", "CPX") + # The same board is fine in a mode its CU count does divide by. + assert layout_for("oddboard", "QPX").cu_per_partition == 75 + + +def test_every_shipped_board_divides_across_every_mode(): + """Guards the assumption the divisibility check exists to catch.""" + for board in gpu_partition.AMD_GPU_DISPATCH_IDENTITIES: + for mode in gpu_partition.MODE_PARTITION_COUNTS: + layout = layout_for(board, mode) + assert layout.cu_per_partition > 0 diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 370444db7b..9327f51061 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1661,7 +1661,14 @@ def _export_partition_lever( Returns: The canonical modes, empty when the lever is off. """ - from hyperloom.common.gpu_partition import PartitionError, parse_modes + from hyperloom.common.gpu_partition import ( + PARTITION_SUDO_ENV, + PartitionError, + parse_modes, + partition_count_conflicts, + supported_modes, + unsupported_modes, + ) budget = float(max_latency_ms or 0.0) if budget > 0: @@ -1686,6 +1693,49 @@ def _export_partition_lever( file=sys.stderr, ) sys.exit(2) + + # Scope the request to what this card says it can do. The name being one of + # the four known modes does not mean this board offers it, and the + # alternative to checking here is discovering it at the apply site -- a + # privileged mutation partway through a session, which is the whole reason + # the rest of this validation happens at launch. + # + # A card that reports nothing is the ordinary unprivileged case, not an + # error: the profile query needs the same elevation as the set. So an + # unanswerable query warns and proceeds rather than blocking a session that + # may be perfectly able to run -- and says so, because "not validated" and + # "validated as fine" must not look alike in a log. + available = supported_modes(0) + if not available: + print( + "WARN: could not read this card's supported partition profiles, so " + f"{','.join(modes)} is unvalidated. amd-smi reports profiles only when " + f"elevated; set {PARTITION_SUDO_ENV}=1 to check the request here rather " + "than at the first mode change.", + file=sys.stderr, + ) + else: + rejected = unsupported_modes(modes) + if rejected: + print( + f"ERROR: this card does not support {','.join(rejected)}. " + f"It reports: {','.join(available)}.", + file=sys.stderr, + ) + sys.exit(2) + # The card is also the authority on how many partitions a mode makes, + # and that number drives the CU arithmetic device selection matches on. + # A disagreement means the sizing is wrong, so it stops the session here + # rather than surfacing as a benchmark that finds no device. + conflicts = partition_count_conflicts(0) + if conflicts: + print( + "ERROR: this card's partition ladder disagrees with the built-in " + "table, so partition sizing would be wrong:\n " + "\n ".join(conflicts), + file=sys.stderr, + ) + sys.exit(2) + print(f"Compute partitioning : {','.join(modes)} (card reports {','.join(available)})") if budget <= 0: # Not fatal: maximizing offline throughput regardless of per-request # latency is a legitimate goal. But it is not usually what someone diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index 77b30ee4f8..e3d0fecd4d 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -298,7 +298,11 @@ def _build_parser() -> argparse.ArgumentParser: "and pays only in aggregate throughput at concurrency -- see " "--streams-per-partition. Requires a privileged amd-smi on the " "host: the mode is a property of the card, not of the process, so a " - "session confined to an unprivileged container cannot set it. The " + "session confined to an unprivileged container cannot set it. When " + "elevated, the request is checked against the profiles the card " + "actually reports and a mode it does not offer is refused at launch; " + "unelevated, the check is skipped with a warning because amd-smi " + "reports profiles only to root. The " "session restores the mode it found on the way out. Pair with " "--max-latency-ms; without a budget the search picks the narrowest " "partition on offer, which is the slowest one per request.", From e6cf834a8279432b8f09943920be255e027e8c76 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 01:26:24 +0000 Subject: [PATCH 04/10] feat(partition): make the mode list drive the search, not just bound it The lever validated and recorded --compute-partition-modes but nothing turned that list into variants, so the modes were only ever reachable by hand. partition_lever_grid expands the validated list into one env-only variant per mode and seeds them ahead of the grid. Ordering them first is what makes the rest of the round useful: explore already stacks a KEEP'd variant's envs onto everything after it, so the winning mode becomes the topology the remaining knobs are explored inside, and each later mode has to beat the best mode so far rather than the original baseline. The prepend-and-dedupe the framework levers already did is factored out and shared, keeping its payload in its own name because the attribution pass needs it. Scriptable frameworks only. plan_partition_run is called by the scriptable runner and not the serving path, so a serving framework would take the env, change nothing, and file the number under a mode the card was never in -- refused at launch, with the generator declining as a second line. Co-authored-by: Cursor (cherry picked from commit 69c4731fc647207f2c5495cc2d5dca53361db732) --- .../inference_optimizer/cli/__init__.py | 20 +++ .../inference_optimizer/cli/parser.py | 8 +- .../tests/test_partition_lever.py | 129 ++++++++++++++++++ .../actions/executors/_partition_lever.py | 98 +++++++++++++ .../orchestrator/actions/executors/explore.py | 69 ++++++++-- 5 files changed, 314 insertions(+), 10 deletions(-) diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 9327f51061..2c5b8398ad 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1642,6 +1642,7 @@ def _export_partition_lever( modes_raw: str | None, streams_per_partition: int, max_latency_ms: float | None, + framework: str | None = None, ) -> tuple[str, ...]: """Validate and project the compute-partition lever into env. @@ -1657,6 +1658,9 @@ def _export_partition_lever( modes_raw: The raw ``--compute-partition-modes`` value. streams_per_partition: The resolved ``--streams-per-partition``. max_latency_ms: The resolved ``--max-latency-ms``, if any. + framework: The session's framework. Checked because only the scriptable + runner applies a mode; unset skips the check for callers that run + before the framework is resolved. Returns: The canonical modes, empty when the lever is off. @@ -1694,6 +1698,20 @@ def _export_partition_lever( ) sys.exit(2) + # Only the scriptable runner establishes a mode. On a serving framework the + # env would be delivered, no partition would be set, and the number would be + # filed under a mode the card was never in -- a wrong answer that looks + # exactly like a right one. Refused here rather than dropped quietly at grid + # time, because the operator asked for something this session cannot do. + if framework and not framework_registry.is_scriptable(framework): + print( + f"ERROR: --compute-partition-modes needs a scriptable framework; {framework!r} " + "runs a server, and its benchmarks would report the card's current topology " + "under the requested mode's name.", + file=sys.stderr, + ) + sys.exit(2) + # Scope the request to what this card says it can do. The name being one of # the four known modes does not mean this board offers it, and the # alternative to checking here is discovering it at the apply site -- a @@ -1884,6 +1902,7 @@ async def _run_optimize(args: argparse.Namespace) -> int: modes_raw=getattr(args, "compute_partition_modes", None), streams_per_partition=int(getattr(args, "streams_per_partition", 2) or 2), max_latency_ms=getattr(args, "max_latency_ms", None), + framework=str(getattr(args, "framework", "") or "").strip().lower(), ) # Project resolved workload knobs into env for the fresh-launch path only. # A resume must NOT export here: ``args.tp``/etc. are still unresolved @@ -2111,6 +2130,7 @@ async def _run_optimize(args: argparse.Namespace) -> int: modes_raw=getattr(args, "compute_partition_modes", None), streams_per_partition=int(getattr(args, "streams_per_partition", None) or 2), max_latency_ms=getattr(args, "max_latency_ms", None), + framework=str(getattr(args, "framework", "") or "").strip().lower(), ) _persist_partition_lever(state) if state.compute_partition_modes: diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index e3d0fecd4d..d32f874ff4 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -293,7 +293,13 @@ def _build_parser() -> argparse.ArgumentParser: default=None, metavar="MODES", help="Comma-separated compute-partition modes to evaluate, e.g. " - "'spx,dpx,cpx'. Off by default. Partitioning a card only ever gives " + "'spx,dpx,cpx'. Off by default. Each mode becomes one explore variant, " + "tried in the order given and ahead of the rest of the grid, so a mode " + "that is kept becomes the topology the remaining variants are then " + "explored inside. Scriptable frameworks only: the mode is established " + "by the scriptable runner, and on a serving framework the request " + "would be refused at launch rather than measured under the wrong " + "name. Partitioning a card only ever gives " "one stream fewer CUs, so it cannot improve single-stream latency " "and pays only in aggregate throughput at concurrency -- see " "--streams-per-partition. Requires a privileged amd-smi on the " diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py index 007068ebdf..77eba14511 100644 --- a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -21,6 +21,7 @@ from hyperloom.inference_optimizer import cli from hyperloom.orchestrator.actions.executors import _partition_lever as pl from hyperloom.orchestrator.actions.executors import bypass_scriptable as bs +from hyperloom.orchestrator.actions.executors import explore from hyperloom.orchestrator.actions.executors._latency_budget import ( REASON_OVER_BUDGET, REASON_UNMEASURED, @@ -316,3 +317,131 @@ def test_read_session_lever_tolerates_a_malformed_budget(monkeypatch): assert modes == ("DPX", "CPX") assert budget == 0.0 assert streams >= 1 + + +class TestModeAxis: + """Turning the session's mode list into variants the search actually runs. + + Declaring the modes is not the same as trying them. These cover the + expansion, the order it has to keep, and the one framework combination that + must not produce variants at all. + """ + + @staticmethod + def _state(modes): + return SimpleNamespace(compute_partition_modes=list(modes)) + + def test_the_session_list_becomes_one_variant_per_mode(self): + grid = pl.partition_lever_grid( + {"compute_partition_modes": "spx,dpx,cpx"}, None, framework="custom" + ) + assert [v["name"] for v in grid] == [ + "partition_spx", + "partition_dpx", + "partition_cpx", + ] + # Env-only, and carrying nothing but the selector: a mode variant that + # also moved a server flag would confound the two. + assert [v["extra_envs"] for v in grid] == [ + {pl.PARTITION_MODE_ENV: "SPX"}, + {pl.PARTITION_MODE_ENV: "DPX"}, + {pl.PARTITION_MODE_ENV: "CPX"}, + ] + assert {v["extra_args"] for v in grid} == {""} + assert {v["provenance"] for v in grid} == {"partition_lever"} + + def test_the_operators_order_is_preserved(self): + # The list is a search order, not a set: the stack advances mode by mode, + # so reordering changes which mode each later one has to beat. + grid = pl.partition_lever_grid({"compute_partition_modes": "cpx,spx"}, None, framework="xdit") + assert [v["name"] for v in grid] == ["partition_cpx", "partition_spx"] + + def test_no_modes_seeds_nothing(self): + assert pl.partition_lever_grid({}, None, framework="custom") == [] + assert pl.partition_lever_grid(None, self._state([]), framework="custom") == [] + + def test_a_serving_framework_gets_no_variants(self): + """The mislabelling guard, at the grid rather than the card. + + Only the scriptable runner calls ``plan_partition_run``. On a serving + framework the env would ride along, no partition would be established, + and the result would be filed under the requested mode -- a number that + is wrong in a way nothing downstream can detect. + """ + for framework in ("sglang", "vllm", ""): + assert ( + pl.partition_lever_grid( + {"compute_partition_modes": "dpx"}, None, framework=framework + ) + == [] + ) + + def test_an_unusable_mode_list_seeds_nothing_rather_than_raising(self, monkeypatch): + # Grid assembly is not the place to end a session; the CLI already + # refused this at launch, so reaching here means the env was edited. + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "spx,nope") + assert pl.partition_lever_grid(None, None, framework="custom") == [] + + def test_mode_precedence_is_most_specific_first(self, monkeypatch): + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "cpx") + state = self._state(["QPX"]) + assert pl.resolve_session_modes({"compute_partition_modes": "dpx"}, state) == ("DPX",) + assert pl.resolve_session_modes({}, state) == ("QPX",) + assert pl.resolve_session_modes({}, None) == ("CPX",) + + +class TestModeAxisOrdering: + """The mode has to be decided before the knobs that are tuned against it. + + Explore stacks a KEEP'd variant's envs onto every variant after it, so + putting the mode axis first is what makes the rest of the grid get + re-explored inside the winning partition. Landing it at the end instead + would measure every knob against the old topology and then change the + topology afterwards. + """ + + def test_seeded_variants_land_in_front(self): + grid, fresh = explore._prepend_fresh_variants( + [{"name": "llm_flag_a"}, {"name": "llm_flag_b"}], + [{"name": "partition_dpx"}], + ) + assert [v["name"] for v in grid] == [ + "partition_dpx", + "llm_flag_a", + "llm_flag_b", + ] + assert [v["name"] for v in fresh] == ["partition_dpx"] + + def test_a_name_the_grid_already_uses_is_left_alone(self): + # An operator or specialist who named the variant said something more + # specific than the generated default; the generated one steps aside. + pinned = {"name": "partition_dpx", "extra_envs": {"CUSTOM": "1"}} + grid, fresh = explore._prepend_fresh_variants( + [pinned], [{"name": "partition_dpx", "extra_envs": {}}] + ) + assert grid == [pinned] + assert fresh == [] + + def test_nothing_to_add_leaves_the_grid_untouched(self): + original = [{"name": "llm_flag_a"}] + grid, fresh = explore._prepend_fresh_variants(original, []) + assert grid == original + assert fresh == [] + + +def test_launch_refuses_the_lever_on_a_serving_framework(capsys): + """Fail at launch, not by quietly seeding an empty axis. + + An operator who passed the flag expects modes to be tried. Dropping them at + grid time with only a log line would look like the lever ran and found + nothing worth keeping. + """ + with pytest.raises(SystemExit) as excinfo: + cli._export_partition_lever( + modes_raw="dpx", + streams_per_partition=2, + max_latency_ms=400.0, + framework="sglang", + ) + assert excinfo.value.code == 2 + assert "scriptable framework" in capsys.readouterr().err diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py index eeb20db759..b1e1bfb350 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -38,6 +38,7 @@ PartitionError, PartitionLayout, layout_for, + parse_modes, partitioned, ) @@ -108,6 +109,101 @@ def read_session_lever() -> tuple[tuple[str, ...], int, float]: return modes, streams_per_partition(), budget +def resolve_session_modes( + params: dict[str, Any] | None = None, + shared_state: Any = None, +) -> tuple[str, ...]: + """Resolve the modes this session is allowed to explore. + + Precedence is most-specific-first, matching + :func:`_latency_budget.resolve_latency_budget_ms`: an explicit task + parameter, then the session state the CLI seeded, then the environment. + + Args: + params: Task params, which may carry ``compute_partition_modes``. + shared_state: Live SharedState, which may carry the persisted list. + + Returns: + Canonical modes, empty when the lever is off. + """ + for candidate in ( + (params or {}).get("compute_partition_modes"), + getattr(shared_state, "compute_partition_modes", None), + os.environ.get(PARTITION_MODES_ENV), + ): + if not candidate: + continue + try: + modes = parse_modes(candidate) + except PartitionError as exc: + log.warning("ignoring unusable compute-partition modes %r: %s", candidate, exc) + continue + if modes: + return modes + return () + + +def partition_lever_grid( + params: dict[str, Any] | None = None, + shared_state: Any = None, + *, + framework: str | None, +) -> list[dict[str, Any]]: + """Expand the session's mode list into one explore variant per mode. + + Each variant is env-only, carrying nothing but + :data:`PARTITION_MODE_ENV`. That keeps the mode a plain point in the search + space: it is fingerprinted, deduplicated, gated and journalled by the same + code as every other variant, and a mode that loses is reverted like any + other losing knob. + + Every listed mode is emitted, including one that may already match the + card. This function stays pure -- no hardware read -- and a mode the + operator named explicitly is worth a measurement under the same harness as + its rivals rather than an inherited baseline number taken on trust. + + Only scriptable frameworks get variants. The mode is applied by + :func:`plan_partition_run`, which the scriptable runner calls and the + serving path does not, so emitting these for a server framework would + deliver the env, change nothing, and record the result under a mode the + card was never in -- the exact mislabelling this module refuses elsewhere. + The CLI rejects that combination at launch; this is the second line. + + Args: + params: Task params, consulted for the mode list. + shared_state: Live SharedState, consulted for the persisted list. + framework: Framework this round runs under. + + Returns: + Variant payload dicts ready for ``_grid_variants_from_payload``, or + ``[]`` when the lever is off or the framework cannot apply it. + """ + from hyperloom.inference_optimizer.framework_registry import is_scriptable + + modes = resolve_session_modes(params, shared_state) + if not modes: + return [] + if not is_scriptable(framework): + log.warning( + "compute-partition modes %s ignored: framework %r does not apply " + "partition modes, so the variants would measure the card's current " + "topology under another mode's name", + ",".join(modes), + framework or "", + ) + return [] + return [ + { + "name": f"partition_{mode.lower()}", + "extra_args": "", + "extra_envs": {PARTITION_MODE_ENV: mode}, + "note": f"compute-partition {mode}", + "provenance": "partition_lever", + } + for mode in modes + ] + + def requested_mode(envs: dict[str, Any] | None) -> str: """Return the partition mode this run asks for, or ``""`` when none. @@ -210,8 +306,10 @@ def maybe_hold_partition_mode(mode: str, *, gpu_type: str | None): "hold_partition_mode", "maybe_hold_partition_mode", "partition_gpu_id", + "partition_lever_grid", "plan_partition_run", "requested_mode", + "resolve_session_modes", "runtime_env", "streams_per_partition", ] diff --git a/src/hyperloom/orchestrator/actions/executors/explore.py b/src/hyperloom/orchestrator/actions/executors/explore.py index fe7abf4ef2..d52f97a9c4 100644 --- a/src/hyperloom/orchestrator/actions/executors/explore.py +++ b/src/hyperloom/orchestrator/actions/executors/explore.py @@ -85,6 +85,7 @@ latency_keep_block, resolve_latency_budget_ms, ) +from ._partition_lever import partition_lever_grid from ._ray_serving import maybe_serving_lease # DEFAULT_STACK_STABLE_PCT: post-KEEP confirmation floor; override via @@ -247,6 +248,36 @@ def _grid_variants_from_payload(payload: list[Any]) -> list[GridVariant]: return out +def _prepend_fresh_variants( + grid_payload: list[Any], + addition: list[dict[str, Any]], +) -> tuple[list[Any], list[dict[str, Any]]]: + """Prepend the entries of ``addition`` that the grid does not already name. + + Seeded axes go in front rather than on the end so a long LLM-supplied grid + cannot crowd them out of the round's budget, and -- for anything the later + variants are measured on top of -- so the stack resolves it first. + + Name collisions defer to the grid. An operator or specialist who already + named a variant has said something specific about it, and overwriting that + with a generated default would discard the more informed entry. + + Args: + grid_payload: The grid so far. + addition: Candidate payload dicts to prepend. + + Returns: + ``(grid, fresh)`` where ``fresh`` is what was actually added. + """ + if not addition: + return list(grid_payload), [] + existing = {str(v.get("name") or "") for v in grid_payload if isinstance(v, dict)} + fresh = [v for v in addition if str(v.get("name") or "") not in existing] + if not fresh: + return list(grid_payload), [] + return fresh + list(grid_payload), fresh + + def framework_lever_grid(shared_state: Any) -> list[dict[str, Any]]: """Build explore variants that attribute each registered rewrite lever. @@ -898,16 +929,36 @@ async def __call__(self, ctx) -> dict[str, Any]: # judge and better evidenced than a proposed config knob. Prepending also # means an LLM-supplied grid does not crowd the attribution out of the # round's budget. + # Kept in its own name: the round's attribution pass needs the payload it + # seeded from, not just what survived the merge. lever_payload = framework_lever_grid(extra.get("shared_state") or extra.get("state")) - if lever_payload: - existing_names = {str(v.get("name") or "") for v in grid_payload if isinstance(v, dict)} - fresh = [v for v in lever_payload if str(v.get("name") or "") not in existing_names] - if fresh: - log.info( - "explore: seeding %d framework-rewrite lever variant(s) for attribution", - len(fresh), - ) - grid_payload = fresh + list(grid_payload) + grid_payload, _lever_fresh = _prepend_fresh_variants(grid_payload, lever_payload) + if _lever_fresh: + log.info( + "explore: seeding %d framework-rewrite lever variant(s) for attribution", + len(_lever_fresh), + ) + # Compute-partition modes ahead of even those, because the mode decides + # the topology every other knob is then tuned against. Ordering it first + # is what makes the stack do the useful thing: a KEEP'd mode joins + # ``stack_extra_envs``, so the variants after it are measured inside the + # winning partition rather than against a shape the session has already + # moved off. It also means each mode has to beat the best mode so far + # rather than the original baseline. + grid_payload, _partition_fresh = _prepend_fresh_variants( + grid_payload, + partition_lever_grid( + params, + extra.get("shared_state") or extra.get("state"), + framework=framework, + ), + ) + if _partition_fresh: + log.info( + "explore: seeding %d compute-partition variant(s) ahead of the grid: %s", + len(_partition_fresh), + ",".join(str(v.get("name") or "?") for v in _partition_fresh), + ) if not grid_payload: # No LLM variants: fall through to the framework's programmatic # seed grid instead of failing the task. From 892f19aa2d14d857b33b3d85d8a37e866246b011 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 01:53:07 +0000 Subject: [PATCH 05/10] feat(partition): drop the modes whose partitions cannot hold the workload A narrow mode divides the card's memory while every stream on it keeps a full copy of the weights, so CPX at two streams is the first to run out. Measuring that costs a run and returns an OOM instead of a number: the variant is scored as a failure and the operator learns from a stack trace what arithmetic could have said for free. fits_in_partition already encoded the predicate and was called from nowhere; prune_infeasible_modes now calls it per requested mode and the generator emits what survives. Two inputs had to exist. read_hbm_gib asks the card its capacity, since boards sharing an ISA do not share one -- and answers only in SPX, because amd-smi reports VRAM per device and under a split mode a device is a partition, with nothing in the payload saying which the number is. Dividing an already-divided figure again would understate capacity eightfold and prune every mode that fits, so an ambiguous reading is reported as unknown. The footprint has two sources. peak_gib_per_stream, read from a report that carries it and carried onto current_best from the baseline, is the real one. Absent that, the checkpoint's own weight bytes serve as a lower bound and are used strictly as one: each stream holds its own copy, so "does not fit by the weights alone" is a proof while "fits" is no evidence. Pruning acts only on the former, so it cannot drop a mode that would have won -- and an unknown capacity or unknown footprint drops nothing at all rather than guessing. Every drop is warned with its arithmetic, because a list that quietly comes back shorter than the operator's is how this gets rediscovered as a bug. Co-authored-by: Cursor (cherry picked from commit a6cee1cb7b4df75823762f586f1ea0978abd4982) --- src/hyperloom/common/gpu_partition.py | 76 ++++++ .../tests/test_partition_lever.py | 223 ++++++++++++++++-- .../actions/executors/_partition_lever.py | 180 +++++++++++++- .../actions/executors/benchmark_result.py | 14 ++ src/hyperloom/orchestrator/loop/writeback.py | 5 + 5 files changed, 469 insertions(+), 29 deletions(-) diff --git a/src/hyperloom/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py index dca11baeb2..1612f4699d 100644 --- a/src/hyperloom/common/gpu_partition.py +++ b/src/hyperloom/common/gpu_partition.py @@ -422,6 +422,81 @@ def fits_in_partition( return required_gib * max(1, int(streams_per_partition)) <= layout.gib_per_partition +#: GiB per unit amd-smi may label a VRAM size with. The binary divisor is the +#: right one for its "MB": an MI355X with 288 GiB reports 294896 "MB", which is +#: mebibytes. "GB" is treated the same way for consistency with that. +_VRAM_UNIT_GIB: dict[str, float] = { + "B": 1.0 / (1024**3), + "KB": 1.0 / (1024**2), + "KIB": 1.0 / (1024**2), + "MB": 1.0 / 1024, + "MIB": 1.0 / 1024, + "GB": 1.0, + "GIB": 1.0, + "TB": 1024.0, + "TIB": 1024.0, +} + + +def read_hbm_gib(gpu_id: int = 0) -> float | None: + """Read a card's total HBM in GiB, or ``None`` when it cannot be determined. + + This is the capacity :func:`layout_for` needs to fill in + ``gib_per_partition``, and it is read from the device rather than tabled + because boards sharing an ISA do not share a capacity. + + Only answers for a card currently in a single partition. ``amd-smi`` reports + VRAM per device, and under a split mode a device is a partition, so the same + field means "the whole card" in SPX and "one eighth of it" in CPX with + nothing in the payload to say which. Dividing an already-divided figure by + the partition count would understate capacity eightfold and prune every mode + that actually fits. Rather than guess the field's scope, this reports + *unknown* outside SPX and lets the caller proceed without a memory opinion -- + the lever restores the card to SPX between runs, so the normal case answers. + + Args: + gpu_id: GPU to interrogate. + + Returns: + Total HBM in GiB, or ``None`` when unreadable, unparseable, or when the + card is partitioned and the figure's scope would be ambiguous. + """ + try: + mode = read_partition_mode(gpu_id) + except PartitionError as exc: + log.debug("cannot scope GPU %d HBM without its partition mode: %s", gpu_id, exc) + return None + if MODE_PARTITION_COUNTS.get(mode) != 1: + log.debug("not reading GPU %d HBM: card is in %s, so a per-device size is not the card's", gpu_id, mode) + return None + + try: + payload = _amd_smi_json(["static", "-g", str(gpu_id), "--vram"], _READ_TIMEOUT_S) + except PartitionError as exc: + log.debug("could not read HBM capacity for GPU %d: %s", gpu_id, exc) + return None + + rows = payload.get("gpu_data") if isinstance(payload, dict) else payload + if not isinstance(rows, list): + return None + for row in rows: + if not isinstance(row, dict): + continue + vram = row.get("vram") + size = vram.get("size") if isinstance(vram, dict) else None + if not isinstance(size, dict): + continue + try: + value = float(size.get("value")) + except (TypeError, ValueError): + continue + scale = _VRAM_UNIT_GIB.get(str(size.get("unit") or "").strip().upper()) + if scale is None or value <= 0: + continue + return value * scale + return None + + def _amd_smi_json(args: Sequence[str], timeout_s: float, privileged: bool = False) -> object: """Run an ``amd-smi`` subcommand with ``--json`` and parse its output. @@ -637,6 +712,7 @@ def partitioned(gpu_id: int, mode: str, restore_to: str | None = None) -> Iterat "parse_modes", "partition_device_predicate", "partitioned", + "read_hbm_gib", "read_partition_mode", "read_partition_modes", "read_partition_profiles", diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py index 77eba14511..f0d229acc4 100644 --- a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -18,6 +18,7 @@ import pytest +from hyperloom.common import gpu_partition as gp from hyperloom.inference_optimizer import cli from hyperloom.orchestrator.actions.executors import _partition_lever as pl from hyperloom.orchestrator.actions.executors import bypass_scriptable as bs @@ -47,10 +48,7 @@ def __call__(self, cmd, **kwargs): self.modes[gpu_id] = mode self.history.append(mode) return subprocess.CompletedProcess(cmd, 0, "", "") - rows = [ - {"gpu_id": gid, "memory": "NPS1", "accelerator_type": m} - for gid, m in sorted(self.modes.items()) - ] + rows = [{"gpu_id": gid, "memory": "NPS1", "accelerator_type": m} for gid, m in sorted(self.modes.items())] return subprocess.CompletedProcess(cmd, 0, json.dumps({"current_partition": rows}), "") @@ -209,9 +207,7 @@ class TestResumeContract: @staticmethod def _args(**over): - ns = argparse.Namespace( - compute_partition_modes=None, streams_per_partition=None, max_latency_ms=None - ) + ns = argparse.Namespace(compute_partition_modes=None, streams_per_partition=None, max_latency_ms=None) for key, value in over.items(): setattr(ns, key, value) return ns @@ -302,9 +298,7 @@ def test_streams_per_partition_parses_to_none_when_not_passed(): assert args.compute_partition_modes is None assert args.max_latency_ms is None - passed = _build_parser().parse_args( - ["optimize", "--model", "/tmp/m", "--streams-per-partition", "2"] - ) + passed = _build_parser().parse_args(["optimize", "--model", "/tmp/m", "--streams-per-partition", "2"]) assert passed.streams_per_partition == 2 @@ -332,9 +326,7 @@ def _state(modes): return SimpleNamespace(compute_partition_modes=list(modes)) def test_the_session_list_becomes_one_variant_per_mode(self): - grid = pl.partition_lever_grid( - {"compute_partition_modes": "spx,dpx,cpx"}, None, framework="custom" - ) + grid = pl.partition_lever_grid({"compute_partition_modes": "spx,dpx,cpx"}, None, framework="custom") assert [v["name"] for v in grid] == [ "partition_spx", "partition_dpx", @@ -369,12 +361,7 @@ def test_a_serving_framework_gets_no_variants(self): is wrong in a way nothing downstream can detect. """ for framework in ("sglang", "vllm", ""): - assert ( - pl.partition_lever_grid( - {"compute_partition_modes": "dpx"}, None, framework=framework - ) - == [] - ) + assert pl.partition_lever_grid({"compute_partition_modes": "dpx"}, None, framework=framework) == [] def test_an_unusable_mode_list_seeds_nothing_rather_than_raising(self, monkeypatch): # Grid assembly is not the place to end a session; the CLI already @@ -416,9 +403,7 @@ def test_a_name_the_grid_already_uses_is_left_alone(self): # An operator or specialist who named the variant said something more # specific than the generated default; the generated one steps aside. pinned = {"name": "partition_dpx", "extra_envs": {"CUSTOM": "1"}} - grid, fresh = explore._prepend_fresh_variants( - [pinned], [{"name": "partition_dpx", "extra_envs": {}}] - ) + grid, fresh = explore._prepend_fresh_variants([pinned], [{"name": "partition_dpx", "extra_envs": {}}]) assert grid == [pinned] assert fresh == [] @@ -429,6 +414,200 @@ def test_nothing_to_add_leaves_the_grid_untouched(self): assert fresh == [] +class TestHbmCapacity: + """Reading how much memory a partition would actually get. + + ``layout_for`` cannot size a partition's memory without the card's total, + and the total is not tabled because boards sharing an ISA do not share a + capacity. These cover the parse and the one state where the reading means + something other than what it says. + """ + + @staticmethod + def _smi(monkeypatch, *, mode: str = "SPX", vram: dict | None = None): + def fake_run(cmd, **kwargs): + if "--vram" in cmd: + payload = {"gpu_data": [{"gpu": 0, "vram": {"size": vram} if vram else {}}]} + else: + payload = {"current_partition": [{"gpu_id": 0, "accelerator_type": mode}]} + return subprocess.CompletedProcess(cmd, 0, json.dumps(payload), "") + + monkeypatch.setattr(gp.subprocess, "run", fake_run) + + def test_mebibytes_labelled_mb_are_read_as_mebibytes(self, monkeypatch): + # An MI355X with 288 GiB reports 294896 and calls it "MB". Taking that + # decimally would understate the card by 7% and prune modes that fit. + self._smi(monkeypatch, vram={"value": 294896, "unit": "MB"}) + assert gp.read_hbm_gib(0) == pytest.approx(288.0, abs=0.1) + + def test_a_partitioned_card_reports_unknown(self, monkeypatch): + """The reading is per device, and under a split mode a device is a partition. + + Nothing in the payload says whether the figure is the card or a slice of + it, and dividing an already-divided number by the partition count again + would understate capacity eightfold -- pruning every mode that fits. + """ + self._smi(monkeypatch, mode="CPX", vram={"value": 36864, "unit": "MB"}) + assert gp.read_hbm_gib(0) is None + + def test_an_unparseable_size_is_unknown_rather_than_zero(self, monkeypatch): + # Zero would read as "nothing fits" and silently empty the mode list. + self._smi(monkeypatch, vram={"value": "N/A", "unit": "MB"}) + assert gp.read_hbm_gib(0) is None + self._smi(monkeypatch, vram={"value": 288, "unit": "furlongs"}) + assert gp.read_hbm_gib(0) is None + + def test_no_amd_smi_is_unknown(self, monkeypatch): + monkeypatch.setattr(gp.subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError())) + assert gp.read_hbm_gib(0) is None + + +class TestFootprintPruning: + """Refusing a mode whose partitions cannot hold what would run on them. + + A partition gets its share of HBM while every stream on it keeps a full copy + of the weights, so the narrow modes exhaust memory first. Measuring that + costs a run per mode and returns an OOM instead of a number. + """ + + def test_a_mode_too_small_for_its_streams_is_dropped(self): + # 288 GiB card: CPX gives 36 GiB/partition, and two 20.7 GiB streams + # need 41.4 GiB. QPX's 72 GiB holds them. + kept, reasons = pl.prune_infeasible_modes( + ("SPX", "DPX", "QPX", "CPX"), + gpu_type="mi355x", + hbm_gib=288.0, + footprint_gib=20.7, + streams=2, + ) + assert kept == ("SPX", "DPX", "QPX") + assert len(reasons) == 1 + # The arithmetic travels with the verdict; "CPX dropped" alone does not + # tell an operator whether to lower streams or pick a bigger mode. + assert "CPX" in reasons[0] and "41.4" in reasons[0] and "36.0" in reasons[0] + + def test_streams_per_partition_decides_it(self): + """One stream fitting says nothing about two. + + This is the case that makes gating on the single-stream figure worse than + not gating: the configuration is declared feasible and dies at the second + worker, after the first has already been measured. + """ + args = dict(gpu_type="mi355x", hbm_gib=288.0, footprint_gib=20.7) + assert pl.prune_infeasible_modes(("CPX",), streams=1, **args)[0] == ("CPX",) + assert pl.prune_infeasible_modes(("CPX",), streams=2, **args)[0] == () + + def test_an_unknown_capacity_drops_nothing(self): + kept, reasons = pl.prune_infeasible_modes( + ("SPX", "CPX"), gpu_type="mi355x", hbm_gib=None, footprint_gib=20.7, streams=2 + ) + assert kept == ("SPX", "CPX") + assert reasons == () + + def test_an_unknown_footprint_drops_nothing(self): + # Wrongly dropping a mode costs the optimization the configuration that + # would have won, and leaves no trace that it was ever a candidate. + kept, _ = pl.prune_infeasible_modes( + ("SPX", "CPX"), gpu_type="mi355x", hbm_gib=288.0, footprint_gib=0.0, streams=2 + ) + assert kept == ("SPX", "CPX") + + def test_an_unsizeable_board_keeps_the_mode(self): + # Refusing here would report an unknown board as a memory verdict; the + # apply path raises with the real reason attached. + kept, reasons = pl.prune_infeasible_modes( + ("CPX",), gpu_type="nvidia-h100", hbm_gib=288.0, footprint_gib=999.0, streams=2 + ) + assert kept == ("CPX",) + assert reasons == () + + def test_the_grid_drops_the_infeasible_mode(self, monkeypatch): + monkeypatch.setenv(pl.STREAMS_PER_PARTITION_ENV, "2") + grid = pl.partition_lever_grid( + { + "compute_partition_modes": "spx,cpx", + "gpu_type": "mi355x", + "peak_gib_per_stream": 20.7, + }, + None, + framework="custom", + hbm_gib=288.0, + ) + assert [v["name"] for v in grid] == ["partition_spx"] + + def test_a_model_that_fits_nowhere_seeds_nothing(self, monkeypatch): + monkeypatch.setenv(pl.STREAMS_PER_PARTITION_ENV, "2") + grid = pl.partition_lever_grid( + { + "compute_partition_modes": "dpx,cpx", + "gpu_type": "mi355x", + "peak_gib_per_stream": 200.0, + }, + None, + framework="custom", + hbm_gib=288.0, + ) + assert grid == [] + + +class TestFootprintSource: + """Which number the partitions get sized against, and how sure it is.""" + + def test_a_measured_peak_wins(self): + state = SimpleNamespace(current_best={"peak_gib_per_stream": 20.7}) + assert pl.per_stream_footprint_gib(None, state) == (20.7, "measured") + + def test_params_override_the_baseline(self): + state = SimpleNamespace(current_best={"peak_gib_per_stream": 20.7}) + assert pl.per_stream_footprint_gib({"peak_gib_per_stream": 30.0}, state) == (30.0, "measured") + + def test_weights_are_the_fallback_and_a_lower_bound(self, tmp_path): + """Read from the checkpoint, so it needs no run to be known. + + A lower bound is the right thing to prune on: the true footprint adds + activations and never subtracts weights, so "does not fit by the weights + alone" is a proof while "fits" is no evidence at all. + """ + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text( + json.dumps({"num_hidden_layers": 4, "hidden_size": 512, "torch_dtype": "float16"}), + encoding="utf-8", + ) + (model / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 150 * 1024**3}}), encoding="utf-8" + ) + gib, source = pl.per_stream_footprint_gib(None, SimpleNamespace(model_path=str(model))) + assert source == "weights" + assert gib == pytest.approx(150.0) + + def test_no_measurement_and_no_readable_model_is_unknown(self, tmp_path): + assert pl.per_stream_footprint_gib(None, SimpleNamespace(model_path=str(tmp_path / "gone"))) == (0.0, "") + assert pl.per_stream_footprint_gib(None, None) == (0.0, "") + + def test_a_large_checkpoint_rules_out_the_narrow_modes(self, tmp_path, monkeypatch): + """The case the weights bound is actually for. + + Two streams of a 60 GiB checkpoint need 120 GiB per partition, which + QPX's 72 GiB cannot hold whatever the activations do -- and that is + knowable from the checkpoint before the first run. + """ + monkeypatch.setenv(pl.STREAMS_PER_PARTITION_ENV, "2") + model = tmp_path / "model" + model.mkdir() + (model / "config.json").write_text(json.dumps({"torch_dtype": "float16"}), encoding="utf-8") + (model / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": 60 * 1024**3}}), encoding="utf-8" + ) + grid = pl.partition_lever_grid( + {"compute_partition_modes": "spx,dpx,qpx,cpx", "gpu_type": "mi355x"}, + SimpleNamespace(model_path=str(model)), + framework="custom", + hbm_gib=288.0, + ) + assert [v["name"] for v in grid] == ["partition_spx", "partition_dpx"] + + def test_launch_refuses_the_lever_on_a_serving_framework(capsys): """Fail at launch, not by quietly seeding an empty axis. diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py index b1e1bfb350..505210821b 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -32,14 +32,16 @@ import logging import os from contextlib import contextmanager, nullcontext -from typing import Any, Iterator +from typing import Any, Iterator, Sequence from hyperloom.common.gpu_partition import ( PartitionError, PartitionLayout, + fits_in_partition, layout_for, parse_modes, partitioned, + read_hbm_gib, ) log = logging.getLogger(__name__) @@ -143,11 +145,130 @@ def resolve_session_modes( return () +def per_stream_footprint_gib( + params: dict[str, Any] | None = None, + shared_state: Any = None, +) -> tuple[float, str]: + """Resolve the per-stream HBM footprint to size partitions against. + + Two sources, tightest first: + + * **Measured.** ``peak_gib_per_stream`` from the baseline run, when the + harness reported it. This is the real footprint -- weights, activations + and workspace -- so it is the only source that can rule out a mode the + weights alone would fit. + * **Weights.** ``weight_bytes`` read byte-exact from the checkpoint's + safetensors index. A *lower bound*, and deliberately used as one: each + stream holds its own copy of the weights, so the true footprint is never + smaller. That makes a "does not fit" verdict from this source a proof and + a "fits" verdict no evidence at all -- which is exactly the asymmetry + pruning needs, since it only ever acts on the former. + + Args: + params: Task params, consulted for an explicit override. + shared_state: Live SharedState, consulted for the baseline measurement + and the model identity. + + Returns: + ``(gib, source)``, or ``(0.0, "")`` when neither source can answer. + ``source`` names the origin for the log line that reports a drop, since + "too big by the weights alone" and "too big as measured" call for + different responses from an operator. + """ + measured = (params or {}).get("peak_gib_per_stream") + if measured is None: + best = getattr(shared_state, "current_best", None) + if isinstance(best, dict): + measured = best.get("peak_gib_per_stream") + try: + if measured is not None and float(measured) > 0: + return float(measured), "measured" + except (TypeError, ValueError): + pass + + model_path = str((params or {}).get("model_path") or getattr(shared_state, "model_path", "") or "").strip() + if not model_path: + return 0.0, "" + # Lazy: the kernel package pulls in the analytical stack, and this module is + # imported by the executors on every run whether the lever is on or not. + from hyperloom.orchestrator.kernel.roofline_ceiling import load_model_meta + + try: + meta = load_model_meta( + model_path, + precision_hint=str(getattr(shared_state, "precision", "") or ""), + ) + except Exception as exc: # noqa: BLE001 — an unreadable checkpoint is "unknown", not fatal + log.debug("cannot size partitions from %s: %s", model_path, exc) + return 0.0, "" + if meta is None or meta.weight_bytes <= 0: + return 0.0, "" + return meta.weight_bytes / float(1024**3), "weights" + + +def prune_infeasible_modes( + modes: Sequence[str], + *, + gpu_type: str | None, + hbm_gib: float | None, + footprint_gib: float, + streams: int, +) -> tuple[tuple[str, ...], tuple[str, ...]]: + """Drop modes whose partitions provably cannot hold the streams placed on them. + + A partition gets its fraction of the card's HBM while every stream on it + keeps a full copy of the weights, so the narrow modes run out of memory + first. Measuring that is expensive and uninformative: the run OOMs, the + variant is scored as a failure, and the operator learns from a stack trace + what arithmetic could have said for free. + + Silent when it cannot compute -- an unknown capacity or an unknown footprint + yields no drops rather than a guess, because the cost of wrongly dropping a + mode is an optimization that never considers the configuration that would + have won. + + Args: + modes: Canonical modes the operator asked for, in their order. + gpu_type: Board name, needed to size partitions. + hbm_gib: Card capacity, or ``None`` when unknown. + footprint_gib: Per-stream footprint; ``0`` when unknown. + streams: Streams that will share each partition. + + Returns: + ``(kept, reasons)``. ``kept`` preserves the input order; ``reasons`` + holds one human-readable line per dropped mode. + """ + if not hbm_gib or footprint_gib <= 0: + return tuple(modes), () + kept: list[str] = [] + reasons: list[str] = [] + for mode in modes: + try: + layout = layout_for(gpu_type, mode, hbm_gib=hbm_gib) + except PartitionError as exc: + # Undescribable here means undescribable at apply time too, but this + # is not the place that gets to refuse it: keep the mode and let the + # apply path raise with its own context. + log.debug("not sizing %s: %s", mode, exc) + kept.append(mode) + continue + if fits_in_partition(footprint_gib, layout, streams): + kept.append(mode) + continue + reasons.append( + f"{mode}: {streams} x {footprint_gib:.1f} GiB = {footprint_gib * streams:.1f} GiB " + f"needed per partition, {layout.gib_per_partition:.1f} GiB available " + f"({layout.partitions} x {layout.cu_per_partition} CU)" + ) + return tuple(kept), tuple(reasons) + + def partition_lever_grid( params: dict[str, Any] | None = None, shared_state: Any = None, *, framework: str | None, + hbm_gib: float | None = None, ) -> list[dict[str, Any]]: """Expand the session's mode list into one explore variant per mode. @@ -157,10 +278,14 @@ def partition_lever_grid( code as every other variant, and a mode that loses is reverted like any other losing knob. - Every listed mode is emitted, including one that may already match the - card. This function stays pure -- no hardware read -- and a mode the - operator named explicitly is worth a measurement under the same harness as - its rivals rather than an inherited baseline number taken on trust. + A mode that may already match the card is still emitted: it is worth a + measurement under the same harness as its rivals rather than an inherited + baseline number taken on trust. + + A mode whose partitions cannot hold the streams destined for them is not + emitted, because that variant has only one outcome and it is an OOM. See + :func:`prune_infeasible_modes`; the arithmetic needs a card capacity and a + per-stream footprint, and drops nothing when either is unknown. Only scriptable frameworks get variants. The mode is applied by :func:`plan_partition_run`, which the scriptable runner calls and the @@ -173,10 +298,14 @@ def partition_lever_grid( params: Task params, consulted for the mode list. shared_state: Live SharedState, consulted for the persisted list. framework: Framework this round runs under. + hbm_gib: Card capacity for the feasibility check. Read from the managed + GPU when omitted; injectable so the decision is testable without a + card and overridable by a caller that already knows. Returns: Variant payload dicts ready for ``_grid_variants_from_payload``, or - ``[]`` when the lever is off or the framework cannot apply it. + ``[]`` when the lever is off, the framework cannot apply it, or no + requested mode has room for the workload. """ from hyperloom.inference_optimizer.framework_registry import is_scriptable @@ -192,6 +321,41 @@ def partition_lever_grid( framework or "", ) return [] + + streams = streams_per_partition() + footprint_gib, footprint_source = per_stream_footprint_gib(params, shared_state) + # Capacity is only read once a footprint exists to compare it against. With + # no footprint the check cannot reach a verdict either way, and this spares + # every mode-less-of-a-workload caller an amd-smi subprocess. + if hbm_gib is None and footprint_gib > 0: + hbm_gib = read_hbm_gib(partition_gpu_id()) + feasible, dropped = prune_infeasible_modes( + modes, + gpu_type=str((params or {}).get("gpu_type") or getattr(shared_state, "gpu_type", "") or ""), + hbm_gib=hbm_gib, + footprint_gib=footprint_gib, + streams=streams, + ) + for reason in dropped: + # Warned rather than debugged: the operator asked for this mode by name, + # and a list that quietly comes back shorter than it went in is the kind + # of thing that gets rediscovered as a bug. + log.warning( + "compute-partition %s dropped at %d stream(s)/partition, footprint from %s", + reason, + streams, + footprint_source, + ) + if not feasible: + log.warning( + "no requested compute-partition mode has room for a %.1f GiB/stream " + "footprint (%s) at %d stream(s)/partition; the lever contributes no " + "variants this round", + footprint_gib, + footprint_source, + streams, + ) + return [] return [ { "name": f"partition_{mode.lower()}", @@ -200,7 +364,7 @@ def partition_lever_grid( "note": f"compute-partition {mode}", "provenance": "partition_lever", } - for mode in modes + for mode in feasible ] @@ -307,7 +471,9 @@ def maybe_hold_partition_mode(mode: str, *, gpu_type: str | None): "maybe_hold_partition_mode", "partition_gpu_id", "partition_lever_grid", + "per_stream_footprint_gib", "plan_partition_run", + "prune_infeasible_modes", "requested_mode", "resolve_session_modes", "runtime_env", diff --git a/src/hyperloom/orchestrator/actions/executors/benchmark_result.py b/src/hyperloom/orchestrator/actions/executors/benchmark_result.py index 800dfa370a..c92248b3b6 100644 --- a/src/hyperloom/orchestrator/actions/executors/benchmark_result.py +++ b/src/hyperloom/orchestrator/actions/executors/benchmark_result.py @@ -722,6 +722,8 @@ def _merge_raw_result( raw.get("p99_e2el_ms"), raw.get("p99_latency_ms"), ) + if measurement.get("peak_gib_per_stream") is None: + measurement["peak_gib_per_stream"] = to_float(raw.get("peak_gib_per_stream")) if measurement.get("raw_result_path") is None: measurement["raw_result_path"] = str(source_path) # AgentX scenario verdict. The KEY'S PRESENCE marks the result as AgentX @@ -765,6 +767,7 @@ def extract_benchmark_measurement( ttft = latency.get("ttft") or {} tpot = latency.get("tpot") or {} e2el = latency.get("e2el") or {} + memory = report.get("memory") or {} measurement: dict[str, Any] = { "reported_success": report.get("success") if report else None, @@ -793,6 +796,17 @@ def extract_benchmark_measurement( "tpot_mean_ms": to_float(tpot.get("mean_ms")), "e2el_mean_ms": to_float(e2el.get("mean_ms")), "e2el_p99_ms": to_float(e2el.get("p99_ms")), + # Peak HBM one stream held, when the harness reports it. "per_stream" is + # in the key because the scope is the whole value of the number: memory + # feasibility is asked per partition, and a total-across-workers figure + # substituted here would overstate the footprint by the worker count. + # Optional -- absent from every harness that does not measure it, and + # read as "unknown" rather than "zero" by anyone who wants it. + "peak_gib_per_stream": first_float( + memory.get("peak_gib_per_stream"), + report.get("peak_gib_per_stream"), + throughput.get("peak_gib_per_stream"), + ), "raw_result_path": None, "nonfatal_warnings": [], } diff --git a/src/hyperloom/orchestrator/loop/writeback.py b/src/hyperloom/orchestrator/loop/writeback.py index 55f70a6c47..4b10e37df3 100644 --- a/src/hyperloom/orchestrator/loop/writeback.py +++ b/src/hyperloom/orchestrator/loop/writeback.py @@ -3065,6 +3065,11 @@ async def _promote_baseline( "ttft_mean_ms": result.get("ttft_mean_ms"), "e2el_mean_ms": result.get("e2el_mean_ms"), "tpot_mean_ms": result.get("tpot_mean_ms"), + # Measured on the reference config, so it is the footprint the + # compute-partition lever sizes its partitions against. Carried + # here rather than re-derived because the baseline is the one run + # that measures the workload before any knob has changed it. + "peak_gib_per_stream": result.get("peak_gib_per_stream"), "workspace": result.get("workspace"), } changed = True From fbc0cca9f7382ce2f2e8fe9f2151b8531266910e Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 02:34:11 +0000 Subject: [PATCH 06/10] feat(report): tell the operator the partition lever exists Every other lever is reachable by the optimizer on its own, so an operator meets it in the results. This one stays off unless a flag names the modes, which means the only people who can use it are the ones who already knew about it. The report is the artefact read after every run, so it is where the lever gets introduced. Advertising is not recommending. Partitioning only ever gives a single stream fewer CUs, so the copy carries the cost as plainly as the 20%: the same measurement took per-request latency from 183 ms to 1211 ms, it needs a privileged amd-smi, and repartitioning evicts every process on the card. An operator who turns this on because a report suggested it and then finds latency quadrupled was misled by that report. Silent where the advice would not apply -- serving frameworks, which refuse the lever at launch, and multi-node sessions, where it manages one card. When the lever did run the section reports what happened instead, including the two states that read as gaps if left blank: no mode kept, which is a measured loss rather than a lever that never ran, and no latency budget, which is how a session ends up on the narrowest mode on offer. Co-authored-by: Cursor (cherry picked from commit 93b06162a3ed1da996172b8080135645a0241696) --- .../inference_optimizer/tests/test_report.py | 67 ++++++++++++ .../orchestrator/actions/executors/report.py | 100 ++++++++++++++++++ 2 files changed, 167 insertions(+) diff --git a/src/hyperloom/inference_optimizer/tests/test_report.py b/src/hyperloom/inference_optimizer/tests/test_report.py index a812116b72..963f928452 100644 --- a/src/hyperloom/inference_optimizer/tests/test_report.py +++ b/src/hyperloom/inference_optimizer/tests/test_report.py @@ -52,6 +52,73 @@ def test_degraded_mode_section_empty(): assert rp._format_degraded_mode_section({}) == [] +# ---- _format_compute_partition_section ---- +def _partition_summary(**over): + base = {"framework": "xdit", "compute_partition_modes": [], "streams_per_partition": 2} + base.update(over) + return base + + +def test_partition_section_advertises_the_lever_when_it_was_not_used(): + """The whole point of the section: an operator cannot ask for a lever nobody mentions.""" + body = "\n".join(rp._format_compute_partition_section(_partition_summary())) + assert "not exercised" in body + # The flag has to be there literally -- naming the capability without the + # incantation leaves the reader exactly as stuck. + assert "--compute-partition-modes spx,dpx,qpx,cpx" in body + assert "--streams-per-partition 2" in body + # And the three things that bite: latency, privilege, other tenants. + assert "--max-latency-ms" in body and "1211 ms" in body + assert "HYPERLOOM_PARTITION_SUDO=1" in body + assert "evicts every process" in body + + +def test_partition_section_is_silent_on_serving_frameworks(): + # The lever is refused at launch there, so advertising it would send an + # operator to a flag that exits 2. + for framework in ("sglang", "vllm", "atom", ""): + assert rp._format_compute_partition_section(_partition_summary(framework=framework)) == [] + + +def test_partition_section_is_silent_on_a_multi_node_session(): + # The lever manages one card; the advice would not apply to the run made. + summary = _partition_summary(platform={"multi_node_session": True}) + assert rp._format_compute_partition_section(summary) == [] + + +def test_partition_section_reports_the_mode_that_won(): + summary = _partition_summary( + compute_partition_modes=["SPX", "DPX"], + latency_budget_ms=400.0, + current_best={"action": "explore", "extra_envs": {"HYPERLOOM_PARTITION_MODE": "DPX"}}, + ) + body = "\n".join(rp._format_compute_partition_section(summary)) + assert "not exercised" not in body + assert "`SPX`, `DPX`" in body + assert "`DPX` — the best configuration runs on a partitioned card" in body + assert "`400.0` ms" in body + + +def test_partition_section_says_so_when_no_mode_was_kept(): + # A measured loss is a result. Reporting it as a blank would read as a + # lever that never ran. + summary = _partition_summary( + compute_partition_modes=["CPX"], + latency_budget_ms=400.0, + current_best={"action": "baseline", "tput": 100.0}, + ) + body = "\n".join(rp._format_compute_partition_section(summary)) + assert "none — no mode beat the unpartitioned card" in body + + +def test_partition_section_flags_a_session_that_partitioned_without_a_budget(): + """The unbounded case, which is how a run ends up on the slowest mode on offer.""" + summary = _partition_summary(compute_partition_modes=["CPX"], latency_budget_ms=0.0) + body = "\n".join(rp._format_compute_partition_section(summary)) + assert "⚠ none set" in body + assert "--max-latency-ms" in body + + def test_format_md_shows_validated_gain_when_timestamp_missing(): md = rp._format_md( { diff --git a/src/hyperloom/orchestrator/actions/executors/report.py b/src/hyperloom/orchestrator/actions/executors/report.py index b27ed96969..4bf2ee7771 100644 --- a/src/hyperloom/orchestrator/actions/executors/report.py +++ b/src/hyperloom/orchestrator/actions/executors/report.py @@ -562,6 +562,13 @@ def _build_summary_dict( # Degraded-mode advisory: benchmark numbers reflect the text path only. "degraded_mode": bool(getattr(state, "degraded_mode", False)), "model_warnings": list(getattr(state, "model_warnings", None) or []), + # The compute-partition lever's session contract. Recorded even when + # the lever was off, because "off" is what the advisory section keys on: + # an empty list is the difference between a lever that lost and a lever + # the operator never knew was there. + "compute_partition_modes": list(getattr(state, "compute_partition_modes", None) or []), + "streams_per_partition": int(getattr(state, "streams_per_partition", 0) or 0), + "latency_budget_ms": float(getattr(state, "latency_budget_ms", 0.0) or 0.0), } if external_baseline: summary["external_baseline"] = external_baseline @@ -717,6 +724,7 @@ def _q(value: Any) -> str: lines.append("") lines.extend(_format_degraded_mode_section(summary)) + lines.extend(_format_compute_partition_section(summary)) roofline_cmp = summary.get("roofline_comparison") if roofline_cmp: @@ -765,6 +773,98 @@ def _format_degraded_mode_section(summary: dict[str, Any]) -> list[str]: return lines +def _format_compute_partition_section(summary: dict[str, Any]) -> list[str]: + """Render the compute-partition lever's state, or advertise it when unused. + + The unused case is the reason this section exists. Every other lever is + reachable by the optimizer on its own, so an operator learns about it from + the results; this one is off unless a flag names the modes, needs privilege + the session may not have, and reconfigures hardware other tenants share. A + lever nobody is told about is a lever nobody uses, and the report is the one + artefact an operator reads after every run. + + Advertising it is not the same as recommending it. Partitioning a card only + ever gives a single stream fewer CUs, so the copy has to carry the cost as + plainly as the benefit -- an operator who enables this because a report + suggested it, and discovers per-request latency quadrupled, was misled by + that report. + + Silent on serving frameworks, which the lever refuses at launch, and on + multi-node sessions, where it manages a single card and the advice would not + apply to the run that was made. + + Args: + summary: The summary payload built by :func:`_build_summary_dict`. + + Returns: + Markdown lines, or ``[]`` when the lever could not have applied. + """ + from hyperloom.inference_optimizer import framework_registry + + from ._partition_lever import PARTITION_MODE_ENV + + if not framework_registry.is_scriptable(summary.get("framework")): + return [] + platform = summary.get("platform") or {} + if platform.get("multi_node_session"): + return [] + + modes = [str(m) for m in (summary.get("compute_partition_modes") or [])] + streams = int(summary.get("streams_per_partition") or 0) or 2 + budget = float(summary.get("latency_budget_ms") or 0.0) + + if not modes: + return [ + "## Compute partitioning (not exercised)", + "", + "This session left the GPU's compute partitioning alone. An AMD card can be split " + "into independent partitions -- `SPX` (whole card), `DPX` (2), `QPX` (4), `CPX` (8) -- " + "and the optimizer can search those modes as a lever, but only when asked.", + "", + "It is worth asking for when the workload runs many concurrent streams and is " + "throughput-bound. On one MI355X with a 1.26B-parameter vision model, `CPX` at two " + "streams per partition carried ~20% more aggregate throughput than the best `SPX` " + "configuration.", + "", + "- **Enable**: `--compute-partition-modes spx,dpx,qpx,cpx` " + f"(each mode becomes one explore variant) and `--streams-per-partition {streams}`.", + "- **Bound the cost first**: pair it with `--max-latency-ms `. Partitioning " + "only ever gives a single stream fewer CUs, so it cannot improve per-request latency " + "and always worsens it -- in that same measurement, from 183 ms to 1211 ms. Without a " + "budget the search is free to pick the narrowest partition on offer, which is the " + "slowest one per request.", + "- **Needs privilege**: the mode belongs to the card, not the process. Set " + "`HYPERLOOM_PARTITION_SUDO=1` with a NOPASSWD sudoers entry for `amd-smi`; an " + "unprivileged session cannot set a mode and will not pretend to.", + "- **Blast radius**: repartitioning evicts every process resident on the card and " + "renumbers its devices. The session restores the mode it found on the way out.", + "", + ] + + lines = ["## Compute partitioning", ""] + lines.append(f"- modes offered : {', '.join(f'`{m}`' for m in modes)}") + lines.append(f"- streams/partition : {streams}") + cb = summary.get("current_best") or {} + envs = cb.get("extra_envs") if isinstance(cb, dict) else None + won = str((envs or {}).get(PARTITION_MODE_ENV) or "").strip().upper() + if won: + lines.append(f"- kept : `{won}` — the best configuration runs on a partitioned card") + else: + # Absence is a result, not a gap: the modes were measured and none beat + # the unpartitioned card, which is the expected outcome below the + # concurrency where partitioning pays. + lines.append("- kept : none — no mode beat the unpartitioned card") + if budget > 0: + lines.append(f"- latency budget : `{budget:.1f}` ms (candidates over it were refused)") + else: + lines.append( + "- latency budget : ⚠ none set — the search was free to trade per-request latency " + "for throughput without limit. Pass `--max-latency-ms ` to bound it." + ) + lines.append("") + return lines + + def _format_completeness_annotations(summary: dict[str, Any]) -> list[str]: """Render honesty annotations for work left unfinished (unvalidated KEEPs, untried hot kernels, KEEPs awaiting integrate). From bc6825e98190df2f79319173ee85cfd692d64d68 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 04:54:54 +0000 Subject: [PATCH 07/10] feat(breakdown): carry the partition lever into capabilities not attempted The session report already tells an operator which capabilities never ran. The partition lever belongs on that list more than anything else on it: every other capability is reachable by the optimizer on its own, so not_attempted there means the search declined to go somewhere. This one is off unless a flag names the modes, so not_attempted means nobody offered it -- and an operator who has not heard of the lever cannot discover it from results that never mention it. The reason carries the flag, so the row is actionable rather than a reproach. Absent, not not_attempted, on a framework that cannot apply it. Listing it there would be false in the way that matters: it reads as a missed opportunity when the launch would in fact have been refused, sending the reader to a flag that exits 2. Modes that were offered and lost read as tried, because a measured loss is evidence and filing it under not_attempted would claim the lever never ran. reason was in the documented contract and rendered nowhere, which left a row able to explain itself with no way to say it; the capability table now prints it. No existing row sets it, so nothing else changes. Co-authored-by: Cursor (cherry picked from commit 406c332ffbc8b468fa18d3d9e7bb198d469921df) --- docs/reference/session-breakdown.md | 6 ++ .../breakdown/collectors/timeline.py | 68 ++++++++++++- .../_renderers/capability_summary.py | 6 ++ .../inference_optimizer/breakdown/schema.py | 7 ++ .../tests/test_reporters_v1_1.py | 97 +++++++++++++++++++ 5 files changed, 183 insertions(+), 1 deletion(-) diff --git a/docs/reference/session-breakdown.md b/docs/reference/session-breakdown.md index 1a31a201fa..4d3466330a 100644 --- a/docs/reference/session-breakdown.md +++ b/docs/reference/session-breakdown.md @@ -467,6 +467,12 @@ For the kernel lanes (`geak`, `forge`) these counts are not interchangeable: The `specialist` row uses `keeps` / `attempts` differently: see `CapabilitySummary` in `schema.py`. +`compute_partition` (the AMD SPX/DPX/QPX/CPX lever) is present only on +frameworks that can apply it, and is the one row whose `not_attempted` +means *nobody offered it* rather than *the search declined it* — the +lever is off unless `--compute-partition-modes` names the modes. Its +`reason` carries the flag needed to turn it on. + --- ## `kernel_lifecycle` — `KernelLifecycle` diff --git a/src/hyperloom/inference_optimizer/breakdown/collectors/timeline.py b/src/hyperloom/inference_optimizer/breakdown/collectors/timeline.py index bd56e301eb..ab67538823 100644 --- a/src/hyperloom/inference_optimizer/breakdown/collectors/timeline.py +++ b/src/hyperloom/inference_optimizer/breakdown/collectors/timeline.py @@ -536,7 +536,7 @@ def _from_invocations(invs: list[dict[str, Any]]) -> dict[str, Any]: # Specialist row derived from ``specialist_rounds``. specialist_row = _specialist_capability_row(state) - return { + summary: dict[str, Any] = { "geak": geak_cap, "forge": forge_cap, # Primary post-merge row; backends/params/validate_stack are compat rows. @@ -547,6 +547,72 @@ def _from_invocations(invs: list[dict[str, Any]]) -> dict[str, Any]: "validate_stack": validate, "specialist": specialist_row, } + # Omitted rather than reported not_attempted where it could not have run; + # see :func:`_compute_partition_capability_row`. + partition_row = _compute_partition_capability_row(state) + if partition_row is not None: + summary["compute_partition"] = partition_row + return summary + + +def _compute_partition_capability_row(state: dict[str, Any]) -> dict[str, Any] | None: + """Derive ``capability_summary.compute_partition`` from the session lever. + + Unlike every other capability here, this one cannot be reached by the + optimizer on its own: it is off unless ``--compute-partition-modes`` names + the modes. So ``not_attempted`` on this row means something the other rows + never mean -- not "the search declined to go there" but "nobody offered it" + -- and it is worth surfacing precisely because an operator who has not heard + of the lever has no way to discover it from results that never mention it. + The ``reason`` carries the flag, so the row is actionable and not merely a + reproach. + + Returns ``None`` on a framework that cannot apply the lever, which omits the + row entirely. Listing it as ``not_attempted`` there would be false in the + way that matters: it would read as a missed opportunity, when the launch + would in fact have been refused. + + Args: + state: Parsed ``state.json``. + + Returns: + The capability row, or ``None`` when the lever could not have applied. + """ + from hyperloom.inference_optimizer.framework_registry import is_scriptable + from hyperloom.orchestrator.actions.executors._partition_lever import PARTITION_MODE_ENV + + if not is_scriptable(state.get("framework")): + return None + modes = [str(m).strip().upper() for m in (state.get("compute_partition_modes") or []) if str(m).strip()] + if not modes: + return { + "status": "not_attempted", + "attempts": 0, + "keeps": 0, + "reason": ( + "never offered — enable with `--compute-partition-modes spx,dpx,qpx,cpx`; " + "needs a privileged amd-smi, and pair it with `--max-latency-ms` because " + "partitioning always costs single-stream latency" + ), + } + + current_best = state.get("current_best") or {} + envs = current_best.get("extra_envs") if isinstance(current_best, dict) else None + kept_mode = str((envs or {}).get(PARTITION_MODE_ENV) or "").strip().upper() + return { + # A mode that lost is evidence, not absence: the modes were measured and + # the unpartitioned card won, which is the expected result below the + # concurrency where partitioning pays. + "status": "kept" if kept_mode else "tried", + "attempts": len(modes), + "keeps": 1 if kept_mode else 0, + "tested": len(modes), + "reason": ( + f"{kept_mode} in the best configuration" + if kept_mode + else f"offered {', '.join(modes)}; none beat the unpartitioned card" + ), + } def _specialist_capability_row(state: dict[str, Any]) -> dict[str, Any]: diff --git a/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/capability_summary.py b/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/capability_summary.py index 79323402a2..4c013a7249 100644 --- a/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/capability_summary.py +++ b/src/hyperloom/inference_optimizer/breakdown/reporters/_renderers/capability_summary.py @@ -15,6 +15,7 @@ "backends", "params", "sweep", + "compute_partition", "geak", "validate_stack", ) @@ -80,6 +81,11 @@ def render(breakdown: dict[str, Any]) -> RenderedSection: extras.append(f"keep_unstable={v['keep_unstable_count']}") if "winners_history" in v and v["winners_history"]: extras.append(f"history={v['winners_history']}") + # Last, and unlabelled: it is prose rather than a metric. Part of the + # documented contract but rendered nowhere until now, which left a row + # able to explain itself with no way to say it. + if v.get("reason"): + extras.append(str(v["reason"])) extras_str = " · ".join(extras) if extras else "" rows.append([name, status, attempts, keeps, extras_str]) diff --git a/src/hyperloom/inference_optimizer/breakdown/schema.py b/src/hyperloom/inference_optimizer/breakdown/schema.py index 1f966b5e94..9299a61983 100644 --- a/src/hyperloom/inference_optimizer/breakdown/schema.py +++ b/src/hyperloom/inference_optimizer/breakdown/schema.py @@ -441,6 +441,12 @@ class CapabilitySummary(TypedDict, total=False): specialist (CapabilityEntry): Specialist sub-agent capability; ``tested`` = total proposals across rounds, ``keeps`` = proposals kept, ``attempts`` = number of dispatch rounds. + compute_partition (CapabilityEntry): AMD compute-partition lever + (SPX/DPX/QPX/CPX); ``tested`` = modes offered, ``keeps`` = 1 when a + mode is in the final configuration. **Absent** on frameworks that + cannot apply the lever, so a present ``not_attempted`` row means an + operator could have used it and did not -- unlike the other rows, + this one is unreachable unless a flag names the modes. """ geak: CapabilityEntry @@ -452,6 +458,7 @@ class CapabilitySummary(TypedDict, total=False): sweep: CapabilityEntry validate_stack: CapabilityEntry specialist: CapabilityEntry + compute_partition: CapabilityEntry # Kernel backend invocations diff --git a/src/hyperloom/inference_optimizer/tests/test_reporters_v1_1.py b/src/hyperloom/inference_optimizer/tests/test_reporters_v1_1.py index de49ce3b5b..283291aa50 100644 --- a/src/hyperloom/inference_optimizer/tests/test_reporters_v1_1.py +++ b/src/hyperloom/inference_optimizer/tests/test_reporters_v1_1.py @@ -7,7 +7,11 @@ from pathlib import Path +from hyperloom.inference_optimizer.breakdown.collectors import collect_capability_summary from hyperloom.inference_optimizer.breakdown.reporters import render_session_report +from hyperloom.inference_optimizer.breakdown.reporters._renderers.capability_summary import ( + render as render_capability_summary, +) from hyperloom.inference_optimizer.breakdown.reporters._renderers.decision_journal import render as render_dj from hyperloom.inference_optimizer.breakdown.reporters._renderers.invocations import render_forge, render_geak from hyperloom.inference_optimizer.breakdown.reporters._renderers.kernel_profiling import render as render_kp @@ -250,3 +254,96 @@ def test_compose_includes_v1_1_sections_in_report() -> None: assert "### Kernel Profiling" in md assert "single_shot" in md assert "magpie_torch_profiler" in md + + +# ---- capability_summary.compute_partition ---- +def test_an_unused_partition_lever_is_reported_as_never_offered() -> None: + """The row exists to be discoverable: the operator could have used it and did not.""" + cap = collect_capability_summary({"framework": "xdit", "compute_partition_modes": []}, [], []) + row = cap["compute_partition"] + assert row["status"] == "not_attempted" + assert row["attempts"] == 0 and row["keeps"] == 0 + # No tested=0 next to attempts=0; the row should not pad itself with zeroes. + assert "tested" not in row + # The flag has to travel with the verdict, or the row only says "no". + assert "--compute-partition-modes spx,dpx,qpx,cpx" in row["reason"] + assert "--max-latency-ms" in row["reason"] + + +def test_a_framework_that_cannot_partition_gets_no_row_at_all() -> None: + """Absent, not not_attempted. + + ``not_attempted`` reads as a missed opportunity, and it flows into + "Capabilities not attempted" in the report. On a serving framework the + launch would have been refused, so that would send an operator to a flag + that exits 2. + """ + for framework in ("sglang", "vllm", "atom", ""): + cap = collect_capability_summary({"framework": framework, "compute_partition_modes": ["DPX"]}, [], []) + assert "compute_partition" not in cap + # The empty-state callers every other collector test uses must stay unaffected. + assert "compute_partition" not in collect_capability_summary({}, [], []) + + +def test_a_kept_mode_is_credited_to_the_lever() -> None: + cap = collect_capability_summary( + { + "framework": "custom", + "compute_partition_modes": ["SPX", "DPX"], + "current_best": {"action": "explore", "extra_envs": {"HYPERLOOM_PARTITION_MODE": "DPX"}}, + }, + [], + [], + ) + row = cap["compute_partition"] + assert row["status"] == "kept" + assert (row["attempts"], row["keeps"], row["tested"]) == (2, 1, 2) + assert "DPX" in row["reason"] + + +def test_modes_that_all_lost_read_as_tried_not_untried() -> None: + # A measured loss is evidence. Reporting it as not_attempted would put the + # lever in "Capabilities not attempted" after it had actually run. + cap = collect_capability_summary( + { + "framework": "custom", + "compute_partition_modes": ["CPX"], + "current_best": {"action": "baseline", "tput": 100.0}, + }, + [], + [], + ) + row = cap["compute_partition"] + assert row["status"] == "tried" + assert row["keeps"] == 0 + assert "none beat the unpartitioned card" in row["reason"] + + +def test_the_capability_table_shows_the_reason() -> None: + """``reason`` is in the documented contract but was rendered nowhere.""" + sec = render_capability_summary( + {"capability_summary": {"compute_partition": {"status": "not_attempted", "reason": "never offered — enable X"}}} + ) + assert "never offered — enable X" in sec.markdown_block + + +def test_an_unused_lever_lands_in_capabilities_not_attempted() -> None: + """End to end: the section group an operator actually reads.""" + bd = _base_breakdown( + session={"session_id": "cp", "session_dir": "/tmp/s"}, + workload={"model_name": "m", "framework_name": "xdit"}, + capability_summary={ + "compute_partition": { + "status": "not_attempted", + "attempts": 0, + "keeps": 0, + "reason": "never offered — enable with `--compute-partition-modes spx,dpx,qpx,cpx`", + } + }, + attribution={"method": "missing"}, + source_files={}, + ) + md = render_session_report(bd).markdown + assert "**Capabilities not attempted**: `compute_partition`" in md + assert "Capabilities never invoked: `compute_partition`" in md + assert "--compute-partition-modes spx,dpx,qpx,cpx" in md From 9f06ad1118f2ad5be3f3c28290a0e21dc9e948b6 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 17:16:58 +0000 Subject: [PATCH 08/10] fix(partition): close the three gaps a review of this branch found Each is a way the lever could mislabel a measurement or mutate something it was never scoped to, reached from a direction that had no check. * A failed restore made every later baseline a silent lie. ``partitioned`` logs a restore failure rather than raising, which is right -- it must not mask the exception that caused the exit -- but the session then carries on with the card left split, and the runs requesting *no* mode are precisely the ones with nothing to notice. They would measure a split card and be recorded as the unpartitioned baseline: the mislabelling this module refuses everywhere else, arrived at from the one side that never looked. ``plan_partition_run`` now reads the mode before a mode-less run and refuses a split card. Silent while the lever is off, because then nothing here has touched the hardware and a split card is the operator's own arrangement; silent too on an unreadable mode, since "cannot tell" is not evidence of a problem. * Launch validation checked card 0 while the session mutated another. ``supported_modes``, ``unsupported_modes`` and ``partition_count_conflicts`` were all asked about GPU 0, but the apply path manages ``partition_gpu_id()``. On a heterogeneous node that validated a different board than the one that gets repartitioned -- defeating the whole reason this check happens at launch rather than at the privileged mode change. * Multi-node was silent in the report but ungated in the lever. The report skips ``multi_node_session``, on the grounds that the lever manages one card; the launch path refused serving frameworks for the same class of reason but never refused a cluster. So a multi-node scriptable session could repartition one node's GPU, keep the mode, and have nothing in the report mention that a privileged change happened at all. Refused at launch now, beside the serving-framework check. Also hardens the test file's own lever-env fixture. ``_export_partition_lever`` writes ``os.environ`` directly, which ``monkeypatch`` cannot undo on its behalf, so the first test to complete an export leaked a mode list and a latency budget into every test that ran after it -- which reaches the KEEP gate, and turned 15 unrelated explore cases red. The fixture now saves and restores the five variables outright. Co-authored-by: Cursor --- .../inference_optimizer/cli/__init__.py | 34 +++++- .../tests/test_partition_lever.py | 109 ++++++++++++++++-- .../actions/executors/_partition_lever.py | 46 +++++++- 3 files changed, 175 insertions(+), 14 deletions(-) diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 2c5b8398ad..b15d0a7dae 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1643,6 +1643,7 @@ def _export_partition_lever( streams_per_partition: int, max_latency_ms: float | None, framework: str | None = None, + nodes: int = 1, ) -> tuple[str, ...]: """Validate and project the compute-partition lever into env. @@ -1661,6 +1662,9 @@ def _export_partition_lever( framework: The session's framework. Checked because only the scriptable runner applies a mode; unset skips the check for callers that run before the framework is resolved. + nodes: The resolved node count. The lever manages one card, so a + multi-node session is refused here for the same reason a serving + framework is. Returns: The canonical modes, empty when the lever is off. @@ -1673,6 +1677,7 @@ def _export_partition_lever( supported_modes, unsupported_modes, ) + from hyperloom.orchestrator.actions.executors._partition_lever import partition_gpu_id budget = float(max_latency_ms or 0.0) if budget > 0: @@ -1712,6 +1717,19 @@ def _export_partition_lever( ) sys.exit(2) + # The lever repartitions a single card, and nothing coordinates that across + # a cluster. Left to run, it would mutate one node's GPU and record the + # result as a property of the whole topology -- and the report stays silent + # on multi-node sessions, so the privileged change would go unmentioned too. + if nodes >= 2: + print( + f"ERROR: --compute-partition-modes manages one card and does not " + f"coordinate across nodes; this session has --nodes {nodes}. Partition " + f"the card in a single-node session, or drop the flag.", + file=sys.stderr, + ) + sys.exit(2) + # Scope the request to what this card says it can do. The name being one of # the four known modes does not mean this board offers it, and the # alternative to checking here is discovering it at the apply site -- a @@ -1723,7 +1741,11 @@ def _export_partition_lever( # unanswerable query warns and proceeds rather than blocking a session that # may be perfectly able to run -- and says so, because "not validated" and # "validated as fine" must not look alike in a log. - available = supported_modes(0) + # The card the session will actually mutate, not card 0: validating a + # different board than the one the apply path touches would defeat the + # reason this check happens at launch at all. + gpu_id = partition_gpu_id() + available = supported_modes(gpu_id) if not available: print( "WARN: could not read this card's supported partition profiles, so " @@ -1733,10 +1755,10 @@ def _export_partition_lever( file=sys.stderr, ) else: - rejected = unsupported_modes(modes) + rejected = unsupported_modes(modes, gpu_id) if rejected: print( - f"ERROR: this card does not support {','.join(rejected)}. " + f"ERROR: GPU {gpu_id} does not support {','.join(rejected)}. " f"It reports: {','.join(available)}.", file=sys.stderr, ) @@ -1745,7 +1767,7 @@ def _export_partition_lever( # and that number drives the CU arithmetic device selection matches on. # A disagreement means the sizing is wrong, so it stops the session here # rather than surfacing as a benchmark that finds no device. - conflicts = partition_count_conflicts(0) + conflicts = partition_count_conflicts(gpu_id) if conflicts: print( "ERROR: this card's partition ladder disagrees with the built-in " @@ -1903,6 +1925,7 @@ async def _run_optimize(args: argparse.Namespace) -> int: streams_per_partition=int(getattr(args, "streams_per_partition", 2) or 2), max_latency_ms=getattr(args, "max_latency_ms", None), framework=str(getattr(args, "framework", "") or "").strip().lower(), + nodes=nodes_resolved, ) # Project resolved workload knobs into env for the fresh-launch path only. # A resume must NOT export here: ``args.tp``/etc. are still unresolved @@ -2131,6 +2154,9 @@ async def _run_optimize(args: argparse.Namespace) -> int: streams_per_partition=int(getattr(args, "streams_per_partition", None) or 2), max_latency_ms=getattr(args, "max_latency_ms", None), framework=str(getattr(args, "framework", "") or "").strip().lower(), + # A resume must re-pass --nodes, so the persisted count is the one + # that says whether this session was ever multi-node. + nodes=max(int(getattr(args, "nodes", 1) or 1), int(getattr(state, "nodes", 1) or 1)), ) _persist_partition_lever(state) if state.compute_partition_modes: diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py index f0d229acc4..d361042025 100644 --- a/src/hyperloom/inference_optimizer/tests/test_partition_lever.py +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -12,6 +12,7 @@ import argparse import json +import os import subprocess from pathlib import Path from types import SimpleNamespace @@ -24,6 +25,7 @@ from hyperloom.orchestrator.actions.executors import bypass_scriptable as bs from hyperloom.orchestrator.actions.executors import explore from hyperloom.orchestrator.actions.executors._latency_budget import ( + LATENCY_BUDGET_ENV, REASON_OVER_BUDGET, REASON_UNMEASURED, latency_keep_block, @@ -87,16 +89,33 @@ def _run(tmp_path: Path, monkeypatch, envs: dict | None = None): return rc, error, (workspace / "ran.marker").is_file(), seen +#: Every variable the lever reads or publishes. Restored around each case +#: because ``_export_partition_lever`` writes ``os.environ`` directly, which +#: ``monkeypatch`` cannot undo on its behalf -- and a budget or a mode list left +#: behind here silently changes the KEEP gate for every test that runs after. +_LEVER_ENV = ( + pl.PARTITION_MODE_ENV, + pl.PARTITION_MODES_ENV, + pl.STREAMS_PER_PARTITION_ENV, + pl.PARTITION_GPU_ENV, + LATENCY_BUDGET_ENV, +) + + @pytest.fixture(autouse=True) -def _clean_lever_env(monkeypatch): - """No ambient lever: these cases each state their own.""" - for name in ( - pl.PARTITION_MODE_ENV, - pl.PARTITION_MODES_ENV, - pl.STREAMS_PER_PARTITION_ENV, - pl.PARTITION_GPU_ENV, - ): - monkeypatch.delenv(name, raising=False) +def _clean_lever_env(): + """No ambient lever: these cases each state their own, and leak none.""" + saved = {name: os.environ.get(name) for name in _LEVER_ENV} + for name in _LEVER_ENV: + os.environ.pop(name, None) + try: + yield + finally: + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value def test_lever_off_leaves_the_run_untouched(tmp_path, monkeypatch): @@ -624,3 +643,75 @@ def test_launch_refuses_the_lever_on_a_serving_framework(capsys): ) assert excinfo.value.code == 2 assert "scriptable framework" in capsys.readouterr().err + + +def test_launch_refuses_the_lever_on_a_multi_node_session(capsys): + """One card, no cluster coordination -- and the report is silent there. + + Left to run it would mutate one node's GPU and file the result as a property + of the whole topology, with the report saying nothing about the privileged + change because it skips multi-node sessions. + """ + with pytest.raises(SystemExit) as excinfo: + cli._export_partition_lever( + modes_raw="dpx", + streams_per_partition=2, + max_latency_ms=400.0, + framework="custom", + nodes=2, + ) + assert excinfo.value.code == 2 + assert "one card" in capsys.readouterr().err + + +def test_launch_validates_the_card_the_session_will_actually_mutate(monkeypatch): + """Not card 0, when the session manages another. + + Validating a different board than the apply path touches would defeat the + reason this check happens at launch instead of at the mode change. + """ + monkeypatch.setenv(pl.PARTITION_GPU_ENV, "3") + asked: list[int] = [] + monkeypatch.setattr(gp, "supported_modes", lambda gpu_id: (asked.append(gpu_id), ("SPX", "DPX"))[1]) + monkeypatch.setattr(gp, "unsupported_modes", lambda modes, gpu_id=0: (asked.append(gpu_id), ())[1]) + monkeypatch.setattr(gp, "partition_count_conflicts", lambda gpu_id=0: (asked.append(gpu_id), ())[1]) + + cli._export_partition_lever( + modes_raw="dpx", + streams_per_partition=2, + max_latency_ms=400.0, + framework="custom", + ) + assert asked == [3, 3, 3] + + +class TestUnpartitionedRunOnASplitCard: + """A failed restore must not turn later baselines into silent lies. + + ``partitioned`` logs a restore failure rather than raising, so it cannot mask + the exception that caused the exit. The cost is that the card can be left + split, and the runs requesting *no* mode are the ones with nothing to notice. + """ + + def test_a_split_card_refuses_a_run_that_asked_for_no_mode(self, tmp_path, monkeypatch): + monkeypatch.setattr(subprocess, "run", _FakeSmi("QPX")) + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "spx,dpx") + rc, error, ran, _ = _run(tmp_path, monkeypatch) + # Refused before the benchmark started: the number would have been filed + # as the unpartitioned baseline. + assert (rc, ran) == (2, False) + assert "QPX" in str(error) and "no partition mode" in str(error) + + def test_an_unpartitioned_card_runs_as_before(self, tmp_path, monkeypatch): + monkeypatch.setattr(subprocess, "run", _FakeSmi("SPX")) + monkeypatch.setenv(pl.PARTITION_MODES_ENV, "spx,dpx") + rc, error, ran, seen = _run(tmp_path, monkeypatch) + assert (rc, error, ran) == (0, None, True) + assert seen == {} + + def test_the_check_is_silent_when_the_lever_is_off(self, tmp_path, monkeypatch): + # A split card with no lever is the operator's own arrangement, and + # nothing here has touched the hardware. + monkeypatch.setattr(subprocess, "run", _FakeSmi("CPX")) + rc, error, ran, _ = _run(tmp_path, monkeypatch) + assert (rc, error, ran) == (0, None, True) diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py index 505210821b..7257400824 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -35,6 +35,7 @@ from typing import Any, Iterator, Sequence from hyperloom.common.gpu_partition import ( + MODE_PARTITION_COUNTS, PartitionError, PartitionLayout, fits_in_partition, @@ -42,6 +43,7 @@ parse_modes, partitioned, read_hbm_gib, + read_partition_mode, ) log = logging.getLogger(__name__) @@ -400,6 +402,45 @@ def runtime_env(layout: PartitionLayout, streams: int) -> dict[str, str]: } +def _refuse_split_card_for_unpartitioned_run() -> None: + """Refuse a mode-less run on a card that is still split. + + :func:`partitioned` logs a failed restore instead of raising, so it cannot + mask the exception that caused the exit. That is right, but it means a + session can carry on with the card left in a mode nobody asked for -- and + the runs that request *no* mode are the ones with nothing to notice it. They + would measure a split card and be recorded as the unpartitioned baseline, + which is the mislabelling this module refuses everywhere else, arrived at + from the one direction that had no check. + + Only speaks when the session lever is engaged. With the lever off nothing + here has touched the hardware, so a split card is the operator's own + arrangement and none of this module's business. + + Raises: + PartitionError: If the managed card is in a split mode. An unreadable + mode is left alone: "cannot tell" is not evidence of a problem, and + the apply path already refuses what it cannot verify. + """ + if not resolve_session_modes(): + return + gpu_id = partition_gpu_id() + try: + current = read_partition_mode(gpu_id) + except PartitionError as exc: + log.debug("cannot check GPU %d topology before an unpartitioned run: %s", gpu_id, exc) + return + if MODE_PARTITION_COUNTS.get(current) == 1: + return + raise PartitionError( + f"GPU {gpu_id} is in {current}, but this run requested no partition mode. " + f"Measuring it now would file a split card's number as the unpartitioned " + f"baseline. A restore this session failed, or the card was already split " + f"when the session started; return it to a single partition before " + f"continuing." + ) + + def plan_partition_run( envs: dict[str, Any] | None, *, @@ -419,10 +460,13 @@ def plan_partition_run( PartitionError: If a mode was requested but cannot be described -- an unknown mode or an unrecognised board. Raised rather than ignored: the request was explicit, so silently not honouring it would - mislabel the measurement. + mislabel the measurement. Also if no mode was requested but the + managed card is still split; see + :func:`_refuse_split_card_for_unpartitioned_run`. """ mode = requested_mode(envs) if not mode: + _refuse_split_card_for_unpartitioned_run() return "", {} layout = layout_for(gpu_type, mode) streams = streams_per_partition() From 03a65a2de49f30e41ff77890370d59ab9093a60f Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 17:20:13 +0000 Subject: [PATCH 09/10] style(partition): satisfy the repo's ruff format gate ``ruff format --check .`` is a hard gate in lint.yml alongside ``ruff check``, and three files on this branch fail it. All three are line joins the configured line length allows, so nothing about the code changes. Two of them predate the review pass and would have failed CI on the original branch too: the ``layout_for`` odd-CU case in ``test_gpu_partition.py`` and the over-budget reason string in ``_latency_budget.py``. The third is the ``unsupported_modes`` error string this branch just rewrote to name the GPU. Worth recording why they were missed: the branch was verified with ``pytest`` and with ``ruff check`` on the changed files, and neither sees a formatting difference. Only ``ruff format`` does, and only the repo-wide invocation CI actually runs. Co-authored-by: Cursor --- src/hyperloom/common/tests/test_gpu_partition.py | 4 +--- src/hyperloom/inference_optimizer/cli/__init__.py | 3 +-- .../orchestrator/actions/executors/_latency_budget.py | 3 +-- 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/src/hyperloom/common/tests/test_gpu_partition.py b/src/hyperloom/common/tests/test_gpu_partition.py index c064daa2f1..4a1c5feebe 100644 --- a/src/hyperloom/common/tests/test_gpu_partition.py +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -478,9 +478,7 @@ def test_layout_refuses_a_cu_count_that_does_not_divide(monkeypatch): # Flooring would be silent, and then fatal much later and for an apparently # unrelated reason: device selection matches the per-partition CU count # exactly, so a floored value matches no device at all. - monkeypatch.setitem( - gpu_partition.AMD_GPU_DISPATCH_IDENTITIES, "oddboard", ("gfx950", 300, "x") - ) + monkeypatch.setitem(gpu_partition.AMD_GPU_DISPATCH_IDENTITIES, "oddboard", ("gfx950", 300, "x")) with pytest.raises(PartitionError, match="does not divide"): layout_for("oddboard", "CPX") # The same board is fine in a mode its CU count does divide by. diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index b15d0a7dae..f8d2048eec 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1758,8 +1758,7 @@ def _export_partition_lever( rejected = unsupported_modes(modes, gpu_id) if rejected: print( - f"ERROR: GPU {gpu_id} does not support {','.join(rejected)}. " - f"It reports: {','.join(available)}.", + f"ERROR: GPU {gpu_id} does not support {','.join(rejected)}. It reports: {','.join(available)}.", file=sys.stderr, ) sys.exit(2) diff --git a/src/hyperloom/orchestrator/actions/executors/_latency_budget.py b/src/hyperloom/orchestrator/actions/executors/_latency_budget.py index 2d3affee29..56b2e59493 100644 --- a/src/hyperloom/orchestrator/actions/executors/_latency_budget.py +++ b/src/hyperloom/orchestrator/actions/executors/_latency_budget.py @@ -121,8 +121,7 @@ def latency_keep_block( ) if observed > budget: return True, ( - f"{REASON_OVER_BUDGET}: {observed:.0f} ms exceeds the " - f"{budget:.0f} ms budget ({observed / budget:.2f}x)" + f"{REASON_OVER_BUDGET}: {observed:.0f} ms exceeds the {budget:.0f} ms budget ({observed / budget:.2f}x)" ) return False, "" From 1fed743874e3a0ab1ff050c376cf3e4714a9968a Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Tue, 25 Aug 2026 17:30:10 +0000 Subject: [PATCH 10/10] fix(partition): clear the CodeQL findings on this branch's own files Three findings, all on files this branch adds or edits. The earlier ones the bots reported were against ``prompt_builder.py``, ``codex_session.py`` and the TraceLens tools -- files this PR never meant to touch, which came from the stale lineage the branch was rebased off. These are the ones that are actually ours. * **Empty except** in ``per_stream_footprint_gib``. The ``pass`` was deliberate but said so nowhere, which is the same finding the repo already answered with a comment in ``_steal_stale_claim``. A report carrying ``peak_gib_per_stream`` as a non-number is treated as not carrying it, which is what the caller's "prunes nothing when the footprint is unknown" contract already promises. * **Implicit string concatenation in a list**, five times over, in the not-exercised branch of the partition report section. Adjacent literals in a list display are ambiguous with a forgotten comma -- a real class of bug, and here it would have split one markdown bullet into two. Each paragraph is now bound to a name before the list. The rendered section is byte-for-byte identical, checked by diffing the output against the previous commit's. * **Unreachable code**, twice, in the ``partitioned`` restore tests. A false positive with a real cost: CodeQL does not model ``pytest.raises`` as suppressing, so an inline ``raise`` as the block's last statement makes everything after it read as dead -- and what follows is the entire point of those two cases, which is what the context manager did on its way out. A ``_raise`` helper keeps the flow analysable without weakening either assertion. Not changed: the same concatenation pattern at ``explore.py:593`` is main's code, not this branch's, and the one in ``_latency_budget.py`` is already parenthesized, which is why the rule does not fire on it. Co-authored-by: Cursor --- .../common/tests/test_gpu_partition.py | 17 ++++++- .../actions/executors/_partition_lever.py | 4 ++ .../orchestrator/actions/executors/report.py | 44 ++++++++++++++----- 3 files changed, 52 insertions(+), 13 deletions(-) diff --git a/src/hyperloom/common/tests/test_gpu_partition.py b/src/hyperloom/common/tests/test_gpu_partition.py index 4a1c5feebe..51b7cd60db 100644 --- a/src/hyperloom/common/tests/test_gpu_partition.py +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -31,6 +31,19 @@ ) +def _raise(exc: BaseException) -> None: + """Raise ``exc`` from inside a ``with`` body. + + Called rather than raised inline so the assertions *after* an enclosing + ``pytest.raises`` stay reachable to a static analyser. CodeQL does not model + ``pytest.raises`` as suppressing what the block raises, so an inline + ``raise`` as the body's last statement makes every following line read as + dead code -- and those lines are the actual subject of these cases: what the + context manager did on its way out. + """ + raise exc + + class _FakeSmi: """An ``amd-smi`` whose partition state a test can drive. @@ -261,7 +274,7 @@ def test_partitioned_restores_after_a_failure_inside_the_block(smi): with pytest.raises(ZeroDivisionError): with partitioned(0, "QPX"): assert smi.modes[0] == "QPX" - raise ZeroDivisionError + _raise(ZeroDivisionError()) assert smi.modes[0] == "SPX" @@ -279,7 +292,7 @@ def flaky(cmd, **kwargs): monkeypatch.setattr(subprocess, "run", flaky) with pytest.raises(ValueError, match="workload blew up"): with partitioned(0, "DPX"): - raise ValueError("workload blew up") + _raise(ValueError("workload blew up")) assert "left in DPX" in caplog.text diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py index 7257400824..2a03530607 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_lever.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -186,6 +186,10 @@ def per_stream_footprint_gib( if measured is not None and float(measured) > 0: return float(measured), "measured" except (TypeError, ValueError): + # A report that carries the field as something other than a number is + # treated as not carrying it: fall through to the weight-bytes bound + # rather than fail, since the caller's contract is "prunes nothing when + # the footprint is unknown" and a malformed reading is unknown. pass model_path = str((params or {}).get("model_path") or getattr(shared_state, "model_path", "") or "").strip() diff --git a/src/hyperloom/orchestrator/actions/executors/report.py b/src/hyperloom/orchestrator/actions/executors/report.py index 4bf2ee7771..643e5dadca 100644 --- a/src/hyperloom/orchestrator/actions/executors/report.py +++ b/src/hyperloom/orchestrator/actions/executors/report.py @@ -814,30 +814,52 @@ def _format_compute_partition_section(summary: dict[str, Any]) -> list[str]: budget = float(summary.get("latency_budget_ms") or 0.0) if not modes: - return [ - "## Compute partitioning (not exercised)", - "", + # Each paragraph is bound to a name before the list rather than wrapped + # inside it: adjacent string literals in a list display are ambiguous + # with a forgotten comma, which is both a real class of bug and a + # standing CodeQL finding on this file. + what_it_is = ( "This session left the GPU's compute partitioning alone. An AMD card can be split " "into independent partitions -- `SPX` (whole card), `DPX` (2), `QPX` (4), `CPX` (8) -- " - "and the optimizer can search those modes as a lever, but only when asked.", - "", + "and the optimizer can search those modes as a lever, but only when asked." + ) + when_to_ask = ( "It is worth asking for when the workload runs many concurrent streams and is " "throughput-bound. On one MI355X with a 1.26B-parameter vision model, `CPX` at two " "streams per partition carried ~20% more aggregate throughput than the best `SPX` " - "configuration.", - "", + "configuration." + ) + how_to_enable = ( "- **Enable**: `--compute-partition-modes spx,dpx,qpx,cpx` " - f"(each mode becomes one explore variant) and `--streams-per-partition {streams}`.", + f"(each mode becomes one explore variant) and `--streams-per-partition {streams}`." + ) + bound_the_cost = ( "- **Bound the cost first**: pair it with `--max-latency-ms `. Partitioning " "only ever gives a single stream fewer CUs, so it cannot improve per-request latency " "and always worsens it -- in that same measurement, from 183 ms to 1211 ms. Without a " "budget the search is free to pick the narrowest partition on offer, which is the " - "slowest one per request.", + "slowest one per request." + ) + needs_privilege = ( "- **Needs privilege**: the mode belongs to the card, not the process. Set " "`HYPERLOOM_PARTITION_SUDO=1` with a NOPASSWD sudoers entry for `amd-smi`; an " - "unprivileged session cannot set a mode and will not pretend to.", + "unprivileged session cannot set a mode and will not pretend to." + ) + blast_radius = ( "- **Blast radius**: repartitioning evicts every process resident on the card and " - "renumbers its devices. The session restores the mode it found on the way out.", + "renumbers its devices. The session restores the mode it found on the way out." + ) + return [ + "## Compute partitioning (not exercised)", + "", + what_it_is, + "", + when_to_ask, + "", + how_to_enable, + bound_the_cost, + needs_privilege, + blast_radius, "", ]