config: borrowed-record reads follow the config bags (#35908)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

test/registered/unit/model_executor/model_runner_components/test_attention_backend_setup.py

@@ -24,49 +24,60 @@ class _FakeBackend:
 
 
 def test_split_full_attention_applies_model_wrapper_once():
-    runner = SimpleNamespace(
-        server_args=SimpleNamespace(speculative_attention_mode="prefill"),
-        model_config=SimpleNamespace(context_len=2048),
-        kv_cache_dtype=None,
-        token_to_kv_pool=object(),
-        req_to_token_pool=object(),
-        init_new_workspace=None,
-    )
-    wrapper_inputs = []
-    wrapped_backend = object()
+    # The hybrid backend takes the speculative attention mode from the
+    # published configuration.
+    from sglang.srt.runtime_context import get_context
 
-    def wrap_once(model_runner, backend):
-        assert model_runner is runner
-        wrapper_inputs.append(backend)
-        return wrapped_backend
+    override = get_context().override_server_args(speculative_attention_mode="prefill")
+    override.install()
+    try:
+        runner = SimpleNamespace(
+            server_args=SimpleNamespace(speculative_attention_mode="prefill"),
+            model_config=SimpleNamespace(context_len=2048),
+            kv_cache_dtype=None,
+            token_to_kv_pool=object(),
+            req_to_token_pool=object(),
+            init_new_workspace=None,
+        )
+        wrapper_inputs = []
+        wrapped_backend = object()
 
-    constructors = {
-        "decode-test": lambda model_runner: _FakeBackend("decode"),
-        "prefill-test": lambda model_runner: _FakeBackend("prefill"),
-    }
-    resolved = ResolvedAttentionBackendStr(decode="decode-test", prefill="prefill-test")
+        def wrap_once(model_runner, backend):
+            assert model_runner is runner
+            wrapper_inputs.append(backend)
+            return wrapped_backend
 
-    with (
-        patch.dict(attention_backend_setup.ATTENTION_BACKENDS, constructors),
-        patch.object(
-            attention_backend_setup,
-            "attn_backend_wrapper",
-            side_effect=wrap_once,
-        ),
-    ):
-        result = attention_backend_setup._build_resolved_backend(
-            model_runner=runner,
-            resolved=resolved,
-            init_new_workspace=True,
+        constructors = {
+            "decode-test": lambda model_runner: _FakeBackend("decode"),
+            "prefill-test": lambda model_runner: _FakeBackend("prefill"),
+        }
+        resolved = ResolvedAttentionBackendStr(
+            decode="decode-test", prefill="prefill-test"
         )
 
-    assert result is wrapped_backend
-    assert len(wrapper_inputs) == 1
-    split_backend = wrapper_inputs[0]
-    assert isinstance(split_backend, HybridAttnBackend)
-    assert split_backend.decode_backend.name == "decode"
-    assert split_backend.prefill_backend.name == "prefill"
-    assert runner.init_new_workspace is True
+        with (
+            patch.dict(attention_backend_setup.ATTENTION_BACKENDS, constructors),
+            patch.object(
+                attention_backend_setup,
+                "attn_backend_wrapper",
+                side_effect=wrap_once,
+            ),
+        ):
+            result = attention_backend_setup._build_resolved_backend(
+                model_runner=runner,
+                resolved=resolved,
+                init_new_workspace=True,
+            )
+
+        assert result is wrapped_backend
+        assert len(wrapper_inputs) == 1
+        split_backend = wrapper_inputs[0]
+        assert isinstance(split_backend, HybridAttnBackend)
+        assert split_backend.decode_backend.name == "decode"
+        assert split_backend.prefill_backend.name == "prefill"
+        assert runner.init_new_workspace is True
+    finally:
+        override.restore()
 
 
 if __name__ == "__main__":

test/registered/unit/model_executor/model_runner_components/test_cuda_graph_setup.py

@@ -13,6 +13,15 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
 def test_model_runner_can_override_decode_graph_runner(monkeypatch):
+    from sglang.srt.runtime_context import get_context
+
+    # The capture decision reads the graph configuration and the MoE backends
+    # out of the bags.
+    override = get_context().override_server_args(
+        cuda_graph_config=SimpleNamespace(decode=SimpleNamespace(backend="default")),
+    )
+    override.install()
+
     class CustomGraphRunner:
         def __init__(self, model_runner):
             self.model_runner = model_runner
@@ -23,12 +32,7 @@ def test_model_runner_can_override_decode_graph_runner(monkeypatch):
         gpu_id = 0
         is_draft_worker = False
         spec_algorithm = SimpleNamespace(is_speculative=lambda: False)
-        server_args = SimpleNamespace(
-            model_impl="auto",
-            cuda_graph_config=SimpleNamespace(
-                decode=SimpleNamespace(backend="default")
-            ),
-        )
+        server_args = SimpleNamespace(model_impl="auto")
 
         def _decode_cuda_graph_runner_cls(self):
             return CustomGraphRunner
@@ -43,10 +47,13 @@ def test_model_runner_can_override_decode_graph_runner(monkeypatch):
         cuda_graph_setup.current_platform, "is_out_of_tree", lambda: False
     )
 
-    capture = capture_decode_graph(model_runner=model_runner)
+    try:
+        capture = capture_decode_graph(model_runner=model_runner)
 
-    assert isinstance(capture.runner, CustomGraphRunner)
-    assert capture.runner.model_runner is model_runner
+        assert isinstance(capture.runner, CustomGraphRunner)
+        assert capture.runner.model_runner is model_runner
+    finally:
+        override.restore()
 
 
 if __name__ == "__main__":

test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py

@@ -65,6 +65,16 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
     def test_low_free_memory_still_captures_prefill_graph(self):
         eager_runner = object()
         prefill_runner = object()
+        # The capture decision reads the graph configuration and the LoRA flag
+        # out of the bags.
+        override = get_context().override_server_args(
+            enable_lora=False,
+            cuda_graph_config=SimpleNamespace(
+                prefill=SimpleNamespace(bs=[1], backend=Backend.BREAKABLE)
+            ),
+        )
+        override.install()
+        self.addCleanup(override.restore)
         model_runner = SimpleNamespace(
             device="cuda",
             gpu_id=0,
@@ -73,11 +83,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             # reads it rather than the process-wide LoRA config.
             lora_manager=None,
             spec_algorithm=SimpleNamespace(is_eagle=lambda: False),
-            server_args=SimpleNamespace(
-                cuda_graph_config=SimpleNamespace(
-                    prefill=SimpleNamespace(bs=[1], backend=Backend.BREAKABLE)
-                ),
-            ),
+            server_args=SimpleNamespace(),
             model=SimpleNamespace(),
             model_config=SimpleNamespace(context_len=8192, num_hidden_layers=1),
             req_to_token_pool=SimpleNamespace(size=1),
@@ -140,15 +146,18 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         self.assertIs(capture.runner, eager_runner)
 
     def test_prefix_chunk_capacity_is_aggregate_and_can_be_overridden(self):
+        graph_config = SimpleNamespace(
+            prefill=SimpleNamespace(full_prefill_prefix_chunk_tokens=None, max_bs=8)
+        )
+        # Both leaves come from the bags; the published object is this one, so
+        # the cases below still drive them by mutating it.
+        override = get_context().override_server_args(
+            chunked_prefill_size=16, cuda_graph_config=graph_config
+        )
+        published = override.install()
+        self.addCleanup(override.restore)
         model_runner = SimpleNamespace(
-            server_args=SimpleNamespace(
-                chunked_prefill_size=16,
-                cuda_graph_config=SimpleNamespace(
-                    prefill=SimpleNamespace(
-                        full_prefill_prefix_chunk_tokens=None, max_bs=8
-                    )
-                ),
-            ),
+            server_args=SimpleNamespace(),
             # Wider than the token table, so the table is the binding limit.
             model_config=SimpleNamespace(context_len=4096),
             req_to_token_pool=SimpleNamespace(
@@ -161,24 +170,20 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
             (4, 16),
         )
 
-        model_runner.server_args.chunked_prefill_size = -1
+        get_context().override("test", chunked_prefill_size=-1)
         self.assertEqual(
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
             (2, 8),
         )
-        model_runner.server_args.chunked_prefill_size = 16
+        get_context().override("test", chunked_prefill_size=16)
 
-        model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
-            24
-        )
+        graph_config.prefill.full_prefill_prefix_chunk_tokens = 24
         self.assertEqual(
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
             (6, 24),
         )
 
-        model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
-            256
-        )
+        graph_config.prefill.full_prefill_prefix_chunk_tokens = 256
         self.assertEqual(
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
             (32, 128),
@@ -186,9 +191,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
 
         # At least one token is reserved per request lane even if the requested
         # aggregate capacity is smaller than the fixed request-slot count.
-        model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
-            2
-        )
+        graph_config.prefill.full_prefill_prefix_chunk_tokens = 2
         self.assertEqual(
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
             (1, 4),
@@ -197,17 +200,13 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
         # A context shorter than the token table binds instead: a draft runner
         # capped at the target's context, or a short --context-length.
         model_runner.model_config.context_len = 8
-        model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
-            256
-        )
+        graph_config.prefill.full_prefill_prefix_chunk_tokens = 256
         self.assertEqual(
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4),
             (8, 32),
         )
 
-        model_runner.server_args.cuda_graph_config.prefill.full_prefill_prefix_chunk_tokens = (
-            0
-        )
+        graph_config.prefill.full_prefill_prefix_chunk_tokens = 0
         with self.assertRaisesRegex(ValueError, "must be positive"):
             PrefillCudaGraphRunner._resolve_prefix_chunk_shape(model_runner, 4)