From ac179eec119a7253d1b482f4b591e417ed4b3492 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Sat, 22 Aug 2026 11:51:16 +0800 Subject: [PATCH] fix(disagg): PD transfer-failure injection was silently inert (#35890) --- python/sglang/srt/disaggregation/utils.py | 11 +---------- .../amd/disaggregation/test_disaggregation_basic.py | 9 ++++++--- .../disaggregation/test_disaggregation_basic.py | 9 ++++++--- .../disaggregation/test_disaggregation_nixl.py | 1 - 4 files changed, 13 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 0fe4076b2..c7da0bf3c 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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 diff --git a/test/registered/amd/disaggregation/test_disaggregation_basic.py b/test/registered/amd/disaggregation/test_disaggregation_basic.py index ece633fe8..97bfbf208 100644 --- a/test/registered/amd/disaggregation/test_disaggregation_basic.py +++ b/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 diff --git a/test/registered/disaggregation/test_disaggregation_basic.py b/test/registered/disaggregation/test_disaggregation_basic.py index 2812ce216..ea5a49cad 100644 --- a/test/registered/disaggregation/test_disaggregation_basic.py +++ b/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): diff --git a/test/registered/disaggregation/test_disaggregation_nixl.py b/test/registered/disaggregation/test_disaggregation_nixl.py index ca5bf09f9..de43e596d 100644 --- a/test/registered/disaggregation/test_disaggregation_nixl.py +++ b/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()