config: pin that resolution is reproducible from the raw input (#34094)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
co-authored by
Claude Fable 5
Claude Fable 5
parent
63833f8034
commit
a2199c1dee
@@ -0,0 +1,375 @@
|
||||
"""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 shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
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, {}),)
|
||||
|
||||
# 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 TestResolutionIsReproducible(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 _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)
|
||||
|
||||
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)
|
||||
return ServerArgs(model_path=model_path, **kwargs)
|
||||
|
||||
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
|
||||
value = getattr(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)),
|
||||
)
|
||||
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()
|
||||
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)),
|
||||
)
|
||||
if label == "multimodal":
|
||||
# And the documented exception, asserted rather than
|
||||
# assumed: the multimodal handler's env write does reach
|
||||
# the next resolution. What it carries is the
|
||||
# intermediate's own device-dependent selection — cuda_ipc
|
||||
# on single-node CUDA, cpu on the CPU/ROCm runners (the
|
||||
# same `is_cuda()` gate the handler branches on).
|
||||
expected = "cuda_ipc" if is_cuda() else "cpu"
|
||||
self.assertEqual(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_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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user