Expose the declared sglang env vars of a scheduler in its internal state (#35928)
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
import base64
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
@@ -24,6 +25,9 @@ def _default_hip() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
_NON_UTF8_PREFIX = "base64:"
|
||||
|
||||
|
||||
def _default_cache_subdir(name: str) -> str:
|
||||
"""A directory under SGLANG_CACHE_DIR, for env defaults that track it.
|
||||
|
||||
@@ -36,11 +40,12 @@ def _default_cache_subdir(name: str) -> str:
|
||||
class EnvField:
|
||||
_allow_set_name = True
|
||||
|
||||
def __init__(self, default: Any):
|
||||
def __init__(self, default: Any, secret: bool = False):
|
||||
self.default = default
|
||||
# NOTE: environ can only accept str values, so we need a flag to indicate
|
||||
# whether the env var is explicitly set to None.
|
||||
self._set_to_none = False
|
||||
self.secret = secret
|
||||
|
||||
def __set_name__(self, owner, name):
|
||||
assert EnvField._allow_set_name, "Usage like `a = envs.A` is not allowed"
|
||||
@@ -156,8 +161,8 @@ class _DeprecatedEnvFallback:
|
||||
SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias(True, deprecated_name="SGLANG_NSA_FUSE_TOPK")
|
||||
"""
|
||||
|
||||
def __init__(self, default: Any, deprecated_name: str):
|
||||
super().__init__(default)
|
||||
def __init__(self, default: Any, deprecated_name: str, secret: bool = False):
|
||||
super().__init__(default, secret=secret)
|
||||
self.deprecated_name = deprecated_name
|
||||
|
||||
def get(self) -> Any:
|
||||
@@ -317,6 +322,7 @@ class Envs:
|
||||
# too short when many workers cold-start and load tokenizers in parallel.
|
||||
SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT = EnvInt(10)
|
||||
SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION = EnvBool(True)
|
||||
SGLANG_EXPOSE_OWN_ENV_VARS = EnvBool(False)
|
||||
|
||||
# ===================================================================
|
||||
# Logging
|
||||
@@ -670,7 +676,7 @@ class Envs:
|
||||
# Native web search (Exa). EXA_API_KEY is the vendor BYOK credential
|
||||
# (kept as-is, not renamed to SGLANG_*); the SGLANG_EXA_* knobs tune the
|
||||
# request defaults for the built-in GPT-OSS web_search tool.
|
||||
EXA_API_KEY = EnvStr(None)
|
||||
EXA_API_KEY = EnvStr(None, secret=True)
|
||||
SGLANG_EXA_NUM_RESULTS = EnvInt(10)
|
||||
SGLANG_EXA_SEARCH_TYPE = EnvStr("auto")
|
||||
SGLANG_EXA_INCLUDE_HIGHLIGHTS = EnvBool(True)
|
||||
@@ -1539,6 +1545,28 @@ envs = Envs()
|
||||
EnvField._allow_set_name = False
|
||||
|
||||
|
||||
def exportable_env_vars() -> dict[str, str]:
|
||||
return {
|
||||
field.name: _exportable_value(os.environ[field.name])
|
||||
for field in sorted(
|
||||
(value for value in vars(Envs).values() if isinstance(value, EnvField)),
|
||||
key=lambda field: field.name,
|
||||
)
|
||||
if not field.secret and field.name in os.environ
|
||||
}
|
||||
|
||||
|
||||
def _exportable_value(value: str) -> str:
|
||||
try:
|
||||
value.encode()
|
||||
except UnicodeEncodeError:
|
||||
return (
|
||||
_NON_UTF8_PREFIX
|
||||
+ base64.b64encode(value.encode(errors="surrogateescape")).decode()
|
||||
)
|
||||
return value
|
||||
|
||||
|
||||
class _DeprecatedEnv:
|
||||
"""One deprecated env var: warn if it is set, and optionally forward its
|
||||
(possibly transformed) value to a replacement env var."""
|
||||
|
||||
@@ -102,7 +102,7 @@ from sglang.srt.distributed import get_pp_group, get_world_group
|
||||
from sglang.srt.distributed.parallel_state import get_tp_group
|
||||
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
||||
from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.environ import envs, exportable_env_vars
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.hardware_backend.mlx.runtime import use_mlx
|
||||
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
|
||||
@@ -4463,6 +4463,9 @@ class Scheduler(
|
||||
if info_record is not None:
|
||||
ret["dspark_info_record"] = info_record
|
||||
|
||||
if envs.SGLANG_EXPOSE_OWN_ENV_VARS.get():
|
||||
ret["env_vars"] = exportable_env_vars()
|
||||
|
||||
# These fields are not msgpack-serializable (a config object and a bound
|
||||
# signal handler); no reader consumes them.
|
||||
ret.pop("model_config", None)
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
import json
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestSchedulerInternalStateEnvVars(unittest.TestCase):
|
||||
def _get_internal_state(self) -> dict:
|
||||
scheduler = Scheduler.__new__(Scheduler)
|
||||
scheduler.metrics_reporter = SimpleNamespace(
|
||||
last_gen_throughput=1.0,
|
||||
spec_total_num_forward_ct=0,
|
||||
spec_total_num_accept_tokens=0,
|
||||
step_time_dict={},
|
||||
)
|
||||
scheduler.tp_worker = SimpleNamespace(
|
||||
model_runner=SimpleNamespace(weight_load_mem_usage=1.0),
|
||||
graph_memory_usage=None,
|
||||
)
|
||||
scheduler.token_to_kv_pool_allocator = SimpleNamespace(
|
||||
get_kvcache=lambda: SimpleNamespace(mem_usage=3.0)
|
||||
)
|
||||
scheduler.startup_available_gpu_memory_gb = 4.0
|
||||
scheduler.startup_time = 1.0
|
||||
scheduler.max_total_num_tokens = 100
|
||||
scheduler.swa_tokens_per_layer = None
|
||||
scheduler.max_running_requests = 8
|
||||
scheduler.spec_algorithm = SimpleNamespace(
|
||||
is_none=lambda: True,
|
||||
is_dspark=lambda: False,
|
||||
)
|
||||
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)),
|
||||
):
|
||||
output = scheduler.get_internal_state(recv_req=GetInternalStateReq())
|
||||
|
||||
return output.internal_state
|
||||
|
||||
def test_the_gate_is_declared_off(self):
|
||||
"""Nothing is exposed unless an operator opts in, so the declared default is the safety net."""
|
||||
self.assertIs(envs.SGLANG_EXPOSE_OWN_ENV_VARS.default, False)
|
||||
|
||||
def test_env_vars_absent_when_disabled(self):
|
||||
"""The env_vars key must not exist at all when the gate is off."""
|
||||
with envs.SGLANG_EXPOSE_OWN_ENV_VARS.override(False):
|
||||
internal_state = self._get_internal_state()
|
||||
|
||||
self.assertNotIn("env_vars", internal_state)
|
||||
|
||||
def test_declared_env_vars_exposed_when_enabled(self):
|
||||
"""Enabling the gate exposes the declared, non secret environment of the scheduler."""
|
||||
with envs.SGLANG_EXPOSE_OWN_ENV_VARS.override(True):
|
||||
with patch.dict(
|
||||
"os.environ",
|
||||
{"SGLANG_LOG_SCHEDULER_STATUS_TARGET": "some-value"},
|
||||
):
|
||||
internal_state = self._get_internal_state()
|
||||
|
||||
self.assertIn("env_vars", internal_state)
|
||||
self.assertEqual(
|
||||
internal_state["env_vars"]["SGLANG_LOG_SCHEDULER_STATUS_TARGET"],
|
||||
"some-value",
|
||||
)
|
||||
|
||||
def test_undeclared_env_vars_never_exposed(self):
|
||||
"""Only what Envs declares is auditable; the SGLANG_ namespace also holds real credentials."""
|
||||
with envs.SGLANG_EXPOSE_OWN_ENV_VARS.override(True):
|
||||
with patch.dict(
|
||||
"os.environ",
|
||||
{
|
||||
"SGLANG_LOG_SCHEDULER_STATUS_TARGET": "some-value",
|
||||
"SGLANG_S3_SECRET_ACCESS_KEY": "a-cloud-credential",
|
||||
"SGLANG_DIFFUSION_SLACK_TOKEN": "a-slack-token",
|
||||
},
|
||||
):
|
||||
internal_state = self._get_internal_state()
|
||||
|
||||
exported = internal_state["env_vars"]
|
||||
self.assertNotIn("SGLANG_S3_SECRET_ACCESS_KEY", exported)
|
||||
self.assertNotIn("SGLANG_DIFFUSION_SLACK_TOKEN", exported)
|
||||
self.assertNotIn("a-cloud-credential", json.dumps(exported))
|
||||
|
||||
def test_a_field_marked_secret_is_never_exposed(self):
|
||||
"""A declared credential is still a credential, so the declaration has to be able to say so."""
|
||||
with envs.SGLANG_EXPOSE_OWN_ENV_VARS.override(True):
|
||||
with patch.dict("os.environ", {"EXA_API_KEY": "a-credential"}):
|
||||
internal_state = self._get_internal_state()
|
||||
|
||||
self.assertNotIn("EXA_API_KEY", internal_state["env_vars"])
|
||||
|
||||
def test_a_non_utf8_value_is_encoded_rather_than_dropped(self):
|
||||
"""An undecodable byte in one variable must not cost the whole response its json encoding."""
|
||||
with envs.SGLANG_EXPOSE_OWN_ENV_VARS.override(True):
|
||||
with patch.dict(
|
||||
"os.environ", {"SGLANG_LOG_SCHEDULER_STATUS_TARGET": "bad-\udcff"}
|
||||
):
|
||||
internal_state = self._get_internal_state()
|
||||
|
||||
exported = internal_state["env_vars"]["SGLANG_LOG_SCHEDULER_STATUS_TARGET"]
|
||||
self.assertTrue(exported.startswith("base64:"))
|
||||
json.dumps(internal_state["env_vars"]).encode()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user