[Core] Refactor server argument choices (#36586)

This commit is contained in:
Lianmin Zheng
2026-08-27 01:56:03 -07:00
committed by GitHub
parent a126a5fa31
commit 3402265989
5 changed files with 192 additions and 218 deletions
+17 -16
View File
@@ -21,29 +21,30 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../python"))
class TestDSAChoicesAndFields(unittest.TestCase):
"""Verify DSA_CHOICES constant and ServerArgs field renaming."""
"""Verify DSA CLI choices and ServerArgs field renaming."""
def setUp(self):
from sglang.srt.server_args import (
DSA_CHOICES,
NSA_CHOICES,
ServerArgs,
)
from sglang.srt.server_args import ServerArgs
self.ServerArgs = ServerArgs
self.DSA_CHOICES = DSA_CHOICES
self.NSA_CHOICES = NSA_CHOICES
parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(parser)
self.actions = {
option: action
for action in parser._actions
for option in action.option_strings
}
def test_dsa_choices_is_canonical(self):
self.assertIn("fa3", self.DSA_CHOICES)
self.assertIn("tilelang", self.DSA_CHOICES)
self.assertIn("flashinfer_sparse_mla", self.DSA_CHOICES)
choices = self.actions["--dsa-prefill-backend"].choices
self.assertIn("fa3", choices)
self.assertIn("tilelang", choices)
self.assertIn("flashinfer_sparse_mla", choices)
def test_nsa_choices_is_alias(self):
self.assertIs(
self.NSA_CHOICES,
self.DSA_CHOICES,
"NSA_CHOICES must be the same object as DSA_CHOICES",
def test_nsa_choices_match_dsa_choices(self):
self.assertEqual(
self.actions["--nsa-prefill-backend"].choices,
self.actions["--dsa-prefill-backend"].choices,
)
def test_serverargs_has_dsa_fields(self):
@@ -104,8 +104,8 @@ def test_epd_rejection_reads_the_resolved_transfer_backend():
The record is produced by actual resolution -- a language-only Kimi-K3
launch at TP2, whose `encoder_transfer_backend` starts at the argument
default `"auto"` (`ENCODER_TRANSFER_BACKEND_CHOICES[0]`) and is filled in
by `resolve_encoder_transfer_backend` to `"zmq_to_tokenizer"`. The guard
default `"auto"` and is filled in by `resolve_encoder_transfer_backend` to
`"zmq_to_tokenizer"`. The guard
reads that resolved value out of the published bags, so the rejection
survives the record going raw: what a reader must never do is go back to
the record for this field.
@@ -1,5 +1,5 @@
import argparse
import dataclasses
import importlib
import json
import os
import socket
@@ -2067,55 +2067,62 @@ class TestCutedslMoeMaxNumTokens(CustomTestCase):
class TestSamplingBackendTokenOracleEnvGate(CustomTestCase):
"""The 'token_oracle' choice is gated on SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.
The choice set is built once at server_args.py import time, so each subtest
reloads the module with the env var set to the desired value.
The choice set is finalized when CLI arguments are registered, so each
parser must reflect the environment at construction time.
"""
def _reload_server_args_with_env(self, *, enabled: bool):
previous = os.environ.get("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE")
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1" if enabled else "0"
try:
return importlib.reload(server_args_module)
finally:
if previous is None:
os.environ.pop("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE", None)
else:
os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = previous
def test_token_oracle_rejected_when_env_disabled(self):
reloaded = self._reload_server_args_with_env(enabled=False)
self.assertNotIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES)
with patch.dict(os.environ, {"SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "0"}):
with self.assertRaises(SystemExit):
server_args_module.prepare_server_args(
[
"--model-path",
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
"--sampling-backend",
"token_oracle",
]
)
with self.assertRaises(SystemExit):
reloaded.prepare_server_args(
def test_token_oracle_accepted_when_env_enabled(self):
with patch.dict(os.environ, {"SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1"}):
parsed = server_args_module.prepare_server_args(
[
"--model-path",
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
"--sampling-backend",
"token_oracle",
# Explicit device so ServerArgs.__post_init__ does not call
# get_device() (fails on CPU-only CI runners) and does not run
# _handle_cpu_backends (which would override sampling_backend
# to "pytorch", masking what we want to verify).
"--device",
"cuda",
]
)
def test_token_oracle_accepted_when_env_enabled(self):
reloaded = self._reload_server_args_with_env(enabled=True)
self.assertIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES)
parsed = reloaded.prepare_server_args(
[
"--model-path",
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
"--sampling-backend",
"token_oracle",
# Explicit device so ServerArgs.__post_init__ does not call
# get_device() (fails on CPU-only CI runners) and does not run
# _handle_cpu_backends (which would override sampling_backend
# to "pytorch", masking what we want to verify).
"--device",
"cuda",
]
)
self.assertEqual(parsed.sampling_backend, "token_oracle")
def test_gate_is_recomputed_for_each_parser(self):
with patch.dict(os.environ, {"SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1"}):
enabled_parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(enabled_parser)
with patch.dict(os.environ, {"SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "0"}):
disabled_parser = argparse.ArgumentParser()
ServerArgs.add_cli_args(disabled_parser)
enabled_action = next(
action
for action in enabled_parser._actions
if action.dest == "sampling_backend"
)
disabled_action = next(
action
for action in disabled_parser._actions
if action.dest == "sampling_backend"
)
self.assertIn("token_oracle", enabled_action.choices)
self.assertNotIn("token_oracle", disabled_action.choices)
class TestDeepEPv2Args(CustomTestCase):
"""DeepEP v2 server-argument resolution and validation."""