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

rlix/pipeline/miles_coordinator.py

@@ -18,6 +18,7 @@
 import asyncio
 import logging
 import math
+import os
 import threading
 import time
 from copy import deepcopy
@@ -37,7 +38,7 @@
     get_pipeline_namespace,
 )
 from rlix.protocol.validation import validate_pipeline_id
-from rlix.utils.env import pipeline_identity_env_vars
+from rlix.utils.env import parse_env_positive_float, pipeline_identity_env_vars
 from rlix.utils.ray import get_actor_or_raise
 
 logger = logging.getLogger(__name__)
@@ -60,9 +61,13 @@ def _build_pipeline_env_vars(*, pipeline_id: str, ray_namespace: str) -> Dict[st
     runtime_env. Reads ``RLIX_CONTROL_PLANE`` from the environment so
     actors inside an existing pipeline preserve the inherited value.
     """
-    return pipeline_identity_env_vars(
+    env_vars = pipeline_identity_env_vars(
         pipeline_id=str(pipeline_id), ray_namespace=str(ray_namespace)
     )
+    for key in ("MILES_MAX_RESIDUAL_GPU_MEM_GB",):
+        if (value := os.environ.get(key)) is not None:
+            env_vars[key] = value
+    return env_vars
 
 
 class MilesCoordinator(Coordinator):
@@ -431,7 +436,30 @@ def _shrink_workers(self, engine_indices: Set[int]) -> None:
             if rollout_manager is None:
                 raise RuntimeError("resource registration missing for shrink")
         # RPC outside the lock.
-        ray.get(rollout_manager.shrink_engines.remote(sorted(engine_indices)))
+        # Whole-GPU residual threshold (GiB). Miles shrink_engines logs
+        # SGLang per-process/server_info attribution diagnostics with this
+        # value; the hard whole-GPU gate runs in MilesPipeline after the
+        # engines report offloaded. Default 13.0 is a temporary smoke-safe
+        # value based on observed whole-GPU residuals with the known Megatron
+        # train-offload gap; lower it after that follow-up is fixed and
+        # re-measured.
+        residual_threshold_gb = parse_env_positive_float(
+            "MILES_MAX_RESIDUAL_GPU_MEM_GB", 13.0
+        )
+        shrunk = ray.get(
+            rollout_manager.shrink_engines.remote(
+                sorted(engine_indices),
+                post_sleep_vram_threshold_gb=residual_threshold_gb,
+            )
+        )
+        logger.info(
+            "[MilesCoordinator] shrink_engines complete pipeline_id=%s "
+            "engine_indices=%s whole_gpu_residual_threshold=%.1f GB "
+            "(whole-GPU hard gate runs in MilesPipeline; Miles shrink logs SGLang diagnostics)",
+            self._pipeline_id,
+            sorted(shrunk),
+            residual_threshold_gb,
+        )
         # Commit under lock.
         with self._resize_sync_lock:
             self._active_engine_indices -= engine_indices

rlix/pipeline/miles_pipeline.py

@@ -51,6 +51,7 @@
     SCHEDULER_ACTOR_NAME,
     get_pipeline_namespace,
 )
+from rlix.utils.env import parse_env_positive_float
 from rlix.utils.ray import get_actor_or_raise
 
 logger = logging.getLogger(__name__)
@@ -503,16 +504,15 @@ def _init_phase_b_infer(self) -> None:
 
     def _wait_for_overlap_engines_offloaded(self, allocated_train_gpus, *, timeout_s: float = 60.0) -> None:
         """After scheduler grants actor_train, poll the rollout manager
-        until the engines on overlap GPUs have transitioned to ``offloaded``
-        AND the OS-reported GPU memory is actually free. SGLang's HTTP
-        ``/release_memory_occupation`` 200 OK + state="offloaded" do not
-        by themselves guarantee the CUDA driver has returned the memory
-        to the OS pool — the wake_up in the next-process train actor
-        would then OOM. Verify actual GPU mem free by parsing
-        ``nvidia-smi --query-gpu=memory.free`` on the same node, since
-        miles' single-node smoke topology has driver+actors+engines all
-        on the head node and ``CUDA_VISIBLE_DEVICES`` is the per-actor
-        slice of the shared physical pool.
+        until the engines on overlap GPUs have transitioned to ``offloaded``.
+
+        After the rollout engines report ``offloaded`` / ``shell``, this
+        method enforces the whole-GPU residual availability gate using
+        ``nvidia-smi memory.used`` over the overlap GPUs. This intentionally
+        counts non-SGLang co-tenants (Megatron/Miles/vLLM/orphan processes)
+        so residual GPU pressure is not missed; Miles logs SGLang
+        per-process/server_info diagnostics during ``shrink_engines`` for
+        attribution.
         """
         rollout_manager = getattr(self, "_rollout_manager", None)
         if rollout_manager is None:
@@ -573,52 +573,39 @@ def _wait_for_overlap_engines_offloaded(self, allocated_train_gpus, *, timeout_s
                 timeout_s, target_indices, uniq,
             )
 
-        # Phase 2: probe nvidia-smi for OS-level free memory on the
-        # overlap GPU IDs. The train actor will need ~3.7 GB for the
-        # 0.5B model + a few GB for activations; aim for ≥20 GB free
-        # before we let _before_training proceed to wake_up.
-        target_free_gb = 20.0
-        deadline2 = time.time() + float(timeout_s)
-        last_min_free_gb: Optional[float] = None
-        nvidia_smi_unavail_count = 0
-        while time.time() < deadline2:
-            min_free_gb = self._probe_min_free_gpu_mem_gb(target_gpu_ids)
-            if min_free_gb is None:
-                # F5 (m11-review.review-report.md §2): nvidia-smi unavailable
-                # or unparseable. Was logged at DEBUG only — promoted to INFO
-                # so operators see the fallback without flipping log levels.
-                # If this fires repeatedly across sessions, it's a hardware
-                # / image regression worth investigating (driver missing,
-                # nvidia-smi path changed, etc.).
-                nvidia_smi_unavail_count += 1
-                logger.info(
-                    "_wait_for_overlap_engines_offloaded: nvidia-smi probe "
-                    "unavailable (count=%d); falling back to 3s grace sleep",
-                    nvidia_smi_unavail_count,
-                )
-                time.sleep(3.0)
-                return
-            last_min_free_gb = min_free_gb
-            if min_free_gb >= target_free_gb:
-                logger.info(
-                    "_wait_for_overlap_engines_offloaded: OS-level GPU mem free "
-                    "min=%.2f GB across overlap GPUs %s (target=%.1f GB)",
-                    min_free_gb, target_gpu_ids, target_free_gb,
-                )
-                return
-            time.sleep(0.5)
-        logger.warning(
-            "_wait_for_overlap_engines_offloaded: free-mem timeout after %.1fs; "
-            "min_free_gb=%.2f below %.1f GB target on GPUs %s — wake_up may OOM",
-            timeout_s,
-            last_min_free_gb if last_min_free_gb is not None else float("nan"),
-            target_free_gb,
+        # Phase 2: whole-GPU residual availability hard gate. This is
+        # intentionally broader than the SGLang process tree: it catches
+        # non-child/orphan/co-tenant VRAM (Megatron/Miles/vLLM/etc.) on the
+        # GPUs that must be clear before actor_train wakes up.
+        threshold_gb = parse_env_positive_float("MILES_MAX_RESIDUAL_GPU_MEM_GB", 13.0)
+        max_used_gb = self._probe_max_used_gpu_mem_gb(target_gpu_ids)
+        if max_used_gb is None:
+            logger.warning(
+                "_wait_for_overlap_engines_offloaded: nvidia-smi probe unavailable; "
+                "skipping whole-GPU residual hard gate (fail-open)"
+            )
+            return
+        logger.info(
+            "_wait_for_overlap_engines_offloaded: whole-GPU mem used max=%.2f GB "
+            "across overlap GPUs %s threshold=%.2f GB",
+            max_used_gb,
             target_gpu_ids,
+            threshold_gb,
         )
+        if max_used_gb > threshold_gb:
+            raise RuntimeError(
+                f"Post-offload whole-GPU residual {max_used_gb:.2f} GiB exceeds "
+                f"threshold {threshold_gb:.2f} GiB on overlap GPUs {target_gpu_ids}. "
+                "GPU was not sufficiently cleared before wake_up; this whole-GPU "
+                "check may be caused by non-SGLang co-tenants such as "
+                "Megatron/Miles/vLLM/orphan processes. SGLang per-engine "
+                "process-resident and server_info residual diagnostics are logged "
+                "engine-side for attribution."
+            )
 
     @staticmethod
-    def _probe_min_free_gpu_mem_gb(gpu_ids: list[int]) -> Optional[float]:
-        """Return the minimum free GPU memory (GB) across ``gpu_ids`` as
+    def _probe_max_used_gpu_mem_gb(gpu_ids: list[int]) -> Optional[float]:
+        """Return the maximum used GPU memory (GB) across ``gpu_ids`` as
         reported by ``nvidia-smi``. Returns ``None`` if nvidia-smi is
         not available or output cannot be parsed.
         """
@@ -634,7 +621,7 @@ def _probe_min_free_gpu_mem_gb(gpu_ids: list[int]) -> Optional[float]:
                 [
                     "nvidia-smi",
                     f"--id={','.join(str(g) for g in gpu_ids)}",
-                    "--query-gpu=memory.free",
+                    "--query-gpu=memory.used",
                     "--format=csv,noheader,nounits",
                 ],
                 stderr=subprocess.STDOUT,
@@ -643,18 +630,18 @@ def _probe_min_free_gpu_mem_gb(gpu_ids: list[int]) -> Optional[float]:
         except (subprocess.SubprocessError, OSError) as exc:
             logger.debug("nvidia-smi probe failed: %r", exc)
             return None
-        free_mibs: list[float] = []
+        used_mibs: list[float] = []
         for line in out.strip().splitlines():
             line = line.strip()
             if not line:
                 continue
             try:
-                free_mibs.append(float(line))
+                used_mibs.append(float(line))
             except ValueError:
                 continue
-        if not free_mibs:
+        if not used_mibs:
             return None
-        return min(free_mibs) / 1024.0
+        return max(used_mibs) / 1024.0
 
     def _before_training(self, step: int) -> None:
         if not self._initialized:

rlix/utils/env.py

@@ -52,3 +52,21 @@ def parse_env_timeout_s(env_key: str, default_s: Optional[float] = None) -> Opti
     except ValueError as exc:
         raise RuntimeError(f"{env_key} must be a number, got: {raw!r}") from exc
     return None if value <= 0 else value
+
+
+def parse_env_positive_float(env_key: str, default: float) -> float:
+    """Read a positive float from an env var; fail-fast on invalid values.
+
+    Returns *default* when the env var is unset. Raises RuntimeError if the
+    value cannot be parsed as a number, or if the parsed value is <= 0.
+    """
+    raw = os.environ.get(env_key)
+    if raw is None:
+        return float(default)
+    try:
+        value = float(raw)
+    except ValueError as exc:
+        raise RuntimeError(f"{env_key} must be a number, got: {raw!r}") from exc
+    if value <= 0.0:
+        raise RuntimeError(f"{env_key} must be > 0, got: {value!r}")
+    return value

tests/test_env_utils.py

@@ -0,0 +1,56 @@
+import importlib
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+RLIX_ROOT = REPO_ROOT / "rlix"
+
+
+def _load_env_module(monkeypatch):
+    for module_name in list(sys.modules):
+        if module_name == "rlix" or module_name.startswith("rlix."):
+            monkeypatch.delitem(sys.modules, module_name, raising=False)
+
+    package_roots = {
+        "rlix": RLIX_ROOT,
+        "rlix.utils": RLIX_ROOT / "utils",
+    }
+    for module_name, module_path in package_roots.items():
+        package_module = types.ModuleType(module_name)
+        package_module.__path__ = [str(module_path)]  # type: ignore[attr-defined]
+        monkeypatch.setitem(sys.modules, module_name, package_module)
+
+    return importlib.import_module("rlix.utils.env")
+
+
+def test_parse_env_positive_float_uses_default_when_unset(monkeypatch):
+    env = _load_env_module(monkeypatch)
+    monkeypatch.delenv("MILES_MAX_RESIDUAL_GPU_MEM_GB", raising=False)
+
+    assert env.parse_env_positive_float("MILES_MAX_RESIDUAL_GPU_MEM_GB", 2.0) == 2.0
+
+
+def test_parse_env_positive_float_reads_override(monkeypatch):
+    env = _load_env_module(monkeypatch)
+    monkeypatch.setenv("MILES_MAX_RESIDUAL_GPU_MEM_GB", "40.5")
+
+    assert env.parse_env_positive_float("MILES_MAX_RESIDUAL_GPU_MEM_GB", 2.0) == 40.5
+
+
+def test_parse_env_positive_float_rejects_non_positive(monkeypatch):
+    env = _load_env_module(monkeypatch)
+    monkeypatch.setenv("MILES_MAX_RESIDUAL_GPU_MEM_GB", "0")
+
+    with pytest.raises(RuntimeError, match="must be > 0"):
+        env.parse_env_positive_float("MILES_MAX_RESIDUAL_GPU_MEM_GB", 2.0)
+
+
+def test_parse_env_positive_float_rejects_non_numeric(monkeypatch):
+    env = _load_env_module(monkeypatch)
+    monkeypatch.setenv("MILES_MAX_RESIDUAL_GPU_MEM_GB", "not-a-number")
+
+    with pytest.raises(RuntimeError, match="must be a number"):
+        env.parse_env_positive_float("MILES_MAX_RESIDUAL_GPU_MEM_GB", 2.0)

tests/test_miles_residual_threshold_wiring.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import ast
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+
+
+def _is_name(node: ast.AST, name: str) -> bool:
+    return isinstance(node, ast.Name) and node.id == name
+
+
+def _is_attr(node: ast.AST, attr: str) -> bool:
+    return isinstance(node, ast.Attribute) and node.attr == attr
+
+
+def test_miles_uses_whole_gpu_residual_threshold() -> None:
+    source = (REPO_ROOT / "rlix" / "pipeline" / "miles_coordinator.py").read_text(
+        encoding="utf-8"
+    )
+    tree = ast.parse(source)
+
+    shrink_fn = next(
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == "_shrink_workers"
+    )
+
+    assert any(
+        isinstance(node, ast.Call)
+        and _is_name(node.func, "parse_env_positive_float")
+        and len(node.args) >= 2
+        and isinstance(node.args[0], ast.Constant)
+        and node.args[0].value == "MILES_MAX_RESIDUAL_GPU_MEM_GB"
+        and isinstance(node.args[1], ast.Constant)
+        and node.args[1].value == 13.0
+        for node in ast.walk(shrink_fn)
+    ), "_shrink_workers must parse the residual threshold env var with 13GB default"
+
+    assert any(
+        isinstance(node, ast.Call)
+        and _is_attr(node.func, "remote")
+        and any(
+            kw.arg == "post_sleep_vram_threshold_gb"
+            and _is_name(kw.value, "residual_threshold_gb")
+            for kw in node.keywords
+        )
+        for node in ast.walk(shrink_fn)
+    ), "shrink_engines must receive post_sleep_vram_threshold_gb"
+
+
+def test_miles_coordinator_forwards_residual_threshold_env_var() -> None:
+    source = (REPO_ROOT / "rlix" / "pipeline" / "miles_coordinator.py").read_text(
+        encoding="utf-8"
+    )
+    tree = ast.parse(source)
+
+    build_env_fn = next(
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == "_build_pipeline_env_vars"
+    )
+
+    assert any(
+        isinstance(node, ast.Constant)
+        and node.value == "MILES_MAX_RESIDUAL_GPU_MEM_GB"
+        for node in ast.walk(build_env_fn)
+    ), "_build_pipeline_env_vars must forward MILES_MAX_RESIDUAL_GPU_MEM_GB"
+
+
+def test_miles_pipeline_enforces_whole_gpu_residual_gate() -> None:
+    source = (REPO_ROOT / "rlix" / "pipeline" / "miles_pipeline.py").read_text(
+        encoding="utf-8"
+    )
+    tree = ast.parse(source)
+
+    wait_fn = next(
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.FunctionDef) and node.name == "_wait_for_overlap_engines_offloaded"
+    )
+
+    assert "_probe_max_used_gpu_mem_gb" in source
+    assert "whole-GPU residual hard gate" in source
+    assert "non-SGLang co-tenants" in source
+    assert any(
+        isinstance(node, ast.Call)
+        and _is_name(node.func, "parse_env_positive_float")
+        and len(node.args) >= 2
+        and isinstance(node.args[0], ast.Constant)
+        and node.args[0].value == "MILES_MAX_RESIDUAL_GPU_MEM_GB"
+        and isinstance(node.args[1], ast.Constant)
+        and node.args[1].value == 13.0
+        for node in ast.walk(wait_fn)
+    ), "MilesPipeline hard gate must parse MILES_MAX_RESIDUAL_GPU_MEM_GB with 13GB default"
+    assert any(isinstance(node, ast.Raise) for node in ast.walk(wait_fn)), (
+        "MilesPipeline must raise when whole-GPU residual exceeds the threshold"
+    )