fix(disagg): PD transfer-failure injection was silently inert (#35890)

python/sglang/srt/disaggregation/utils.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 import random
 from collections import deque
 from contextlib import nullcontext
@@ -160,16 +159,8 @@ def unified_memory_disagg_move_gate(scheduler):
 #########################
 
 
-def _get_failure_prob() -> float:
-    try:
-        return float(envs.SGLANG_TEST_DISAGG_FAILURE_PROB.get())
-    except Exception:
-        # fallback to legacy env var
-        return float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", "0"))
-
-
 def _poll_with_failure_injection(pollers) -> List[int]:
-    if (failure_prob := _get_failure_prob()) > 0:
+    if (failure_prob := envs.SGLANG_TEST_DISAGG_FAILURE_PROB.get()) > 0:
         return [
             int(KVPoll.Failed) if random.random() < failure_prob else int(poller.poll())
             for poller in pollers

test/registered/amd/disaggregation/test_disaggregation_basic.py

@@ -7,6 +7,7 @@ import openai
 import requests
 from transformers import AutoTokenizer
 
+from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_amd_ci
 from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
 from sglang.test.server_fixtures.disaggregation_fixture import (
@@ -231,8 +232,10 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
             print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.")
             cls.rdma_devices = []
 
-        # set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
-        os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
+        # Inject transfer failures so the retry path is actually exercised.
+        # Entered before the servers start so they inherit it.
+        cls._disagg_failure_ctx = envs.SGLANG_TEST_DISAGG_FAILURE_PROB.override(0.05)
+        cls._disagg_failure_ctx.__enter__()
 
         cls.model = "Qwen/Qwen3-8B"
 
@@ -248,7 +251,7 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
 
     @classmethod
     def tearDownClass(cls):
-        os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB")
+        cls._disagg_failure_ctx.__exit__(None, None, None)
         super().tearDownClass()
 
     @classmethod

test/registered/disaggregation/test_disaggregation_basic.py

@@ -14,6 +14,7 @@ import openai
 import requests
 from transformers import AutoTokenizer
 
+from sglang.srt.environ import envs
 from sglang.srt.mem_cache.kv_cache_builder import BACKUP_ONLY_HICACHE_RATIO
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
@@ -188,14 +189,16 @@ class TestDisaggregationMooncakeFailure(PDDisaggregationServerBase):
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
-        # set DISAGGREGATION_TEST_FAILURE_PROB to simulate failure
-        os.environ["DISAGGREGATION_TEST_FAILURE_PROB"] = "0.05"
+        # Inject transfer failures so the retry path is actually exercised.
+        # Entered before launch_all() so the server subprocesses inherit it.
+        cls._disagg_failure_ctx = envs.SGLANG_TEST_DISAGG_FAILURE_PROB.override(0.05)
+        cls._disagg_failure_ctx.__enter__()
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.launch_all()
 
     @classmethod
     def tearDownClass(cls):
-        os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB")
+        cls._disagg_failure_ctx.__exit__(None, None, None)
         super().tearDownClass()
 
     def test_gsm8k(self):

test/registered/disaggregation/test_disaggregation_nixl.py

@@ -126,7 +126,6 @@ def _clear_disagg_failure_env():
     # Failure injection is scoped to the failure test; clear inherited env so
     # Basic/Accuracy cannot run with injected transfer failures.
     os.environ.pop("SGLANG_TEST_DISAGG_FAILURE_PROB", None)
-    os.environ.pop("DISAGGREGATION_TEST_FAILURE_PROB", None)
 
 
 _HAS_CONFIGURED_NIXL_BACKEND = is_in_ci() or _has_configured_nixl_backend()