[Bugfix][ROCm] Make Qwen OOM recovery pressure deterministic

tests/buildkite/test_reliability_rocm_marks.py

@@ -7,6 +7,7 @@
 
 import ast
 from pathlib import Path
+from typing import Any, cast
 
 import pytest
 
@@ -55,16 +56,33 @@ def test_fault_injection_cases_declare_mi325(filename: str, num_cards: int) -> N
     assert "connection" in keywords
 
 
-def test_qwen_hard_oom_uses_free_memory_ceiling_and_eager_params() -> None:
+def test_qwen_oom_uses_free_memory_ceiling_and_eager_params() -> None:
     case_source = (RELIABILITY_DIR / "test_qwen3_omni.py").read_text(encoding="utf-8")
     case_tree = ast.parse(case_source)
+    case_functions = {node.name: node for node in case_tree.body if isinstance(node, ast.FunctionDef)}
     hard_config = next(
         node for node in case_tree.body if isinstance(node, ast.FunctionDef) and node.name == "_hard_oom_config"
     )
     assert any(isinstance(node, ast.Constant) and node.value == "max_free_bytes" for node in ast.walk(hard_config))
     recover = _assignment_value(case_tree, "OOM_RECOVER_INJECTION_CONFIG")
     assert isinstance(recover, ast.Dict)
-    assert all(not (isinstance(key, ast.Constant) and key.value == "max_free_bytes") for key in recover.keys)
+    recover_items = {
+        key.value: value
+        for key, value in zip(recover.keys, recover.values)
+        if isinstance(key, ast.Constant) and isinstance(key.value, str)
+    }
+    recover_ceiling = recover_items["max_free_bytes"]
+    assert isinstance(recover_ceiling, ast.Call)
+    assert isinstance(recover_ceiling.func, ast.Name)
+    assert recover_ceiling.func.id == "_recover_oom_max_free_bytes"
+
+    ceiling_helper = case_functions["_recover_oom_max_free_bytes"]
+    return_node = next(node for node in ast.walk(ceiling_helper) if isinstance(node, ast.Return))
+    assert isinstance(return_node.value, ast.IfExp)
+    assert ast.unparse(return_node.value.body) == "32 * 1024 * 1024"
+    assert isinstance(return_node.value.orelse, ast.Constant)
+    assert return_node.value.orelse.value is None
+    assert "torch.version" in ast.unparse(return_node.value.test)
     assert "reliability_scenarios_for_eager_oom" in (RELIABILITY_DIR / "helpers.py").read_text(encoding="utf-8")
     assert "eager_oom_stage_ids = (0, 1)" in case_source
 
@@ -88,7 +106,7 @@ def _server_set(class_name: str) -> str:
 
     for class_name in ("GpuOomErrorContract", "GpuOomLarge", "GpuOomConcurrent"):
         assert _server_set(class_name) == "oom"
-    assert _server_set("GpuOomRecovery") == "default"
+    assert _server_set("GpuOomRecovery") == "oom"
 
     functions = {node.name: node for node in oom_tree.body if isinstance(node, ast.FunctionDef)}
     starter = functions["_start_hard_oom"]
@@ -102,7 +120,100 @@ def _server_set(class_name: str) -> str:
         if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "inject_gpu_oom"
     ]
     assert len(inject_calls) == 1
-    assert all(keyword.arg != "max_free_bytes" for keyword in inject_calls[0].keywords)
+    max_free_keyword = next(keyword for keyword in inject_calls[0].keywords if keyword.arg == "max_free_bytes")
+    assert isinstance(max_free_keyword.value, ast.Subscript)
+    assert isinstance(max_free_keyword.value.value, ast.Name)
+    assert max_free_keyword.value.value.id == "OOM_RECOVER_INJECTION_CONFIG"
+    assert isinstance(max_free_keyword.value.slice, ast.Constant)
+    assert max_free_keyword.value.slice.value == "max_free_bytes"
+
+
+def _class_assignment(path: Path, class_name: str, name: str) -> ast.AST | None:
+    tree = ast.parse(path.read_text(encoding="utf-8"))
+    for node in tree.body:
+        if not isinstance(node, ast.ClassDef) or node.name != class_name:
+            continue
+        for statement in node.body:
+            if isinstance(statement, ast.Assign) and any(
+                isinstance(target, ast.Name) and target.id == name for target in statement.targets
+            ):
+                return statement.value
+    return None
+
+
+def _literal_class_assignment(path: Path, class_name: str, name: str) -> object:
+    value = _class_assignment(path, class_name, name)
+    assert value is not None, f"{class_name}.{name} is missing"
+    return ast.literal_eval(value)
+
+
+def test_qwen_oom_recovery_forces_hip_eager_and_leaves_cuda_and_wan_unchanged(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    import torch
+    import yaml
+
+    from tests.dfx.reliability import helpers
+    from tests.helpers.stage_config import modify_stage_config
+
+    qwen_path = RELIABILITY_DIR / "test_qwen3_omni.py"
+    qwen_tree = ast.parse(qwen_path.read_text(encoding="utf-8"))
+    scenarios = _literal_class_assignment(qwen_path, "Qwen3Omni", "scenarios")
+    stage_ids = _literal_class_assignment(qwen_path, "Qwen3Omni", "eager_oom_stage_ids")
+    recover = _assignment_value(qwen_tree, "OOM_RECOVER_INJECTION_CONFIG")
+    assert isinstance(recover, ast.Dict)
+    recover_items = {
+        key.value: value
+        for key, value in zip(recover.keys, recover.values)
+        if isinstance(key, ast.Constant) and isinstance(key.value, str)
+    }
+    oom_source = RELIABILITY_DIR / "scenarios" / "oom.py"
+    server_set = _literal_class_assignment(oom_source, "GpuOomRecovery", "server_set")
+    assert isinstance(scenarios, list)
+    assert isinstance(server_set, str)
+    assert stage_ids == (0, 1)
+    assert isinstance(stage_ids, tuple)
+    assert isinstance(recover_items["max_free_bytes"], ast.Call)
+    ratio = recover_items["target_mem_ratio"]
+    assert isinstance(ratio, ast.Constant)
+    assert ratio.value == 0.95
+    qwen_scenarios = cast(list[dict[str, Any]], scenarios)
+    qwen_stage_ids = cast(tuple[int, ...], stage_ids)
+    assert {item["test_name"] for item in qwen_scenarios} == {
+        "qwen3_omni_reliability_async_chunk",
+        "qwen3_omni_reliability_default",
+    }
+
+    model = helpers.FaultModel()
+    model.model_id = "qwen3_omni"
+    model.scenarios = qwen_scenarios
+    model.eager_oom_stage_ids = qwen_stage_ids
+    deploy = Path("vllm_omni/deploy/qwen3_omni_moe.yaml")
+
+    monkeypatch.setattr(torch.version, "hip", "6.0.0")
+    hip_scenarios = model.scenarios_for_server_set(server_set)
+    assert [item["test_name"] for item in hip_scenarios] == [item["test_name"] for item in qwen_scenarios]
+    for scenario in hip_scenarios:
+        rendered = modify_stage_config(str(deploy), scenario["server_params"]["update"], None)
+        loaded = yaml.safe_load(Path(rendered).read_text(encoding="utf-8"))
+        eager = {int(stage["stage_id"]): stage.get("enforce_eager") for stage in loaded["stages"]}
+        assert eager[0] is True
+        assert eager[1] is True
+        assert eager[2] is False
+
+    monkeypatch.setattr(torch.version, "hip", None)
+    cuda_scenarios = model.scenarios_for_server_set(server_set)
+    assert cuda_scenarios == qwen_scenarios
+    assert all("update" not in item["server_params"] for item in cuda_scenarios)
+
+    wan_path = RELIABILITY_DIR / "test_wan22.py"
+    assert _class_assignment(wan_path, "Wan22", "eager_oom_stage_ids") is None
+    wan = helpers.FaultModel()
+    wan.model_id = "wan22"
+    wan_scenarios = _literal_class_assignment(wan_path, "Wan22", "scenarios")
+    assert isinstance(wan_scenarios, list)
+    wan.scenarios = cast(list[dict[str, Any]], wan_scenarios)
+    assert wan.scenarios_for_server_set("oom") == wan.scenarios_for_server_set("default")
 
 
 def test_nightly_pipeline_does_not_own_fault_injection_files() -> None:

tests/dfx/reliability/helpers.py

@@ -1431,7 +1431,12 @@ class FaultModel:
     oom_large_fail_message: str = "expected request failure during GPU OOM injection"
     oom_concurrent_fail_message: str = "expected concurrent request failure under OOM injection"
 
-    def servers(self, server_set: str) -> list[tuple[str, OmniServerParams]]:
+    def scenarios_for_server_set(self, server_set: str) -> list[dict[str, Any]]:
+        """Return deploy scenarios for one fault server set.
+
+        The ``oom`` set forces ``enforce_eager`` on ``eager_oom_stage_ids``
+        under HIP. Models that do not declare those stages keep their deploy.
+        """
         scenarios = list(self.scenarios)
         if server_set == "lifecycle":
             if self.lifecycle_scenario is not None:
@@ -1446,6 +1451,10 @@ def servers(self, server_set: str) -> list[tuple[str, OmniServerParams]]:
             )
         elif server_set not in ("default", "oom"):
             raise ValueError(f"unknown server set {server_set!r} for {self.model_id}")
+        return scenarios
+
+    def servers(self, server_set: str) -> list[tuple[str, OmniServerParams]]:
+        scenarios = self.scenarios_for_server_set(server_set)
         from tests.dfx.conftest import create_reliability_omni_server_params
 
         params = create_reliability_omni_server_params(scenarios, self.deploy_dir)

tests/dfx/reliability/scenarios/oom.py

@@ -77,7 +77,10 @@ def run(self, model: FaultModel, server: Any, client: Any, signal_name: str | No
 
 
 class GpuOomRecovery(FaultScenario):
-    server_set = "default"
+    # Captured graph pools can serve the probe without a new allocation, so
+    # recovery must use the eager OOM server. Models with no eager stage ids
+    # keep their normal deploy when this set is selected.
+    server_set = "oom"
 
     def run(self, model: FaultModel, server: Any, client: Any, signal_name: str | None = None) -> None:
         model.run_oom_recovery(server, client)

tests/dfx/reliability/test_qwen3_omni.py

@@ -78,9 +78,18 @@ def _hard_oom_config() -> dict[str, Any]:
     }
 
 
+def _recover_oom_max_free_bytes() -> int | None:
+    import torch
+
+    # A proportional target leaves several GiB free on high-memory MI300 cards.
+    # Match the proven HIP hard-OOM ceiling while preserving ratio-only CUDA.
+    return 32 * 1024 * 1024 if getattr(torch.version, "hip", None) else None
+
+
 OOM_RECOVER_INJECTION_CONFIG: dict[str, Any] = {
     "device": _default_oom_device_spec(),
     "target_mem_ratio": 0.95,
+    "max_free_bytes": _recover_oom_max_free_bytes(),
     "hold_seconds": 0,
     "startup_timeout_sec": 20,
     "strict": False,
@@ -443,9 +452,9 @@ def _mix_prompt() -> str:
 def _run_oom_recovery(omni_server_function: Any, online_client_function: Any) -> None:
     """Black-box: under bounded GPU pressure, a heavy mix chat fails; after hog stops, the same request succeeds.
 
-    Uses ``OOM_RECOVER_INJECTION_CONFIG``: same ratio as the hard OOM sidecar, but
-    without ``max_free_bytes``, so pressure can starve the multimodal forward pass
-    without pinning the device down to the hard ceiling.
+    Uses ``OOM_RECOVER_INJECTION_CONFIG``: HIP applies the proven hard-OOM free-memory
+    ceiling because a proportional target alone leaves several GiB free on MI300.
+    CUDA keeps the existing ratio-only behavior.
 
     Fault and recovery phases use raw ``/v1/chat/completions`` HTTP so tests can assert
     on status codes and JSON error payloads (or fault keywords in streamed bodies).
@@ -460,6 +469,7 @@ def _run_oom_recovery(omni_server_function: Any, online_client_function: Any) ->
         hold_seconds=OOM_RECOVER_INJECTION_CONFIG["hold_seconds"],
         startup_timeout_sec=OOM_RECOVER_INJECTION_CONFIG["startup_timeout_sec"],
         strict=OOM_RECOVER_INJECTION_CONFIG["strict"],
+        max_free_bytes=OOM_RECOVER_INJECTION_CONFIG["max_free_bytes"],
     )
     host = omni_server_function.host
     port = omni_server_function.port