"""Unit tests for runtime_context: delegation, singletons, and override().""" from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=33, suite="base-a-test-cpu") import dataclasses import json import os import pathlib as _pathlib import shutil import tempfile import unittest import warnings from unittest.mock import patch import msgspec import msgspec.structs import sglang as _sglang import sglang.srt.server_args as server_args_module from sglang.srt.arg_groups import prefill_buffer_ceiling from sglang.srt.arg_groups.arg_utils import NS, A, Arg from sglang.srt.arg_groups.model_override_base import resolving_view from sglang.srt.arg_groups.overrides import ( attention_backends_of, ) from sglang.srt.arg_groups.overrides import ( mamba_cache_chunk_size as mamba_cache_chunk_size_of, ) from sglang.srt.arg_groups.overrides import ( max_prefill_buffer_tokens as max_prefill_buffer_tokens_of, ) from sglang.srt.arg_groups.overrides import ( resolution_result, resolved_view, ) from sglang.srt.runtime_context import ( Flags, ParallelContext, RuntimeContext, SpawnRanks, _FlagGroupBase, _validate_parallel, assert_published, derive_parallel_widths, get_context, get_exec, get_flags, get_parallel, get_schedule, get_server_args, max_prefill_buffer_tokens, max_speculative_num_draft_tokens, publish, publish_role, reset_context, ) from sglang.srt.server_args import ServerArgs from sglang.test.test_utils import CustomTestCase _SRT = _pathlib.Path(next(iter(_sglang.__path__))).resolve() / "srt" _PACKAGE = _pathlib.Path(next(iter(_sglang.__path__))).resolve() def _sources(): """Every Python file this checkout ships, package and siblings alike. The package alone is the wrong subject set for anything about entries or public names: `benchmark/`, `examples/` and the top-level `test/` call the same doors and are not covered by any suite that would notice them break. An installed package has no siblings, and then this is the package alone.""" roots = [_PACKAGE] checkout = _PACKAGE.parents[1] roots += [ checkout / name for name in ("benchmark", "examples", "scripts", "test") if (checkout / name).is_dir() ] for root in roots: for path in root.rglob("*.py"): yield path def _scope_entries_that_say_nothing(paths): """Draft-scope entries that do not state `owns_attention`, as `path:line`. The scope either narrows the draft's attention and expert identity or leaves the target's in place, and only the worker knows which -- so the keyword has no default. Omitting it is a `TypeError`, but only on the path that runs, and those paths want a GPU and a draft model. """ import ast missing = [] for path in paths: try: tree = ast.parse(path.read_text(encoding="utf-8")) except (SyntaxError, UnicodeDecodeError): continue for node in ast.walk(tree): if not isinstance(node, ast.Call): continue func = node.func name = getattr(func, "attr", None) or getattr(func, "id", None) if name not in ("draft_tp_context", "patch_tensor_parallel_group"): continue if not any(kw.arg == "owns_attention" for kw in node.keywords): missing.append(f"{path}:{node.lineno}") return missing _PS = "sglang.srt.distributed.parallel_state" _DP = "sglang.srt.layers.dp_attention" # Ranks and the launch width are asked of the group: they are not implied by # anything, so there is nothing to derive them from. The quotients are not # here -- `attn_tp_size` and its siblings are functions of the configured # leaves, and `TestDerivedWidths` pins them. `attn_dp_rank` is not here either: no group coordinator # knows it, so it is stamped when the attention topology is initialized and # `TestStampedRanks` is what pins it. The other world width is not here # because the group does not know it; `TestTheTwoWorldWidths` pins it. SIZE_RANK_DELEGATIONS = [ ("launch_world_size", f"{_PS}.get_world_size"), ("launch_world_rank", f"{_PS}.get_world_rank"), ("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), ("dcp_rank", f"{_PS}.get_dcp_rank"), ("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), ("moe_ep_rank", f"{_PS}.get_moe_expert_parallel_rank"), ("moe_dp_rank", f"{_PS}.get_moe_data_parallel_rank"), ("moe_tp_rank", f"{_PS}.get_moe_tensor_parallel_rank"), ("attn_tp_rank", f"{_PS}.get_attn_tensor_model_parallel_rank"), ("attn_cp_rank", f"{_PS}.get_attn_context_model_parallel_rank"), ] GROUP_DELEGATIONS = [ ("world_group", f"{_PS}.get_world_group"), ("tp_group", f"{_PS}.get_tp_group"), ("dcp_group", f"{_PS}.get_dcp_group"), ("pp_group", f"{_PS}.get_pp_group"), ("moe_ep_group", f"{_PS}.get_moe_ep_group"), ("moe_dp_group", f"{_PS}.get_moe_dp_group"), ("moe_tp_group", f"{_PS}.get_moe_tp_group"), ("attn_tp_group", f"{_PS}.get_attn_tp_group"), ("attn_cp_group", f"{_PS}.get_attn_cp_group"), ] class TestRuntimeContextSingletons(CustomTestCase): def test_singletons(self): self.assertIs(get_parallel(), get_parallel()) self.assertIsInstance(get_parallel(), ParallelContext) self.assertIsInstance(get_context(), RuntimeContext) self.assertIs(get_context().parallel, get_parallel()) class _IsolatedOverrides(CustomTestCase): """Give each test a clean override map, restoring afterward only the overrides installed outside it (e.g. by another test file sharing the process).""" def setUp(self): super().setUp() p = get_parallel() self._saved_overrides = dict(p._overrides) p._overrides.clear() def tearDown(self): p = get_parallel() p._overrides.clear() p._overrides.update(self._saved_overrides) super().tearDown() class TestParallelDelegation(_IsolatedOverrides): def test_size_rank_delegate_to_canonical_getters(self): # Patch each getter to a distinct sentinel: a miswired attribute would read # a different (unpatched) getter and fail. for i, (attr, target) in enumerate(SIZE_RANK_DELEGATIONS): sentinel = 1000 + i with patch(target, return_value=sentinel): self.assertEqual( getattr(get_parallel(), attr), sentinel, msg=f"{attr} must delegate to {target}", ) def test_groups_delegate_to_canonical_getters(self): for attr, target in GROUP_DELEGATIONS: sentinel = object() with patch(target, return_value=sentinel): self.assertIs( getattr(get_parallel(), attr), sentinel, msg=f"{attr} must delegate to {target}", ) def test_wrapper_holds_no_resolved_state(self): # __slots__: no __dict__; the only instance state is the override hook. self.assertFalse(hasattr(get_parallel(), "__dict__")) # tp_group IS exposed: live delegation handles PD-multiplexing / the tp patch. self.assertTrue(hasattr(ParallelContext, "tp_group")) # local_attn_dp is intentionally not part of the wrapper surface. self.assertFalse(hasattr(ParallelContext, "local_attn_dp_size")) class TestTheTwoWorldWidths(_IsolatedOverrides): """Two questions about the WORLD group: what it was built at, and what it has room for. Neither is stored here. How much of that room is serving after a scale-up is elastic-EP state, and is asked of the manager that owns it rather than mirrored onto this namespace. """ def test_the_launch_width_is_what_the_group_was_built_at(self): with patch(f"{_PS}.get_world_size", return_value=4): self.assertEqual(get_parallel().launch_world_size, 4) def test_the_ceiling_is_the_configured_one_when_there_is_one(self): parallel = get_parallel() with ( parallel.override(max_ep_size=32), patch( f"{_PS}.get_world_size", side_effect=AssertionError("the built group must not be asked"), ), ): self.assertEqual(parallel.max_world_size, 32) def test_without_a_configured_ceiling_the_room_is_the_launch_width(self): parallel = get_parallel() with ( parallel.override(max_ep_size=None), patch(f"{_PS}.get_world_size", return_value=8), ): self.assertEqual(parallel.max_world_size, 8) def test_each_width_can_be_stated_on_its_own(self): """Stating one must not answer for the other: they are two names.""" parallel = get_parallel() with ( parallel.override(launch_world_size=2, max_ep_size=None), patch( f"{_PS}.get_world_size", side_effect=AssertionError("the built group must not be asked"), ), ): self.assertEqual(parallel.launch_world_size, 2) self.assertEqual(parallel.max_world_size, 2) with parallel.override(max_ep_size=6): self.assertEqual(parallel.max_world_size, 6) self.assertEqual(parallel.launch_world_size, 2) class TestSpawnIdentities(_IsolatedOverrides): """`dp_rank` and `gpu_id` come from the spawn, because nothing else has them. Both vary per process while the record is identical across them, and neither is a position in any process group -- no group has one member per data-parallel replica. So the process entry states them at publish. """ def setUp(self): super().setUp() parallel = get_parallel() self._saved_stamp = dict(parallel._stamp) self.addCleanup( lambda: ( parallel.clear_stamp(), parallel.override_permanently(**self._saved_stamp), ) ) reset_context() self.addCleanup(reset_context) def test_one_rank_fixes_the_rest(self): """Every other rank is a position in a group laid out from the widths, so `world_rank` is the whole placement: rank 5 of a `tp=4, pp=2` world is the second stage's second device.""" publish( ServerArgs(model_path="dummy", tp_size=4, pp_size=2), role="test", ranks=SpawnRanks(world_rank=5, dp_rank=2), ) parallel = get_parallel() self.assertEqual(parallel.launch_world_rank, 5) self.assertEqual(parallel.tp_rank, 1) self.assertEqual(parallel.pp_rank, 1) self.assertEqual(parallel.dp_rank, 2) def test_no_controller_is_an_answer_not_a_failure(self): """`dp_rank=None` means "not under a data parallel controller", which is a fact about the deployment, unlike never having been told. The replicas are separate WORLD groups, so no rank implies it.""" publish( ServerArgs(model_path="dummy", tp_size=2), role="test", ranks=SpawnRanks(world_rank=0, dp_rank=None), ) self.assertIsNone(get_parallel().dp_rank) def test_publishing_without_a_bundle_names_what_is_missing(self): publish(ServerArgs(model_path="dummy", tp_size=2), role="test") with self.assertRaises(RuntimeError) as caught: get_parallel().dp_rank self.assertIn("rank bundle", str(caught.exception)) def test_the_attention_rank_keeps_its_own_explanation(self): """Two stamp-only names, two different reasons to be missing.""" publish(ServerArgs(model_path="dummy", tp_size=2), role="test") with self.assertRaises(RuntimeError) as caught: get_parallel().attn_dp_rank self.assertIn("initialize_dp_attention", str(caught.exception)) class TestAttentionRanksComeFromPublish(_IsolatedOverrides): """With a spawn bundle, a rank read works before any group exists. This is what `ParallelState` provided by being a plain frozen record, and what the topology init could not: it needs the groups. Deriving at publish is what lets a reader ask the context in a process that never initialises distributed -- every unit test that builds a scheduler component, for one. """ def setUp(self): super().setUp() parallel = get_parallel() self._saved_stamp = dict(parallel._stamp) self.addCleanup( lambda: ( parallel.clear_stamp(), parallel.override_permanently(**self._saved_stamp), ) ) reset_context() self.addCleanup(reset_context) def test_it_matches_the_topology_init_for_every_shape(self): """Cross-checked against the function the groups use, not restated. Same inputs, two callers: one has them from the configuration and the spawn, the other from the groups it just built. """ from sglang.srt.layers.dp_attention import compute_dp_attention_world_info shapes = [ (8, 1, 1, False), (8, 2, 1, True), (8, 4, 1, True), (8, 2, 2, True), (16, 4, 2, True), ] for tp_size, dp_size, attn_cp_size, dp_attn in shapes: for tp_rank in range(tp_size): reset_context() publish( ServerArgs( model_path="dummy", tp_size=tp_size, dp_size=dp_size, attn_cp_size=attn_cp_size, enable_dp_attention=dp_attn, ), role="test", ranks=SpawnRanks(world_rank=tp_rank), ) want_tp, _, want_dp, _ = compute_dp_attention_world_info( dp_attn, tp_rank, tp_size, dp_size, attn_cp_size ) msg = f"tp={tp_size} dp={dp_size} cp={attn_cp_size} rank={tp_rank}" self.assertEqual(get_parallel().attn_tp_rank, want_tp, msg) self.assertEqual(get_parallel().attn_dp_rank, want_dp, msg) def test_the_rank_reads_without_a_process_group(self): """No distributed init, no patching of any getter.""" publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True ), role="test", ranks=SpawnRanks(world_rank=5), ) with patch( f"{_PS}.get_attn_tensor_model_parallel_rank", side_effect=AssertionError("no group must be consulted"), ): self.assertEqual(get_parallel().attn_tp_rank, 1) self.assertEqual(get_parallel().attn_dp_rank, 1) def test_without_a_bundle_it_still_asks_the_group(self): """Unchanged for every process that publishes without a placement.""" publish(ServerArgs(model_path="dummy", tp_size=8), role="test") with patch(f"{_PS}.get_attn_tensor_model_parallel_rank", return_value=3): self.assertEqual(get_parallel().attn_tp_rank, 3) class TestStampedRanks(_IsolatedOverrides): """`attn_dp_rank` comes from the stamp, and says so when there is none. It is the one rank no group answers with: `initialize_dp_attention` computes it from this process's `tp_rank`. Falling back to anything would be inventing a placement for this process. """ def setUp(self): super().setUp() parallel = get_parallel() self._saved_derived = dict(parallel._stamp) parallel.clear_stamp() self.addCleanup( lambda: ( parallel.clear_stamp(), parallel.override_permanently(**self._saved_derived), ) ) def test_the_stamp_is_the_answer(self): parallel = get_parallel() parallel.override_permanently(attn_dp_rank=3) self.assertEqual(parallel.attn_dp_rank, 3) # An elastic scale-up restamps it; the newest stamp wins. parallel.override_permanently(attn_dp_rank=9) self.assertEqual(parallel.attn_dp_rank, 9) def test_a_scope_still_wins_over_the_stamp(self): parallel = get_parallel() parallel.override_permanently(attn_dp_rank=3) with parallel.override(attn_dp_rank=0): self.assertEqual(parallel.attn_dp_rank, 0) self.assertEqual(parallel.attn_dp_rank, 3) def test_unstamped_names_the_cause(self): with self.assertRaises(RuntimeError) as caught: get_parallel().attn_dp_rank self.assertIn("initialize_dp_attention", str(caught.exception)) def test_a_stated_width_reaches_the_padding_mode(self): """The reason this PR exists, from a reader's side. `get_dp_padding_mode` reads the attention-DP width. Before the width had one home, a scoped `override` moved the context and left the module global answering, so stating a topology moved only half the runtime: this asserted `SUM_LEN` with the width stated as 1. """ from sglang.srt.layers.dp_attention import DpPaddingMode with get_parallel().override(attn_dp_size=1): mode = DpPaddingMode.get_dp_padding_mode( is_extend_in_batch=True, global_num_tokens=[3, 5] ) self.assertIs(mode, DpPaddingMode.MAX_LEN) # And the branch it would have taken with the target's width. with get_parallel().override(attn_dp_size=2): mode = DpPaddingMode.get_dp_padding_mode( is_extend_in_batch=True, global_num_tokens=[3, 5] ) self.assertIs(mode, DpPaddingMode.SUM_LEN) def test_the_gather_slot_follows_the_list_that_was_gathered(self): """The DP sync gathers over the attention-DP replicas, or over the expanded WORLD once a scale-up has moved the gather there. The index into that list is a property of the gather, so it is read beside the flag that says which one happened rather than kept on the topology.""" from sglang.srt.layers.dp_attention import dp_gather_slot self.addCleanup(reset_context) dp_flags = get_flags().dp saved = ( dp_flags.use_world_group_for_gather, dp_flags.joiner_skip_all_gather, ) def restore(): ( dp_flags.use_world_group_for_gather, dp_flags.joiner_skip_all_gather, ) = saved self.addCleanup(restore) publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=8, enable_dp_attention=True ), role="test", ranks=SpawnRanks(world_rank=3), ) parallel = get_parallel() dp_flags.use_world_group_for_gather = False self.assertEqual(dp_gather_slot(), parallel.attn_dp_rank) # After a scale-up the gather spans the expanded WORLD, and the joining # cohort is numbered from its offset. dp_flags.use_world_group_for_gather = True dp_flags.joiner_skip_all_gather = False parallel.override_permanently(ep_join_rank_offset=8) self.assertEqual(dp_gather_slot(), 8 + parallel.tp_rank) # and the topology it was read off is untouched self.assertEqual(parallel.attn_dp_size, 8) self.assertEqual(parallel.tp_size, 8) def test_a_scale_up_writes_no_width(self): """The identities stay unconditional because nothing overrides them: the scale-up only points the gather at the expanded WORLD.""" from sglang.srt.layers.dp_attention import update_dp_attention_post_scale dp_flags = get_flags().dp saved_gather = dp_flags.use_world_group_for_gather self.addCleanup(setattr, dp_flags, "use_world_group_for_gather", saved_gather) self.addCleanup(reset_context) publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=8, enable_dp_attention=True ), role="test", ranks=SpawnRanks(world_rank=3), ) parallel = get_parallel() before = (parallel.attn_dp_size, parallel.attn_dp_rank, parallel.tp_size) update_dp_attention_post_scale(new_dp_size=16, new_dp_rank=11) self.assertTrue(dp_flags.use_world_group_for_gather) self.assertEqual( (parallel.attn_dp_size, parallel.attn_dp_rank, parallel.tp_size), before ) class TestEveryDeclaredParallelNameIsStatable(_IsolatedOverrides): """The overridable set is read from the declarations, not maintained by hand. A hand-kept list can hold a name the class does not answer, or miss one it does; either way `override()` refuses or accepts the wrong thing with nothing to say so. The three tests below check the set against the declarations from both sides. """ def test_every_declared_name_can_be_stated_and_reads_back(self): from sglang.srt.runtime_context import _parallel_fields names = sorted(_parallel_fields()) # Sizes, ranks, groups and the configured leaves of the namespace. self.assertGreater(len(names), 30) parallel = get_parallel() for name in names: sentinel = object() with parallel.override(**{name: sentinel}): self.assertIs(getattr(parallel, name), sentinel, msg=name) def test_every_name_the_class_answers_for_is_in_the_set(self): """Cross-check from the other side: the class's own surface. Derived from the class rather than from the same declarations the set is built from, so a source dropped out of `_parallel_fields` shows up here instead of agreeing with itself. """ from sglang.srt.runtime_context import _parallel_fields answered = { name for name, value in vars(ParallelContext).items() if isinstance(value, property) } self.assertTrue(answered) self.assertEqual(answered - _parallel_fields(), set()) def test_a_live_name_is_never_also_answered_from_the_bag(self): """The two answer differently, so a name in both would make the read order -- not the declaration -- decide which one a caller gets. The bag carries the declared quotients as well as the operator's leaves, and both are ahead of the live getter once a configuration is published: a name in `_LIVE_READS` and in either of them would answer from the getter before publish and from the bag after.""" from sglang.srt.runtime_context import ( _LIVE_READS, _derived_widths, _parallel_config_leaves, ) self.assertEqual(set(_LIVE_READS) & _parallel_config_leaves(), set()) self.assertEqual(set(_LIVE_READS) & set(_derived_widths()), set()) def test_an_undeclared_name_is_refused(self): with self.assertRaises(ValueError): with get_parallel().override(not_a_parallel_name=1): pass class TestReadsWithoutAPublishedConfig(_IsolatedOverrides): """The namespace has to answer in a process that publishes nothing. `multimodal_gen` lends its own TP group to shared `srt` layers from a process with no `srt` config to publish against, and those layers ask for `attn_tp_size` anyway -- through code `multimodal_gen` does not own, which is why grepping that package for `get_parallel()` finds nothing while the read plainly happens. """ def setUp(self): super().setUp() parallel = get_parallel() self._saved_stamp = dict(parallel._stamp) self.addCleanup( lambda: ( parallel.clear_stamp(), parallel.override_permanently(**self._saved_stamp), ) ) reset_context() self.addCleanup(reset_context) def test_a_stamped_width_reads_with_nothing_published(self): parallel = get_parallel() self.assertIsNone(parallel._config) parallel.override_permanently( **derive_parallel_widths( tp_size=2, attn_cp_size=1, attn_dp_size=1, moe_ep_size=1, moe_dp_size=1, dcp_size=1, dcp_enabled=False, ) ) self.assertEqual(parallel.attn_tp_size, 2) self.assertEqual(parallel.moe_tp_size, 2) def test_an_unstamped_width_still_names_the_cause(self): """Without a stamp there is nothing to answer with, and the failure has to say so rather than invent a width.""" with self.assertRaisesRegex(RuntimeError, r"not available"): get_parallel().attn_tp_size class TestPrivateAttributeProbing(_IsolatedOverrides): def test_probing_a_private_name_does_not_recurse(self): """`copy` and `pickle` probe for hooks before `__init__` has run. `__getattr__` reaches for `self._config`, so if it did not refuse underscore names outright, probing one on a half-built instance would recurse until the stack ran out. """ fresh = ParallelContext.__new__(ParallelContext) # slots unset for probe in ("_config", "_stamp", "_overrides", "__deepcopy__"): with self.assertRaises(AttributeError, msg=probe): getattr(fresh, probe) def test_a_built_context_survives_a_copy(self): import copy self.assertIsInstance(copy.copy(get_parallel()), ParallelContext) class TestAWidthReadStaysTraceable(_IsolatedOverrides): """A width read inside compiled model code must stay inside the graph. Shared layers read widths inside a compiled forward. A graph break there is a performance regression and nothing else -- every suite stays green through it -- so `fullgraph=True` is what turns it into a failure. This pins the read path, whichever form it takes: the sibling leaf test compiles names served by `__getattr__` and they trace too. """ def test_a_width_read_compiles_into_the_graph(self): import torch reset_context() self.addCleanup(reset_context) publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True ), role="test", ) def read(x): return x * get_parallel().attn_tp_size # backend="eager": this pins tracing, not code generation, and stays # runnable on a box with no inductor toolchain. compiled = torch.compile(read, fullgraph=True, backend="eager") self.assertEqual(compiled(torch.ones(3)).tolist(), [4.0, 4.0, 4.0]) class TestParallelOverride(_IsolatedOverrides): def test_override_takes_precedence(self): p = get_parallel() with p.override(tp_size=99, tp_rank=3, attn_dp_size=8): self.assertEqual(p.tp_size, 99) self.assertEqual(p.tp_rank, 3) self.assertEqual(p.attn_dp_size, 8) # same singleton: a fresh get_parallel() sees the override too self.assertEqual(get_parallel().tp_size, 99) self.assertEqual(p._overrides, {}) def test_override_can_force_groups(self): sentinel = object() with get_parallel().override(tp_group=sentinel): self.assertIs(get_parallel().tp_group, sentinel) def test_override_nests_and_restores(self): p = get_parallel() with p.override(tp_size=2): self.assertEqual(p.tp_size, 2) with p.override(tp_size=4, pp_size=2): self.assertEqual(p.tp_size, 4) self.assertEqual(p.pp_size, 2) self.assertEqual(p.tp_size, 2) self.assertNotIn("pp_size", p._overrides) def test_override_unknown_key_raises_and_does_not_mutate(self): p = get_parallel() with self.assertRaises(ValueError): with p.override(tp_sizee=1): # typo pass self.assertEqual(p._overrides, {}) class TestParallelDCP(_IsolatedOverrides): """The DCP width is a quotient; the DCP rank is a live reading. They used to be tested the same way, by mocking the group getters, because the width read the group too. It does not: `attn_dcp_size` is `dcp_size if dcp_enabled else 1`, so the way to state it is to state the leaves. """ def _published(self, **fields): reset_context() self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy", **fields), role="test") return get_parallel() def test_attn_dcp_is_one_when_dcp_is_off(self): parallel = self._published(tp_size=8, dcp_size=1) self.assertFalse(parallel.dcp_enabled) self.assertEqual(parallel.attn_dcp_size, 1) def test_attn_dcp_is_the_configured_width_when_on(self): parallel = self._published(tp_size=8, dcp_size=8) self.assertTrue(parallel.dcp_enabled) self.assertEqual(parallel.attn_dcp_size, 8) def test_the_dcp_rank_still_reads_the_group(self): """A rank is not implied by the configuration, so it reads the group -- gated on a width that is.""" with ( get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=False), patch(f"{_PS}.get_dcp_rank", side_effect=AssertionError), ): self.assertEqual(get_parallel().attn_dcp_rank, 0) with ( get_parallel().override(tp_size=8, dcp_size=8, dcp_enabled=True), patch(f"{_PS}.get_dcp_rank", return_value=3), ): self.assertEqual(get_parallel().attn_dcp_rank, 3) def test_the_width_does_not_consult_the_platform(self): with patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda: parallel = self._published(tp_size=8, dcp_size=8) self.assertTrue(parallel.dcp_enabled) self.assertEqual(parallel.attn_dcp_size, 8) is_cuda.assert_not_called() class _IsolatedServerArgs(CustomTestCase): """Save/restore the published ServerArgs around each test (the slot is process-global; another test file sharing the process may have published).""" def setUp(self): super().setUp() self._saved_server_args = get_context()._server_args def tearDown(self): if self._saved_server_args is None: reset_context() else: get_context().set_server_args(self._saved_server_args) super().tearDown() class TestServerArgsOwnership(_IsolatedServerArgs): """V2b: the context owns the slot; the legacy getters are identity shims.""" def test_legacy_setter_publishes_into_context(self): # Identity, not equality: the slot holds the very object published. sentinel = ServerArgs(model_path="dummy") server_args_module.set_global_server_args_for_scheduler(sentinel) self.assertIs(get_server_args(), sentinel) self.assertIs(get_context().server_args, sentinel) def test_the_retired_accessor_raises_and_names_the_replacement(self): """`get_global_server_args` is retired: it answered with the record, so a caller reading a field resolution had decided got a stale value and no error at all. `RuntimeError` unconditionally, not a warning first: a `DeprecationWarning` is filtered by default outside `__main__`, so no production caller would have seen it, and under `-W error::DeprecationWarning` it would have changed the exception a caller catches. The message has to name where to read instead, since the answer differs by what the caller wanted. """ with self.assertRaises(RuntimeError) as cm: server_args_module.get_global_server_args() message = str(cm.exception) self.assertIn("runtime_context", message) self.assertIn("get_server_args()", message) # And the type does not change when warnings are errors. with warnings.catch_warnings(): warnings.simplefilter("error") with self.assertRaises(RuntimeError): server_args_module.get_global_server_args() def test_tokenizer_alias_is_distinct_role_shim(self): # Deliberately NOT an alias: the two legacy setters publish with # different process roles (scheduler vs tokenizer). self.assertIsNot( server_args_module.set_global_server_args_for_tokenizer, server_args_module.set_global_server_args_for_scheduler, ) def test_pre_publish_error_verbatim(self): reset_context() with self.assertRaises(ValueError) as cm: get_server_args() self.assertEqual(str(cm.exception), "Global server args is not set yet!") def test_republish_overwrite_allowed(self): first = ServerArgs(model_path="dummy") second = ServerArgs(model_path="dummy") server_args_module.set_global_server_args_for_scheduler(first) server_args_module.set_global_server_args_for_scheduler(second) self.assertIs(get_server_args(), second) def test_reset_context_clears_owned_store(self): server_args_module.set_global_server_args_for_scheduler( ServerArgs(model_path="dummy") ) reset_context() with self.assertRaises(ValueError): get_server_args() class TestAssertPublished(_IsolatedServerArgs): """Publishing is the process entry's job; the constructors only check. `ModelRunner`, `TokenizerManager` and `MMEncoder` assert. A publish inside a process that has already published re-projects the bags, discarding every `override()` taken since and the provenance log with it, so a constructor that finds nothing published fails loud. """ def _record(self, **fields): return ServerArgs(model_path="dummy", **fields) def test_the_check_leaves_a_live_process_alone(self): record = self._record(grammar_backend="xgrammar") publish(record, role="scheduler") get_context().override("grammar.import_fallback", grammar_backend="none") assert_published(record, role="scheduler") self.assertEqual( get_exec().kernel.grammar_backend, "none", "the check re-projected the bags, so the import fallback was " "discarded and the process reports a backend it is not using", ) self.assertEqual( len(get_context().overrides_log()), 1, "the provenance of the override went with it", ) def test_a_different_record_fails(self): first = self._record(grammar_backend="xgrammar") publish(first, role="scheduler") second = self._record(grammar_backend="llguidance") with self.assertRaisesRegex(RuntimeError, "a different record is published"): assert_published(second, role="scheduler") self.assertIs( get_server_args(), first, "the failing check published anyway", ) def test_an_empty_slot_fails(self): """An empty slot fails.""" reset_context() record = self._record(grammar_backend="xgrammar") with self.assertRaisesRegex( RuntimeError, "nothing is published in this process" ): assert_published(record, role="scheduler") def test_the_same_record_under_a_different_role_fails(self): """The role decides which namespaces this process may read.""" record = self._record() publish(record, role="tokenizer") with self.assertRaisesRegex(RuntimeError, "published under role 'tokenizer'"): assert_published(record, role="scheduler") self.assertEqual(publish_role(), "tokenizer") class TestServerArgsScopedOverride(_IsolatedServerArgs): """ctx.override_server_args: the config tier's scoped test override — tests force execution paths by overriding the context, not by hand-building and publishing config objects.""" def test_install_publishes_fresh_config_with_fields(self): reset_context() override = get_context().override_server_args( attention_backend="triton", chunked_prefill_size=-1 ) published = override.install() self.assertIs(get_server_args(), published) # The hook declares; the record keeps the operator's input, so the # values are read where resolution puts them. self.assertEqual(resolution_result(published, "attention_backend"), "triton") self.assertEqual(resolution_result(published, "chunked_prefill_size"), -1) # unnamed fields keep their dataclass defaults self.assertEqual(resolution_result(published, "tp_size"), 1) def test_unknown_fields_are_rejected(self): with self.assertRaises(ValueError): get_context().override_server_args(not_a_config_field=1).install() def test_restore_reinstates_previous_publish(self): previous = object() get_context().set_server_args(previous) override = get_context().override_server_args(tp_size=8) override.install() self.assertEqual(get_parallel().tp_size, 8) override.restore() self.assertIs(get_server_args(), previous) def test_restore_reinstates_the_empty_slot(self): reset_context() with get_context().override_server_args(): get_server_args() # published inside the scope with self.assertRaises(ValueError): get_server_args() def test_nesting_restores_in_order(self): reset_context() with get_context().override_server_args(tp_size=2) as outer: with get_context().override_server_args(tp_size=4): self.assertEqual(get_parallel().tp_size, 4) self.assertIs(get_server_args(), outer) self.assertEqual(get_parallel().tp_size, 2) def test_private_attribute_seeding(self): # Property caches (e.g. _mamba_cache_chunk_size) are seeded through # the same call; the strict guard exempts underscore names. published = ( get_context().override_server_args(_mamba_cache_chunk_size=64).install() ) self.assertEqual(mamba_cache_chunk_size_of(published), 64) def test_an_underscore_field_is_declared_like_any_other(self): """The split is fields vs not-fields, not the leading underscore. `_speculative_draft_quantization_explicitly_set` is a real field published under `spec`. Seeding it as a raw attribute instead of declaring it would leave the earlier resolution authoritative, so both the resolution and the bag would keep answering the pre-override value while the record said otherwise. """ from sglang.srt.arg_groups.overrides import resolution_result from sglang.srt.runtime_context import get_spec name = "_speculative_draft_quantization_explicitly_set" self.assertIn(name, ServerArgs.__struct_fields__) published = get_context().override_server_args(**{name: True}).install() # The record keeps the operator's input, as it does for every other # field; the override travels as a declaration. self.assertIsNone(getattr(published, name)) self.assertIs(resolution_result(published, name), True) self.assertIs(getattr(get_spec(), name), True) def test_installed_config_arms_the_strict_guard(self): # The published dummy must behave like a resolved config: bare writes # raise. published = get_context().override_server_args(tp_size=2).install() with self.assertRaises(AttributeError): published.tp_size = 4 self.assertEqual(resolution_result(published, "tp_size"), 2) def test_restore_resets_the_capture_seed(self): # install() seeds flags.capture from the published dummy; restore() # must put back the pre-install runtime state on both restore paths. reset_context() self.assertFalse(get_flags().capture.enable_torch_compile) override = get_context().override_server_args(enable_torch_compile=True) override.install() self.assertTrue(get_flags().capture.enable_torch_compile) override.restore() self.assertFalse(get_flags().capture.enable_torch_compile) def test_double_install_rejected(self): override = get_context().override_server_args() override.install() with self.assertRaises(AssertionError): override.install() class _FakeCaptureGroup(_FlagGroupBase): gamma: int = 0 class TestFlagsTier(_IsolatedServerArgs): """Runtime-flags tier: typed groups, typo-safe writes, override primitive. Resolved configuration lives on server_args fields (materialized at the end of __post_init__); the flags tier only carries runtime state (today: the capture lifecycle).""" def test_wiring_and_groups(self): flags = get_flags() self.assertIs(flags, get_context().flags) self.assertIsInstance(flags, Flags) self.assertTrue(hasattr(flags, "capture")) def test_typo_safety(self): group = _FakeCaptureGroup() with self.assertRaises(AttributeError): group.gamma_misspelled = 2 # undeclared leaf with self.assertRaises(AttributeError): get_flags().not_a_flag = 1 def test_override_is_transactional(self): group = _FakeCaptureGroup() with group.override(gamma=99): self.assertEqual(group.gamma, 99) self.assertEqual(group.gamma, 0) with self.assertRaises(ValueError): with group.override(gamma=2, delta=3): # delta undeclared pass self.assertEqual(group.gamma, 0) # validated before any write def test_reset_context_installs_fresh_flags(self): old = get_flags() old.capture.enable_torch_compile = True reset_context() self.assertIsNot(get_flags(), old) self.assertFalse(get_flags().capture.enable_torch_compile) @dataclasses.dataclass class _FakeResolvedArgs: """Publishable fixture with a resolvable whitelist (real flat leaves).""" page_size: A[int | None, Arg(help="p", resolvable=True), NS("schedule")] = None sampling_backend: A[ str | None, Arg(help="s", resolvable=True), NS("exec.kernel") ] = None attention_backend: A[str | None, Arg(help="ab"), NS("exec.kernel")] = None prefill_attention_backend: A[str | None, Arg(help="pab"), NS("exec.kernel")] = None decode_attention_backend: A[str | None, Arg(help="dab"), NS("exec.kernel")] = None disable_radix_cache: A[bool, Arg(help="drc"), NS("memory")] = False mamba_radix_cache_strategy: A[str, Arg(help="mrcs"), NS("exec.mamba")] = "auto" speculative_algorithm: A[str | None, Arg(help="sa"), NS("spec")] = None speculative_num_draft_tokens: A[int | None, Arg(help="d"), NS("spec")] = None speculative_adaptive: A[bool, Arg(help="a"), NS("spec")] = False speculative_adaptive_config: A[str | None, Arg(help="c"), NS("spec")] = None load_format: A[str, Arg(help="lf"), NS("model")] = "auto" remote_instance_weight_loader_backend: A[str, Arg(help="rb"), NS("model")] = "nccl" remote_instance_weight_loader_start_seed_via_transfer_engine: A[ bool, Arg(help="rs"), NS("model") ] = False modelexpress_config: A[str | None, Arg(help="mx"), NS("model")] = None disaggregation_mode: A[str, Arg(help="dm"), NS("disagg")] = "null" max_running_requests: A[int | None, Arg(help="mrr"), NS("schedule")] = None chunked_prefill_size: A[int, Arg(help="cps"), NS("schedule")] = -1 max_prefill_tokens: A[int, Arg(help="mpt"), NS("schedule")] = 16384 enable_dynamic_chunking: A[bool, Arg(help="edc"), NS("schedule")] = False cuda_graph_config: A[object | None, Arg(help="cgc"), NS("exec.graph")] = None tp_size: A[int, Arg(help="tp"), NS("parallel")] = 1 pp_size: A[int, Arg(help="pp"), NS("parallel")] = 1 _resolved_overrides: list = dataclasses.field(default_factory=list) class TestMoeFlagsGroup(_IsolatedServerArgs): """flags.moe: materialized by initialize_moe_config; the ACTIVE backends swap under the speculative contexts and restore on exit.""" def _init(self, **kw): from sglang.srt.layers.moe.utils import initialize_moe_config defaults = dict( moe_a2a_backend="none", moe_runner_backend="auto", speculative_moe_runner_backend=None, speculative_moe_a2a_backend=None, deepep_mode="auto", deepep_config=None, enable_two_batch_overlap=False, enable_single_batch_overlap=False, tbo_token_distribution_threshold=0.48, disable_flashinfer_cutlass_moe_fp4_allgather=False, quantization=None, disable_shared_experts_fusion=False, ) defaults.update(kw) # The flags are seeded from the bags, so the test publishes a config # carrying these values. override = get_context().override_server_args(**defaults) override.install() self.addCleanup(override.restore) initialize_moe_config() def test_lazy_defaults_before_initialize(self): from sglang.srt.layers.moe.utils import ( get_moe_a2a_backend, get_moe_runner_backend, is_tbo_enabled, ) reset_context() self.assertTrue(get_moe_a2a_backend().is_none()) self.assertEqual(get_moe_runner_backend().name, "AUTO") self.assertFalse(is_tbo_enabled()) def test_initialize_materializes_group(self): from sglang.srt.layers.moe.utils import get_moe_a2a_backend, is_tbo_enabled self._init(moe_a2a_backend="deepep", enable_two_batch_overlap=True) self.assertTrue(get_moe_a2a_backend().is_deepep()) self.assertTrue(is_tbo_enabled()) self.assertEqual(get_flags().moe.deepep_config, "") def test_speculative_swap_and_restore(self): from sglang.srt.layers.moe.utils import ( get_moe_a2a_backend, get_moe_runner_backend, speculative_moe_a2a_backend_context, speculative_moe_backend_context, ) self._init( moe_a2a_backend="deepep", moe_runner_backend="triton", speculative_moe_runner_backend="auto", speculative_moe_a2a_backend="none", ) with speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): self.assertEqual(get_moe_runner_backend().name, "AUTO") self.assertTrue(get_moe_a2a_backend().is_none()) # MTP layers are unquantized: fp4 allgather is forced off self.assertTrue(get_flags().moe.disable_fp4_allgather) self.assertTrue(get_flags().moe.speculative_context) self.assertEqual(get_moe_runner_backend().name, "TRITON") self.assertTrue(get_moe_a2a_backend().is_deepep()) self.assertFalse(get_flags().moe.disable_fp4_allgather) self.assertFalse(get_flags().moe.speculative_context) def test_swap_restores_on_exception(self): from sglang.srt.layers.moe.utils import ( get_moe_runner_backend, speculative_moe_backend_context, ) self._init(moe_runner_backend="triton", speculative_moe_runner_backend="auto") with self.assertRaises(RuntimeError): with speculative_moe_backend_context(): raise RuntimeError("boom") self.assertEqual(get_moe_runner_backend().name, "TRITON") class TestDpFlagsGroup(_IsolatedServerArgs): """flags.dp: the DP-attention runtime flags; is_dp_attention_enabled is a thin shim over the group leaf.""" def test_shim_reads_the_leaf(self): from sglang.srt.layers.dp_attention import is_dp_attention_enabled reset_context() self.assertFalse(is_dp_attention_enabled()) get_flags().dp.enabled = True self.assertTrue(is_dp_attention_enabled()) def test_scoped_override_forces_the_predicate(self): from sglang.srt.layers.dp_attention import is_dp_attention_enabled reset_context() with get_flags().dp.override(enabled=True): self.assertTrue(is_dp_attention_enabled()) self.assertFalse(is_dp_attention_enabled()) class TestResources(_IsolatedServerArgs): """ctx.resources: named slots for process-level resource handles with one reset lifecycle; owning accessors keep their creation/publish semantics.""" def test_graph_pool_lazy_create_and_reuse(self): from types import SimpleNamespace from sglang.srt.model_executor.runner_utils.pool import ( get_global_graph_memory_pool, get_or_create_global_graph_memory_pool, ) reset_context() self.assertIsNone(get_global_graph_memory_pool()) dev = SimpleNamespace(graph_pool_handle=lambda: object()) handle = get_or_create_global_graph_memory_pool(dev) self.assertIs(get_or_create_global_graph_memory_pool(dev), handle) def test_expert_recorder_noop_default_and_injection(self): from sglang.srt.eplb.expert_distribution import ( get_global_expert_distribution_recorder, ) from sglang.srt.runtime_context import get_resources reset_context() self.assertEqual( type(get_global_expert_distribution_recorder()).__name__, "_ExpertDistributionRecorderNoop", ) with get_resources().override(expert_distribution_recorder="mock"): self.assertEqual(get_global_expert_distribution_recorder(), "mock") def test_expert_location_metadata_publish_once_until_reset(self): from sglang.srt.eplb.expert_location import ( get_global_expert_location_metadata, set_global_expert_location_metadata, ) reset_context() self.assertIsNone(get_global_expert_location_metadata()) set_global_expert_location_metadata("meta") with self.assertRaises(AssertionError): set_global_expert_location_metadata("again") reset_context() self.assertIsNone(get_global_expert_location_metadata()) class TestNamedStreams(_IsolatedServerArgs): """ctx.get_stream(name): keyed get-or-create (the persistent-buffer pattern); set_stream installs explicitly.""" def test_get_or_create_shares_by_name(self): from unittest.mock import patch reset_context() created = [] class _FakeStream: def __init__(self): created.append(self) with patch("torch.cuda.Stream", _FakeStream): a = get_context().get_stream("alt") b = get_context().get_stream("alt") c = get_context().get_stream("other") self.assertIs(a, b) self.assertIsNot(a, c) self.assertEqual(len(created), 2) def test_get_buffer_keyed_lazy(self): reset_context() created = [] def factory(): created.append(object()) return created[-1] a = get_context().get_buffer("ws", factory) b = get_context().get_buffer("ws", factory) self.assertIs(a, b) self.assertEqual(len(created), 1) self.assertIsNot(get_context().get_buffer("other", factory), a) def test_set_stream_installs_explicitly(self): reset_context() sentinel = object() get_context().set_stream("alt", sentinel) self.assertIs(get_context().get_stream("alt"), sentinel) def test_reset_clears_the_registry(self): reset_context() get_context().set_stream("alt", object()) reset_context() self.assertEqual(get_context().resources.streams, {}) def test_capturer_slots_roundtrip_and_reset(self): from sglang.srt.state_capturer.indexer_topk import ( get_global_indexer_capturer, set_global_indexer_capturer, ) from sglang.srt.state_capturer.routed_experts import ( get_global_experts_capturer, set_global_experts_capturer, ) reset_context() self.assertIsNone(get_global_indexer_capturer()) self.assertIsNone(get_global_experts_capturer()) indexer, experts = object(), object() set_global_indexer_capturer(indexer) set_global_experts_capturer(experts) self.assertIs(get_global_indexer_capturer(), indexer) self.assertIs(get_global_experts_capturer(), experts) reset_context() self.assertIsNone(get_global_indexer_capturer()) self.assertIsNone(get_global_experts_capturer()) def test_tcp_store_slot_roundtrip_and_reset(self): from sglang.srt.distributed.utils import ( get_global_tcp_store, set_global_tcp_store, ) reset_context() self.assertIsNone(get_global_tcp_store()) store = object() set_global_tcp_store(store) self.assertIs(get_global_tcp_store(), store) reset_context() self.assertIsNone(get_global_tcp_store()) def test_trace_level_env_seeded_lazy_default(self): from sglang.srt.observability.trace import ( get_global_trace_level, set_global_trace_level, ) reset_context() with patch.dict(os.environ, {}, clear=False): os.environ.pop("SGLANG_TRACE_LEVEL", None) self.assertEqual(get_global_trace_level(), 3) set_global_trace_level(5) self.assertEqual(get_global_trace_level(), 5) reset_context() with patch.dict(os.environ, {"SGLANG_TRACE_LEVEL": "1"}): self.assertEqual(get_global_trace_level(), 1) class TestEpBufferState(_IsolatedServerArgs): """EP dispatcher buffer managers: state lives on ctx.resources; the facade keeps the mode-transition and clean semantics.""" def test_deepep_dispatch_mode_transitions_and_reset(self): try: from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer except ImportError: self.skipTest("deep_ep not installed") reset_context() cleans = [] class _FakeBuffer: low_latency_mode = True def clean_low_latency_buffer(self, *args): cleans.append(args) state = DeepEPBuffer._state() state.buffer = _FakeBuffer() state.hidden_size = 7168 state.num_max_dispatch_tokens_per_rank = 128 state.num_experts = 256 DeepEPBuffer.set_dispatch_mode_as_normal() # NORMAL -> LOW_LATENCY must clean the low-latency buffer once. DeepEPBuffer.set_dispatch_mode_as_low_latency() self.assertEqual(cleans, [(128, 7168, 256)]) # LOW_LATENCY -> LOW_LATENCY must not clean again. DeepEPBuffer.set_dispatch_mode_as_low_latency() self.assertEqual(len(cleans), 1) reset_context() self.assertIsNone(DeepEPBuffer._state().buffer) class TestForwardFlags(_IsolatedServerArgs): """ctx.forward: contextvar-backed per-forward flags; scoped() restores, threads see defaults.""" def test_scoped_set_restore_and_nesting(self): from sglang.srt.runtime_context import get_forward reset_context() fwd = get_forward() self.assertFalse(fwd.multi_stream) with fwd.scoped(multi_stream=True): self.assertTrue(fwd.multi_stream) with fwd.scoped(multi_stream=False): self.assertFalse(fwd.multi_stream) self.assertTrue(fwd.multi_stream) self.assertFalse(fwd.multi_stream) def test_scoped_restores_on_exception_and_validates_keys(self): from sglang.srt.runtime_context import get_forward reset_context() fwd = get_forward() with self.assertRaises(RuntimeError): with fwd.scoped(moe_output_buffer="buf"): raise RuntimeError("boom") self.assertIsNone(fwd.moe_output_buffer) with self.assertRaises(ValueError): with fwd.scoped(nope=1): pass with self.assertRaises(AttributeError): fwd.multi_stream = True # attribute writes are rejected def test_threads_see_defaults(self): import threading from sglang.srt.runtime_context import get_forward reset_context() fwd = get_forward() seen = {} with fwd.scoped(multi_stream=True): def probe(): seen["value"] = get_forward().multi_stream worker = threading.Thread(target=probe) worker.start() worker.join() self.assertFalse(seen["value"]) # a new thread sees the default def test_graph_visible_flags_trace_under_torch_compile(self): # Regression: dynamo cannot trace ContextVar.get, and these flags are # read inside compiled model code (vocab embedding, communicator, DP # gather) — they must stay plain-slot backed. fullgraph=True turns # any graph break back into a failure. import torch from sglang.srt.runtime_context import get_forward reset_context() @torch.compile(fullgraph=True, backend="eager", dynamic=False) def probe(x): fwd = get_forward() if fwd.attn_input_scattered: x = x + 1 if fwd.is_extend_in_batch: x = x + 2 if fwd.fuse_mlp_allreduce: x = x + 4 if fwd.mlp_reduce_scatter: x = x + 8 if fwd.flashinfer_trtllm_bypass: x = x + 16 return x self.assertEqual(probe(torch.zeros(())).item(), 0) with get_forward().scoped(attn_input_scattered=True): self.assertEqual(probe(torch.zeros(())).item(), 1) get_forward().set("is_extend_in_batch", True) self.assertEqual(probe(torch.zeros(())).item(), 2) get_forward().set("is_extend_in_batch", False) with get_forward().scoped( fuse_mlp_allreduce=True, mlp_reduce_scatter=True, flashinfer_trtllm_bypass=True, ): self.assertEqual(probe(torch.zeros(())).item(), 28) self.assertEqual(probe(torch.zeros(())).item(), 0) def test_parallel_config_leaves_trace_under_torch_compile(self): # Regression: gate helpers such as ``enable_moe_dense_fully_dp()`` read # parallel config leaves inside compiled model forwards, which must # stay dynamo-traceable (``object.__getattribute__`` graph-breaks). # fullgraph=True turns any graph break back into a failure. import torch from sglang.srt.runtime_context import get_parallel reset_context() with get_context().override_server_args(moe_dense_tp_size=1, dwdp_size=4): @torch.compile(fullgraph=True, backend="eager", dynamic=False) def probe(x): par = get_parallel() if par.enable_prefill_cp: x = x + 1 if par.moe_dense_tp_size == 1: x = x + 2 if par.dwdp_size > 1: x = x + 4 return x self.assertEqual(probe(torch.zeros(())).item(), 6) def test_graph_visible_flags_are_process_visible_across_threads(self): # Documented divergence from the contextvar-backed flags: plain slots # are process-global (the storage form these flags had before the # tier), so another thread sees the current value, not the default. import threading from sglang.srt.runtime_context import get_forward reset_context() seen = {} with get_forward().scoped(attn_input_scattered=True): def probe(): seen["value"] = get_forward().attn_input_scattered worker = threading.Thread(target=probe) worker.start() worker.join() self.assertTrue(seen["value"]) self.assertFalse(get_forward().attn_input_scattered) def test_multi_stream_shims(self): from sglang.srt.utils.multi_stream_utils import ( do_multi_stream, with_multi_stream, ) reset_context() self.assertFalse(do_multi_stream()) with with_multi_stream(True): self.assertTrue(do_multi_stream()) self.assertFalse(do_multi_stream()) def test_attn_tp_context_per_forward_slots(self): from types import SimpleNamespace from sglang.srt.layers.communicator import get_attn_tp_context from sglang.srt.runtime_context import get_forward reset_context() ctx = get_attn_tp_context() self.assertFalse(ctx.input_scattered) fb = SimpleNamespace( forward_mode=SimpleNamespace( is_extend=lambda: False, is_target_verify=lambda: False ), input_ids=None, can_run_tbo=False, ) sentinel = SimpleNamespace(fetch_qkv_latent=lambda: "qkv") with ctx.maybe_input_scattered(fb): ctx.set_attn_inputs(sentinel) self.assertEqual(ctx.fetch_qkv_latent(), "qkv") # attn inputs are cleared at scope exit, flag restored self.assertIsNone(get_forward().attn_inputs) self.assertFalse(ctx.input_scattered) def test_dp_buffer_state_split(self): import torch from sglang.srt.layers.dp_attention import _DpGatheredBufferWrapper as wrapper from sglang.srt.layers.dp_attention import ( get_dp_dtype, get_dp_global_num_tokens, get_global_dp_buffer_len, is_dp_max_padding, set_dp_buffer_len, ) reset_context() # metadata is init-static (flags.dp); sizing is per-forward sticky wrapper.set_metadata(64, torch.float16, torch.device("cpu")) self.assertEqual(get_dp_dtype(), torch.float16) set_dp_buffer_len(128, 32, True, [64, 64]) self.assertEqual(get_global_dp_buffer_len(), 128) self.assertTrue(is_dp_max_padding()) self.assertEqual(get_dp_global_num_tokens(), [64, 64]) set_dp_buffer_len(256, 64, False) # sticky until the next write self.assertEqual(get_global_dp_buffer_len(), 256) self.assertFalse(is_dp_max_padding()) self.assertIsNone(get_dp_global_num_tokens()) reset_context() self.assertIsNone(get_dp_dtype()) def test_is_extend_in_batch_sticky_within_thread(self): from sglang.srt.layers.dp_attention import ( get_is_extend_in_batch, set_is_extend_in_batch, ) reset_context() self.assertFalse(get_is_extend_in_batch()) set_is_extend_in_batch(True) self.assertTrue(get_is_extend_in_batch()) # sticky until next write set_is_extend_in_batch(False) self.assertFalse(get_is_extend_in_batch()) def test_moe_output_buffer_ctx(self): from sglang.srt.layers.moe.moe_runner.base import moe_output_buffer_ctx from sglang.srt.runtime_context import get_forward reset_context() sentinel = object() with moe_output_buffer_ctx(sentinel): self.assertIs(get_forward().moe_output_buffer, sentinel) self.assertIsNone(get_forward().moe_output_buffer) def test_mlp_comm_forward_flags(self): """Decoder-published MLP collective flags: scoped restore + skip helpers.""" from sglang.srt.layers.moe.utils import ( should_skip_mlp_all_reduce, should_skip_post_experts_all_reduce, ) from sglang.srt.runtime_context import get_forward reset_context() fwd = get_forward() self.assertFalse(fwd.fuse_mlp_allreduce) self.assertFalse(fwd.mlp_reduce_scatter) self.assertFalse(fwd.flashinfer_trtllm_bypass) self.assertFalse(should_skip_mlp_all_reduce()) with fwd.scoped(fuse_mlp_allreduce=True): self.assertTrue(fwd.fuse_mlp_allreduce) self.assertTrue(should_skip_mlp_all_reduce()) # Fusion alone is enough to skip post-experts AR. self.assertTrue(should_skip_post_experts_all_reduce(is_tp_path=True)) self.assertFalse(fwd.fuse_mlp_allreduce) self.assertFalse(should_skip_mlp_all_reduce()) with fwd.scoped(mlp_reduce_scatter=True): self.assertTrue(fwd.mlp_reduce_scatter) self.assertTrue(should_skip_mlp_all_reduce()) self.assertFalse(fwd.mlp_reduce_scatter) with fwd.scoped(flashinfer_trtllm_bypass=True): self.assertTrue(fwd.flashinfer_trtllm_bypass) self.assertFalse(fwd.flashinfer_trtllm_bypass) def test_dp_reduce_scatterv_requires_single_rank_attention_dp_shards(self): from sglang.srt.layers.moe.utils import should_use_dp_reduce_scatterv reset_context() with patch( "sglang.srt.layers.moe.utils.is_dp_attention_enabled", return_value=True, ): # The optimized path is valid when the collective group and the # variable-split list have the same number of entries. with get_parallel().override(tp_size=8, attn_dp_size=8, moe_ep_size=8): self.assertTrue(should_use_dp_reduce_scatterv()) # Otherwise the standard all-reduce plus scatter path must be used. with get_parallel().override(tp_size=8, attn_dp_size=2, moe_ep_size=2): self.assertFalse(should_use_dp_reduce_scatterv()) class TestPublishLifecycle(_IsolatedServerArgs): """Publish installs the resolved server_args and seeds the capture tier.""" def _publish(self, **kw): args = _FakeResolvedArgs(**kw) get_context().set_server_args(args) return args def test_capture_tier_seeded_at_publish(self): args = self._publish(page_size=1) args.enable_torch_compile = True get_context().set_server_args(args) # re-publish picks up the value self.assertTrue(get_flags().capture.enable_torch_compile) # capture-time write (B4) targets the capture leaf get_flags().capture.enable_torch_compile = False self.assertFalse(get_flags().capture.enable_torch_compile) def test_capture_tier_defaults_for_sentinel_publish(self): get_context().set_server_args(object()) self.assertFalse(get_flags().capture.enable_torch_compile) class TestDerivedPredicatesAgreeAcrossTiers(_IsolatedServerArgs): """One definition per predicate, checked rather than asserted in prose. Each of these exists twice by construction -- once over a config-shaped object (the resolution pipeline's `*_of` helper, which `ServerArgs` delegates to) and once over the published bags. The pair must agree on every input, or a decision made before publish differs from the same decision made after it. """ _STRATEGIES = ("auto", "no_buffer", "extra_buffer", "extra_buffer_lazy") def test_the_mamba_extra_buffer_predicate_has_one_answer(self): """It used to be asserted that two spellings agreed. There is one now: the declaration computes it at publish, and the bag carries it.""" for disable_radix_cache in (False, True): for strategy in self._STRATEGIES: with self.subTest(radix=disable_radix_cache, strategy=strategy): reset_context() publish( ServerArgs( model_path="dummy", disable_radix_cache=disable_radix_cache, mamba_radix_cache_strategy=strategy, ), role="test", ) expected = disable_radix_cache is False and strategy in ( "extra_buffer", "extra_buffer_lazy", ) self.assertEqual( get_exec().mamba.enable_mamba_extra_buffer, expected ) self.assertEqual( get_exec().mamba.enable_mamba_extra_buffer_lazy, disable_radix_cache is False and strategy == "extra_buffer_lazy", ) def test_prefill_buffer_ceiling_matches_the_member(self): from sglang.srt.runtime_context import max_prefill_buffer_tokens for chunked in (-1, 0, 1024, 8192): for dynamic in (False, True): for pp in (1, 4): for max_prefill in (0, 2048, 16384): with self.subTest( chunked=chunked, dynamic=dynamic, pp=pp, max_prefill=max_prefill, ): args = _FakeResolvedArgs( chunked_prefill_size=chunked, enable_dynamic_chunking=dynamic, pp_size=pp, max_prefill_tokens=max_prefill, ) get_context().set_server_args(args) self.assertEqual( max_prefill_buffer_tokens_of(args), max_prefill_buffer_tokens(), ) def test_prefill_buffer_ceiling_hook_honored_across_tiers(self): args = _FakeResolvedArgs( chunked_prefill_size=8192, enable_dynamic_chunking=True, pp_size=4, max_prefill_tokens=16384, ) def provider(record, default_ceiling): self.assertIs(record, args) return default_ceiling + 5 with patch.object(prefill_buffer_ceiling, "_prefill_buffer_ceiling_fn", None): register = prefill_buffer_ceiling.register_prefill_buffer_ceiling self.assertEqual(max_prefill_buffer_tokens_of(args), 16384) self.assertIs(register(provider), provider) register(provider) with self.assertRaisesRegex(RuntimeError, "already registered"): register(lambda record, default_ceiling: default_ceiling) for record_or_view in (args, resolving_view(args), resolved_view(args)): self.assertEqual(max_prefill_buffer_tokens_of(record_or_view), 16389) get_context().set_server_args(args) self.assertEqual(max_prefill_buffer_tokens(), 16389) with get_schedule().override(max_prefill_tokens=32768): self.assertEqual(max_prefill_buffer_tokens(), 32773) self.assertEqual(args.max_prefill_tokens, 16384) def test_prefill_buffer_ceiling_provider_can_preserve_defaults(self): args = _FakeResolvedArgs(chunked_prefill_size=4096) def provider(record, default_ceiling): return default_ceiling with patch.object(prefill_buffer_ceiling, "_prefill_buffer_ceiling_fn", None): prefill_buffer_ceiling.register_prefill_buffer_ceiling(provider) for record_or_view in (args, resolving_view(args), resolved_view(args)): self.assertEqual(max_prefill_buffer_tokens_of(record_or_view), 4096) get_context().set_server_args(args) self.assertEqual(max_prefill_buffer_tokens(), 4096) def test_activation_reserve_matches_the_member(self): from types import SimpleNamespace from sglang.srt.arg_groups.overrides import ( pre_capture_activation_reserve_mb_of, ) from sglang.srt.runtime_context import pre_capture_activation_reserve_mb graph = SimpleNamespace(decode=SimpleNamespace(max_bs=64)) cases = ( dict(disaggregation_mode="null", chunked_prefill_size=8192), dict(disaggregation_mode="null", chunked_prefill_size=-1), dict( disaggregation_mode="null", chunked_prefill_size=-1, max_prefill_tokens=1024, ), dict(disaggregation_mode="decode", max_running_requests=32), dict(disaggregation_mode="decode", max_running_requests=None), dict( disaggregation_mode="decode", max_running_requests=None, speculative_num_draft_tokens=4, ), dict( disaggregation_mode="null", chunked_prefill_size=8192, tp_size=8, pp_size=2, ), ) for case in cases: for gpu_mem in (None, 20 * 1024, 80 * 1024): with self.subTest(gpu_mem=gpu_mem, **case): args = _FakeResolvedArgs(cuda_graph_config=graph, **case) get_context().set_server_args(args) self.assertEqual( pre_capture_activation_reserve_mb_of(args, gpu_mem), pre_capture_activation_reserve_mb(gpu_mem), ) def test_remote_instance_transfer_engine_matches_the_member(self): from sglang.srt.runtime_context import remote_instance_transfer_engine_enabled backends = ("nccl", "transfer_engine", "modelexpress") transports = (None, '{"transport": "transfer_engine"}', '{"transport": "nixl"}') for seed_via_te in (False, True): for load_format in ("auto", "remote_instance"): for backend in backends: for mx in transports: with self.subTest( seed=seed_via_te, load_format=load_format, backend=backend, modelexpress=mx, ): args = _FakeResolvedArgs( load_format=load_format, remote_instance_weight_loader_backend=backend, remote_instance_weight_loader_start_seed_via_transfer_engine=seed_via_te, modelexpress_config=mx, ) get_context().set_server_args(args) for override in (None, "remote_instance", "auto"): self.assertEqual( ServerArgs.remote_instance_weight_loader_use_transfer_engine( args, override ), remote_instance_transfer_engine_enabled(override), ) def test_attention_backends_match_the_member(self): from sglang.srt.runtime_context import attention_backends backends = (None, "fa3", "triton") for base in backends: for prefill in backends: for decode in backends: with self.subTest(base=base, prefill=prefill, decode=decode): args = _FakeResolvedArgs( attention_backend=base, prefill_attention_backend=prefill, decode_attention_backend=decode, ) get_context().set_server_args(args) self.assertEqual( attention_backends_of(resolved_view(args)), attention_backends(), ) class TestAdaptiveDraftBoundLifecycle(_IsolatedServerArgs): """The adaptive draft-token bound is snapshotted at each publication.""" def _write_config(self, steps): path = os.path.join(tempfile.mkdtemp(prefix="adaptive_cfg_"), "adaptive.json") self.addCleanup(shutil.rmtree, os.path.dirname(path), ignore_errors=True) with open(path, "w") as handle: json.dump({"1": {"candidate_steps": steps}}, handle) return path def test_republishing_recomputes_the_bound(self): path = self._write_config([2]) get_context().set_server_args( _FakeResolvedArgs( speculative_num_draft_tokens=3, speculative_adaptive=True, speculative_adaptive_config=path, ) ) self.assertEqual(max_speculative_num_draft_tokens(), 3) with open(path, "w") as handle: json.dump({"1": {"candidate_steps": [4]}}, handle) # The new publication must not retain the previous capacity. get_context().set_server_args( _FakeResolvedArgs( speculative_num_draft_tokens=3, speculative_adaptive=True, speculative_adaptive_config=path, ) ) self.assertEqual(max_speculative_num_draft_tokens(), 5) def test_reset_clears_the_bound(self): path = self._write_config([2]) get_context().set_server_args( _FakeResolvedArgs( speculative_num_draft_tokens=3, speculative_adaptive=True, speculative_adaptive_config=path, ) ) self.assertEqual(max_speculative_num_draft_tokens(), 3) reset_context() with open(path, "w") as handle: json.dump({"1": {"candidate_steps": [6]}}, handle) get_context().set_server_args( _FakeResolvedArgs( speculative_num_draft_tokens=3, speculative_adaptive=True, speculative_adaptive_config=path, ) ) self.assertEqual(max_speculative_num_draft_tokens(), 7) class TestParallelLeafReads(_IsolatedServerArgs): """The contract ``ParallelContext.__getattr__`` answers a parallel leaf on.""" def test_a_leaf_answers_what_resolution_decided(self): from sglang.srt.arg_groups.overrides import resolution_result with get_context().override_server_args() as server_args: self.assertEqual( resolution_result(server_args, "nccl_port"), get_parallel().nccl_port, "a parallel leaf read off the context disagreed with what " "resolution decided", ) def test_before_publish_the_error_names_the_namespace(self): with self.assertRaisesRegex(ValueError, r"'parallel' not published"): getattr(ParallelContext(), "nccl_port") def test_an_unknown_name_is_still_an_attribute_error(self): with self.assertRaisesRegex(AttributeError, r"has no 'not_a_leaf'"): getattr(ParallelContext(), "not_a_leaf") class TestDerivedWidths(_IsolatedOverrides): """The widths no flag sets are computed from the leaves and permanently overridable. `attn_tp_size` and its siblings used to be read back off the group coordinator that was built from them, which made the answer depend on distributed init and, after an elastic scale, disagree with the leaves. """ def setUp(self): super().setUp() parallel = get_parallel() self._saved_derived = dict(parallel._stamp) parallel.clear_stamp() self.addCleanup( lambda: ( parallel.clear_stamp(), parallel.override_permanently(**self._saved_derived), ) ) def test_the_published_configuration_decides_the_widths(self): """The quotients are computed once, at publish, from the leaves. Every input is a record field, so there is nothing to recompute on a read: `publish` fills the bag and the bag is the answer. """ reset_context() self.addCleanup(reset_context) publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True ), role="test", ) self.assertEqual(get_parallel().attn_tp_size, 4) self.assertEqual(get_parallel().attn_dp_size, 2) self.assertEqual(get_parallel().moe_tp_size, 8) reset_context() publish( ServerArgs(model_path="dummy", tp_size=8, ep_size=4, moe_dp_size=2), role="test", ) self.assertEqual(get_parallel().moe_tp_size, 1) def test_a_topology_is_stated_by_naming_the_width(self): """Overriding a leaf does not move the quotient -- the quotient is not recomputed on read. Naming it is how a caller states one, and naming only some of them is refused: the caller owns the arithmetic, the context only checks it.""" reset_context() self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy", tp_size=8), role="test") self.assertEqual(get_parallel().attn_tp_size, 8) with self.assertRaises(ValueError) as caught: with get_parallel().override(tp_size=2): pass self.assertIn( "tp_size == attn_tp_size * attn_dp_size * attn_cp_size", str(caught.exception), ) with get_parallel().override(tp_size=2, attn_tp_size=2, moe_tp_size=2): self.assertEqual(get_parallel().tp_size, 2) self.assertEqual(get_parallel().attn_tp_size, 2) with get_parallel().override(attn_tp_size=4, tp_size=4, moe_tp_size=4): self.assertEqual(get_parallel().attn_tp_size, 4) def test_an_unstated_topology_still_fails(self): """Neutral leaves are for the dimensions a caller is not using, not for a caller that stated nothing: every width would come back 1, which is a plausible-looking number invented out of nothing.""" with self.assertRaises(RuntimeError) as caught: get_parallel().attn_tp_size self.assertIn("not available", str(caught.exception)) def test_a_permanent_override_and_a_live_group_both_win_over_the_leaves(self): """Order is scoped override, then the stamp, then the published leaf. No group is consulted for a width -- `test_the_group_is_never_consulted` in this class asserts that -- so a stamp is what an elastic scale-up leaves behind, and the leaf answers only where there is none. """ parallel = get_parallel() parallel.override_permanently(attn_tp_size=7) self.addCleanup(parallel.clear_stamp) with parallel.override(tp_size=8, attn_dp_size=2): self.assertEqual(parallel.attn_tp_size, 7) def test_the_quotients_come_from_the_leaves(self): widths = derive_parallel_widths( tp_size=8, attn_cp_size=1, attn_dp_size=2, moe_ep_size=4, moe_dp_size=2, dcp_size=1, dcp_enabled=False, ) self.assertEqual(widths["attn_tp_size"], 8 // 2 // 1) self.assertEqual(widths["moe_tp_size"], 8 // 4 // 2) self.assertEqual(widths["attn_dcp_size"], 1) def test_no_world_width_is_a_quotient_of_the_leaves(self): """Deriving one would answer with the joining cohort's own width on a scale joiner, which lays its groups out at `tp * pp` while WORLD spans `ep_join_rank_offset + tp * pp`. The launch width comes off the group that was actually built; the ceiling is not this function's to give either, and `TestTheTwoWorldWidths` says where each comes from.""" widths = derive_parallel_widths( tp_size=4, attn_cp_size=1, attn_dp_size=1, moe_ep_size=1, moe_dp_size=1, dcp_size=1, dcp_enabled=False, ) self.assertEqual( {name for name in widths if "world" in name}, set(), ) parallel = get_parallel() parallel.override_permanently(attn_tp_size=4) with patch(f"{_PS}.get_world_size", return_value=9): self.assertEqual(parallel.launch_world_size, 9) def test_the_bare_name_is_gone(self): """It answered two questions, so every reader had to remember which. Both spellings fail: reading it, and stating it -- the overridable set is derived from the same declarations the read path is, so a name that cannot be read cannot be stated either. """ with self.assertRaisesRegex(AttributeError, r"has no 'world_size'"): get_parallel().world_size with self.assertRaisesRegex(ValueError, r"unknown parallel field"): with get_parallel().override(world_size=4): pass def test_a_permanently_overridden_width_is_what_the_reader_answers_with(self): parallel = get_parallel() parallel.override_permanently(attn_tp_size=4, moe_tp_size=1) with patch( f"{_PS}.get_attn_tensor_model_parallel_world_size", side_effect=AssertionError("the group must not be asked"), ): self.assertEqual(parallel.attn_tp_size, 4) def test_a_scoped_override_still_wins_over_the_permanent_one(self): parallel = get_parallel() parallel.override_permanently(attn_tp_size=4) with parallel.override(attn_tp_size=1): self.assertEqual(parallel.attn_tp_size, 1) self.assertEqual(parallel.attn_tp_size, 4) def test_the_group_is_never_consulted(self): """There is no third source. A quotient comes from a scoped override, a permanent override, or the published leaf -- never from a group coordinator.""" reset_context() self.addCleanup(reset_context) with patch( f"{_PS}.get_attn_tensor_model_parallel_world_size", side_effect=AssertionError("the group must not be consulted"), ): publish( ServerArgs( model_path="dummy", tp_size=8, dp_size=2, enable_dp_attention=True ), role="test", ) self.assertEqual(get_parallel().attn_tp_size, 4) def test_with_neither_the_failure_names_the_cause(self): with patch( f"{_PS}.get_attn_tensor_model_parallel_world_size", side_effect=AssertionError("attention tp group is not initialized"), ): with self.assertRaisesRegex(RuntimeError, r"derived parallel width"): get_parallel().attn_tp_size def test_the_permanent_override_is_cleared_and_reset(self): parallel = get_parallel() parallel.override_permanently(attn_dp_size=2) self.assertEqual(parallel.attn_dp_size, 2) # Elastic scaling overrides again where it updates the live width. parallel.override_permanently(attn_dp_size=4) self.assertEqual(parallel.attn_dp_size, 4) parallel.clear_stamp() with parallel.override(tp_size=8, attn_dp_size=1): self.assertEqual(parallel.attn_dp_size, 1) def test_reset_context_drops_the_permanent_override(self): """The permanent override belongs to the lifecycle that made it. `_read` prefers it over the published leaf, so one that outlived `reset_context()` would let the next test read the previous topology. """ parallel = get_parallel() parallel.override_permanently(attn_tp_size=4) self.assertEqual(parallel.attn_tp_size, 4) reset_context() self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy", tp_size=1), role="test") self.assertEqual(get_parallel().attn_tp_size, 1) def test_the_rank_helper_agrees_with_the_override(self): """`compute_dp_attention_world_info` keeps the ranks and takes the widths from the same derivation `override_permanently`'s callers use.""" from sglang.srt.layers.dp_attention import compute_dp_attention_world_info for tp_size, dp_size, attn_cp_size in ((8, 2, 1), (8, 2, 2), (16, 4, 2)): _, attn_tp_size, _, attn_dp_size = compute_dp_attention_world_info( True, 0, tp_size, dp_size, attn_cp_size ) widths = derive_parallel_widths( tp_size=tp_size, attn_cp_size=attn_cp_size, attn_dp_size=attn_dp_size, moe_ep_size=1, moe_dp_size=1, dcp_size=1, dcp_enabled=False, ) self.assertEqual(attn_tp_size, widths["attn_tp_size"]) self.assertEqual(attn_dp_size, widths["attn_dp_size"]) def test_recomputing_from_published_leaves_matches_the_publish_bag(self): """`initialize_model_parallel` no longer overrides anything -- see `test_initialize_model_parallel_no_longer_touches_the_bag` below -- so every real caller must forward leaves that already match its own published config, because nothing corrects a mismatch anymore. `scheduler.py`'s `ps.attn_dp_size`/`ps.moe_ep_size`/etc, and the weight-cache daemon's own already-published config, both do -- this pins that the formula they'd recompute from those leaves (`derive_attention_widths`, `derive_parallel_widths`, the same ones `publish` itself used) agrees with what's already in the bag, across the widths `test_the_rank_helper_agrees_with_the_override` does not vary -- moe_ep_size, moe_dp_size, and dcp_size -- using real `publish()`. A caller that does NOT keep the two in sync is a bug in that caller, not something this framework silently corrects: two real ones existed (`test/registered/eplb/test_lplb_distributed.py` and `test/manual/ep/test_flashinfer_dispatcher.py`, both publishing a placeholder config and then building real groups at a width it never reflected) and were fixed by publishing the actual width instead of relying on a correction to paper over the mismatch. """ shapes = ( dict(tp_size=8), dict(tp_size=8, dp_size=2, enable_dp_attention=True), dict(tp_size=8, ep_size=4, moe_dp_size=2), dict(tp_size=8, dcp_size=8), ) for shape in shapes: with self.subTest(shape=shape): reset_context() self.addCleanup(reset_context) publish(ServerArgs(model_path="dummy", **shape), role="test") parallel = get_parallel() published = { "attn_tp_size": parallel.attn_tp_size, "attn_dp_size": parallel.attn_dp_size, "moe_ep_size": parallel.moe_ep_size, "moe_tp_size": parallel.moe_tp_size, "dcp_enabled": parallel.dcp_enabled, "attn_dcp_size": parallel.attn_dcp_size, } # What every real `initialize_model_parallel` caller forwards: # its own already-published leaves, through the same two # functions the bag was projected with. recomputed = derive_parallel_widths( tp_size=parallel.tp_size, attn_cp_size=parallel.attn_cp_size, attn_dp_size=( parallel.dp_size if parallel.enable_dp_attention else 1 ), moe_ep_size=parallel.ep_size, moe_dp_size=parallel.moe_dp_size, dcp_size=parallel.dcp_size, dcp_enabled=parallel.dcp_size > 1, ) self.assertEqual(published, recomputed) def test_initialize_model_parallel_builds_at_the_published_widths(self): """The build takes every width from the context rather than from an argument, so "published one width, built another" is no longer a state a caller can reach -- there is nothing left to translate, and nothing to correct afterwards either. """ from unittest.mock import Mock from sglang.srt.distributed import parallel_state reset_context() self.addCleanup(reset_context) world_size = 8 publish(ServerArgs(model_path="dummy", tp_size=world_size), role="test") self.assertEqual(get_parallel().attn_tp_size, world_size) built_at = [] with ( patch.object(parallel_state, "_WORLD", None), patch.object(parallel_state, "_TP", None), patch.object(parallel_state, "_DCP", None), patch.object(parallel_state, "_ATTN_CP", None), patch.object(parallel_state, "_ATTN_TP", None), patch.object(parallel_state, "_MOE_DP", None), patch.object(parallel_state, "_MOE_EP", None), patch.object(parallel_state, "_MOE_TP", None), patch.object(parallel_state, "_PP", None), patch.object(parallel_state, "_SELF_PP", None), patch("torch.distributed.is_initialized", return_value=True), patch("torch.distributed.get_world_size", return_value=world_size), patch("torch.distributed.get_rank", return_value=0), patch("torch.distributed.get_backend", return_value="nccl"), patch.object( parallel_state, "init_model_parallel_group", side_effect=lambda group_ranks, *a, **k: ( built_at.append(group_ranks), Mock(device_group=Mock()), )[1], ), patch.object(parallel_state, "get_world_group") as mock_world_group, ): mock_world_group.return_value = Mock(device_group=Mock(), local_rank=0) parallel_state.initialize_model_parallel() self.addCleanup(parallel_state.destroy_model_parallel) # The first group built is TP, one group spanning the published width. self.assertEqual(built_at[0], [list(range(world_size))]) self.assertEqual(get_parallel().attn_tp_size, world_size) class TestTheDerivedHalfIsDeclared(CustomTestCase): """The quotients are declared beside the leaves, in the same class. A namespace is one file and one class. `Parallel` says both what an operator can set and what that decides; the quotients are unannotated, so they are not dataclass fields and never reach the record. `ParallelContext` installs a property per declaration rather than carrying its own list, so the two cannot drift. """ def test_every_declared_quotient_has_a_property(self): from sglang.srt.arg_groups.arg_utils import Derived from sglang.srt.arg_groups.fields.parallel import Parallel declared = { name for name, value in vars(Parallel).items() if isinstance(value, Derived) } self.assertTrue(declared, "the derived half is empty") for name in declared: self.assertIsInstance( getattr(type(get_context().parallel), name, None), property, f"{name} is declared but no property was installed", ) def test_the_declared_set_is_what_derive_parallel_widths_produces(self): """The declaration is not a second list to keep in step: it names exactly the quotients the derivation returns.""" from sglang.srt.arg_groups.arg_utils import Derived from sglang.srt.arg_groups.fields.parallel import Parallel declared = { name for name, value in vars(Parallel).items() if isinstance(value, Derived) } produced = set( derive_parallel_widths( tp_size=8, attn_cp_size=1, attn_dp_size=2, moe_ep_size=1, moe_dp_size=1, dcp_size=1, dcp_enabled=False, ) ) self.assertEqual(declared, produced) def test_a_declared_quotient_is_not_a_record_field(self): """It has no operator input to preserve, and the record is what crosses a process boundary.""" from sglang.srt.arg_groups.arg_utils import Derived from sglang.srt.arg_groups.fields.parallel import Parallel from sglang.srt.server_args import ServerArgs fields = {f.name for f in msgspec.structs.fields(ServerArgs)} for name, value in vars(Parallel).items(): if isinstance(value, Derived): self.assertNotIn(name, fields) class TestAnEntryThatBuildsARunnerHandsOverItsPlacement(CustomTestCase): """`ModelRunner.__init__` reads a recorded identity, so an entry that publishes without a bundle and then builds one fails at construction. Every such entry is an `__main__`-reachable path, so nothing in the unit suite exercises it; the benchmark entry was found this way rather than by a test. This walks the sources instead: a module that publishes and builds a runner has to pass `ranks=`. """ def test_every_publisher_that_builds_a_runner_passes_a_bundle(self): import ast as _ast root = _pathlib.Path(next(iter(_sglang.__path__))).resolve() offenders = [] for path in root.rglob("*.py"): text = path.read_text(encoding="utf-8-sig") if "ModelRunner(" not in text or "publish(" not in text: continue tree = _ast.parse(text) builds = any( isinstance(n, _ast.Call) and getattr(n.func, "id", getattr(n.func, "attr", None)) == "ModelRunner" for n in _ast.walk(tree) ) if not builds: continue for node in _ast.walk(tree): if ( isinstance(node, _ast.Call) and getattr(node.func, "id", None) == "publish" and not any(kw.arg == "ranks" for kw in node.keywords) ): offenders.append(f"{path.relative_to(root)}:{node.lineno}") self.assertEqual( offenders, [], "these publish without a spawn bundle and then build a ModelRunner, " "whose construction reads a recorded identity:\n " + "\n ".join(offenders), ) class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase): """`parallel_state`'s getters are the definition, not a second spelling. Business code asks `get_parallel()`; a call that goes straight to the getter is a read the context cannot redirect, which is what a scope needs it to be able to do. The package that defines them is exempt -- a read there would go through the context back into itself -- and so is `multimodal_gen`, which has its own parallel state. """ #: Not topology. `get_self_pp_group` builds the single-rank group a draft #: pipeline scope installs, so there is nothing for the context to answer #: with until the scope has installed it. ALLOWED = { "get_self_pp_group", "get_default_distributed_backend", "get_mooncake_transfer_engine", } def _accessors(self): """Derived from the source, not listed here: a guard whose subject set is written by hand stops watching whatever gets added next.""" from sglang.srt.distributed import parallel_state as parallel_state_module source = _pathlib.Path(parallel_state_module.__file__).read_text().splitlines() return { line[len("def ") : line.index("(")] for line in source if line.startswith("def get_") or line.startswith("def is_") } def _callers(self, name): import re from sglang.srt.distributed import parallel_state as parallel_state_module root = _pathlib.Path(parallel_state_module.__file__).parents[2] pattern = re.compile(rf"(?