Take the parallel getters off the package's public surface (#40344)

python/sglang/srt/distributed/bootstrap.py

@@ -7,7 +7,9 @@ import torch
 import torch.distributed as dist
 
 from sglang.srt.configs.model_config import ModelConfig
-from sglang.srt.distributed import (
+from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
+from sglang.srt.distributed.parallel_state import (
+    _tag_groups_for_flashinfer_allreduce_only,
     get_default_distributed_backend,
     get_tp_group,
     get_world_group,
@@ -18,10 +20,6 @@ from sglang.srt.distributed import (
     set_mscclpp_all_reduce,
     set_torch_symm_mem_all_reduce,
 )
-from sglang.srt.distributed.gated_launch import maybe_wait_for_gated_launch
-from sglang.srt.distributed.parallel_state import (
-    _tag_groups_for_flashinfer_allreduce_only,
-)
 from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
 from sglang.srt.layers.layernorm_sp import initialize_layernorm_sp

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -427,7 +427,7 @@ def prealloc_symmetric_memory_pool(
     ):
         return
 
-    from sglang.srt.distributed import get_tp_group
+    from sglang.srt.distributed.parallel_state import get_tp_group
 
     # Memory allocation is tied to a cuda stream, use the forward stream
     with torch.get_device_module(device).stream(forward_stream):

python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py

@@ -469,8 +469,10 @@ class MultimemAllGatherer:
         self._state = self._UNINIT if enabled else None
         if self._state is self._UNINIT:
             # Lazy import avoids a module-load dependency on the distributed facade.
-            from sglang.srt.distributed import get_tp_group
-            from sglang.srt.distributed.parallel_state import in_the_same_node_as
+            from sglang.srt.distributed.parallel_state import (
+                get_tp_group,
+                in_the_same_node_as,
+            )
 
             tp_group = get_tp_group()
             # Only probe node topology when the deployment can actually span
@@ -527,7 +529,7 @@ class MultimemAllGatherer:
         if x.shape[-1] % _NUMEL_PER_THREAD != 0:
             return None
         try:
-            from sglang.srt.distributed import get_tp_group
+            from sglang.srt.distributed.parallel_state import get_tp_group
 
             tp_group = get_tp_group()
             if tp_group.world_size <= 1:

python/sglang/srt/distributed/gated_launch.py

@@ -9,7 +9,7 @@ import uvicorn
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 
-from sglang.srt.distributed import get_world_group
+from sglang.srt.distributed.parallel_state import get_world_group
 
 logger = logging.getLogger(__name__)
 

python/sglang/srt/distributed/parallel_state.py

@@ -3534,3 +3534,14 @@ for _name, _replacement in _CONTEXT_NAME_OF.items():
     if _fn is not None:
         globals()[_name] = _warn_if_called_from_outside(_name, _replacement)(_fn)
 del _name, _replacement, _fn
+
+
+# What `from sglang.srt.distributed import *` re-exports: everything public
+# except the deprecated getters. Business code reaches them through
+# `get_parallel()`, and the package that defines them imports them from this
+# module by name, so nothing needs the package path to reach one.
+__all__ = [
+    _public
+    for _public in list(globals())
+    if not _public.startswith("_") and _public not in _CONTEXT_NAME_OF
+]

python/sglang/srt/layers/dp_attention.py

@@ -15,9 +15,6 @@ from sglang.srt.arg_groups.model_override_base import (
 )
 from sglang.srt.distributed import (
     GroupCoordinator,
-)
-from sglang.srt.distributed import get_moe_dp_group as _get_moe_dp_group
-from sglang.srt.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -1069,15 +1066,15 @@ def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
 
 def get_moe_cp_group() -> GroupCoordinator:
     """Returns the MOE_DP group, which includes CP partners when attn_cp_size > moe_dp_size."""
-    return _get_moe_dp_group()
+    return get_parallel().moe_dp_group
 
 
 def get_moe_cp_rank() -> int:
-    return _get_moe_dp_group().rank_in_group
+    return get_parallel().moe_dp_group.rank_in_group
 
 
 def get_moe_cp_size() -> int:
-    return _get_moe_dp_group().world_size
+    return get_parallel().moe_dp_group.world_size
 
 
 def is_enable_moe_cp_allgather() -> bool:
@@ -1090,7 +1087,7 @@ def is_enable_moe_cp_allgather() -> bool:
 
 
 def moe_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
-    return _get_moe_dp_group().all_gather_into_tensor(output, input)
+    return get_parallel().moe_dp_group.all_gather_into_tensor(output, input)
 
 
 def attn_tp_all_gather(output_list: List[torch.Tensor], input: torch.Tensor):

test/registered/unit/test_runtime_context.py

@@ -2596,18 +2596,30 @@ class TestTheAccessorsHaveNoCallersOutsideTheirPackage(CustomTestCase):
         }
 
     def _callers(self, name):
+        """Every call in business code, including one hiding behind an import
+        alias -- `from ... import get_moe_dp_group as _g` then `_g()` is the
+        same reach past the context, and searching for the original spelling
+        alone reports zero while it is right there."""
         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"(?<![.\w]){re.escape(name)}\(")
         hits = []
         for path in root.rglob("*.py"):
             rel = path.relative_to(root).as_posix()
             if rel.startswith(("srt/distributed/", "multimodal_gen/", "test/")):
                 continue
-            for number, line in enumerate(path.read_text().splitlines(), 1):
+            text = path.read_text()
+            spellings = (
+                {name}
+                | set(re.findall(rf"import\s+{re.escape(name)}\s+as\s+(\w+)", text))
+                | set(re.findall(rf"^\s*{re.escape(name)}\s+as\s+(\w+),?$", text, re.M))
+            )
+            pattern = re.compile(
+                r"(?<![.\w])(?:" + "|".join(re.escape(s) for s in spellings) + r")\("
+            )
+            for number, line in enumerate(text.splitlines(), 1):
                 if line.lstrip().startswith(("def ", "#")):
                     continue
                 if pattern.search(line):
@@ -3150,6 +3162,72 @@ class TestTheRecordIsNeverWrittenTo(CustomTestCase):
         )
 
 
+class TestTheRetiredNamesAreGoneEverywhere(CustomTestCase):
+    """The package stopped re-exporting the getters and the build stopped
+    taking widths. Both are import-time or call-time failures in whatever tree
+    they survive in, and the trees beside the package have no suite to notice.
+    """
+
+    def _retired(self):
+        from sglang.srt.distributed.parallel_state import _CONTEXT_NAME_OF
+
+        return set(_CONTEXT_NAME_OF)
+
+    def test_nothing_imports_a_retired_name_from_the_package(self):
+        import ast as _ast
+
+        retired = self._retired()
+        offenders = []
+        for path in _sources():
+            for node in _ast.walk(_ast.parse(path.read_text(encoding="utf-8-sig"))):
+                if (
+                    isinstance(node, _ast.ImportFrom)
+                    and node.module == "sglang.srt.distributed"
+                ):
+                    for alias in node.names:
+                        if alias.name in retired:
+                            offenders.append(f"{path}:{node.lineno} {alias.name}")
+        self.assertEqual(
+            offenders,
+            [],
+            "these import a name the package no longer re-exports; import it "
+            "from parallel_state, or read get_parallel():\n  " + "\n  ".join(offenders),
+        )
+
+    def test_nothing_passes_a_width_to_the_build(self):
+        """`multimodal_gen` is out: it has a function of this name that builds
+        its own parallelism from its own degrees."""
+        import ast as _ast
+        import inspect
+
+        from sglang.srt.distributed.parallel_state import initialize_model_parallel
+
+        takes = set(inspect.signature(initialize_model_parallel).parameters)
+        offenders = []
+        for path in _sources():
+            if "multimodal_gen" in path.parts:
+                continue
+            for node in _ast.walk(_ast.parse(path.read_text(encoding="utf-8-sig"))):
+                if (
+                    isinstance(node, _ast.Call)
+                    and getattr(node.func, "id", getattr(node.func, "attr", None))
+                    == "initialize_model_parallel"
+                ):
+                    stale = [
+                        kw.arg for kw in node.keywords if kw.arg and kw.arg not in takes
+                    ]
+                    if stale or node.args:
+                        offenders.append(
+                            f"{path}:{node.lineno} {stale or 'positional'}"
+                        )
+        self.assertEqual(
+            offenders,
+            [],
+            "the build reads every width from the context; publish the "
+            "topology instead of passing it:\n  " + "\n  ".join(offenders),
+        )
+
+
 class TestNothingReadsThePlacementBeforeItIsFrozen(CustomTestCase):
     """`ModelRunner.__init__` freezes its placement partway through.