1089 lines
48 KiB
Python
1089 lines
48 KiB
Python
"""Resolution is a pure function of the raw input plus this node's environment.
|
|
|
|
The end state for the configuration tier keeps ``ServerArgs`` at the user's raw
|
|
input and lets every process that publishes derive the resolved values itself
|
|
(bags do not cross a process boundary — a child projects its own from the record
|
|
it is handed). That is only sound if resolving the same raw input twice gives
|
|
the same answer, so this pins it:
|
|
|
|
- twice in this process, from equal raw inputs, every field agrees;
|
|
- the resolution is not order-dependent on a shared registry (a second config
|
|
resolved after the first does not inherit its declarations);
|
|
- the raw record the two started from is itself unchanged by resolving a
|
|
sibling.
|
|
|
|
A failure here means some resolution step reads state it also writes, and the
|
|
"re-derive in the child" contract would silently diverge between the launcher
|
|
and its schedulers.
|
|
|
|
Scope: this pins reproducibility *within one process*, which is the fork case
|
|
(the child inherits the parent's environment and its module-level caches). A
|
|
spawn child starts with cold module state instead -- ``functools`` memos in
|
|
``runtime_context`` among them -- so a divergence that needs a cold cache to
|
|
show up is outside what these cases can see.
|
|
"""
|
|
|
|
import copy
|
|
import dataclasses
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
import unittest.mock
|
|
|
|
import torch
|
|
|
|
import sglang
|
|
from sglang.srt.arg_groups.overrides import resolution_result
|
|
from sglang.srt.environ import EnvField, envs
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils import is_cuda
|
|
from sglang.test.ci.ci_register import (
|
|
register_amd_ci,
|
|
register_cpu_ci,
|
|
register_cuda_ci,
|
|
)
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
|
# Also on a GPU runner: the resolution branches that matter most (backend
|
|
# defaults, DeepSeek handlers, capability gates) go through `is_cuda()` /
|
|
# `is_hip()` / device capability, which inspect the actual hardware -- passing
|
|
# device="cuda" on a CPU box does not reach them, so a leak confined to a GPU
|
|
# handler would never fail the CPU registration alone.
|
|
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
|
|
# ROCm too: `is_hip()` gates its own set of backend and DeepSeek handlers, which
|
|
# neither the CPU suite nor a CUDA runner reaches.
|
|
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
|
|
|
|
_MINI_CONFIG = {
|
|
"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,
|
|
}
|
|
|
|
_DEEPSEEK_MINI_CONFIG = {
|
|
"architectures": ["DeepseekV3ForCausalLM"],
|
|
"model_type": "deepseek_v3",
|
|
"hidden_size": 16,
|
|
"intermediate_size": 32,
|
|
"moe_intermediate_size": 32,
|
|
"num_attention_heads": 2,
|
|
"num_key_value_heads": 2,
|
|
"num_hidden_layers": 2,
|
|
"n_routed_experts": 8,
|
|
"n_shared_experts": 1,
|
|
"num_experts_per_tok": 2,
|
|
"first_k_dense_replace": 1,
|
|
"vocab_size": 128,
|
|
"max_position_embeddings": 2048,
|
|
"kv_lora_rank": 8,
|
|
"q_lora_rank": 8,
|
|
"qk_nope_head_dim": 8,
|
|
"qk_rope_head_dim": 8,
|
|
"v_head_dim": 8,
|
|
"topk_method": "greedy",
|
|
"scoring_func": "softmax",
|
|
# index_topk puts this config on the DSA path, whose handlers are the ones
|
|
# that fan out the most (and write process state on the way through).
|
|
"index_topk": 4,
|
|
"index_head_dim": 8,
|
|
"index_n_heads": 2,
|
|
}
|
|
|
|
_MULTIMODAL_MINI_CONFIG = {
|
|
"architectures": ["Qwen2VLForConditionalGeneration"],
|
|
"model_type": "qwen2_vl",
|
|
"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,
|
|
"vision_config": {
|
|
"depth": 2,
|
|
"hidden_size": 16,
|
|
"num_heads": 2,
|
|
"in_chans": 3,
|
|
"patch_size": 14,
|
|
"spatial_merge_size": 2,
|
|
},
|
|
}
|
|
|
|
# The shapes the step-12 audit calls out as the ones whose resolution branches
|
|
# touch process state: a plain text model, a speculative launch, and a MoE/MLA
|
|
# architecture whose handlers fan out the most.
|
|
_SHAPES = (
|
|
("plain", _MINI_CONFIG, {}),
|
|
(
|
|
"speculative",
|
|
_MINI_CONFIG,
|
|
dict(
|
|
speculative_algorithm="EAGLE",
|
|
speculative_num_steps=2,
|
|
speculative_eagle_topk=1,
|
|
speculative_num_draft_tokens=3,
|
|
),
|
|
),
|
|
# The multimodal transport handler is the one that writes
|
|
# SGLANG_USE_CUDA_IPC_TRANSPORT and reads `is_set()` on the way in, so the
|
|
# shape that exercises it belongs in the dual-resolve matrix.
|
|
("multimodal", _MULTIMODAL_MINI_CONFIG, {}),
|
|
# device="cpu" is the one device every host can resolve, and it is what
|
|
# reaches `_handle_cpu_backends` — the golden device="cuda" default never
|
|
# does, so without this shape a leak confined to the CPU handlers would
|
|
# pass the whole matrix.
|
|
("plain_cpu_device", _MINI_CONFIG, dict(device="cpu")),
|
|
)
|
|
|
|
# Resolving the DSA shape needs a physical device: the DeepSeek-DSA arm of
|
|
# `_handle_model_specific_adjustments` probes `torch.cuda.get_device_capability()`
|
|
# to pick the KV dtype and split backends, and that raises on a driverless
|
|
# host. The GPU registrations are what exercise this shape; a GPU-less runner
|
|
# resolves the other shapes only. (Whether a CPU runner even *reaches* the DSA
|
|
# arm depends on the installed transformers surfacing `index_topk` from the
|
|
# mini config, so without this gate the crash appears runner-dependently.)
|
|
if torch.cuda.is_available():
|
|
_SHAPES = _SHAPES + (("deepseek_dsa", _DEEPSEEK_MINI_CONFIG, {}),)
|
|
|
|
# Since #34662 made cuda_ipc opt-in, no auto-resolution reaches the handler's
|
|
# cuda_ipc arm any more -- an explicit request is the only way in, so without
|
|
# this shape that arm (two raises and the pool-budget logging) is resolved by
|
|
# nothing in this file. It is also the only shape whose env write differs from
|
|
# what the *next* resolution would pick on its own, which is what keeps the
|
|
# sticky-carry assertion in `test_a_resolution_does_not_leak_into_the_next`
|
|
# from being vacuous. Gated on `is_cuda()` rather than
|
|
# `torch.cuda.is_available()`: the handler raises for cuda_ipc off NVIDIA CUDA,
|
|
# ROCm included.
|
|
_CUDA_IPC_SHAPES = ()
|
|
if is_cuda():
|
|
_CUDA_IPC_SHAPES = (
|
|
(
|
|
"multimodal_cuda_ipc",
|
|
_MULTIMODAL_MINI_CONFIG,
|
|
dict(mm_feature_transport="cuda_ipc"),
|
|
),
|
|
)
|
|
_SHAPES = _SHAPES + _CUDA_IPC_SHAPES
|
|
|
|
# The one field a previous resolution genuinely dictates for the next one in
|
|
# this process: `_handle_multimodal_feature_transport` writes
|
|
# SGLANG_USE_CUDA_IPC_TRANSPORT so tokenizer workers inherit the decision, and
|
|
# the next resolution reads `is_set()` and adopts it -- even for a text-only
|
|
# model, and even across Engines. That is main's behaviour (reproduced on the
|
|
# stack's base commit), it is inert for a text model, and pinning it here is
|
|
# deliberate: the assertion below states the exception explicitly so a *new*
|
|
# sticky field fails this case instead of hiding behind it.
|
|
_STICKY_ACROSS_RESOLUTIONS = frozenset({"mm_feature_transport"})
|
|
|
|
# `random_seed` is pinned by `_resolved` so it would compare equal anyway; it
|
|
# stays listed because a case that stops pinning it must not silently start
|
|
# comparing a value resolution randomizes.
|
|
_NOT_COMPARABLE = frozenset({"random_seed"})
|
|
|
|
|
|
class _RestoresProcessState:
|
|
"""Resolution leaves process state behind, so a case that resolves has to
|
|
put it back. `_handle_multimodal_feature_transport` sets
|
|
`SGLANG_USE_CUDA_IPC_TRANSPORT` and the same handler reads `is_set()` on the
|
|
way in, so one resolution is visible to the next one in this process -- and
|
|
`TestMultimodalFeatureTransport` is the case that notices.
|
|
"""
|
|
|
|
def _process_state(self):
|
|
"""What a resolution may leave behind: the environment and the
|
|
descriptor-level flag `EnvField.set()` flips, which `os.environ` does
|
|
not carry."""
|
|
# Walk the MRO: `vars(type(envs))` alone would miss fields declared on
|
|
# a base class.
|
|
fields = {}
|
|
for klass in reversed(type(envs).__mro__):
|
|
for name, field in vars(klass).items():
|
|
if isinstance(field, EnvField):
|
|
fields[name] = field
|
|
return (
|
|
dict(os.environ),
|
|
{name: field._set_to_none for name, field in fields.items()},
|
|
)
|
|
|
|
def _restore_process_state(self, state):
|
|
saved_environ, saved_none_flags = state
|
|
os.environ.clear()
|
|
os.environ.update(saved_environ)
|
|
for name, was_none in saved_none_flags.items():
|
|
getattr(type(envs), name)._set_to_none = was_none
|
|
|
|
def setUp(self):
|
|
# Resolution writes process state on the way through --
|
|
# `_handle_multimodal_feature_transport` sets SGLANG_USE_CUDA_IPC_TRANSPORT
|
|
# so tokenizer workers inherit the decision, and the same handler reads
|
|
# `is_set()` on the way in. One resolution is therefore visible to the
|
|
# next one in this process. These cases restore what they touched, and
|
|
# it is a standing caveat on the determinism pinned here: the guarantee
|
|
# holds per raw input *and* the process state a previous resolution left.
|
|
self._pristine_state = self._process_state()
|
|
self.addCleanup(self._restore_process_state, self._pristine_state)
|
|
|
|
def _callTestMethod(self, method):
|
|
# No retry here. CustomTestCase retries once in CI, but `addCleanup`
|
|
# runs after the last attempt, so a second attempt would start from the
|
|
# state the first one leaked -- exactly the regression these cases exist
|
|
# to catch, turned into a pass.
|
|
unittest.TestCase._callTestMethod(self, method)
|
|
|
|
|
|
class TestResolutionIsReproducible(_RestoresProcessState, CustomTestCase):
|
|
def _config_dir(self, config: dict = None) -> str:
|
|
config_dir = tempfile.mkdtemp(prefix="resolution_repro_")
|
|
self.addCleanup(shutil.rmtree, config_dir, ignore_errors=True)
|
|
with open(os.path.join(config_dir, "config.json"), "w") as handle:
|
|
json.dump(config or _MINI_CONFIG, handle)
|
|
return config_dir
|
|
|
|
def _resolved(self, model_path: str, **kwargs) -> ServerArgs:
|
|
# device="cuda" keeps the golden path host-independent: an
|
|
# accelerator-less runner resolves only the base platform, where
|
|
# get_device() raises.
|
|
kwargs.setdefault("device", "cuda")
|
|
kwargs.setdefault("random_seed", 42)
|
|
server_args = ServerArgs(model_path=model_path, **kwargs)
|
|
server_args.resolve_once()
|
|
return server_args
|
|
|
|
def _comparable(self, server_args: ServerArgs) -> dict:
|
|
"""The dataclass fields, and only those.
|
|
|
|
Non-field artifacts a resolution leaves on the instance (the
|
|
`_resolved_overrides` provenance, a cached `model_config`) are not in
|
|
here; `test_the_declaration_provenance_is_reproducible` is what covers
|
|
the one of those that a shared mutable could corrupt.
|
|
"""
|
|
out = {}
|
|
for field in dataclasses.fields(server_args):
|
|
if field.name in _NOT_COMPARABLE:
|
|
continue
|
|
# The resolution result, not the field: a declaration-only resolver
|
|
# never writes the field, so comparing fields would miss exactly
|
|
# the decisions a leak would shift.
|
|
value = resolution_result(server_args, field.name)
|
|
# Nested dataclasses (cuda_graph_config) compare structurally, and
|
|
# everything else is deep-copied: a snapshot that stored the live
|
|
# list/dict would follow an in-place mutation, which is exactly the
|
|
# regression `test_resolving_a_sibling_leaves_the_first_alone` looks
|
|
# for.
|
|
out[field.name] = (
|
|
dataclasses.asdict(value)
|
|
if dataclasses.is_dataclass(value)
|
|
else copy.deepcopy(value)
|
|
)
|
|
return out
|
|
|
|
def test_two_resolutions_of_the_same_input_agree(self):
|
|
for label, config, kwargs in _SHAPES:
|
|
with self.subTest(shape=label):
|
|
# Each shape starts from the state the test method started in,
|
|
# not from what the previous shape's resolution left behind.
|
|
self._restore_process_state(self._pristine_state)
|
|
model_path = self._config_dir(config)
|
|
first = self._resolved(model_path, **kwargs)
|
|
second = self._resolved(model_path, **kwargs)
|
|
self.assertEqual(self._comparable(first), self._comparable(second))
|
|
|
|
def test_a_resolution_does_not_leak_into_the_next(self):
|
|
# A config resolved with an explicit, non-default backend must not shift
|
|
# what the next one picks. Residual process state (env, caches) is the
|
|
# hazard, not the declaration registry, whose providers are import-time.
|
|
model_path = self._config_dir()
|
|
# The control has to be taken *before* the explicit resolution: if that
|
|
# one contaminated the process, a control read afterwards would inherit
|
|
# the same contamination and the assertion would pass vacuously.
|
|
default_before = self._resolved(model_path)
|
|
# torch_native is not the default on either CPU or CUDA, so the probe
|
|
# really diverges on the suite that runs this.
|
|
explicit = self._resolved(model_path, attention_backend="torch_native")
|
|
self.assertEqual(explicit.attention_backend, "torch_native")
|
|
self.assertNotEqual(
|
|
self._comparable(explicit),
|
|
self._comparable(default_before),
|
|
"the probe resolved to the same config as the default, so this case "
|
|
"would pass without exercising order dependence",
|
|
)
|
|
default_after = self._resolved(model_path)
|
|
# Every field, not just the backend: a declaration registry takes
|
|
# arbitrary field dicts, so a leak can land anywhere.
|
|
self.assertEqual(
|
|
self._comparable(default_after), self._comparable(default_before)
|
|
)
|
|
|
|
# A backend probe only diverges on the kernel fields. The shapes whose
|
|
# handlers write process state on the way through -- the multimodal
|
|
# transport one sets SGLANG_USE_CUDA_IPC_TRANSPORT and reads `is_set()`
|
|
# on the way in, DSA fans out the furthest -- are the ones that can
|
|
# leave something the *next* default reads, so each gets its own turn as
|
|
# the intermediate. Whether those writes actually fire is
|
|
# device-dependent, which is why this case is registered on the GPU
|
|
# runners as well as CPU.
|
|
intermediates = (
|
|
("multimodal", _MULTIMODAL_MINI_CONFIG, {}),
|
|
("torch_compile", _MINI_CONFIG, dict(enable_torch_compile=True)),
|
|
)
|
|
# Auto-resolution picks cpu on every runner this case runs on now that
|
|
# cuda_ipc is opt-in, so the auto shape above writes what the next
|
|
# resolution would have picked anyway; the explicit shape is what makes
|
|
# the carry observable at all.
|
|
intermediates += _CUDA_IPC_SHAPES
|
|
if torch.cuda.is_available():
|
|
# Same device gate as _SHAPES: the DSA arm probes the device
|
|
# capability during resolution.
|
|
intermediates += (("deepseek_dsa", _DEEPSEEK_MINI_CONFIG, {}),)
|
|
for label, config, kwargs in intermediates:
|
|
with self.subTest(intermediate=label):
|
|
# Each intermediate starts from the pristine process, for two
|
|
# reasons: it must not inherit what the previous iteration left,
|
|
# and the handlers under test branch on *unset* state -- the
|
|
# multimodal one auto-selects the transport only when
|
|
# SGLANG_USE_CUDA_IPC_TRANSPORT is not set, and any earlier
|
|
# resolution in this process has already set it. `default_before`
|
|
# is the control precisely because it was taken on this state.
|
|
self._restore_process_state(self._pristine_state)
|
|
# The auto-selection branch under test requires the legacy
|
|
# variable UNSET; a runner that exports it (a supported
|
|
# deployment setting) would otherwise pin every resolution to
|
|
# its value and this subtest would assert the environment
|
|
# rather than the handler. `_STICKY_ACROSS_RESOLUTIONS`
|
|
# already excludes the affected field from the equality
|
|
# against `default_before`, so clearing it here does not skew
|
|
# that comparison.
|
|
envs.SGLANG_USE_CUDA_IPC_TRANSPORT.clear()
|
|
intermediate = self._resolved(self._config_dir(config), **kwargs)
|
|
after = self._resolved(model_path)
|
|
without_sticky = lambda snapshot: {
|
|
k: v
|
|
for k, v in snapshot.items()
|
|
if k not in _STICKY_ACROSS_RESOLUTIONS
|
|
}
|
|
self.assertEqual(
|
|
without_sticky(self._comparable(after)),
|
|
without_sticky(self._comparable(default_before)),
|
|
)
|
|
# And the documented exception, asserted rather than assumed,
|
|
# for every intermediate: each one runs the transport handler,
|
|
# so each one writes the variable the next resolution reads.
|
|
# What carries is the legacy *boolean*, not the tri-state field
|
|
# -- the handler writes 1 only for cuda_ipc -- so every other
|
|
# selection comes back as cpu, which is what keeps this honest
|
|
# if a cuda_vmm shape is ever added (its carry is cpu, not
|
|
# cuda_vmm). The `cuda_ipc` shape is the one whose carry
|
|
# differs from the cpu that `default_before` resolved to.
|
|
expected = (
|
|
"cuda_ipc"
|
|
if resolution_result(intermediate, "mm_feature_transport")
|
|
== "cuda_ipc"
|
|
else "cpu"
|
|
)
|
|
self.assertEqual(
|
|
resolution_result(after, "mm_feature_transport"), expected
|
|
)
|
|
|
|
def test_resolving_a_sibling_leaves_the_first_alone(self):
|
|
for label, config, kwargs in _SHAPES:
|
|
with self.subTest(shape=label):
|
|
self._restore_process_state(self._pristine_state)
|
|
model_path = self._config_dir(config)
|
|
first = self._resolved(model_path, **kwargs)
|
|
snapshot = self._comparable(first)
|
|
self._resolved(
|
|
model_path, tp_size=2, chunked_prefill_size=1024, **kwargs
|
|
)
|
|
self.assertEqual(self._comparable(first), snapshot)
|
|
|
|
def test_the_gate_refuses_a_second_resolution(self):
|
|
"""A record that has been resolved is left exactly as it was.
|
|
|
|
Every publishing process calls the gate, and in a child the record
|
|
arrived already resolved -- so this is the property that keeps the
|
|
child agreeing with the parent. The handlers are not written to survive
|
|
a second pass over their own output (the DP-attention step derives the
|
|
chunked prefill size *from* the chunked prefill size), which is why the
|
|
gate refuses rather than re-deriving.
|
|
"""
|
|
for label, config, kwargs in _SHAPES:
|
|
with self.subTest(shape=label):
|
|
self._restore_process_state(self._pristine_state)
|
|
model_path = self._config_dir(config)
|
|
resolved = self._resolved(model_path, **kwargs)
|
|
snapshot = self._comparable(resolved)
|
|
declarations = list(getattr(resolved, "_resolved_overrides", []))
|
|
resolved.resolve_once()
|
|
self.assertEqual(self._comparable(resolved), snapshot)
|
|
self.assertEqual(
|
|
list(getattr(resolved, "_resolved_overrides", [])), declarations
|
|
)
|
|
|
|
def test_the_gate_closes_on_the_dummy_path_too(self):
|
|
"""The dummy model leaves the pipeline early, and the gate still shuts.
|
|
|
|
That exit is above the materialization the gate reads, so a dummy
|
|
record answered "not resolved yet" forever and every publish of one ran
|
|
the handlers again. Nothing about the early exit makes a second pass
|
|
safe -- the handlers above it declare and apply like any other -- and
|
|
the four that do run happening to be idempotent today is what the gate
|
|
exists to stop depending on. So this counts entries rather than
|
|
comparing values: the values agree either way.
|
|
"""
|
|
self._restore_process_state(self._pristine_state)
|
|
record = ServerArgs(model_path="dummy")
|
|
record.resolve_once()
|
|
|
|
entries = []
|
|
original = ServerArgs._run_resolution_pipeline
|
|
|
|
def counted(self):
|
|
entries.append(1)
|
|
return original(self)
|
|
|
|
with unittest.mock.patch.object(
|
|
ServerArgs, "_run_resolution_pipeline", counted
|
|
):
|
|
record.resolve_once()
|
|
self.assertEqual(
|
|
entries,
|
|
[],
|
|
"a resolved dummy record entered the pipeline again, so every "
|
|
"publish of one re-runs the handlers",
|
|
)
|
|
|
|
def test_the_declaration_provenance_is_reproducible(self):
|
|
model_path = self._config_dir()
|
|
first = self._resolved(model_path)
|
|
# Snapshot before the second resolution: if a regression had the
|
|
# registry hand out a shared mutable list, resolving `second` would
|
|
# mutate what `first` still points at and the two would compare equal.
|
|
first_provenance = copy.deepcopy(getattr(first, "_resolved_overrides", None))
|
|
second = self._resolved(model_path)
|
|
self.assertEqual(
|
|
first_provenance,
|
|
getattr(second, "_resolved_overrides", None),
|
|
)
|
|
# And the first record's own list is untouched by the second
|
|
# resolution -- a shared mutable would show up here.
|
|
self.assertEqual(getattr(first, "_resolved_overrides", None), first_provenance)
|
|
|
|
|
|
class TestProgramsResolveBeforeReadingResolution(CustomTestCase):
|
|
"""A program that builds its own record resolves it before reading what
|
|
resolution decides.
|
|
|
|
Construction is inert, so a program that builds a record and then reads a
|
|
resolution-written field reads the CLI default. Two of these shipped past
|
|
the earlier censuses because those are rooted at the `sglang` package: the
|
|
model gateway's launcher sized its worker plan from a raw `dp_size`
|
|
(`--dwdp-size 4` launched one server instead of four) and a speculative
|
|
benchmark forwarded `--mem-fraction-static None` to the server it spawns.
|
|
So the universe here is the *repository*, not the package.
|
|
"""
|
|
|
|
# Entries that hand the record on instead of reading it. Reason required.
|
|
_EXEMPT: dict = {}
|
|
|
|
def _repo_root(self):
|
|
# <repo>/python/sglang/__init__.py -> <repo>
|
|
root = pathlib.Path(next(iter(sglang.__path__))).resolve().parents[1]
|
|
if root.name == "python":
|
|
root = root.parent
|
|
return root
|
|
|
|
def _written_fields(self):
|
|
"""Fields resolution declares, read out of the pipeline's own source.
|
|
|
|
Deliberately local: the chain ratchet has a wider derivation (it also
|
|
walks the model-override registries), but it arrives later in this
|
|
series, and a check that imports it would fail at this PR's boundary.
|
|
Coarser is fine here -- what this needs is the fields the entries below
|
|
actually read -- and the floor keeps it from drifting narrower.
|
|
"""
|
|
import ast
|
|
import dataclasses as _dataclasses
|
|
|
|
from sglang.srt.server_args import ServerArgs as _ServerArgs
|
|
|
|
srt = pathlib.Path(next(iter(sglang.__path__))).resolve() / "srt"
|
|
declarers = {"_declare", "declare_resolution", "declare_late_resolution"}
|
|
fields = set()
|
|
field_names = {field.name for field in _dataclasses.fields(_ServerArgs)}
|
|
# The record plus every module under `arg_groups/`: a handler declares
|
|
# from whichever of the two it lives in.
|
|
sources = [srt / "server_args.py", *sorted((srt / "arg_groups").rglob("*.py"))]
|
|
for source in sources:
|
|
tree = ast.parse(source.read_text(encoding="utf-8-sig"))
|
|
for node in ast.walk(tree):
|
|
# Registry data: provider dict keys are field names as
|
|
# *data*, invisible to the keyword scan below. Filtered
|
|
# against the real field set.
|
|
if isinstance(node, ast.Dict):
|
|
fields |= {
|
|
key.value
|
|
for key in node.keys
|
|
if isinstance(key, ast.Constant)
|
|
and isinstance(key.value, str)
|
|
and key.value in field_names
|
|
}
|
|
if not isinstance(node, ast.Call):
|
|
continue
|
|
func = node.func
|
|
called = (
|
|
func.attr
|
|
if isinstance(func, ast.Attribute)
|
|
else getattr(func, "id", "")
|
|
)
|
|
if called in declarers or called == "update":
|
|
fields |= {
|
|
kw.arg
|
|
for kw in node.keywords
|
|
if kw.arg and (called != "update" or kw.arg in field_names)
|
|
}
|
|
return fields
|
|
|
|
def _candidates(self, root):
|
|
"""Source files that build a record, with the names they bind it to."""
|
|
import ast
|
|
|
|
skip = {".git", "build", "dist", "node_modules", ".venv", "target"}
|
|
found = {}
|
|
for path in sorted(root.rglob("*.py")):
|
|
parts = set(path.relative_to(root).parts)
|
|
if parts & skip:
|
|
continue
|
|
rel = path.relative_to(root).as_posix()
|
|
# Tests build raw records on purpose.
|
|
if rel.startswith("test/") or "/test/" in rel or "/tests/" in rel:
|
|
continue
|
|
try:
|
|
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
|
except (SyntaxError, UnicodeDecodeError):
|
|
continue
|
|
# Which local names are *the srt record*, by import source: the
|
|
# diffusion runtime has a same-spelled `ServerArgs` with no
|
|
# resolution, so the spelling alone is not enough.
|
|
record_classes, record_helpers = set(), set()
|
|
for node in ast.walk(tree):
|
|
if not isinstance(node, ast.ImportFrom):
|
|
continue
|
|
for alias in node.names:
|
|
bound = alias.asname or alias.name
|
|
if node.module == "sglang" and alias.name == "ServerArgs":
|
|
record_classes.add(bound)
|
|
if node.module == "sglang.srt.server_args":
|
|
if alias.name == "ServerArgs":
|
|
record_classes.add(bound)
|
|
if alias.name == "prepare_server_args":
|
|
record_helpers.add(bound)
|
|
names = set()
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Assign):
|
|
targets = node.targets
|
|
elif isinstance(node, (ast.AnnAssign, ast.NamedExpr)):
|
|
# `x: ServerArgs = ...` is an AnnAssign, not an Assign.
|
|
targets = [node.target]
|
|
else:
|
|
continue
|
|
call = getattr(node, "value", None)
|
|
if not isinstance(call, ast.Call):
|
|
continue
|
|
func = call.func
|
|
# `prepare_server_args(argv)` is the CLI launcher's way.
|
|
builds = (
|
|
isinstance(func, ast.Name)
|
|
and func.id in (record_classes | record_helpers)
|
|
) or (
|
|
isinstance(func, ast.Attribute)
|
|
and func.attr == "from_cli_args"
|
|
and isinstance(func.value, ast.Name)
|
|
and func.value.id in record_classes
|
|
)
|
|
if builds:
|
|
names |= {t.id for t in targets if isinstance(t, ast.Name)}
|
|
if names:
|
|
found[rel] = (tree, names, path)
|
|
return found
|
|
|
|
def test_every_program_that_builds_a_record_resolves_it(self):
|
|
import ast
|
|
|
|
root = self._repo_root()
|
|
candidates = self._candidates(root)
|
|
self.assertGreater(
|
|
len(candidates),
|
|
10,
|
|
f"only {len(candidates)} files build a record under {root}; either "
|
|
"this is not a source checkout or the scan broke",
|
|
)
|
|
written = self._written_fields()
|
|
self.assertGreater(len(written), 50, "the written-field set collapsed")
|
|
# What the escaped entries actually read: a narrower derivation goes
|
|
# quiet on exactly those.
|
|
for field in ("dp_size", "mem_fraction_static"):
|
|
self.assertIn(field, written)
|
|
|
|
offenders = []
|
|
for rel, (tree, names, path) in sorted(candidates.items()):
|
|
source = path.read_text(encoding="utf-8-sig")
|
|
if "resolve_once(" in source or "publish(" in source:
|
|
continue
|
|
reads = sorted(
|
|
{
|
|
f"{node.attr}:{node.lineno}"
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.Attribute)
|
|
and isinstance(node.value, ast.Name)
|
|
and node.value.id in names
|
|
and node.attr in written
|
|
}
|
|
)
|
|
if reads and rel not in self._EXEMPT:
|
|
offenders.append(f"{rel} reads {', '.join(reads[:4])}")
|
|
self.assertEqual(
|
|
offenders,
|
|
[],
|
|
"a program builds its own record and reads what resolution decides "
|
|
"without resolving it, so it reads the CLI default:\n "
|
|
+ "\n ".join(offenders),
|
|
)
|
|
self.assertEqual(
|
|
sorted(set(self._EXEMPT) - set(candidates)),
|
|
[],
|
|
"an exemption names a file that no longer builds a record",
|
|
)
|
|
|
|
|
|
class TestForksResolveFirst(CustomTestCase):
|
|
"""A process that forks a child to run the record resolves it first.
|
|
|
|
The pipeline probes the device (the default attention backend reads the CUDA
|
|
capability), and a forked child cannot initialize CUDA once its parent has.
|
|
Construction used to resolve, so the probe always happened in whoever built
|
|
the record; now it happens at the gate, and the gate must not be reached for
|
|
the first time inside a fork.
|
|
"""
|
|
|
|
# Sites inside the launcher: `_launch_subprocesses` resolves at its top, so
|
|
# every fork below it already has a resolved record.
|
|
_AFTER_LAUNCHER_RESOLVE = {
|
|
"srt/entrypoints/engine.py",
|
|
"srt/managers/data_parallel_controller.py",
|
|
"srt/disaggregation/encoder/grpc_server.py",
|
|
"srt/disaggregation/encoder/runtime.py",
|
|
"srt/elastic_ep/expert_backup_manager.py",
|
|
}
|
|
|
|
def test_every_fork_of_a_record_has_a_resolved_one(self):
|
|
import ast
|
|
|
|
package_root = pathlib.Path(next(iter(sglang.__path__))).resolve()
|
|
offenders, examined = [], 0
|
|
for path in sorted(package_root.rglob("*.py")):
|
|
rel = path.relative_to(package_root).as_posix()
|
|
if rel.startswith("test/") or "/test/" in rel:
|
|
continue
|
|
# The diffusion runtime has its own record with no gate.
|
|
if rel.startswith("multimodal_gen/"):
|
|
continue
|
|
try:
|
|
source = path.read_text(encoding="utf-8-sig")
|
|
if "Process" not in source:
|
|
continue
|
|
tree = ast.parse(source)
|
|
except (SyntaxError, UnicodeDecodeError):
|
|
continue
|
|
for node in ast.walk(tree):
|
|
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
continue
|
|
forks = [
|
|
call
|
|
for call in ast.walk(node)
|
|
if isinstance(call, ast.Call)
|
|
and (
|
|
(
|
|
isinstance(call.func, ast.Attribute)
|
|
and call.func.attr == "Process"
|
|
)
|
|
or (
|
|
isinstance(call.func, ast.Name)
|
|
and call.func.id == "Process"
|
|
)
|
|
)
|
|
and "server_args" in (ast.get_source_segment(source, call) or "")
|
|
]
|
|
if not forks:
|
|
continue
|
|
body = ast.get_source_segment(source, node) or ""
|
|
examined += 1
|
|
# `spawn` starts a fresh interpreter, so the child may probe.
|
|
if 'get_context("spawn")' in body or "'spawn'" in body:
|
|
continue
|
|
if "resolve_once(" in body or "publish(" in body:
|
|
continue
|
|
if rel in self._AFTER_LAUNCHER_RESOLVE:
|
|
continue
|
|
offenders.append(f"{rel}:{forks[0].lineno} {node.name}")
|
|
self.assertGreater(
|
|
examined, 5, f"only {examined} fork sites found; the scan broke"
|
|
)
|
|
self.assertEqual(
|
|
offenders,
|
|
[],
|
|
"these fork a child that will resolve the record, without resolving "
|
|
"it first -- the child cannot initialize CUDA if this process "
|
|
f"already has:\n " + "\n ".join(offenders),
|
|
)
|
|
|
|
|
|
class TestACopyStaysResolved(_RestoresProcessState, CustomTestCase):
|
|
"""A resolved record copied with `dataclasses.replace` loses what makes it
|
|
resolved, and the next publish resolves it a second time -- over values it
|
|
already decided. The Ray paths copy a resolved record to set
|
|
`dist_init_addr`, which is how they reach this.
|
|
"""
|
|
|
|
def _resolved(self):
|
|
config_dir = tempfile.mkdtemp(prefix="replace_resolved_")
|
|
self.addCleanup(shutil.rmtree, config_dir, ignore_errors=True)
|
|
with open(os.path.join(config_dir, "config.json"), "w") as handle:
|
|
json.dump(_MINI_CONFIG, handle)
|
|
# Two steps that are not repeatable on their own output.
|
|
server_args = ServerArgs(
|
|
model_path=config_dir,
|
|
device="cuda",
|
|
dp_size=2,
|
|
tp_size=2,
|
|
enable_dp_attention=True,
|
|
random_seed=42,
|
|
)
|
|
server_args.resolve_once()
|
|
return server_args
|
|
|
|
def test_a_bare_replace_resolves_again_and_lands_in_the_same_place(self):
|
|
"""A bare copy resolves to the same place: the fields are the raw input.
|
|
|
|
`dataclasses.replace` copies the fields, so a bare copy re-runs
|
|
resolution over the *same input* the parent got -- the DP-attention
|
|
halving and the conservativeness scaling apply once. `replace_resolved`
|
|
buys something else: it carries the parent's declarations and its
|
|
`model_config`, so the copy answers without resolving at all.
|
|
"""
|
|
parent = self._resolved()
|
|
bare = dataclasses.replace(parent, dist_init_addr="1.2.3.4:5000")
|
|
self.assertFalse(
|
|
getattr(bare, "_resolution_finished", False),
|
|
"a bare replace carried the flag; then this test proves nothing",
|
|
)
|
|
bare.resolve_once()
|
|
drifted = {
|
|
field.name: (
|
|
resolution_result(parent, field.name),
|
|
resolution_result(bare, field.name),
|
|
)
|
|
for field in dataclasses.fields(parent)
|
|
if field.name not in ("dist_init_addr", "random_seed")
|
|
and repr(resolution_result(parent, field.name))
|
|
!= repr(resolution_result(bare, field.name))
|
|
}
|
|
self.assertEqual(
|
|
drifted,
|
|
{},
|
|
"resolving a bare copy landed somewhere else, so the pipeline is "
|
|
"reading its own output again",
|
|
)
|
|
|
|
def test_replace_resolved_keeps_the_parents_resolution(self):
|
|
parent = self._resolved()
|
|
copy_ = parent.replace_resolved("ray.test", dist_init_addr="1.2.3.4:5000")
|
|
self.assertTrue(getattr(copy_, "_resolution_finished", False))
|
|
drifted = {
|
|
field.name: (getattr(parent, field.name), getattr(copy_, field.name))
|
|
for field in dataclasses.fields(parent)
|
|
if field.name != "dist_init_addr"
|
|
and getattr(parent, field.name) != getattr(copy_, field.name)
|
|
}
|
|
self.assertEqual(
|
|
drifted,
|
|
{},
|
|
f"the copy differs from its parent beyond the change: {drifted}",
|
|
)
|
|
self.assertEqual(copy_.dist_init_addr, "1.2.3.4:5000")
|
|
|
|
def test_the_copy_carries_what_resolution_left_on_the_record(self):
|
|
"""Not just the stash and the flag.
|
|
|
|
`get_model_config()` memoizes on the record, and that cache is filled
|
|
during resolution. A copy that is marked resolved but arrives without it
|
|
cannot fill it -- the read-only guard refuses the cache write -- so the
|
|
first `get_model_config()` raises. That is what killed the Ray
|
|
schedulers, and it is why the carry is enumerated from the instance
|
|
rather than from a list of names.
|
|
"""
|
|
parent = self._resolved()
|
|
copy_ = parent.replace_resolved("ray.test", dist_init_addr="1.2.3.4:5000")
|
|
fields = {field.name for field in dataclasses.fields(parent)}
|
|
missing = sorted(
|
|
name
|
|
for name in vars(parent)
|
|
if name not in fields and name not in vars(copy_)
|
|
)
|
|
self.assertEqual(
|
|
missing,
|
|
[],
|
|
f"the copy did not carry what resolution left on the record: {missing}",
|
|
)
|
|
self.assertIsNotNone(copy_.get_model_config())
|
|
# Containers are copied, so the copy's declaration stays with it.
|
|
self.assertEqual(
|
|
len(parent._resolved_overrides) + 1, len(copy_._resolved_overrides)
|
|
)
|
|
|
|
def test_the_change_reaches_the_bags(self):
|
|
"""The projection reads the raw snapshot plus the declarations, so a
|
|
change the copy only wrote to the field would publish the parent's raw
|
|
value."""
|
|
from sglang.srt.runtime_context import (
|
|
get_parallel,
|
|
get_schedule,
|
|
publish,
|
|
reset_context,
|
|
)
|
|
|
|
parent = self._resolved()
|
|
copy_ = parent.replace_resolved("ray.test", dist_init_addr="1.2.3.4:5000")
|
|
self.addCleanup(reset_context)
|
|
reset_context()
|
|
publish(copy_, role="scheduler")
|
|
self.assertEqual(get_parallel().dist_init_addr, "1.2.3.4:5000")
|
|
self.assertEqual(
|
|
get_schedule().chunked_prefill_size,
|
|
resolution_result(parent, "chunked_prefill_size"),
|
|
"publishing the copy re-ran resolution; the bag disagrees with what "
|
|
"the parent's resolution decided",
|
|
)
|
|
|
|
def test_no_bare_replace_of_a_record_outside_the_helper(self):
|
|
"""`dataclasses.replace` on a record is the helper's job now.
|
|
|
|
Derived, not listed: any `dataclasses.replace` whose first argument is
|
|
named for a record. The helper's own call is the positive control -- if
|
|
the scan stops seeing it, the scan broke rather than the tree.
|
|
"""
|
|
import ast
|
|
|
|
# The repository, not the package: the gateway is outside `sglang/`.
|
|
package_root = pathlib.Path(next(iter(sglang.__path__))).resolve().parents[1]
|
|
if package_root.name == "python":
|
|
package_root = package_root.parent
|
|
helper = "python/sglang/srt/server_args.py"
|
|
bare, inside_helper = [], 0
|
|
|
|
def replaces_a_record(node, record_names):
|
|
if not (
|
|
isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Attribute)
|
|
and node.func.attr == "replace"
|
|
and isinstance(node.func.value, ast.Name)
|
|
and node.func.value.id == "dataclasses"
|
|
and node.args
|
|
):
|
|
return False
|
|
first = node.args[0]
|
|
name = (
|
|
first.id if isinstance(first, ast.Name) else getattr(first, "attr", "")
|
|
)
|
|
return name in record_names or "server_args" in name
|
|
|
|
for path in sorted(package_root.rglob("*.py")):
|
|
try:
|
|
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
|
except SyntaxError:
|
|
continue
|
|
rel = path.relative_to(package_root).as_posix()
|
|
# `self` is a record only inside the record's own class body.
|
|
in_record_class = [
|
|
node
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.ClassDef) and node.name == "ServerArgs"
|
|
]
|
|
for scope, record_names in [(tree, set())] + [
|
|
(klass, {"self"}) for klass in in_record_class
|
|
]:
|
|
for node in ast.walk(scope):
|
|
if not replaces_a_record(node, record_names):
|
|
continue
|
|
if rel == helper and record_names:
|
|
inside_helper += 1
|
|
elif not record_names:
|
|
bare.append(f"{rel}:{node.lineno}")
|
|
self.assertEqual(
|
|
inside_helper,
|
|
1,
|
|
"the scan no longer finds `replace_resolved`'s own call; it broke",
|
|
)
|
|
self.assertEqual(
|
|
bare,
|
|
[],
|
|
"a record is copied with a bare `dataclasses.replace`, so the copy "
|
|
"loses the parent's resolution and the next publish resolves it "
|
|
"again: " + ", ".join(bare),
|
|
)
|
|
|
|
|
|
class TestTheResolutionSeamHasOneCaller(CustomTestCase):
|
|
"""The pipeline is entered from exactly one place, and that place decides
|
|
whether it runs at all.
|
|
|
|
``resolve_once`` is the gate: the handlers are not written to survive a
|
|
second pass over their own output, so a record must go through the pipeline
|
|
at most once. Keeping the pipeline itself down to a single caller is what
|
|
makes that gate impossible to bypass -- and what keeps the remaining move
|
|
(construction time to publish time) a matter of who calls the gate.
|
|
"""
|
|
|
|
def test_only_the_gate_runs_the_pipeline(self):
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
import sglang
|
|
|
|
package_root = Path(next(iter(sglang.__path__)))
|
|
callers = []
|
|
for path in sorted(package_root.rglob("*.py")):
|
|
try:
|
|
source = path.read_text()
|
|
if "_run_resolution_pipeline" not in source:
|
|
continue
|
|
tree = ast.parse(source)
|
|
except SyntaxError:
|
|
continue
|
|
# The full (class, function, ...) scope chain, so the assertion can
|
|
# say "the one caller is ServerArgs.__post_init__" -- not merely
|
|
# that nothing outside a function named __post_init__ calls it.
|
|
scopes = {}
|
|
for node in ast.walk(tree):
|
|
own = scopes.get(id(node), ())
|
|
if isinstance(
|
|
node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)
|
|
):
|
|
own = own + (node.name,)
|
|
for child in ast.iter_child_nodes(node):
|
|
scopes[id(child)] = own
|
|
for node in ast.walk(tree):
|
|
if (
|
|
isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Attribute)
|
|
and node.func.attr == "_run_resolution_pipeline"
|
|
):
|
|
rel = path.relative_to(package_root).as_posix()
|
|
callers.append((rel, ".".join(scopes.get(id(node), ()))))
|
|
# Every call, compared whole: a removed call, a duplicate inside
|
|
# __post_init__, or another class growing a same-named __post_init__
|
|
# all show up here.
|
|
self.assertEqual(
|
|
[("srt/server_args.py", "ServerArgs.resolve_once")],
|
|
callers,
|
|
"the resolution pipeline must be entered exactly once, from "
|
|
f"ServerArgs.resolve_once; found: {callers}",
|
|
)
|
|
|
|
def test_the_gate_is_reached_from_the_launcher_and_from_publish(self):
|
|
"""Both entries go through the gate, so neither can resolve twice.
|
|
|
|
The launcher resolves the engine's record before reading any resolved
|
|
value from it; every publishing process asks the gate on the way in and
|
|
finds nothing left to do when the record arrived resolved.
|
|
"""
|
|
import ast
|
|
from pathlib import Path
|
|
|
|
import sglang
|
|
|
|
package_root = Path(next(iter(sglang.__path__)))
|
|
callers = []
|
|
for path in sorted(package_root.rglob("*.py")):
|
|
try:
|
|
source = path.read_text()
|
|
if "resolve_once" not in source:
|
|
continue
|
|
tree = ast.parse(source)
|
|
except SyntaxError:
|
|
continue
|
|
for node in ast.walk(tree):
|
|
# `self.resolve_once()` at construction; publish looks the
|
|
# attribute up first, so it appears as a bare name call.
|
|
called = (
|
|
isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Attribute)
|
|
and node.func.attr == "resolve_once"
|
|
) or (
|
|
isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Name)
|
|
and node.func.id == "resolve_once"
|
|
)
|
|
if called:
|
|
callers.append(path.relative_to(package_root).as_posix())
|
|
machinery = {"srt/entrypoints/engine.py", "srt/runtime_context.py"}
|
|
self.assertEqual(
|
|
[
|
|
# Program entries: each builds a record from its own
|
|
# arguments and then reads effective configuration, or hands it
|
|
# to a fork that must not be the first to probe the device.
|
|
"benchmark/endpoint.py",
|
|
"benchmark/offline_throughput.py",
|
|
"benchmark/one_batch.py",
|
|
"benchmark/one_batch_server.py",
|
|
"compile_deep_gemm.py",
|
|
"lang/backend/runtime_endpoint.py",
|
|
"launch_server.py",
|
|
# The mechanism.
|
|
"srt/entrypoints/engine.py",
|
|
"srt/entrypoints/http_server_engine.py",
|
|
"srt/runtime_context.py",
|
|
],
|
|
sorted(set(callers)),
|
|
f"the resolution gate grew or lost a caller: {sorted(set(callers))}",
|
|
)
|
|
# The rule the list stands for: a caller that is not the mechanism
|
|
# resolves a record it built itself from argv. Anything else was handed
|
|
# one someone already resolved, or should publish.
|
|
for caller in sorted(set(callers) - machinery):
|
|
source = (package_root / caller).read_text()
|
|
# Either the module turned argv into the record -- the dataclass,
|
|
# the CLI classmethod, or the argv helper `launch_server.py` uses
|
|
# -- or it hands the record to a fork, which has to resolve first:
|
|
# the pipeline probes the device and a forked child cannot
|
|
# re-initialize CUDA. A worker handed a resolved record is neither.
|
|
builds_its_own = any(
|
|
spelling in source
|
|
for spelling in (
|
|
"ServerArgs(",
|
|
".from_cli_args(",
|
|
"prepare_server_args(",
|
|
)
|
|
) or ("Process(" in source and "server_args" in source)
|
|
# `assertTrue`, not `assertIn`: the container is a whole module.
|
|
self.assertTrue(
|
|
builds_its_own,
|
|
f"{caller} calls the resolution gate but does not build the "
|
|
"record it resolves; a record it was handed is already "
|
|
"resolved by whoever built it, and publish resolves what it "
|
|
"is handed",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|