Retire the per-runner parallel record (#40343)

This commit is contained in:
Cheng Wan
2026-09-21 12:26:40 -07:00
committed by GitHub
parent 73f071db52
commit 970e946e4f
79 changed files with 395 additions and 463 deletions
@@ -15,7 +15,6 @@ import torch
from sglang.benchmark.one_batch import TreeCacheNamespace
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import (
@@ -65,7 +64,6 @@ class TestForwardSplitPrefill(CustomTestCase):
model_config=cls.model_config,
mem_fraction_static=cls.server_args.mem_fraction_static,
gpu_id=0,
ps=ParallelState.trivial(tp_size=cls.tp_size),
nccl_port=cls.port_args.nccl_port,
server_args=cls.server_args,
)
-2
View File
@@ -9,7 +9,6 @@ import torch.nn.functional as F
from transformers import AutoModel, AutoProcessor, AutoTokenizer
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
from sglang.srt.managers.mm_utils import embed_mm_inputs, init_mm_embedding_cache
from sglang.srt.managers.schedule_batch import (
@@ -151,7 +150,6 @@ class VisionLLMLogitsBase(unittest.IsolatedAsyncioTestCase):
model_config=ModelConfig(self.model_path, model_override_args="{}"),
mem_fraction_static=0.8,
gpu_id=0,
ps=ParallelState.trivial(),
nccl_port=12435,
server_args=server_args,
)
@@ -11,7 +11,6 @@ from sglang.srt.disaggregation.decode import (
)
from sglang.srt.disaggregation.fake.conn import FakeKVManager, FakeKVReceiver
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.schedule_batch import FINISH_ABORT
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.runtime_context import get_context, publish, reset_context
@@ -440,7 +439,6 @@ class TestDecodeQueueCleanup(CustomTestCase):
scheduler.last_batch = None
scheduler.cur_batch_for_debug = None
scheduler.enable_overlap = False
scheduler.ps = ParallelState.trivial()
scheduler.running_mbs = []
scheduler.waiting_queue = []
scheduler.grammar_manager = SimpleNamespace(grammar_queue=[])
@@ -131,7 +131,6 @@ class TestHandlePdRoleSwitch(unittest.TestCase):
def test_rejected_when_decode_graph_headroom_is_insufficient(self):
s = self._scheduler(DisaggregationMode.PREFILL)
s.device = "cuda"
s.ps = SimpleNamespace(gpu_id=0)
s.tp_worker.get_decode_cuda_graph_bs.return_value = []
with patch.object(role_switch, "get_available_gpu_memory", return_value=0.5):
out = Scheduler.handle_pd_role_switch(
@@ -151,7 +150,6 @@ class TestHandlePdRoleSwitch(unittest.TestCase):
def test_decode_graph_headroom_allows_flip(self):
s = self._scheduler(DisaggregationMode.PREFILL)
s.device = "cuda"
s.ps = SimpleNamespace(gpu_id=0)
s.tp_worker.get_decode_cuda_graph_bs.return_value = []
with patch.object(role_switch, "get_available_gpu_memory", return_value=1.0):
out = Scheduler.handle_pd_role_switch(
@@ -43,18 +43,18 @@ def test_dp_leaders_reuse_node_local_ports(
for dp_rank, tp_rank in enumerate(ranks):
scheduler = SimpleNamespace(
server_args=SimpleNamespace(),
ps=SimpleNamespace(
tp_rank=tp_rank,
tp_size=parallel.tp_size,
pp_size=parallel.pp_size,
attn_tp_size=parallel.attn_tp_size,
attn_cp_size=parallel.attn_cp_size,
attn_dp_rank=dp_rank,
dp_size=dp_size,
),
model_config=SimpleNamespace(is_multimodal=False),
)
ports.append(rust_server.RustServer.launch(scheduler).http_port)
# Where this rank sits, stated whole: the attention rank follows
# from the TP rank and the attention-TP width, and the identities
# refuse the combination if it describes no real layout.
with parallel.override(
tp_rank=tp_rank,
attn_dp_rank=dp_rank,
attn_tp_rank=tp_rank % parallel.attn_tp_size,
attn_cp_rank=0,
):
ports.append(rust_server.RustServer.launch(scheduler).http_port)
calls = extension.return_value.Server.call_args_list
assert [c.kwargs["port_offset"] for c in calls] == expected
@@ -23,7 +23,6 @@ _HAS_MLX = importlib.util.find_spec("mlx") is not None
_SKIP_REASON = "requires mlx"
if _HAS_MLX:
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.hardware_backend.mlx.model_runner_stub import (
MLX_AUX_STATE_SIZE_MAX_RUNNING_REQUESTS_RATIO as RATIO,
)
@@ -62,7 +61,7 @@ def _stub_for_initialize(
stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
stub._mlx_pool_size = pool_size
stub.device = "cpu"
stub.ps = ParallelState.trivial(dp_size=dp_size, attn_dp_size=attn_dp_size)
stub.attn_dp_size = attn_dp_size
stub.server_args = server_args
stub.model_config = SimpleNamespace(
is_hybrid_swa=False,
@@ -208,14 +208,7 @@ class TestSchedulerProfilerManagerMPS(unittest.TestCase):
SchedulerProfilerManager,
)
class FakePS:
tp_rank = dp_rank = pp_rank = moe_ep_rank = 0
dp_size = pp_size = moe_ep_size = 1
gpu_id = 0
mgr = SchedulerProfilerManager(
ps=FakePS(), dp_tp_cpu_group=None, get_forward_ct=lambda: 0
)
mgr = SchedulerProfilerManager(dp_tp_cpu_group=None, get_forward_ct=lambda: 0)
mgr._init_profile(output_dir, None, None, None, None, None, False, "test")
return mgr
@@ -228,7 +228,7 @@ class TestMambaPrefillTrackMetadata(unittest.TestCase):
prefill_attention_backend_str="torch_native",
ngram_embedding_manager=SimpleNamespace(enabled=False),
lora_manager=None,
ps=SimpleNamespace(attn_dcp_size=1),
attn_dcp_size=1,
attn_backend=SimpleNamespace(
get_cpu_graph_seq_len_fill_value=lambda: 1,
get_cuda_graph_seq_len_fill_value=lambda: 1,
@@ -20,7 +20,11 @@ from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
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=1, suite="base-a-test-cpu")
@@ -58,6 +62,12 @@ def load_mlx_scheduler_module():
class TestSchedulerIdleStepCounters(CustomTestCase):
def setUp(self):
super().setUp()
# The loop asks the context where this process sits; nothing here
# builds a process group, so the placement arrives by publishing one.
enter_scope(self, published_topology(role="scheduler"))
@parameterized.expand(
[
(
@@ -406,7 +416,6 @@ class TestSchedulerIdleStepCounters(CustomTestCase):
scheduler.forward_ct = 0
scheduler.processed_tokens_counter = 0
scheduler.spec_algorithm = SpeculativeAlgorithm.NONE
scheduler.ps = SimpleNamespace(pp_rank=0, attn_tp_rank=0, attn_cp_rank=0)
scheduler._poll_timeout_aborts = Mock(return_value=[])
scheduler.scheduler_stage_metrics = None
scheduler.metrics_reporter = SimpleNamespace(record_scheduler_active=Mock())
@@ -455,7 +464,7 @@ class TestSchedulerIdleStepCounters(CustomTestCase):
return scheduler
def prepare_pp_scheduler(self, scheduler):
scheduler.ps.pp_size = 2
enter_scope(self, get_parallel().override(pp_size=2, pp_rank=0))
scheduler.pp_group = SimpleNamespace(is_last_rank=True)
scheduler.forward_stream_ctx = nullcontext()
scheduler.forward_stream = Mock()
@@ -148,7 +148,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(
stream_interval=1,
enable_request_time_stats_logging=False,
@@ -181,7 +180,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(
stream_interval=1,
enable_request_time_stats_logging=False,
@@ -218,7 +216,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(
stream_interval=1,
enable_request_time_stats_logging=False,
@@ -258,7 +255,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(
stream_interval=1,
enable_request_time_stats_logging=False,
@@ -286,7 +282,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(),
is_generation=True,
spec_algorithm=SpeculativeAlgorithm.NONE,
@@ -312,7 +307,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
streamer = Streamer(
send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
tree_cache=None,
ps=SimpleNamespace(dp_rank=0, attn_tp_rank=0),
server_args=SimpleNamespace(),
is_generation=True,
spec_algorithm=SpeculativeAlgorithm.NONE,
@@ -334,7 +328,6 @@ class TestOutputStreamerCustomizedInfo(unittest.TestCase):
Streamer(
send_to_detokenizer=SimpleNamespace(),
tree_cache=None,
ps=SimpleNamespace(),
server_args=SimpleNamespace(),
is_generation=True,
spec_algorithm=SpeculativeAlgorithm.NONE,
@@ -11,7 +11,6 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.runtime_context import (
get_memory,
get_parallel,
@@ -36,10 +35,17 @@ def mock_cpu_env(kv_size=2, tp_size=1, swa_eviction_interval=4):
with (
patch("torch._utils._element_size", return_value=kv_size),
# A width is a whole topology: state the TP siblings the identities
# relate it to, not the attention share alone.
# The whole attention triple, not just one leaf: a width that does not
# factor describes no layout.
get_parallel().override(
tp_size=tp_size, attn_tp_size=tp_size, moe_tp_size=tp_size
tp_size=tp_size,
attn_tp_size=tp_size,
attn_dp_size=1,
attn_cp_size=1,
moe_ep_size=1,
moe_ep_group=None,
moe_dp_size=1,
moe_tp_size=tp_size,
),
envs.SGLANG_SWA_EVICTION_INTERVAL.override(swa_eviction_interval),
):
@@ -174,7 +180,8 @@ def _make_model_runner(
mr.layer_info = SimpleNamespace(
start_layer=0, end_layer=num_layers, num_effective_layers=num_layers
)
mr.ps = ParallelState.trivial()
mr.attn_dp_size = 1
mr.pp_size = 1
mr.pp_group = SimpleNamespace(rank_in_group=0)
mr.spec_aux_config = SimpleNamespace(
eagle_draft_num_layers=None,
@@ -1271,7 +1278,8 @@ class TestSWAPoolFloor(CustomTestCase):
kv_cache_dtype_str="fp8_e4m3",
model_config=cfg,
layer_info=SimpleNamespace(start_layer=0, end_layer=40),
ps=SimpleNamespace(pp_size=1, attn_dp_size=1),
pp_size=1,
attn_dp_size=1,
sliding_window_size=128,
page_size=256,
spec_algorithm=spec,
@@ -29,14 +29,6 @@ class TestRustServerExtension(CustomTestCase):
self.server.start_mm_workers(sentinel.spec, 8)
scheduler = SimpleNamespace(
ps=SimpleNamespace(
dp_size=2,
attn_dp_rank=1,
tp_size=2,
tp_rank=1,
attn_tp_size=1,
attn_cp_size=1,
),
model_config=SimpleNamespace(is_multimodal=True),
)
with (
@@ -50,7 +42,16 @@ class TestRustServerExtension(CustomTestCase):
patch.object(
server_module,
"get_parallel",
return_value=SimpleNamespace(nnodes=1, pp_size=1),
return_value=SimpleNamespace(
nnodes=1,
pp_size=1,
dp_size=2,
attn_dp_rank=1,
tp_size=2,
tp_rank=1,
attn_tp_size=1,
attn_cp_size=1,
),
),
patch.object(ModelServer, "_partition_cores", return_value=(None, None)),
patch.object(
@@ -9,7 +9,6 @@ import unittest
from unittest.mock import patch
from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.scheduler_components.metrics_reporter import (
PrefillStats,
SchedulerMetricsReporter,
@@ -18,16 +17,6 @@ from sglang.srt.managers.scheduler_components.metrics_reporter import (
from sglang.test.test_utils import CustomTestCase, enter_scope
def _make_ps(**overrides) -> ParallelState:
"""Build a ParallelState with reasonable defaults for tests; override fields via kwargs."""
defaults = dict(
dp_rank=None,
moe_dp_rank=None,
)
defaults.update(overrides)
return ParallelState.trivial(**defaults)
class _FakeReq:
def __init__(
self,
@@ -75,11 +64,30 @@ class _DummyPublisherThread:
def _publish_server_args(test, **fields):
"""Publish a config for the reporter under test and return the instance."""
"""Publish a config for the reporter under test and return the instance.
The collector asks the context where this process sits, so the ranks are
stated too: without them a rank read falls through to a process group that
a unit test has not built.
"""
fields.setdefault("decode_log_interval", 40)
override = get_context().override_server_args(**fields)
server_args = override.install()
test.addCleanup(override.restore)
enter_scope(
test,
get_parallel().override(
tp_rank=0,
attn_tp_rank=0,
attn_cp_rank=0,
moe_ep_rank=0,
attn_dp_rank=0,
dp_rank=0,
moe_ep_size=1,
moe_dp_size=1,
moe_tp_size=1,
),
)
return server_args
@@ -93,8 +101,6 @@ def _make_reporter(test, scheduler) -> SchedulerMetricsReporter:
enable_mfu_metrics=False,
enable_forward_pass_metrics=False,
)
if not hasattr(scheduler, "ps"):
scheduler.ps = ParallelState.trivial()
if not hasattr(scheduler, "kv_events_publisher"):
scheduler.kv_events_publisher = types.SimpleNamespace(
init_kv_events=lambda *a, **kw: None,
@@ -291,9 +297,9 @@ class TestForwardPassMetrics(unittest.TestCase):
forward_pass_metrics_ipc_name=None,
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))
# The reporter asks the context whether this is the last stage, and
# which replica it is reporting for.
enter_scope(self, get_parallel().override(pp_rank=0, pp_size=1, dp_rank=2))
scheduler.enable_kv_cache_events = False
with patch(
@@ -330,7 +336,6 @@ class TestForwardPassMetrics(unittest.TestCase):
forward_pass_metrics_ipc_name=None,
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
@@ -175,7 +175,6 @@ class TestDraftPerRunnerConfig(CustomTestCase):
scheduler.tp_worker = SimpleNamespace(
model_runner=SimpleNamespace(model_config=SimpleNamespace(context_len=4096))
)
scheduler.ps = SimpleNamespace(gpu_id=0)
scheduler.nccl_port = 0
scheduler.spec_algorithm = SimpleNamespace(
is_none=lambda: False,
@@ -6,6 +6,7 @@ import torch
from sglang.srt.layers.aux_hidden_states import pack_aux_hidden_states
from sglang.srt.models.dspark import DSparkDraftMixin
from sglang.srt.runtime_context import get_parallel
from sglang.srt.speculative.dspark_components.dspark_kv_inject import (
TargetHiddenKvInjector,
)
@@ -48,16 +49,13 @@ class DSparkTargetHiddenProjectionTest(CustomTestCase):
),
)
with (
mock.patch(
"sglang.srt.speculative.dspark_components.dspark_worker_v2.get_pp_group",
return_value=SimpleNamespace(is_last_rank=False),
),
get_parallel().override(pp_group=SimpleNamespace(is_last_rank=False)),
mock.patch(
"sglang.srt.speculative.dspark_components.dspark_worker_v2.get_schedule",
return_value=SimpleNamespace(page_size=1),
),
):
worker = DSparkWorkerV2(None, 0, None, 0, target)
worker = DSparkWorkerV2(None, 0, 0, target)
worker.alloc_memory_pool()
worker.init_attention_backends()
worker.init_cuda_graphs()
+83 -1
View File
@@ -46,6 +46,7 @@ from sglang.srt.runtime_context import (
assert_published,
derive_parallel_widths,
get_context,
get_device,
get_exec,
get_flags,
get_parallel,
@@ -366,6 +367,35 @@ class TestSpawnIdentities(_IsolatedOverrides):
self.assertEqual(parallel.pp_rank, 1)
self.assertEqual(parallel.dp_rank, 2)
def test_the_spawn_states_the_device_and_the_record_stays_clean(self):
"""The parent picks the device, so it arrives with the rest of the
placement. It is stamped onto the bag: the record is the startup
input and stays as the caller handed it over."""
server_args = ServerArgs(model_path="dummy")
publish(
server_args,
role="test",
ranks=SpawnRanks(world_rank=0, gpu_id=3),
)
self.assertEqual(get_device().gpu_id, 3)
# Not on the record at all. An `Arg` is the operator's input and is
# collected into `ServerArgs`; nobody types this one, so it is
# declared rather than carried, and the startup input has no field
# for the spawn to have to leave alone.
self.assertNotIn(
"gpu_id", {f.name for f in msgspec.structs.fields(type(server_args))}
)
def test_a_process_on_no_device_is_told_nothing(self):
"""Most roles run on no device at all, so the bundle leaves it out and
the bag keeps the declared default rather than inventing a zero."""
publish(
ServerArgs(model_path="dummy"),
role="test",
ranks=SpawnRanks(world_rank=0),
)
self.assertIsNone(get_device().gpu_id)
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
@@ -394,7 +424,7 @@ class TestSpawnIdentities(_IsolatedOverrides):
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
This is what the per-runner record provided by being a plain frozen object, 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.
@@ -3068,6 +3098,58 @@ class TestWhoAnswersDuringADraftScope(CustomTestCase):
self.assertEqual((info.pp_rank, info.pp_size), (0, 1))
class TestTheRecordIsNeverWrittenTo(CustomTestCase):
"""`server_args` is the startup record; the bags are the truth afterwards.
Writing a field onto it after `resolve_once()` has sealed it puts a second
answer where there is supposed to be one, and it is invisible to anything
reading the bag. The sanctioned writer is `RuntimeContext.override`, which
writes the bag and says so in its own contract. `arg_groups/` is exempt: it
is the resolution pipeline, so building the record is its job.
"""
#: Assignments here are the record being built, not mutated behind a reader.
EXEMPT = ("srt/arg_groups/",)
def test_nothing_assigns_a_field_of_the_record(self):
import ast as _ast
from sglang.srt.arg_groups.arg_utils import namespace_of
from sglang.srt.server_args import ServerArgs
fields = set(namespace_of(ServerArgs))
offenders = []
for path in _sources():
rel = path.as_posix()
if "sglang/srt/" not in rel and "sglang/benchmark/" not in rel:
continue
if any(part in rel for part in self.EXEMPT):
continue
for node in _ast.walk(_ast.parse(path.read_text(encoding="utf-8-sig"))):
targets = (
node.targets
if isinstance(node, _ast.Assign)
else [node.target]
if isinstance(node, (_ast.AugAssign, _ast.AnnAssign))
else []
)
for target in targets:
if not isinstance(target, _ast.Attribute):
continue
base = target.value
name = getattr(base, "id", getattr(base, "attr", None))
if target.attr.startswith("_"):
continue
if name == "server_args" and target.attr in fields:
offenders.append(f"{rel}:{target.lineno} .{target.attr}")
self.assertEqual(
offenders,
[],
"write the bag through get_context().override(source, ...) instead "
"-- the record is not a channel:\n " + "\n ".join(offenders),
)
class TestNothingReadsThePlacementBeforeItIsFrozen(CustomTestCase):
"""`ModelRunner.__init__` freezes its placement partway through.
@@ -20,7 +20,6 @@ from unittest.mock import patch
import torch
from torch import nn
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.quantization.fp8_utils import (
quant_weight_ue8m0,
transform_scale_ue8m0,
@@ -181,15 +180,6 @@ class _FakeModelRunner:
attn_dp_size: int | None = None,
):
self.model = model
self.ps = ParallelState.trivial(
tp_rank=tp_rank,
tp_size=tp_size,
dp_rank=dp_rank,
dp_size=dp_size,
attn_dp_size=attn_dp_size if attn_dp_size is not None else dp_size,
pp_rank=pp_rank,
pp_size=pp_size,
)
# ---------------------------------------------------------------------------