diff --git a/docs/components/geak.md b/docs/components/geak.md index fb2ea52f5d..b1a7f457f9 100644 --- a/docs/components/geak.md +++ b/docs/components/geak.md @@ -38,6 +38,67 @@ GEAK's `e2e_workflow` recursively drives `kernel_workflow` to author and tune th individual hot kernels worth fixing. See [Hyperloom optimization loop](../conceptual/optimization-loop.md). +## GPU pinning in the handoff + +GEAK launches full servers out-of-process (baseline, profile, config-tuning +validation) and writes a visible-devices mask for each one, so the handoff has +to say which cards the run owns. Two fields carry that, in two different +coordinate systems: + +| Field | Coordinate system | Value | +|-------|-------------------|-------| +| `gpu_ids` | HIP-level device list — HIP indexes into the ROCr-visible set | logical positions inside an *inherited* ROCr-level mask, capped at `tp` and at the mask width (`ROCR=6` → `"0"`); a HIP-level mask nested inside that slice is already logical and is forwarded instead (`ROCR=4,5,6,7` + `HIP=2,3` → `"2,3"`, i.e. cards 6 and 7); any other mask passes through uncapped (`HIP=4,5` → `"4,5"`); `0..tp-1` when the run is unpinned. Never empty — a falsy `gpu_ids` sends GEAK back to its own `0..tp-1` fallback | +| `gpu_ids_space` | — | `"logical"` when `gpu_ids` indexes into an inherited ROCr mask, `"absolute"` when they are whole-machine ids, `"none"` when the mask is set but empty. The one field that makes the coordinate systems distinguishable from the payload alone | +| `gpu_pin` | absolute device ids | `{"var", "value", "ids", "count", "source"}` for the winning mask, plus `"inner"` (the same shape) when a HIP-level mask is nested inside a ROCr-level one. Omitted entirely only when no mask is set anywhere, which means "whole machine visible", not "pinned to card 0"; a mask that is *set but empty* is reported with `count: 0` (zero devices visible). `value` is the mask as authored, only whitespace-trimmed and (for a YAML sequence) comma-joined, so a UUID mask can be re-exported as-is; `ids` are its numeric ids and `count` how many devices it exposes — both derived from the same effective token list, with duplicate and negative ordinals dropped, so `count >= len(ids)` always and they differ only for a (partly) non-numeric mask | + +The mask is resolved variable-major, over the full ROCm precedence chain in +`common/visible_devices.py`: the ROCr-level masks (`ROCR_VISIBLE_DEVICES`, then +its legacy spelling `HSA_VISIBLE_DEVICES`) before the HIP-level ones +(`HIP_VISIBLE_DEVICES`, `CUDA_VISIBLE_DEVICES`, `GPU_DEVICE_ORDINAL`). Within +each variable the process environment comes before the materialized baseline +recipe's `benchmark.envs`. + +That module is also the single definition of the tuple and the mask parser that +`orchestrator/bus/gpu_pool.py`, `orchestrator/policy/gate.py`, +`actions/executors/_ray_serving.py` and `common/env_safety.py` use. Those layers +read the narrower `COUNTING_VISIBLE_DEVICE_VARS` on purpose — they answer "how +many GPUs does this process have", and widening that would change GPU accounting +repo-wide — while this resolver answers "where is this run pinned", and a run +pinned with a legacy spelling is really pinned. + +The process env comes first because the recipe's ROCR key is not evidence of a +pin: `materialize_config_with_envs` autofills `ROCR_VISIBLE_DEVICES=0..tp-1` +into every materialized recipe when the mask is absent or narrower than `TP`. +A recipe ROCR value byte-identical to that default is therefore ignored, so a +`HIP`-pinned or genuinely unpinned run is not silently re-pinned to cards +`0..tp-1`. A recipe mask that differs from the default *is* honoured — but its +ids are forwarded absolute, not logical, because the GEAK child inherits the +process environment and never sees that mask. + +`tp` in the handoff is read from the same resolved recipe as `gpu_ids`, so the +two cannot disagree when the materializer clamps `TP` to the visible GPU count. + +A consumer that re-exports `HIP_VISIBLE_DEVICES` and leaves ROCr alone should +use `gpu_ids`; that is correct in both coordinate systems, because the child +inherits the same ROCr mask this process runs under. + +A consumer that writes `ROCR_VISIBLE_DEVICES` itself must use +`gpu_pin["value"]` — writing `gpu_ids` there resets the child to physical card +0 regardless of the run's pin — and must then renumber: it has just made the +device set `0..count-1` from the child's point of view, so the inner HIP mask +is `0..count-1`, **not** `gpu_ids`. Re-applying an absolute `gpu_ids` on top of +a ROCr mask it also wrote yields out-of-range ordinals (`ROCR=4,5` plus +`HIP=4,5` indexes positions 4 and 5 of a two-element set). `gpu_ids_space` is +how a consumer tells the cases apart without inferring it from `source`. + +`gpu_ids_space: "none"` is the third case: the mask is set but empty, so the run +has no visible devices and no id list can be truthful. `gpu_ids` still carries +`0..tp-1` — a falsy `gpu_ids` is read as "unset" and falls back to exactly those +ids anyway (`interface/run_e2e.py`), so an empty string would buy nothing and +lose the ability to say why. A consumer must treat those ids as placeholders and +not launch on them; the orchestrator also logs the condition at ERROR, because a +GEAK run that dies on an invalid device ordinal is otherwise unexplainable. + ## GEAK documentation For detailed documentation on GEAK, see [GEAK on ROCm Docs](https://rocm.docs.amd.com/projects/geak/en/latest/). diff --git a/src/hyperloom/common/env_safety.py b/src/hyperloom/common/env_safety.py index 53c0a28b93..9730ea4911 100644 --- a/src/hyperloom/common/env_safety.py +++ b/src/hyperloom/common/env_safety.py @@ -10,6 +10,8 @@ from __future__ import annotations +from hyperloom.common.visible_devices import GPU_MASK_ENV_NAMES as _GPU_MASK_ENV_NAMES + import os import re from collections.abc import Mapping @@ -208,15 +210,8 @@ ) # GPU visibility masks: setting one selects the hardware rather than tuning it. -GPU_MASK_ENV_NAMES: frozenset[str] = frozenset( - { - "CUDA_VISIBLE_DEVICES", - "GPU_DEVICE_ORDINAL", - "HIP_VISIBLE_DEVICES", - "HSA_VISIBLE_DEVICES", - "ROCR_VISIBLE_DEVICES", - } -) +# Single definition in ``hyperloom.common.visible_devices``. +GPU_MASK_ENV_NAMES = _GPU_MASK_ENV_NAMES # Env names an untrusted external source (reference recipe, framework-switch # manifest) may never set: shell-unsafe vars plus the workload/benchmark keys the diff --git a/src/hyperloom/common/visible_devices.py b/src/hyperloom/common/visible_devices.py new file mode 100644 index 0000000000..7abec2a4e1 --- /dev/null +++ b/src/hyperloom/common/visible_devices.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +"""Visible-device mask names and parsing, shared by every layer that reads one. + +The same tuple of mask variables and the same "split on ``,``/``;``, keep the +first occurrence of each non-negative int" parser had accumulated five separate +copies (``bus/gpu_pool``, ``policy/gate``, ``actions/executors/_ray_serving``, +``common/env_safety``, ``loop/coordinator_helpers``), and their empty-mask +semantics had already drifted apart. This module is the single definition; it +imports nothing outside the standard library so the pure-helper layers can use +it without dragging in the SQLite connection ``gpu_pool`` owns. + +Two var tuples, deliberately distinct: + +* :data:`COUNTING_VISIBLE_DEVICE_VARS` — what ``gpu_pool`` / ``gate`` consult to + answer "how many GPUs does this process have". Identical to what those layers + always used; widening it would change GPU accounting repo-wide. It is DERIVED + from the chain below by subtracting an explicit exclusion set, so the two can + never drift: a new var is counted unless someone names it as uncounted, and a + test asserts every chain member is classified. +* :data:`VISIBLE_DEVICE_VARS` — the full ROCm pin-resolution chain, used when + answering "where is this run pinned". ``HSA_VISIBLE_DEVICES`` is ROCr's legacy + name and ``GPU_DEVICE_ORDINAL`` is the legacy HIP-level filter; a run pinned + with either is really pinned, and omitting them left it reported as unpinned. + +The ROCr-level and HIP-level groups are exposed separately because the +distinction is load-bearing: a ROCr-level mask renumbers the devices the child +sees (so ids inside it are LOGICAL), while a HIP-level mask indexes into +whatever ROCr already exposed (so its ids are absolute unless a ROCr mask is +also in force). +""" + +from __future__ import annotations + +from typing import Any + +__all__ = [ + "COUNTING_VISIBLE_DEVICE_VARS", + "GPU_MASK_ENV_NAMES", + "HIP_LEVEL_VARS", + "ROCR_LEVEL_VARS", + "VISIBLE_DEVICE_VARS", + "effective_mask_tokens", + "is_rocr_level", + "mask_tokens", + "parse_device_list", +] + +#: ROCr-level masks: these slice the device set and renumber it ``0..N-1``. +#: ``HSA_VISIBLE_DEVICES`` is the legacy spelling; ``ROCR_VISIBLE_DEVICES`` +#: wins when both are set. +ROCR_LEVEL_VARS: tuple[str, ...] = ( + "ROCR_VISIBLE_DEVICES", + "HSA_VISIBLE_DEVICES", +) + +#: HIP-level masks: these index INTO whatever ROCr exposed. +HIP_LEVEL_VARS: tuple[str, ...] = ( + "HIP_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", +) + +#: Full pin-resolution precedence: ROCr level before HIP level, canonical +#: spelling before legacy within each level. +VISIBLE_DEVICE_VARS: tuple[str, ...] = ROCR_LEVEL_VARS + HIP_LEVEL_VARS + +#: Vars the capacity-counting layers deliberately do NOT read. +#: +#: Both are legacy aliases of a var already in the counting set, and both are +#: honoured only when their modern spelling is absent — so counting them would +#: not find a GPU the modern spelling missed, it would only change the answer on +#: hosts that happen to export the legacy name. That is a repo-wide GPU +#: accounting change, not a bugfix, so it stays out until someone makes it +#: deliberately. +_UNCOUNTED_VISIBLE_DEVICE_VARS: frozenset[str] = frozenset( + { + "HSA_VISIBLE_DEVICES", + "GPU_DEVICE_ORDINAL", + } +) + +#: The subset the capacity-counting layers read — ``gpu_pool``, ``policy.gate``, +#: ``actions.executors._ray_serving``. DERIVED from +#: :data:`VISIBLE_DEVICE_VARS` rather than re-listed, so a var added to the +#: precedence chain is counted by default and can only be left out by naming it +#: in :data:`_UNCOUNTED_VISIBLE_DEVICE_VARS`. The two tuples cannot silently +#: drift apart: ``test_visible_devices.py`` asserts every chain member is +#: classified exactly once. +COUNTING_VISIBLE_DEVICE_VARS: tuple[str, ...] = tuple( + var for var in VISIBLE_DEVICE_VARS if var not in _UNCOUNTED_VISIBLE_DEVICE_VARS +) + +#: Every name that selects hardware rather than tuning it (superset of the +#: precedence chain), for env scrubbing. +GPU_MASK_ENV_NAMES: frozenset[str] = frozenset(VISIBLE_DEVICE_VARS) + + +def is_rocr_level(var: str) -> bool: + """Does ``var`` slice and renumber the device set (rather than index it)? + + Args: + var: A visible-devices env var name. + + Returns: + ``True`` for the ROCr-level masks, whose member ids are logical + positions from the child's point of view. + """ + return str(var or "") in ROCR_LEVEL_VARS + + +def mask_tokens(raw: Any) -> list[str]: + """Split a visible-devices mask into its device tokens. + + Tokens are NOT required to be numeric: ROCm accepts GPU UUID masks + (``ROCR_VISIBLE_DEVICES=GPU-a1b2c3,GPU-d4e5f6``), and those still say how + MANY devices the child will see, which is all the logical-index arithmetic + needs. + + Args: + raw: A ``,``/``;``-separated mask, or an already-parsed YAML sequence. + + Returns: + Non-empty tokens in order, duplicates preserved. + """ + if isinstance(raw, (list, tuple)): + parts = [str(p) for p in raw] + else: + parts = str(raw if raw is not None else "").replace(";", ",").split(",") + return [tok for tok in (p.strip() for p in parts) if tok] + + +def _is_negative_ordinal(tok: str) -> bool: + """Is ``tok`` a negative device ordinal, i.e. a token that names no device? + + Written as a positive test rather than ``int(tok) < 0`` in a ``try`` so the + non-numeric case (a GPU UUID) is an ordinary ``False`` rather than a + swallowed ``ValueError``. + + Args: + tok: One already-stripped mask token. + + Returns: + ``True`` only for a leading ``-`` followed by digits. + """ + return tok.startswith("-") and tok[1:].isdigit() + + +def effective_mask_tokens(raw: Any) -> list[str]: + """The devices a mask actually exposes, in the order the runtime sees them. + + :func:`mask_tokens` is the LITERAL split; this is the *effective* set. ROCm + exposes ``ROCR_VISIBLE_DEVICES="3,3,2"`` as two devices, not three, and + drops a negative ordinal. Deriving both the device COUNT and the forwarded + id list from this one function is what keeps them from disagreeing: counting + literal tokens inflates the count (logical index 2 of a 2-device set is + invalid), while re-serializing from the parsed ints deflates it. + + Non-numeric tokens are kept — a UUID mask names real devices — so this + cannot filter out a genuinely invalid non-numeric entry; it removes only the + two forms that are unambiguously not extra devices. + + Args: + raw: A ``,``/``;``-separated mask, or an already-parsed YAML sequence. + + Returns: + Tokens with duplicates and negative ordinals removed, first-seen order. + """ + out: list[str] = [] + for tok in mask_tokens(raw): + if tok in out or _is_negative_ordinal(tok): + continue + out.append(tok) + return out + + +def parse_device_list(raw: Any) -> list[int]: + """Parse a visible-devices mask into absolute NUMERIC device ids. + + Args: + raw: A ``,``/``;``-separated mask (``"4,5,6,7"``) or a YAML sequence; + ``None`` and malformed entries are tolerated. + + Returns: + Unique non-negative ids in first-seen order; ``[]`` for an empty mask + and for a well-formed but non-numeric one (e.g. a UUID mask), which is + why callers that need a device COUNT must use + :func:`effective_mask_tokens`. + """ + out: list[int] = [] + for tok in effective_mask_tokens(raw): + try: + out.append(int(tok)) + except ValueError: + continue + return out diff --git a/src/hyperloom/inference_optimizer/breakdown/collectors/geak.py b/src/hyperloom/inference_optimizer/breakdown/collectors/geak.py index 761e2a0055..ae6a67a958 100644 --- a/src/hyperloom/inference_optimizer/breakdown/collectors/geak.py +++ b/src/hyperloom/inference_optimizer/breakdown/collectors/geak.py @@ -11,6 +11,7 @@ from __future__ import annotations +from collections.abc import Mapping from datetime import datetime, timezone from pathlib import Path from typing import Any @@ -614,6 +615,34 @@ def _geak_accepted_kernels_from_integrate_results( return accepted +def _handoff_gpu_fields(handoff: Mapping[str, Any] | None) -> dict[str, Any]: + """Pull the device set + absolute GPU pin out of a handoff, if it has them. + + A GEAK baseline that reads ``no_gain``/``incomplete`` because its servers + landed on a foreign tenant's card is otherwise indistinguishable from a + real result, so the breakdown records what the handoff told GEAK about the + run's cards (issue #1312). + + Keys are inserted only when present, mirroring the writer's ``if gpu_pin:`` + guard. An explicit ``null`` would conflate three different things: a + genuinely unpinned run, a pin that resolved empty, and a pre-v3 handoff on + disk (still produced by any session resumed from before this change). + + Args: + handoff: The parsed ``geak/handoff.json``, or ``None``. + + Returns: + A mapping with ``gpu_ids`` / ``gpu_ids_space`` / ``gpu_pin`` for + whichever keys the handoff carries; ``{}`` when it carries none. + """ + out: dict[str, Any] = {} + for key in ("gpu_ids", "gpu_ids_space", "gpu_pin"): + val = (handoff or {}).get(key) + if val is not None: + out[key] = val + return out + + def _geak_reconstruct_from_disk( session_dir: Path, warnings: list[str], @@ -671,6 +700,7 @@ def _load_json(p: Path) -> dict[str, Any]: "workload": handoff.get("workload"), "accepted_flags": handoff.get("accepted_flags"), "raw_baseline_tput": _to_float(handoff.get("raw_baseline_tput")), + **_handoff_gpu_fields(handoff), } # 2) a flushed-but-unpromoted result.json (absent or non-ok status). @@ -1003,9 +1033,30 @@ def _rel_if_under(p: Any) -> Any: else: accepted_kernels = [] + # The cards GEAK was told to use. Read from the handoff on disk rather than + # from ``geak_result``, which never carried them — and recorded on THIS + # path, not just the crash-recovery one, because the outcome that needs + # disambiguating (`no_gain`) is a completed run. + # An absent handoff is the normal shape for a run that never reached the + # handoff write (and for every pre-v3 session on disk), so it must not cost + # a stat+read or raise a warning — only a handoff that EXISTS and cannot be + # parsed is worth reporting. + _handoff_path = session_dir / "geak" / "handoff.json" + _gpu_fields: dict[str, Any] = {} + if _handoff_path.is_file(): + _gpu_fields = _handoff_gpu_fields( + read_json( + _handoff_path, + default={}, + require_dict=True, + on_error=lambda exc: warnings.append(f"geak: handoff read failed: {exc}"), + ) + ) + section: dict[str, Any] = { "engaged": True, "status": status, + **_gpu_fields, # Failure provenance (None on success). "error_class": result.get("error_class"), "error": result.get("error"), diff --git a/src/hyperloom/inference_optimizer/tests/test_geak_handoff_gpu_pin.py b/src/hyperloom/inference_optimizer/tests/test_geak_handoff_gpu_pin.py new file mode 100644 index 0000000000..d90a3b85ef --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_geak_handoff_gpu_pin.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +"""GPU-pin forwarding in the GEAK handoff (issue #1312). + +GEAK launches full servers out-of-process and writes a visible-devices mask for +each one. When the handoff carries no pin it falls back to ``0..tp-1``, so every +server lands on physical GPU 0 no matter where the run was pinned — on a shared +host that collides with a foreign tenant and the resulting OOM reads like a real +regression. + +These tests guard both halves of the contract: + +* ``gpu_ids`` stays in the coordinate system the consumer applies it in (HIP + indexes into the ROCr-visible set), so existing pins keep working; +* ``gpu_pin`` carries the ABSOLUTE mask plus the variable it came from, so a + consumer that writes ``ROCR_VISIBLE_DEVICES`` re-applies the pin instead of + resetting the child to card 0. +""" + +from __future__ import annotations + +import pytest + +from hyperloom.orchestrator.loop.coordinator_helpers import ( + _coerce_tp, + _is_autofilled_rocr, + _parse_device_list, + _resolve_gpu_pin, + _resolve_handoff_gpu_ids, + _resolve_handoff_gpu_ids_space, + _resolve_handoff_tp, +) +from hyperloom.common.visible_devices import VISIBLE_DEVICE_VARS + + +def _autofilled(tp: int) -> dict[str, object]: + """The ``benchmark.envs`` every materialized recipe carries. + + ``materialize_config_with_envs`` writes ``ROCR_VISIBLE_DEVICES=0..tp-1`` + unconditionally when the mask is absent or narrower than TP, so this shape + — not an empty mapping — is what the resolver sees in production. + """ + return {"TP": tp, "ROCR_VISIBLE_DEVICES": ",".join(str(i) for i in range(tp))} + + +_MASK_VARS = VISIBLE_DEVICE_VARS + + +@pytest.fixture(autouse=True) +def _clear_masks(monkeypatch: pytest.MonkeyPatch) -> None: + """Run every case against a known-unpinned environment.""" + for var in _MASK_VARS: + monkeypatch.delenv(var, raising=False) + + +# --------------------------------------------------------------------------- # +# _parse_device_list +# --------------------------------------------------------------------------- # + + +def test_parse_device_list_forms() -> None: + assert _parse_device_list("4,5,6,7") == [4, 5, 6, 7] + assert _parse_device_list(" 6 ") == [6] + assert _parse_device_list("0;1") == [0, 1] + assert _parse_device_list("3,3,2") == [3, 2] + + +def test_parse_device_list_tolerates_junk_and_empty() -> None: + assert _parse_device_list("") == [] + assert _parse_device_list(None) == [] + assert _parse_device_list("a,,-1,2") == [2] + + +# --------------------------------------------------------------------------- # +# _resolve_gpu_pin +# --------------------------------------------------------------------------- # + + +def test_pin_unset_everywhere_is_empty() -> None: + """No mask anywhere means "whole machine visible", NOT "pinned to 0".""" + assert _resolve_gpu_pin(recipe_envs={}, environ={}) == {} + + +def test_pin_from_process_rocr() -> None: + """The case issue #1312 hit: ROCm's canonical mask, previously ignored.""" + out = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "7"}) + assert out == { + "var": "ROCR_VISIBLE_DEVICES", + "value": "7", + "ids": [7], + "count": 1, + "source": "process_env", + } + + +def test_pin_prefers_rocr_over_hip_and_cuda() -> None: + env = { + "CUDA_VISIBLE_DEVICES": "0", + "HIP_VISIBLE_DEVICES": "1", + "ROCR_VISIBLE_DEVICES": "4,5", + } + out = _resolve_gpu_pin(recipe_envs={}, environ=env) + assert out["var"] == "ROCR_VISIBLE_DEVICES" + assert out["ids"] == [4, 5] + + +def test_pin_prefers_process_env_over_recipe() -> None: + """The process mask is the one the GEAK child actually inherits.""" + out = _resolve_gpu_pin( + recipe_envs={"TP": 1, "ROCR_VISIBLE_DEVICES": "6"}, + environ={"ROCR_VISIBLE_DEVICES": "3"}, + ) + assert out["source"] == "process_env" + assert out["ids"] == [3] + + +def test_pin_uses_recipe_when_the_process_is_unmasked() -> None: + """A hand-authored recipe mask is still a pin when nothing else says otherwise.""" + out = _resolve_gpu_pin(recipe_envs={"TP": 2, "ROCR_VISIBLE_DEVICES": "6,7"}, environ={}) + assert out["source"] == "baseline_recipe" + assert out["ids"] == [6, 7] + + +def test_a_blank_value_does_not_shadow_a_real_pin_further_down_the_chain() -> None: + """A blank ROCR must not hide a HIP pin — but it is still reported, see below. + + ``gpu_pool._visible_device_mask`` and ``gate.detect_gpu_count`` read + ``VAR=""`` as "zero devices visible"; this resolver used to skip it + entirely, so with a blank ROCR and a stale ``HIP=2,3`` those layers saw + zero GPUs while the handoff advertised two. The blank is now recorded (see + :func:`test_an_empty_mask_reports_zero_devices_rather_than_unpinned`) but + only as the fallback, so a real pin still wins. + """ + out = _resolve_gpu_pin( + recipe_envs={"ROCR_VISIBLE_DEVICES": " "}, + environ={"HIP_VISIBLE_DEVICES": "2,3"}, + ) + assert out["var"] == "HIP_VISIBLE_DEVICES" + assert out["ids"] == [2, 3] + + +# --------------------------------------------------------------------------- # +# The materializer's autofilled ROCR mask (PR #1321 review) +# --------------------------------------------------------------------------- # + + +def test_autofilled_recipe_rocr_does_not_override_a_hip_pin() -> None: + """Regression: recipe-first made every HIP-pinned run report cards 0..tp-1. + + ``materialize_config_with_envs`` synthesizes ``ROCR_VISIBLE_DEVICES=0,1`` + into the recipe for a ``TP=2`` run that has no ROCR anywhere. Honouring + that as a pin overrode the real ``HIP_VISIBLE_DEVICES=4,5`` and told a + ROCR-writing consumer to hard-pin physical cards 0 and 1 — recreating the + card-0 collision this whole change exists to remove. + """ + out = _resolve_gpu_pin( + recipe_envs=_autofilled(2), + environ={"HIP_VISIBLE_DEVICES": "4,5"}, + ) + assert out["var"] == "HIP_VISIBLE_DEVICES" + assert out["ids"] == [4, 5] + assert _resolve_handoff_gpu_ids(gpu_pin=out, tp=2) == "4,5" # pre-PR value, preserved + + +def test_autofilled_recipe_rocr_leaves_an_unpinned_run_unpinned() -> None: + """The documented ``{}`` contract has to be reachable in production.""" + assert _resolve_gpu_pin(recipe_envs=_autofilled(4), environ={}) == {} + assert _resolve_handoff_gpu_ids(gpu_pin={}, tp=4) == "0,1,2,3" + + +def test_a_real_recipe_rocr_pin_survives_the_autofill_check() -> None: + assert not _is_autofilled_rocr(value="4,5", recipe_envs={"TP": 2}) + assert _is_autofilled_rocr(value="0,1", recipe_envs={"TP": 2}) + # A recipe that records no TP still gets the materializer's mask written + # into it, so fall back to the SHAPE of the autofill (0..n-1). Requiring a + # recipe TP here left the synthetic mask posing as a pin for exactly the + # recipes that never recorded one. + assert _is_autofilled_rocr(value="0,1", recipe_envs={}) + assert _is_autofilled_rocr(value="0", recipe_envs={"TP": "not-a-number"}) + # A real pin is still a pin, with or without a TP to compare against. + assert not _is_autofilled_rocr(value="4,5", recipe_envs={}) + assert not _is_autofilled_rocr(value="6", recipe_envs={}) + assert not _is_autofilled_rocr(value="1,0", recipe_envs={}) + + +def test_variable_precedence_is_global_not_per_source() -> None: + """A leftover recipe CUDA key must not outrank a real process ROCR pin.""" + out = _resolve_gpu_pin( + recipe_envs={"TP": 2, "CUDA_VISIBLE_DEVICES": "0"}, + environ={"ROCR_VISIBLE_DEVICES": "6,7"}, + ) + assert out["var"] == "ROCR_VISIBLE_DEVICES" + assert out["ids"] == [6, 7] + + +def test_pin_reads_process_env_by_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "5") + assert _resolve_gpu_pin()["ids"] == [5] + + +# --------------------------------------------------------------------------- # +# _resolve_handoff_gpu_ids +# --------------------------------------------------------------------------- # + + +def test_gpu_ids_unpinned_is_range_tp() -> None: + """Unchanged legacy behaviour for an unpinned run.""" + assert _resolve_handoff_gpu_ids(gpu_pin={}, tp=4) == "0,1,2,3" + assert _resolve_handoff_gpu_ids(gpu_pin=None, tp=1) == "0" + assert _resolve_handoff_gpu_ids(gpu_pin={}, tp=0) == "0" + + +def test_gpu_ids_rocr_pin_is_logical() -> None: + """HIP indexes into the ROCr-visible set, so ROCR=6 is HIP index 0.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "6"}) + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=1) == "0" + + pin4 = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "4,5,6,7"}) + assert _resolve_handoff_gpu_ids(gpu_pin=pin4, tp=4) == "0,1,2,3" + # Capped at tp, as the unpinned path always was. + assert _resolve_handoff_gpu_ids(gpu_pin=pin4, tp=2) == "0,1" + # ...and at the mask when tp overshoots it: you cannot serve on cards you + # cannot see. + assert _resolve_handoff_gpu_ids(gpu_pin=pin4, tp=8) == "0,1,2,3" + + +def test_gpu_ids_hip_pin_is_verbatim() -> None: + """No ROCr mask => ROCr shows every card, so HIP ids are absolute.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"HIP_VISIBLE_DEVICES": "4,5"}) + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "4,5" + + +def test_gpu_ids_cuda_pin_is_verbatim() -> None: + pin = _resolve_gpu_pin(recipe_envs={}, environ={"CUDA_VISIBLE_DEVICES": "3"}) + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=1) == "3" + + +def test_gpu_ids_forwards_a_non_numeric_hip_mask_instead_of_recentring_on_card_0() -> None: + """A UUID HIP/CUDA mask parses to no numeric ids but is still a real pin. + + Re-serializing from ``ids`` alone yielded ``0..tp-1`` here, which moves the + servers onto cards ``0..tp-1`` — the #1312 failure, reintroduced for anyone + who pins by UUID. The tokens are forwarded as-is instead. + """ + pin = _resolve_gpu_pin(recipe_envs={}, environ={"HIP_VISIBLE_DEVICES": "GPU-a1b2c3,GPU-d4e5f6"}) + assert pin["ids"] == [] + assert pin["count"] == 2 + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "GPU-a1b2c3,GPU-d4e5f6" + + +def test_handoff_tp_never_exceeds_the_advertised_device_count() -> None: + """``tp`` follows ``gpu_ids`` down so the two cannot disagree. + + ``gpu_ids`` is capped at the mask width, so a stale ``$TP`` used to ship + alongside fewer ids and GEAK would launch ``--tp N`` against fewer visible + cards and fail to load weights. + """ + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "6"}) + ids = _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) + assert ids == "0" + assert _resolve_handoff_tp(gpu_ids=ids, tp=2) == 1 + # A four-card pin with a stale TP=8 clamps to the four cards it can see. + pin4 = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "4,5,6,7"}) + ids4 = _resolve_handoff_gpu_ids(gpu_pin=pin4, tp=8) + assert ids4 == "0,1,2,3" + assert _resolve_handoff_tp(gpu_ids=ids4, tp=8) == 4 + # An unpinned run is unaffected. + assert _resolve_handoff_tp(gpu_ids="0,1", tp=2) == 2 + + +def test_coerce_tp_never_raises_out_of_its_own_fallback() -> None: + """A non-numeric ``$TP`` must not escape as a ValueError. + + The previous form called a bare ``int()`` inside the ``except`` that was + handling the identical failure, so a junk ``$TP`` raised during handling. + """ + assert _coerce_tp("2", "8") == 2 + assert _coerce_tp(None, "8") == 8 + assert _coerce_tp("", " ") == 1 + assert _coerce_tp("not-a-number", "also-junk") == 1 + assert _coerce_tp("0", "-3", "4") == 4 + + +def test_gpu_ids_never_empty_for_a_blank_mask() -> None: + """A present-but-empty mask must not produce an empty device list.""" + pin = {"var": "ROCR_VISIBLE_DEVICES", "value": "", "ids": [], "count": 0, "source": "process_env"} + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "0,1" + + +def test_gpu_ids_counts_a_uuid_mask_instead_of_falling_back_to_card_0() -> None: + """ROCm accepts UUID masks; they parse to zero numeric ids but N devices. + + Counting ``ids`` here would see an empty list, read the run as unpinned and + emit ``0..tp-1`` — landing every GEAK server on card 0, the exact default + this change exists to eliminate. + """ + pin = _resolve_gpu_pin( + recipe_envs={}, + environ={"ROCR_VISIBLE_DEVICES": "GPU-a1b2c3,GPU-d4e5f6"}, + ) + assert pin["ids"] == [] + assert pin["count"] == 2 + assert pin["value"] == "GPU-a1b2c3,GPU-d4e5f6" # re-exportable as-is + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "0,1" + + +def test_a_yaml_sequence_mask_is_not_stringified_into_junk() -> None: + """``ROCR_VISIBLE_DEVICES: [4, 5]`` in a recipe is a list, not a string.""" + pin = _resolve_gpu_pin(recipe_envs={"TP": 2, "ROCR_VISIBLE_DEVICES": [4, 5]}, environ={}) + assert pin["value"] == "4,5" + assert pin["ids"] == [4, 5] + + +def test_gpu_ids_are_absolute_for_a_recipe_only_rocr_pin() -> None: + """A mask the child does not inherit cannot be indexed logically. + + The phase launches GEAK with ``dict(os.environ)``, so a mask that exists + only in the recipe never reaches the child; ROCr shows it every card and + the absolute ids are the correct HIP indices. + """ + pin = _resolve_gpu_pin(recipe_envs={"TP": 2, "ROCR_VISIBLE_DEVICES": "6,7"}, environ={}) + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "6,7" + + +# --------------------------------------------------------------------------- # +# Legacy mask spellings, empty masks, nested masks, coordinate space +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("var", "expect_logical"), + [ + ("ROCR_VISIBLE_DEVICES", True), + ("HSA_VISIBLE_DEVICES", True), + ("HIP_VISIBLE_DEVICES", False), + ("CUDA_VISIBLE_DEVICES", False), + ("GPU_DEVICE_ORDINAL", False), + ], +) +def test_every_mask_spelling_counts_as_a_pin(monkeypatch: pytest.MonkeyPatch, var: str, expect_logical: bool) -> None: + """A run pinned with a legacy spelling is pinned; omitting it read as unpinned.""" + monkeypatch.setenv(var, "6") + pin = _resolve_gpu_pin(recipe_envs=_autofilled(1)) + assert pin["var"] == var + assert pin["ids"] == [6] + # ROCr-level masks renumber the child's devices, HIP-level ones index into them. + assert (_resolve_handoff_gpu_ids_space(gpu_pin=pin) == "logical") is expect_logical + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=1) == ("0" if expect_logical else "6") + + +def test_a_recipe_autofill_is_ignored_under_the_legacy_rocr_spelling_too() -> None: + """``HSA_VISIBLE_DEVICES`` gets the same autofill test as its modern name.""" + pin = _resolve_gpu_pin( + recipe_envs={"TP": 2, "HSA_VISIBLE_DEVICES": "0,1"}, + environ={}, + ) + assert pin == {} + + +def test_an_empty_mask_reports_zero_devices_rather_than_unpinned() -> None: + """``ROCR_VISIBLE_DEVICES=""`` is "no cards", not "whole machine".""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": ""}) + assert pin["var"] == "ROCR_VISIBLE_DEVICES" + assert pin["count"] == 0 + assert pin["ids"] == [] + # gpu_ids still must not be blank: GEAK reads a falsy gpu_ids as "unset" and + # falls straight back to 0..tp-1 (interface/run_e2e.py), so an empty string + # buys nothing. The ids are declared placeholders instead — that is what + # the third coordinate space exists for. + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "0,1" + assert _resolve_handoff_gpu_ids_space(gpu_pin=pin) == "none" + + +def test_a_mask_with_no_valid_ordinal_is_also_reported_as_zero_devices() -> None: + """``ROCR="-1"`` is non-blank but exposes nothing; it must not read as a pin.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "-1"}) + assert pin["count"] == 0 + assert _resolve_handoff_gpu_ids_space(gpu_pin=pin) == "none" + + +def test_a_uuid_mask_is_not_mistaken_for_zero_devices() -> None: + """``ids`` is empty for a UUID mask, but it exposes real cards.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "GPU-a1b2c3,GPU-d4e5f6"}) + assert pin["ids"] == [] + assert pin["count"] == 2 + assert _resolve_handoff_gpu_ids_space(gpu_pin=pin) == "logical" + + +def test_a_real_pin_outranks_an_earlier_empty_mask() -> None: + """An empty ROCR must not shadow a real HIP pin further down the chain.""" + pin = _resolve_gpu_pin( + recipe_envs={}, + environ={"ROCR_VISIBLE_DEVICES": "", "HIP_VISIBLE_DEVICES": "4,5"}, + ) + assert pin["var"] == "HIP_VISIBLE_DEVICES" + assert pin["ids"] == [4, 5] + + +def test_a_hip_mask_nested_in_a_rocr_pin_is_forwarded_not_overwritten() -> None: + """ROCR=4,5,6,7 + HIP=2,3 is cards 6,7 — advertising 0,1 would move the servers.""" + pin = _resolve_gpu_pin( + recipe_envs={}, + environ={"ROCR_VISIBLE_DEVICES": "4,5,6,7", "HIP_VISIBLE_DEVICES": "2,3"}, + ) + assert pin["var"] == "ROCR_VISIBLE_DEVICES" + assert pin["inner"]["var"] == "HIP_VISIBLE_DEVICES" + assert pin["inner"]["ids"] == [2, 3] + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "2,3" + + +def test_a_nested_hip_mask_pointing_outside_the_rocr_set_is_dropped() -> None: + """HIP ids beyond the ROCr width name devices the child cannot see.""" + pin = _resolve_gpu_pin( + recipe_envs={}, + environ={"ROCR_VISIBLE_DEVICES": "6", "HIP_VISIBLE_DEVICES": "3"}, + ) + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=1) == "0" + + +def test_no_inner_mask_is_recorded_for_a_hip_level_pin() -> None: + """``inner`` only means "nested inside a ROCr slice"; a HIP pin has no inside.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"HIP_VISIBLE_DEVICES": "4,5"}) + assert "inner" not in pin + + +def test_gpu_ids_space_is_absolute_when_unpinned() -> None: + assert _resolve_handoff_gpu_ids_space(gpu_pin={}) == "absolute" + assert _resolve_handoff_gpu_ids_space(gpu_pin=None) == "absolute" + + +def test_a_recipe_rocr_pin_is_not_inherited_so_its_ids_stay_absolute() -> None: + """The child inherits the process env, not the recipe's envs.""" + pin = _resolve_gpu_pin(recipe_envs={"TP": 1, "ROCR_VISIBLE_DEVICES": "6"}, environ={}) + assert pin["source"] == "baseline_recipe" + assert _resolve_handoff_gpu_ids_space(gpu_pin=pin) == "absolute" + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=1) == "6" + + +def test_a_yaml_sequence_mask_is_joined_not_stringified() -> None: + """``ROCR_VISIBLE_DEVICES: [4, 5]`` in YAML must not become ``"[4, 5]"``.""" + pin = _resolve_gpu_pin(recipe_envs={"TP": 2, "ROCR_VISIBLE_DEVICES": [4, 5]}, environ={}) + assert pin["value"] == "4,5" + assert pin["ids"] == [4, 5] + assert pin["count"] == 2 + + +# --------------------------------------------------------------------------- # +# count / ids cardinality agreement +# --------------------------------------------------------------------------- # + + +def test_a_repeated_ordinal_does_not_inflate_the_device_count() -> None: + """ROCR="3,3,2" exposes two devices; logical index 2 would abort the server.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "3,3,2"}) + assert pin["count"] == 2 + assert pin["ids"] == [3, 2] + assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=3) == "0,1" + + +def test_a_negative_ordinal_is_not_counted_as_a_device() -> None: + pin = _resolve_gpu_pin(recipe_envs={}, environ={"ROCR_VISIBLE_DEVICES": "-1,2"}) + assert pin["count"] == 1 + assert pin["ids"] == [2] + + +def test_a_repeated_hip_ordinal_keeps_ids_and_count_in_agreement() -> None: + """The mirror of the count case: gpu_ids must not deflate below ``count``.""" + pin = _resolve_gpu_pin(recipe_envs={}, environ={"HIP_VISIBLE_DEVICES": "4,4"}) + gpu_ids = _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) + assert gpu_ids == "4" + assert len(gpu_ids.split(",")) == pin["count"] == 1 + # tp follows gpu_ids down, so GEAK is never told "tp=2" alongside one device. + assert _resolve_handoff_tp(gpu_ids=gpu_ids, tp=2) == 1 + + +def test_a_yaml_sequence_element_is_stripped_before_it_reaches_value() -> None: + """``[' 4', 5]`` must not produce ``" 4,5"`` — ROCm's parser rejects the token.""" + pin = _resolve_gpu_pin(recipe_envs={"TP": 2, "ROCR_VISIBLE_DEVICES": [" 4", 5]}, environ={}) + assert pin["value"] == "4,5" diff --git a/src/hyperloom/inference_optimizer/tests/test_geak_resume_recovery.py b/src/hyperloom/inference_optimizer/tests/test_geak_resume_recovery.py index f4508e33b5..1a914be57b 100644 --- a/src/hyperloom/inference_optimizer/tests/test_geak_resume_recovery.py +++ b/src/hyperloom/inference_optimizer/tests/test_geak_resume_recovery.py @@ -226,3 +226,104 @@ def _runner_resolved(_name: str) -> Path: assert handoff["accepted_flags"] == "--no-enable-prefix-caching" assert handoff["raw_baseline_tput"] == 100.0 assert handoff["e2e_metric"] == "output" + + +@pytest.mark.asyncio +async def test_geak_handoff_forwards_the_actual_gpu_pin( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The handoff must carry the run's real pin, not the literal card 0 (#1312). + + GEAK writes its own visible-devices mask for every server it launches. With + no pin in the handoff it defaults to physical GPU 0, so a run pinned to the + last card silently benchmarks on card 0 and OOMs against whatever else holds + it. ``gpu_ids`` stays HIP-logical (the consumer inherits the ROCR mask); + ``gpu_pin`` carries the absolute mask for consumers that write ROCR + themselves. + """ + coord = Coordinator.__new__(Coordinator) + coord.session_dir = tmp_path + coord.shared_state = SharedState(baseline_tput=100.0, model_path="/models/m", gpu_type="mi355x") + coord.phase_kernel._record_geak_kernel_journey = lambda _result: None + + monkeypatch.setenv("TP", "1") + monkeypatch.setenv("ROCR_VISIBLE_DEVICES", "7") + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + def _runner_resolved(_name: str) -> Path: + raise RuntimeError("stop after handoff write") + + monkeypatch.setattr( + "hyperloom.orchestrator.kernel.request_handlers._kernel_agent_tool_path", + _runner_resolved, + ) + + await coord._run_geak_kernel_phase(from_phase="KERNEL") + + handoff = json.loads((tmp_path / "geak" / "handoff.json").read_text(encoding="utf-8")) + assert handoff["schema_version"] >= 3 + assert handoff["gpu_pin"] == { + "var": "ROCR_VISIBLE_DEVICES", + "value": "7", + "ids": [7], + "count": 1, + "source": "process_env", + } + # Logical inside the inherited mask: index 0 IS physical card 7. + assert handoff["gpu_ids"] == "0" + + +@pytest.mark.asyncio +async def test_geak_handoff_keeps_a_hip_pin_against_the_recipe_autofill( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A materialized recipe always carries an autofilled ROCR mask (#1321 review). + + ``materialize_config_with_envs`` writes ``ROCR_VISIBLE_DEVICES=0..tp-1`` + into ``benchmark.envs`` whenever the mask is absent, and that file is what + ``state.baseline_config_path`` points at by the time KERNEL runs. Reading + the recipe first therefore overrode every HIP-pinned run with cards + ``0..tp-1`` — a new card-0 collision, in the change meant to remove one. + This is the production shape, so it is asserted end to end. + """ + recipe = tmp_path / "baseline.yaml" + recipe.write_text( + "benchmark:\n envs:\n TP: 2\n ROCR_VISIBLE_DEVICES: '0,1'\n NUM_PROMPTS: 192\n", + encoding="utf-8", + ) + + coord = Coordinator.__new__(Coordinator) + coord.session_dir = tmp_path + coord.shared_state = SharedState( + baseline_tput=100.0, + model_path="/models/m", + gpu_type="mi355x", + baseline_config_path=str(recipe), + ) + coord.phase_kernel._record_geak_kernel_journey = lambda _result: None + + monkeypatch.setenv("TP", "2") + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "4,5") + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising=False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + + def _runner_resolved(_name: str) -> Path: + raise RuntimeError("stop after handoff write") + + monkeypatch.setattr( + "hyperloom.orchestrator.kernel.request_handlers._kernel_agent_tool_path", + _runner_resolved, + ) + + await coord._run_geak_kernel_phase(from_phase="KERNEL") + + handoff = json.loads((tmp_path / "geak" / "handoff.json").read_text(encoding="utf-8")) + assert handoff["gpu_pin"]["var"] == "HIP_VISIBLE_DEVICES" + assert handoff["gpu_pin"]["ids"] == [4, 5] + # The pre-PR value. The whole point is that this change did not move it. + assert handoff["gpu_ids"] == "4,5" + # tp comes from the same recipe as gpu_ids, so the two cannot disagree. + assert handoff["tp"] == 2 diff --git a/src/hyperloom/inference_optimizer/tests/test_visible_devices.py b/src/hyperloom/inference_optimizer/tests/test_visible_devices.py new file mode 100644 index 0000000000..9e5098d330 --- /dev/null +++ b/src/hyperloom/inference_optimizer/tests/test_visible_devices.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT +"""The shared visible-devices module (PR #1321 review). + +Five layers used to carry their own copy of the mask tuple and the mask parser, +and their empty-mask semantics had already drifted apart. +``hyperloom.common.visible_devices`` is now the single definition, and it holds +two deliberately different tuples. These tests exist so the difference stays +deliberate: the counting subset is derived from the precedence chain, so a var +added to one cannot silently miss the other. +""" + +from __future__ import annotations + +import pytest + +from hyperloom.common.visible_devices import ( + COUNTING_VISIBLE_DEVICE_VARS, + GPU_MASK_ENV_NAMES, + HIP_LEVEL_VARS, + ROCR_LEVEL_VARS, + VISIBLE_DEVICE_VARS, + _UNCOUNTED_VISIBLE_DEVICE_VARS, + effective_mask_tokens, + is_rocr_level, + mask_tokens, + parse_device_list, +) + + +# --------------------------------------------------------------------------- # +# The two tuples cannot drift apart +# --------------------------------------------------------------------------- # + + +def test_every_chain_member_is_classified_as_counted_or_not() -> None: + """A var added to the chain must be a deliberate decision, not an omission. + + ``COUNTING_VISIBLE_DEVICE_VARS`` is derived, so a new var is counted by + default; leaving it out requires naming it in the exclusion set. This + asserts the two halves partition the chain exactly. + """ + assert set(COUNTING_VISIBLE_DEVICE_VARS) | _UNCOUNTED_VISIBLE_DEVICE_VARS == set(VISIBLE_DEVICE_VARS) + assert set(COUNTING_VISIBLE_DEVICE_VARS) & _UNCOUNTED_VISIBLE_DEVICE_VARS == set() + + +def test_the_exclusion_set_names_nothing_outside_the_chain() -> None: + """A typo'd or removed exclusion would silently widen the counting set.""" + assert _UNCOUNTED_VISIBLE_DEVICE_VARS <= set(VISIBLE_DEVICE_VARS) + + +def test_the_counting_subset_is_exactly_what_the_counting_layers_always_read() -> None: + """Deriving the tuple must not have changed GPU accounting.""" + assert COUNTING_VISIBLE_DEVICE_VARS == ( + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + ) + + +def test_the_counting_subset_keeps_the_chain_order() -> None: + """Precedence is the whole point of an ordered tuple.""" + kept = [var for var in VISIBLE_DEVICE_VARS if var in COUNTING_VISIBLE_DEVICE_VARS] + assert list(COUNTING_VISIBLE_DEVICE_VARS) == kept + + +def test_the_chain_is_rocr_level_before_hip_level() -> None: + """A ROCr mask slices the set a HIP mask then indexes into, so it wins.""" + assert VISIBLE_DEVICE_VARS == ROCR_LEVEL_VARS + HIP_LEVEL_VARS + assert all(is_rocr_level(var) for var in ROCR_LEVEL_VARS) + assert not any(is_rocr_level(var) for var in HIP_LEVEL_VARS) + + +def test_the_scrub_set_covers_the_whole_chain() -> None: + """Env scrubbing must not leave a mask behind that pin resolution honours.""" + assert GPU_MASK_ENV_NAMES == frozenset(VISIBLE_DEVICE_VARS) + + +# --------------------------------------------------------------------------- # +# The parsers +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize( + ("raw", "tokens"), + [ + ("4,5", ["4", "5"]), + ("4;5", ["4", "5"]), + (" 4 , 5 ", ["4", "5"]), + ("4,,5", ["4", "5"]), + ([" 4", 5], ["4", "5"]), + ("", []), + (None, []), + ], +) +def test_mask_tokens_is_the_literal_split(raw: object, tokens: list[str]) -> None: + assert mask_tokens(raw) == tokens + + +def test_mask_tokens_keeps_duplicates_but_effective_tokens_drops_them() -> None: + """The literal split and the effective device set are different questions.""" + assert mask_tokens("3,3,2") == ["3", "3", "2"] + assert effective_mask_tokens("3,3,2") == ["3", "2"] + + +def test_effective_tokens_drops_negative_ordinals_and_keeps_uuids() -> None: + """A negative ordinal is not a device; a UUID is.""" + assert effective_mask_tokens("-1,2") == ["2"] + assert effective_mask_tokens("GPU-a1b2c3,GPU-d4e5f6") == ["GPU-a1b2c3", "GPU-d4e5f6"] + + +def test_parse_device_list_agrees_with_effective_tokens() -> None: + """The ids are the numeric members of the effective set, never a wider one. + + This is the invariant that keeps ``gpu_pin["count"]`` and ``gpu_pin["ids"]`` + from disagreeing: counting the literal tokens inflates the count, and + re-serializing the parsed ints deflates the id list. + """ + for raw in ("4,5", "3,3,2", "-1,2", "GPU-a1b2c3,4", "", " 4, 4 ,5"): + ids = parse_device_list(raw) + tokens = effective_mask_tokens(raw) + assert len(ids) <= len(tokens) + assert [str(i) for i in ids] == [t for t in tokens if t.lstrip("-").isdigit()] diff --git a/src/hyperloom/orchestrator/actions/executors/_ray_serving.py b/src/hyperloom/orchestrator/actions/executors/_ray_serving.py index 0d9b077c7e..d9622713fe 100644 --- a/src/hyperloom/orchestrator/actions/executors/_ray_serving.py +++ b/src/hyperloom/orchestrator/actions/executors/_ray_serving.py @@ -13,6 +13,7 @@ from typing import Any from hyperloom.common.env_safety import scrub_benchmark_process_env +from hyperloom.common.visible_devices import COUNTING_VISIBLE_DEVICE_VARS from ._subprocess_kill import COOPERATIVE_REAP_BUDGET_SEC @@ -55,11 +56,9 @@ # the very round it is meant to stop. _SERVING_ACTOR_CONCURRENCY: int = 2 -_VISIBLE_DEVICE_ENV_KEYS: tuple[str, ...] = ( - "ROCR_VISIBLE_DEVICES", - "HIP_VISIBLE_DEVICES", - "CUDA_VISIBLE_DEVICES", -) +#: The masks Ray owns for its serving children. Single definition lives in +#: ``hyperloom.common.visible_devices``. +_VISIBLE_DEVICE_ENV_KEYS: tuple[str, ...] = COUNTING_VISIBLE_DEVICE_VARS class RayInfeasibleError(RuntimeError): diff --git a/src/hyperloom/orchestrator/bus/gpu_pool.py b/src/hyperloom/orchestrator/bus/gpu_pool.py index 7ef76bc7ab..c252a720c1 100644 --- a/src/hyperloom/orchestrator/bus/gpu_pool.py +++ b/src/hyperloom/orchestrator/bus/gpu_pool.py @@ -29,6 +29,7 @@ from datetime import datetime, timezone from hyperloom.common.timeutil import now_iso +from hyperloom.common.visible_devices import COUNTING_VISIBLE_DEVICE_VARS, parse_device_list from .storage.connection import SqliteConnection @@ -60,6 +61,10 @@ def _parse_gpu_list(raw: str) -> list[int]: """Parse a comma/semicolon-separated GPU id list. + Thin alias for :func:`hyperloom.common.visible_devices.parse_device_list`, + the single definition of this parse; kept as a module-local name because + call sites and tests already reference it. + Args: raw: Raw string of GPU ids (``,`` or ``;`` separated). @@ -67,18 +72,7 @@ def _parse_gpu_list(raw: str) -> list[int]: Unique non-negative GPU ids in first-seen order; malformed entries are skipped. """ - out: list[int] = [] - for part in (raw or "").replace(";", ",").split(","): - p = part.strip() - if not p: - continue - try: - idx = int(p) - except ValueError: - continue - if idx >= 0 and idx not in out: - out.append(idx) - return out + return parse_device_list(raw) def _explicit_pool() -> list[int] | None: @@ -110,7 +104,7 @@ def _visible_device_mask() -> tuple[list[int], bool]: set (even to an empty string, which means "no visible GPUs" → ``[]``); ``present`` is False only when none of the masks is set. """ - for env_name in ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + for env_name in COUNTING_VISIBLE_DEVICE_VARS: raw = os.environ.get(env_name) if raw is None: continue diff --git a/src/hyperloom/orchestrator/loop/coordinator_helpers.py b/src/hyperloom/orchestrator/loop/coordinator_helpers.py index 57e3ad7194..73cae52107 100644 --- a/src/hyperloom/orchestrator/loop/coordinator_helpers.py +++ b/src/hyperloom/orchestrator/loop/coordinator_helpers.py @@ -21,6 +21,14 @@ from typing import Any from hyperloom.common.env_safety import filter_untrusted_env_mapping, is_allowed_variant_env_key +from hyperloom.common.visible_devices import ( + HIP_LEVEL_VARS, + effective_mask_tokens, + VISIBLE_DEVICE_VARS, + is_rocr_level, + mask_tokens, + parse_device_list, +) from ..specialists.patch_safety import ( ADVISE_VERDICT, @@ -1546,6 +1554,380 @@ def _geak_sweep_measured_tput(res: dict[str, Any]) -> float | None: return None +#: Visible-device env masks, in the repo's ROCm precedence order. +#: The pin-resolution chain, imported rather than re-declared: the same tuple +#: and the same parser had five copies in this repo (``bus/gpu_pool``, +#: ``policy/gate``, ``actions/executors/_ray_serving``, ``common/env_safety``, +#: and this module) and their empty-mask semantics had already drifted apart. +#: ``hyperloom.common.visible_devices`` is now the single definition and is +#: dependency-free, so this pure-helper layer can use it without dragging in +#: the SQLite connection ``gpu_pool`` owns. +#: +#: Note this resolver uses the FULL chain, not the three vars the +#: capacity-counting layers read: it answers "where is this run pinned", and a +#: run pinned with ``HSA_VISIBLE_DEVICES`` or ``GPU_DEVICE_ORDINAL`` is really +#: pinned. Those layers keep their narrower :data:`COUNTING_VISIBLE_DEVICE_VARS` +#: because widening them would change GPU accounting repo-wide. +_VISIBLE_DEVICE_VARS: tuple[str, ...] = VISIBLE_DEVICE_VARS + +_mask_tokens = mask_tokens +_parse_device_list = parse_device_list + + +def _is_autofilled_rocr(*, value: str, recipe_envs: Mapping[str, Any]) -> bool: + """Is this recipe's ROCR mask the materializer's autofill rather than a pin? + + ``materialize_config_with_envs`` unconditionally writes + ``ROCR_VISIBLE_DEVICES=0..tp-1`` into ``benchmark.envs`` whenever the mask + is absent or narrower than TP (``_workload_envs.py``). Every materialized + recipe therefore carries the key, so a recipe ROCR value that is + byte-identical to that default carries no information about where the run + is actually pinned — treating it as a pin is what made this resolver + override a real ``HIP_VISIBLE_DEVICES`` and re-pin GEAK to cards ``0..tp-1``. + + A hand-authored ``ROCR_VISIBLE_DEVICES: "0,1"`` at ``TP=2`` is + indistinguishable from the autofill and is also treated as "not a pin"; + that is harmless, because the unpinned path emits the same ``gpu_ids`` and + merely omits ``gpu_pin``. + + When the recipe carries no usable ``TP`` — a hand-written or pre-clamp YAML + — there is no width to compare against, so the test falls back to the SHAPE + the materializer always produces: a mask that is exactly ``0..n-1`` for its + own length. Returning ``False`` there instead would let the synthetic mask + pose as a pin for precisely the recipes that never recorded a TP, which is + the hole this function exists to close. + + Args: + value: The recipe's ROCR mask, already stripped. + recipe_envs: The recipe's ``benchmark.envs`` (read for its resolved TP). + + Returns: + ``True`` when the value equals the ``0..tp-1`` the materializer would + have synthesized — or, absent a recipe TP, the ``0..n-1`` shape of one. + """ + tokens = _mask_tokens(value) + if not tokens: + return False + try: + tp = int(str(recipe_envs.get("TP") or 0)) + except (TypeError, ValueError): + tp = 0 + if tp <= 0: + tp = len(tokens) + return tokens == [str(i) for i in range(tp)] + + +def _mask_value(raw: Any) -> str: + """Normalize a raw mask (string or YAML sequence) to its string form. + + A YAML ``ROCR_VISIBLE_DEVICES: [4, 5]`` reaches us as a list, and + ``str([4, 5])`` would produce ``"[4, 5]"`` — a value no consumer can export. + + Args: + raw: The value as read from the env mapping or the recipe. + + Returns: + The comma-joined, stripped mask; ``""`` for an empty or blank one. + """ + if isinstance(raw, (list, tuple)): + return ",".join(str(p).strip() for p in raw if str(p).strip()) + return str(raw if raw is not None else "").strip() + + +def _resolve_inner_hip_mask( + *, + var: str, + env: Mapping[str, str], + recipe: Mapping[str, Any], +) -> dict[str, Any]: + """The HIP-level mask nested inside a winning ROCr-level pin, if any. + + ``ROCR_VISIBLE_DEVICES=4,5,6,7`` with ``HIP_VISIBLE_DEVICES=2,3`` does not + mean "cards 2 and 3": HIP indexes INTO what ROCr exposed, so the run is on + absolute cards 6 and 7. Dropping the inner mask and advertising + ``0..tp-1`` would move the servers to cards 4 and 5 — a quieter version of + the same #1312 bug, so the inner mask travels with the pin. + + Args: + var: The winning mask variable. + env: Process environment mapping. + recipe: The baseline recipe's ``benchmark.envs``. + + Returns: + ``{"var", "value", "ids", "count", "source"}`` for the innermost + HIP-level mask, or ``{}`` when the winner is not ROCr-level or no + HIP-level mask is set. + """ + if not is_rocr_level(var): + return {} + for hip_var in HIP_LEVEL_VARS: + for source, table in (("process_env", env), ("baseline_recipe", recipe)): + value = _mask_value(table.get(hip_var)) + if not value: + continue + return { + "var": hip_var, + "value": value, + "ids": _parse_device_list(value), + "count": len(effective_mask_tokens(value)), + "source": source, + } + return {} + + +def _resolve_gpu_pin( + *, + recipe_envs: Mapping[str, Any] | None = None, + environ: Mapping[str, str] | None = None, +) -> dict[str, Any]: + """Resolve the run's ACTUAL GPU pin for the geak handoff. + + GEAK launches full servers out-of-process and re-writes a visible-devices + mask for each one. Without the pin it can only guess, and the guess + (``0..tp-1``) silently lands on physical GPU 0 — see issue #1312, where a + run pinned elsewhere collided with a foreign tenant on card 0. Forwarding + the pin lets the consumer compose masks instead of clobbering them. + + Precedence is VARIABLE-major: ``ROCR_VISIBLE_DEVICES`` before ``HIP`` + before ``CUDA`` — the repo-wide order — and within each variable the + process env before the baseline recipe. Source-major ordering was wrong in + both directions: a leftover recipe ``CUDA_VISIBLE_DEVICES`` would outrank a + real process ROCR pin, and the recipe's autofilled ROCR (see + :func:`_is_autofilled_rocr`) would outrank everything. + + Args: + recipe_envs: The baseline recipe's ``benchmark.envs`` mapping (may be + ``None`` when no recipe is materialized yet). + environ: Environment mapping to read; defaults to ``os.environ``. + + Returns: + ``{"var", "value", "ids", "count", "source"}`` for the winning mask. + ``ids`` are the ABSOLUTE NUMERIC device ids and ``count`` is how many + devices the mask exposes; both derive from + :func:`effective_mask_tokens`, so ``count >= len(ids)`` always, and + they differ only when the mask is (partly) non-numeric — a UUID mask + gives ``ids == []`` with a non-zero ``count``. ``source`` is + ``"process_env"`` or ``"baseline_recipe"``. + A mask that is SET BUT EMPTY yields ``count == 0`` (zero devices + visible) rather than ``{}``; when the winner is a ROCr-level mask and a + HIP-level mask is also in force, the latter travels under ``"inner"`` + because it selects a subset *within* the ROCr-visible set. + ``{}`` only when no mask is set anywhere — meaning "whole machine + visible", not "pinned to 0". + """ + env = os.environ if environ is None else environ + recipe = dict(recipe_envs or {}) + blank: dict[str, Any] = {} + for var in _VISIBLE_DEVICE_VARS: + for source, table in (("process_env", env), ("baseline_recipe", recipe)): + raw = table.get(var) + if raw is None: + continue + value = _mask_value(raw) + if not value: + # Present but empty: a real "zero devices visible" state, not an + # absent mask. Remember the first one and keep looking — a real + # pin further down the chain still outranks it — but if nothing + # else is set, report it as a zero-device pin rather than as + # "unpinned", which reads as "whole machine". + if not blank: + blank = {"var": var, "value": "", "ids": [], "count": 0, "source": source} + continue + if ( + source == "baseline_recipe" + and is_rocr_level(var) + and _is_autofilled_rocr(value=value, recipe_envs=recipe) + ): + continue + pin: dict[str, Any] = { + "var": var, + "value": value, + "ids": _parse_device_list(value), + "count": len(effective_mask_tokens(value)), + "source": source, + } + inner = _resolve_inner_hip_mask(var=var, env=env, recipe=recipe) + if inner: + pin["inner"] = inner + return pin + return blank + + +def _resolve_handoff_gpu_ids(*, gpu_pin: Mapping[str, Any] | None, tp: int) -> str: + """Resolve the handoff's ``gpu_ids`` in the coordinate system GEAK applies it in. + + ``gpu_ids`` is a HIP-level device list: the consumer exports it as + ``HIP_VISIBLE_DEVICES``/``CUDA_VISIBLE_DEVICES`` for the servers it + launches, and HIP indexes into the ROCr-visible set. So: + + * pinned with a ROCr-level mask the child INHERITS — that mask renumbers + the child's devices, so the ids must be LOGICAL positions inside it + (``ROCR=6`` → ``"0"``), capped at ``tp`` (``ROCR=4,5,6,7`` with + ``tp=2`` → ``"0,1"``) and at the mask width when ``tp`` overshoots it. + Counted from :func:`effective_mask_tokens`, so a UUID mask resolves to + the right number of logical slots and a repeated ordinal does not + invent one. A HIP-level mask nested inside the ROCr slice is already in + logical coordinates and is forwarded instead (``ROCR=4,5,6,7`` + + ``HIP=2,3`` is cards 6 and 7, so ``"2,3"``); + * any other pin — ROCr still shows every card, so the mask's own tokens + pass through uncapped (``HIP=4,5`` → ``"4,5"``). They come from + :func:`effective_mask_tokens`, the same list ``gpu_pin["count"]`` is + derived from, so whitespace is normalized without the id list and the + advertised device count ever disagreeing. A NON-NUMERIC mask (a UUID + list) is forwarded token for token rather than collapsed to + ``0..tp-1``, which would silently move the servers onto cards + ``0..tp-1`` — the #1312 failure this resolver exists to prevent; + * not pinned — ``0..tp-1``, unchanged. + + The absolute pin travels separately in ``handoff["gpu_pin"]``, and + :func:`_resolve_handoff_gpu_ids_space` says which of the two coordinate + systems the result is in. A consumer that exports the result as + ``HIP_VISIBLE_DEVICES`` without touching ROCr is correct in both; a + consumer that re-applies ``gpu_pin["value"]`` as ``ROCR_VISIBLE_DEVICES`` + has just renumbered the devices itself and must use ``0..count-1``, NOT + these ids, for the inner HIP mask. + + Args: + gpu_pin: The :func:`_resolve_gpu_pin` result (``{}``/``None`` = unpinned). + tp: Tensor-parallel size; ``<= 1`` is treated as 1. + + Returns: + A comma-separated device list, never empty. + """ + width = max(int(tp or 1), 1) + pin = gpu_pin or {} + ids = list(pin.get("ids") or []) + # Logical remapping applies only to a mask the GEAK child actually + # INHERITS. The phase launches it with ``dict(os.environ)``, so a + # process-env ROCR mask is inherited and its ids are logical; a mask that + # only exists in the recipe is not, ROCr shows every card, and the absolute + # ids are the correct HIP indices. + if _pin_is_inherited_rocr(pin): + # Token count, not len(ids): a UUID mask parses to zero numeric ids but + # still exposes that many cards to the child. + visible = int(pin.get("count") or len(ids) or 0) + if visible > 0: + # A HIP-level mask nested inside the ROCr pin is ALREADY expressed + # in the child's logical coordinates, so it is forwarded as-is + # rather than overwritten with ``0..n-1``. Out-of-range entries are + # dropped: they name devices the ROCr mask never exposed. + inner = _mask_tokens((pin.get("inner") or {}).get("value")) + kept = [tok for tok in inner if not tok.isdigit() or int(tok) < visible] + if kept: + return ",".join(kept[:width]) + return ",".join(str(i) for i in range(min(visible, width))) + # Forward the EFFECTIVE tokens, not a re-serialization of the parsed ints: + # a UUID mask has no ints to re-serialize and would otherwise collapse to + # ``0..tp-1`` (the #1312 failure), and ``pin["count"]`` is derived from this + # same list, so the id list and the advertised device count cannot disagree. + tokens = effective_mask_tokens(pin.get("value")) + if tokens: + return ",".join(tokens) + return ",".join(str(i) for i in range(width)) + + +def _pin_is_inherited_rocr(pin: Mapping[str, Any] | None) -> bool: + """Will the GEAK child inherit this pin as a ROCr-level device slice? + + Only then are the handoff's ``gpu_ids`` logical. The phase launches GEAK + with ``dict(os.environ)``, so a process-env ROCr mask is inherited and + renumbers the child's devices; a mask that only exists in the recipe is + not, ROCr shows every card, and absolute ids are the correct HIP indices. + + Args: + pin: The :func:`_resolve_gpu_pin` result. + + Returns: + ``True`` for a process-env ROCr-level pin. + """ + pin = pin or {} + return is_rocr_level(str(pin.get("var") or "")) and str(pin.get("source") or "") == "process_env" + + +def _resolve_handoff_gpu_ids_space(*, gpu_pin: Mapping[str, Any] | None) -> str: + """Which coordinate system the handoff's ``gpu_ids`` are expressed in. + + ``gpu_ids`` alone is ambiguous: ``"0,1"`` is either "the first two cards of + the inherited ROCr mask" or "absolute cards 0 and 1", and a consumer that + guesses wrong re-pins the servers onto physical GPU 0 — issue #1312. This + field makes the distinction explicit so a consumer that composes masks + itself (rather than exporting ``gpu_ids`` into HIP) can tell which it was + handed. Consumers that ignore it keep the old, correct behaviour of + exporting ``gpu_ids`` as ``HIP_VISIBLE_DEVICES``, which is a HIP-level + variable in both spaces. + + ``"none"`` is the third case and the reason this is a tri-state rather + than a boolean: the mask is SET BUT EMPTY, so the run has no visible + devices and NO id list can be truthful. ``gpu_ids`` still carries + ``0..tp-1`` because the consumer reads a falsy ``gpu_ids`` as "unset" and + falls back to exactly those ids anyway (``interface/run_e2e.py``) — an + empty string would buy nothing and lose the ability to say why. The ids are + placeholders in that case and a consumer must not launch on them. + + Args: + gpu_pin: The :func:`_resolve_gpu_pin` result (``{}``/``None`` = unpinned). + + Returns: + ``"none"`` when the pin exposes zero devices, ``"logical"`` when the + ids index into an inherited ROCr mask, ``"absolute"`` otherwise + (including unpinned). + """ + pin = gpu_pin or {} + if pin and int(pin.get("count") or 0) <= 0: + return "none" + return "logical" if _pin_is_inherited_rocr(pin) else "absolute" + + +def _coerce_tp(*args: Any, default: int = 1) -> int: + """First positional that parses as a positive int, else ``default``. + + Every candidate is guarded, so no caller has to wrap ``int()`` in a + ``try`` whose handler then calls ``int()`` again on a value that can raise + the same exception it is handling. + + Args: + *args: Candidate TP values in precedence order (``None``/blank skipped). + default: Returned when nothing parses; floored at 1. + + Returns: + A TP of at least 1. + """ + for cand in args: + text = str(cand if cand is not None else "").strip() + if not text: + continue + try: + val = int(text) + except (TypeError, ValueError): + continue + if val > 0: + return val + return max(int(default), 1) + + +def _resolve_handoff_tp(*, gpu_ids: str, tp: int) -> int: + """Clamp ``tp`` to the number of devices the handoff actually advertises. + + ``gpu_ids`` is capped at the pin's mask width, so a run whose ``$TP`` + overshoots its pin (``ROCR=6`` with ``TP=2``, or a stale ``TP=8`` against a + materializer-clamped 4-card recipe) would otherwise ship ``tp`` and + ``gpu_ids`` that disagree — and GEAK would launch ``--tp N`` against fewer + visible cards and fail to load weights. Deriving both from the same resolved + mask makes that state unrepresentable. + + Args: + gpu_ids: The resolved handoff ``gpu_ids`` string. + tp: The TP resolved from the recipe/process env. + + Returns: + ``min(tp, len(gpu_ids))``, never below 1. + """ + advertised = len(_mask_tokens(gpu_ids)) + if advertised <= 0: + return max(int(tp or 1), 1) + return max(min(int(tp or 1), advertised), 1) + + def _parse_server_arg_value(server_args: str, flag: str) -> str | None: """Extract a CLI flag's value from a server-args string. diff --git a/src/hyperloom/orchestrator/phases/kernel.py b/src/hyperloom/orchestrator/phases/kernel.py index 9b455a7161..e7832660af 100644 --- a/src/hyperloom/orchestrator/phases/kernel.py +++ b/src/hyperloom/orchestrator/phases/kernel.py @@ -37,6 +37,11 @@ _geak_has_accepted_kernel, _resolve_roofline_watermark_ratio, _accepted_config_as_variant, + _coerce_tp, + _resolve_gpu_pin, + _resolve_handoff_gpu_ids, + _resolve_handoff_gpu_ids_space, + _resolve_handoff_tp, _resolve_serving_fidelity, ) from .base import PhaseHandler @@ -629,23 +634,42 @@ def _is_bf16_dense_gemm_fallback_attempt(entry: dict[str, Any]) -> bool: return False @staticmethod - def _resolve_bench_protocol(recipe_path: str) -> dict[str, Any]: - """Extract Hyperloom's bench measurement protocol for the GEAK handoff. + def _read_recipe_bench_envs(recipe_path: str) -> dict[str, Any]: + """Read the materialized baseline recipe's ``benchmark.envs``. Never raises. - Reads the materialized baseline recipe's ``benchmark.envs`` (falling back - to the process env) and returns only the keys that resolve, so absent - values leave GEAK on its standalone defaults. Never raises. + Args: + recipe_path: Path to the baseline recipe YAML (may be empty/missing). + + Returns: + The ``benchmark.envs`` mapping, or ``{}`` when the recipe is absent + or unreadable. """ - envs: dict[str, Any] = {} try: import yaml if recipe_path and Path(recipe_path).is_file(): cfg = yaml.safe_load(Path(recipe_path).read_text(encoding="utf-8")) or {} envs = ((cfg.get("benchmark") or {}).get("envs")) or {} + return dict(envs) if isinstance(envs, dict) else {} except Exception: # noqa: BLE001 - log.warning("bench_protocol: could not read recipe %r", recipe_path, exc_info=True) - envs = {} + log.warning("geak handoff: could not read recipe %r", recipe_path, exc_info=True) + return {} + + @classmethod + def _resolve_bench_protocol(cls, recipe_path: str, *, envs: dict[str, Any] | None = None) -> dict[str, Any]: + """Extract Hyperloom's bench measurement protocol for the GEAK handoff. + + Reads the materialized baseline recipe's ``benchmark.envs`` (falling back + to the process env) and returns only the keys that resolve, so absent + values leave GEAK on its standalone defaults. Never raises. + + Args: + recipe_path: Path to the baseline recipe YAML. + envs: An already-parsed ``benchmark.envs`` for that path. Pass it + when the caller needs the envs too, so the YAML is read once and + both consumers see the same snapshot. + """ + envs = cls._read_recipe_bench_envs(recipe_path) if envs is None else envs def _pick(key: str, cast: Callable[[str], Any]) -> Any: raw = envs.get(key) @@ -746,7 +770,40 @@ async def _run_geak_kernel_phase(self, *, from_phase: str) -> None: # Forward the SAME bench knobs Hyperloom benched with so GEAK's internal # e2e measures identically; source = the baseline recipe's benchmark.envs # (process-env fallback). Only resolved keys are sent. - bench_protocol = self._resolve_bench_protocol(str(getattr(state, "baseline_config_path", "") or "")) + _recipe_path = str(getattr(state, "baseline_config_path", "") or "") + # One read, two consumers: the recipe is parsed once so bench_protocol, + # the GPU pin and tp below all see the same snapshot. + _recipe_envs = self._read_recipe_bench_envs(_recipe_path) + bench_protocol = self._resolve_bench_protocol(_recipe_path, envs=_recipe_envs) + # The run's ACTUAL GPU pin (issue #1312). GEAK launches full servers + # out-of-process and writes its own visible-devices mask for each one; + # with no pin in the handoff it defaults to physical GPU 0 and collides + # with whatever else holds that card. Process mask first, recipe as the + # fallback; {} means "whole machine", not "card 0". + gpu_pin = _resolve_gpu_pin(recipe_envs=_recipe_envs) + # tp must come from the SAME place as gpu_ids or the two can disagree. + # The materializer clamps TP to the visible GPU count, so a stale + # $TP=8 on a 4-card pod would otherwise ship `tp: 8` alongside four + # gpu_ids and GEAK would launch sglang with --tp 8 and fail to load. + # Recipe TP first, process env as the fallback; both guarded, so a + # non-numeric $TP cannot raise out of the fallback itself. + _tp = _coerce_tp(_recipe_envs.get("TP"), os.environ.get("TP")) + # gpu_ids is capped at the pin's mask width; tp follows it down so the + # two can never disagree (e.g. ROCR=6 with TP=2 ships tp=1, not tp=2). + _gpu_ids = _resolve_handoff_gpu_ids(gpu_pin=gpu_pin, tp=_tp) + _tp = _resolve_handoff_tp(gpu_ids=_gpu_ids, tp=_tp) + _gpu_ids_space = _resolve_handoff_gpu_ids_space(gpu_pin=gpu_pin) + if _gpu_ids_space == "none": + # Set-but-empty mask: the run has no visible devices, so no id list + # in this handoff can be truthful. The payload says so via + # gpu_ids_space, but say it in the log too — a GEAK run that dies + # on "invalid device ordinal" is otherwise unexplainable. + log.error( + "geak handoff: %s is set but empty; the run has no visible GPUs. " + "gpu_ids=%s is a placeholder (gpu_ids_space=none), not a device set.", + gpu_pin.get("var"), + _gpu_ids, + ) # Serving-launch fidelity: forward the SAME max-model-len / gpu-mem-util # the baseline served with so GEAK launches the identical engine and its # baseline matches raw_baseline_tput. Resolver parses these from the raw @@ -765,11 +822,14 @@ async def _run_geak_kernel_phase(self, *, from_phase: str) -> None: handoff = { # v2 adds ``baseline_env_spec`` (the full layered env of current_best); # v1-only consumers ignore it and degrade to the flags/env-only baseline. - "schema_version": 2, + # v3 adds ``gpu_pin`` (the run's actual visible-devices mask) and + # ``gpu_ids_space``; older consumers ignore both and keep reading + # ``gpu_ids`` exactly as before. + "schema_version": 3, "model_path": str(getattr(state, "model_path", "") or os.environ.get("MODEL_PATH", "")), "framework": str(os.environ.get("FRAMEWORK", "") or "sglang"), "gpu_type": str(getattr(state, "gpu_type", "") or os.environ.get("GPU_TYPE", "")), - "tp": int(os.environ.get("TP", "1") or 1), + "tp": _tp, "workload": workload, "accepted_flags": accepted_flags, "accepted_env": accepted_env, @@ -795,13 +855,23 @@ async def _run_geak_kernel_phase(self, *, from_phase: str) -> None: "bench_client": "auto", "e2e_metric": "output", "inferencex_path": str(os.environ.get("INFERENCEX_PATH", "")), - # Pin the serving GPU set: explicit visibility mask, else 0..tp-1. - "gpu_ids": ( - os.environ.get("HIP_VISIBLE_DEVICES") - or os.environ.get("CUDA_VISIBLE_DEVICES") - or ",".join(str(i) for i in range(int(os.environ.get("TP", "1") or 1))) - ), + # The serving/optimization device set, as HIP-level ids (what the + # consumer exports as HIP_VISIBLE_DEVICES). Logical positions inside + # an inherited ROCR mask, a HIP/CUDA mask as-is, else 0..tp-1. + "gpu_ids": _gpu_ids, + # Which coordinate system ``gpu_ids`` is in: "logical" (positions + # inside the ROCR mask the child inherits), "absolute" (whole- + # machine device ids), or "none" (the mask is set but empty — the + # ids are placeholders and must not be launched on). Exporting them + # as HIP_VISIBLE_DEVICES is correct in the first two; a consumer + # that instead writes ROCR itself needs to know which it holds. + "gpu_ids_space": _gpu_ids_space, } + if gpu_pin: + # ABSOLUTE ids + the var they came from, so a consumer that writes + # ROCR_VISIBLE_DEVICES itself re-applies the same pin instead of + # resetting the child to card 0. + handoff["gpu_pin"] = gpu_pin if bench_protocol: handoff["bench_protocol"] = bench_protocol # Only forward resolved fidelity knobs; absence => GEAK adapter default. diff --git a/src/hyperloom/orchestrator/policy/gate.py b/src/hyperloom/orchestrator/policy/gate.py index 449b2431e3..285f63a210 100644 --- a/src/hyperloom/orchestrator/policy/gate.py +++ b/src/hyperloom/orchestrator/policy/gate.py @@ -21,6 +21,7 @@ resolve_gpu_specialist_devices, resolve_whole_machine_devices, ) +from hyperloom.common.visible_devices import COUNTING_VISIBLE_DEVICE_VARS from hyperloom.inference_optimizer.protocol.intent import Intent, IntentType from hyperloom.inference_optimizer.protocol.action_surfaces import ( COORDINATOR_INTERNAL_ACTIONS, @@ -179,7 +180,7 @@ def detect_gpu_count() -> int: ``CUDA_VISIBLE_DEVICES`` env masks (first one set wins), else the count parsed from ``rocm-smi``; 0 when nothing can be probed. """ - for env_name in ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + for env_name in COUNTING_VISIBLE_DEVICE_VARS: raw = os.environ.get(env_name) if raw is None: continue