diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 9f113c3a7..889917d4f 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/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 diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 283a07130..8adb6e261 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/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): diff --git a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py b/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py index a7cf40bd7..f60b65ed7 100644 --- a/python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py +++ b/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: diff --git a/python/sglang/srt/distributed/gated_launch.py b/python/sglang/srt/distributed/gated_launch.py index da36301ae..09fd229d4 100644 --- a/python/sglang/srt/distributed/gated_launch.py +++ b/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__) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index f127f2fa5..e9e1e7f97 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 +] diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index ce6666375..7e947f44d 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/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): diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 5c1f8f081..08616231f 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/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"(?