A runner and the objects it builds freeze the placement they describe (#40341)

test/registered/layers/mamba/test_mamba2_mixer.py

@@ -129,8 +129,27 @@ def mixer2_gated_norm_tensor_parallel(
         )
         mixer.weight.weight_loader(mixer.weight, weight)
 
-    # m2 reads tp via get_parallel().tp_size/rank — force it through the context.
-    with get_parallel().override(tp_size=1, tp_rank=0):
+    # m2 reads tp via get_parallel().tp_size/rank — state a single-rank topology
+    # through the context. Every width that follows from `tp_size` is named:
+    # narrowing one leaf and leaving the quotients behind describes no layout.
+    with get_parallel().override(
+        tp_size=1,
+        tp_rank=0,
+        tp_group=None,
+        attn_tp_size=1,
+        attn_tp_rank=0,
+        attn_tp_group=None,
+        attn_dp_size=1,
+        attn_dp_rank=0,
+        attn_cp_size=1,
+        attn_cp_rank=0,
+        attn_cp_group=None,
+        moe_ep_size=1,
+        moe_ep_rank=0,
+        moe_ep_group=None,
+        moe_dp_size=1,
+        moe_tp_size=1,
+    ):
         # create gated-norm without TP to compute reference
         mixer_single_gpu = m2.Mixer2RMSNormGated(
             full_hidden_size=hidden_size,

test/registered/unit/hardware_backend/mlx/test_max_running_requests.py

@@ -84,7 +84,7 @@ def _stub(
     stub._max_mamba_cache_size = max_mamba_cache_size
     stub._disable_radix_cache = disable_radix_cache
     stub.max_total_num_tokens = max_total_num_tokens
-    stub.ps = SimpleNamespace(attn_dp_size=dp_size)
+    stub.attn_dp_size = dp_size
     return stub
 
 
@@ -94,7 +94,7 @@ def _hybrid_stub_for_initialize(
     """A stub carrying what the real initialize() reads (hybrid path)."""
     stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
     stub._mlx_pool_size = pool
-    stub.ps = SimpleNamespace(attn_dp_size=1)
+    stub.attn_dp_size = 1
     stub.device = "cpu"  # read by init_ngram_embedding_manager
     # Evaluated as a call argument in init_ngram_embedding_manager before
     # the use_ngram_embedding short-circuit; never read.

test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py

@@ -91,7 +91,8 @@ def _make_prefill_aware_swa_runner(
         kv_cache_dtype=torch.float16,
         kv_cache_dtype_str="auto",
         page_size=1,
-        ps=SimpleNamespace(attn_cp_size=1, tp_size=1),
+        attn_cp_size=1,
+        tp_size=1,
         is_draft_worker=False,
         server_args=server_args,
         attention_chunk_size=None,

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

@@ -671,7 +671,7 @@ class TestModelRunnerStartupWeightLoadOwnership(CustomTestCase):
             elastic_ep_backend=None,
             is_ep_joiner=False,
         )
-        runner.ps = SimpleNamespace(tp_rank=0)
+        runner.tp_rank = 0
         return runner
 
     def test_start_delegates_to_the_manager(self):

test/registered/unit/spec/test_dflash_logits.py

@@ -248,7 +248,7 @@ def test_worker_folds_a_gate_admitted_quantized_selector_head(monkeypatch):
     worker = SimpleNamespace(
         block_size=8,
         selector=object(),
-        ps=SimpleNamespace(tp_rank=0),
+        model_runner=SimpleNamespace(tp_rank=0),
         draft_model=SimpleNamespace(lm_head=None),
         device="cpu",
         _selector_sampling_enabled=True,
@@ -282,7 +282,7 @@ def test_worker_warns_once_when_selector_sampling_is_disabled(monkeypatch):
         selector=object(),
         _selector_sampling_enabled=False,
         _warned_sampling_fallback=False,
-        ps=SimpleNamespace(tp_rank=0),
+        model_runner=SimpleNamespace(tp_rank=0),
     )
     batch = SimpleNamespace(sampling_info=SimpleNamespace(is_all_greedy=False))
 

test/registered/unit/test_runtime_context.py

@@ -2781,17 +2781,154 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
     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
+        report has to name the runner it was built for, which is why it reads
+        the placement once, where it is built, instead of asking again when the
+        request arrives."""
+        from sglang.srt.distributed import parallel_state
         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()
+        self.assertEqual(get_parallel().pp_size, 2)
+        group = self._single_member_group()
+        with patch.object(parallel_state, "_PP", group):
+            with parallel_state.patch_pipeline_parallel_group(group):
+                checker = WeightChecker(get_model=lambda: None)
+
+        # The scope has closed and the context answers the target's shape again.
+        self.assertEqual(get_parallel().pp_size, 2)
         info = checker._parallelism_info()
         self.assertEqual((info.pp_rank, info.pp_size), (0, 1))
 
 
+class TestNothingReadsThePlacementBeforeItIsFrozen(CustomTestCase):
+    """`ModelRunner.__init__` freezes its placement partway through.
+
+    A method called before that point reads an attribute that does not exist
+    yet, and only on the configuration that reaches it -- a remote weight
+    transporter, a NUMA binding -- so the suites say nothing and a GPU job is
+    where it surfaces. The order is what makes it wrong, so the order is what
+    is checked.
+    """
+
+    def _model_runner(self):
+        import ast as _ast
+
+        source = (_SRT / "model_executor" / "model_runner.py").read_text()
+        for node in _ast.parse(source).body:
+            if isinstance(node, _ast.ClassDef) and node.name == "ModelRunner":
+                return {m.name: m for m in node.body if isinstance(m, _ast.FunctionDef)}
+        raise AssertionError("ModelRunner not found")
+
+    def _frozen_names(self, methods):
+        import ast as _ast
+
+        return {
+            target.attr
+            for node in _ast.walk(methods["init_torch_distributed"])
+            if isinstance(node, _ast.Assign)
+            for target in node.targets
+            if isinstance(target, _ast.Attribute)
+            and isinstance(target.value, _ast.Name)
+            and target.value.id == "self"
+        }
+
+    def _reads(self, methods, fn, frozen, depth=0):
+        import ast as _ast
+
+        found = set()
+        for node in _ast.walk(fn):
+            if (
+                isinstance(node, _ast.Attribute)
+                and isinstance(node.value, _ast.Name)
+                and node.value.id == "self"
+                and isinstance(node.ctx, _ast.Load)
+                and node.attr in frozen
+            ):
+                found.add(node.attr)
+            if (
+                depth < 2
+                and isinstance(node, _ast.Call)
+                and isinstance(node.func, _ast.Attribute)
+                and isinstance(node.func.value, _ast.Name)
+                and node.func.value.id == "self"
+                and node.func.attr in methods
+                and node.func.attr != "init_torch_distributed"
+            ):
+                found |= self._reads(
+                    methods, methods[node.func.attr], frozen, depth + 1
+                )
+        return found
+
+    def _calls_before_the_freeze(self, methods):
+        import ast as _ast
+
+        calls = []
+        for statement in methods["__init__"].body:
+            for node in _ast.walk(statement):
+                if (
+                    isinstance(node, _ast.Call)
+                    and isinstance(node.func, _ast.Attribute)
+                    and isinstance(node.func.value, _ast.Name)
+                    and node.func.value.id == "self"
+                ):
+                    calls.append((node.lineno, node.func.attr))
+        calls.sort()
+        names = [name for _, name in calls]
+        self.assertIn(
+            "init_torch_distributed",
+            names,
+            "the freeze moved; this census is keyed on where it happens",
+        )
+        return calls[: names.index("init_torch_distributed")]
+
+    def test_no_method_called_before_the_freeze_reads_what_it_freezes(self):
+        methods = self._model_runner()
+        frozen = self._frozen_names(methods)
+        self.assertGreater(len(frozen), 5, "found no frozen names; census is broken")
+        offenders = []
+        for lineno, name in self._calls_before_the_freeze(methods):
+            fn = methods.get(name)
+            if fn is None:
+                continue
+            read = self._reads(methods, fn, frozen)
+            if read:
+                offenders.append(
+                    f"__init__:{lineno} self.{name}() reads {sorted(read)}"
+                )
+        self.assertEqual(
+            offenders,
+            [],
+            "these run before init_torch_distributed and read what it sets; "
+            "ask get_parallel() there, or move the call after the freeze:\n  "
+            + "\n  ".join(offenders),
+        )
+
+    def test_the_census_would_notice_one(self):
+        """The positive control: the walk has to find a read that is there."""
+        import ast as _ast
+        import textwrap
+
+        methods = {
+            m.name: m
+            for m in _ast.parse(
+                textwrap.dedent(
+                    """
+                    class R:
+                        def init_torch_distributed(self):
+                            self.tp_rank = 0
+
+                        def early(self):
+                            return self.tp_rank
+                    """
+                )
+            )
+            .body[0]
+            .body
+        }
+        frozen = self._frozen_names(methods)
+        self.assertEqual(frozen, {"tp_rank"})
+        self.assertEqual(self._reads(methods, methods["early"], frozen), {"tp_rank"})
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/utils/test_weight_checker.py

@@ -44,7 +44,7 @@ from sglang.srt.utils.weight_checker_comparator import (
     RawComparable,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import CustomTestCase, enter_scope, published_topology
 
 register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
 register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
@@ -551,7 +551,8 @@ class _WeightCheckerTestBase(CustomTestCase):
         torch.manual_seed(0)
         self.model = _TinyModel().cuda()
         runner = _FakeModelRunner(self.model)
-        self.checker = WeightChecker(get_model=lambda: runner.model, ps=runner.ps)
+        enter_scope(self, published_topology())
+        self.checker = WeightChecker(get_model=lambda: runner.model)
 
 
 class TestSnapshot(_WeightCheckerTestBase):
@@ -760,9 +761,16 @@ class _ChecksumTestBase(CustomTestCase):
             pp_rank=0,
             pp_size=1,
         )
-        self.checker = WeightChecker(
-            get_model=lambda: self.runner.model, ps=self.runner.ps
+        enter_scope(
+            self,
+            published_topology(
+                tp_size=4,
+                dp_size=2,
+                enable_dp_attention=True,
+                ranks={"world_rank": 2, "dp_rank": 1},
+            ),
         )
+        self.checker = WeightChecker(get_model=lambda: self.runner.model)
 
 
 class TestComputeChecksum(_ChecksumTestBase):