Record a process's placement at publish, not at group build (#40071)

test/registered/unit/constrained/test_grammar_manager.py

@@ -26,13 +26,19 @@ from sglang.srt.constrained.base_grammar_backend import (
 from sglang.srt.constrained.grammar_manager import GrammarManager
 from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject
 from sglang.srt.distributed.communication_tags import P2PTag
-from sglang.srt.runtime_context import get_context, publish, reset_context
+from sglang.srt.runtime_context import (
+    SpawnRanks,
+    get_context,
+    get_parallel,
+    publish,
+    reset_context,
+)
 from sglang.srt.sampling.sampling_params import (
     REQUEST_REASONING_END_TOKEN_IDS_KEY,
 )
 from sglang.srt.server_args import ServerArgs
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import enter_override
+from sglang.test.test_utils import enter_override, enter_scope
 
 register_cpu_ci(2.0, "base-a-test-cpu")
 
@@ -43,9 +49,10 @@ register_cpu_ci(est_time=5, suite="stage-b-test-cpu-intel")
 def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
     """Create a mock scheduler with necessary attributes.
 
-    The grammar manager reads its config from the bags, so the settings that
-    used to be hung off the mock are published instead. The caller resets the
-    context; every test here goes through `_GrammarFixture`.
+    The grammar manager reads its config and its place in the pipeline from
+    the context, so the settings that used to be hung off the mock are
+    published instead. The caller resets the context; every test here goes
+    through `_GrammarFixture`.
     """
     reset_context()
     server_args = ServerArgs(
@@ -56,7 +63,11 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
         constrained_json_whitespace_pattern=None,
         constrained_json_disable_any_whitespace=False,
     )
-    publish(server_args, role="scheduler")
+    publish(
+        server_args,
+        role="scheduler",
+        ranks=SpawnRanks(world_rank=0),
+    )
     scheduler = MagicMock()
     scheduler.server_args = server_args
     scheduler.model_config.request_selectable_think_end_id_sequences = None
@@ -66,8 +77,6 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False):
     scheduler.dp_tp_group.world_size = 1
     scheduler.dp_tp_group.first_rank = 0
     scheduler.dp_tp_group.is_first_rank = True
-    scheduler.ps.pp_rank = 0
-    scheduler.ps.pp_size = 1
     scheduler.pp_group = None
 
     return scheduler
@@ -769,8 +778,10 @@ class TestGrammarManagerPPSync(unittest.TestCase):
         enter_override(
             self, get_context().override_server_args(skip_tokenizer_init=True)
         )
-        scheduler.ps.pp_rank = pp_rank
-        scheduler.ps.pp_size = pp_size
+        # After that override, not before: installing a server-args override
+        # re-resolves the parallel bag from defaults, which puts `pp_size`
+        # back to 1 whatever was published.
+        enter_scope(self, get_parallel().override(pp_size=pp_size, pp_rank=pp_rank))
         scheduler.pp_group = pp_group
         mgr = GrammarManager(scheduler)
         mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend)

test/registered/unit/managers/test_auxiliary_output.py

@@ -27,6 +27,7 @@ from sglang.srt.runtime_context import publish, reset_context
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import published_topology
 
 register_cpu_ci(est_time=14, suite="base-a-test-cpu")
 
@@ -453,7 +454,6 @@ def test_active_observer_uses_observer_logits_preprocessing():
 def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
     event = object()
     scheduler = object.__new__(Scheduler)
-    scheduler.ps = SimpleNamespace(pp_size=1)
     scheduler.device_module = SimpleNamespace(Event=Mock(return_value=event))
     result = SimpleNamespace(
         logits_output=SimpleNamespace(auxiliary_device_output=object()),
@@ -463,7 +463,8 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
     )
     batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
 
-    Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
+    with published_topology():
+        Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
 
     assert result.copy_done is event
     result.copy_to_cpu.assert_called_once_with(
@@ -474,7 +475,6 @@ def test_scheduler_copies_auxiliary_output_for_non_overlap_results():
 
 def test_scheduler_preserves_pipeline_parallel_output_for_transport():
     scheduler = object.__new__(Scheduler)
-    scheduler.ps = SimpleNamespace(pp_size=2)
     scheduler.device_module = SimpleNamespace(Event=Mock())
     result = SimpleNamespace(
         logits_output=SimpleNamespace(auxiliary_device_output=object()),
@@ -484,7 +484,8 @@ def test_scheduler_preserves_pipeline_parallel_output_for_transport():
     )
     batch = SimpleNamespace(return_logprob=False, return_hidden_states=False)
 
-    Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
+    with published_topology(pp_size=2):
+        Scheduler._copy_auxiliary_output_to_cpu(scheduler, batch, result)
 
     assert result.copy_done is None
     result.copy_to_cpu.assert_not_called()
@@ -514,7 +515,6 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
     scheduler.is_generation = True
     scheduler.enable_overlap = False
     scheduler.enable_pdmux = True
-    scheduler.ps = SimpleNamespace(pp_size=1)
     scheduler.tp_worker = SimpleNamespace(
         forward_batch_split_prefill=Mock(return_value=result)
     )
@@ -535,9 +535,12 @@ def test_pdmux_split_prefill_schedules_auxiliary_output_copy():
         return_hidden_states=False,
     )
 
-    with patch(
-        "sglang.srt.managers.scheduler.resolve_forward_inputs"
-    ) as resolve_forward_inputs:
+    with (
+        published_topology(),
+        patch(
+            "sglang.srt.managers.scheduler.resolve_forward_inputs"
+        ) as resolve_forward_inputs,
+    ):
         output_result = Scheduler.run_batch(scheduler, batch)
 
     resolve_forward_inputs.assert_called_once_with(batch, scheduler.future_map)

test/registered/unit/managers/test_output_streamer_customized_info.py

@@ -14,6 +14,7 @@ from sglang.srt.utils.weight_versions import (
     record_weight_version_events,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import enter_scope, published_topology
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
@@ -102,6 +103,8 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
         )
         serving_patch.start()
         observability_patch.start()
+        # The streamer asks the context which rank it is streaming from.
+        enter_scope(self, published_topology(ranks={"dp_rank": 0}))
         self.addCleanup(serving_patch.stop)
         self.addCleanup(observability_patch.stop)
 

test/registered/unit/managers/test_pp_cp_rank_offsets.py

@@ -3,7 +3,11 @@ from types import SimpleNamespace
 from unittest.mock import patch
 
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import maybe_stub_sgl_kernel
+from sglang.test.test_utils import (
+    enter_scope,
+    maybe_stub_sgl_kernel,
+    published_topology,
+)
 
 maybe_stub_sgl_kernel()
 
@@ -32,6 +36,25 @@ def _make_ps(**overrides) -> ParallelState:
     return ParallelState.trivial(**defaults)
 
 
+def _published_topology():
+    """The topology `_make_ps` describes, published instead of stood in.
+
+    World rank 12 of a `tp=8, pp=2` world is `tp_rank=4` on the second stage,
+    which puts this process at `attn_dp_rank=1` with `attn_tp_rank=0`: the
+    context derives all of them from that one number and the widths, where the
+    record above had to be handed each.
+    """
+    return published_topology(
+        role="scheduler",
+        ranks={"world_rank": 12, "dp_rank": 1},
+        tp_size=8,
+        pp_size=2,
+        dp_size=2,
+        attn_cp_size=2,
+        enable_dp_attention=True,
+    )
+
+
 def _fake_group() -> SimpleNamespace:
     return SimpleNamespace(rank=0, ranks=[0], cpu_group=object())
 
@@ -158,6 +181,7 @@ class TestRequestReceiverBroadcast(unittest.TestCase):
 class TestPPCPRankOffsets(unittest.TestCase):
     def test_request_receiver_uses_cp_size_for_pp_recv_rank(self):
         ps = _make_ps()
+        enter_scope(self, _published_topology())
         calls = []
 
         def fake_point_to_point_pyobj(data, rank, group, src, dst, **kwargs):
@@ -176,6 +200,7 @@ class TestPPCPRankOffsets(unittest.TestCase):
 
     def test_pp_mixin_uses_cp_size_for_pyobj_send_and_recv_rank(self):
         ps = _make_ps()
+        enter_scope(self, _published_topology())
         scheduler = SchedulerPPMixin()
         scheduler.ps = ps
         scheduler.world_group = _fake_group()

test/registered/unit/managers/test_scheduler_chunked_abort_race.py

@@ -5,7 +5,12 @@ from types import SimpleNamespace
 from unittest.mock import Mock
 
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+from sglang.test.test_utils import (
+    CustomTestCase,
+    enter_scope,
+    maybe_stub_sgl_kernel,
+    published_topology,
+)
 
 maybe_stub_sgl_kernel()
 
@@ -38,13 +43,16 @@ def _make_scheduler(pending_req, *, chunked_req, running_reqs) -> Scheduler:
     sched.disaggregation_mode = None
     sched.enable_hicache_storage = False
     sched.mm_receiver = None
-    sched.ps = SimpleNamespace(pp_size=1)
     sched.running_batch = SimpleNamespace(reqs=running_reqs)
     sched.last_batch = None
     return sched
 
 
 class TestPendingChunkedAbortRace(CustomTestCase):
+    def setUp(self):
+        # The abort path asks the context for the pipeline width.
+        enter_scope(self, published_topology())
+
     def test_req_left_chunked_slot_is_aborted(self):
         req = _FakeReq("zombie_rid")
         sched = _make_scheduler(req, chunked_req=None, running_reqs=[req])

test/registered/unit/managers/test_scheduler_hicache_events.py

@@ -5,7 +5,11 @@ from types import SimpleNamespace
 from unittest.mock import Mock, call, patch
 
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import maybe_stub_sgl_kernel
+from sglang.test.test_utils import (
+    enter_scope,
+    maybe_stub_sgl_kernel,
+    published_topology,
+)
 
 maybe_stub_sgl_kernel()
 
@@ -114,7 +118,7 @@ class TestSchedulerHiCacheEvents(unittest.TestCase):
         s = self.scheduler
         s.init_pp_loop_state = Mock()
         s.pp_loop_size = 1
-        s.ps = SimpleNamespace(pp_size=2)
+        enter_scope(self, published_topology(pp_size=2))
         s.pp_group = SimpleNamespace(is_last_rank=True)
         s.running_mbs = [self.running_batch]
         s.last_mbs = [None]

test/registered/unit/managers/test_scheduler_timeouts.py

@@ -12,7 +12,12 @@ from unittest.mock import MagicMock, patch
 
 from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
+from sglang.test.test_utils import (
+    CustomTestCase,
+    enter_scope,
+    maybe_stub_sgl_kernel,
+    published_topology,
+)
 
 maybe_stub_sgl_kernel()
 
@@ -62,7 +67,6 @@ def _scheduler(waiting_queue, running_reqs=(), last_batch_reqs=()):
     s.enable_unified_cache_external_linker = False
     s.ipc_channels = SimpleNamespace(send_to_tokenizer=MagicMock())
     s.beam_coordinator = MagicMock()
-    s.ps = SimpleNamespace(pp_size=1)
     s.running_batch = _batch(list(running_reqs))
     s.last_batch = _batch(list(last_batch_reqs)) if last_batch_reqs else None
     return s
@@ -126,6 +130,10 @@ class TestWaitingTimeout(CustomTestCase):
 
 
 class TestRunningTimeout(CustomTestCase):
+    def setUp(self):
+        # The poll asks the context for the pipeline width.
+        enter_scope(self, published_topology())
+
     def test_emits_only_stale_unfinished_reqs_without_marking(self):
         now = time.perf_counter()
         stale = _req("stale", forward_entry=now - 10)

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

@@ -736,12 +736,15 @@ class TestStartupWeightLoadSchedulerRouting(CustomTestCase):
         # publishing a record rather than by standing one in.
         reset_context()
         publish(
-            ServerArgs(model_path="dummy", startup_weight_load_mode=mode),
+            ServerArgs(
+                model_path="dummy",
+                startup_weight_load_mode=mode,
+                pp_size=pp_size,
+            ),
             role="scheduler",
         )
         scheduler = Scheduler.__new__(Scheduler)
         scheduler.enable_overlap = enable_overlap
-        scheduler.ps = SimpleNamespace(pp_size=pp_size)
         scheduler.init_tp_model_worker = lambda: setattr(scheduler, "tp_worker", worker)
         scheduler.maybe_init_draft_worker = lambda: setattr(
             scheduler, "draft_worker", draft_worker

test/registered/unit/multimodal/rust/shared/test_rust_server_extension.py

@@ -34,7 +34,6 @@ class TestRustServerExtension(CustomTestCase):
                 attn_dp_rank=1,
                 tp_size=2,
                 tp_rank=1,
-                pp_size=1,
                 attn_tp_size=1,
                 attn_cp_size=1,
             ),
@@ -49,7 +48,9 @@ class TestRustServerExtension(CustomTestCase):
                 ),
             ),
             patch.object(
-                server_module, "get_parallel", return_value=SimpleNamespace(nnodes=1)
+                server_module,
+                "get_parallel",
+                return_value=SimpleNamespace(nnodes=1, pp_size=1),
             ),
             patch.object(ModelServer, "_partition_cores", return_value=(None, None)),
             patch.object(

test/registered/unit/observability/test_forward_pass_metrics.py

@@ -1,4 +1,4 @@
-from sglang.srt.runtime_context import get_context, get_observability
+from sglang.srt.runtime_context import get_context, get_observability, get_parallel
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=11, suite="base-a-test-cpu")
@@ -15,7 +15,7 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import (
     SchedulerMetricsReporter,
     _CacheHitRateWindow,
 )
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, enter_scope
 
 
 def _make_ps(**overrides) -> ParallelState:
@@ -292,6 +292,8 @@ class TestForwardPassMetrics(unittest.TestCase):
             kv_events_config=None,
         )
         scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=2, pp_rank=0, pp_size=1)
+        # The reporter asks the context whether this is the last stage.
+        enter_scope(self, get_parallel().override(pp_rank=0, pp_size=1))
         scheduler.enable_kv_cache_events = False
 
         with patch(
@@ -329,6 +331,8 @@ class TestForwardPassMetrics(unittest.TestCase):
             kv_events_config=None,
         )
         scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=0, pp_rank=0, pp_size=2)
+        # The reporter asks the context whether this is the last stage.
+        enter_scope(self, get_parallel().override(pp_rank=0, pp_size=2))
         scheduler.enable_kv_cache_events = False
 
         with patch(

test/registered/unit/test_runtime_context.py

@@ -39,6 +39,7 @@ from sglang.srt.runtime_context import (
     Flags,
     ParallelContext,
     RuntimeContext,
+    SpawnRanks,
     _FlagGroupBase,
     assert_published,
     derive_parallel_widths,
@@ -200,6 +201,148 @@ class TestTheTwoWorldWidths(_IsolatedOverrides):
                 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.
 
@@ -2107,5 +2250,108 @@ class TestTheDerivedHalfIsDeclared(CustomTestCase):
                 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 TestWhoAnswersDuringADraftScope(CustomTestCase):
+    """A draft worker runs in one process with the target, under a scope.
+
+    Two things have to hold for that to be workable, and neither is visible
+    from a single read: inside the scope every source agrees on the draft's
+    shape, and a reader that runs *outside* it still gets the draft's answer
+    from whatever it carried out.
+    """
+
+    def _single_member_group(self):
+        from sglang.srt.distributed.parallel_state import GroupCoordinator
+
+        group = GroupCoordinator.__new__(GroupCoordinator)
+        group.world_size = 1
+        group.rank_in_group = 0
+        return group
+
+    def _two_stage_pipeline(self):
+        """This process is stage 1 of 2, published the way a spawn states it."""
+        reset_context()
+        self.addCleanup(reset_context)
+        publish(
+            ServerArgs(model_path="dummy", pp_size=2),
+            role="scheduler",
+            ranks=SpawnRanks(world_rank=1),
+        )
+
+    def test_the_pipeline_swap_states_every_member_it_installs(self):
+        """`pp_size` is a configured leaf: unlike `pp_rank` it does not follow
+        the group being swapped underneath, so a scope that installs a group
+        without stating its width reports the target's."""
+        from sglang.srt.distributed import parallel_state
+
+        group = self._single_member_group()
+        self._two_stage_pipeline()
+        self.assertEqual(get_parallel().pp_size, 2)
+        with patch.object(parallel_state, "_PP", group):
+            with parallel_state.patch_pipeline_parallel_group(group):
+                self.assertEqual(get_parallel().pp_size, 1)
+                self.assertEqual(get_parallel().pp_rank, 0)
+                self.assertIs(get_parallel().pp_group, group)
+        self.assertEqual(get_parallel().pp_size, 2)
+        self.assertEqual(get_parallel().pp_rank, 1)
+
+    def test_a_report_built_for_a_runner_follows_that_runner(self):
+        """A weight check is an on-demand request served from the scheduler
+        loop, so it runs outside the scope that describes a draft runner. Its
+        report has to name the runner it was built for, which is why it holds
+        a record instead of asking the context."""
+        from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+        from sglang.srt.utils.weight_checker import WeightChecker
+
+        draft = ParallelState.trivial(pp_rank=0, pp_size=1)
+        checker = WeightChecker(get_model=lambda: None, ps=draft)
+        self._two_stage_pipeline()
+        info = checker._parallelism_info()
+        self.assertEqual((info.pp_rank, info.pp_size), (0, 1))
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/utils/test_weight_versions.py

@@ -15,7 +15,11 @@ from sglang.srt.utils.weight_versions import (
     truncate_weight_version_events,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import (
+    CustomTestCase,
+    enter_scope,
+    published_topology,
+)
 
 register_cpu_ci(est_time=12, suite="base-a-test-cpu")
 
@@ -300,11 +304,9 @@ class _SchedulerStub:
         waiting,
         chunked=None,
         last_batch=None,
-        pp_size=1,
         hisparse=None,
     ):
         self.serving = _ServingStub(version)
-        self.ps = SimpleNamespace(pp_size=pp_size)
         self.running_batch = SimpleNamespace(reqs=running)
         self.last_batch = last_batch
         self.waiting_queue = waiting
@@ -313,7 +315,9 @@ class _SchedulerStub:
 
 
 class TestSchedulerRecordWeightVersionChange(CustomTestCase):
-    def _scheduler(self, *args, **kwargs):
+    def _scheduler(self, *args, pp_size=1, **kwargs):
+        # The recording path asks the context for the pipeline width.
+        enter_scope(self, published_topology(pp_size=pp_size))
         scheduler = _SchedulerStub(*args, **kwargs)
         for name, value in (
             ("get_serving", scheduler.serving),