fix(rlix): gate MILES wake on whole-GPU residual memory (R02-01)

[howard989]

Note (2026-07-05, maintainer edit): the title/body of this PR had been accidentally swapped with the paired MILES-side PR rlops/miles#5 (each described the other's diff, making both cross-references self-referential). Corrected below to describe the actual rlix-side change, with new audit evidence appended. The original text is preserved in the PR edit history. No commits were added to the contributor branch.

What

Replace the hardcoded free-memory gate in MilesPipeline._wait_for_overlap_engines_offloaded() with a configurable whole-GPU residual used-memory hard gate.

The threshold is controlled by:

MILES_MAX_RESIDUAL_GPU_MEM_GB

Default: 13.0 GiB

Sender side: rlops/miles PR #5.

Why

Per @taoluo review (R02-01, on the superseded #11): "free memory is gpu-model dependent e.g. 24gb vs 80gb gpu. it would be more robust to check the residual memory allocation."

The base used target_free_gb = 20.0 against nvidia-smi --query-gpu=memory.free, which is GPU-capacity dependent and not portable — and it was warn-only (never raised). The condition before wake_up should be: "the previous tenant released enough GPU memory", i.e. residual used memory.

The gate intentionally uses whole-GPU memory.used over the overlap GPUs. This is broader than SGLang process-tree memory: it also catches non-SGLang co-tenants such as Megatron, Miles, vLLM, or orphan processes that can still occupy VRAM and block the next tenant.

The paired MILES PR logs SGLang-specific process-resident and /server_info diagnostics for attribution, but those diagnostics are not the hard gate.

What This PR Does

• rlix/utils/env.py
  • adds parse_env_positive_float
• rlix/pipeline/miles_coordinator.py
  • forwards MILES_MAX_RESIDUAL_GPU_MEM_GB into per-pipeline runtime env
  • parses it with default 13.0
  • passes it to shrink_engines(post_sleep_vram_threshold_gb=...)
• rlix/pipeline/miles_pipeline.py
  • removes the old target_free_gb = 20.0 free-memory polling gate (warn-only)
  • keeps state == offloaded polling as the liveness gate
  • probes whole-GPU nvidia-smi --query-gpu=memory.used over overlap GPUs
  • raises RuntimeError if whole-GPU residual exceeds MILES_MAX_RESIDUAL_GPU_MEM_GB
  • fail-open (warning) if the nvidia-smi probe is unavailable
• new tests: tests/test_env_utils.py, tests/test_miles_residual_threshold_wiring.py

Gate Semantics

max(nvidia-smi memory.used over overlap GPUs) <= MILES_MAX_RESIDUAL_GPU_MEM_GB

This is a whole-GPU availability gate, not a SGLang-only attribution gate. On failure the error message explicitly notes the residual may come from non-SGLang co-tenants (Megatron / Miles / vLLM / orphan processes).

Default 13.0 Rationale

13.0 GiB is a temporary smoke-safe whole-GPU threshold, not a model-derived final value.

Smokes showed:

RTX PRO 6000 96GB:
  whole-GPU residual peak: 11.95-11.97 GiB
  SGLang process-resident diagnostic: 2.516-2.535 GiB

RTX 5090 smoke:
  whole-GPU residual peak: ~12.5 GiB

A 10.0 GiB threshold would false-fail current smokes. 13.0 GiB keeps the gate enabled without failing on the known Megatron train-offload residual. Intentionally temporary — see roadmap below.

Offload Audit Evidence (added 2026-07-05)

Independent audit on 4× RTX 5090 32GB (cu130 image), same dual overlap topology (P1 train=[0] infer=[0,1,2] / P2 train=[3] infer=[1,2,3]), with 1 Hz whole-GPU + per-process nvidia-smi sampling. Both smokes EXIT_CODE=0. Units in this section: GiB (nvidia-smi MiB / 1024), matching the gate's own log output.

1. The ~11-12 GiB residual is root-caused and quantified (current MILES_TMS_HOOK_MODE=torch):

whole-GPU, from actor print_memory (the train GPU also hosts one sleeping SGLang engine):
Megatron post-INIT  sleep: 5.91 -> 3.15 GiB   (tms pause releases exactly 2.76 GiB: init-time weight buffers)
Megatron post-TRAIN sleep: 13.83 -> 11.07 GiB (pause STILL releases only the same 2.76 GiB)

Sampler decomposition of that post-train 11.07 GiB (simultaneous sample, MiB): whole-GPU 11,337 = Megatron process 9,556 + sleeping SGLang engine 1,760 (+ ~21 unattributed). The Megatron process residual grows 3.64 → 9.33 GiB across the first train step: ~5.7 GiB of lazily created optimizer state (Adam m/v, fp32 master weights, grad buffers; the torch-allocator share of the growth is allocated_GB 4.62 → 8.7) sits outside tms-pausable allocations and never offloads. Same growth pattern as this PR's own smokes (6.25 → 11.95-11.97 GiB); exact figures shift with hardware/stack.

2. The offload machinery itself is healthy. Isolated micro-tests: SGLang release_memory_occupation residual 1.6-1.7 GiB (CUDA context + NCCL + scheduler overhead); bare torch_memory_saver.pause() residual ~0.6 GiB (CUDA context).

3. preload hook no longer segfaults on Blackwell with the cu130 image (mock + full dual smoke both pass). Zero-code-change experiment — MILES_TMS_HOOK_MODE=preload + the existing guard escape hatch MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1 (actor_group.py already injects LD_PRELOAD + TMS_INIT_ENABLE_CPU_BACKUP=1 under --offload-train):

Megatron post-TRAIN sleep releases: 2.76 -> 8.80 GiB
Megatron process residual:          9.33 -> 3.29 GiB (stable across cycles, no growth)
Gate-GPU whole residual (cycle 2):  11.07 -> 5.02 GiB (sampler MiB: 5,141 = 3,368 Megatron + 1,740 engine + ~33)
sleep cost: 0.5 s -> ~3 s (CPU-backup copy; still seconds-level)

4. Remaining tail: even under preload, the Megatron process residual (sampler MiB: 1,764 → 3,368) grows 1.72 → 3.29 GiB across the first train step — ~1.6 GiB of train-created state still escapes pause. The suspected source — not yet confirmed — is allocations made inside the torch_memory_saver.disable() weight-sync window (miles actor.py). Separate follow-up.

5. Threshold roadmap for MILES_MAX_RESIDUAL_GPU_MEM_GB:

13.0  today (torch hook; tolerates the un-offloaded optimizer state)
~7.0  once the preload hook mode lands (measured floor ~5.0 GiB + margin)
~3.0  after the disable-window tail is fixed (floor ≈ engine ~1.7 + Megatron ctx ~0.6 + margin)

6. Signal-choice validation: on containerized hosts (e.g. vast.ai) nvidia-smi --query-compute-apps reports host-namespace PIDs, so the MILES per-process diagnostic (miles#5 gpu_probe) fail-opens there. Of the two candidate signals, whole-GPU memory.used — as implemented here — remains reliable in containers, unlike the per-process diagnostic. This PR's design choice is confirmed.

Diff Baseline Note

This is a clean branch off latest zhenyu/miles-mvp-e2e. The closed #11 used intermediate thresholds while we investigated signal choice. This PR's effective change is:

20.0 GiB free-memory gate (warn-only)
  -> 13.0 GiB whole-GPU residual used-memory hard gate

Tests

python -m pytest -q tests/test_env_utils.py tests/test_miles_residual_threshold_wiring.py

Result:

6 passed

E2E Verification

Vast dual smoke with paired MILES branch:

SGLang diagnostic:
process_resident_max=2.516-2.535 GiB
whole_gpu_threshold=13.000 GiB

Whole-GPU hard gate:
whole-GPU mem used max=6.25 / 6.27 GiB across overlap GPUs [0]/[3]
whole-GPU mem used max=11.95 / 11.97 GiB across overlap GPUs [0]/[3]
threshold=13.00 GiB

mp2 training loop complete
mp1 training loop complete
shutdown_hard complete for both pipelines
EXIT_CODE=0

Plus the 2026-07-05 audit smokes above (torch-hook and preload variants, both EXIT_CODE=0).

Known SharedStorage actor unavailable warnings and shutdown-time RolloutManager 500 / RemoteProtocolError teardown noise may appear. Training completed, both pipelines reached shutdown_hard, and EXIT_CODE=0.

Scope

Gate signal + configurability only. No model-size-derived threshold. No Megatron train-offload fix in this PR.

Follow-ups (audit-informed): (1) enable preload hook mode on cu130+ Blackwell and lower the default to ~7.0; (2) fix the torch_memory_saver.disable() window tail, re-measure, lower to ~3.0.

Refs: plans/m11-review.review-report/R02.md (R02-01, MEDIUM).

[zhenyulincs]

@howard989 Thanks for the gate — the whole-GPU memory.used signal choice turned out to be exactly right (see the last bullet). Two maintainer updates on this PR, no commits touched on your branch:

1. Description fix: this PR's title/body had been swapped with the paired rlops/miles#5 (each described the other's diff). I restored the correct rlix-side description here; the original text remains in the edit history.

2. Audit evidence merged into the description (2026-07-05, 4× RTX 5090 cu130, same dual topology, both smokes EXIT_CODE=0; units GiB = nvidia-smi MiB / 1024) — highlights:

• The ~11-12 GiB residual is root-caused: with MILES_TMS_HOOK_MODE=torch, tms.pause() only ever releases the 2.76 GiB of init-time weight buffers; the ~5.7 GiB of optimizer state (Adam m/v, fp32 master, grad buffers) created lazily at the first optimizer.step() is never pausable. Post-train sleep on the train GPU: 13.83 → 11.07 GiB whole-GPU; simultaneous sampler decomposition 11,337 MiB = 9,556 (Megatron process) + 1,760 (sleeping engine) + ~21.
• preload hook mode no longer segfaults on Blackwell with the cu130 image. A zero-code-change run (MILES_TMS_HOOK_MODE=preload + MILES_TMS_ALLOW_PRELOAD_ON_BLACKWELL=1) cuts the gate-GPU residual from 11.07 to 5.02 GiB, stable across cycles.
• Suggested threshold roadmap: 13.0 today → ~7.0 once preload lands → ~3.0 after the remaining ~1.6 GiB tail — suspected (not yet confirmed) to be allocations inside the torch_memory_saver.disable() weight-sync window — is eliminated.
• On containerized hosts, --query-compute-apps reports host-namespace PIDs, so the per-process diagnostic in miles#5 fail-opens there — of the two candidate signals, the whole-GPU gate in this PR is the one that remains reliable in containers.

Full numbers and methodology are in the updated PR description ("Offload Audit Evidence" section). The gate as implemented needs no code change for this PR to land; the roadmap items are follow-ups.