fix(geak): forward the run's actual GPU pin in the handoff - #1321
fix(geak): forward the run's actual GPU pin in the handoff#1321zihaoanllm wants to merge 2 commits into
Conversation
`geak/handoff.json` never carried the run's visible-devices mask. `gpu_ids` was resolved from `HIP_VISIBLE_DEVICES` / `CUDA_VISIBLE_DEVICES` only, so a run pinned the ROCm-canonical way (`ROCR_VISIBLE_DEVICES`) fell through to `0..tp-1`. GEAK writes its own mask for every full server it launches (baseline, profile, config-tuning validation), so every one of them landed on physical GPU 0 regardless of the pin. On a shared host that collides with whatever else holds card 0 and surfaces as `torch.OutOfMemoryError` in the baseline/profile server logs, or as a specialist declining to A/B because the "serving GPU" is held by a foreign co-tenant — indistinguishable from a real `no_gain` result. The two coordinate systems are now explicit: * `gpu_ids` stays a HIP-level device list (HIP indexes into the ROCr-visible set): logical positions inside an inherited ROCR mask, a HIP/CUDA mask verbatim, `0..tp-1` when unpinned. Behaviour is unchanged for every mask that worked before; the ROCR case is now defined instead of accidental. * `gpu_pin` (new, schema_version 3) carries the ABSOLUTE ids plus the variable and source they came from, for consumers that write `ROCR_VISIBLE_DEVICES` themselves. Omitted entirely when no mask is set anywhere — that means "whole machine visible", not "pinned to card 0". The mask is sourced from the materialized baseline recipe's `benchmark.envs` first (what Hyperloom actually benched with), then the process environment, ROCR before HIP/CUDA — the same precedence as `bus/gpu_pool.py` and `policy/gate.py`. The GEAK breakdown collector now records `gpu_ids` / `gpu_pin` so a degraded baseline is diagnosable from the session artifacts instead of by hand. Fixes #1312
xiaofei-zheng
left a comment
xiaofei-zheng
left a comment
There was a problem hiding this comment.
Code review of the GPU-pin handoff change. The direction is right — ROCR_VISIBLE_DEVICES genuinely was the missing link, and forwarding an absolute pin is the correct shape. But as written I believe the recipe-first precedence defeats the fix for HIP/CUDA-pinned runs and makes the documented "unpinned" contract unreachable in production, because materialize_config_with_envs autofills ROCR_VISIBLE_DEVICES=0..tp-1 into every materialized recipe. Details inline; the first two are the ones I would block on.
One cross-cutting note not anchorable to the diff: handoff["tp"] is still read raw from $TP while gpu_ids is now clamped to the mask width, so the two fields can disagree (see the comment on kernel.py).
| 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)): |
There was a problem hiding this comment.
Recipe-first precedence defeats the fix for HIP/CUDA-pinned runs.
materialize_config_with_envs unconditionally synthesizes ROCR_VISIBLE_DEVICES=0..tp-1 into benchmark.envs (_workload_envs.py:943-956), and that materialized recipe is what state.baseline_config_path points at by the time KERNEL runs. So the baseline_recipe source always wins, and it always carries a synthetic 0..tp-1 mask.
Repro: run with HIP_VISIBLE_DEVICES=4,5, TP=2, no ROCR anywhere. _resolve_gpu_pin returns {'var': 'ROCR_VISIBLE_DEVICES', 'value': '0,1', 'ids': [0, 1], 'source': 'baseline_recipe'} (verified by executing the helper). handoff["gpu_ids"] becomes "0,1" where the pre-PR code emitted "4,5", and gpu_pin["value"] = "0,1" tells a ROCR-writing consumer to hard-pin physical cards 0 and 1 — recreating exactly the foreign-tenant card-0 collision #1312 is meant to fix, as a new regression for HIP users.
The recipe mask is only meaningful as a pin when it was authored, not when it was autofilled. Either skip the autofilled ROCR key, or consult the process env first and use the recipe only as a fallback.
| 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} |
There was a problem hiding this comment.
The documented {} ("whole machine") case is unreachable in production.
Same root cause as above: because the materialized recipe always carries an autofilled ROCR_VISIBLE_DEVICES=0..tp-1, a genuinely unpinned run still produces a truthy pin. With no mask anywhere in the process env and TP=4, this returns {'ids': [0,1,2,3], 'source': 'baseline_recipe', ...} (verified by executing the helper), so handoff["gpu_pin"] is emitted.
That contradicts the docstring above and the table in docs/components/geak.md ("omitted entirely when no mask is set anywhere, which means whole machine visible, not pinned to card 0"). A GEAK launcher that writes ROCR_VISIBLE_DEVICES = gpu_pin["value"] now restricts a whole-machine run to cards 0-3, where before it inherited everything.
| """ | ||
| 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: |
There was a problem hiding this comment.
Loop nesting makes source precedence dominate variable precedence.
The source loop is outer and the variable loop inner, so all three vars are checked under baseline_recipe before process_env is reached. That means a recipe-level CUDA_VISIBLE_DEVICES/HIP_VISIBLE_DEVICES silently overrides a real process-level ROCR_VISIBLE_DEVICES pin.
Example: a hand-authored recipe carries benchmark.envs.CUDA_VISIBLE_DEVICES: "0" (a common leftover on a CUDA-derived YAML) while the run is launched with ROCR_VISIBLE_DEVICES=6,7. This returns {'var': 'CUDA_VISIBLE_DEVICES', 'ids': [0]}, and _resolve_handoff_gpu_ids takes the non-ROCR branch and emits gpu_ids="0" verbatim — pinning GEAK to physical card 0 while the run owns 6 and 7.
The docstring's claimed "ROCR_VISIBLE_DEVICES before HIP/CUDA" precedence only holds within a single source. If that ordering is meant to be global, the loops need to be swapped (or the precedence documented as source-major).
| """ | ||
| width = max(int(tp or 1), 1) | ||
| ids = list((gpu_pin or {}).get("ids") or []) | ||
| if not ids: |
There was a problem hiding this comment.
A non-numeric but non-blank mask yields a truthy pin with empty ids, and falls back to card 0.
_resolve_gpu_pin returns early on any non-blank string, but _parse_device_list returns [] for anything non-numeric. Two real cases:
- ROCm accepts GPU UUID masks:
ROCR_VISIBLE_DEVICES=GPU-a1b2c3,GPU-d4e5f6. - YAML parses
ROCR_VISIBLE_DEVICES: [0,1]as a list, andstr()of it is"[0, 1]".
In both, handoff["gpu_pin"] is present with ids: [], this branch sees not ids and returns "0,...,tp-1" — landing every GEAK server on card 0, the exact default the PR exists to eliminate. Meanwhile a consumer reading gpu_pin["ids"] gets nothing, and one writing gpu_pin["value"] exports the literal "[0, 1]".
test_gpu_ids_never_empty_for_a_blank_mask enshrines the fallback as intended, but "never empty" and "never silently card 0" are different goals. Suggest treating a non-blank-but-unparseable mask as an error/warning rather than as unpinned.
| #: ``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, ...] = ( |
There was a problem hiding this comment.
The "every layer agrees on the pin" claim does not hold for a present-but-empty mask.
gpu_pool._visible_device_mask returns ([], present=True) for VAR="" and gate.detect_gpu_count() derives 0 GPUs from it. This resolver instead skips a blank value and falls through to the next variable.
Divergence: ROCR_VISIBLE_DEVICES="" together with a stale HIP_VISIBLE_DEVICES="2,3". The orchestrator believes it has zero GPUs, while _resolve_gpu_pin returns {'var': 'HIP_VISIBLE_DEVICES', 'ids': [2, 3]} and hands GEAK two cards. test_pin_skips_blank_values locks the divergent behaviour in.
Either match the other layers (blank = "zero visible", stop) or reword this comment so it doesn't claim an agreement that isn't there.
| # 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 verbatim, else 0..tp-1. | ||
| "gpu_ids": _resolve_handoff_gpu_ids(gpu_pin=gpu_pin, tp=int(os.environ.get("TP", "1") or 1)), |
There was a problem hiding this comment.
gpu_ids logical indices are computed against the recipe mask, but the child inherits the process mask.
The GEAK subprocess is launched with runner_env = dict(os.environ) (line ~1171); nothing overrides runner_env["ROCR_VISIBLE_DEVICES"] from the resolved pin. So whenever the recipe mask and the process mask differ, the logical indices resolve against the wrong set.
This is exactly the case the new test_pin_prefers_recipe_over_process_env encodes: recipe ROCR_VISIBLE_DEVICES="6", process ROCR_VISIBLE_DEVICES="0". The handoff advertises gpu_pin.ids=[6] and gpu_ids="0", but the child inherits ROCR=0, so HIP index 0 is physical card 0 — the handoff claims card 6 while every server GEAK launches sits on card 0. If the recipe mask is to be authoritative, the phase should also export it into runner_env.
Separately, tp and gpu_ids can now disagree. handoff["tp"] a few lines up is still int(os.environ.get("TP")) raw, while gpu_ids is clamped to the mask width. With TP=8 exported on a 4-GPU pod the materializer clamps the recipe to TP=4 / ROCR="0,1,2,3" (test_baseline_param_overrides.py:385), so gpu_ids becomes "0,1,2,3" while tp stays 8 — GEAK launches sglang with --tp 8 and four visible cards and fails to load weights. Pre-PR both fields derived from the same $TP and could not disagree; tp should now come from the same resolved mask/recipe as gpu_ids.
| # 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"), |
There was a problem hiding this comment.
These fields are only recorded on the crash-recovery path, so the stated goal is not met.
_geak_reconstruct_from_disk has exactly one call site (line ~906), guarded by if not has_result. A GEAK run that finishes normally and writes geak_result={'status': 'no_gain'} takes the has_result path at line ~900 and returns without ever calling it.
So for every completed run — including the no_gain outcome this comment names — the breakdown contains no gpu_ids/gpu_pin, and a foreign-tenant collision still can't be told apart from a real no_gain. Only a crashed run with no committed result gets the fields. The same two keys need to be recorded on the has_result path too.
| # `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"), |
There was a problem hiding this comment.
Writes an explicit null for unpinned runs and for every v1/v2 handoff.
handoff.get("gpu_pin") yields None for any pre-v3 handoff on disk — still produced by any session resumed from before this deploy — and also for a genuinely unpinned v3 run. A breakdown reader then cannot distinguish "no pin was set" from "this handoff predates the field" from "the pin resolved empty".
This also contradicts the writer's own if gpu_pin: guard in kernel.py and the "omitted when nothing is pinned" contract documented in docs/components/geak.md. Suggest inserting the keys only when present, mirroring the writer.
| 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 |
There was a problem hiding this comment.
The end-to-end case never exercises the recipe branch — the branch that wins in production.
This SharedState leaves baseline_config_path unset, so _read_recipe_bench_envs returns {}, _resolve_gpu_pin falls through to the process env, and the source == "process_env" assertion passes.
In every real run baseline_config_path points at a materialized recipe whose benchmark.envs always contains an autofilled ROCR_VISIBLE_DEVICES, so the asserted process_env path is effectively dead code in production — which is why the HIP-pin regression flagged on _resolve_gpu_pin is invisible to the suite.
A case that writes a real materialized YAML (with the autofilled ROCR mask) alongside a HIP_VISIBLE_DEVICES process pin, and asserts the HIP pin survives, would fail today. That is the test this PR most needs.
|
|
||
| | 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 | |
There was a problem hiding this comment.
Two inaccuracies in this row.
- "capped at
tp" is only true for the ROCR branch (min(len(ids), width)in_resolve_handoff_gpu_ids). The HIP/CUDA branch forwards the mask with no cap:HIP_VISIBLE_DEVICES=4,5,6,7withTP=2yieldsgpu_ids="4,5,6,7"— four devices for a two-way tensor-parallel launch. - "verbatim" is not accurate either — the ids go through
_parse_device_list, which deduplicates and re-serializes.HIP_VISIBLE_DEVICES=" 4, 4 ,5"produces"4,5", not the original string. The_resolve_handoff_gpu_idsdocstring makes the same "VERBATIM" claim.
Review of #1321 found the resolver defeated by the very autofill it had to account for. `materialize_config_with_envs` writes `ROCR_VISIBLE_DEVICES=0..tp-1` into `benchmark.envs` whenever the mask is absent or narrower than TP, and that materialized recipe is what `state.baseline_config_path` points at by the time KERNEL runs. Reading the recipe first therefore meant the synthetic mask always won: a `HIP`-pinned run shipped `gpu_ids="0,1"` where it used to ship `"4,5"`, and the documented "no mask anywhere => omit gpu_pin" case was unreachable in production. Both re-created the card-0 collision this change exists to remove. - Resolve variable-major (ROCR -> HIP -> CUDA), process env before recipe within each variable. Source-major was wrong in both directions: a leftover recipe CUDA key outranked a real process ROCR pin, and the autofill outranked everything. - Ignore a recipe ROCR value byte-identical to the `0..tp-1` the materializer would have synthesized. - Forward a recipe-only pin as ABSOLUTE ids: the child is launched with `dict(os.environ)` and never inherits that mask, so logical indices would resolve against the wrong set. - Count mask TOKENS, not parsed ids, so a UUID mask maps to the right number of logical slots instead of silently falling back to card 0. Accept a YAML sequence mask. Carry the count in `gpu_pin`. - Take `tp` from the same resolved recipe as `gpu_ids`, so a stale `$TP=8` on a 4-card pod can no longer ship `tp: 8` beside four `gpu_ids`. - Parse the recipe once and share it between `bench_protocol` and the pin. - Record `gpu_ids`/`gpu_pin` on the collector's `has_result` path too — the `no_gain` outcome that needs disambiguating is a COMPLETED run, and the fields were only being written on crash recovery. Insert the keys only when present, matching the writer, so a pre-v3 handoff is not reported as null. - Correct the docs row: the `tp` cap applies to the ROCR branch only, and ids are re-serialized rather than passed through verbatim. Tests: the end-to-end case now builds a real materialized recipe carrying the autofilled mask alongside a HIP process pin; it fails on the previous commit with `gpu_pin.var == ROCR_VISIBLE_DEVICES` and passes here.
CI E2E report — ❌ Timeout
|
2 participants
Fixes #1312.
Problem
geak/handoff.jsonnever told GEAK which cards the run owns.phases/kernel.pyresolvedgpu_idsfromHIP_VISIBLE_DEVICES/CUDA_VISIBLE_DEVICESonly —ROCR_VISIBLE_DEVICES,the canonical ROCm pin honoured everywhere else in this repo (
bus/gpu_pool.py,policy/gate.py,cli/preflight.py), was not consulted. So a ROCR-pinned run fell back to0..tp-1and every server GEAK launches (baseline, profile, config-tuning validation) went tophysical GPU 0, OOM-ing against whatever else holds that card and reporting a plausible-looking
no_gain.Fix
ROCR→HIP→CUDAprecedence, recipebenchmark.envsbefore process env.
gpu_pinto the handoff (schema_version3): the absolute ids plus the var they camefrom, for consumers that write
ROCR_VISIBLE_DEVICESthemselves. Omitted when nothing ispinned — that means "whole machine", not "card 0".
gpu_idskeeps its existing HIP-level meaning (HIP indexes into the ROCr-visible set), sinceGEAK's sglang/vllm adapters export it as
HIP_VISIBLE_DEVICES. Only new behaviour: it is nowclamped to the mask when
tpovershoots it.GEAK-side follow-up (not this repo): launch paths that write
ROCR_VISIBLE_DEVICESshould readgpu_pin["value"], asadapters/launchers/magpie.shalready does.Tests
New
test_geak_handoff_gpu_pin.pyplus an end-to-end case intest_geak_resume_recovery.py(
ROCR=7→gpu_pin.ids=[7],gpu_ids="0").pytest src/hyperloom/orchestrator91 passed /14 skipped; ruff and
pylint --errors-onlyclean. Not verified on hardware — I have no repro ofthe xDiT run, so the "server actually lands on the pinned card" leg is untested.