From 0b50c79d5a748ec221a012c1cda0cc7bdbbf3924 Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Wed, 26 Aug 2026 18:33:48 +0000 Subject: [PATCH 1/2] feat(platform): record, check and publish the compute-partition shape An MI300-series card can be split into independent partitions (SPX, DPX, QPX, CPX), and splitting one trades per-request latency for aggregate throughput. Nothing recorded which shape a number came from, so the same configuration measured on the same card in SPX and in CPX was indistinguishable in the history: two different experiments filed under one name. The observed mode now joins NPS in the platform fingerprint, the session report names it on partitioned runs, and the shape is published for the benchmark entrypoint that places work across partitions. The optimizer does not change the mode. Setting it is privileged, evicts every process holding a context on the card, and renumbers its devices -- not something an optimization loop should do between benchmark rounds, and not something to hand agent-authored code. The card must be in its mode before optimize starts: the shape is checked and recorded at launch, so a mode applied later is too late to be either. Every probe added here is an unprivileged read, and a host without amd-smi behaves exactly as before. That leaves two things worth doing at the boundary, both at launch. --compute-partition-mode asserts the mode the card is already in and refuses the session when it is in another, or when the card cannot be read at all: the flag exists to catch an external set that did not take, so an unverifiable assertion is not a satisfied one. And the per-stream footprint is checked against one partition's memory, sized from the checkpoint's weight bytes -- a lower bound, since each stream holds its own copy of the weights, which is why a "does not fit" verdict from it is a proof and a "fits" verdict is no evidence. The arithmetic costs milliseconds and replaces an out-of-memory crash three hours in. When the checkpoint cannot be sized the session runs and says so. The footprint refusal applies only where streams will actually share a partition. Without a fan-out nothing places a second stream, and nothing pins the benchmark to a partition at all -- whole cards enumerate before partitions, so on a node with one card of eight split, device 0 is a whole card. Refusing a serving session that merely started on a card someone else left split would be arithmetic about a shape it was never going to run in, so the mode is recorded there and nothing is refused. For the same reason the published env is split by reader: mode, count and CU describe the card and are always published, since the platform fingerprint reads them back on the crash path, while streams and total streams are directions to a benchmark that fans out and are published only when one will. CU per partition is read from the device rather than divided out of a board table, because partition devices are selected by matching that count exactly: an index list computed at launch would be wrong in the one case that matters and wrong invisibly. The table remains a fallback and the recorded shape says which of the two it came from, so a derived count is never presented as a measurement -- and an unknown provenance is reported as unknown rather than as the table. Multi-node sessions record no shape. The card this process can read is not the card the benchmark runs on, and a shape recorded from the wrong node is the mislabelling this exists to prevent. Co-authored-by: Cursor --- CHANGELOG.md | 42 ++ docs/reference/environment-variables.md | 76 +++ .../agents/robustness/role/envelope.py | 1 + src/hyperloom/common/gpu_partition.py | 566 ++++++++++++++++ src/hyperloom/common/platform_probe.py | 10 + .../common/tests/test_gpu_partition.py | 302 +++++++++ .../inference_optimizer/cli/__init__.py | 245 +++++++ .../inference_optimizer/cli/bootstrap.py | 9 + .../inference_optimizer/cli/parser.py | 36 + .../tests/test_cli_bootstrap.py | 42 ++ .../tests/test_partition_shape.py | 615 ++++++++++++++++++ .../actions/executors/_partition_shape.py | 398 ++++++++++++ .../orchestrator/actions/executors/report.py | 69 ++ src/hyperloom/orchestrator/policy/gate.py | 6 + .../orchestrator/state/shared_state.py | 6 + 15 files changed, 2423 insertions(+) create mode 100644 src/hyperloom/common/gpu_partition.py create mode 100644 src/hyperloom/common/tests/test_gpu_partition.py create mode 100644 src/hyperloom/inference_optimizer/tests/test_partition_shape.py create mode 100644 src/hyperloom/orchestrator/actions/executors/_partition_shape.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8523365861..7852bb936b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,48 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Added + +- **The card's compute-partition shape is now recorded, checked, and published.** + An MI300-series card can be split into independent partitions (`SPX`, `DPX`, + `QPX`, `CPX`), and splitting one trades per-request latency for aggregate + throughput. Until now nothing in a session recorded which shape a + number came from, so two runs of the same configuration on the same card in + `SPX` and in `CPX` were indistinguishable in the history — different + experiments filed under one name.
+ The observed mode now goes into the platform fingerprint alongside NPS, the + session report names it on partitioned runs, and the shape is published to + the environment for the benchmark entrypoint to fan work out across + partitions. That entrypoint lives outside this repository, so until it reads + them a session on a split card measures one partition rather than the total; + the recorded shape is still what stops a `CPX` number being filed as though + it were `SPX`. The published variables are set only for the scriptable + frameworks whose benchmarks can fan out; a serving session records the shape + but is handed no fan-out contract, and its report says the figure cannot be + read as an aggregate.
+ Two optional flags configure it. `--compute-partition-mode` **asserts** the + mode the card is already in and refuses the session if it is in another one, + or if the card cannot be read — the flag exists to catch an external set that + did not take, so an unverifiable assertion is treated as a failed one. + `--streams-per-partition` (default `2`) is how many concurrent streams go on + each partition.
+ **The optimizer does not change the mode.** Setting it is privileged and + disrupts every process holding a GPU context, which is not something an + optimization loop should do between benchmark rounds. The card must be in its + mode before `optimize` starts: the shape is checked and recorded at launch, so + a mode applied later — by the benchmark entrypoint, for instance — is too late + to be either. Nothing added here needs privilege: every probe is an + unprivileged read, and a host without `amd-smi` behaves exactly as before.
+ **Operator note**: launch now refuses a session whose streams provably will + not fit one partition, sized from the checkpoint's weight bytes as a lower + bound. The arithmetic costs milliseconds and the failure it replaces is an + out-of-memory crash hours in. When the checkpoint cannot be sized the session + runs and says so. The refusal applies where streams will actually share a + partition — a scriptable framework, or an operator who named the flags — and + not to a serving session that merely happens to start on a card someone else + left split. Multi-node sessions record no shape, since the readable card is + not the benchmark's. + ### Changed - **BREAKING: the EXPLORE phase is merged into FRAMEWORK_AGENT.** The chain is diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index 20494f9448..823014ee62 100644 --- a/docs/reference/environment-variables.md +++ b/docs/reference/environment-variables.md @@ -425,6 +425,82 @@ multi-node runs or when the Ray backend is disabled. --- +## Compute partitioning (AMD) + +An MI300-series card can be split into independent partitions (`SPX`, `DPX`, +`QPX`, `CPX`). Splitting it trades per-request latency for aggregate throughput, +so a partitioned measurement is not comparable with a whole-card one. + +**The optimizer does not set the mode.** Changing it is privileged, disruptive to +every process holding a GPU context, and not something an optimization loop +should be doing between benchmark rounds. The card is put in its mode before +launch — by the operator or the provisioning platform — and `optimize` only +observes what it is in, refuses a session that cannot work in that shape, and +hands the shape to the benchmark entrypoint that places work across partitions. + +Two CLI flags configure this, both optional: + +- `--compute-partition-mode {SPX,DPX,QPX,CPX}` **asserts** the mode the card is + already in. It is a check, not a request: if the card is in a different mode + the session is refused rather than silently measuring the wrong topology. If + the card cannot be read at all, a declared mode is also a refusal — the flag + exists to catch an external set that did not take, and an unverifiable + assertion is not a satisfied one. +- `--streams-per-partition N` (default `2`) is how many concurrent streams the + benchmark places on each partition. One stream leaves each partition idle + through the fixed per-pass cost; beyond two, on the workloads measured so far, + only queueing is added. + +At launch the per-stream HBM footprint is checked against one partition's +memory. A workload that provably will not fit is refused in milliseconds instead +of failing out of memory hours in. The footprint is the checkpoint's weight bytes +— a lower bound, since each stream holds its own copy of the weights, which is +why a "does not fit" verdict from it is trustworthy and a "fits" verdict proves +nothing. When the checkpoint cannot be sized the session runs with a warning. + +The check only applies where streams will actually share a partition: with a +serving framework and no partition flags, the shape is recorded and nothing is +refused, because nothing in that session places work per partition and, since +whole cards enumerate before partitions, its benchmark may not even land on one. + +Multi-node sessions (`--nodes >= 2`) record no shape at all. The card this +process can read is not the card the benchmark runs on, and a shape recorded from +the wrong node is the exact mislabelling this feature exists to prevent. A +declared mode there is a usage error rather than a silently unchecked assertion. + +| Variable | Default | Description | +|----------|---------|-------------| +| `HYPERLOOM_PARTITION_GPU` | `0` | Which GPU's partition state describes this session. | + +### Runtime hand-off + +Published once at launch by the CLI and read by the benchmark entrypoint. Do not +set these by hand: they are overwritten at every launch, and clearing them first +is what stops a second session in the same shell from inheriting a shape that +was not asked for. + +The first three describe the card and are published for every session on a +readable card, because `platform_fingerprint()` reads them back from here — it +runs on the crash path, where spawning `amd-smi` is not acceptable. The last two +are instructions to a benchmark that places work on each partition, so they are +published only when one will. + +| Variable | Published | Description | +|----------|-----------|-------------| +| `HYPERLOOM_PARTITION_MODE` | Always | The observed mode. Also recorded in the platform fingerprint, so a result is never filed under a topology it was not measured on. | +| `HYPERLOOM_PARTITION_COUNT` | Always | Partitions the card presents in that mode. | +| `HYPERLOOM_PARTITION_CU` | Always | Compute units per partition. The entrypoint selects partition devices by matching this exactly — HIP enumerates whole cards before partitions, so an index list computed at launch would be wrong in the one case that matters, and wrong invisibly. | +| `HYPERLOOM_PARTITION_STREAMS_PER_PARTITION` | Fan-out only | Streams to place on each partition. | +| `HYPERLOOM_PARTITION_TOTAL_STREAMS` | Fan-out only | `count × streams`; the total concurrency the entrypoint should drive if it fans out across every partition. | + +Only scriptable frameworks (`xdit`, `custom`) place work per partition. A serving +session is handed no fan-out instruction rather than a concurrency nothing will +drive, and passing the flags with one warns. Its shape is still recorded in the +report and the fingerprint — that is provenance, not a hand-off — and the report +states plainly that the figure cannot be read as an aggregate. + +--- + ## Multi-node / prefill-decode (PD) Use CLI flags for multi-node topology and prefill-decode configuration: diff --git a/src/hyperloom/agents/robustness/role/envelope.py b/src/hyperloom/agents/robustness/role/envelope.py index ceffb32fb7..ccfbd8e164 100644 --- a/src/hyperloom/agents/robustness/role/envelope.py +++ b/src/hyperloom/agents/robustness/role/envelope.py @@ -96,6 +96,7 @@ class IntentType(str, Enum): "model_path", "model_name", "model_class", + "compute_partition", "start_ts", "resumed_ts", "max_minutes", diff --git a/src/hyperloom/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py new file mode 100644 index 0000000000..b0fdac8d0f --- /dev/null +++ b/src/hyperloom/common/gpu_partition.py @@ -0,0 +1,566 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""AMD compute-partition modes (SPX/DPX/QPX/CPX) as an observed session property. + +The mode a card is in is a property of the *card*, not of a process: changing it +is privileged, it evicts every process resident on the card, and it renumbers +the devices underneath anything already running. Nothing here changes it. The +mode must be established before ``optimize`` starts -- by the operator or the +provisioning platform, since the shape is checked and recorded at launch, before +any benchmark process exists to establish it -- and this module reads it, +describes it, and answers one question about it before a session commits three +hours to a shape that cannot work. + +That division is deliberate and matches how the rest of the repo treats +high-risk host state: NPS and the CPU governor are probed and warned about, not +set, and BIOS-level knobs are deferred to ``scripts/platform_audit.py``. An +optimization loop that runs agent-authored code is the wrong place to hold a +privileged hardware mutation, so every entry point here is a read. + +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 names one. Measured on one +MI355X with a 1.26B-parameter vision model at six views per forward pass, 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 belongs to whoever owns the SLA, which is another reason it is an input +here rather than something the optimizer decides. + +Two streams per partition is where every mode peaked on that workload, but it is +a ceiling only a light footprint can reach: the same model at 25 views per pass +peaks at 20.7 GiB per stream, which two streams cannot hold in a 36 GiB CPX +partition. Whether a session can run in the mode it was given is therefore a +memory question, and :func:`fits_in_partition` is what answers it -- at launch, +where a refusal costs seconds instead of hours. + +Two invariants this module exists to enforce: + +* **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. +* **Ask the device, do not derive from a table.** Under a split mode a device + *is* a partition, so ``amd-smi`` reports that partition's own CU count and + memory. Reading them is exact and needs no assumption about how a board + divides; deriving them from a per-board table means a board whose entry is + stale produces a plausible wrong answer. The table remains only as a fallback + for when the device cannot be reached at all, and says so when it is used. +""" + +from __future__ import annotations + +import json +import logging +import os +import subprocess +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, 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. Used to interpret a +#: mode name; the per-partition CU count and memory are read from the device. +MODE_PARTITION_COUNTS: dict[str, int] = { + "SPX": 1, + "DPX": 2, + "QPX": 4, + "CPX": 8, +} + +#: The unpartitioned mode: one partition spanning the whole card. The only mode +#: every board supports, and what a card is in unless someone changed it. +UNPARTITIONED_MODE = "SPX" + +#: The observed shape, published once at launch for the benchmark entrypoint to +#: fan out across and for the provenance record to quote. Named here, in +#: ``common``, because the platform fingerprint reads them and must not import +#: the orchestrator -- and must not spawn a probe of its own, since it also runs +#: on the crash path where a subprocess is the last thing wanted. +PARTITION_MODE_ENV = "HYPERLOOM_PARTITION_MODE" +PARTITION_COUNT_ENV = "HYPERLOOM_PARTITION_COUNT" +PARTITION_CU_ENV = "HYPERLOOM_PARTITION_CU" +PARTITION_STREAMS_ENV = "HYPERLOOM_PARTITION_STREAMS_PER_PARTITION" +PARTITION_TOTAL_STREAMS_ENV = "HYPERLOOM_PARTITION_TOTAL_STREAMS" + +_READ_TIMEOUT_S = 30.0 + + +class PartitionError(RuntimeError): + """Raised when a partition mode cannot be read or interpreted.""" + + +@dataclass(frozen=True) +class PartitionLayout: + """What one compute-partition mode means for 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 it could + not be determined. + probed: ``True`` when ``cu_per_partition`` came from the device, + ``False`` when it was derived from the per-board table. Carried so a + caller can say which, rather than presenting a fallback as a + measurement. + """ + + mode: str + partitions: int + cu_per_partition: int + gib_per_partition: float | None = None + probed: bool = False + + @property + def partitioned(self) -> bool: + """Whether this mode splits the card at all.""" + return self.partitions > 1 + + 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 "" + source = "" if self.probed else " (derived)" + return f"{self.mode} ({self.partitions} x {self.cu_per_partition} CU{mem}){source}" + + +def parse_mode(raw: str | None) -> str: + """Parse an operator-supplied mode name into its canonical form. + + Args: + raw: A mode name (``"cpx"``, ``"CPX"``). Empty or ``None`` means the + operator named no expectation. + + Returns: + The canonical upper-case mode name, or ``""`` when nothing was named. + + Raises: + PartitionError: If the value is not one of the known modes. Refused at + parse time so a typo is a usage error rather than a session that + runs to completion under an assertion that could never hold. + """ + name = str(raw or "").strip().upper() + if not name: + return "" + if name not in MODE_PARTITION_COUNTS: + raise PartitionError( + f"unknown compute-partition mode {raw!r}; expected one of {', '.join(MODE_PARTITION_COUNTS)}" + ) + return name + + +def _amd_smi_json(args: Sequence[str], timeout_s: float = _READ_TIMEOUT_S) -> object: + """Run a read-only ``amd-smi`` subcommand with ``--json`` and parse it. + + Every call this module makes is unprivileged. The one query that would need + elevation -- ``amd-smi partition -a``, which enumerates the profiles a board + *could* enter -- is deliberately absent: it exists to validate a mode before + setting one, and nothing here sets one. + + Args: + args: Subcommand and its flags. + timeout_s: Per-call timeout. + + Returns: + The parsed JSON payload. + + Raises: + PartitionError: If ``amd-smi`` is missing, fails, times out, or returns + output that is not JSON. + """ + cmd = ["amd-smi", *args, "--json"] + try: + proc = subprocess.run( # noqa: S603 — fixed argv, no shell + cmd, + capture_output=True, + text=True, + timeout=timeout_s, + check=False, + ) + except FileNotFoundError as exc: + raise PartitionError("amd-smi not found; reading 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 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"]) + 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 = 0) -> 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] + + +#: Field names ``amd-smi static --asic`` has used for a device's compute-unit +#: count. Tried in order; the first that parses wins. Several are accepted +#: because the schema has moved between releases and a missing CU count silently +#: sends the caller to the per-board table, which is the answer this probe +#: exists to avoid. +_CU_KEYS: tuple[str, ...] = ("num_compute_units", "compute_units", "num_cu", "cu_count") + + +def read_device_cu(gpu_id: int = 0) -> int | None: + """Read the compute-unit count of one device. + + Under a split mode a device *is* a partition, so this is that partition's CU + count -- exactly the number :func:`partition_device_predicate` matches on, + with no assumption about how the board divides. In an unpartitioned mode it + is the whole card's. + + Args: + gpu_id: GPU to interrogate. + + Returns: + The device's CU count, or ``None`` when it could not be read. + """ + try: + payload = _amd_smi_json(["static", "-g", str(gpu_id), "--asic"]) + except PartitionError as exc: + # Warned, not debugged: falling back to the per-board table is a + # downgrade in accuracy, and the consequence lands far from here as a + # benchmark that finds no device of the expected width. + log.warning("could not read GPU %d compute units, will fall back to the board table: %s", gpu_id, exc) + return None + rows = payload.get("gpu_data") if isinstance(payload, dict) else payload + candidates = rows if isinstance(rows, list) else [payload] + for row in candidates: + if not isinstance(row, dict): + continue + asic = row.get("asic") if isinstance(row.get("asic"), dict) else row + for key in _CU_KEYS: + try: + value = int(asic[key]) + except (KeyError, TypeError, ValueError): + continue + if value > 0: + return value + log.warning("GPU %d reported no compute-unit count, will fall back to the board table", gpu_id) + return None + + +#: 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_device_gib(gpu_id: int = 0) -> float | None: + """Read the HBM available to one device, in GiB. + + ``amd-smi`` reports VRAM per device, and under a split mode a device is a + partition -- so this is the memory one partition has, which is the figure + :func:`fits_in_partition` needs. Read rather than derived by dividing a card + total, which is what made the same field ambiguous when the mode was + something this process changed underneath itself. + + Args: + gpu_id: GPU to interrogate. + + Returns: + The device's HBM in GiB, or ``None`` when unreadable or unparseable. + """ + try: + payload = _amd_smi_json(["static", "-g", str(gpu_id), "--vram"]) + except PartitionError as exc: + # Warned rather than debugged: this figure is the only input to the + # feasibility check, so losing it means the session proceeds without one + # -- a quieter log than the drop it is meant to prevent would invert the + # severities. + log.warning("could not read HBM capacity for GPU %d; feasibility will not be checked: %s", gpu_id, exc) + return None + + rows = payload.get("gpu_data") if isinstance(payload, dict) else payload + if not isinstance(rows, list): + log.warning("GPU %d returned no VRAM rows; feasibility will not be checked", gpu_id) + 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 + log.warning("GPU %d reported no usable VRAM size; feasibility will not be checked", gpu_id) + return None + + +def layout_for( + mode: str, + *, + gpu_type: str | None = None, + cu_per_partition: int | None = None, + gib_per_partition: float | None = None, +) -> PartitionLayout: + """Describe what ``mode`` means, preferring probed numbers to tabled ones. + + Args: + mode: Compute-partition mode. + gpu_type: Board name, used only for the fallback CU derivation. + cu_per_partition: The partition's CU count as read from the device. When + given, it is used as-is and no board table is consulted. + gib_per_partition: The partition's HBM as read from the device. + + Returns: + The layout, with ``probed`` recording which source the CU count came + from. + + Raises: + PartitionError: If the mode is unknown, or if no CU count was probed and + the board cannot be sized from the table. + """ + canonical = parse_mode(mode) + if not canonical: + raise PartitionError("no compute-partition mode given") + partitions = MODE_PARTITION_COUNTS[canonical] + if cu_per_partition and int(cu_per_partition) > 0: + return PartitionLayout( + mode=canonical, + partitions=partitions, + cu_per_partition=int(cu_per_partition), + gib_per_partition=gib_per_partition, + probed=True, + ) + + # Fallback: derive from the board's total. This is the path that made a + # stale table dangerous, so it is only reached when the device could not be + # asked, and the result is marked as derived. + identity = AMD_GPU_DISPATCH_IDENTITIES.get(str(gpu_type or "").strip().lower()) + if identity is None: + raise PartitionError( + f"cannot size {canonical} partitions: GPU {gpu_type!r} reported no CU count and is not in the board table" + ) + 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. + raise PartitionError( + f"{gpu_type} has {cu_total} CU by the board table, which does not divide into " + f"{partitions} {canonical} partitions; the per-partition CU count would be wrong " + f"and device selection matches on it exactly" + ) + return PartitionLayout( + mode=canonical, + partitions=partitions, + cu_per_partition=cu_total // partitions, + gib_per_partition=gib_per_partition, + probed=False, + ) + + +def observe_partition(gpu_id: int = 0, *, gpu_type: str | None = None) -> PartitionLayout | None: + """Read the live partition topology of one card. + + The single entry point a caller needs to answer "what shape is this card in, + and how big is one partition". Everything comes from the device except the + CU fallback, which is marked as such. + + Args: + gpu_id: GPU to interrogate. + gpu_type: Board name, used only if the CU probe fails. + + Returns: + The live layout, or ``None`` when the mode itself could not be read -- + which is the ordinary case on a host without ``amd-smi``, and means the + caller should proceed without a partitioning opinion rather than fail. + """ + try: + mode = read_partition_mode(gpu_id) + except PartitionError as exc: + log.debug("no compute-partition state for GPU %d: %s", gpu_id, exc) + return None + if mode not in MODE_PARTITION_COUNTS: + log.warning("GPU %d reports compute-partition mode %r, which this build does not know", gpu_id, mode) + return None + try: + return layout_for( + mode, + gpu_type=gpu_type, + cu_per_partition=read_device_cu(gpu_id), + gib_per_partition=read_device_gib(gpu_id), + ) + except PartitionError as exc: + log.warning("GPU %d is in %s but its partitions could not be sized: %s", gpu_id, mode, exc) + return None + + +def published_shape(env: Mapping[str, str] | None = None) -> dict[str, Any] | None: + """Read the partition shape a launch published, without touching the device. + + The provenance record needs to state which topology produced a number, and + it is written in places where spawning ``amd-smi`` is not acceptable: the + crash-safe ``final.json`` writer, and unit tests that must not run a probe. + The launch already established the shape and put it in the environment, so + this reads that rather than asking again. + + Args: + env: Environment to read; defaults to the process environment. + + Returns: + The shape, or ``None`` when this session published none -- which is the + ordinary case, and must stay distinguishable from an unpartitioned card. + """ + source = os.environ if env is None else env + mode = str(source.get(PARTITION_MODE_ENV, "") or "").strip().upper() + if not mode: + return None + + def _int(key: str) -> int | None: + try: + return int(str(source.get(key, "")).strip()) + except (TypeError, ValueError): + return None + + shape: dict[str, Any] = {"mode": mode} + for key, name in ( + (PARTITION_COUNT_ENV, "partitions"), + (PARTITION_CU_ENV, "cu_per_partition"), + (PARTITION_STREAMS_ENV, "streams_per_partition"), + ): + value = _int(key) + if value is not None: + shape[name] = value + return shape + + +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 + + +__all__ = [ + "MODE_PARTITION_COUNTS", + "PARTITION_COUNT_ENV", + "PARTITION_CU_ENV", + "PARTITION_MODE_ENV", + "PARTITION_STREAMS_ENV", + "PARTITION_TOTAL_STREAMS_ENV", + "UNPARTITIONED_MODE", + "PartitionError", + "PartitionLayout", + "fits_in_partition", + "layout_for", + "observe_partition", + "parse_mode", + "partition_device_predicate", + "published_shape", + "read_device_cu", + "read_device_gib", + "read_partition_mode", + "read_partition_modes", +] diff --git a/src/hyperloom/common/platform_probe.py b/src/hyperloom/common/platform_probe.py index d89098f0d8..3be5165561 100644 --- a/src/hyperloom/common/platform_probe.py +++ b/src/hyperloom/common/platform_probe.py @@ -37,6 +37,7 @@ from pathlib import Path from typing import Any +from hyperloom.common.gpu_partition import published_shape from hyperloom.common.provenance import detect_gfx_arch, detect_stack_fingerprint log = logging.getLogger(__name__) @@ -251,6 +252,15 @@ def platform_fingerprint( "gfx_arch": detect_gfx_arch(os.environ, gpu_type=gpu_type, probe=False) or "unknown", "amdgpu_driver": read_kernel_file("/sys/module/amdgpu/version") or "unknown", } + # The card's compute-partition shape, when this session established + # one. Recorded for the same reason NPS is: it changes what the + # numbers mean, and without it two runs of the same configuration on + # the same card in SPX and in CPX are indistinguishable in the + # history. Read from the env the launch published rather than probed, + # so this stays subprocess-free on the crash path. + partition = published_shape() + if partition: + record["gpu"]["compute_partition"] = partition except Exception: # noqa: BLE001 - one degraded field, not a dropped record log.warning("platform fingerprint: GPU block unreadable", exc_info=True) record["gpu"] = {"status": "error"} 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..0c61a811e1 --- /dev/null +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Tests for reading a card's compute-partition shape. + +Every entry point in the module under test is a read, so the tests supply +``amd-smi`` payloads rather than asserting on commands issued. The payload +shapes are the real ones: a partitioned device reports its *partition's* CU +count and VRAM, which is the fact the module is built around. +""" + +from __future__ import annotations + +import pytest + +from hyperloom.common import gpu_partition as gp + + +@pytest.fixture +def smi(monkeypatch): + """Route every amd-smi read through a scripted payload table.""" + payloads: dict[str, object] = {} + calls: list[list[str]] = [] + + def fake(args, timeout_s=gp._READ_TIMEOUT_S): + argv = list(args) + calls.append(argv) + for key, payload in payloads.items(): + if key in " ".join(argv): + if isinstance(payload, Exception): + raise payload + return payload + raise gp.PartitionError(f"no scripted payload for {argv}") + + monkeypatch.setattr(gp, "_amd_smi_json", fake) + return type("Smi", (), {"payloads": payloads, "calls": calls})() + + +def _partition_payload(mode: str, gpu_id: int = 0) -> dict: + return {"current_partition": [{"gpu_id": gpu_id, "accelerator_type": mode}]} + + +def _vram_payload(value: float, unit: str = "MB") -> dict: + return {"gpu_data": [{"vram": {"size": {"value": value, "unit": unit}}}]} + + +def _asic_payload(cu: int) -> dict: + return {"gpu_data": [{"asic": {"num_compute_units": cu}}]} + + +class TestParseMode: + def test_canonicalizes_case_and_whitespace(self): + assert gp.parse_mode(" cpx ") == "CPX" + + @pytest.mark.parametrize("raw", [None, "", " "]) + def test_nothing_named_is_empty(self, raw): + assert gp.parse_mode(raw) == "" + + def test_an_unknown_mode_is_a_parse_error(self): + """A typo must fail at parse time, not as an assertion that never holds.""" + with pytest.raises(gp.PartitionError, match="unknown compute-partition mode"): + gp.parse_mode("opx") + + def test_the_error_lists_what_is_valid(self): + with pytest.raises(gp.PartitionError, match="SPX"): + gp.parse_mode("octuple") + + +class TestReadPartitionMode: + def test_reads_the_mode_of_the_named_card(self, smi): + smi.payloads["partition"] = { + "current_partition": [ + {"gpu_id": 0, "accelerator_type": "SPX"}, + {"gpu_id": 3, "accelerator_type": "CPX"}, + ] + } + assert gp.read_partition_mode(3) == "CPX" + + def test_a_card_absent_from_the_report_raises(self, smi): + smi.payloads["partition"] = _partition_payload("SPX", gpu_id=0) + with pytest.raises(gp.PartitionError, match="no state for GPU 7"): + gp.read_partition_mode(7) + + def test_an_empty_report_raises_rather_than_returning_nothing(self, smi): + smi.payloads["partition"] = {"current_partition": []} + with pytest.raises(gp.PartitionError, match="reported no compute-partition state"): + gp.read_partition_modes() + + def test_rows_without_a_usable_gpu_id_are_skipped(self, smi): + smi.payloads["partition"] = { + "current_partition": [ + {"gpu_id": "n/a", "accelerator_type": "SPX"}, + {"gpu_id": 1, "accelerator_type": "DPX"}, + ] + } + assert gp.read_partition_modes() == {1: "DPX"} + + +class TestReadDeviceCu: + def test_reads_the_devices_own_cu_count(self, smi): + smi.payloads["--asic"] = _asic_payload(32) + assert gp.read_device_cu(0) == 32 + + @pytest.mark.parametrize("key", ["num_compute_units", "compute_units", "num_cu", "cu_count"]) + def test_each_field_name_in_use_is_accepted(self, smi, key): + """The schema has moved between releases; a miss sends us to the table.""" + smi.payloads["--asic"] = {"gpu_data": [{"asic": {key: 64}}]} + assert gp.read_device_cu(0) == 64 + + def test_a_flat_payload_without_an_asic_block_still_parses(self, smi): + smi.payloads["--asic"] = {"gpu_data": [{"num_compute_units": 256}]} + assert gp.read_device_cu(0) == 256 + + def test_an_unreadable_probe_is_none_and_warns(self, smi, caplog): + smi.payloads["--asic"] = gp.PartitionError("amd-smi absent") + with caplog.at_level("WARNING"): + assert gp.read_device_cu(0) is None + assert "fall back to the board table" in caplog.text + + def test_a_missing_count_is_none_and_warns(self, smi, caplog): + smi.payloads["--asic"] = {"gpu_data": [{"asic": {"market_name": "MI355X"}}]} + with caplog.at_level("WARNING"): + assert gp.read_device_cu(0) is None + assert "no compute-unit count" in caplog.text + + @pytest.mark.parametrize("bad", [0, -8, "n/a", None]) + def test_an_unusable_value_is_not_taken(self, smi, bad): + smi.payloads["--asic"] = {"gpu_data": [{"asic": {"num_compute_units": bad}}]} + assert gp.read_device_cu(0) is None + + +class TestReadDeviceGib: + def test_mebibytes_are_the_binary_kind(self, smi): + """An MI355X's 288 GiB is reported as 294896 "MB".""" + smi.payloads["--vram"] = _vram_payload(294896, "MB") + assert gp.read_device_gib(0) == pytest.approx(288.0, abs=0.1) + + def test_a_partitioned_device_reports_its_partitions_memory(self, smi): + """Under CPX a device *is* a partition, so this needs no dividing.""" + smi.payloads["--vram"] = _vram_payload(36862, "MB") + assert gp.read_device_gib(0) == pytest.approx(36.0, abs=0.1) + + @pytest.mark.parametrize( + ("value", "unit", "expected"), + [(288, "GIB", 288.0), (288, "GB", 288.0), (1, "TIB", 1024.0)], + ) + def test_other_units_scale(self, smi, value, unit, expected): + smi.payloads["--vram"] = _vram_payload(value, unit) + assert gp.read_device_gib(0) == pytest.approx(expected) + + def test_an_unreadable_probe_warns_rather_than_debugs(self, smi, caplog): + """Losing this figure skips the feasibility check, which is the louder event.""" + smi.payloads["--vram"] = gp.PartitionError("amd-smi failed") + with caplog.at_level("WARNING"): + assert gp.read_device_gib(0) is None + assert "feasibility will not be checked" in caplog.text + + def test_an_unknown_unit_is_not_guessed(self, smi): + smi.payloads["--vram"] = _vram_payload(288, "parsecs") + assert gp.read_device_gib(0) is None + + +class TestLayoutFor: + def test_a_probed_cu_count_is_used_verbatim(self): + layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) + assert (layout.partitions, layout.cu_per_partition, layout.probed) == (8, 32, True) + + def test_a_probed_count_does_not_consult_the_board_table(self): + """The table conflates boards that share an ISA; the device does not.""" + layout = gp.layout_for("DPX", gpu_type="mi308x", cu_per_partition=40) + assert layout.cu_per_partition == 40 + assert layout.probed is True + + def test_the_table_is_the_fallback_and_says_so(self): + layout = gp.layout_for("DPX", gpu_type="mi300x") + assert (layout.partitions, layout.cu_per_partition) == (2, 152) + assert layout.probed is False + assert "derived" in layout.describe() + + def test_an_unsizable_board_without_a_probe_raises(self): + with pytest.raises(gp.PartitionError, match="not in the board table"): + gp.layout_for("CPX", gpu_type="some-new-board") + + def test_an_uneven_division_raises_rather_than_flooring(self, monkeypatch): + """A floored CU count matches no device, and reports the wrong cause. + + Every board in the table divides evenly today, so the guard is exercised + against an injected entry -- it exists for the one that does not. + """ + monkeypatch.setitem(gp.AMD_GPU_DISPATCH_IDENTITIES, "oddball", ("gfx950", 300)) + with pytest.raises(gp.PartitionError, match="does not divide"): + gp.layout_for("CPX", gpu_type="oddball") + + def test_an_unknown_mode_raises(self): + with pytest.raises(gp.PartitionError, match="unknown compute-partition mode"): + gp.layout_for("OPX", cu_per_partition=32) + + def test_spx_is_one_partition_and_not_partitioned(self): + layout = gp.layout_for("SPX", cu_per_partition=256) + assert layout.partitions == 1 + assert layout.partitioned is False + + +class TestObservePartition: + def test_describes_the_live_topology_from_the_device(self, smi): + smi.payloads["partition"] = _partition_payload("CPX") + smi.payloads["--asic"] = _asic_payload(32) + smi.payloads["--vram"] = _vram_payload(36862, "MB") + + layout = gp.observe_partition(0) + + assert layout is not None + assert (layout.mode, layout.partitions, layout.cu_per_partition) == ("CPX", 8, 32) + assert layout.gib_per_partition == pytest.approx(36.0, abs=0.1) + assert layout.probed is True + + def test_an_unreadable_card_is_none_not_an_error(self, smi): + """A host without amd-smi is the ordinary case, not a failure.""" + smi.payloads["partition"] = gp.PartitionError("amd-smi not found") + assert gp.observe_partition(0) is None + + def test_a_mode_this_build_does_not_know_is_refused_loudly(self, smi, caplog): + smi.payloads["partition"] = _partition_payload("OPX") + with caplog.at_level("WARNING"): + assert gp.observe_partition(0) is None + assert "does not know" in caplog.text + + def test_falls_back_to_the_table_when_the_cu_probe_fails(self, smi): + smi.payloads["partition"] = _partition_payload("DPX") + smi.payloads["--asic"] = gp.PartitionError("no asic block") + smi.payloads["--vram"] = _vram_payload(294896, "MB") + + layout = gp.observe_partition(0, gpu_type="mi300x") + + assert layout is not None + assert layout.cu_per_partition == 152 + assert layout.probed is False + + +class TestFitsInPartition: + def test_the_measured_mi355x_case_does_not_fit_cpx_in_pairs(self): + """20.7 GiB per stream, two streams, 36 GiB partition.""" + layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) + assert gp.fits_in_partition(20.7, layout, 2) is False + + def test_the_same_footprint_fits_a_single_stream(self): + """Which is why gating on the one-stream figure is the trap.""" + layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) + assert gp.fits_in_partition(20.7, layout, 1) is True + + def test_unknown_capacity_does_not_refuse(self): + layout = gp.layout_for("CPX", cu_per_partition=32) + assert gp.fits_in_partition(999.0, layout, 2) is True + + def test_unknown_footprint_does_not_refuse(self): + layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) + assert gp.fits_in_partition(0.0, layout, 2) is True + + def test_exactly_filling_a_partition_fits(self): + layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) + assert gp.fits_in_partition(18.0, layout, 2) is True + + +class TestPartitionDevicePredicate: + def test_matches_a_partition_and_rejects_a_whole_card(self): + """HIP lists whole cards first, so index selection measures the wrong device.""" + is_partition = gp.partition_device_predicate(32) + assert is_partition(32) is True + assert is_partition(256) is False + + +class TestPublishedShape: + def test_reads_back_what_a_launch_published(self): + env = { + gp.PARTITION_MODE_ENV: "CPX", + gp.PARTITION_COUNT_ENV: "8", + gp.PARTITION_CU_ENV: "32", + gp.PARTITION_STREAMS_ENV: "2", + } + assert gp.published_shape(env) == { + "mode": "CPX", + "partitions": 8, + "cu_per_partition": 32, + "streams_per_partition": 2, + } + + def test_no_published_mode_is_none(self): + """Distinct from an unpartitioned card, which is a known SPX.""" + assert gp.published_shape({}) is None + + def test_unparseable_numbers_are_omitted_not_zeroed(self): + env = {gp.PARTITION_MODE_ENV: "DPX", gp.PARTITION_COUNT_ENV: "lots"} + assert gp.published_shape(env) == {"mode": "DPX"} + + def test_it_reads_no_device(self, monkeypatch): + """This runs on the crash path, where spawning a probe is unacceptable.""" + + def explode(*args, **kwargs): + raise AssertionError("published_shape must not run amd-smi") + + monkeypatch.setattr(gp, "_amd_smi_json", explode) + assert gp.published_shape({gp.PARTITION_MODE_ENV: "SPX"}) == {"mode": "SPX"} diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 5993e7638e..11b4930b7f 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1608,6 +1608,205 @@ def _export_operator_launch_shape( os.environ.pop("INFERENCE_OPTIMIZER_EXTRA_ENV", None) +def _partition_fanout_supported(framework: str | None) -> tuple[bool, str]: + """Whether this framework's runner can place work per partition. + + Split out and returned rather than tested inline because the interesting + case is the third one. ``--framework`` defaults to ``None`` and is resolved + later, so a truthiness guard here reads as "checked and fine" while actually + meaning "not checked" -- and the operator sees nothing either way. + + Args: + framework: The session framework, possibly unresolved. + + Returns: + ``(supported, detail)``. ``detail`` is empty when supported, and + otherwise says whether the answer is *no* or *not yet known*. + """ + name = str(framework or "").strip().lower() + if not name: + return False, ( + "the framework is not resolved yet, so whether its runner places work " + "per partition could not be checked here" + ) + if framework_registry.is_scriptable(name): + return True, "" + return False, ( + f"{name!r} runs a server, and its benchmark does not place work per partition, " + f"so the streams-per-partition setting would be ignored" + ) + + +def _export_partition_shape( + *, + declared_mode: str | None, + streams_per_partition: int | None, + framework: str | None = None, + gpu_type: str | None = None, + nodes: int = 1, + model_path: str | None = None, + precision: str | None = None, + shared_state: Any = None, +) -> dict[str, Any]: + """Validate the session's compute-partition shape and publish it. + + Read-only. The mode belongs to the card and is set outside the optimizer; + this observes what the card is in, refuses a session that cannot work in + that shape, and hands the shape to the benchmark entrypoint that fans work + out across partitions. + + Validating at launch is the whole point. Every failure this catches -- a mode + that was never applied, a partition too small for the workload -- otherwise + surfaces hours later as an out-of-memory crash or, worse, as a perfectly + good measurement filed under a topology the card was never in. + + Args: + declared_mode: The raw ``--compute-partition-mode`` value, if any. + streams_per_partition: The raw ``--streams-per-partition`` value, if any. + framework: The resolved session framework. Decides whether anything will + place work per partition, which gates both the footprint refusal and + the runtime hand-off. Must be resolved before this is called: an + unresolved framework reads as "cannot fan out". + gpu_type: Board name, used only if the device's CU count cannot be read. + nodes: Resolved node count. The shape describes the card this process + can see, which on a multi-node session is not the benchmark's, so + nothing is observed or recorded there. + model_path: Checkpoint to size the workload from. Without it the + feasibility check has nothing to weigh and can only warn. + precision: Resolved precision, which sets the bytes per weight. + shared_state: Persisted state on a resume, carrying any measured + per-stream peak. Tighter than the weight-bytes bound, so preferred + when present. + + Returns: + The published shape, or ``{}`` when this session has none. + """ + from hyperloom.common.gpu_partition import PartitionError, parse_mode + from hyperloom.orchestrator.actions.executors._partition_shape import ( + DEFAULT_STREAMS_PER_PARTITION, + PARTITION_COUNT_ENV, + PARTITION_CU_ENV, + PARTITION_MODE_ENV, + PARTITION_STREAMS_ENV, + PARTITION_TOTAL_STREAMS_ENV, + runtime_env, + session_shape_summary, + validate_session_shape, + ) + + # Cleared first so a second session in the same shell cannot inherit a shape + # the operator did not ask for this time. + for key in ( + PARTITION_MODE_ENV, + PARTITION_COUNT_ENV, + PARTITION_CU_ENV, + PARTITION_STREAMS_ENV, + PARTITION_TOTAL_STREAMS_ENV, + ): + os.environ.pop(key, None) + + try: + mode = parse_mode(declared_mode) + except PartitionError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(2) + + streams_named = streams_per_partition is not None + # Tested against None rather than falsiness: `0 or DEFAULT` is DEFAULT, which + # would quietly honour an invalid request as the default instead of refusing + # it, and leave the guard below unreachable for the one value most likely to + # be passed by mistake. + streams = DEFAULT_STREAMS_PER_PARTITION if streams_per_partition is None else int(streams_per_partition) + if streams < 1: + print( + f"ERROR: --streams-per-partition must be >= 1, got {streams_per_partition}", + file=sys.stderr, + ) + sys.exit(2) + + if nodes >= 2: + # Unconditional: the misleading record is produced by observing at all, + # not by asking for a mode. A controller that happens to have GPUs would + # otherwise publish its own topology as the session's while the + # benchmark ran somewhere else entirely. + print( + f"WARN: this session has --nodes {nodes}, and a compute-partition shape " + f"describes one card. The benchmark node's topology cannot be read from here, " + f"so no shape is recorded for this session.", + file=sys.stderr, + ) + if mode: + print( + f"ERROR: --compute-partition-mode {mode} cannot be checked on a " + f"--nodes {nodes} session: the assertion is about the benchmark node's " + f"card, which this process cannot read. An unverifiable assertion is not " + f"a satisfied one. Drop the flag to run multi-node.", + file=sys.stderr, + ) + sys.exit(2) + return {} + + fanout, fanout_detail = _partition_fanout_supported(framework) + if (mode or streams_named) and not fanout: + print(f"WARN: {fanout_detail}.", file=sys.stderr) + + verdict = validate_session_shape( + declared_mode=mode, + streams=streams, + gpu_type=gpu_type, + params={ + "model_path": str(model_path or ""), + "precision": str(precision or ""), + }, + shared_state=shared_state, + # The footprint refusal is arithmetic about streams sharing a partition. + # An operator who named the flags has asserted that shape and is held to + # it; otherwise it only applies where something will actually fan out. + fanout_expected=fanout or bool(mode) or streams_named, + ) + for warning in verdict.warnings: + print(f"WARN: {warning}", file=sys.stderr) + if not verdict.ok: + print(f"ERROR: {verdict.refusal}", file=sys.stderr) + sys.exit(2) + for note in verdict.notes: + print(note) + + if verdict.layout is None: + return {} + # The observed shape is published either way -- the platform fingerprint + # reads it back from here, and provenance is the point. Only the fan-out + # instruction is withheld from a session whose benchmark cannot act on it. + os.environ.update(runtime_env(verdict.layout, streams, fanout=fanout)) + return session_shape_summary(verdict.layout, streams, fanout_expected=fanout) + + +def _restore_partition_shape_from_state(args: Any, state: SharedState) -> None: + """Fill the partition flags from the archive when this resume omitted them. + + Priority is CLI flag > 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_shape` stays the only writer of the env it owns, + and so a restored declaration is re-checked against the live card rather + than trusted -- the card may have been repartitioned while the session was + stopped, which is exactly the case the declaration exists to catch. + + Args: + args: Parsed CLI namespace, updated in place. + state: Resumed session state. + """ + archived = dict(getattr(state, "compute_partition", None) or {}) + if not str(getattr(args, "compute_partition_mode", None) or "").strip(): + args.compute_partition_mode = str(archived.get("mode") or "") + # `is None` rather than falsiness, so an explicit `--streams-per-partition 0` + # still reaches the guard that refuses it instead of being read as "omitted" + # and quietly replaced by the archived value. + if getattr(args, "streams_per_partition", None) is None: + streams = archived.get("streams_per_partition") + args.streams_per_partition = int(streams) if streams else None + + # 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. @@ -1692,6 +1891,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), ) + # The partition shape is deliberately NOT exported here. It needs the + # resolved framework, GPU type and post-quantization model path, none of + # which exist yet, and each fresh-launch and resume branch calls it once at + # the point where they do. Exporting here as well made a resume validate + # twice -- the first time against an unresolved model it could only warn + # about. + # 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 @@ -1929,6 +2135,28 @@ async def _run_optimize(args: argparse.Namespace) -> int: gpu_type=os.environ.get("GPU_TYPE") or state.gpu_type, ) _persist_operator_supplied_paths(state) + # The partition shape is part of the measurement contract, so it resumes + # on the same restore / apply / persist path as the paths above. The + # re-check is not ceremony: a card can be repartitioned while a session + # is stopped, and resuming into a different topology would compare + # candidates measured under one shape against a baseline from another. + _restore_partition_shape_from_state(args, state) + state.compute_partition = _export_partition_shape( + declared_mode=getattr(args, "compute_partition_mode", None), + streams_per_partition=getattr(args, "streams_per_partition", None), + framework=state.framework or getattr(args, "framework", None), + gpu_type=os.environ.get("GPU_TYPE") or state.gpu_type, + # 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)), + model_path=state.model_path or str(getattr(args, "model", "") or ""), + precision=state.precision or getattr(args, "precision", None), + # A resume can size the workload against what the last session + # actually measured, which rules out a mode the weights alone fit. + shared_state=state, + ) + if state.compute_partition.get("mode"): + print(f" re-exported partition shape: {state.compute_partition['mode']}") if state.framework_repo_path: print(f" re-exported FRAMEWORK_REPO_PATH: {state.framework_repo_path}") if state.bypass_scripts_dir: @@ -2219,10 +2447,27 @@ async def _run_optimize(args: argparse.Namespace) -> int: model=str(args.model) if args.model else "", launch_info_file=getattr(args, "launch_info_file", None), ) + # Placed here, not at the top of _run_optimize, because everything it + # weighs is resolved by now and none of it was then: the framework + # (whether anything fans out), args.gpu_type (the CU fallback), and + # args.model, which --quantize rewrites to the exported checkpoint -- + # sizing partitions against the source model would weigh the wrong + # weights. Still before the seed, so the shape it returns is the one + # persisted rather than a lossy re-read from the environment. + compute_partition = _export_partition_shape( + declared_mode=getattr(args, "compute_partition_mode", None), + streams_per_partition=getattr(args, "streams_per_partition", None), + framework=framework, + gpu_type=args.gpu_type, + nodes=nodes_resolved, + model_path=str(args.model or os.environ.get("MODEL_PATH") or ""), + precision=getattr(args, "precision", None), + ) state = _seed_shared_state( session_dir, args, session_id=manifest["session_id"], + compute_partition=compute_partition, ) # Unsupported-model preflight: reject multimodal/vision configs (runs after seed, before heavy bring-up). if _preflight_unsupported_model_arch(args, session_dir): diff --git a/src/hyperloom/inference_optimizer/cli/bootstrap.py b/src/hyperloom/inference_optimizer/cli/bootstrap.py index acde7ee188..e10e6d08d1 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.gpu_partition import published_shape from hyperloom.common.timeutil import now_iso from hyperloom.orchestrator.actions.executors._workload_envs import ( agentx_enabled as _agentx_enabled, @@ -193,6 +194,7 @@ def _seed_shared_state( args: argparse.Namespace, *, session_id: str, + compute_partition: dict[str, Any] | None = None, ) -> SharedState: """Construct and persist the initial :class:`SharedState` for a run. @@ -203,6 +205,12 @@ def _seed_shared_state( session_dir: Directory for the new session. args: Parsed CLI arguments. session_id: Identifier assigned to the session. + compute_partition: The shape the launch validated, passed in rather than + re-read because the environment carries a lossy subset of it: the + published variables cannot express where the CU count came from, and + an absent provenance flag would be reported as a board-table guess + when the device was in fact probed. Falls back to the published + variables when a caller has no verdict to hand over. Returns: The seeded :class:`SharedState` instance. @@ -395,6 +403,7 @@ 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=dict(compute_partition if compute_partition is not None else (published_shape() or {})), 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 722630b5f6..79303018b4 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -287,6 +287,42 @@ 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-mode", + type=str, + default=None, + metavar="MODE", + help="Declare the compute-partition mode the GPU is already in: SPX " + "(whole card), DPX (2), QPX (4) or CPX (8). This is an assertion, not " + "a request -- nothing in the optimizer changes the mode, because doing " + "so is privileged, evicts every process on the card, and renumbers its " + "devices. Set the mode with amd-smi before launching this command: the " + "shape is checked and recorded at launch, so a mode applied later is " + "too late to be either. Passing this flag makes the session " + "refuse to start if the card is in a different mode, which is what " + "catches an external set that did not take effect before three hours " + "are spent recording the numbers under the wrong topology. The mode is " + "recorded in the platform fingerprint whether or not this flag is " + "passed. See also --streams-per-partition.", + ) + 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 it is resolved. + default=None, + metavar="N", + help="Concurrent streams to place on each partition when the card is " + "partitioned. 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 the workload has to fit, and the session refuses to start " + "when it provably will not. Only a scriptable framework's benchmark " + "places work per partition; passing this with a serving framework " + "warns, because nothing would act on it.", + ) opt.add_argument( "--framework", choices=list(framework_registry.names()), diff --git a/src/hyperloom/inference_optimizer/tests/test_cli_bootstrap.py b/src/hyperloom/inference_optimizer/tests/test_cli_bootstrap.py index f6d820f1f7..62b278bdd2 100644 --- a/src/hyperloom/inference_optimizer/tests/test_cli_bootstrap.py +++ b/src/hyperloom/inference_optimizer/tests/test_cli_bootstrap.py @@ -164,6 +164,48 @@ def test_seed_shared_state_records_custom_workload_paths( assert state.benchmark_backend == "bypass" +def _neutralize_seed_io(monkeypatch): + """Stub the model/recipe reads so a seed can be asserted on one field.""" + monkeypatch.setattr(cb, "_load_model_config_tags", lambda _p: {}) + monkeypatch.setattr(cb, "_load_model_arch", lambda *_a, **_k: {}) + monkeypatch.setattr(cb, "_resolve_reference_recipe", lambda _args: ("", {}, "", "")) + from hyperloom.orchestrator.policy import gate as policy + + monkeypatch.setattr(policy, "detect_gpu_count", lambda: 1) + monkeypatch.setattr(policy, "research_lane_ceiling", lambda: 1) + + +def test_seed_records_the_launch_verdict_for_the_partition_shape( + tmp_path: Path, + monkeypatch, +) -> None: + """The verdict carries provenance the published env cannot express. + + ``published_shape()`` reads back mode, count, CU and streams, but nothing + that says the CU count was probed from the device rather than derived from + the board table. Re-reading the env therefore reported a fresh launch's + probed count as a table guess, which is the one thing the section is for. + """ + _neutralize_seed_io(monkeypatch) + monkeypatch.setattr(cb, "published_shape", lambda: {"mode": "CPX", "cu_per_partition": 32}) + + verdict = {"mode": "CPX", "partitions": 8, "cu_per_partition": 32, "cu_probed": True} + state = cb._seed_shared_state(tmp_path, _args(), session_id="s-shape", compute_partition=verdict) + + assert state.compute_partition == verdict + + +def test_seed_falls_back_to_the_published_shape_when_handed_no_verdict( + tmp_path: Path, + monkeypatch, +) -> None: + _neutralize_seed_io(monkeypatch) + monkeypatch.setattr(cb, "published_shape", lambda: {"mode": "DPX"}) + + state = cb._seed_shared_state(tmp_path, _args(), session_id="s-fallback") + assert state.compute_partition == {"mode": "DPX"} + + def test_seed_shared_state_exact_forge_records_native_kernel_optimizer( tmp_path: Path, monkeypatch, diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_shape.py b/src/hyperloom/inference_optimizer/tests/test_partition_shape.py new file mode 100644 index 0000000000..3dfd44efa5 --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_partition_shape.py @@ -0,0 +1,615 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Tests for the launch-time check on a session's compute-partition shape. + +The module under test never changes the card, so these tests are about what it +concludes from what it reads. Two behaviours carry the weight: a declared mode +that cannot be verified must refuse the session rather than assume it holds, and +a workload that provably will not fit must be refused in milliseconds at launch +rather than discovered as an out-of-memory crash hours in. +""" + +from __future__ import annotations + +import pytest + +from hyperloom.common.gpu_partition import PartitionLayout, layout_for +from hyperloom.orchestrator.actions.executors import _partition_shape as ps + + +@pytest.fixture(autouse=True) +def _clear_env(monkeypatch): + for key in ( + ps.PARTITION_GPU_ENV, + ps.PARTITION_MODE_ENV, + ps.PARTITION_COUNT_ENV, + ps.PARTITION_CU_ENV, + ps.PARTITION_STREAMS_ENV, + ps.PARTITION_TOTAL_STREAMS_ENV, + ): + monkeypatch.delenv(key, raising=False) + + +@pytest.fixture +def card(monkeypatch): + """Present a card in a given shape, or none at all.""" + + def present(layout): + monkeypatch.setattr(ps, "observe_partition", lambda gpu_id, gpu_type=None: layout) + + return present + + +def _layout(mode, cu, gib=None, probed=True): + """Build a layout directly, so a table-derived one can be posed as such.""" + if probed: + return layout_for(mode, cu_per_partition=cu, gib_per_partition=gib) + derived = layout_for(mode, cu_per_partition=cu, gib_per_partition=gib) + return PartitionLayout( + mode=derived.mode, + partitions=derived.partitions, + cu_per_partition=derived.cu_per_partition, + gib_per_partition=derived.gib_per_partition, + probed=False, + ) + + +CPX_36 = _layout("CPX", 32, 36.0) +SPX_288 = _layout("SPX", 256, 288.0) + + +class TestUnreadableCard: + def test_a_declared_mode_that_cannot_be_verified_refuses(self, card): + """The flag exists to catch a set that did not take; unverified is not satisfied.""" + card(None) + verdict = ps.validate_session_shape(declared_mode="CPX") + + assert verdict.ok is False + assert "could not be read" in verdict.refusal + assert "it does not set it" in verdict.refusal + + def test_nothing_declared_on_an_unreadable_card_runs_as_before(self, card): + """A host without amd-smi is the ordinary case, not a failure.""" + card(None) + verdict = ps.validate_session_shape() + + assert verdict.ok is True + assert verdict.layout is None + assert verdict.warnings == () + + +class TestDeclaredMode: + def test_a_mismatch_refuses_and_names_both_modes(self, card): + card(SPX_288) + verdict = ps.validate_session_shape(declared_mode="CPX") + + assert verdict.ok is False + assert "is in SPX, not the declared CPX" in verdict.refusal + + def test_the_refusal_says_the_optimizer_will_not_fix_it(self, card): + """Otherwise the natural reading is that the flag applies the mode.""" + card(SPX_288) + verdict = ps.validate_session_shape(declared_mode="CPX") + assert "Nothing in the optimizer changes the mode" in verdict.refusal + + def test_a_match_proceeds(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(declared_mode="CPX", params={"peak_gib_per_stream": 4.0}) + assert verdict.ok is True + + def test_no_declaration_accepts_whatever_the_card_is_in(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(params={"peak_gib_per_stream": 4.0}) + assert verdict.ok is True + assert verdict.layout is not None + assert verdict.layout.mode == "CPX" + + +class TestUnpartitionedCard: + def test_spx_needs_no_feasibility_check(self, card): + """One partition is the whole card, which every session already assumes.""" + card(SPX_288) + verdict = ps.validate_session_shape(params={"peak_gib_per_stream": 200.0}, streams=4) + + assert verdict.ok is True + assert not any("Streams per partition" in note for note in verdict.notes) + + def test_the_shape_is_still_recorded(self, card): + card(SPX_288) + verdict = ps.validate_session_shape() + assert any("SPX" in note for note in verdict.notes) + + +class TestFeasibility: + def test_the_measured_mi355x_case_is_refused_at_launch(self, card): + """20.7 GiB per stream, two streams, a 36 GiB CPX partition.""" + card(CPX_36) + verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 20.7}) + + assert verdict.ok is False + assert "2 x 20.7 GiB = 41.4 GiB needed per partition" in verdict.refusal + assert "36.0 GiB available" in verdict.refusal + + def test_a_measured_refusal_says_it_was_measured(self, card): + """Measured and weights-only call for different responses from an operator.""" + card(CPX_36) + verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 20.7}) + assert "a measured per-stream peak" in verdict.refusal + + def test_the_same_workload_at_one_stream_is_allowed(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(streams=1, params={"peak_gib_per_stream": 20.7}) + assert verdict.ok is True + + def test_a_fit_reports_the_headroom_it_checked(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 10.0}) + + assert verdict.ok is True + assert any("20.0 GiB needed of 36.0 GiB" in note for note in verdict.notes) + + def test_streams_are_multiplied_in_not_ignored(self, card): + """Each stream holds its own copy, so the one-stream figure is the trap.""" + card(CPX_36) + assert ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 17.0}).ok is True + assert ps.validate_session_shape(streams=3, params={"peak_gib_per_stream": 17.0}).ok is False + + def test_an_unknown_footprint_warns_and_runs(self, card): + """Refusing here would ground every session that cannot be sized.""" + card(CPX_36) + verdict = ps.validate_session_shape(streams=2) + + assert verdict.ok is True + assert any("out-of-memory failure rather than a refusal" in w for w in verdict.warnings) + + def test_unreported_partition_memory_warns_and_runs(self, card): + card(_layout("CPX", 32, gib=None)) + verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 20.7}) + + assert verdict.ok is True + assert any("did not report its per-partition memory" in w for w in verdict.warnings) + + +class TestFitCheckNeedsAFanOut: + """A partitioned card met by a session that will not place work on partitions. + + The refusal multiplies a footprint by streams sharing one partition. With no + fan-out that premise is false twice over: nothing places a second stream, + and nothing pins the benchmark to a partition at all -- whole cards + enumerate first, so the run may land on a whole card. + """ + + def test_a_workload_too_big_for_a_partition_is_not_refused(self, card): + card(CPX_36) + verdict = ps.validate_session_shape( + streams=2, + params={"peak_gib_per_stream": 20.7}, + fanout_expected=False, + ) + assert verdict.ok is True + + def test_the_shape_is_still_recorded(self, card): + """Provenance is the point; withholding the refusal must not lose the mode.""" + card(CPX_36) + verdict = ps.validate_session_shape(fanout_expected=False) + + assert verdict.layout is not None + assert verdict.layout.mode == "CPX" + assert any("CPX" in note for note in verdict.notes) + + def test_it_says_the_numbers_belong_to_an_unknown_fraction_of_the_card(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(fanout_expected=False) + assert any("unknown fraction of the card" in w for w in verdict.warnings) + + def test_no_stream_note_is_made_for_streams_nothing_will_place(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(streams=4, fanout_expected=False) + assert not any("Streams per partition" in note for note in verdict.notes) + + def test_a_mode_mismatch_is_still_refused(self, card): + """The assertion is about the card, so no fan-out is needed to check it.""" + card(SPX_288) + verdict = ps.validate_session_shape(declared_mode="CPX", fanout_expected=False) + assert verdict.ok is False + + def test_an_unpartitioned_card_is_unaffected(self, card): + card(SPX_288) + assert ps.validate_session_shape(fanout_expected=False).ok is True + + +class TestTableFallbackIsSurfaced: + def test_a_derived_cu_count_warns_because_selection_matches_on_it(self, card): + card(_layout("DPX", 152, 144.0, probed=False)) + verdict = ps.validate_session_shape(params={"peak_gib_per_stream": 4.0}) + + assert verdict.ok is True + assert any("built-in board table" in w for w in verdict.warnings) + + def test_a_probed_count_warns_about_nothing(self, card): + card(CPX_36) + verdict = ps.validate_session_shape(params={"peak_gib_per_stream": 4.0}) + assert verdict.warnings == () + + +class TestPerStreamFootprint: + def test_an_explicit_param_is_taken_as_measured(self): + assert ps.per_stream_footprint_gib({"peak_gib_per_stream": 12.5}) == (12.5, "measured") + + def test_a_prior_run_supplies_it(self): + state = type("S", (), {"current_best": {"peak_gib_per_stream": 9.0}, "model_path": ""})() + assert ps.per_stream_footprint_gib(None, state) == (9.0, "measured") + + def test_a_param_overrides_a_prior_run(self): + state = type("S", (), {"current_best": {"peak_gib_per_stream": 9.0}, "model_path": ""})() + assert ps.per_stream_footprint_gib({"peak_gib_per_stream": 3.0}, state)[0] == 3.0 + + @pytest.mark.parametrize("bad", ["n/a", None, 0, -1, [1]]) + def test_an_unusable_reading_is_treated_as_unknown(self, bad): + """The contract is "refuse nothing when unknown", and malformed is unknown.""" + assert ps.per_stream_footprint_gib({"peak_gib_per_stream": bad}) == (0.0, "") + + def test_no_model_and_no_measurement_is_unknown(self): + assert ps.per_stream_footprint_gib({}, None) == (0.0, "") + + def test_an_unreadable_checkpoint_is_unknown_not_fatal(self, monkeypatch): + state = type("S", (), {"current_best": {}, "model_path": "/nonexistent/model"})() + assert ps.per_stream_footprint_gib(None, state) == (0.0, "") + + +class TestRuntimeEnv: + def test_publishes_the_shape_the_entrypoint_fans_out_on(self): + env = ps.runtime_env(CPX_36, 2) + assert env == { + ps.PARTITION_MODE_ENV: "CPX", + ps.PARTITION_COUNT_ENV: "8", + ps.PARTITION_CU_ENV: "32", + ps.PARTITION_STREAMS_ENV: "2", + ps.PARTITION_TOTAL_STREAMS_ENV: "16", + } + + def test_without_a_fan_out_the_topology_is_still_published(self): + """The platform fingerprint reads it back from here, on the crash path.""" + env = ps.runtime_env(CPX_36, 2, fanout=False) + assert env == { + ps.PARTITION_MODE_ENV: "CPX", + ps.PARTITION_COUNT_ENV: "8", + ps.PARTITION_CU_ENV: "32", + } + + def test_without_a_fan_out_no_concurrency_is_stated(self): + """Streams are directions to a benchmark that places work per partition.""" + env = ps.runtime_env(CPX_36, 2, fanout=False) + assert ps.PARTITION_STREAMS_ENV not in env + assert ps.PARTITION_TOTAL_STREAMS_ENV not in env + + def test_it_publishes_no_device_list(self): + """HIP enumerates whole cards before partitions; only the GPU process can check.""" + assert not any("DEVICE" in key or "VISIBLE" in key for key in ps.runtime_env(CPX_36, 2)) + + def test_total_streams_is_the_concurrency_a_fanned_out_entrypoint_should_drive(self): + assert ps.runtime_env(CPX_36, 3)[ps.PARTITION_TOTAL_STREAMS_ENV] == "24" + + +class TestSessionShapeSummary: + def test_an_unknown_shape_is_distinguishable_from_an_unpartitioned_one(self): + assert ps.session_shape_summary(None, 2)["mode"] == "" + assert ps.session_shape_summary(SPX_288, 2)["mode"] == "SPX" + + def test_it_records_whether_the_cu_count_was_probed(self): + assert ps.session_shape_summary(CPX_36, 2)["cu_probed"] is True + assert ps.session_shape_summary(_layout("DPX", 152, probed=False), 2)["cu_probed"] is False + + def test_it_is_json_safe(self): + import json + + json.dumps(ps.session_shape_summary(CPX_36, 2)) + + def test_it_records_whether_anything_will_fan_out(self): + assert ps.session_shape_summary(CPX_36, 2, fanout_expected=False)["fanout_expected"] is False + assert ps.session_shape_summary(CPX_36, 2)["fanout_expected"] is True + + +class TestReportedProvenance: + """What the report says about where the CU count came from. + + A false provenance line is the one failure this feature exists to prevent, + so the rendering gets the same scrutiny as the probe. + """ + + @staticmethod + def _render(shape): + from hyperloom.orchestrator.actions.executors.report import ( + _format_compute_partition_section, + ) + + return "\n".join(_format_compute_partition_section({"compute_partition": shape})) + + def test_a_probed_count_is_reported_as_probed(self): + assert "32 (from the device)" in self._render(ps.session_shape_summary(CPX_36, 2)) + + def test_a_derived_count_is_reported_as_derived(self): + shape = ps.session_shape_summary(_layout("CPX", 32, 36.0, probed=False), 2) + assert "32 (derived from the board table)" in self._render(shape) + + def test_an_unknown_provenance_claims_neither(self): + """A shape recovered from the published env knows the count, not its origin. + + Truthiness on an absent key read that as the board table, so a fresh + launch that probed the device reported a guess it never made. + """ + rendered = self._render({"mode": "CPX", "partitions": 8, "cu_per_partition": 32}) + + assert "CU per partition : 32" in rendered + assert "board table" not in rendered + assert "from the device" not in rendered + + def test_a_session_that_cannot_fan_out_is_told_the_figure_is_one_device(self): + shape = ps.session_shape_summary(CPX_36, 2, fanout_expected=False) + assert "does not place work on individual partitions" in self._render(shape) + + def test_a_session_that_cannot_fan_out_claims_no_stream_placement(self): + """A stated total concurrency above a paragraph denying it is a contradiction.""" + shape = ps.session_shape_summary(CPX_36, 2, fanout_expected=False) + assert "streams/partition" not in self._render(shape) + assert "streams/partition" in self._render(ps.session_shape_summary(CPX_36, 2)) + + def test_an_unpartitioned_card_renders_nothing(self): + assert self._render(ps.session_shape_summary(SPX_288, 2)) == "" + + +class TestRecordedShapeIsProvenanceNotADecision: + def test_the_topology_cannot_be_rewritten_by_update_state(self): + """Locked for the same reason as model_path: the report prints whatever it says.""" + from hyperloom.orchestrator.policy.gate import CORE_STATE_FIELDS + + assert "compute_partition" in CORE_STATE_FIELDS + + def test_the_published_env_is_a_lossy_subset_of_the_verdict(self): + """Why the seed is handed the verdict instead of re-reading the environment.""" + from hyperloom.common.gpu_partition import published_shape + + verdict = ps.session_shape_summary(CPX_36, 2) + from_env = published_shape(ps.runtime_env(CPX_36, 2)) or {} + + assert set(verdict) - set(from_env) == {"cu_probed", "gib_per_partition", "fanout_expected"} + + +class TestEnvReaders: + @pytest.mark.parametrize(("raw", "expected"), [("3", 3), ("-1", 0), ("", 0), ("x", 0)]) + def test_the_gpu_id_is_clamped_and_never_raises(self, monkeypatch, raw, expected): + monkeypatch.setenv(ps.PARTITION_GPU_ENV, raw) + assert ps.partition_gpu_id() == expected + + def test_the_module_exports_no_reader_without_a_caller(self): + """Two env readers shipped with only their own tests as callers; both are gone.""" + assert not hasattr(ps, "expected_mode") + assert not hasattr(ps, "streams_per_partition") + assert "EXPECTED_MODE_ENV" not in ps.__all__ + assert "STREAMS_PER_PARTITION_ENV" not in ps.__all__ + + +class TestFrameworkFanout: + """The finding this fixes: ``--framework`` defaults to ``None``.""" + + @pytest.mark.parametrize("scriptable", ["xdit", "custom", "XDiT"]) + def test_a_scriptable_framework_can_place_work_per_partition(self, scriptable): + from hyperloom.inference_optimizer.cli import _partition_fanout_supported + + assert _partition_fanout_supported(scriptable) == (True, "") + + @pytest.mark.parametrize("serving", ["vllm", "sglang", "atom"]) + def test_a_serving_framework_is_refused_with_a_reason(self, serving): + from hyperloom.inference_optimizer.cli import _partition_fanout_supported + + supported, detail = _partition_fanout_supported(serving) + assert supported is False + assert "runs a server" in detail + + @pytest.mark.parametrize("unresolved", [None, "", " "]) + def test_an_unresolved_framework_says_so_rather_than_passing_silently(self, unresolved): + """A truthiness guard here reads as "checked and fine" while meaning "not checked".""" + from hyperloom.inference_optimizer.cli import _partition_fanout_supported + + supported, detail = _partition_fanout_supported(unresolved) + assert supported is False + assert "not resolved yet" in detail + + +class TestCliExport: + """The launch entry point: validates, publishes, or exits 2.""" + + @pytest.fixture + def export(self, card, monkeypatch): + from hyperloom.inference_optimizer import cli + + monkeypatch.setattr(cli.os, "environ", dict(cli.os.environ), raising=False) + + def call(layout=CPX_36, **kw): + card(layout) + kw.setdefault("declared_mode", None) + kw.setdefault("streams_per_partition", None) + kw.setdefault("framework", "xdit") + return cli._export_partition_shape(**kw) + + return call + + def test_a_readable_card_is_published_without_any_flag(self, export): + """The shape is a measurement property, so it is recorded either way.""" + assert export()["mode"] == "CPX" + + def test_an_unreadable_card_without_a_flag_publishes_nothing(self, export): + assert export(layout=None) == {} + + def test_a_mode_mismatch_exits_two(self, export): + with pytest.raises(SystemExit) as exit_info: + export(layout=SPX_288, declared_mode="CPX") + assert exit_info.value.code == 2 + + def test_a_declared_mode_on_an_unreadable_card_exits_two(self, export): + with pytest.raises(SystemExit) as exit_info: + export(layout=None, declared_mode="CPX") + assert exit_info.value.code == 2 + + def test_a_misspelled_mode_exits_two_at_launch(self, export): + with pytest.raises(SystemExit) as exit_info: + export(declared_mode="OPX") + assert exit_info.value.code == 2 + + def test_an_infeasible_workload_exits_two(self, export, monkeypatch): + """The whole point: a refusal at launch instead of an OOM three hours in.""" + monkeypatch.setattr(ps, "per_stream_footprint_gib", lambda *a, **k: (20.7, "measured")) + with pytest.raises(SystemExit) as exit_info: + export(streams_per_partition=2) + assert exit_info.value.code == 2 + + def test_the_workload_is_actually_sized_at_launch(self, export, monkeypatch): + """Without the model path reaching the resolver, the check can only ever warn.""" + seen: dict = {} + + def spy(params=None, shared_state=None): + seen.update(params or {}) + return 0.0, "" + + monkeypatch.setattr(ps, "per_stream_footprint_gib", spy) + export(model_path="/models/flux", precision="bf16") + assert seen.get("model_path") == "/models/flux" + assert seen.get("precision") == "bf16" + + def test_a_resume_can_use_the_measured_peak(self, export, monkeypatch): + """A prior measurement rules out a mode the weights alone would fit.""" + state = type("S", (), {"current_best": {"peak_gib_per_stream": 20.7}, "model_path": "", "precision": ""})() + with pytest.raises(SystemExit) as exit_info: + export(streams_per_partition=2, shared_state=state) + assert exit_info.value.code == 2 + + @pytest.mark.parametrize("invalid", [0, -1]) + def test_a_non_positive_stream_count_exits_rather_than_defaulting(self, export, invalid): + """``0 or DEFAULT`` is DEFAULT, which would honour a mistake as the default.""" + with pytest.raises(SystemExit) as exit_info: + export(streams_per_partition=invalid) + assert exit_info.value.code == 2 + + def test_an_omitted_stream_count_takes_the_default(self, export): + assert export()["streams_per_partition"] == ps.DEFAULT_STREAMS_PER_PARTITION + + def test_a_resume_re_checks_rather_than_trusts_the_archive(self, export): + """A card can be repartitioned while a session is stopped.""" + from hyperloom.inference_optimizer import cli + + args = type("A", (), {"compute_partition_mode": None, "streams_per_partition": None})() + state = type("S", (), {"compute_partition": {"mode": "CPX", "streams_per_partition": 4}})() + cli._restore_partition_shape_from_state(args, state) + + assert (args.compute_partition_mode, args.streams_per_partition) == ("CPX", 4) + with pytest.raises(SystemExit): + export(layout=SPX_288, declared_mode=args.compute_partition_mode) + + def test_a_resume_flag_overrides_the_archive(self, export): + from hyperloom.inference_optimizer import cli + + args = type("A", (), {"compute_partition_mode": "DPX", "streams_per_partition": 1})() + state = type("S", (), {"compute_partition": {"mode": "CPX", "streams_per_partition": 4}})() + cli._restore_partition_shape_from_state(args, state) + + assert (args.compute_partition_mode, args.streams_per_partition) == ("DPX", 1) + + def test_a_resume_of_an_unpartitioned_session_restores_nothing(self, export): + from hyperloom.inference_optimizer import cli + + args = type("A", (), {"compute_partition_mode": None, "streams_per_partition": None})() + cli._restore_partition_shape_from_state(args, type("S", (), {"compute_partition": {}})()) + + assert (args.compute_partition_mode, args.streams_per_partition) == ("", None) + + def test_a_serving_session_on_a_split_card_is_not_refused_without_flags(self, export, monkeypatch): + """The reported bug: a plain sglang run exited 2 on a card someone else left in CPX. + + No flags, so no fan-out was ever asked for, and a serving benchmark + cannot do one. Refusing on ``2 x footprint`` was arithmetic about a + shape the session was never going to run in. + """ + monkeypatch.setattr(ps, "per_stream_footprint_gib", lambda *a, **k: (20.7, "weights")) + shape = export(framework="sglang") + + assert shape["mode"] == "CPX" + assert shape["fanout_expected"] is False + + def test_the_same_session_is_still_refused_when_the_operator_asks_for_the_shape(self, export, monkeypatch): + """Naming the flags asserts the shape, and an assertion is held to.""" + monkeypatch.setattr(ps, "per_stream_footprint_gib", lambda *a, **k: (20.7, "weights")) + with pytest.raises(SystemExit) as exit_info: + export(framework="sglang", streams_per_partition=2) + assert exit_info.value.code == 2 + + def test_a_scriptable_session_is_still_refused_without_flags(self, export, monkeypatch): + """Where the fan-out is real, the default of two streams is a real premise.""" + monkeypatch.setattr(ps, "per_stream_footprint_gib", lambda *a, **k: (20.7, "weights")) + with pytest.raises(SystemExit) as exit_info: + export(framework="xdit") + assert exit_info.value.code == 2 + + def test_the_runtime_handoff_is_published_only_where_something_reads_it(self, export): + from hyperloom.inference_optimizer import cli + + export(framework="xdit") + assert cli.os.environ.get(ps.PARTITION_MODE_ENV) == "CPX" + + def test_a_serving_framework_is_handed_no_fan_out_instruction(self, export): + """Stating a concurrency nothing will drive is the contradiction here.""" + from hyperloom.inference_optimizer import cli + + assert export(framework="sglang")["mode"] == "CPX" + assert ps.PARTITION_STREAMS_ENV not in cli.os.environ + assert ps.PARTITION_TOTAL_STREAMS_ENV not in cli.os.environ + + def test_a_serving_framework_still_publishes_the_topology(self, export): + """Otherwise the platform fingerprint loses the mode -- the whole point.""" + from hyperloom.inference_optimizer import cli + + export(framework="sglang") + assert cli.os.environ.get(ps.PARTITION_MODE_ENV) == "CPX" + + def test_the_summary_carries_what_the_env_cannot(self, export): + """The published variables are a lossy subset: no provenance, no memory.""" + shape = export(framework="xdit") + assert shape["cu_probed"] is True + assert shape["gib_per_partition"] == 36.0 + + def test_a_multi_node_session_records_no_shape(self, export): + """The card this process can read is not the card the benchmark ran on.""" + assert export(nodes=2) == {} + + def test_a_multi_node_session_publishes_nothing(self, export): + """Not even the topology: it would be the wrong node's.""" + from hyperloom.inference_optimizer import cli + + export(nodes=4) + assert ps.PARTITION_MODE_ENV not in cli.os.environ + + def test_a_multi_node_assertion_exits_rather_than_going_unchecked(self, export): + """Same rule as an unreadable card: unverifiable is not satisfied.""" + with pytest.raises(SystemExit) as exit_info: + export(nodes=2, declared_mode="CPX") + assert exit_info.value.code == 2 + + def test_an_explicit_zero_on_resume_still_reaches_the_guard(self, export): + """Falsiness here would read the mistake as "omitted" and paper over it.""" + from hyperloom.inference_optimizer import cli + + args = type("A", (), {"compute_partition_mode": None, "streams_per_partition": 0})() + state = type("S", (), {"compute_partition": {"mode": "CPX", "streams_per_partition": 4}})() + cli._restore_partition_shape_from_state(args, state) + + assert args.streams_per_partition == 0 + with pytest.raises(SystemExit) as exit_info: + export(streams_per_partition=args.streams_per_partition) + assert exit_info.value.code == 2 + + def test_a_stale_shape_cannot_be_inherited_from_the_shell(self, export, monkeypatch): + """A second session in the same shell must not adopt the first one's shape.""" + monkeypatch.setenv(ps.PARTITION_MODE_ENV, "CPX") + monkeypatch.setenv(ps.PARTITION_COUNT_ENV, "8") + export(layout=None) + assert ps.PARTITION_MODE_ENV not in ps.os.environ diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_shape.py b/src/hyperloom/orchestrator/actions/executors/_partition_shape.py new file mode 100644 index 0000000000..9e63483991 --- /dev/null +++ b/src/hyperloom/orchestrator/actions/executors/_partition_shape.py @@ -0,0 +1,398 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""The compute-partition shape a session runs in, checked once at launch. + +The mode is fixed for the session and established outside it. This module does +three things with that fact, and nothing else: + +* **Records it.** Whatever mode the card is in becomes part of the session's + platform fingerprint, so a number is never filed under a topology it was not + measured on. Two runs of the same configuration on the same card in SPX and in + CPX are different experiments, and without this they are indistinguishable in + the history. +* **Checks it.** A session given a mode whose partitions cannot hold the + workload is going to fail, and it is going to take hours to find out. The + arithmetic that says so costs milliseconds, so it runs at launch. +* **Publishes it.** The benchmark entrypoint already has to fan work out across + partitions. It gets the shape -- mode, partition count, CU per partition, + streams -- so it can place work and, crucially, verify what it is running on. + +What it deliberately does not do is change the mode. The card must already be in +its mode before ``optimize`` starts: this runs at launch, and the entrypoint that +places work across partitions does not start until the first benchmark, long +after the shape has been checked and recorded. So the apply belongs to the +operator or the provisioning platform, not to anything downstream of here and +not to an optimization loop running agent-authored code. Nothing here needs +privilege. + +The check fails closed on the one thing it can be wrong about. If the operator +declares an expected mode and the card cannot be read, the session is refused +rather than run: the declaration exists precisely to catch an external set that +did not take, and an unverifiable assertion is not a satisfied one. +""" + +from __future__ import annotations + +import logging +import os +from dataclasses import dataclass, field +from typing import Any + +from hyperloom.common.gpu_partition import ( + PARTITION_COUNT_ENV, + PARTITION_CU_ENV, + PARTITION_MODE_ENV, + PARTITION_STREAMS_ENV, + PARTITION_TOTAL_STREAMS_ENV, + PartitionLayout, + fits_in_partition, + observe_partition, +) + +log = logging.getLogger(__name__) + +#: GPU whose partition state describes this session. Single-card: this +#: session's numbers describe partitions of that one physical GPU. +PARTITION_GPU_ENV = "HYPERLOOM_PARTITION_GPU" + +#: Streams per partition when unset. Two is where every mode measured on MI355X +#: peaked: one leaves each partition idle through the fixed per-pass cost, a +#: second fills it, a third only adds queueing. A default rather than a constant +#: because it is a property of the workload, not of the hardware. +DEFAULT_STREAMS_PER_PARTITION = 2 + + +@dataclass(frozen=True) +class ShapeVerdict: + """The launch-time verdict on a session's partition shape. + + Attributes: + layout: The live topology, or ``None`` when it could not be read. + refusal: Why the session must not start. Empty when it may. + warnings: Things the operator should know that do not stop the run. + notes: Lines describing the shape, for the launch banner. + """ + + layout: PartitionLayout | None = None + refusal: str = "" + warnings: tuple[str, ...] = () + notes: tuple[str, ...] = field(default=()) + + @property + def ok(self) -> bool: + """Whether the session may proceed.""" + return not self.refusal + + +def partition_gpu_id() -> int: + """Return the GPU whose partition state describes this session (default 0).""" + try: + return max(0, int(os.environ.get(PARTITION_GPU_ENV, "0").strip() or 0)) + except ValueError: + return 0 + + +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. + + In practice today there is one source: the checkpoint's weight bytes. + + * **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 a + refusal needs, since it only ever acts on the former. + * **Measured.** ``peak_gib_per_stream`` would be the real footprint -- + weights, activations and workspace -- and is the only thing that could + rule out a mode the weights alone fit. Nothing in this repository writes + it: the branch below reads it from task params and from ``current_best`` + so a harness that starts reporting it is honoured without a change here, + but no in-tree producer exists, so the weights bound is what every + refusal is actually made on. Do not read the two bullets as a fallback + chain that is exercised. + + Args: + params: Task params, consulted for an explicit override. + shared_state: Live SharedState, consulted for a prior measurement and + the model identity. + + Returns: + ``(gib, source)``, or ``(0.0, "")`` when neither source can answer. + ``source`` names the origin for the message that reports a refusal, + 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 carrying 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 contract is "refuses 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, "" + # Params first, then state: the launch path has resolved flags but no + # SharedState yet, and it is the path where a refusal is worth the most. + precision = str((params or {}).get("precision") or getattr(shared_state, "precision", "") or "") + # Lazy: the kernel package pulls in the analytical stack, and this module is + # imported whether a partition shape is in play or not. + from hyperloom.orchestrator.kernel.roofline_ceiling import load_model_meta + + try: + meta = load_model_meta(model_path, precision_hint=precision) + 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 validate_session_shape( + *, + declared_mode: str = "", + streams: int = DEFAULT_STREAMS_PER_PARTITION, + gpu_type: str | None = None, + params: dict[str, Any] | None = None, + shared_state: Any = None, + gpu_id: int | None = None, + fanout_expected: bool = True, +) -> ShapeVerdict: + """Check, at launch, that this session can run in the shape it was given. + + Args: + declared_mode: The mode the operator asserts the card is in. Empty means + no assertion, in which case an unreadable card is merely unrecorded + rather than a refusal. + streams: Concurrent streams intended per partition. + gpu_type: Board name, used only if the CU probe fails. + params: Task params, for the footprint resolution. + shared_state: Live SharedState, for the footprint resolution. + gpu_id: GPU to interrogate; defaults to the configured one. + fanout_expected: Whether anything in this session will actually place + streams on partitions. ``False`` suppresses the footprint refusal, + because the question it answers does not arise: see below. + + Returns: + The verdict. ``refusal`` is non-empty only for a shape that cannot work: + a declared mode that does not match the card, and a workload whose + streams provably do not fit one partition. + + The footprint check is gated on ``fanout_expected`` rather than run + unconditionally because "N streams share one partition" is a premise, not a + fact about the card. Without a fan-out, nothing puts a second stream on a + partition, and worse, nothing pins the benchmark to a partition at all -- + HIP enumerates whole cards before partitions, so on a node where one card of + eight is split, device 0 is a *whole* card. Refusing such a session on + ``streams x footprint`` would be arithmetic about a shape it was never going + to run in. The mode is still observed, recorded and published; only the + refusal is withheld. + """ + warnings: list[str] = [] + notes: list[str] = [] + gpu = partition_gpu_id() if gpu_id is None else gpu_id + streams = max(1, int(streams or DEFAULT_STREAMS_PER_PARTITION)) + + layout = observe_partition(gpu, gpu_type=gpu_type) + if layout is None: + if declared_mode: + return ShapeVerdict( + refusal=( + f"--compute-partition-mode {declared_mode} was declared, but GPU {gpu}'s " + f"compute-partition state could not be read, so the claim cannot be " + f"checked. The flag asserts what the card is already in -- it does not " + f"set it. Drop the flag to run without the assertion." + ), + ) + # Nothing declared and nothing readable is the ordinary case on a host + # without amd-smi. The session runs exactly as it always has. + return ShapeVerdict() + + if declared_mode and layout.mode != declared_mode: + return ShapeVerdict( + layout=layout, + refusal=( + f"GPU {gpu} is in {layout.mode}, not the declared {declared_mode}. Nothing in " + f"the optimizer changes the mode, so this session would measure {layout.mode} " + f"and record it under a name the operator did not intend. Set the card to " + f"{declared_mode} before launching, or drop the flag." + ), + ) + + notes.append(f"Compute partitioning : {layout.describe()}") + if not layout.probed: + warnings.append( + f"GPU {gpu}'s CU count came from the built-in board table, not the device. " + f"Partition devices are selected by matching that count exactly, so if the " + f"table is wrong for this board the benchmark will find no partitions." + ) + + if not layout.partitioned: + # An unpartitioned card is the default and needs no further checking: + # one partition is the whole card, which is what every other session + # already assumes. + return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + + if not fanout_expected: + # Record the shape, refuse nothing. This is the card someone else left + # split, met by a session that will not place work per partition. + warnings.append( + f"GPU {gpu} is in {layout.mode}, but nothing in this session places work on " + f"individual partitions, so the shape is recorded and no per-partition memory " + f"check is made. Which device the benchmark lands on is not decided here -- HIP " + f"enumerates whole cards before partitions -- so treat the numbers as belonging " + f"to an unknown fraction of the card until the fan-out is known." + ) + return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + + notes.append(f"Streams per partition: {streams} ({streams * layout.partitions} total)") + + footprint_gib, source = per_stream_footprint_gib(params, shared_state) + if footprint_gib <= 0: + warnings.append( + f"Cannot size this workload against a {layout.mode} partition: the checkpoint's " + f"weight bytes could not be read. The session will run, but a partition too small " + f"for it will surface as an out-of-memory failure rather than a refusal here." + ) + return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + + if layout.gib_per_partition is None: + warnings.append( + f"GPU {gpu} did not report its per-partition memory, so the {footprint_gib:.1f} GiB " + f"per stream could not be checked against a {layout.mode} partition." + ) + return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + + needed = footprint_gib * streams + if not fits_in_partition(footprint_gib, layout, streams): + return ShapeVerdict( + layout=layout, + refusal=( + f"this workload does not fit GPU {gpu}'s {layout.mode} partitions: " + f"{streams} x {footprint_gib:.1f} GiB = {needed:.1f} GiB needed per partition, " + f"{layout.gib_per_partition:.1f} GiB available " + f"({layout.partitions} x {layout.cu_per_partition} CU). " + f"The {footprint_gib:.1f} GiB is " + + ( + "a measured per-stream peak." + if source == "measured" + else "the weights alone, so the real footprint is larger." + ) + + " Use a wider partition mode, fewer streams per partition, or a smaller model." + ), + warnings=tuple(warnings), + notes=tuple(notes), + ) + + notes.append( + f"Per-partition memory : {needed:.1f} GiB needed of {layout.gib_per_partition:.1f} GiB " + f"({streams} x {footprint_gib:.1f} GiB, from {source})" + ) + return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + + +def runtime_env(layout: PartitionLayout, streams: int, *, fanout: bool = True) -> dict[str, str]: + """Build the env published at launch for the shape this session runs in. + + Two readers, and they want different things, which is why ``fanout`` splits + the block rather than suppressing it: + + * **The provenance record.** ``platform_fingerprint`` reads the observed mode, + partition count and CU per partition back out of here, because it is + written on the crash path where spawning ``amd-smi`` is not acceptable. + Those three describe the card and are true whatever the benchmark does, so + they are always published -- withholding them is how a CPX number gets + filed as though it were SPX, which is the failure this module exists to + prevent. + * **The fan-out instruction.** Streams per partition and the total + concurrency are directions to a benchmark that places work on each + partition. Only a scriptable entrypoint does that, so publishing them to a + serving session would state a concurrency nothing was going to drive. + + The entrypoint is given the shape rather than a device list on purpose. HIP + enumerates whole cards before partitions, so an index list computed here + would be wrong in the one case that matters and wrong invisibly; the process + holding the GPU context is the one positioned to check a device's CU count + against :data:`PARTITION_CU_ENV` and refuse what does not match. + + Args: + layout: The observed topology. + streams: Streams to place on each partition. + fanout: Whether this session's benchmark places work per partition. + + Returns: + Env mapping for the benchmark process. + """ + streams = max(1, int(streams)) + env = { + PARTITION_MODE_ENV: layout.mode, + PARTITION_COUNT_ENV: str(layout.partitions), + PARTITION_CU_ENV: str(layout.cu_per_partition), + } + if fanout: + env[PARTITION_STREAMS_ENV] = str(streams) + env[PARTITION_TOTAL_STREAMS_ENV] = str(layout.partitions * streams) + return env + + +def session_shape_summary( + layout: PartitionLayout | None, + streams: int, + *, + fanout_expected: bool = True, +) -> dict[str, Any]: + """Summarize the session's partition shape for the report and the manifest. + + Args: + layout: The observed topology, or ``None`` when unknown. + streams: Streams placed on each partition. + fanout_expected: Whether this session's benchmark places work on each + partition. Recorded because it decides whether the throughput can be + an aggregate at all, which the report has to state and cannot infer. + + Returns: + A JSON-safe mapping. ``mode`` is empty when the shape is unknown, which + the report distinguishes from an unpartitioned card. + """ + if layout is None: + return {"mode": "", "partitions": 0, "cu_per_partition": 0, "streams_per_partition": 0} + return { + "mode": layout.mode, + "partitions": layout.partitions, + "cu_per_partition": layout.cu_per_partition, + "gib_per_partition": layout.gib_per_partition, + "streams_per_partition": max(1, int(streams)), + "cu_probed": layout.probed, + "fanout_expected": bool(fanout_expected), + } + + +__all__ = [ + "DEFAULT_STREAMS_PER_PARTITION", + "PARTITION_COUNT_ENV", + "PARTITION_CU_ENV", + "PARTITION_GPU_ENV", + "PARTITION_MODE_ENV", + "PARTITION_STREAMS_ENV", + "PARTITION_TOTAL_STREAMS_ENV", + "ShapeVerdict", + "partition_gpu_id", + "per_stream_footprint_gib", + "runtime_env", + "session_shape_summary", + "validate_session_shape", +] diff --git a/src/hyperloom/orchestrator/actions/executors/report.py b/src/hyperloom/orchestrator/actions/executors/report.py index 79236a6d14..0f9c72c3b3 100644 --- a/src/hyperloom/orchestrator/actions/executors/report.py +++ b/src/hyperloom/orchestrator/actions/executors/report.py @@ -568,6 +568,9 @@ 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 card's partition shape these numbers were measured in. A property + # of the session, not a result of it. + "compute_partition": dict(getattr(state, "compute_partition", None) or {}), } if external_baseline: summary["external_baseline"] = external_baseline @@ -723,6 +726,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: @@ -771,6 +775,71 @@ def _format_degraded_mode_section(summary: dict[str, Any]) -> list[str]: return lines +def _format_compute_partition_section(summary: dict[str, Any]) -> list[str]: + """State the compute-partition shape these numbers were measured in. + + Only rendered for a partitioned card. An unpartitioned card is what every + other report in the corpus describes, so saying so on all of them would be + noise; a split card is the exception that changes how the numbers compare, + and it says so where someone reading two reports side by side will see it. + + Args: + summary: The summary payload built by :func:`_build_summary_dict`. + + Returns: + Markdown lines, or ``[]`` when the card was whole or unknown. + """ + shape = summary.get("compute_partition") or {} + partitions = int(shape.get("partitions") or 0) + if not shape.get("mode") or partitions <= 1: + return [] + streams = int(shape.get("streams_per_partition") or 0) + lines = ["## Compute partitioning", ""] + lines.append( + f"This card was split into {partitions} partitions (`{shape['mode']}`), so these " + f"numbers are not comparable with an unpartitioned run of the same configuration." + ) + lines.append("") + lines.append(f"- mode : `{shape['mode']}` ({partitions} partitions)") + if shape.get("cu_per_partition"): + # Absent is its own answer. The published environment cannot carry the + # provenance flag, so a shape recovered from it knows the count but not + # where it came from -- and reporting that as the board table would be + # the exact false provenance this section exists to prevent. + probed = shape.get("cu_probed") + origin = "" if probed is None else (" (from the device)" if probed else " (derived from the board table)") + lines.append(f"- CU per partition : {shape['cu_per_partition']}{origin}") + if shape.get("gib_per_partition"): + lines.append(f"- HBM per partition : `{float(shape['gib_per_partition']):.1f}` GiB") + # Omitted where nothing fans out: the number would describe a placement that + # never happened, directly above a paragraph saying it did not. + if streams and shape.get("fanout_expected") is not False: + lines.append(f"- streams/partition : {streams} ({streams * partitions} concurrent streams total)") + lines.append("") + if shape.get("fanout_expected") is False: + lines.append( + f"**This session's benchmark does not place work on individual partitions.** The " + f"throughput is therefore one device's, not the total across all {partitions}, and " + f"which device it was is not recorded here -- whole cards enumerate before " + f"partitions, so it may be a whole card or a single partition." + ) + else: + lines.append( + f"Whether the throughput is one partition's or the total across all {partitions} " + f"depends on the benchmark placing work on each of them. The shape above is read " + f"from the card, but the fan-out is not this process's to do, and it cannot be " + f"verified from here -- so do not read the figure as an aggregate unless the " + f"benchmark entrypoint is known to fan out." + ) + lines.append("") + lines.append( + "Partitioning only ever gives a single stream fewer CUs, so per-request latency " + "is worse here than on the whole card by construction." + ) + 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/policy/gate.py b/src/hyperloom/orchestrator/policy/gate.py index 74d7d95413..378d13d82f 100644 --- a/src/hyperloom/orchestrator/policy/gate.py +++ b/src/hyperloom/orchestrator/policy/gate.py @@ -541,6 +541,12 @@ def _trace_path_allowlist() -> tuple[str, ...]: "model_path", "model_name", "model_class", + # The topology every number in the session was measured on, established + # once at launch from a read of the card. Locked for the same reason as + # model_path: it is provenance, not a decision, and a rewrite would file + # the results under a shape the card was never in -- silently, since the + # report prints whatever this says. + "compute_partition", "start_ts", # Where the current run leg begins; a forged value hands a previous # leg's CLOSE transition back the right to speak for this one. diff --git a/src/hyperloom/orchestrator/state/shared_state.py b/src/hyperloom/orchestrator/state/shared_state.py index 592f6cd721..135a5b0b56 100644 --- a/src/hyperloom/orchestrator/state/shared_state.py +++ b/src/hyperloom/orchestrator/state/shared_state.py @@ -697,6 +697,12 @@ class SharedState(_RenderMixin, _ExploreStateMixin): framework_repo_path: str = "" # ``HYPERLOOM_BENCHMARK_BACKEND`` at seed time (``bypass`` for custom). benchmark_backend: str = "" + # The card's compute-partition shape this session was measured in, as + # observed at launch: mode, partition count, CU and memory per partition, + # streams per partition. Empty when the card reported nothing. Part of the + # measurement contract, not a tuning knob -- the same configuration in SPX + # and in CPX is two different experiments. + compute_partition: dict[str, Any] = field(default_factory=dict) # ``--nodes``, feeding the robustness defaults and the IR-8 check. NOT the # cluster hand-off, which is resolved from argv before this state loads. nodes: int = 1 From 3044f8168537e761480e13291dfdfbe4b46efadc Mon Sep 17 00:00:00 2001 From: Rajesh Poornachandran Date: Fri, 28 Aug 2026 03:53:30 +0000 Subject: [PATCH 2/2] fix(platform): refuse a zero stream count, warn on an unusable partition GPU id Second review round on the compute-partition shape check. Still read-only, still no mutation anywhere. Two changed behaviours. validate_session_shape used `streams or DEFAULT`, so 0 silently became 2 and reported success for the same value the CLI exits 2 on; it now refuses, and refuses before the card is read, since a bad request needs no probe to judge. partition_gpu_id swallowed an unparseable HYPERLOOM_PARTITION_GPU and filed card 0's topology as the session's in silence; it now warns and names the consequence. The unknown-capacity guards in validate_session_shape and fits_in_partition asked one question with two different tests, so a zero capacity skipped the arithmetic *and* skipped the warning that explains why. Both now share PartitionLayout.capacity_known. Cleanup in the same pass: drops the dead UNPARTITIONED_MODE, takes the four probe helpers behind observe_partition out of __all__ so it describes the interface rather than the call graph, requires a layout in session_shape_summary instead of answering None with a second schema whose absent provenance key read as a positive claim, and retires the last two "measured peak, preferred when present" claims for a field nothing in this repository writes. Co-authored-by: Cursor --- CHANGELOG.md | 3 +- docs/reference/environment-variables.md | 5 +- src/hyperloom/common/gpu_partition.py | 52 ++++++--- .../common/tests/test_gpu_partition.py | 52 +++++++++ .../inference_optimizer/cli/__init__.py | 13 ++- .../inference_optimizer/cli/parser.py | 7 +- .../tests/test_partition_shape.py | 108 +++++++++++++++++- .../actions/executors/_partition_shape.py | 76 +++++++++--- 8 files changed, 273 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7852bb936b..bc0f5c457c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,7 +29,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), or if the card cannot be read — the flag exists to catch an external set that did not take, so an unverifiable assertion is treated as a failed one. `--streams-per-partition` (default `2`) is how many concurrent streams go on - each partition.
+ each partition; a value below `1` is refused rather than quietly replaced by + the default, since `0` is far more likely to be a mistake than a request.
**The optimizer does not change the mode.** Setting it is privileged and disrupts every process holding a GPU context, which is not something an optimization loop should do between benchmark rounds. The card must be in its diff --git a/docs/reference/environment-variables.md b/docs/reference/environment-variables.md index 823014ee62..58c82ec41c 100644 --- a/docs/reference/environment-variables.md +++ b/docs/reference/environment-variables.md @@ -449,7 +449,8 @@ Two CLI flags configure this, both optional: - `--streams-per-partition N` (default `2`) is how many concurrent streams the benchmark places on each partition. One stream leaves each partition idle through the fixed per-pass cost; beyond two, on the workloads measured so far, - only queueing is added. + only queueing is added. A value below `1` is refused rather than replaced by + the default, so `0` is a usage error instead of a silent `2`. At launch the per-stream HBM footprint is checked against one partition's memory. A workload that provably will not fit is refused in milliseconds instead @@ -470,7 +471,7 @@ declared mode there is a usage error rather than a silently unchecked assertion. | Variable | Default | Description | |----------|---------|-------------| -| `HYPERLOOM_PARTITION_GPU` | `0` | Which GPU's partition state describes this session. | +| `HYPERLOOM_PARTITION_GPU` | `0` | Which GPU's partition state describes this session. A value that is not a GPU id falls back to `0` with a warning rather than silently — the fallback reads a different card, and every number the session files afterwards would carry that card's shape. | ### Runtime hand-off diff --git a/src/hyperloom/common/gpu_partition.py b/src/hyperloom/common/gpu_partition.py index b0fdac8d0f..052ef19e90 100644 --- a/src/hyperloom/common/gpu_partition.py +++ b/src/hyperloom/common/gpu_partition.py @@ -75,10 +75,6 @@ "CPX": 8, } -#: The unpartitioned mode: one partition spanning the whole card. The only mode -#: every board supports, and what a card is in unless someone changed it. -UNPARTITIONED_MODE = "SPX" - #: The observed shape, published once at launch for the benchmark entrypoint to #: fan out across and for the provenance record to quote. Named here, in #: ``common``, because the platform fingerprint reads them and must not import @@ -124,6 +120,18 @@ def partitioned(self) -> bool: """Whether this mode splits the card at all.""" return self.partitions > 1 + @property + def capacity_known(self) -> bool: + """Whether :attr:`gib_per_partition` is a figure worth checking against. + + One predicate for the two places that ask, so an unreadable capacity + cannot be treated as unknown by the caller that reports it and as a real + limit by the arithmetic that acts on it. A non-positive capacity is + unknown rather than tiny: no partition has zero memory, so the reading + is wrong rather than restrictive. + """ + return self.gib_per_partition is not None and self.gib_per_partition > 0 + 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 "" @@ -491,6 +499,11 @@ def _int(key: str) -> int | None: def partition_device_predicate(cu_per_partition: int): """Return a predicate selecting partition devices by CU count. + **No caller in this repository.** It is the reference implementation of a + rule the out-of-tree benchmark entrypoint is most likely to get wrong, kept + beside the documentation that states the rule rather than left to be + re-derived downstream; do not go looking for the consumer. + 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. @@ -535,14 +548,30 @@ def fits_in_partition( streams_per_partition: Concurrent streams intended per partition. Returns: - ``True`` when they fit, or when capacity is unknown -- an unknown is + ``True`` when they fit, or when either figure is unknown -- an unknown is reported by the caller that has the number, not guessed at here. + + The unknown-capacity guard shares :attr:`PartitionLayout.capacity_known` + with :func:`~hyperloom.orchestrator.actions.executors._partition_shape.validate_session_shape`, + which tests it first so it can warn -- something a bool return cannot do. + That makes this guard unreachable from the in-tree caller by design, and it + stays because a predicate that silently multiplies by ``None`` for anyone + else is worse than one redundant test. """ - if not layout.gib_per_partition or required_gib <= 0: + if not layout.capacity_known or required_gib <= 0: return True - return required_gib * max(1, int(streams_per_partition)) <= layout.gib_per_partition - - + # ``capacity_known`` has already established this is a positive float; the + # coercion is only so the shared predicate can carry the decision without + # the arithmetic having to re-assert the type. + capacity_gib = float(layout.gib_per_partition or 0.0) + return required_gib * max(1, int(streams_per_partition)) <= capacity_gib + + +# The probe helpers behind observe_partition -- read_partition_mode(s), +# read_device_cu, read_device_gib -- are deliberately absent. observe_partition's +# docstring calls itself "the single entry point a caller needs", and listing the +# steps it takes here would contradict that: they are the call graph, not the +# interface. They stay importable for the tests that exercise each payload shape. __all__ = [ "MODE_PARTITION_COUNTS", "PARTITION_COUNT_ENV", @@ -550,7 +579,6 @@ def fits_in_partition( "PARTITION_MODE_ENV", "PARTITION_STREAMS_ENV", "PARTITION_TOTAL_STREAMS_ENV", - "UNPARTITIONED_MODE", "PartitionError", "PartitionLayout", "fits_in_partition", @@ -559,8 +587,4 @@ def fits_in_partition( "parse_mode", "partition_device_predicate", "published_shape", - "read_device_cu", - "read_device_gib", - "read_partition_mode", - "read_partition_modes", ] diff --git a/src/hyperloom/common/tests/test_gpu_partition.py b/src/hyperloom/common/tests/test_gpu_partition.py index 0c61a811e1..5be2f8efa3 100644 --- a/src/hyperloom/common/tests/test_gpu_partition.py +++ b/src/hyperloom/common/tests/test_gpu_partition.py @@ -260,6 +260,29 @@ def test_exactly_filling_a_partition_fits(self): layout = gp.layout_for("CPX", cu_per_partition=32, gib_per_partition=36.0) assert gp.fits_in_partition(18.0, layout, 2) is True + def test_a_zero_capacity_is_unknown_rather_than_a_limit(self): + """Shared with the caller's guard, which used to test ``is None`` instead.""" + layout = gp.PartitionLayout(mode="CPX", partitions=8, cu_per_partition=32, gib_per_partition=0.0) + assert layout.capacity_known is False + assert gp.fits_in_partition(999.0, layout, 2) is True + + +class TestCapacityKnown: + """One predicate for "is this capacity worth checking against". + + Two call sites asked the question with different tests -- ``is None`` in the + validator and falsiness in the arithmetic -- so a zero took opposite paths + through them. + """ + + @pytest.mark.parametrize( + ("gib", "known"), + [(36.0, True), (0.1, True), (None, False), (0.0, False), (-1.0, False)], + ) + def test_only_a_positive_capacity_is_known(self, gib, known): + layout = gp.PartitionLayout(mode="CPX", partitions=8, cu_per_partition=32, gib_per_partition=gib) + assert layout.capacity_known is known + class TestPartitionDevicePredicate: def test_matches_a_partition_and_rejects_a_whole_card(self): @@ -268,6 +291,35 @@ def test_matches_a_partition_and_rejects_a_whole_card(self): assert is_partition(32) is True assert is_partition(256) is False + def test_its_docstring_admits_it_has_no_in_tree_caller(self): + """Otherwise the next reader goes looking for the consumer and finds none.""" + assert "No caller in this repository" in (gp.partition_device_predicate.__doc__ or "") + + +class TestPublicSurface: + def test_the_dead_unpartitioned_mode_constant_is_gone(self): + """Defined and exported, referenced nowhere -- including by its own tests.""" + assert not hasattr(gp, "UNPARTITIONED_MODE") + assert "UNPARTITIONED_MODE" not in gp.__all__ + + @pytest.mark.parametrize( + "name", + ["read_partition_mode", "read_partition_modes", "read_device_cu", "read_device_gib"], + ) + def test_the_probe_helpers_behind_observe_partition_are_not_advertised(self, name): + """``observe_partition`` calls itself the single entry point a caller needs. + + Listing the four steps it takes contradicted that: they are its call + graph, not the module's interface. They stay importable for the tests + that exercise each amd-smi payload shape. + """ + assert name not in gp.__all__ + assert callable(getattr(gp, name)) + + def test_everything_advertised_exists(self): + for name in gp.__all__: + assert hasattr(gp, name), name + class TestPublishedShape: def test_reads_back_what_a_launch_published(self): diff --git a/src/hyperloom/inference_optimizer/cli/__init__.py b/src/hyperloom/inference_optimizer/cli/__init__.py index 11b4930b7f..5fbbb1cc69 100644 --- a/src/hyperloom/inference_optimizer/cli/__init__.py +++ b/src/hyperloom/inference_optimizer/cli/__init__.py @@ -1674,9 +1674,10 @@ def _export_partition_shape( model_path: Checkpoint to size the workload from. Without it the feasibility check has nothing to weigh and can only warn. precision: Resolved precision, which sets the bytes per weight. - shared_state: Persisted state on a resume, carrying any measured - per-stream peak. Tighter than the weight-bytes bound, so preferred - when present. + shared_state: Persisted state on a resume, consulted for the model + identity and for a ``peak_gib_per_stream`` that nothing in this + repository writes -- so in practice the footprint is the + weight-bytes bound either way. Returns: The published shape, or ``{}`` when this session has none. @@ -2151,8 +2152,10 @@ async def _run_optimize(args: argparse.Namespace) -> int: nodes=max(int(getattr(args, "nodes", 1) or 1), int(getattr(state, "nodes", 1) or 1)), model_path=state.model_path or str(getattr(args, "model", "") or ""), precision=state.precision or getattr(args, "precision", None), - # A resume can size the workload against what the last session - # actually measured, which rules out a mode the weights alone fit. + # Passed for the persisted model identity. It is also where a + # measured per-stream peak would be read from, but nothing writes + # one, so a resume sizes against the same weight-bytes bound as a + # fresh launch. shared_state=state, ) if state.compute_partition.get("mode"): diff --git a/src/hyperloom/inference_optimizer/cli/parser.py b/src/hyperloom/inference_optimizer/cli/parser.py index 79303018b4..50a76236b5 100644 --- a/src/hyperloom/inference_optimizer/cli/parser.py +++ b/src/hyperloom/inference_optimizer/cli/parser.py @@ -319,9 +319,10 @@ def _build_parser() -> argparse.ArgumentParser: "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 the workload has to fit, and the session refuses to start " - "when it provably will not. Only a scriptable framework's benchmark " - "places work per partition; passing this with a serving framework " - "warns, because nothing would act on it.", + "when it provably will not. A value below 1 is refused rather than " + "quietly replaced by the default. Only a scriptable framework's " + "benchmark places work per partition; passing this with a serving " + "framework warns, because nothing would act on it.", ) opt.add_argument( "--framework", diff --git a/src/hyperloom/inference_optimizer/tests/test_partition_shape.py b/src/hyperloom/inference_optimizer/tests/test_partition_shape.py index 3dfd44efa5..cd2f4e778a 100644 --- a/src/hyperloom/inference_optimizer/tests/test_partition_shape.py +++ b/src/hyperloom/inference_optimizer/tests/test_partition_shape.py @@ -168,7 +168,64 @@ def test_unreported_partition_memory_warns_and_runs(self, card): verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 20.7}) assert verdict.ok is True - assert any("did not report its per-partition memory" in w for w in verdict.warnings) + assert any("reported no usable per-partition memory" in w for w in verdict.warnings) + + def test_a_zero_capacity_is_unknown_rather_than_a_limit_that_refuses_everything(self, card): + """The guard here and the one in ``fits_in_partition`` must not disagree. + + A ``0.0`` capacity used to pass this side's ``is None`` test and then hit + the other side's falsiness test, so the arithmetic was skipped while the + warning that explains why was not printed. Both now ask + ``capacity_known``. + """ + card(_layout("CPX", 32, gib=0.0)) + verdict = ps.validate_session_shape(streams=2, params={"peak_gib_per_stream": 20.7}) + + assert verdict.ok is True + assert any("reported no usable per-partition memory" in w for w in verdict.warnings) + assert not any("GiB needed of" in n for n in verdict.notes) + + +class TestStreamsAreRefusedNotDefaulted: + """``streams=0`` must not become the default, here or at the CLI. + + The CLI already refuses it, and says in a comment why falsiness is the wrong + test. This entry point used ``streams or DEFAULT`` anyway, so the same value + exited 2 through one door and reported "Streams per partition: 2" through the + other. + """ + + @pytest.mark.parametrize("streams", [0, -1, -8]) + def test_a_value_below_one_refuses(self, card, streams): + card(CPX_36) + verdict = ps.validate_session_shape(streams=streams, params={"peak_gib_per_stream": 1.0}) + + assert verdict.ok is False + assert "must be >= 1" in verdict.refusal + assert str(streams) in verdict.refusal + + def test_it_refuses_before_the_card_is_read(self, monkeypatch): + """A usage error about the request needs no probe to decide.""" + + def _fail(*_args, **_kwargs): + raise AssertionError("the card must not be read to refuse streams=0") + + monkeypatch.setattr(ps, "observe_partition", _fail) + assert ps.validate_session_shape(streams=0).ok is False + + def test_omitted_streams_still_take_the_default(self, card): + """``None`` means "not named" and is the only thing that may default.""" + card(CPX_36) + for verdict in ( + ps.validate_session_shape(params={"peak_gib_per_stream": 1.0}), + ps.validate_session_shape(streams=None, params={"peak_gib_per_stream": 1.0}), + ): + assert verdict.ok is True + assert any(f"Streams per partition: {ps.DEFAULT_STREAMS_PER_PARTITION}" in n for n in verdict.notes) + + def test_a_non_numeric_request_refuses_rather_than_raises(self, card): + card(CPX_36) + assert ps.validate_session_shape(streams="two").ok is False # type: ignore[arg-type] class TestFitCheckNeedsAFanOut: @@ -293,10 +350,31 @@ def test_total_streams_is_the_concurrency_a_fanned_out_entrypoint_should_drive(s class TestSessionShapeSummary: - def test_an_unknown_shape_is_distinguishable_from_an_unpartitioned_one(self): - assert ps.session_shape_summary(None, 2)["mode"] == "" + def test_an_unknown_shape_is_absent_rather_than_a_second_schema(self): + """No ``layout is None`` branch: an unknown shape is ``{}`` at the caller. + + The branch that used to answer ``None`` returned four keys where the live + path returns seven, so ``cu_probed`` and ``fanout_expected`` were missing + rather than false -- and a missing provenance key is how the report came + to claim a board-table derivation it had not made. + """ + with pytest.raises(AttributeError): + ps.session_shape_summary(None, 2) # type: ignore[arg-type] assert ps.session_shape_summary(SPX_288, 2)["mode"] == "SPX" + def test_every_key_is_present_on_every_call(self): + expected = { + "mode", + "partitions", + "cu_per_partition", + "gib_per_partition", + "streams_per_partition", + "cu_probed", + "fanout_expected", + } + assert set(ps.session_shape_summary(SPX_288, 2)) == expected + assert set(ps.session_shape_summary(CPX_36, 1, fanout_expected=False)) == expected + def test_it_records_whether_the_cu_count_was_probed(self): assert ps.session_shape_summary(CPX_36, 2)["cu_probed"] is True assert ps.session_shape_summary(_layout("DPX", 152, probed=False), 2)["cu_probed"] is False @@ -382,6 +460,30 @@ def test_the_gpu_id_is_clamped_and_never_raises(self, monkeypatch, raw, expected monkeypatch.setenv(ps.PARTITION_GPU_ENV, raw) assert ps.partition_gpu_id() == expected + @pytest.mark.parametrize("raw", ["abc", "-1", "0x3", "2.5"]) + def test_an_unusable_gpu_id_is_warned_about_not_swallowed(self, monkeypatch, caplog, raw): + """Card 0's topology filed as the session's is the mislabelling this prevents. + + The reader returns 0 so a bad value cannot crash a launch, but it has to + say so: the fallback reads a different card, and every number the session + files afterwards carries that card's shape. + """ + monkeypatch.setenv(ps.PARTITION_GPU_ENV, raw) + with caplog.at_level("WARNING"): + assert ps.partition_gpu_id() == 0 + + messages = [r.getMessage() for r in caplog.records] + assert any(ps.PARTITION_GPU_ENV in m for m in messages) + assert any(repr(raw) in m for m in messages) + assert any("may not be this session's" in m for m in messages) + + @pytest.mark.parametrize("raw", ["0", "7", ""]) + def test_a_usable_gpu_id_says_nothing(self, monkeypatch, caplog, raw): + monkeypatch.setenv(ps.PARTITION_GPU_ENV, raw) + with caplog.at_level("WARNING"): + ps.partition_gpu_id() + assert caplog.records == [] + def test_the_module_exports_no_reader_without_a_caller(self): """Two env readers shipped with only their own tests as callers; both are gone.""" assert not hasattr(ps, "expected_mode") diff --git a/src/hyperloom/orchestrator/actions/executors/_partition_shape.py b/src/hyperloom/orchestrator/actions/executors/_partition_shape.py index 9e63483991..2b24f787b1 100644 --- a/src/hyperloom/orchestrator/actions/executors/_partition_shape.py +++ b/src/hyperloom/orchestrator/actions/executors/_partition_shape.py @@ -86,11 +86,29 @@ def ok(self) -> bool: def partition_gpu_id() -> int: - """Return the GPU whose partition state describes this session (default 0).""" + """Return the GPU whose partition state describes this session (default 0). + + An unusable value is warned about rather than quietly replaced. Reading card + 0 in silence is the exact mislabelling this module exists to prevent: the + session would file card 0's topology as its own while the benchmark ran on + the card the operator meant to name. + """ + raw = os.environ.get(PARTITION_GPU_ENV, "").strip() + if not raw: + return 0 try: - return max(0, int(os.environ.get(PARTITION_GPU_ENV, "0").strip() or 0)) + gpu = int(raw) except ValueError: + gpu = -1 + if gpu < 0: + log.warning( + "%s=%r is not a usable GPU id; reading compute partitioning from GPU 0 instead, " + "whose topology may not be this session's.", + PARTITION_GPU_ENV, + raw, + ) return 0 + return gpu def per_stream_footprint_gib( @@ -165,7 +183,7 @@ def per_stream_footprint_gib( def validate_session_shape( *, declared_mode: str = "", - streams: int = DEFAULT_STREAMS_PER_PARTITION, + streams: int | None = None, gpu_type: str | None = None, params: dict[str, Any] | None = None, shared_state: Any = None, @@ -178,7 +196,11 @@ def validate_session_shape( declared_mode: The mode the operator asserts the card is in. Empty means no assertion, in which case an unreadable card is merely unrecorded rather than a refusal. - streams: Concurrent streams intended per partition. + streams: Concurrent streams intended per partition. ``None`` means the + caller named none and takes the default; a value below one is + refused, not quietly replaced by it. Tested against ``None`` rather + than falsiness for that reason -- ``0 or DEFAULT`` is ``DEFAULT``, + which would honour the one value most likely to be a mistake. gpu_type: Board name, used only if the CU probe fails. params: Task params, for the footprint resolution. shared_state: Live SharedState, for the footprint resolution. @@ -204,8 +226,24 @@ def validate_session_shape( """ warnings: list[str] = [] notes: list[str] = [] + try: + requested_streams = DEFAULT_STREAMS_PER_PARTITION if streams is None else int(streams) + except (TypeError, ValueError): + requested_streams = 0 + if requested_streams < 1: + # Refused rather than defaulted, and refused before the card is touched: + # this is a usage error about the request, not a fact about the hardware, + # so it needs no probe to decide and must not differ from the CLI's own + # verdict on the same value. + return ShapeVerdict( + refusal=( + f"streams per partition must be >= 1, got {streams!r}. One stream per " + f"partition is the floor: a partition with no stream on it measures nothing, " + f"and a mode is only worth setting at two." + ), + ) + streams = requested_streams gpu = partition_gpu_id() if gpu_id is None else gpu_id - streams = max(1, int(streams or DEFAULT_STREAMS_PER_PARTITION)) layout = observe_partition(gpu, gpu_type=gpu_type) if layout is None: @@ -270,13 +308,18 @@ def validate_session_shape( ) return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) - if layout.gib_per_partition is None: + if not layout.capacity_known: + # Shares its predicate with fits_in_partition rather than testing + # `is None` here and falsiness there, which sent a zero capacity down + # opposite paths. Checked here, ahead of the arithmetic, because only + # this side can say so out loud. warnings.append( - f"GPU {gpu} did not report its per-partition memory, so the {footprint_gib:.1f} GiB " + f"GPU {gpu} reported no usable per-partition memory, so the {footprint_gib:.1f} GiB " f"per stream could not be checked against a {layout.mode} partition." ) return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) + capacity_gib = float(layout.gib_per_partition or 0.0) needed = footprint_gib * streams if not fits_in_partition(footprint_gib, layout, streams): return ShapeVerdict( @@ -284,7 +327,7 @@ def validate_session_shape( refusal=( f"this workload does not fit GPU {gpu}'s {layout.mode} partitions: " f"{streams} x {footprint_gib:.1f} GiB = {needed:.1f} GiB needed per partition, " - f"{layout.gib_per_partition:.1f} GiB available " + f"{capacity_gib:.1f} GiB available " f"({layout.partitions} x {layout.cu_per_partition} CU). " f"The {footprint_gib:.1f} GiB is " + ( @@ -299,7 +342,7 @@ def validate_session_shape( ) notes.append( - f"Per-partition memory : {needed:.1f} GiB needed of {layout.gib_per_partition:.1f} GiB " + f"Per-partition memory : {needed:.1f} GiB needed of {capacity_gib:.1f} GiB " f"({streams} x {footprint_gib:.1f} GiB, from {source})" ) return ShapeVerdict(layout=layout, warnings=tuple(warnings), notes=tuple(notes)) @@ -350,26 +393,29 @@ def runtime_env(layout: PartitionLayout, streams: int, *, fanout: bool = True) - def session_shape_summary( - layout: PartitionLayout | None, + layout: PartitionLayout, streams: int, *, fanout_expected: bool = True, ) -> dict[str, Any]: """Summarize the session's partition shape for the report and the manifest. + A layout is required. An unknown shape is ``{}`` at the call site, not a + record of zeroes here: this used to accept ``None`` and answer with a + four-key mapping missing ``cu_probed``, ``gib_per_partition`` and + ``fanout_expected``, so a consumer that met it saw a second schema for the + same field -- one whose absent provenance key reads as a positive claim. + Args: - layout: The observed topology, or ``None`` when unknown. + layout: The observed topology. streams: Streams placed on each partition. fanout_expected: Whether this session's benchmark places work on each partition. Recorded because it decides whether the throughput can be an aggregate at all, which the report has to state and cannot infer. Returns: - A JSON-safe mapping. ``mode`` is empty when the shape is unknown, which - the report distinguishes from an unpartitioned card. + A JSON-safe mapping, always carrying every key. """ - if layout is None: - return {"mode": "", "partitions": 0, "cu_per_partition": 0, "streams_per_partition": 0} return { "mode": layout.mode, "partitions": layout.partitions,