From 6dd79576cd97b5caf0e98161c9fc0201f75f88dc Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:20:29 +0800 Subject: [PATCH] Expose the declared sglang env vars of a scheduler in its internal state (#35928) --- python/sglang/srt/environ.py | 36 +++++- python/sglang/srt/managers/scheduler.py | 5 +- .../test_scheduler_internal_state_env_vars.py | 122 ++++++++++++++++++ 3 files changed, 158 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/managers/test_scheduler_internal_state_env_vars.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ec1365875..8e055bbf9 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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.""" diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f99724e87..8baa6be6e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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) diff --git a/test/registered/unit/managers/test_scheduler_internal_state_env_vars.py b/test/registered/unit/managers/test_scheduler_internal_state_env_vars.py new file mode 100644 index 000000000..6fdd6eaad --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_internal_state_env_vars.py @@ -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()