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/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py new file mode 100644 index 0000000000..1612f4699d --- /dev/null +++ b/src/hyperloom/common/gpu_partition.py @@ -0,0 +1,722 @@ +# 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. +* **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 + +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 +#: 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 +_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})" + + +@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. + + 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] + 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, + 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 + + +#: 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. + + 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, + 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", + "PartitionProfile", + "fits_in_partition", + "layout_for", + "parse_modes", + "partition_device_predicate", + "partitioned", + "read_hbm_gib", + "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 new file mode 100644 index 0000000000..51b7cd60db --- /dev/null +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -0,0 +1,506 @@ +# 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, +) + + +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. + + ``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") + + +#: 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/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/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 775d439b38..f8d2048eec 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -65,6 +65,12 @@ 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._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, ) @@ -1631,6 +1637,204 @@ def _export_operator_launch_shape( os.environ.pop("INFERENCE_OPTIMIZER_EXTRA_ENV", None) +def _export_partition_lever( + *, + modes_raw: str | None, + 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. + + 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. + 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. + """ + from hyperloom.common.gpu_partition import ( + PARTITION_SUDO_ENV, + PartitionError, + parse_modes, + partition_count_conflicts, + 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: + 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) + + # 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) + + # 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 + # 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. + # 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 " + 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, gpu_id) + if rejected: + print( + f"ERROR: GPU {gpu_id} does not support {','.join(rejected)}. 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(gpu_id) + 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 + # 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 + + +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. @@ -1715,6 +1919,13 @@ 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), + 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 # (``None`` -> 1) because the persisted SharedState is loaded later; the @@ -1934,6 +2145,23 @@ 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, 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), + 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: + 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 6aa5d5400e..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(), @@ -355,6 +360,13 @@ 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=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 6de63ab826..d32f874ff4 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -287,6 +287,48 @@ 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. 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 " + "host: the mode is a property of the card, not of the process, so a " + "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.", + ) + opt.add_argument( + "--streams-per-partition", + type=int, + # 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 " + "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 +630,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..d361042025 --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_partition_lever.py @@ -0,0 +1,717 @@ +# 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 argparse +import json +import os +import subprocess +from pathlib import Path +from types import SimpleNamespace + +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 +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, + 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 + + +#: 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(): + """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): + 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 + + +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 + + +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 == [] + + +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. + + 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 + + +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/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/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 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..56b2e59493 --- /dev/null +++ b/src/hyperloom/orchestrator/actions/executors/_latency_budget.py @@ -0,0 +1,142 @@ +# 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 {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..2a03530607 --- /dev/null +++ b/src/hyperloom/orchestrator/actions/executors/_partition_lever.py @@ -0,0 +1,529 @@ +# 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, Sequence + +from hyperloom.common.gpu_partition import ( + MODE_PARTITION_COUNTS, + PartitionError, + PartitionLayout, + fits_in_partition, + layout_for, + parse_modes, + partitioned, + read_hbm_gib, + read_partition_mode, +) + +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 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 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 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): + # 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() + 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. + + 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. + + 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 + 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. + 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, the framework cannot apply it, or no + requested mode has room for the workload. + """ + 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 [] + + 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()}", + "extra_args": "", + "extra_envs": {PARTITION_MODE_ENV: mode}, + "note": f"compute-partition {mode}", + "provenance": "partition_lever", + } + for mode in feasible + ] + + +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 _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, + *, + 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. 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() + 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", + "partition_lever_grid", + "per_stream_footprint_gib", + "plan_partition_run", + "prune_infeasible_modes", + "requested_mode", + "resolve_session_modes", + "runtime_env", + "streams_per_partition", +] 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/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..d52f97a9c4 100644 --- a/src/hyperloom/orchestrator/actions/executors/explore.py +++ b/src/hyperloom/orchestrator/actions/executors/explore.py @@ -80,6 +80,12 @@ 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 ._partition_lever import partition_lever_grid from ._ray_serving import maybe_serving_lease # DEFAULT_STACK_STABLE_PCT: post-KEEP confirmation floor; override via @@ -242,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. @@ -799,6 +835,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", @@ -890,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. @@ -1514,6 +1573,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 +1587,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/actions/executors/report.py b/src/hyperloom/orchestrator/actions/executors/report.py index b27ed96969..643e5dadca 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,120 @@ 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: + # 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." + ) + 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." + ) + how_to_enable = ( + "- **Enable**: `--compute-partition-modes spx,dpx,qpx,cpx` " + 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." + ) + 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." + ) + 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." + ) + return [ + "## Compute partitioning (not exercised)", + "", + what_it_is, + "", + when_to_ask, + "", + how_to_enable, + bound_the_cost, + needs_privilege, + blast_radius, + "", + ] + + 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). 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 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