From 7071cfb873a403bba48bf1d700129eec358735cf Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 1 Aug 2026 08:59:49 -0700 Subject: [PATCH] runtime_context: per-role namespace enforcement behind SGLANG_ROLE_NAMESPACES (#33172) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit publish(role=...) has recorded provenance since the namespace split; this wires the enforcement the role was reserved for. SGLANG_ROLE_NAMESPACES selects the mode: - off (default): no bookkeeping; the mode gate in config_bag stays a dead-branch-prunable check under dynamo (bag reads run inside compiled forwards — pinned by a fullgraph test). - record: audit mode — collect (role, namespace) pairs per process and persist each new pair immediately to SGLANG_ROLE_NAMESPACES_OUT (worker teardown skips atexit), plus a per-process stderr summary at exit. - enforce: a bag read outside the role's ROLE_NAMESPACE_SETS entry fails closed with an actionable error; None entries mean full tree. Sets are filled only where audits back them: dp_controller reads only exec (record-mode plain + DP-attention smokes agree with the module's static read set — the elastic-EP gate). tokenizer observed zero bag reads (per-instance managers read self.server_args by design) but keeps the full tree until the multi-tokenizer disagg shape (TokenizerWorker's get_disagg read) is audited; encoder / expert_backup / weight_cache_daemon likewise await their deployment shapes. Verified end-to-end: DP-attention smoke under enforce boots and serves with zero violations. --- python/sglang/srt/environ.py | 8 + python/sglang/srt/runtime_context.py | 184 +++++++++++++++++- .../unit/test_runtime_context_config_bags.py | 110 +++++++++++ 3 files changed, 299 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 029a917a4..5d21605de 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -232,6 +232,14 @@ class Envs: # must go through ServerArgs.override() (enabled by the test harness). SGLANG_STRICT_CONFIG_MUTATION = EnvBool(False) + # Per-role config-namespace bookkeeping: off / record / enforce (value is + # validated fail-loud in runtime_context, which resolves it once at import + # so the read stays dynamo-prunable). + SGLANG_ROLE_NAMESPACES = EnvStr("off") + # record mode: append each newly observed (role, namespace) pair to this + # file so the audit survives signal-killed workers. + SGLANG_ROLE_NAMESPACES_OUT = EnvStr(None) + # Model & File Download SGLANG_USE_MODELSCOPE = EnvBool(False) # Controls weight-file ordering for load-time I/O optimization. diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 1d0dbab38..617d96489 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -47,6 +47,8 @@ test-only ``override(**kw)``. from __future__ import annotations import dataclasses +import os +import sys from contextlib import contextmanager from typing import TYPE_CHECKING, Any @@ -826,8 +828,34 @@ class RuntimeContext: bags = self._config_bags if not bags or name not in bags: raise ValueError(f"config namespace {name!r} not published") + if _ROLE_NS_MODE != "off": + self._check_role_namespace(name) return bags[name] + def _check_role_namespace(self, name: str) -> None: + # Out of line so the mode gate above stays one dead-branch-prunable + # check under dynamo in the default "off" mode (config_bag runs inside + # compiled model forwards). + role = self._publish_role + if _ROLE_NS_MODE == "record": + if not _is_compiling(): + _record_namespace_read(role, name) + elif _ROLE_NS_MODE == "enforce" and role is not None: + if role not in ROLE_NAMESPACE_SETS: + raise ValueError( + f"publish role {role!r} has no ROLE_NAMESPACE_SETS entry; " + "declare its namespace set (None for the full tree)." + ) + allowed = ROLE_NAMESPACE_SETS[role] + if allowed is not None and name not in allowed: + raise ValueError( + f"config namespace {name!r} is outside the declared set " + f"for publish role {role!r} ({sorted(allowed)}). If this " + "read is legitimate for the process type, extend " + "ROLE_NAMESPACE_SETS; if not, the read belongs in a " + "different process or behind a per-instance boundary." + ) + def override(self, source: str, **fields) -> None: """The business mutation entry: write resolved config leaves onto the namespace bags — the single source of truth. It does @@ -1102,6 +1130,138 @@ def get_observability() -> _ConfigBag: return _CONTEXT.config_bag("observability") +# --- Per-role namespace sets (2c) ------------------------------------------- +# +# ``publish(role=...)`` records which process type installed the config; this +# table declares which top-level config namespaces each role reads. ``None`` +# means the full tree — either the role genuinely needs everything (scheduler) +# or its deployment shape has not been audited yet (restrict only what smoke +# coverage can verify). ``parallel`` is served by ``get_parallel()`` and every +# process legitimately reads topology config, so it is not part of this table. +# +# ``SGLANG_ROLE_NAMESPACES`` selects the mode (read once at import): +# off (default) no bookkeeping, zero overhead; +# record audit mode — collect (role, namespace) reads per process and dump +# them at exit (the data that seeds this table). Reads made inside +# torch.compile-traced code are NOT observed (recording is pruned +# under tracing to keep capture legal) — run audits with +# compilation disabled before restricting a role. +# enforce fail closed — a bag read outside the role's declared set raises. +ROLE_NAMESPACE_SETS: dict[str, frozenset[str] | None] = { + # Reads (almost) everything by design — the model-executing process. + "scheduler": None, + "launcher": None, + "test": None, + # Audited (record-mode smokes, plain + DP-attention): the DP controller + # reads only the elastic-EP gate; its module's static read set agrees. + "dp_controller": frozenset({"exec"}), + # Zero bag reads observed (per-instance managers read self.server_args by + # design). Keep full: it may need namespaces (e.g. disagg) once tokenizer + # paths migrate off self.server_args, and restricting on a zero-read + # audit would be guesswork. + "tokenizer": None, + # Deployment shapes not exercised locally; audit before restricting. + "encoder": None, + "expert_backup": None, + "weight_cache_daemon": None, +} + + +def _validated_role_ns_mode(value: str) -> str: + mode = value.strip().lower() + if mode not in ("off", "record", "enforce"): + raise ValueError( + f"SGLANG_ROLE_NAMESPACES={value!r} is not one of off / record / " + "enforce — refusing to guess (a typo here would silently disable " + "enforcement)." + ) + return mode + + +def _role_ns_mode_from_env() -> str: + # Resolved once at import so the config_bag gate stays a dynamo-prunable + # constant; validated fail-loud here (EnvField's warn-and-default parse + # would silently turn a typo into "off"). + from sglang.srt.environ import envs + + return _validated_role_ns_mode(envs.SGLANG_ROLE_NAMESPACES.get()) + + +_ROLE_NS_MODE = _role_ns_mode_from_env() +_RECORDED_NS_READS: set[tuple[str | None, str]] = set() +_RECORD_DUMP_REGISTERED = False + + +def _is_compiling() -> bool: + # Recording has Python side effects (set mutation, file I/O, atexit) that + # must never run under tracing; torch.compiler.is_compiling() is dynamo's + # sanctioned probe. The lazy lookup keeps this module import-light. + torch = sys.modules.get("torch") + return torch is not None and torch.compiler.is_compiling() + + +def _ensure_record_dump_registered() -> None: + global _RECORD_DUMP_REGISTERED + if not _RECORD_DUMP_REGISTERED: + _RECORD_DUMP_REGISTERED = True + import atexit + + atexit.register(_dump_recorded_namespace_reads) + + +def _append_role_ns_out(role: str | None, name: str) -> None: + # Persist immediately: worker processes are routinely torn down with + # signals that skip atexit, and the audit must survive that. + from sglang.srt.environ import envs + + out = envs.SGLANG_ROLE_NAMESPACES_OUT.get() + if not out: + return + try: + with open(out, "a") as f: + f.write(f"{role} {name}\n") + except OSError as e: + # The entry stays in the in-memory set; the exit summary still covers it. + print( + f"[role-namespaces] pid={os.getpid()} failed to append " + f"({role}, {name}) to {out!r}: {e}", + file=sys.stderr, + flush=True, + ) + + +def _record_namespace_read(role: str | None, name: str) -> None: + if (role, name) in _RECORDED_NS_READS: + return + _RECORDED_NS_READS.add((role, name)) + _append_role_ns_out(role, name) + _ensure_record_dump_registered() + + +def _dump_recorded_namespace_reads() -> None: + """Emit the record-mode audit: one line per role with the namespaces its + process actually read (multi-process runs dump once per process). The + process's own publish role is always included, so a zero-read role emits + an (empty) line rather than being indistinguishable from a process where + recording never ran.""" + by_role: dict = {} + own_role = _CONTEXT._publish_role + if own_role is not None: + by_role.setdefault(own_role, set()) + for role, name in _RECORDED_NS_READS: + if name == "-": # publish-time marker, not a namespace read + by_role.setdefault(role, set()) + continue + by_role.setdefault(role, set()).add(name) + for role in sorted(by_role, key=str): + print( + f"[role-namespaces] pid={os.getpid()} role={role} " + f"read={','.join(sorted(by_role[role]))}", + file=sys.stderr, + flush=True, + ) + + def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: """Install process-wide config for this OS process. @@ -1109,9 +1269,10 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: ``dp_controller`` / ``encoder`` / ``expert_backup`` / ``weight_cache_daemon`` / ``launcher`` / ``test``) and projects the config bags. Draft workers skip publish (they must not clobber - the target). ``role`` is provenance today — per-role namespace projection - and fail-closed enforcement is a later unit. ``hf_config`` is accepted for - forward-compat and currently unused. + the target). ``role`` is provenance, and — when ``SGLANG_ROLE_NAMESPACES`` + is ``enforce`` — the key into ``ROLE_NAMESPACE_SETS`` for fail-closed + namespace-read enforcement (``record`` audits the reads instead). + ``hf_config`` is accepted for forward-compat and currently unused. Normally one call per process, but re-publish is allowed and is **last-publish-wins** (bags re-projected, provenance reset, role @@ -1122,8 +1283,25 @@ def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext: must read ``self.server_args`` for anything engine-specific rather than the process-global bags. """ + if _ROLE_NS_MODE == "enforce" and role not in ROLE_NAMESPACE_SETS: + # Fail closed at publish time, not at the first stray read. + raise ValueError( + f"publish role {role!r} has no ROLE_NAMESPACE_SETS entry; declare " + "its namespace set (None for the full tree)." + ) _CONTEXT.set_server_args(server_args) _CONTEXT._publish_role = role + if _ROLE_NS_MODE == "record": + # The '-' marker distinguishes a zero-read role from a process where + # recording never ran (signal teardown skips atexit). + _record_namespace_read(role, "-") + print( + f"[role-namespaces] pid={os.getpid()} role={role} recording; note: " + "reads inside torch.compile-traced code are not observed — audit " + "with compilation disabled before restricting a role.", + file=sys.stderr, + flush=True, + ) return _CONTEXT diff --git a/test/registered/unit/test_runtime_context_config_bags.py b/test/registered/unit/test_runtime_context_config_bags.py index 66525ec09..4db5ddb22 100644 --- a/test/registered/unit/test_runtime_context_config_bags.py +++ b/test/registered/unit/test_runtime_context_config_bags.py @@ -7,6 +7,7 @@ by bare assignment and fail closed until published. import dataclasses import unittest +from unittest import mock from sglang.srt import runtime_context as rc from sglang.srt.arg_groups.arg_utils import NS, A @@ -127,5 +128,114 @@ class TestConfigBagTree(CustomTestCase): rc._build_config_bags(_CollisionFake()) +class TestRoleNamespaceEnforcement(CustomTestCase): + """SGLANG_ROLE_NAMESPACES: off (default) is free; record collects the + per-role read audit; enforce fails closed on reads outside the role's + declared set.""" + + def setUp(self): + rc.reset_context() + + def tearDown(self): + rc.reset_context() + + def _publish(self, role): + rc.publish(ServerArgs(model_path="dummy"), role=role) + + def test_off_mode_ignores_declared_sets(self): + self._publish("test") + with mock.patch.dict(rc.ROLE_NAMESPACE_SETS, {"test": frozenset({"serving"})}): + rc.get_exec() # off mode: no enforcement despite the narrow set + + def test_enforce_blocks_reads_outside_the_declared_set(self): + self._publish("test") + with mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"), mock.patch.dict( + rc.ROLE_NAMESPACE_SETS, {"test": frozenset({"serving", "schedule"})} + ): + rc.get_serving() + rc.get_schedule() + with self.assertRaisesRegex(ValueError, "outside the declared set"): + rc.get_exec() + + def test_enforce_full_tree_role_and_roleless_install_are_unrestricted(self): + with mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"): + self._publish("scheduler") # None in the table = full tree + rc.get_exec() + rc.get_mm() + # A direct set_server_args install is roleless; enforcement only + # keys off a recorded publish role. + rc.get_context().set_server_args(ServerArgs(model_path="dummy")) + rc.get_exec() + + def test_off_mode_bag_read_traces_under_torch_compile(self): + # config_bag runs inside compiled model forwards; the mode gate must + # stay a dead-branch-prunable check in the default "off" mode. + import torch + + self._publish("test") + + @torch.compile(fullgraph=True, backend="eager", dynamic=False) + def probe(x): + if rc.get_schedule().max_running_requests is None: + return x + 1 + return x * 2 + + self.assertEqual(probe(torch.zeros(())).item(), 1.0) + + def test_record_mode_collects_the_audit(self): + self._publish("test") + with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( + rc, "_RECORDED_NS_READS", set() + ): + rc.get_exec() + rc.get_disagg() + self.assertIn(("test", "exec"), rc._RECORDED_NS_READS) + self.assertIn(("test", "disagg"), rc._RECORDED_NS_READS) + + def test_enforce_rejects_roles_missing_from_the_table(self): + # Fail closed: an unknown/misspelled role must not silently inherit + # the full tree — rejected at publish, and defensively at read time. + with mock.patch.object(rc, "_ROLE_NS_MODE", "enforce"): + with self.assertRaisesRegex(ValueError, "no ROLE_NAMESPACE_SETS entry"): + self._publish("not_a_registered_role") + self._publish("scheduler") + with mock.patch.dict(rc.ROLE_NAMESPACE_SETS, {}, clear=True): + with self.assertRaisesRegex(ValueError, "no ROLE_NAMESPACE_SETS entry"): + rc.get_exec() + + def test_mode_env_value_is_validated(self): + self.assertEqual(rc._validated_role_ns_mode(" Enforce "), "enforce") + with self.assertRaisesRegex(ValueError, "SGLANG_ROLE_NAMESPACES"): + rc._validated_role_ns_mode("bogus") + + def test_record_mode_registers_the_exit_summary_at_publish(self): + # A role that reads no bags must still emit its audit line; the exit + # hook therefore registers at publish, not at the first read. + with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( + rc, "_RECORD_DUMP_REGISTERED", False + ): + self._publish("test") + self.assertTrue(rc._RECORD_DUMP_REGISTERED) + + def test_record_mode_bag_read_traces_under_torch_compile(self): + # Recording has side effects dynamo must never trace; the + # is_compiling() probe prunes them, keeping fullgraph capture legal + # even in record mode. + import torch + + self._publish("test") + with mock.patch.object(rc, "_ROLE_NS_MODE", "record"), mock.patch.object( + rc, "_RECORDED_NS_READS", set() + ): + + @torch.compile(fullgraph=True, backend="eager", dynamic=False) + def probe(x): + if rc.get_schedule().max_running_requests is None: + return x + 1 + return x * 2 + + self.assertEqual(probe(torch.zeros(())).item(), 1.0) + + if __name__ == "__main__": unittest.main()