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

config: a parallel size has one spelling; a patched scope declares its own (#36621)

Browse Source 
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-08-27 12:56:42 -07:00
committed by GitHub
co-authored by Claude Opus 5
parent ca1d7ed8e6
commit fd40a331bf
62 changed files with 439 additions and 1313 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/layers/mamba/test_mamba2_mixer.py
+6 -3
View File
@@ -109,9 +109,12 @@ def mixer2_gated_norm_tensor_parallel(
import sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated as m2
# Force attn-TP rank through the context (the weight loader reads it via
# get_parallel().attn_tp_rank); avoids calling initialize_dp_attention.
with get_parallel().override(attn_tp_rank=local_rank):
# Force the TP topology through the context (the weight loader reads
# get_parallel().attn_tp_rank, Mixer2RMSNormGated reads tp_size / tp_rank);
# avoids calling initialize_dp_attention.
with get_parallel().override(
attn_tp_rank=local_rank, tp_size=world_size, tp_rank=local_rank
):
# create gated-norm with TP
mixer = m2.Mixer2RMSNormGated(
full_hidden_size=hidden_size,
test/registered/unit/layers/quantization/test_int8_linear_methods.py
+4 -1
View File
@@ -12,6 +12,7 @@ import torch
from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config
from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import get_device_sm
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.layer_ut_utils import (
@@ -112,7 +113,9 @@ class TestBlockInt8Linear(_Int8LinearCheck):
activation_scheme="dynamic",
weight_block_size=[128, 128],
)
layer = make_tp1_column_parallel_linear(quant_config, n, k)
# create_weights reads get_parallel().tp_size, not the layer's argument.
with get_parallel().override(tp_size=1, tp_rank=0):
layer = make_tp1_column_parallel_linear(quant_config, n, k)
w = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) / 10
w_int8, scale_inv, w_dequant = _quantize_int8_block(w)
load_linear_weights(layer, weight=w_int8, weight_scale_inv=scale_inv)
test/registered/unit/managers/test_scheduler_internal_state_env_vars.py
+2 -12
View File
@@ -11,6 +11,7 @@ maybe_stub_sgl_kernel()
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import GetInternalStateReq
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.runtime_context import get_context
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
@@ -42,18 +43,7 @@ class TestSchedulerInternalStateEnvVars(unittest.TestCase):
)
scheduler.draft_worker = None
with patch(
"sglang.srt.managers.scheduler.get_context",
return_value=SimpleNamespace(resolved_server_args_dict=dict),
), patch(
"sglang.srt.managers.scheduler.get_exec",
return_value=SimpleNamespace(moe=SimpleNamespace(elastic_ep_backend=None)),
), patch(
"sglang.srt.managers.scheduler.compute_world_size", return_value=1
), patch(
"sglang.srt.managers.scheduler.get_parallel",
return_value=SimpleNamespace(config=SimpleNamespace()),
):
with get_context().override_server_args():
output = scheduler.get_internal_state(recv_req=GetInternalStateReq())
return output.internal_state
test/registered/unit/managers/test_scheduler_internal_state_world_size.py
+24 -45
View File
@@ -1,6 +1,5 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import maybe_stub_sgl_kernel
@@ -9,59 +8,52 @@ maybe_stub_sgl_kernel()
from sglang.srt.managers.io_struct import GetInternalStateReq
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import compute_world_size
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _make_parallel_config(
def _shape(
*, tp_size: int, pp_size: int, dp_size: int, enable_dp_attention: bool
) -> SimpleNamespace:
) -> dict:
"""The four `parallel` leaves the world size is computed from."""
return SimpleNamespace(
tp_size=tp_size,
pp_size=pp_size,
dp_size=dp_size,
enable_dp_attention=enable_dp_attention,
)
return {
"tp_size": tp_size,
"pp_size": pp_size,
"dp_size": dp_size,
"enable_dp_attention": enable_dp_attention,
}
class TestComputeWorldSize(unittest.TestCase):
def test_a_single_gpu_server_holds_one_gpu(self):
"""The default shape has to come out as one, or every consumer is off by a factor."""
config = _make_parallel_config(
tp_size=1, pp_size=1, dp_size=1, enable_dp_attention=False
)
shape = _shape(tp_size=1, pp_size=1, dp_size=1, enable_dp_attention=False)
self.assertEqual(compute_world_size(config), 1)
self.assertEqual(compute_world_size(**shape), 1)
def test_tensor_and_pipeline_stages_multiply(self):
"""Each (pp_rank, tp_rank) pair is its own scheduler process on its own gpu."""
config = _make_parallel_config(
tp_size=2, pp_size=3, dp_size=1, enable_dp_attention=False
)
shape = _shape(tp_size=2, pp_size=3, dp_size=1, enable_dp_attention=False)
self.assertEqual(compute_world_size(config), 6)
self.assertEqual(compute_world_size(**shape), 6)
def test_plain_data_parallel_replicas_each_hold_their_own_gpus(self):
"""Without dp attention every replica launches a full tensor-parallel group of its own."""
config = _make_parallel_config(
tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False
)
shape = _shape(tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False)
self.assertEqual(compute_world_size(config), 4)
self.assertEqual(compute_world_size(**shape), 4)
def test_data_parallel_attention_shares_the_tensor_parallel_gpus(self):
"""With dp attention the dp ranks live inside the tensor-parallel world, not beside it."""
config = _make_parallel_config(
tp_size=4, pp_size=1, dp_size=2, enable_dp_attention=True
)
shape = _shape(tp_size=4, pp_size=1, dp_size=2, enable_dp_attention=True)
self.assertEqual(compute_world_size(config), 4)
self.assertEqual(compute_world_size(**shape), 4)
class TestSchedulerInternalStateWorldSize(unittest.TestCase):
def _get_internal_state(self, config: SimpleNamespace) -> dict:
def _get_internal_state(self, shape: dict) -> dict:
scheduler = Scheduler.__new__(Scheduler)
scheduler.metrics_reporter = SimpleNamespace(
last_gen_throughput=1.0,
@@ -87,40 +79,27 @@ class TestSchedulerInternalStateWorldSize(unittest.TestCase):
)
scheduler.draft_worker = None
with patch(
"sglang.srt.managers.scheduler.get_context",
return_value=SimpleNamespace(resolved_server_args_dict=dict),
), patch(
"sglang.srt.managers.scheduler.get_exec",
return_value=SimpleNamespace(moe=SimpleNamespace(elastic_ep_backend=None)),
), patch(
"sglang.srt.managers.scheduler.get_parallel",
return_value=SimpleNamespace(config=config),
):
with get_context().override_server_args(**shape):
output = scheduler.get_internal_state(recv_req=GetInternalStateReq())
return output.internal_state
def test_the_internal_state_reports_the_whole_server(self):
"""A consumer sizing an external fleet reads the gpus the server occupies, not the declared sizes."""
config = _make_parallel_config(
tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False
)
shape = _shape(tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False)
internal_state = self._get_internal_state(config)
internal_state = self._get_internal_state(shape)
self.assertEqual(internal_state["world_size"], 4)
def test_the_reported_size_is_not_one_replica_of_a_data_parallel_server(self):
"""Each plain dp replica has its own process group, so no scheduler can report the whole server from it."""
config = _make_parallel_config(
tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False
)
shape = _shape(tp_size=2, pp_size=1, dp_size=2, enable_dp_attention=False)
internal_state = self._get_internal_state(config)
internal_state = self._get_internal_state(shape)
self.assertNotEqual(
internal_state["world_size"], config.tp_size * config.pp_size
internal_state["world_size"], shape["tp_size"] * shape["pp_size"]
)
test/registered/unit/mem_cache/test_dsa_layer_split_broadcast.py
+75 -69
View File
@@ -63,79 +63,85 @@ def _run(rank: int, world: int, port: int):
LayerSplitDSATokenToKVPool,
)
cp_rank = get_parallel().attn_cp_rank
cp_size = get_parallel().attn_cp_size
assert cp_size == world
# This worker builds the groups but never publishes, so the scope declares
# the size it runs at.
with get_parallel().override(attn_cp_size=world):
cp_rank = get_parallel().attn_cp_rank
cp_size = get_parallel().attn_cp_size
assert cp_size == world
pool = LayerSplitDSATokenToKVPool(
SIZE,
page_size=PAGE_SIZE,
kv_lora_rank=KV_LORA_RANK,
dtype=torch.bfloat16,
qk_rope_head_dim=QK_ROPE,
layer_num=LAYER_NUM,
device=f"cuda:{rank}",
index_head_dim=INDEX_HEAD_DIM,
enable_memory_saver=False,
kv_cache_dim=KV_LORA_RANK + QK_ROPE,
layer_shard_rank=cp_rank,
layer_shard_size=cp_size,
)
# Owner writes a layer-distinct constant into each owned kv_buffer layer.
for layer_id in range(LAYER_NUM):
if pool._is_layer_owned(layer_id):
pool.kv_buffer[layer_id].fill_(float(layer_id + 1))
torch.cuda.synchronize()
torch.distributed.barrier()
# Every rank reads every layer; broadcast must surface the owner's value.
ok = True
for layer_id in range(LAYER_NUM):
buf = pool._get_broadcastable_kv_buffer(layer_id)
expected = float(layer_id + 1)
got = buf.float().mean().item()
if abs(got - expected) > 1e-3:
print(f"[rank {rank}] layer {layer_id}: expected {expected}, got {got}")
ok = False
assert ok, f"rank {rank} read stale/incorrect broadcast contents"
# Indexer buffer owner-broadcast: owner writes a layer-distinct value, then
# every rank must read it back for every layer.
for layer_id in range(LAYER_NUM):
store_buf = pool.get_index_k_with_scale_buffer(layer_id)
assert (
store_buf.data_ptr() == pool.index_k_with_scale_buffer[layer_id].data_ptr()
pool = LayerSplitDSATokenToKVPool(
SIZE,
page_size=PAGE_SIZE,
kv_lora_rank=KV_LORA_RANK,
dtype=torch.bfloat16,
qk_rope_head_dim=QK_ROPE,
layer_num=LAYER_NUM,
device=f"cuda:{rank}",
index_head_dim=INDEX_HEAD_DIM,
enable_memory_saver=False,
kv_cache_dim=KV_LORA_RANK + QK_ROPE,
layer_shard_rank=cp_rank,
layer_shard_size=cp_size,
)
if pool._is_layer_owned(layer_id):
store_buf.fill_(layer_id + 10)
torch.cuda.synchronize()
torch.distributed.barrier()
for layer_id in range(LAYER_NUM):
# invalidate any cached remote copy so the read forces a fresh broadcast
pool.invalidate_index_buffer_for_layer(layer_id)
buf = pool._get_broadcastable_index_buffer(layer_id)
expected = layer_id + 10
got = buf.float().mean().item()
if abs(got - expected) > 1e-3:
print(f"[rank {rank}] index layer {layer_id}: exp {expected}, got {got}")
ok = False
assert ok, f"rank {rank} read stale/incorrect index broadcast contents"
# Async prefetch path: prefetch layer, then read must return owner value.
for layer_id in range(LAYER_NUM):
pool.remote_kv_layer_id = None # force a fresh broadcast
pool.prefetch_kv_buffer(layer_id)
buf = pool._get_broadcastable_kv_buffer(layer_id)
got = buf.float().mean().item()
if abs(got - float(layer_id + 1)) > 1e-3:
print(f"[rank {rank}] prefetch layer {layer_id}: got {got}")
ok = False
assert ok, f"rank {rank} prefetch path returned incorrect contents"
# Owner writes a layer-distinct constant into each owned kv_buffer layer.
for layer_id in range(LAYER_NUM):
if pool._is_layer_owned(layer_id):
pool.kv_buffer[layer_id].fill_(float(layer_id + 1))
print(f"[rank {rank}] OK: all {LAYER_NUM} layers read correct owner contents")
torch.distributed.barrier()
torch.cuda.synchronize()
torch.distributed.barrier()
# Every rank reads every layer; broadcast must surface the owner's value.
ok = True
for layer_id in range(LAYER_NUM):
buf = pool._get_broadcastable_kv_buffer(layer_id)
expected = float(layer_id + 1)
got = buf.float().mean().item()
if abs(got - expected) > 1e-3:
print(f"[rank {rank}] layer {layer_id}: expected {expected}, got {got}")
ok = False
assert ok, f"rank {rank} read stale/incorrect broadcast contents"
# Indexer buffer owner-broadcast: owner writes a layer-distinct value, then
# every rank must read it back for every layer.
for layer_id in range(LAYER_NUM):
store_buf = pool.get_index_k_with_scale_buffer(layer_id)
assert (
store_buf.data_ptr()
== pool.index_k_with_scale_buffer[layer_id].data_ptr()
)
if pool._is_layer_owned(layer_id):
store_buf.fill_(layer_id + 10)
torch.cuda.synchronize()
torch.distributed.barrier()
for layer_id in range(LAYER_NUM):
# invalidate any cached remote copy so the read forces a fresh broadcast
pool.invalidate_index_buffer_for_layer(layer_id)
buf = pool._get_broadcastable_index_buffer(layer_id)
expected = layer_id + 10
got = buf.float().mean().item()
if abs(got - expected) > 1e-3:
print(
f"[rank {rank}] index layer {layer_id}: exp {expected}, got {got}"
)
ok = False
assert ok, f"rank {rank} read stale/incorrect index broadcast contents"
# Async prefetch path: prefetch layer, then read must return owner value.
for layer_id in range(LAYER_NUM):
pool.remote_kv_layer_id = None # force a fresh broadcast
pool.prefetch_kv_buffer(layer_id)
buf = pool._get_broadcastable_kv_buffer(layer_id)
got = buf.float().mean().item()
if abs(got - float(layer_id + 1)) > 1e-3:
print(f"[rank {rank}] prefetch layer {layer_id}: got {got}")
ok = False
assert ok, f"rank {rank} prefetch path returned incorrect contents"
print(f"[rank {rank}] OK: all {LAYER_NUM} layers read correct owner contents")
torch.distributed.barrier()
class TestLayerSplitDSABroadcast(CustomTestCase):
test/registered/unit/models/test_kimi_k25.py
+2 -3
View File
@@ -67,7 +67,6 @@ from sglang.srt.multimodal.transport.cuda_ipc import (
CudaIpcTensorTransportProxy,
)
from sglang.srt.runtime_context import (
ParallelContext,
get_context,
get_parallel,
publish,
@@ -911,7 +910,7 @@ def test_kimi_k3_normal_cache_path_connects_real_producer_to_model_consumer():
hot_items = pickle.loads(pickle.dumps(hot.mm_items))
with (
patch.object(ParallelContext, "config", SimpleNamespace(tp_size=1)),
get_parallel().override(tp_size=1),
patch(
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
return_value=(
@@ -975,7 +974,7 @@ def test_kimi_k3_model_accepts_mixed_cached_eager_and_deferred_artifacts():
)
with (
patch.object(ParallelContext, "config", SimpleNamespace(tp_size=1)),
get_parallel().override(tp_size=1),
patch(
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
return_value=(torch.full((1, 3), 2.0), torch.tensor([[1, 1, 1]])),
test/registered/unit/models/test_zaya_cca.py
+11 -4
View File
@@ -29,6 +29,7 @@ from typing import List, Optional
import torch
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.layer_ut_utils import init_single_process_dist
from sglang.test.test_utils import CustomTestCase
@@ -36,10 +37,16 @@ from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=30, suite="base-a-test-cpu")
def _ensure_dist_initialized() -> None:
def _ensure_dist_initialized(cls) -> None:
"""CCA reads the TP rank / world size inside ``__init__`` to size its
head-parallel projections, so the groups must exist before construction."""
head-parallel projections. The rank is the live group's, so the groups must
exist before construction; the size answers from the published ``parallel``
bag, so the case has to publish a context as well.
"""
init_single_process_dist()
override = get_context().override_server_args(tp_size=1)
override.install()
cls.addClassCleanup(override.restore)
@dataclass(frozen=True)
@@ -243,7 +250,7 @@ def _make_tiny_cca(
class TestZayaCCA(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
_ensure_dist_initialized()
_ensure_dist_initialized(cls)
def test_single_chunk_matches_reference(self):
"""A single-chunk extend with empty prefix matches the no-state path."""
@@ -540,7 +547,7 @@ class TestZayaCCATensorParallel(CustomTestCase):
@classmethod
def setUpClass(cls) -> None:
_ensure_dist_initialized()
_ensure_dist_initialized(cls)
def _slice_full_state_dict_into_rank(self, ref_cca, tp_cca, tp_rank: int):
"""Copy the reference's full weights into the per-rank CCA, using the
test/registered/unit/server_args/test_resolution_declarations.py
+5 -17
View File
@@ -433,19 +433,11 @@ class TestResolutionDeclarations(CustomTestCase):
mapping = namespace_of(ServerArgs)
self.assertGreater(len(mapping), 400, "the namespace mapping collapsed")
# The five sizes keep a live property shadowing the bare name; the
# comparison below reaches them anyway, through `get_parallel().config`.
self.assertGreaterEqual(
_live_topology_leaves()
& {
"tp_size",
"pp_size",
"moe_dp_size",
"attn_cp_size",
"dcp_size",
},
{"tp_size", "pp_size", "moe_dp_size", "attn_cp_size", "dcp_size"},
"a parallel size stopped being served from the live topology",
self.assertEqual(
set(),
_live_topology_leaves() & set(mapping),
"a parallel leaf gained a live member of the same name, so the "
"comparison below reads the group rather than the published leaf",
)
compared = 0
@@ -461,10 +453,6 @@ class TestResolutionDeclarations(CustomTestCase):
unreachable.append(f"no get_{groups[0]}() for {path}.{field}")
continue
node = accessor()
if groups[0] == "parallel":
# Bare names there are the live topology; the published
# leaves are one hop down, so the reader takes that hop.
node = node.config
try:
for group in groups[1:]:
node = getattr(node, group)
test/registered/unit/server_args/test_resolution_reads_the_declarations.py
+4 -4
View File
@@ -420,7 +420,7 @@ def reads_a_leaf_through_the_alias(runner):
def hands_the_accessor_to_a_helper():
return compute_world_size(get_server_args())
return attention_backends_of(get_server_args())
def reads_the_view(runner):
@@ -580,7 +580,7 @@ class TestResolutionReadsTheDeclarations(CustomTestCase):
"""
helpers = _config_reading_helpers()
decided = _declared_fields()
for name in ("m3_fp8_attn_gemm_enabled", "compute_world_size"):
for name in ("m3_fp8_attn_gemm_enabled", "attention_backends_of"):
self.assertIn(name, helpers, f"the helper derivation lost {name}")
for field in ("speculative_num_draft_tokens", "attention_backend"):
self.assertIn(field, decided, f"the declared set lost {field}")
@@ -598,8 +598,8 @@ class TestResolutionReadsTheDeclarations(CustomTestCase):
"sa_local.attention_backend",
"self._server_args.attention_backend",
"engine_args.attention_backend",
"compute_world_size(get_server_args())"
" reads " + ", ".join(helpers["compute_world_size"]),
"attention_backends_of(get_server_args())"
" reads " + ", ".join(helpers["attention_backends_of"]),
},
"the scan lost a spelling, or started flagging a legal one:\n "
+ "\n ".join(sorted(flagged)),
test/registered/unit/test_global_config_read_ratchet.py
-463
View File
@@ -13,9 +13,6 @@ slot.
The reads that remain live in ``runtime_context.py`` (exempt by module): the
``@property`` / method members computed from several fields plus the HF config,
which are not namespace leaves and have no home but ``ServerArgs``.
Separately, ``_CONFIGURED_SIZE_CALL_SITES`` registers every business read of
``get_parallel().config.<size>`` — the config tier of a size whose bare name is
the live topology — with the reason the live property cannot serve it.
What the scan sees: ``get_server_args().field``, an alias (``sa =
get_server_args()`` then ``sa.field`` -- function-local, module-level, or parked
@@ -48,250 +45,6 @@ _PACKAGE_ROOT = Path(next(iter(sglang.__path__)))
# resolution pipeline.
_SLOT_OWNERS = ("srt/runtime_context.py", "srt/server_args.py", "srt/arg_groups/")
# Every configured read of a live-shadowed size (``get_parallel().config.pp_size``
# and its four siblings), with the reason the live topology cannot answer there.
# The test below asserts this map is exactly the set of such reads, so the
# reasons cannot drift away from the code.
_CONFIGURED_SIZE_CALL_SITES = {
("srt/layers/cp/base.py", "attn_cp_size"): (
"the lazy strategy bind in a worker: the CP group is what the strategy "
"is being built for, and the configured width is what describes it"
),
("benchmark/one_batch.py", "pp_size"): (
"CPU affinity for this rank, computed right after the work function "
"publishes and before dist init, so the groups do not exist yet"
),
("benchmark/one_batch.py", "tp_size"): (
"the same affinity computation: the layout is the configured one, and "
"the live group is not up at this point in the work function"
),
("srt/entrypoints/engine.py", "pp_size"): (
"the launch path decides how many scheduler processes to spawn; it runs "
"before any of them exists, so there is no group to ask"
),
("srt/entrypoints/engine.py", "attn_cp_size"): (
"the launcher's per-TP-rank layout, computed while deciding what to "
"spawn -- the groups it is laying out do not exist yet"
),
("srt/entrypoints/engine.py", "moe_dp_size"): (
"the MoE factor of that same pre-spawn layout"
),
("srt/ray/engine.py", "pp_size"): (
"the Ray driver sizes the actor placement group; the actors it is about "
"to create are the ones that will hold the process groups"
),
("srt/ray/engine.py", "tp_size"): (
"the same placement arithmetic as the stage count: the driver sizes "
"the actors that will hold the process groups"
),
("srt/ray/data_parallel_controller.py", "tp_size"): (
"the same arithmetic on the DP path, also in the driver"
),
("srt/ray/data_parallel_controller.py", "pp_size"): (
"same placement arithmetic on the DP path -- ranks per TP group, "
"computed in the driver before the actors start"
),
("srt/ray/data_parallel_controller.py", "attn_cp_size"): (
"the attention-CP factor of that same placement arithmetic, and the one "
"size whose live value cannot express the configured intent when "
"attn_cp_size > moe_dp_size aliases the groups"
),
("srt/layers/attention/dsa/dsa_indexer.py", "pp_size"): (
"gates `pp_size > 1 and not get_pp_group()...`; the short circuit is the "
"point, since with PP off the group is never touched, which is what lets "
"the Indexer be constructed before distributed init"
),
("srt/managers/scheduler.py", "pp_size"): (
"dispatch_event_loop picks the PP event loop; the MLX runner stub never "
"initializes torch.distributed, so the live property asserts before the "
"MLX loop can start -- the configured leaf answers the same value "
"wherever the live groups exist"
),
("srt/mem_cache/kv_cache_configurator.py", "pp_size"): (
"decides whether the token capacity needs a cross-PP all-reduce at all; "
"asking the configured size keeps that decision independent of whether a "
"PP group is installed in this process"
),
("srt/layers/dp_attention.py", "attn_cp_size"): (
"compared against the configured moe_dp_size below"
),
("srt/layers/dp_attention.py", "moe_dp_size"): (
"the configuration this predicate detects (attn_cp_size > moe_dp_size) is "
"the one where initialize_model_parallel aliases _MOE_DP to _ATTN_CP, so "
"the live sizes are equal there and a live comparison is always false"
),
("srt/managers/scheduler.py", "tp_size"): (
"configure_scheduler_process runs before the scheduler's own process "
"groups exist -- configuring the process is what it is for -- so there "
"is nothing live to ask yet"
),
("srt/managers/scheduler.py", "moe_dp_size"): (
"same pre-distributed-init arithmetic in configure_scheduler_process"
),
("srt/managers/scheduler.py", "attn_cp_size"): (
"same pre-distributed-init arithmetic in configure_scheduler_process"
),
("srt/managers/scheduler.py", "dcp_size"): (
"same pre-distributed-init arithmetic in configure_scheduler_process"
),
("srt/model_executor/runner/base_runner.py", "tp_size"): (
"the same window as the stage count next to it: a draft runner shares "
"the target's groups, so the live property would answer for the wrong "
"runner"
),
("srt/model_executor/cpu_graph_runner.py", "tp_size"): (
"the same window, on the CPU graph path"
),
("srt/entrypoints/v1_loads.py", "tp_size"): (
"the accelerator count is arithmetic over the launch shape, reported "
"from the tokenizer process, which holds no model groups"
),
("srt/disaggregation/nixl/conn.py", "tp_size"): (
"the NIXL rank arithmetic runs on the transfer path, which the CPU-only "
"conn tests exercise without starting torch.distributed"
),
("srt/managers/tokenizer_control_mixin.py", "tp_size"): (
"the tokenizer divides its worker count by the launch width; it holds "
"no model groups"
),
("srt/model_executor/runner/base_runner.py", "pp_size"): (
"the runner's layer window is arithmetic over the configured stage "
"count; a draft runner shares the target's groups, so the live "
"property would answer for the wrong runner"
),
("srt/model_executor/cpu_graph_runner.py", "pp_size"): (
"the same window, on the CPU graph path"
),
(
"srt/managers/scheduler_components/metrics_reporter.py",
"pp_size",
): (
"the reporter labels its metrics with the stage count it was launched "
"with, which is configuration; the live group answers per process"
),
("srt/speculative/eagle_draft_cuda_graph_runner.py", "pp_size"): (
"the draft runner's window over the target's stages: its own groups are "
"the target's, so the configured count is the one that describes it"
),
(
"srt/speculative/eagle_draft_extend_cuda_graph_runner.py",
"pp_size",
): ("the same draft window, on the extend path"),
(
"srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py",
"pp_size",
): ("the same draft window, multi-layer extend"),
("srt/speculative/frozen_kv_mtp_cuda_graph_runner.py", "pp_size"): (
"the same draft window, frozen-KV MTP"
),
("srt/managers/data_parallel_controller.py", "pp_size"): (
"the controller lays out its schedulers' ranks before spawning them, so "
"the groups it is sizing for do not exist yet"
),
("srt/managers/data_parallel_controller.py", "attn_cp_size"): (
"the same pre-spawn rank arithmetic"
),
("srt/managers/data_parallel_controller.py", "moe_dp_size"): (
"the same pre-spawn rank arithmetic"
),
("srt/entrypoints/v1_loads.py", "pp_size"): (
"the /v1/loads accelerator count is arithmetic over the launch shape, "
"reported from the tokenizer process, which holds no model groups"
),
("srt/disaggregation/common/conn.py", "pp_size"): (
"the bootstrap connection is built by the KV manager on the transfer "
"path, which the CPU-only conn tests exercise without ever starting "
"torch.distributed"
),
("srt/elastic_ep/elastic_ep.py", "tp_size"): (
"the joiner's rank window is computed against the size the process was "
"configured with, not the size of the group it is about to join"
),
("srt/elastic_ep/expert_backup_manager.py", "tp_size"): (
"the backup server counts the clients it expects to report in, which "
"is how many the launch configured -- the live group is what they are "
"still joining"
),
(
"srt/model_executor/model_runner_components/startup_weight_load.py",
"tp_size",
): (
"the load options are assembled in ModelRunner.__init__ for a runner "
"that may be a draft, whose groups are the target's; the configured "
"sizes are what the record answered before"
),
(
"srt/model_executor/model_runner_components/startup_weight_load.py",
"pp_size",
): ("same options object, same reason"),
(
"srt/model_executor/model_runner_components/startup_weight_load.py",
"attn_cp_size",
): ("same options object, same reason"),
(
"srt/model_executor/model_runner_components/startup_weight_load.py",
"dcp_size",
): ("same options object, same reason"),
(
"srt/model_executor/model_runner_components/spec_aux_hidden_state.py",
"tp_size",
): (
"the draft KV bytes/token estimate sizes the memory pool before the "
"draft runner exists, so its shard count is configuration"
),
("srt/eplb/expert_location.py", "tp_size"): (
"the elastic-EP joiner window, used to size the expert layout: the "
"size the process was configured with, not the group it is joining"
),
("srt/utils/cuda_vmm_transport_utils.py", "tp_size"): (
"the consumer count is configured fan-out arithmetic (tp_size // "
"dp_size), which is what the record answered before"
),
("srt/disaggregation/encoder/runtime.py", "tp_size"): (
"the encode server's launch entry sizes its workers before it has "
"spawned any of them"
),
("srt/disaggregation/encoder/grpc_server.py", "tp_size"): (
"the same worker-count arithmetic on the gRPC entry: it spawns the TP "
"workers, so their groups do not exist yet"
),
("srt/disaggregation/encoder/server.py", "tp_size"): (
"`MMEncoder` builds its own TP group from this size -- "
"`initialize_model_parallel` is the call being handed it, so there is "
"nothing live to ask"
),
("srt/disaggregation/encoder/receiver.py", "tp_size"): (
"the receiver labels and shards by the launch width; it runs in the "
"tokenizer process, which holds no encoder groups"
),
("srt/managers/rust_server.py", "tp_size"): (
"the rust server decides its transport from the launch width, in the "
"tokenizer process, which holds no model groups"
),
("compile_deep_gemm.py", "tp_size"): (
"the warm-up request fans bootstrap rooms across the launch's ranks; it "
"runs in the tokenizer process, which holds no model groups"
),
("srt/utils/common.py", "tp_size"): (
"the require_*_tp_gather predicates compared the configured tp_size "
"when they read the record; the live property answers a different "
"question wherever the groups alias, so the configured accessor is the "
"mechanical substitution and the live one would be a semantic change"
),
("srt/model_loader/loader.py", "moe_dp_size"): (
"the same dict already carries the live moe_dp_size under 'dp'; this entry "
"is the configured intent"
),
("srt/models/kimi_k25.py", "tp_size"): (
"the IPC refcount must match the configured TP consumer count captured "
"when the tokenizer creates MmItemMemoryPool; a live attention subgroup "
"size could strand leases in the bounded pool"
),
("srt/models/kimi_k3.py", "tp_size"): (
"same as kimi_k25: the IPC refcount must agree with the recycler's waiter"
),
}
_DIRECT_BASELINE = 0
_ALIAS_BASELINE = 0
@@ -594,222 +347,6 @@ class TestGlobalConfigReadRatchet(CustomTestCase):
self._check("alias-form", alias, _ALIAS_BASELINE)
def _live_shadowed_sizes() -> frozenset:
"""Names that are BOTH a live ``ParallelContext`` property and a ``parallel``
config leaf.
Derived from the two sides themselves: a size that gains a live property, or
a live property that gains a leaf, joins the registry's subject set without a
list here.
"""
from sglang.srt.arg_groups.arg_utils import namespace_of
from sglang.srt.runtime_context import ParallelContext
from sglang.srt.server_args import ServerArgs
live = {
name
for name, value in vars(ParallelContext).items()
if isinstance(value, property)
}
leaves = {
field for field, path in namespace_of(ServerArgs).items() if path == "parallel"
}
shadowed = frozenset(live & leaves)
assert shadowed, "no live-shadowed size found; the derivation is broken"
return shadowed
def _parallel_config_reads(tree, subjects):
"""Names in ``subjects`` read through the parallel bag's ``config`` hop.
Sees ``get_parallel().config.pp_size``, the module-qualified spelling, a
local bound to either hop (``p = get_parallel()`` / ``cfg = p.config``), and
the ``getattr`` form of each.
"""
fns, modules = set(), set()
for node in ast.walk(tree):
if (
isinstance(node, ast.ImportFrom)
and node.module
and node.module.endswith("runtime_context")
):
fns |= {a.asname or a.name for a in node.names if a.name == "get_parallel"}
elif isinstance(node, ast.ImportFrom) and node.module:
# `from sglang.srt import runtime_context as rc` binds the module.
for a in node.names:
if f"{node.module}.{a.name}".endswith("runtime_context"):
modules.add(a.asname or a.name)
elif isinstance(node, ast.Import):
for a in node.names:
if a.name.endswith("runtime_context"):
# Unaliased, the call site spells the whole dotted path.
modules.add(a.asname or a.name)
def dotted(node):
parts = []
while isinstance(node, ast.Attribute):
parts.append(node.attr)
node = node.value
if not isinstance(node, ast.Name):
return None
parts.append(node.id)
return ".".join(reversed(parts))
def is_bag_call(node):
if not isinstance(node, ast.Call):
return False
func = node.func
if isinstance(func, ast.Name):
return func.id in fns
return (
isinstance(func, ast.Attribute)
and func.attr == "get_parallel"
and dotted(func.value) in modules
)
bag_aliases, config_aliases = set(), set()
for _ in range(2): # a local copy of a local is still the same object
for node in ast.walk(tree):
if not isinstance(node, ast.Assign):
continue
value = node.value
if is_bag_call(value) or (
isinstance(value, ast.Name) and value.id in bag_aliases
):
bucket = bag_aliases
elif (
isinstance(value, ast.Attribute)
and value.attr == "config"
and (
is_bag_call(value.value)
or (
isinstance(value.value, ast.Name)
and value.value.id in bag_aliases
)
)
) or (isinstance(value, ast.Name) and value.id in config_aliases):
bucket = config_aliases
else:
continue
bucket |= {t.id for t in node.targets if isinstance(t, ast.Name)}
def is_config_hop(node):
return (
isinstance(node, ast.Attribute)
and node.attr == "config"
and (
is_bag_call(node.value)
or (isinstance(node.value, ast.Name) and node.value.id in bag_aliases)
)
) or (isinstance(node, ast.Name) and node.id in config_aliases)
found = set()
for node in ast.walk(tree):
if isinstance(node, ast.Attribute) and node.attr in subjects:
base, name = node.value, node.attr
elif (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "getattr"
and len(node.args) >= 2
and isinstance(node.args[1], ast.Constant)
and node.args[1].value in subjects
):
base, name = node.args[0], node.args[1].value
else:
continue
if is_config_hop(base):
found.add(name)
return found
_READ_SPELLINGS = (
"from sglang.srt.runtime_context import get_parallel\nx = get_parallel().config.tp_size",
"from sglang.srt.runtime_context import get_parallel as gp\nx = gp().config.tp_size",
"from sglang.srt import runtime_context as rc\nx = rc.get_parallel().config.tp_size",
"import sglang.srt.runtime_context\nx = sglang.srt.runtime_context.get_parallel().config.tp_size",
"from sglang.srt.runtime_context import get_parallel\np = get_parallel()\nx = p.config.tp_size",
"from sglang.srt.runtime_context import get_parallel\nc = get_parallel().config\nx = c.tp_size",
'from sglang.srt.runtime_context import get_parallel\nx = getattr(get_parallel().config, "tp_size")',
)
class TestParallelConfigReadSpellings(CustomTestCase):
"""``_parallel_config_reads`` resolves every spelling it claims to.
The scan below decides what the documented set is compared against, so a
spelling it cannot resolve does not fail anything -- it drops the read.
"""
def test_every_documented_spelling_resolves(self):
for source in _READ_SPELLINGS:
with self.subTest(source=source):
found = _parallel_config_reads(ast.parse(source), {"tp_size"})
self.assertEqual({"tp_size"}, set(found))
def test_the_live_property_is_not_a_config_read(self):
source = (
"from sglang.srt.runtime_context import get_parallel\n"
"x = get_parallel().tp_size"
)
self.assertEqual(
set(), set(_parallel_config_reads(ast.parse(source), {"tp_size"}))
)
class TestConfiguredSizeCallSites(CustomTestCase):
"""The configured-vs-live exceptions are enumerated, with reasons.
``get_parallel().config.tp_size`` answers what the process was configured
with where the bare ``get_parallel().tp_size`` answers what the process ended
up with. Each site that needs the former is listed above with why the live
property cannot serve it, and this case fails if the code and that list
disagree.
The unit is **(file, size)**, not the individual read: a second
``.config.pp_size`` in a file already registered for it collapses into the
same entry, so the reason has to cover the file's use of that size rather
than one line. A new file, or a new size in a listed file, is what this
catches -- through any spelling of the hop.
"""
def test_the_call_sites_match_the_documented_set(self):
subjects = _live_shadowed_sizes()
found = set()
scanned = 0
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
if rel.startswith(_SLOT_OWNERS):
continue
source = path.read_text()
# Every spelling `_parallel_config_reads` resolves -- the direct
# call, an aliased import, a module-qualified call, a local bound to
# either hop -- needs the name in the source, so skipping the rest is
# free. Filtering on anything narrower silently empties the scan.
if "get_parallel" not in source:
continue
scanned += 1
try:
tree = ast.parse(source)
except SyntaxError:
continue
found |= {(rel, name) for name in _parallel_config_reads(tree, subjects)}
self.assertGreater(
scanned,
50,
f"the pre-filter left only {scanned} files to scan; the derivation "
"is broken, not the tree",
)
documented = set(_CONFIGURED_SIZE_CALL_SITES)
self.assertEqual(
documented,
found,
"configured-size reads drifted from their documented reasons.\n"
f" undocumented: {sorted(found - documented)}\n"
f" stale entries: {sorted(documented - found)}",
)
class TestNoRenamedAccessorImports(CustomTestCase):
"""The baseline scanner matches ``get_server_args`` by its literal name, so
an ``import ... as`` rename would walk a read straight past the zero
test/registered/unit/test_launch_path_reads_configured_sizes.py
-419
View File
@@ -1,419 +0,0 @@
"""Launch paths read the configured parallel sizes, not the live ones.
`get_parallel().pp_size` and its four siblings are read-through properties over
the process groups, so they answer only after distributed init. The launcher
decides how many processes to spawn *before* that, and a live read there raises
`Distributed environment is not initialized` -- a startup crash no unit test
reaches, because nothing short of booting a server runs the launcher. The
configured answer is one hop away on the same object,
`get_parallel().config.pp_size`, which reads the published `parallel` bag.
"""
import ast
import pathlib
import unittest
import sglang
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=9, suite="base-a-test-cpu")
_PACKAGE_ROOT = pathlib.Path(sglang.__file__).resolve().parent
def _live_shadowed() -> dict:
"""{name: remedy} for every name that is BOTH a live ParallelContext
property and a `parallel` config leaf.
Derived from the two sides themselves, so a new size that gains a live
property (or a live property that gains a leaf) is watched without a second
list here. ParallelContext shadows more properties than these -- every
`_v(name, ...)` one raises the same "Distributed environment is not
initialized" -- but only a shadowed name has a configured answer to point a
launcher at.
"""
from sglang.srt.arg_groups.arg_utils import namespace_of
from sglang.srt.runtime_context import ParallelContext
from sglang.srt.server_args import ServerArgs
live = {
name
for name, value in vars(ParallelContext).items()
if isinstance(value, property)
}
leaves = {
field for field, path in namespace_of(ServerArgs).items() if path == "parallel"
}
shadowed = live & leaves
assert shadowed, (
"no live-shadowed parallel size found; the derivation is broken, not "
"the tree"
)
return {name: f"get_parallel().config.{name}" for name in sorted(shadowed)}
_LIVE_SHADOWED = _live_shadowed()
# Launch paths that decide how many children to spawn are derived below
# from the spawn itself. These launch without a size-driven spawn, so no
# derivation reaches them and they are carried by hand.
_HAND_CARRIED = (
"srt/entrypoints/http_server.py",
"srt/entrypoints/sidecar.py",
"srt/ray/data_parallel_controller.py",
"srt/ray/engine.py",
"srt/ray/http_server.py",
)
def _multiprocessing_names(tree):
"""Names bound to multiprocessing, to one of its start contexts, or to the
process constructors themselves."""
modules, constructors = set(), set()
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for a in node.names:
if a.name == "multiprocessing" or a.name.startswith("multiprocessing."):
modules.add(a.asname or a.name.split(".")[0])
elif a.name == "torch.multiprocessing":
modules.add(a.asname or "torch")
elif isinstance(node, ast.ImportFrom):
if node.module in (
"multiprocessing",
"multiprocessing.context",
"torch.multiprocessing",
):
constructors |= {
a.asname or a.name for a in node.names if a.name == "Process"
}
elif node.module == "concurrent.futures":
constructors |= {
a.asname or a.name
for a in node.names
if a.name == "ProcessPoolExecutor"
}
for node in ast.walk(tree):
if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
func = node.value.func
if (
isinstance(func, ast.Attribute)
and func.attr == "get_context"
and isinstance(func.value, ast.Name)
and func.value.id in modules
):
modules |= {t.id for t in node.targets if isinstance(t, ast.Name)}
return modules, constructors
def _spawns_from_a_size(tree) -> bool:
"""Does a function here spawn a child *and* read a live-shadowed size?
Both tiers count: deriving on the live read alone drops a launcher from the
scan the moment it is converted, so the guard would only watch the ones
that already fail it.
"""
modules, constructors = _multiprocessing_names(tree)
for fn in ast.walk(tree):
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
continue
spawns = False
for node in ast.walk(fn):
if not isinstance(node, ast.Call):
continue
func = node.func
if isinstance(func, ast.Attribute) and func.attr in (
"Process",
"ProcessPoolExecutor",
"Popen",
"spawn",
):
# `mp.Process(`, `mp.get_context("spawn").Process(` and
# `subprocess.Popen(` all reach a child process; the receiver of
# a chained call is itself a call, so this cannot require a bare
# Name.
spawns = True
elif isinstance(func, ast.Name) and func.id in constructors:
spawns = True
if not spawns:
continue
# A record read (`server_args.tp_size`) sizes a spawn too, but it cannot
# raise pre-dist; only a bag read is this guard's subject.
live, configured = _shadowed_size_reads(tree, scope=fn)
if live or configured:
return True
return False
def _parallel_bag_names(tree):
"""What this module calls `get_parallel`, plus any runtime_context alias.
A literal-name match reads only one spelling; an aliased import or a
module-qualified call is the same read with a different surface.
"""
names, modules = set(), set()
for node in ast.walk(tree):
if (
isinstance(node, ast.ImportFrom)
and node.module
and node.module.endswith("runtime_context")
):
names |= {
a.asname or a.name for a in node.names if a.name == "get_parallel"
}
elif isinstance(node, ast.ImportFrom) and node.module:
# `from sglang.srt import runtime_context as rc` binds the module,
# so `rc.get_parallel()` is the same call under another spelling.
for a in node.names:
if f"{node.module}.{a.name}".endswith("runtime_context"):
modules.add(a.asname or a.name)
elif isinstance(node, ast.Import):
for a in node.names:
if a.name.endswith("runtime_context"):
modules.add(a.asname or a.name.split(".")[0])
return names, modules
def _is_parallel_bag_call(node, names, modules) -> bool:
if not isinstance(node, ast.Call):
return False
if isinstance(node.func, ast.Name):
return node.func.id in names
return (
isinstance(node.func, ast.Attribute)
and node.func.attr == "get_parallel"
and isinstance(node.func.value, ast.Name)
and node.func.value.id in modules
)
def _bag_aliases(tree, names, qualified):
"""Locals bound to either tier: `p = get_parallel()` then `p.pp_size` is the
same live read one line later, and `cfg = get_parallel().config` then
`cfg.pp_size` is the same configured read."""
live, config = set(), set()
for node in ast.walk(tree):
if not isinstance(node, ast.Assign):
continue
value = node.value
if _is_parallel_bag_call(value, names, qualified):
bucket = live
elif (
isinstance(value, ast.Attribute)
and value.attr == "config"
and _is_parallel_bag_call(value.value, names, qualified)
):
bucket = config
else:
continue
bucket |= {t.id for t in node.targets if isinstance(t, ast.Name)}
return live, config
def _shadowed_size_reads(module_tree, scope=None):
"""(live, configured) reads of a live-shadowed size in `scope`.
`<parallel bag>.tp_size` is the live group; `<parallel bag>.config.tp_size`
is the published leaf. Both spellings are reported so a caller can tell a
launcher that reads the topology at all from one that reads it live.
What binds the bag -- the import, a module-level alias -- lives at module
scope, so those names always come from `module_tree` even when only one
function is being walked. Deriving them from the function alone finds no
import, reports no reads, and quietly answers "this launcher reads nothing".
"""
names, qualified = _parallel_bag_names(module_tree)
live_aliases, config_aliases = _bag_aliases(module_tree, names, qualified)
def is_live_bag(node):
return _is_parallel_bag_call(node, names, qualified) or (
isinstance(node, ast.Name) and node.id in live_aliases
)
def is_config_bag(node):
return (
isinstance(node, ast.Attribute)
and node.attr == "config"
and is_live_bag(node.value)
) or (isinstance(node, ast.Name) and node.id in config_aliases)
live, configured = [], []
for node in ast.walk(scope if scope is not None else module_tree):
if isinstance(node, ast.Attribute) and node.attr in _LIVE_SHADOWED:
base, name, spelling = node.value, node.attr, "attribute"
elif (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "getattr"
and len(node.args) >= 2
and isinstance(node.args[1], ast.Constant)
and node.args[1].value in _LIVE_SHADOWED
):
base, name, spelling = node.args[0], node.args[1].value, "getattr"
else:
continue
if is_config_bag(base):
configured.append((node.lineno, name, spelling))
elif is_live_bag(base):
live.append((node.lineno, name, spelling))
return live, configured
def _launch_paths():
"""(relative path, tree) per module that runs before its process groups.
A module that sizes a spawn loop from a parallel-bag size is derived from
the spawn itself; `_HAND_CARRIED` holds the launch entries that spawn
nothing, which no derivation can reach.
"""
seen = {}
sizes = frozenset(_LIVE_SHADOWED)
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
source = path.read_text()
# Every spawn shape below names Process, ProcessPoolExecutor or Popen.
if not any(name in source for name in ("Process", "Popen", "spawn")):
continue
if not any(name in source for name in sizes):
continue
try:
tree = ast.parse(source)
except SyntaxError:
continue
if _spawns_from_a_size(tree):
seen[str(path.relative_to(_PACKAGE_ROOT))] = tree
for rel in _HAND_CARRIED:
seen.setdefault(rel, ast.parse((_PACKAGE_ROOT / rel).read_text()))
return sorted(seen.items())
class TestLaunchPathsReadConfiguredSizes(CustomTestCase):
def test_configured_sizes_hold_when_the_live_topology_disagrees(self):
"""The other direction: groups exist and answer something else.
The check above proves nobody reads a live size too early. It says
nothing about what `.config.<size>` returns once the groups *are* up and
answering a different number -- which is not hypothetical: elastic EP
scales the live topology away from what the operator configured, and
that divergence is the entire reason the two tiers are separate. With
only the early-read direction covered, a `config` hop that quietly
delegated to the live property would look correct.
"""
import json
import os
import tempfile
from unittest.mock import patch
from sglang.srt.runtime_context import (
ParallelContext,
get_parallel,
publish,
reset_context,
)
from sglang.srt.server_args import ServerArgs
directory = tempfile.mkdtemp(prefix="configured_sizes_")
with open(os.path.join(directory, "config.json"), "w") as handle:
json.dump(
{
"architectures": ["LlamaForCausalLM"],
"model_type": "llama",
"hidden_size": 16,
"intermediate_size": 32,
"num_attention_heads": 2,
"num_key_value_heads": 2,
"num_hidden_layers": 2,
"vocab_size": 128,
"max_position_embeddings": 2048,
},
handle,
)
# No resolve_once() here: `tp_size` is raw input, so the configured
# value is 2 either way.
server_args = ServerArgs(model_path=directory, device="cuda", tp_size=2)
self.addCleanup(reset_context)
publish(server_args, role="scheduler")
# The live getter behind each property, read out of ParallelContext
# rather than listed here.
context_source = ast.parse(
(_PACKAGE_ROOT / "srt" / "runtime_context.py").read_text(
encoding="utf-8-sig"
)
)
parallel_class = next(
node
for node in ast.walk(context_source)
if isinstance(node, ast.ClassDef) and node.name == "ParallelContext"
)
live_getter = {}
for method in parallel_class.body:
if not isinstance(method, ast.FunctionDef):
continue
for call in ast.walk(method):
if not (
isinstance(call, ast.Call)
and isinstance(call.func, ast.Attribute)
and call.func.attr == "_v"
and call.args
and isinstance(call.args[0], ast.Constant)
):
continue
getter = call.args[1]
if isinstance(getter, ast.Attribute):
live_getter[call.args[0].value] = getter.attr
state = "sglang.srt.distributed.parallel_state"
missing = sorted(set(_LIVE_SHADOWED) - set(live_getter))
self.assertEqual(
missing,
[],
f"these sizes no longer have a live property to diverge from: {missing}",
)
for name in sorted(_LIVE_SHADOWED):
with self.subTest(size=name):
target = f"{state}.{live_getter[name]}"
configured = getattr(get_parallel().config, name)
with patch(target, return_value=configured + 41):
self.assertEqual(
get_parallel().__getattribute__(name),
configured + 41,
f"{name} no longer follows the live topology",
)
self.assertEqual(
getattr(get_parallel().config, name),
configured,
f"get_parallel().config.{name} followed the live topology "
"instead of the published configuration",
)
from sglang.srt.arg_groups.overrides import resolution_result
self.assertEqual(
resolution_result(server_args, "nccl_port"),
getattr(get_parallel(), "nccl_port"),
"a config-only leaf read bare disagreed with what resolution decided",
)
reset_context()
with self.assertRaisesRegex(ValueError, r"'parallel' not published"):
getattr(ParallelContext(), "nccl_port")
with self.assertRaisesRegex(AttributeError, r"has no 'not_a_leaf'"):
getattr(ParallelContext(), "not_a_leaf")
def test_no_live_topology_read_before_distributed_init(self):
offenders = []
for rel, tree in _launch_paths():
live, _ = _shadowed_size_reads(tree)
for lineno, name, spelling in live:
through = " through getattr" if spelling == "getattr" else ""
offenders.append(
f"{rel}:{lineno} reads the live {name}{through}; "
f"use {_LIVE_SHADOWED[name]}"
)
self.assertEqual(
offenders,
[],
"launch paths run before distributed init:\n " + "\n ".join(offenders),
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_ray_driver_reads_the_bags.py
+2 -3
View File
@@ -61,7 +61,7 @@ class TestRayDriverReadsTheBags(CustomTestCase):
self._publish(tp_size=2, pp_size=1, dp_size=1, enable_dp_attention=False)
self.assertEqual(_compute_world_size(), 2)
get_context().override("test.ray_driver", tp_size=8)
self.assertEqual(get_parallel().config.tp_size, 8)
self.assertEqual(get_parallel().tp_size, 8)
self.assertEqual(_compute_world_size(), 8)
def test_the_driver_modules_read_no_field_off_a_record(self):
@@ -112,8 +112,7 @@ class TestRayDriverReadsTheBags(CustomTestCase):
offenders,
[],
"the Ray driver reads a config field off a record; the driver runs "
"after the publish, so read `get_parallel().config`:\n "
+ "\n ".join(offenders),
"after the publish, so read `get_parallel()`:\n " + "\n ".join(offenders),
)
test/registered/unit/test_runtime_context.py
+25 -8
View File
@@ -39,21 +39,16 @@ _DP = "sglang.srt.layers.dp_attention"
SIZE_RANK_DELEGATIONS = [
("world_size", f"{_PS}.get_world_size"),
("world_rank", f"{_PS}.get_world_rank"),
("tp_size", f"{_PS}.get_tensor_model_parallel_world_size"),
("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"),
("dcp_size", f"{_PS}.get_dcp_world_size"),
("dcp_rank", f"{_PS}.get_dcp_rank"),
("pp_size", f"{_PS}.get_pipeline_model_parallel_world_size"),
("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"),
("moe_ep_size", f"{_PS}.get_moe_expert_parallel_world_size"),
("moe_ep_rank", f"{_PS}.get_moe_expert_parallel_rank"),
("moe_dp_size", f"{_PS}.get_moe_data_parallel_world_size"),
("moe_dp_rank", f"{_PS}.get_moe_data_parallel_rank"),
("moe_tp_size", f"{_PS}.get_moe_tensor_parallel_world_size"),
("moe_tp_rank", f"{_PS}.get_moe_tensor_parallel_rank"),
("attn_tp_size", f"{_PS}.get_attn_tensor_model_parallel_world_size"),
("attn_tp_rank", f"{_PS}.get_attn_tensor_model_parallel_rank"),
("attn_cp_size", f"{_PS}.get_attn_context_model_parallel_world_size"),
("attn_cp_rank", f"{_PS}.get_attn_context_model_parallel_rank"),
("attn_dp_size", f"{_DP}.get_attention_dp_size"),
("attn_dp_rank", f"{_DP}.get_attention_dp_rank"),
@@ -897,9 +892,8 @@ class TestForwardFlags(_IsolatedServerArgs):
def test_parallel_config_leaves_trace_under_torch_compile(self):
# Regression: gate helpers such as ``enable_moe_dense_fully_dp()`` read
# parallel config leaves inside compiled model forwards through the
# `config` property, which must stay dynamo-traceable
# (``object.__getattribute__`` graph-breaks).
# parallel config leaves inside compiled model forwards, which must
# stay dynamo-traceable (``object.__getattribute__`` graph-breaks).
# fullgraph=True turns any graph break back into a failure.
import torch
@@ -1368,5 +1362,28 @@ class TestNamedAccessorsCallWhatTheyWrap(CustomTestCase):
self.assertEqual([], wrong, "\n".join(wrong))
class TestParallelLeafReads(_IsolatedServerArgs):
"""The contract ``ParallelContext.__getattr__`` answers a parallel leaf on."""
def test_a_leaf_answers_what_resolution_decided(self):
from sglang.srt.arg_groups.overrides import resolution_result
with get_context().override_server_args() as server_args:
self.assertEqual(
resolution_result(server_args, "nccl_port"),
get_parallel().nccl_port,
"a parallel leaf read off the context disagreed with what "
"resolution decided",
)
def test_before_publish_the_error_names_the_namespace(self):
with self.assertRaisesRegex(ValueError, r"'parallel' not published"):
getattr(ParallelContext(), "nccl_port")
def test_an_unknown_name_is_still_an_attribute_error(self):
with self.assertRaisesRegex(AttributeError, r"has no 'not_a_leaf'"):
getattr(ParallelContext(), "not_a_leaf")
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_runtime_context_override.py
+1 -1
View File
@@ -91,7 +91,7 @@ class TestContextOverride(CustomTestCase):
speculative_accept_threshold_single=0.5,
speculative_accept_threshold_acc=0.9,
)
self.assertEqual(rc.get_parallel().config.pp_max_micro_batch_size, 8)
self.assertEqual(rc.get_parallel().pp_max_micro_batch_size, 8)
self.assertEqual(rc.get_spec().speculative_accept_threshold_single, 0.5)
self.assertEqual(rc.get_spec().speculative_accept_threshold_acc, 0.9)