Files
sglang/test/registered/unit/test_runtime_context_override.py
T

150 lines
6.0 KiB
Python

"""Context-first mutation.
``get_context().override(source, **fields)`` is the business mutation entry: it
writes the resolved config bags (the single source of truth) and never touches
``server_args`` (the pristine startup record). Routing is by NS metadata; a bad
field aborts before any write; provenance is recorded.
"""
import unittest
import msgspec
from sglang.srt import runtime_context as rc
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
class TestContextOverride(CustomTestCase):
def setUp(self):
rc.reset_context()
def tearDown(self):
rc.reset_context()
def _publish(self):
sa = ServerArgs(model_path="dummy")
# Through publish, so the record is resolved the way a process resolves it.
rc.publish(sa, role="test")
return sa
def test_override_writes_bag_not_server_args(self):
sa = self._publish()
# The published leaf, not the field: `hicache_ratio` is resolved by
# declaration, so the field still holds what the caller passed.
before = rc.get_memory().hicache_ratio
pristine = sa.hicache_ratio
rc.get_context().override("test", hicache_ratio=before + 1.0)
self.assertEqual(rc.get_memory().hicache_ratio, before + 1.0)
# server_args stays the pristine startup record: the override does not
# touch it, and neither did resolution.
self.assertEqual(sa.hicache_ratio, pristine)
def test_override_routes_across_namespaces(self):
self._publish()
rc.get_context().override(
"test",
moe_runner_backend="triton",
page_size=64,
disaggregation_mode="decode",
)
self.assertEqual(rc.get_exec().moe.moe_runner_backend, "triton")
self.assertEqual(rc.get_schedule().page_size, 64)
self.assertEqual(rc.get_disagg().disaggregation_mode, "decode")
def test_override_unknown_field_raises_and_is_atomic(self):
self._publish()
before = rc.get_memory().hicache_ratio
with self.assertRaises(ValueError):
rc.get_context().override(
"test", hicache_ratio=before + 5.0, not_a_real_field=1
)
# No partial write: the valid field was not applied.
self.assertEqual(rc.get_memory().hicache_ratio, before)
def test_override_before_publish_raises(self):
with self.assertRaises(ValueError):
rc.get_context().override("test", page_size=32)
def test_override_provenance_recorded(self):
self._publish()
rc.get_context().override("srcA", page_size=16)
log = rc.get_context().overrides_log()
self.assertEqual(log[-1], ("srcA", {"page_size": 16}))
def test_republish_resets_provenance(self):
self._publish()
rc.get_context().override("srcA", page_size=16)
self.assertTrue(rc.get_context().overrides_log())
self._publish()
self.assertEqual(rc.get_context().overrides_log(), [])
def test_set_internal_state_fields_reach_parallel_and_spec(self):
# The fields /set_internal_state overrides must reach the accessors the
# (1e) flipped readers now use: pp via get_parallel(), thresholds via
# get_spec().
self._publish()
rc.get_context().override(
"update_server_args",
pp_max_micro_batch_size=8,
speculative_accept_threshold_single=0.5,
speculative_accept_threshold_acc=0.9,
)
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)
def test_kv_cache_dtype_override_reaches_get_model_not_server_args(self):
# Load-time resolution: the resolved kv-cache dtype is written
# to the model bag; server_args stays the RAW resolver input.
sa = self._publish()
raw = sa.kv_cache_dtype
rc.get_context().override(
"ModelRunner.configure_kv_cache_dtype", kv_cache_dtype="fp8_e4m3"
)
self.assertEqual(rc.get_model().kv_cache_dtype, "fp8_e4m3")
self.assertEqual(sa.kv_cache_dtype, raw)
def test_bare_server_args_write_raises_after_resolution(self):
# server_args is read-only after resolution: resolved config changes go
# to the bags, a per-runner config to a derived variant.
sa = ServerArgs(model_path="dummy")
msgspec.Struct.__setattr__(sa, "_resolution_finished", True)
with self.assertRaises(AttributeError):
sa.page_size = 999
def test_publish_records_role(self):
rc.publish(ServerArgs(model_path="dummy"), role="scheduler")
self.assertEqual(rc.publish_role(), "scheduler")
def test_legacy_shims_record_roles(self):
# Unit 2a: the legacy setters publish with their process role.
from sglang.srt.server_args import (
set_global_server_args_for_scheduler,
set_global_server_args_for_tokenizer,
)
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
self.assertEqual(rc.publish_role(), "scheduler")
set_global_server_args_for_tokenizer(ServerArgs(model_path="dummy"))
self.assertEqual(rc.publish_role(), "tokenizer")
def test_reset_clears_role(self):
rc.publish(ServerArgs(model_path="dummy"), role="test")
rc.reset_context()
self.assertIsNone(rc.publish_role())
def test_direct_install_clears_role(self):
# A role-less set_server_args (test overrides, draft-worker builds)
# must not inherit the previous lifecycle's role.
rc.publish(ServerArgs(model_path="dummy"), role="scheduler")
rc.get_context().set_server_args(ServerArgs(model_path="dummy"))
self.assertIsNone(rc.publish_role())
if __name__ == "__main__":
unittest.main()