fix(geak): forward the run's actual GPU pin in the handoff

docs/components/geak.md

@@ -38,6 +38,29 @@ 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_VISIBLE_DEVICES` mask, capped at `tp` (`ROCR=6` → `"0"`); a `HIP`/`CUDA` mask verbatim (`HIP=4,5` → `"4,5"`); `0..tp-1` when the run is unpinned |
+| `gpu_pin` | absolute device ids | `{"var", "value", "ids", "source"}` for the winning mask — omitted entirely when no mask is set anywhere, which means "whole machine visible", not "pinned to card 0" |
+
+The mask is resolved from the materialized baseline recipe's `benchmark.envs`
+first (the mask Hyperloom actually benched with), then the process environment,
+checking `ROCR_VISIBLE_DEVICES` before `HIP_VISIBLE_DEVICES` /
+`CUDA_VISIBLE_DEVICES` — the same precedence as `orchestrator/bus/gpu_pool.py`
+and `orchestrator/policy/gate.py`.
+
+A consumer that re-exports `HIP_VISIBLE_DEVICES` should use `gpu_ids`; one that
+writes `ROCR_VISIBLE_DEVICES` itself must use `gpu_pin["value"]`, because
+writing `gpu_ids` into `ROCR_VISIBLE_DEVICES` resets the child to physical card
+0 regardless of the run's pin.
+
 ## GEAK documentation
 
 For detailed documentation on GEAK, see [GEAK on ROCm Docs](https://rocm.docs.amd.com/projects/geak/en/latest/).

src/hyperloom/inference_optimizer/breakdown/collectors/geak.py

@@ -671,6 +671,11 @@ 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")),
+            # Device set + the run's absolute pin: 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.
+            "gpu_ids": handoff.get("gpu_ids"),
+            "gpu_pin": handoff.get("gpu_pin"),
         }
 
     # 2) a flushed-but-unpromoted result.json (absent or non-ok status).

src/hyperloom/inference_optimizer/tests/test_geak_handoff_gpu_pin.py

@@ -0,0 +1,154 @@
+# 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 (
+    _parse_device_list,
+    _resolve_gpu_pin,
+    _resolve_handoff_gpu_ids,
+)
+
+_MASK_VARS = ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
+
+
+@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],
+        "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_recipe_over_process_env() -> None:
+    """The recipe mask is what Hyperloom actually benched with."""
+    out = _resolve_gpu_pin(
+        recipe_envs={"ROCR_VISIBLE_DEVICES": "6"},
+        environ={"ROCR_VISIBLE_DEVICES": "0"},
+    )
+    assert out["source"] == "baseline_recipe"
+    assert out["ids"] == [6]
+
+
+def test_pin_skips_blank_values() -> None:
+    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]
+
+
+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_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": [], "source": "process_env"}
+    assert _resolve_handoff_gpu_ids(gpu_pin=pin, tp=2) == "0,1"

src/hyperloom/inference_optimizer/tests/test_geak_resume_recovery.py

@@ -226,3 +226,49 @@ 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],
+        "source": "process_env",
+    }
+    # Logical inside the inherited mask: index 0 IS physical card 7.
+    assert handoff["gpu_ids"] == "0"

src/hyperloom/orchestrator/loop/coordinator_helpers.py

@@ -1546,6 +1546,117 @@ def _geak_sweep_measured_tput(res: dict[str, Any]) -> float | None:
     return None
 
 
+#: Visible-device env masks, in the repo's ROCm precedence order.
+#: ``ROCR_VISIBLE_DEVICES`` is canonical on ROCm (the CLI preflight drops
+#: ``HIP_VISIBLE_DEVICES`` when ROCR is set); HIP/CUDA cover CUDA-style and
+#: legacy pins. Same order as ``gpu_pool._visible_device_mask`` /
+#: ``policy.gate.detect_gpu_count`` so every layer agrees on "the pin".
+#: (Kept local rather than imported from ``gpu_pool``: this module is the pure
+#: helper layer and ``gpu_pool`` drags in the SQLite connection.)
+_VISIBLE_DEVICE_VARS: tuple[str, ...] = (
+    "ROCR_VISIBLE_DEVICES",
+    "HIP_VISIBLE_DEVICES",
+    "CUDA_VISIBLE_DEVICES",
+)
+
+
+def _parse_device_list(raw: Any) -> list[int]:
+    """Parse a visible-devices mask string into absolute GPU ids.
+
+    Args:
+        raw: A ``,``/``;``-separated mask (``"4,5,6,7"``); ``None`` and
+            malformed entries are tolerated.
+
+    Returns:
+        Unique non-negative ids in first-seen order; ``[]`` for an empty or
+        fully malformed mask.
+    """
+    out: list[int] = []
+    for part in str(raw or "").replace(";", ",").split(","):
+        tok = part.strip()
+        if not tok:
+            continue
+        try:
+            idx = int(tok)
+        except ValueError:
+            continue
+        if idx >= 0 and idx not in out:
+            out.append(idx)
+    return out
+
+
+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.
+
+    Source precedence: the materialized baseline recipe's ``benchmark.envs``
+    (the mask Hyperloom actually benched with) before the process env, and
+    ``ROCR_VISIBLE_DEVICES`` before ``HIP``/``CUDA`` within each.
+
+    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", "source"}`` for the winning mask, where
+        ``ids`` are ABSOLUTE device ids and ``source`` is
+        ``"baseline_recipe"`` or ``"process_env"``. ``{}`` when no mask is set
+        anywhere — meaning "whole machine visible", not "pinned to 0".
+    """
+    env = os.environ if environ is None else environ
+    for source, table in (("baseline_recipe", recipe_envs or {}), ("process_env", env)):
+        for var in _VISIBLE_DEVICE_VARS:
+            raw = table.get(var)
+            if raw is None or str(raw).strip() == "":
+                continue
+            value = str(raw).strip()
+            return {"var": var, "value": value, "ids": _parse_device_list(value), "source": source}
+    return {}
+
+
+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 ``ROCR_VISIBLE_DEVICES`` — the child inherits that mask, so
+        the ids must be LOGICAL positions inside it (``ROCR=6`` → ``"0"``),
+        capped at ``tp`` as before (``ROCR=4,5,6,7`` with ``tp=2`` → ``"0,1"``);
+      * pinned with ``HIP``/``CUDA`` — ROCr still shows every card, so the mask
+        is forwarded VERBATIM (``HIP=4,5`` → ``"4,5"``);
+      * not pinned — ``0..tp-1``, unchanged.
+
+    The absolute pin travels separately in ``handoff["gpu_pin"]`` for consumers
+    that write ``ROCR_VISIBLE_DEVICES`` themselves.
+
+    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)
+    ids = list((gpu_pin or {}).get("ids") or [])
+    if not ids:
+        return ",".join(str(i) for i in range(width))
+    if str((gpu_pin or {}).get("var") or "") == "ROCR_VISIBLE_DEVICES":
+        return ",".join(str(i) for i in range(min(len(ids), width)))
+    return ",".join(str(i) for i in ids)
+
+
 def _parse_server_arg_value(server_args: str, flag: str) -> str | None:
     """Extract a CLI flag's value from a server-args string.