minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

State the draft's whole topology in its scope, and read the rest from the context (#40339)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-21 12:19:38 -07:00
committed by GitHub
parent ae7a516ba7
commit 0db1a93adb
22 changed files with 364 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_loadstat_wire.py
+37 -30
View File
@@ -23,13 +23,12 @@ from unittest.mock import MagicMock, patch
import msgspec.msgpack
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.scheduler_components.load_publisher import (
LoadStat,
SchedulerLoadPublisher,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
from sglang.test.test_utils import CustomTestCase, published_topology
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
@@ -97,19 +96,21 @@ class TestLoadPublisherGating(CustomTestCase):
connect-style one.
"""
def _build(
self, *, config=ZMQ_ENDPOINT, dp_size=1, explicit="auto", **ps_overrides
):
def _build(self, *, config=ZMQ_ENDPOINT, explicit="auto", ranks=None, **topology):
"""Construct a publisher with the socket bind stubbed out, returning
(publisher, captured _open_pub_socket mock). Opts in via explicit="auto"
by default (the feature is off without it). dp_size lives on the ps,
which the publisher reads (no separate param to disagree with it)."""
with patch(
"sglang.srt.managers.scheduler_components.load_publisher._open_pub_socket"
) as open_sock:
by default (the feature is off without it). The topology is published
rather than overridden, so the ranks the publisher reads are the ones a
layout of that shape actually produces; every read happens in the
constructor."""
with (
published_topology(ranks=ranks, **topology),
patch(
"sglang.srt.managers.scheduler_components.load_publisher._open_pub_socket"
) as open_sock,
):
pub = SchedulerLoadPublisher(
kv_events_config=config,
ps=ParallelState.trivial(dp_size=dp_size, **ps_overrides),
load_publish_endpoint=explicit,
)
return pub, open_sock
@@ -130,14 +131,17 @@ class TestLoadPublisherGating(CustomTestCase):
def test_disabled_off_pp_rank_zero(self):
# Every PP stage shares attn_tp_rank/attn_cp_rank 0, so without the
# pp_rank gate they all bind the same load port.
pub, open_sock = self._build(pp_rank=1, pp_size=2)
pub, open_sock = self._build(pp_size=2, ranks={"world_rank": 1})
self.assertFalse(pub.enable)
open_sock.assert_not_called()
def test_disabled_off_attn_tp_and_cp_rank_zero(self):
for override in ({"attn_tp_rank": 1}, {"attn_cp_rank": 1}):
with self.subTest(**override):
pub, open_sock = self._build(**override)
for layout in (
{"tp_size": 2},
{"tp_size": 2, "attn_cp_size": 2},
):
with self.subTest(**layout):
pub, open_sock = self._build(ranks={"world_rank": 1}, **layout)
self.assertFalse(pub.enable)
open_sock.assert_not_called()
@@ -145,11 +149,16 @@ class TestLoadPublisherGating(CustomTestCase):
# Pure DP: attn_dp_size == 1 and every worker has attn_dp_rank == 0, so
# the publisher must key off dp_rank or all replicas collide on one
# port. kv 5557 + dp_size 4 => base 5561; rank 2 binds 5563.
_, open_sock = self._build(attn_dp_size=1, attn_dp_rank=0, dp_rank=2, dp_size=4)
_, open_sock = self._build(dp_size=4, ranks={"world_rank": 0, "dp_rank": 2})
open_sock.assert_called_once_with("tcp://*:5563")
def test_dp_attention_keys_the_load_port_by_attn_dp_rank(self):
_, open_sock = self._build(attn_dp_size=4, attn_dp_rank=3, dp_rank=0, dp_size=4)
_, open_sock = self._build(
tp_size=4,
dp_size=4,
enable_dp_attention=True,
ranks={"world_rank": 3, "dp_rank": 0},
)
open_sock.assert_called_once_with("tcp://*:5564")
def test_load_port_is_packed_after_the_kv_range(self):
@@ -259,10 +268,8 @@ class TestLoadPublisherGating(CustomTestCase):
_, open_sock = self._build(
explicit="tcp://*:7000",
attn_dp_size=1,
attn_dp_rank=0,
dp_rank=2,
dp_size=4,
ranks={"world_rank": 0, "dp_rank": 2},
)
open_sock.assert_called_once_with("tcp://*:7002")
@@ -289,11 +296,11 @@ class TestLoadPublisherGating(CustomTestCase):
"sglang.srt.managers.scheduler_components.load_publisher._open_pub_socket",
side_effect=zmq.ZMQError,
) as open_sock:
pub = SchedulerLoadPublisher(
kv_events_config=ZMQ_ENDPOINT,
ps=ParallelState.trivial(),
load_publish_endpoint="auto",
)
with published_topology():
pub = SchedulerLoadPublisher(
kv_events_config=ZMQ_ENDPOINT,
load_publish_endpoint="auto",
)
open_sock.assert_called_once() # the bind was attempted and failed
self.assertFalse(pub.enable)
pub.publish_load_stat(MagicMock(), force=True) # still a no-op
@@ -468,11 +475,11 @@ class TestLoadStatIntegration(CustomTestCase):
with _socket.socket() as probe:
probe.bind(("", 0))
port = probe.getsockname()[1]
pub = SchedulerLoadPublisher(
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
ps=ParallelState.trivial(),
load_publish_endpoint=f"tcp://*:{port}",
)
with published_topology():
pub = SchedulerLoadPublisher(
kv_events_config='{"publisher": "zmq", "endpoint": "tcp://*:5557"}',
load_publish_endpoint=f"tcp://*:{port}",
)
if pub.enable:
break
self.assertTrue(pub.enable, "load socket never bound a free port")
test/registered/unit/managers/test_mm_shm_error_consensus.py
+9 -17
View File
@@ -91,7 +91,7 @@ def _request(feature, rid: str = "vlm-request") -> TokenizedEmbeddingReqInput:
)
def _receiver(tp_size: int = 1) -> SchedulerRequestReceiver:
def _receiver() -> SchedulerRequestReceiver:
group = SimpleNamespace(rank=0, ranks=[0], cpu_group=object())
return SchedulerRequestReceiver(
recv_from_tokenizer=None,
@@ -99,14 +99,6 @@ def _receiver(tp_size: int = 1) -> SchedulerRequestReceiver:
recv_skipper=None,
input_blocker=None,
mm_receiver=None,
ps=SimpleNamespace(
pp_rank=0,
tp_size=tp_size,
attn_tp_rank=0,
attn_cp_rank=0,
attn_tp_size=1,
attn_cp_size=1,
),
tp_group=group,
tp_cpu_group=group,
attn_tp_group=group,
@@ -131,8 +123,8 @@ def _run_consensus_rank(rank: int, world_size: int, init_file: str) -> None:
)
try:
req = _request(_failed_pointer() if rank == 1 else _successful_pointer())
parallel = SimpleNamespace(enable_dp_attention=False)
receiver = _receiver(tp_size=world_size)
parallel = SimpleNamespace(enable_dp_attention=False, tp_size=world_size)
receiver = _receiver()
object.__setattr__(receiver, "tp_cpu_group", torch.distributed.group.WORLD)
with (
patch(
@@ -161,7 +153,7 @@ def _run_image_receiver(rank, init_file, pipe):
backend="gloo", init_method=Path(init_file).as_uri(), rank=rank, world_size=2
)
try:
receiver = _receiver(tp_size=2)
receiver = _receiver()
object.__setattr__(receiver, "tp_cpu_group", torch.distributed.group.WORLD)
torch.distributed.barrier()
torch.distributed.all_reduce(torch.zeros(1))
@@ -179,7 +171,7 @@ def _run_image_receiver(rank, init_file, pipe):
),
patch(
"sglang.srt.managers.scheduler_components.request_receiver.get_parallel",
return_value=SimpleNamespace(enable_dp_attention=False),
return_value=SimpleNamespace(enable_dp_attention=False, tp_size=2),
),
):
for base in [30, 90]:
@@ -382,7 +374,7 @@ class TestShmRequestFailureConsensus(unittest.TestCase):
def test_local_materialization_failure_becomes_request_error(self):
req = _request(_failed_pointer())
parallel = SimpleNamespace(enable_dp_attention=False)
parallel = SimpleNamespace(enable_dp_attention=False, tp_size=1)
with (
patch(
@@ -406,7 +398,7 @@ class TestShmRequestFailureConsensus(unittest.TestCase):
def test_peer_failure_rejects_the_local_request(self):
req = _request(torch.zeros(1))
parallel = SimpleNamespace(enable_dp_attention=False)
parallel = SimpleNamespace(enable_dp_attention=False, tp_size=2)
def inject_peer_failure(mask, **kwargs):
mask.fill_(1)
@@ -429,7 +421,7 @@ class TestShmRequestFailureConsensus(unittest.TestCase):
side_effect=inject_peer_failure,
) as all_reduce,
):
_receiver(tp_size=2)._finalize_shm_features([req])
_receiver()._finalize_shm_features([req])
all_reduce.assert_called_once()
self.assertIsInstance(req.mm_inputs, MMInputsProcessError)
@@ -438,7 +430,7 @@ class TestShmRequestFailureConsensus(unittest.TestCase):
failed_req = _request(torch.zeros(1), rid="failed")
healthy_req = _request(torch.zeros(1), rid="healthy")
batch = BatchTokenizedEmbeddingReqInput(batch=[failed_req, healthy_req])
parallel = SimpleNamespace(enable_dp_attention=False)
parallel = SimpleNamespace(enable_dp_attention=False, tp_size=1)
def materialize(req):
if req.rid == "failed":
test/registered/unit/managers/test_pp_cp_rank_offsets.py
+14 -40
View File
@@ -14,7 +14,6 @@ from sglang.test.test_utils import (
maybe_stub_sgl_kernel()
from sglang.srt.distributed.parallel_state_wrapper import ParallelState # noqa: E402
from sglang.srt.managers.scheduler_components.request_receiver import ( # noqa: E402
SchedulerRequestReceiver,
)
@@ -23,29 +22,12 @@ from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin # noqa: E40
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
def _make_ps(**overrides) -> ParallelState:
defaults = dict(
tp_size=8,
pp_rank=1,
pp_size=2,
dp_rank=None,
attn_tp_size=2,
attn_cp_size=2,
attn_dp_rank=1,
attn_dp_size=2,
moe_dp_rank=None,
)
defaults.update(overrides)
return ParallelState.trivial(**defaults)
def _published_topology():
"""The topology `_make_ps` describes, published instead of stood in.
"""The topology these tests run 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.
context derives all of them from that one number and the widths.
"""
return published_topology(
role="scheduler",
@@ -62,7 +44,7 @@ def _fake_group() -> SimpleNamespace:
return SimpleNamespace(rank=0, ranks=[0], cpu_group=object())
def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver:
def _make_receiver() -> SchedulerRequestReceiver:
tp_group = _fake_group()
attn_tp_group = _fake_group()
attn_cp_group = _fake_group()
@@ -73,7 +55,6 @@ def _make_receiver(ps: ParallelState) -> SchedulerRequestReceiver:
recv_skipper=None,
input_blocker=None,
mm_receiver=None,
ps=ps,
tp_group=tp_group,
tp_cpu_group=tp_group,
attn_tp_group=attn_tp_group,
@@ -97,19 +78,17 @@ class TestRequestReceiverBroadcast(unittest.TestCase):
# Decode uses pure DP attention (attn_tp=attn_cp=1). The DP controller
# sends control requests to every local leader, so no per-tick Gloo
# broadcast should remain in SchedulerRequestReceiver.
ps = SimpleNamespace(
receiver = _make_receiver()
control_req = SimpleNamespace(kind="control")
parallel = SimpleNamespace(
enable_dp_attention=True,
enable_dp_attention_local_control_broadcast=True,
attn_tp_rank=0,
attn_cp_rank=0,
attn_tp_size=1,
attn_cp_size=1,
tp_size=32,
)
receiver = _make_receiver(ps)
control_req = SimpleNamespace(kind="control")
parallel = SimpleNamespace(
enable_dp_attention=True,
enable_dp_attention_local_control_broadcast=True,
)
with (
patch(
@@ -133,19 +112,17 @@ class TestRequestReceiverBroadcast(unittest.TestCase):
broadcast.assert_not_called()
def test_default_control_uses_full_tp_broadcast(self):
ps = SimpleNamespace(
receiver = _make_receiver()
control_req = SimpleNamespace(kind="control")
parallel = SimpleNamespace(
enable_dp_attention=True,
enable_dp_attention_local_control_broadcast=False,
attn_tp_rank=0,
attn_cp_rank=0,
attn_tp_size=1,
attn_cp_size=1,
tp_size=32,
)
receiver = _make_receiver(ps)
control_req = SimpleNamespace(kind="control")
parallel = SimpleNamespace(
enable_dp_attention=True,
enable_dp_attention_local_control_broadcast=False,
)
with (
patch(
@@ -183,7 +160,6 @@ 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 = []
@@ -191,7 +167,7 @@ class TestPPCPRankOffsets(unittest.TestCase):
calls.append((rank, src, dst))
return ["req"]
receiver = _make_receiver(ps)
receiver = _make_receiver()
with patch(
"sglang.srt.managers.scheduler_components.request_receiver."
"point_to_point_pyobj",
@@ -202,10 +178,8 @@ class TestPPCPRankOffsets(unittest.TestCase):
self.assertEqual(calls, [(12, 4, 12)])
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()
scheduler.attn_tp_group = _fake_group()
scheduler.attn_tp_cpu_group = _fake_group()
test/registered/unit/test_runtime_context.py
+113
View File
@@ -2338,6 +2338,119 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
self.assertEqual(get_parallel().pp_size, 2)
self.assertEqual(get_parallel().pp_rank, 1)
def _group(self, world_size, rank):
from sglang.srt.distributed.parallel_state import GroupCoordinator
group = GroupCoordinator.__new__(GroupCoordinator)
group.world_size = world_size
group.rank_in_group = rank
return group
def test_the_tensor_swap_states_the_draft_has_no_attention_replica(self):
"""The draft runs the whole model on the group being installed. Its
attention identity is therefore that group, with one replica -- while
the target this process also serves is attention-DP over four ranks."""
from sglang.srt.distributed import parallel_state
reset_context()
self.addCleanup(reset_context)
publish(
ServerArgs(
model_path="dummy", tp_size=4, dp_size=2, enable_dp_attention=True
),
role="scheduler",
ranks=SpawnRanks(world_rank=0, dp_rank=0),
)
self.assertEqual(get_parallel().attn_dp_size, 2)
self.assertEqual(get_parallel().attn_tp_size, 2)
group = self._group(world_size=2, rank=1)
with patch.object(parallel_state, "_TP", group):
with parallel_state.patch_tensor_parallel_group(group, owns_attention=True):
parallel = get_parallel()
self.assertEqual(parallel.tp_size, 2)
self.assertEqual(parallel.attn_tp_size, 2)
self.assertEqual(parallel.attn_tp_rank, 1)
self.assertEqual(parallel.attn_dp_size, 1)
self.assertEqual(parallel.attn_dp_rank, 0)
self.assertEqual(parallel.attn_cp_size, 1)
self.assertEqual(parallel.attn_cp_rank, 0)
# `dp_size` is the deployment's replica count, not a property
# of the group being installed, so the scope leaves it alone --
# `require_mlp_tp_gather` asserts on it under dp attention.
self.assertEqual(parallel.dp_size, 2)
# The whole point of stating the rest: the identity the
# override path and the group build both check holds in here.
self.assertEqual(
parallel.tp_size,
parallel.attn_tp_size
* parallel.attn_dp_size
* parallel.attn_cp_size,
)
self.assertEqual(get_parallel().attn_dp_size, 2)
self.assertEqual(get_parallel().dp_size, 2)
def test_every_caller_says_whether_the_draft_owns_its_attention(self):
"""The scope cannot work it out from the group it is handed: the same
call site passes an attention-TP slice for one draft and the target's
whole TP group for another, and the two want opposite answers. So the
worker states it, and a caller that forgets is the bug this catches --
`owns_attention` has no default, but a missing one is a TypeError only
on the path that runs, and these paths need a GPU and a draft model."""
import ast
package = _pathlib.Path(next(iter(_sglang.__path__))).resolve()
checkout = package.parents[1]
roots = [package] + [
checkout / name for name in ("test",) if (checkout / name).is_dir()
]
missing = []
for path in (q for root in roots for q in root.rglob("*.py")):
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}")
self.assertEqual(missing, [], "these enter the draft scope without saying")
def test_a_full_width_swap_leaves_the_attention_layout_alone(self):
"""The other caller. A draft built outside any scope carries the
target's whole TP group, and the graph capture installs *that* -- so
the process is still one of two attention-DP replicas and still gathers
with the other one. Narrowing here would claim a replica count it does
not have, and the reader that acts on it is a collective: the DP gather
takes its buffer size from the replica count and its communicator from
this group, so the two stop agreeing and the all-gather is refused."""
from sglang.srt.distributed import parallel_state
reset_context()
self.addCleanup(reset_context)
publish(
ServerArgs(
model_path="dummy", tp_size=4, dp_size=2, enable_dp_attention=True
),
role="scheduler",
ranks=SpawnRanks(world_rank=0, dp_rank=0),
)
whole_tp = self._group(world_size=4, rank=0)
with patch.object(parallel_state, "_TP", whole_tp):
with parallel_state.patch_tensor_parallel_group(
whole_tp, owns_attention=False
):
parallel = get_parallel()
self.assertEqual(parallel.tp_size, 4)
self.assertEqual(parallel.attn_dp_size, 2)
self.assertEqual(parallel.attn_tp_size, 2)
self.assertEqual(parallel.dp_size, 2)
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