runtime_context: per-role namespace enforcement behind SGLANG_ROLE_NAMESPACES (#33172)
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.
This commit is contained in:
@@ -232,6 +232,14 @@ class Envs:
|
|||||||
# must go through ServerArgs.override() (enabled by the test harness).
|
# must go through ServerArgs.override() (enabled by the test harness).
|
||||||
SGLANG_STRICT_CONFIG_MUTATION = EnvBool(False)
|
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
|
# Model & File Download
|
||||||
SGLANG_USE_MODELSCOPE = EnvBool(False)
|
SGLANG_USE_MODELSCOPE = EnvBool(False)
|
||||||
# Controls weight-file ordering for load-time I/O optimization.
|
# Controls weight-file ordering for load-time I/O optimization.
|
||||||
|
|||||||
@@ -47,6 +47,8 @@ test-only ``override(**kw)``.
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
@@ -826,8 +828,34 @@ class RuntimeContext:
|
|||||||
bags = self._config_bags
|
bags = self._config_bags
|
||||||
if not bags or name not in bags:
|
if not bags or name not in bags:
|
||||||
raise ValueError(f"config namespace {name!r} not published")
|
raise ValueError(f"config namespace {name!r} not published")
|
||||||
|
if _ROLE_NS_MODE != "off":
|
||||||
|
self._check_role_namespace(name)
|
||||||
return bags[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:
|
def override(self, source: str, **fields) -> None:
|
||||||
"""The business mutation entry: write resolved config
|
"""The business mutation entry: write resolved config
|
||||||
leaves onto the namespace bags — the single source of truth. It does
|
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")
|
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:
|
def publish(server_args, *, role: str, hf_config: Any = None) -> RuntimeContext:
|
||||||
"""Install process-wide config for this OS process.
|
"""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`` /
|
``dp_controller`` / ``encoder`` / ``expert_backup`` /
|
||||||
``weight_cache_daemon`` / ``launcher`` / ``test``) and
|
``weight_cache_daemon`` / ``launcher`` / ``test``) and
|
||||||
projects the config bags. Draft workers skip publish (they must not clobber
|
projects the config bags. Draft workers skip publish (they must not clobber
|
||||||
the target). ``role`` is provenance today — per-role namespace projection
|
the target). ``role`` is provenance, and — when ``SGLANG_ROLE_NAMESPACES``
|
||||||
and fail-closed enforcement is a later unit. ``hf_config`` is accepted for
|
is ``enforce`` — the key into ``ROLE_NAMESPACE_SETS`` for fail-closed
|
||||||
forward-compat and currently unused.
|
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
|
Normally one call per process, but re-publish is allowed and is
|
||||||
**last-publish-wins** (bags re-projected, provenance reset, role
|
**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
|
must read ``self.server_args`` for anything engine-specific rather than
|
||||||
the process-global bags.
|
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.set_server_args(server_args)
|
||||||
_CONTEXT._publish_role = role
|
_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
|
return _CONTEXT
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ by bare assignment and fail closed until published.
|
|||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
import unittest
|
import unittest
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
from sglang.srt import runtime_context as rc
|
from sglang.srt import runtime_context as rc
|
||||||
from sglang.srt.arg_groups.arg_utils import NS, A
|
from sglang.srt.arg_groups.arg_utils import NS, A
|
||||||
@@ -127,5 +128,114 @@ class TestConfigBagTree(CustomTestCase):
|
|||||||
rc._build_config_bags(_CollisionFake())
|
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__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user