[SRT] Clean up no-op compiler pass, dead helpers, and migration tests (#39295)
This commit is contained in:
Xiaoyu Zhang
@@ -73,11 +73,9 @@ with what the operator typed, not with what resolution decided.**
|
||||
`TokenizerManager.record_config_updates(source, **fields)`, a named wrapper
|
||||
over `get_context().override`. One process keeps one log: the request dumps
|
||||
ship `get_context().overrides_log()`, and `config_value(name)` /
|
||||
`resolved_config_dict(base)` answer from the bags. The exposure ratchet
|
||||
resolves the wrapper, so a field recorded through it joins the post-publish
|
||||
override surface exactly like a direct `override` and needs the same ordering
|
||||
judgment against any supplied-instance read of it
|
||||
(`test_supplied_instance_exposure_ratchet.py`).
|
||||
`resolved_config_dict(base)` answer from the bags. Fields recorded through
|
||||
the wrapper follow the same ordering rules as a direct `override`: a later
|
||||
read must observe the updated bag, not the startup record.
|
||||
- **`model_path` and `served_model_name` are answered off the manager.** Both are
|
||||
`NS` leaves and `override` accepts them, but the tokenizer-side weight reload
|
||||
records only `load_format` and writes the two path fields as `TokenizerManager`
|
||||
@@ -154,8 +152,7 @@ bag to override at all.
|
||||
**retracted** — owner ruling (2026-08-15): a process holds at most one live
|
||||
config at a time (concurrent multi-Engine is unsupported; sequential rebuild
|
||||
stays legal, unit tests rely on it). Nothing in those files reads the instance
|
||||
any more -- the exposure ratchet's pin set is empty, so the next such read is a
|
||||
new entry that has to argue for itself. What
|
||||
any more; review new instance reads against the raw-input contract. What
|
||||
genuinely stays per-instance is what differs per *worker* within one engine:
|
||||
`base_gpu_id` travels as a constructor argument (`MMEncoder(gpu_id=...)`;
|
||||
`BaseMultimodalProcessor._fast_image_processor_device` is the shape to copy).
|
||||
@@ -172,12 +169,9 @@ bag to override at all.
|
||||
attention pair and the encode-server `gpu_id` above are both this). The per-instance
|
||||
boundaries above are **not** exempt from this unless-clause (the multi-Engine
|
||||
exemption is retracted); each one gets its own disposition.
|
||||
`test_supplied_instance_exposure_ratchet.py`
|
||||
pins that set (empty today) — three spellings of the read: `server_args.field`,
|
||||
literal-name `getattr(server_args, "field", default)`, and the parked form
|
||||
(`self.x = server_args` in a method that takes the parameter, read as
|
||||
`self.x.field` anywhere in the class) — and fails on a new one, so the
|
||||
disposition gets picked when the read is written. Two shapes stay parameter-form on purpose: a helper the
|
||||
Check direct attributes, `getattr`, and records stored on `self`; validate
|
||||
the resolved value and any later overrides in behavior tests.
|
||||
Two shapes stay parameter-form on purpose: a helper the
|
||||
*resolution pipeline* calls with a `resolved_view` (its parameter happens to be
|
||||
named `server_args`), and a factory whose contract is "build X from the record
|
||||
you are handed" (`create_kt_config_from_server_args`, `DllmConfig.from_server_args`).
|
||||
@@ -402,12 +396,10 @@ through a view instead:
|
||||
- `resolved_view(server_args)` — snapshots the overlay when built, which is what
|
||||
a post-process pass wants: it reads the state at *its* slot.
|
||||
|
||||
`test_resolution_reads_the_declarations` pins direct field reads at zero over the
|
||||
two scopes it can derive exactly (every `arg_groups` function taking a config,
|
||||
every `ServerArgs` handler the dispatcher reaches). Readers the pipeline calls
|
||||
from elsewhere (`ModelConfig`, the platform defaults, the spec-algo hook) have
|
||||
moved to the view as well — a field read there is the same bug, just one the
|
||||
derivation cannot enumerate.
|
||||
Resolution hooks and the helpers they call (`ModelConfig`, platform defaults,
|
||||
the spec-algo hook) must read through these views too. Keep coverage in
|
||||
`test_resolution_declarations.py`, `test_resolution_is_reproducible.py`, and
|
||||
`test_record_holds_the_raw_input.py` focused on the values callers observe.
|
||||
|
||||
One consequence worth knowing: because the fields are the raw input, resolving a
|
||||
bare `dataclasses.replace` copy lands in the same place as the parent — the
|
||||
|
||||
@@ -174,11 +174,8 @@ def handle_attention_backend_compatibility(server_args: Any):
|
||||
# AMD platforms backends
|
||||
if resolved_view(server_args).attention_backend == "aiter":
|
||||
if model_config.context_len > 8192:
|
||||
# The record, via the input snapshot rather than the field: a
|
||||
# hook may not read a field off the record (the guard in
|
||||
# `test_resolution_reads_the_declarations.py`), and what this
|
||||
# needs is the input anyway -- whether the operator asked for a
|
||||
# memory fraction, not the value in effect.
|
||||
# Check whether the operator supplied a memory fraction using
|
||||
# the raw input snapshot; resolution may have filled the value in.
|
||||
explicit_mem_fraction = (
|
||||
getattr(server_args, "_raw_input", None) or {}
|
||||
).get("mem_fraction_static") is not None
|
||||
|
||||
@@ -128,17 +128,8 @@ def run_post_process_pass(server_args: Any, fn: Callable[..., dict]) -> None:
|
||||
f"got {type(declared).__name__}"
|
||||
)
|
||||
if declared:
|
||||
# Refused only once there is something to record. A pass that declares
|
||||
# nothing is a validation, and `check_server_args` runs those again on
|
||||
# a rebuild: `Engine(server_args=sa)` after `Engine.shutdown()` hands
|
||||
# back the same instance while the context still holds it, and
|
||||
# refusing on identity alone would fail that launch.
|
||||
# Only a non-empty return is a declaration. An empty one is a
|
||||
# validation and may run on the published instance -- see above -- so it
|
||||
# must not reach the guard in `declare_resolution`.
|
||||
if declared:
|
||||
declare_resolution(server_args, fn.__qualname__, **declared)
|
||||
validate_declarations(server_args, [(fn.__qualname__, dict(declared))])
|
||||
declare_resolution(server_args, fn.__qualname__, **declared)
|
||||
validate_declarations(server_args, [(fn.__qualname__, dict(declared))])
|
||||
|
||||
|
||||
def declare_resolution(server_args: Any, source: str, **fields: Any) -> None:
|
||||
@@ -147,8 +138,7 @@ def declare_resolution(server_args: Any, source: str, **fields: Any) -> None:
|
||||
The stash *is* the resolution result: the bags are projected from it,
|
||||
`resolution_result` answers from it, and no field is written. A resolver
|
||||
reading a field another resolver may have decided must read `resolving_view`
|
||||
(or `resolved_view(server_args)`), which
|
||||
`test_resolution_reads_the_declarations` pins.
|
||||
(or `resolved_view(server_args)`).
|
||||
|
||||
Every declaration goes through here, whenever it is made: inside
|
||||
``__post_init__``, at launcher stage (LoRA normalization, the auto-detected
|
||||
|
||||
@@ -396,7 +396,6 @@ class SGLangBackend:
|
||||
self.compile_config = config
|
||||
|
||||
def configure_post_pass(self):
|
||||
self.post_grad_pass_manager.configure()
|
||||
self.inductor_config["post_grad_custom_post_pass"] = self.post_grad_pass_manager
|
||||
|
||||
def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
|
||||
|
||||
@@ -1,136 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fix_functionalization.py
|
||||
|
||||
import logging
|
||||
import operator
|
||||
from collections.abc import Iterable
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
from torch._higher_order_ops.auto_functionalize import auto_functionalized
|
||||
|
||||
from sglang.srt.compilation.fx_utils import is_func
|
||||
from sglang.srt.compilation.inductor_pass import SGLangInductorPass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FixFunctionalizationPass(SGLangInductorPass):
|
||||
"""
|
||||
This pass defunctionalizes certain nodes to avoid redundant tensor copies.
|
||||
After this pass, DCE (dead-code elimination) should never be run,
|
||||
as de-functionalized nodes may appear as dead code.
|
||||
|
||||
To add new nodes to defunctionalize, add to the if-elif chain in __call__.
|
||||
"""
|
||||
|
||||
def __call__(self, graph: torch.fx.Graph):
|
||||
self.begin()
|
||||
self.dump_graph(graph, "before_fix_functionalization")
|
||||
|
||||
self.nodes_to_remove: list[torch.fx.Node] = []
|
||||
count = 0
|
||||
for node in graph.nodes:
|
||||
if not is_func(node, auto_functionalized):
|
||||
continue # Avoid deep if-elif nesting
|
||||
count += 1
|
||||
|
||||
self.dump_graph(graph, "before_fix_functionalization_cleanup")
|
||||
|
||||
# Remove the nodes all at once
|
||||
count_removed = len(self.nodes_to_remove)
|
||||
for node in self.nodes_to_remove:
|
||||
graph.erase_node(node)
|
||||
|
||||
logger.debug(
|
||||
"De-functionalized %s nodes, removed %s nodes", count, count_removed
|
||||
)
|
||||
self.dump_graph(graph, "after_fix_functionalization")
|
||||
self.end_and_log()
|
||||
|
||||
def _remove(self, node_or_nodes: Union[torch.fx.Node, Iterable[torch.fx.Node]]):
|
||||
"""
|
||||
Stage a node (or nodes) for removal at the end of the pass.
|
||||
"""
|
||||
if isinstance(node_or_nodes, torch.fx.Node):
|
||||
self.nodes_to_remove.append(node_or_nodes)
|
||||
else:
|
||||
self.nodes_to_remove.extend(node_or_nodes)
|
||||
|
||||
def defunctionalize(
|
||||
self,
|
||||
graph: torch.fx.Graph,
|
||||
node: torch.fx.Node,
|
||||
mutated_args: dict[int, Union[torch.fx.Node, str]],
|
||||
args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None,
|
||||
):
|
||||
"""
|
||||
De-functionalize a node by replacing it with a call to the original.
|
||||
It also replaces the getitem users with the mutated arguments.
|
||||
See replace_users_with_mutated_args and insert_defunctionalized.
|
||||
"""
|
||||
self.replace_users_with_mutated_args(node, mutated_args)
|
||||
self.insert_defunctionalized(graph, node, args=args)
|
||||
self._remove(node)
|
||||
|
||||
def replace_users_with_mutated_args(
|
||||
self, node: torch.fx.Node, mutated_args: dict[int, Union[torch.fx.Node, str]]
|
||||
):
|
||||
"""
|
||||
Replace all getitem users of the auto-functionalized node with the
|
||||
mutated arguments.
|
||||
:param node: The auto-functionalized node
|
||||
:param mutated_args: The mutated arguments, indexed by getitem index.
|
||||
If the value of an arg is a string, `node.kwargs[arg]` is used.
|
||||
"""
|
||||
for idx, user in self.getitem_users(node).items():
|
||||
arg = mutated_args[idx]
|
||||
arg = node.kwargs[arg] if isinstance(arg, str) else arg
|
||||
user.replace_all_uses_with(arg)
|
||||
self._remove(user)
|
||||
|
||||
def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]:
|
||||
"""
|
||||
Returns the operator.getitem users of the auto-functionalized node,
|
||||
indexed by the index they are getting.
|
||||
"""
|
||||
users = {}
|
||||
for user in node.users:
|
||||
if is_func(user, operator.getitem):
|
||||
idx = user.args[1]
|
||||
users[idx] = user
|
||||
return users
|
||||
|
||||
def insert_defunctionalized(
|
||||
self,
|
||||
graph: torch.fx.Graph,
|
||||
node: torch.fx.Node,
|
||||
args: Optional[tuple[Union[torch.fx.Node, str], ...]] = None,
|
||||
):
|
||||
"""
|
||||
Insert a new defunctionalized node into the graph before node.
|
||||
If one of the kwargs is 'out', provide args directly,
|
||||
as node.kwargs cannot be used.
|
||||
See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351
|
||||
|
||||
:param graph: Graph to insert the defunctionalized node into
|
||||
:param node: The auto-functionalized node to defunctionalize
|
||||
:param args: If we cannot use kwargs, specify args directly.
|
||||
If an arg is a string, `node.kwargs[arg]` is used.
|
||||
""" # noqa: E501
|
||||
assert is_func(node, auto_functionalized), (
|
||||
f"node must be auto-functionalized, is {node} instead"
|
||||
)
|
||||
|
||||
# Create a new call to the original function
|
||||
with graph.inserting_before(node):
|
||||
function = node.args[0]
|
||||
if args is None:
|
||||
graph.call_function(function, kwargs=node.kwargs)
|
||||
else:
|
||||
# Args passed as strings refer to items in node.kwargs
|
||||
args = tuple(
|
||||
node.kwargs[arg] if isinstance(arg, str) else arg for arg in args
|
||||
)
|
||||
graph.call_function(function, args=args)
|
||||
@@ -1,85 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/fx_utils.py
|
||||
|
||||
import operator
|
||||
from collections.abc import Iterable, Iterator
|
||||
from typing import Optional
|
||||
|
||||
from torch import fx
|
||||
from torch._higher_order_ops.auto_functionalize import auto_functionalized
|
||||
from torch._ops import OpOverload
|
||||
|
||||
|
||||
def is_func(node: fx.Node, target) -> bool:
|
||||
return node.op == "call_function" and node.target == target
|
||||
|
||||
|
||||
def is_auto_func(node: fx.Node, op: OpOverload) -> bool:
|
||||
return is_func(node, auto_functionalized) and node.args[0] == op
|
||||
|
||||
|
||||
# Returns the first specified node with the given op (if it exists)
|
||||
def find_specified_fn_maybe(
|
||||
nodes: Iterable[fx.Node], op: OpOverload
|
||||
) -> Optional[fx.Node]:
|
||||
for node in nodes:
|
||||
if node.target == op:
|
||||
return node
|
||||
return None
|
||||
|
||||
|
||||
# Returns the first specified node with the given op
|
||||
def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
|
||||
node = find_specified_fn_maybe(nodes, op)
|
||||
assert node is not None, f"Could not find {op} in nodes {nodes}"
|
||||
return node
|
||||
|
||||
|
||||
# Returns the first auto_functionalized node with the given op (if it exists)
|
||||
def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> Optional[fx.Node]:
|
||||
for node in nodes:
|
||||
if is_func(node, auto_functionalized) and node.args[0] == op: # noqa
|
||||
return node
|
||||
return None
|
||||
|
||||
|
||||
# Returns the first auto_functionalized node with the given op
|
||||
def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
|
||||
node = find_auto_fn_maybe(nodes, op)
|
||||
assert node is not None, f"Could not find {op} in nodes {nodes}"
|
||||
return node
|
||||
|
||||
|
||||
# Returns the getitem node that extracts the idx-th element from node
|
||||
# (if it exists)
|
||||
def find_getitem_maybe(node: fx.Node, idx: int) -> Optional[fx.Node]:
|
||||
for user in node.users:
|
||||
if is_func(user, operator.getitem) and user.args[1] == idx:
|
||||
return user
|
||||
return None
|
||||
|
||||
|
||||
# Returns the getitem node that extracts the idx-th element from node
|
||||
def find_getitem(node: fx.Node, idx: int) -> fx.Node:
|
||||
ret = find_getitem_maybe(node, idx)
|
||||
assert ret is not None, f"Could not find getitem {idx} in node {node}"
|
||||
return ret
|
||||
|
||||
|
||||
# An auto-functionalization-aware utility for finding nodes with a specific op
|
||||
def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]:
|
||||
if not op._schema.is_mutable:
|
||||
yield from graph.find_nodes(op="call_function", target=op)
|
||||
|
||||
for n in graph.find_nodes(op="call_function", target=auto_functionalized):
|
||||
if n.args[0] == op:
|
||||
yield n
|
||||
|
||||
|
||||
# Asserts that the node only has one user and returns it
|
||||
# Even if a node has only 1 user, it might share storage with another node,
|
||||
# which might need to be taken into account.
|
||||
def get_only_user(node: fx.Node) -> fx.Node:
|
||||
assert len(node.users) == 1
|
||||
return next(iter(node.users))
|
||||
@@ -9,10 +9,9 @@ import logging
|
||||
import time
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Callable, Optional, Union
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
import torch
|
||||
from torch import fx
|
||||
from torch._dynamo.utils import lazy_format_graph_code
|
||||
from torch._inductor.custom_graph_pass import CustomGraphPass
|
||||
|
||||
@@ -93,25 +92,6 @@ class InductorPass(CustomGraphPass):
|
||||
return True
|
||||
|
||||
|
||||
class CallableInductorPass(InductorPass):
|
||||
"""
|
||||
This class is a wrapper for a callable that automatically provides an
|
||||
implementation of the UUID.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, callable: Callable[[fx.Graph], None], uuid: Optional[Any] = None
|
||||
):
|
||||
self.callable = callable
|
||||
self._uuid = self.hash_source(callable) if uuid is None else uuid
|
||||
|
||||
def __call__(self, graph: torch.fx.Graph):
|
||||
self.callable(graph)
|
||||
|
||||
def uuid(self) -> Any:
|
||||
return self._uuid
|
||||
|
||||
|
||||
class SGLangInductorPass(InductorPass):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -128,12 +108,3 @@ class SGLangInductorPass(InductorPass):
|
||||
self._end_time = time.perf_counter_ns()
|
||||
duration_ms = float(self._end_time - self._start_time) / 1.0e6
|
||||
logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)
|
||||
|
||||
|
||||
class PrinterInductorPass(SGLangInductorPass):
|
||||
def __init__(self, name: str):
|
||||
super().__init__()
|
||||
self.name = name
|
||||
|
||||
def __call__(self, graph: torch.fx.Graph):
|
||||
self.dump_graph(graph, self.name)
|
||||
|
||||
@@ -2,11 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0/vllm/compilation/pass_manager.py
|
||||
|
||||
import logging
|
||||
|
||||
from torch import fx as fx
|
||||
|
||||
from sglang.srt.compilation.fix_functionalization import FixFunctionalizationPass
|
||||
from sglang.srt.compilation.inductor_pass import (
|
||||
CustomGraphPass,
|
||||
InductorPass,
|
||||
@@ -14,22 +11,11 @@ from sglang.srt.compilation.inductor_pass import (
|
||||
get_pass_context,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostGradPassManager(CustomGraphPass):
|
||||
"""
|
||||
The pass manager for post-grad passes.
|
||||
It handles configuration, adding custom passes, and running passes.
|
||||
It supports uuid for the Inductor code cache. That includes torch<2.6
|
||||
support using pickling (in .inductor_pass.CustomGraphPass).
|
||||
"""Run post-grad passes in insertion order for the current runtime shape.
|
||||
|
||||
The order of the post-grad post-passes is:
|
||||
1. passes (constructor parameter)
|
||||
2. default passes (NoopEliminationPass, FusionPass)
|
||||
3. config["post_grad_custom_post_pass"] (if it exists)
|
||||
4. fix_functionalization
|
||||
This way, all passes operate on a functionalized graph.
|
||||
The ordered pass UUIDs identify the manager in Inductor's code cache.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
@@ -41,15 +27,6 @@ class PostGradPassManager(CustomGraphPass):
|
||||
if pass_.is_applicable_for_shape(shape):
|
||||
pass_(graph)
|
||||
|
||||
# always run fix_functionalization last
|
||||
self.fix_functionalization(graph)
|
||||
|
||||
def configure(
|
||||
self,
|
||||
):
|
||||
self.pass_config = dict()
|
||||
self.fix_functionalization = FixFunctionalizationPass()
|
||||
|
||||
def add(self, pass_: InductorPass):
|
||||
assert isinstance(pass_, InductorPass)
|
||||
self.passes.append(pass_)
|
||||
@@ -58,11 +35,6 @@ class PostGradPassManager(CustomGraphPass):
|
||||
"""
|
||||
The PostGradPassManager is set as a custom pass in the Inductor and
|
||||
affects compilation caching. Its uuid depends on the UUIDs of all
|
||||
dependent passes and the pass config. See InductorPass for more info.
|
||||
dependent passes. See InductorPass for more info.
|
||||
"""
|
||||
pass_manager_uuid = "fshdakhsa"
|
||||
state = {"pass_config": pass_manager_uuid, "passes": []}
|
||||
for pass_ in self.passes:
|
||||
state["passes"].append(pass_.uuid())
|
||||
state["passes"].append(self.fix_functionalization.uuid())
|
||||
return InductorPass.hash_dict(state)
|
||||
return InductorPass.hash_dict({"passes": [p.uuid() for p in self.passes]})
|
||||
|
||||
@@ -756,66 +756,6 @@ def general_mm_embed_routine(
|
||||
return hidden_states
|
||||
|
||||
|
||||
def get_multimodal_data_bounds(
|
||||
input_ids: torch.Tensor, pad_values: List[int], token_pairs: List[Tuple[int, int]]
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Returns a tensor indicating the bounds of multimodal data (images, video, audio, etc.)
|
||||
|
||||
Returns:
|
||||
[bounds_count, 2]
|
||||
"""
|
||||
# All the multimodal data in the batch should share the same special bound token ids.
|
||||
start_tokens = {s for s, _e in token_pairs}
|
||||
end_tokens = {e for _s, e in token_pairs}
|
||||
|
||||
assert all(isinstance(t, int) for t in start_tokens)
|
||||
assert all(isinstance(t, int) for t in end_tokens)
|
||||
|
||||
start_cond = torch.isin(
|
||||
input_ids, torch.as_tensor(start_tokens, device=input_ids.device)
|
||||
)
|
||||
end_cond = torch.isin(
|
||||
input_ids, torch.as_tensor(end_tokens, device=input_ids.device)
|
||||
)
|
||||
|
||||
(data_start_tokens,) = torch.where(start_cond)
|
||||
(data_end_tokens,) = torch.where(end_cond)
|
||||
|
||||
data_start_tokens_cpu = data_start_tokens.cpu().tolist()
|
||||
data_end_tokens_cpu = data_end_tokens.cpu().tolist()
|
||||
|
||||
# the im_start_id sometimes can be cached as prefix, but it is needed for the embedding of the multimodal data
|
||||
if len(data_start_tokens_cpu) != len(data_end_tokens_cpu):
|
||||
if (
|
||||
len(data_start_tokens_cpu) + 1 == len(data_end_tokens_cpu)
|
||||
and input_ids[0].item() in pad_values
|
||||
and data_end_tokens_cpu
|
||||
and data_start_tokens_cpu
|
||||
and data_end_tokens_cpu[0] < data_start_tokens_cpu[0]
|
||||
):
|
||||
data_start_tokens_cpu.insert(0, 0)
|
||||
valid_mm_data_nums = min(len(data_start_tokens_cpu), len(data_end_tokens_cpu))
|
||||
|
||||
if valid_mm_data_nums == 0:
|
||||
return torch.zeros((0, 2), device=input_ids.device)
|
||||
|
||||
# Filter out pairs where start_token >= end_token
|
||||
valid_pairs = []
|
||||
for i in range(valid_mm_data_nums):
|
||||
start_token = data_start_tokens_cpu[i]
|
||||
end_token = data_end_tokens_cpu[i]
|
||||
if start_token < end_token:
|
||||
valid_pairs.append((start_token + 1, end_token - 1))
|
||||
|
||||
if not valid_pairs:
|
||||
return torch.zeros((0, 2), device=input_ids.device)
|
||||
|
||||
# Convert valid pairs to tensor
|
||||
valid_pairs_tensor = torch.as_tensor(valid_pairs, device=input_ids.device)
|
||||
return valid_pairs_tensor
|
||||
|
||||
|
||||
def data_hash(data) -> int:
|
||||
hash_bytes = hashlib.sha256(data).digest()[:8]
|
||||
return int.from_bytes(hash_bytes, byteorder="big", signed=False)
|
||||
|
||||
@@ -1368,24 +1368,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
||||
f"{token_id}; valid range is [0, {vocab_size})."
|
||||
)
|
||||
|
||||
def _validate_input_ids_in_vocab(
|
||||
self, input_ids: Union[List[int], List[List[int]]], vocab_size: int
|
||||
) -> None:
|
||||
# Handle both single sequence and batch of sequences
|
||||
if isinstance(input_ids[0], list):
|
||||
# Batch of sequences
|
||||
for seq in input_ids:
|
||||
if any(id >= vocab_size for id in seq):
|
||||
raise ValueError(
|
||||
f"The input_ids {seq} contains values greater than the vocab size ({vocab_size})."
|
||||
)
|
||||
else:
|
||||
# Single sequence
|
||||
if any(id >= vocab_size for id in input_ids):
|
||||
raise ValueError(
|
||||
f"The input_ids {input_ids} contains values greater than the vocab size ({vocab_size})."
|
||||
)
|
||||
|
||||
def _create_tokenized_object(
|
||||
self,
|
||||
obj: Union[GenerateReqInput, EmbeddingReqInput],
|
||||
|
||||
@@ -1,396 +0,0 @@
|
||||
"""Refcounted content-addressed page cache over the buffer-mode host pool.
|
||||
|
||||
NOT WIRED YET: staged spans register as ``(pool, page_hash) -> (slots,
|
||||
refcount)`` so prefetches can be served zero-copy from local staging
|
||||
(write-around / promote-on-read retention, zero-ref LRU reclaim); keys are
|
||||
the content-chained page hashes, so entries survive node deletion, splits,
|
||||
and recompute. Counterpart of ``StorageExistenceCache`` (beliefs about
|
||||
STORAGE, dedupes writes); this tracks LOCAL HOST RAM and dedupes loads.
|
||||
|
||||
TP determinism: replicas must stay identical across attention ranks — the
|
||||
cache feeds scheduler-visible structure, so divergence is a collective
|
||||
hang, not a soft miss. Preconditions when wiring: (1) mutate only on the
|
||||
scheduler thread at lockstep points; (2) rank-reduce any fold anchored by a
|
||||
per-rank storage outcome (hit count, revoke) before it picks a mutation;
|
||||
(3) controller queues stay FIFO and single-threaded so MIN-count drains
|
||||
process the same prefix on every rank.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
from typing import Callable, Optional, Sequence
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.hicache_storage import (
|
||||
PoolHitPolicy,
|
||||
PoolName,
|
||||
PoolTransfer,
|
||||
)
|
||||
|
||||
|
||||
class _PageRef:
|
||||
"""One cached page: host slot span, reader refcount, and whether to
|
||||
retain at refs==0 (write-around: write staging frees at its storage ack
|
||||
unless a read promoted it). ``first_slot`` mirrors ``slots[0]`` as a
|
||||
plain int because per-page tensor-scalar reads are too slow in
|
||||
``release``."""
|
||||
|
||||
__slots__ = ("slots", "first_slot", "refs", "retain")
|
||||
|
||||
def __init__(
|
||||
self, slots: torch.Tensor, first_slot: int, refs: int = 1, retain: bool = True
|
||||
):
|
||||
self.slots = slots
|
||||
self.first_slot = first_slot
|
||||
self.refs = refs
|
||||
self.retain = retain
|
||||
|
||||
|
||||
class BufferPageCache:
|
||||
def __init__(self) -> None:
|
||||
# (pool, page_hash) -> _PageRef; slots stay allocated in the host
|
||||
# pool for as long as the entry exists.
|
||||
self._entries: dict[tuple[str, str], _PageRef] = {}
|
||||
# Per-pool zero-ref LRU (head = coldest): reclaim victims.
|
||||
self._zero_ref: dict[str, OrderedDict[str, None]] = {}
|
||||
# Per-pool slot tokens held by the cache (refed + zero-ref).
|
||||
self._held_tokens: dict[str, int] = {}
|
||||
# Per-pool slot tokens at refs=0 (reclaimable under pressure).
|
||||
self._zero_ref_tokens: dict[str, int] = {}
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._entries)
|
||||
|
||||
def num_zero_ref_pages(self) -> int:
|
||||
return sum(len(lru) for lru in self._zero_ref.values())
|
||||
|
||||
def held_tokens(self, pool: str) -> int:
|
||||
return self._held_tokens.get(pool, 0)
|
||||
|
||||
def zero_ref_tokens(self, pool: str) -> int:
|
||||
"""Slot tokens reclaimable right now (zero-ref cached pages).
|
||||
Occupancy/rate-limit gates must treat these as free-able, not used:
|
||||
a pool full of zero-ref cache is one reclaim away from empty."""
|
||||
return self._zero_ref_tokens.get(pool, 0)
|
||||
|
||||
def register(
|
||||
self,
|
||||
pool: str,
|
||||
hashes: Sequence[str],
|
||||
host_indices: torch.Tensor,
|
||||
page_size: int,
|
||||
retain: bool = True,
|
||||
) -> int:
|
||||
"""Cache a staged span, one entry per page, refs=1 (the staging op);
|
||||
returns the number of pages newly cached. ``retain=False`` =
|
||||
write-around (freed at last ref unless a read hit promotes it); a
|
||||
duplicate hash keeps the existing entry and the newcomer's slots
|
||||
stay op-owned for a raw free at release."""
|
||||
assert len(host_indices) == len(hashes) * page_size
|
||||
registered = 0
|
||||
entries = self._entries
|
||||
# One batched read of the page-boundary slot ids (see _PageRef).
|
||||
first_slots = host_indices[::page_size].tolist()
|
||||
for i, page_hash in enumerate(hashes):
|
||||
key = (pool, page_hash)
|
||||
existing = entries.get(key)
|
||||
if existing is not None:
|
||||
existing.retain = existing.retain or retain
|
||||
continue
|
||||
entries[key] = _PageRef(
|
||||
host_indices[i * page_size : (i + 1) * page_size],
|
||||
first_slots[i],
|
||||
retain=retain,
|
||||
)
|
||||
registered += 1
|
||||
if registered:
|
||||
self._held_tokens[pool] = (
|
||||
self._held_tokens.get(pool, 0) + registered * page_size
|
||||
)
|
||||
return registered
|
||||
|
||||
def contains(self, pool: str, page_hash: str) -> bool:
|
||||
"""Non-mutating presence probe (no LRU touch)."""
|
||||
return (pool, page_hash) in self._entries
|
||||
|
||||
def peek_run_len(self, pool: str, hashes: Sequence[str]) -> int:
|
||||
"""Length of the leading run of cached pages. Non-mutating."""
|
||||
entries = self._entries
|
||||
run = 0
|
||||
for page_hash in hashes:
|
||||
if (pool, page_hash) not in entries:
|
||||
break
|
||||
run += 1
|
||||
return run
|
||||
|
||||
def acquire(self, pool: str, hashes: Sequence[str]) -> Optional[torch.Tensor]:
|
||||
"""refs++ on every page and return the gathered slot tensor (pages
|
||||
expanded to token slots, in page order). All-or-nothing: returns
|
||||
None without mutating if any page is missing."""
|
||||
entries = self._entries
|
||||
refs = []
|
||||
for page_hash in hashes:
|
||||
entry = entries.get((pool, page_hash))
|
||||
if entry is None:
|
||||
return None
|
||||
refs.append(entry)
|
||||
zero_ref = self._zero_ref.get(pool)
|
||||
for page_hash, entry in zip(hashes, refs):
|
||||
if entry.refs == 0 and zero_ref is not None:
|
||||
if zero_ref.pop(page_hash, None) is not None:
|
||||
self._zero_ref_tokens[pool] -= len(entry.slots)
|
||||
entry.refs += 1
|
||||
# Read demand proven: promote write-around pages to retained.
|
||||
entry.retain = True
|
||||
return torch.cat([entry.slots for entry in refs])
|
||||
|
||||
def release(
|
||||
self,
|
||||
pool: str,
|
||||
hashes: Sequence[str],
|
||||
host_indices: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> Optional[torch.Tensor]:
|
||||
"""Drop one ref per page of a span: at refs==0 retained pages move
|
||||
to the zero-ref LRU tail while write-around pages return their
|
||||
slots for an immediate free. Duplicate-staging slots (canonical
|
||||
entry lives elsewhere) are returned for a raw free without touching
|
||||
the canonical refcount."""
|
||||
assert len(host_indices) == len(hashes) * page_size
|
||||
leftover: list[torch.Tensor] = []
|
||||
entries = self._entries
|
||||
# One batched read of the page-boundary slot ids (see _PageRef).
|
||||
first_slots = host_indices[::page_size].tolist()
|
||||
for i, page_hash in enumerate(hashes):
|
||||
entry = entries.get((pool, page_hash))
|
||||
if entry is None or entry.first_slot != first_slots[i]:
|
||||
leftover.append(host_indices[i * page_size : (i + 1) * page_size])
|
||||
continue
|
||||
assert entry.refs > 0, "release without a matching acquire/register"
|
||||
entry.refs -= 1
|
||||
if entry.refs == 0:
|
||||
if entry.retain:
|
||||
self._zero_ref.setdefault(pool, OrderedDict())[page_hash] = None
|
||||
self._zero_ref_tokens[pool] = self._zero_ref_tokens.get(
|
||||
pool, 0
|
||||
) + len(entry.slots)
|
||||
else:
|
||||
del entries[(pool, page_hash)]
|
||||
self._held_tokens[pool] -= len(entry.slots)
|
||||
leftover.append(entry.slots)
|
||||
if not leftover:
|
||||
return None
|
||||
return torch.cat(leftover)
|
||||
|
||||
def reclaim(
|
||||
self,
|
||||
pool: str,
|
||||
need_tokens: int,
|
||||
free: Callable[[torch.Tensor], int],
|
||||
) -> int:
|
||||
"""Pop zero-ref LRU heads, free their slots back to the host pool,
|
||||
and drop the entries. Called under allocation pressure only. Returns
|
||||
the number of slot tokens freed (may undershoot when everything
|
||||
left is refed)."""
|
||||
zero_ref = self._zero_ref.get(pool)
|
||||
if not zero_ref or need_tokens <= 0:
|
||||
return 0
|
||||
freed = 0
|
||||
batch: list[torch.Tensor] = []
|
||||
while zero_ref and freed < need_tokens:
|
||||
page_hash, _ = zero_ref.popitem(last=False)
|
||||
entry = self._entries.pop((pool, page_hash))
|
||||
batch.append(entry.slots)
|
||||
freed += len(entry.slots)
|
||||
if batch:
|
||||
free(torch.cat(batch))
|
||||
self._held_tokens[pool] -= freed
|
||||
self._zero_ref_tokens[pool] -= freed
|
||||
return freed
|
||||
|
||||
|
||||
class BufferPageCacheOps:
|
||||
"""Pool-facing operations over a :class:`BufferPageCache`: span/hold
|
||||
registration and release keyed the way the storage write keys them,
|
||||
pressure reclaim, and the SWA-folded continuation fold. The caller owns
|
||||
the collectives — rank-reduce any fold anchored by a per-rank storage
|
||||
outcome before acting on it (see the module docstring)."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
page_cache: BufferPageCache,
|
||||
mem_pool_host,
|
||||
sw_window_pages_fn: Callable[[], int],
|
||||
):
|
||||
# Rebound by the owner when the structure is recreated (reset).
|
||||
self.page_cache = page_cache
|
||||
self._mem_pool_host = mem_pool_host
|
||||
# SWA window in KV pages when SWA stages through a host pool
|
||||
# (0 = KV-only: no trailing window in the fold).
|
||||
self._sw_window_pages_fn = sw_window_pages_fn
|
||||
|
||||
def aux_window_keys(
|
||||
self, hash_values: list[str], transfer: PoolTransfer
|
||||
) -> Optional[list[str]]:
|
||||
"""Trailing KV page hashes keying an aux transfer's staged window
|
||||
(one key per aux-pool page), recomputed from the rank-synced span
|
||||
hashes so registration and release always agree across ranks."""
|
||||
if transfer.host_indices is None or transfer.host_indices.numel() == 0:
|
||||
return None
|
||||
if transfer.indices_from_pool is not None:
|
||||
return None # sidecar rides another pool's slots; nothing to key
|
||||
entry = self._mem_pool_host.entry_map.get(transfer.name)
|
||||
if entry is None:
|
||||
return None
|
||||
pool_page_size = entry.host_pool.page_size
|
||||
num_keys = len(transfer.host_indices) // pool_page_size
|
||||
if num_keys == 0 or num_keys > len(hash_values):
|
||||
return None
|
||||
return hash_values[-num_keys:]
|
||||
|
||||
def register_span(
|
||||
self,
|
||||
pool: PoolName,
|
||||
hashes: list[str],
|
||||
host_indices: torch.Tensor,
|
||||
retain: bool = True,
|
||||
) -> None:
|
||||
"""Cache a page-aligned staged span (refs=1 for the staging op)."""
|
||||
if not hashes:
|
||||
return
|
||||
entry = self._mem_pool_host.entry_map.get(pool)
|
||||
if entry is None:
|
||||
return
|
||||
self.page_cache.register(
|
||||
pool, hashes, host_indices, entry.host_pool.page_size, retain=retain
|
||||
)
|
||||
|
||||
def release_span(
|
||||
self,
|
||||
pool: PoolName,
|
||||
hashes: list[str],
|
||||
host_indices: torch.Tensor,
|
||||
) -> None:
|
||||
"""Drop the staging op's ref on a span; zero-ref pages stay cached
|
||||
(servable) until pressure reclaims them. Op-owned duplicate slots
|
||||
(their hash was cached elsewhere) are freed raw, as before."""
|
||||
if host_indices is None or host_indices.numel() == 0:
|
||||
return
|
||||
entry = self._mem_pool_host.entry_map.get(pool)
|
||||
if entry is None:
|
||||
return
|
||||
if not hashes:
|
||||
entry.host_pool.free(host_indices)
|
||||
return
|
||||
leftover = self.page_cache.release(
|
||||
pool, hashes, host_indices, entry.host_pool.page_size
|
||||
)
|
||||
if leftover is not None and leftover.numel() > 0:
|
||||
entry.host_pool.free(leftover)
|
||||
|
||||
def register_hold(
|
||||
self,
|
||||
hash_values: list[str],
|
||||
host_indices: torch.Tensor,
|
||||
aux_xfers: list[PoolTransfer],
|
||||
retain: bool = True,
|
||||
) -> None:
|
||||
"""Register a staged KV span plus its aux windows (SWA/Mamba states
|
||||
keyed by their trailing KV page hashes, same keying the storage
|
||||
write uses). ``retain=False`` = write-around: servable only while
|
||||
the staging op pins the slots, freed at the last release unless a
|
||||
read hit promotes it."""
|
||||
self.register_span(PoolName.KV, hash_values, host_indices, retain=retain)
|
||||
for transfer in aux_xfers:
|
||||
keys = self.aux_window_keys(hash_values, transfer)
|
||||
if keys is not None:
|
||||
self.register_span(
|
||||
transfer.name, keys, transfer.host_indices, retain=retain
|
||||
)
|
||||
|
||||
def release_hold(
|
||||
self,
|
||||
hash_values: list[str],
|
||||
host_indices: torch.Tensor,
|
||||
aux_xfers: list[PoolTransfer],
|
||||
) -> None:
|
||||
"""Mirror of register_hold for every hold retirement path
|
||||
(storage-ack, fill H2D-ack, staged drop, abort)."""
|
||||
self.release_span(PoolName.KV, hash_values, host_indices)
|
||||
for transfer in aux_xfers:
|
||||
if transfer.indices_from_pool is not None:
|
||||
continue
|
||||
keys = self.aux_window_keys(hash_values, transfer)
|
||||
self.release_span(transfer.name, keys or [], transfer.host_indices)
|
||||
|
||||
def reclaim(self, pool: PoolName, num_tokens: int) -> int:
|
||||
"""Free just enough zero-ref cached pages for an allocation of
|
||||
num_tokens to succeed. Scheduler-thread only (lockstep pressure
|
||||
points: staging-hit alloc, prepare_prefetch, cc.write)."""
|
||||
entry = self._mem_pool_host.entry_map.get(pool)
|
||||
if entry is None:
|
||||
return 0
|
||||
shortfall = num_tokens - entry.host_pool.available_size()
|
||||
if shortfall <= 0:
|
||||
return 0
|
||||
return self.page_cache.reclaim(pool, shortfall, entry.host_pool.free)
|
||||
|
||||
def continuation_run(self, chain: list[str], start_pages: int) -> int:
|
||||
"""Longest cached run continuing the span at page ``start_pages``
|
||||
(0 = leading run), folded for SWA: the joint span's trailing window
|
||||
must be fully cache-servable, mirroring batch_exists_v2's
|
||||
trailing_pages fold. Non-mutating and rank-deterministic."""
|
||||
page_cache = self.page_cache
|
||||
kv_run = page_cache.peek_run_len(PoolName.KV, chain[start_pages:])
|
||||
if kv_run == 0:
|
||||
return 0
|
||||
sw_pages = self._sw_window_pages_fn()
|
||||
if sw_pages == 0:
|
||||
return kv_run
|
||||
for cont in range(kv_run, 0, -1):
|
||||
joint = start_pages + cont
|
||||
window = min(sw_pages, joint)
|
||||
if cont < window:
|
||||
# Window straddles into the head; only possible for
|
||||
# anchored runs, and shrinking cont cannot fix it.
|
||||
break
|
||||
if all(
|
||||
page_cache.contains(PoolName.SWA, chain[i])
|
||||
for i in range(joint - window, joint)
|
||||
):
|
||||
return cont
|
||||
return 0
|
||||
|
||||
def acquire_span(
|
||||
self, chain: list[str], start_pages: int, cont_pages: int
|
||||
) -> Optional[tuple[torch.Tensor, list[PoolTransfer]]]:
|
||||
"""Acquire a folded continuation run: its KV pages plus the JOINT
|
||||
span's trailing SWA window (refs++ on every page). Returns
|
||||
(kv_slots, aux_xfers) or None (with no refs held) if a page
|
||||
vanished since the fold — defensive; fold and acquire run in the
|
||||
same lockstep step."""
|
||||
page_cache = self.page_cache
|
||||
cont_hashes = list(chain[start_pages : start_pages + cont_pages])
|
||||
kv_slots = page_cache.acquire(PoolName.KV, cont_hashes)
|
||||
if kv_slots is None:
|
||||
return None
|
||||
aux_xfers: list[PoolTransfer] = []
|
||||
sw_pages = self._sw_window_pages_fn()
|
||||
if sw_pages > 0:
|
||||
joint = start_pages + cont_pages
|
||||
window_hashes = list(chain[joint - min(sw_pages, joint) : joint])
|
||||
swa_slots = page_cache.acquire(PoolName.SWA, window_hashes)
|
||||
if swa_slots is None:
|
||||
self.release_span(PoolName.KV, cont_hashes, kv_slots)
|
||||
return None
|
||||
aux_xfers.append(
|
||||
PoolTransfer(
|
||||
name=PoolName.SWA,
|
||||
host_indices=swa_slots,
|
||||
keys=window_hashes,
|
||||
hit_policy=PoolHitPolicy.TRAILING_PAGES,
|
||||
)
|
||||
)
|
||||
return kv_slots, aux_xfers
|
||||
@@ -1092,29 +1092,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
|
||||
lens_to_track = self.mamba_track_seqlens - self.extend_prefix_lens
|
||||
return (lens_to_track // chunk_size) * chunk_size
|
||||
|
||||
def merge_mm_inputs(self) -> Optional[MultimodalInputs]:
|
||||
"""
|
||||
Merge all multimodal inputs in the batch into a single MultiModalInputs object.
|
||||
|
||||
Returns:
|
||||
if none, current batch contains no multimodal input
|
||||
|
||||
"""
|
||||
if not self.mm_inputs or all(x is None for x in self.mm_inputs):
|
||||
return None
|
||||
# Filter out None values
|
||||
valid_inputs = [x for x in self.mm_inputs if x is not None]
|
||||
|
||||
# TODO: is it expensive?
|
||||
# a workaround to avoid importing `MultimodalInputs`
|
||||
merged = valid_inputs[0].__class__(mm_items=[])
|
||||
|
||||
# Merge remaining inputs
|
||||
for mm_input in valid_inputs:
|
||||
merged.merge(mm_input)
|
||||
|
||||
return merged
|
||||
|
||||
def contains_image_inputs(self) -> bool:
|
||||
if self.mm_inputs is None:
|
||||
return False
|
||||
|
||||
@@ -17,10 +17,8 @@ import re
|
||||
import struct
|
||||
import tempfile
|
||||
import threading
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
Generator,
|
||||
@@ -172,57 +170,6 @@ def get_lock(
|
||||
return lock
|
||||
|
||||
|
||||
def _shared_pointers(tensors):
|
||||
ptrs = defaultdict(list)
|
||||
for k, v in tensors.items():
|
||||
ptrs[v.data_ptr()].append(k)
|
||||
failing = []
|
||||
for _, names in ptrs.items():
|
||||
if len(names) > 1:
|
||||
failing.append(names)
|
||||
return failing
|
||||
|
||||
|
||||
def convert_bin_to_safetensor_file(
|
||||
pt_filename: str,
|
||||
sf_filename: str,
|
||||
) -> None:
|
||||
loaded = torch.load(pt_filename, map_location="cpu", weights_only=True)
|
||||
if "state_dict" in loaded:
|
||||
loaded = loaded["state_dict"]
|
||||
shared = _shared_pointers(loaded)
|
||||
for shared_weights in shared:
|
||||
for name in shared_weights[1:]:
|
||||
loaded.pop(name)
|
||||
|
||||
# For tensors to be contiguous
|
||||
loaded = {k: v.contiguous() for k, v in loaded.items()}
|
||||
|
||||
dirname = os.path.dirname(sf_filename)
|
||||
os.makedirs(dirname, exist_ok=True)
|
||||
|
||||
from safetensors.torch import save_file
|
||||
|
||||
save_file(loaded, sf_filename, metadata={"format": "pt"})
|
||||
|
||||
# check file size
|
||||
sf_size = os.stat(sf_filename).st_size
|
||||
pt_size = os.stat(pt_filename).st_size
|
||||
if (sf_size - pt_size) / pt_size > 0.01:
|
||||
raise RuntimeError(f"""The file size different is more than 1%:
|
||||
- {sf_filename}: {sf_size}
|
||||
- {pt_filename}: {pt_size}
|
||||
""")
|
||||
|
||||
# check if the tensors are the same
|
||||
reloaded = safetensors.torch.load_file(sf_filename)
|
||||
for k in loaded:
|
||||
pt_tensor = loaded[k]
|
||||
sf_tensor = reloaded[k]
|
||||
if not torch.equal(pt_tensor, sf_tensor):
|
||||
raise RuntimeError(f"The output tensors do not match for key {k}")
|
||||
|
||||
|
||||
def replace_prefix(key: str, prefix_mapping: dict[str, str]) -> str:
|
||||
for prefix, new_prefix in prefix_mapping.items():
|
||||
if key.startswith(prefix):
|
||||
@@ -1263,13 +1210,10 @@ def fastsafetensors_weights_iterator(
|
||||
loader.add_filenames(rank_file_map)
|
||||
try:
|
||||
fb = loader.copy_files_to_device()
|
||||
try:
|
||||
keys = list(fb.key_to_rank_lidx.keys())
|
||||
for k in keys:
|
||||
t = fb.get_tensor(k)
|
||||
yield k, t
|
||||
finally:
|
||||
pass
|
||||
keys = list(fb.key_to_rank_lidx.keys())
|
||||
for k in keys:
|
||||
t = fb.get_tensor(k)
|
||||
yield k, t
|
||||
finally:
|
||||
loader.close()
|
||||
if drop_cache_after_load:
|
||||
@@ -1277,50 +1221,6 @@ def fastsafetensors_weights_iterator(
|
||||
_drop_file_cache_after_load(loaded_file)
|
||||
|
||||
|
||||
def multi_thread_safetensors_weights_iterator(
|
||||
hf_weights_files: List[str],
|
||||
max_workers: int,
|
||||
disable_mmap: bool = False,
|
||||
drop_cache_after_load: bool = False,
|
||||
) -> Generator[Tuple[str, torch.Tensor], None, None]:
|
||||
"""Multi-Thread iterate over the weights in the model safetensor files."""
|
||||
enable_tqdm = (
|
||||
not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
|
||||
)
|
||||
|
||||
def _load_file(st_file: str):
|
||||
if disable_mmap:
|
||||
with open(st_file, "rb") as f:
|
||||
result = safetensors.torch.load(f.read())
|
||||
else:
|
||||
with safetensors.safe_open(st_file, framework="pt", device="cpu") as f:
|
||||
result = {k: f.get_tensor(k) for k in f.keys()}
|
||||
|
||||
return st_file, result
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [executor.submit(_load_file, st_file) for st_file in hf_weights_files]
|
||||
|
||||
if enable_tqdm:
|
||||
futures_iter = tqdm(
|
||||
concurrent.futures.as_completed(futures),
|
||||
total=len(hf_weights_files),
|
||||
desc="Multi-thread loading shards",
|
||||
disable=not enable_tqdm,
|
||||
bar_format=BAR_FORMAT,
|
||||
)
|
||||
else:
|
||||
futures_iter = concurrent.futures.as_completed(futures)
|
||||
|
||||
for future in futures_iter:
|
||||
st_file, state_dict = future.result()
|
||||
for name, param in state_dict.items():
|
||||
yield name, param
|
||||
del state_dict
|
||||
if drop_cache_after_load:
|
||||
_drop_file_cache_after_load(st_file)
|
||||
|
||||
|
||||
def buffered_multi_thread_safetensors_weights_iterator(
|
||||
hf_weights_files: List[str],
|
||||
max_workers: int,
|
||||
@@ -1569,55 +1469,19 @@ def gguf_quant_weights_iterator(
|
||||
yield name, param
|
||||
|
||||
|
||||
def convert_pyslice_to_tensor(x: Any) -> torch.Tensor:
|
||||
"""convert PySafeSlice object from safetensors to torch.Tensor
|
||||
|
||||
PySafeSlice object supports indexing, which is done before loading the
|
||||
actual tensor and can reduce the amount of memory being read into the
|
||||
memory. However, it does not support more advanced functionalities
|
||||
like `.view()` or `.t()`. Therefore, if we need to modify the loaded
|
||||
tensor with these more complicated operators, we need to convert to
|
||||
tensor first.
|
||||
"""
|
||||
if not isinstance(x, torch.Tensor):
|
||||
x = x[:]
|
||||
return x
|
||||
|
||||
|
||||
def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
|
||||
"""Default weight loader."""
|
||||
try:
|
||||
if param.numel() == 1 and loaded_weight.numel() == 1:
|
||||
# Sometimes scalar values aren't considered tensors with shapes
|
||||
# so if both param and loaded_weight are a scalar,
|
||||
# "broadcast" instead of copy
|
||||
param.data.fill_(loaded_weight.item())
|
||||
else:
|
||||
assert param.size() == loaded_weight.size(), (
|
||||
f"Attempted to load weight ({loaded_weight.size()}) "
|
||||
f"into parameter ({param.size()})"
|
||||
)
|
||||
|
||||
param.data.copy_(loaded_weight)
|
||||
except Exception:
|
||||
# NOTE: This exception is added for the purpose of setting breakpoint to
|
||||
# debug weight loading issues.
|
||||
raise
|
||||
|
||||
|
||||
def row_parallel_weight_loader(
|
||||
param: torch.Tensor, loaded_weight: torch.Tensor
|
||||
) -> None:
|
||||
"""Load weights that are row-parallelized."""
|
||||
tp_rank = get_parallel().tp_rank
|
||||
shard_dim = 0 if param.dim() != 1 else None
|
||||
|
||||
if shard_dim is not None:
|
||||
shard_size = param.data.shape[shard_dim]
|
||||
start_idx = tp_rank * shard_size
|
||||
loaded_weight = loaded_weight.narrow(shard_dim, start_idx, shard_size)
|
||||
|
||||
return default_weight_loader(param, loaded_weight)
|
||||
if param.numel() == 1 and loaded_weight.numel() == 1:
|
||||
# Sometimes scalar values aren't considered tensors with shapes
|
||||
# so if both param and loaded_weight are a scalar,
|
||||
# "broadcast" instead of copy
|
||||
param.data.fill_(loaded_weight.item())
|
||||
else:
|
||||
assert param.size() == loaded_weight.size(), (
|
||||
f"Attempted to load weight ({loaded_weight.size()}) "
|
||||
f"into parameter ({param.size()})"
|
||||
)
|
||||
param.data.copy_(loaded_weight)
|
||||
|
||||
|
||||
LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor]
|
||||
|
||||
@@ -1145,29 +1145,6 @@ def get_cuda_driver_bindings():
|
||||
return cuda_driver
|
||||
|
||||
|
||||
def get_physical_device_id(pytorch_device_id: int) -> int:
|
||||
"""
|
||||
Convert PyTorch logical device ID to physical device ID.
|
||||
|
||||
When CUDA_VISIBLE_DEVICES is set, maps the logical device ID (as seen by PyTorch)
|
||||
to the actual physical device ID. If CUDA_VISIBLE_DEVICES is not set, returns
|
||||
the device ID unchanged.
|
||||
|
||||
Args:
|
||||
pytorch_device_id: The logical device ID from PyTorch (e.g., torch.cuda.current_device())
|
||||
|
||||
Returns:
|
||||
The physical device ID
|
||||
"""
|
||||
device_idx = int(pytorch_device_id)
|
||||
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
|
||||
if cuda_visible_devices:
|
||||
device_list = cuda_visible_devices.split(",")
|
||||
return int(device_list[device_idx])
|
||||
else:
|
||||
return device_idx
|
||||
|
||||
|
||||
def get_device_sm_nvidia_smi():
|
||||
try:
|
||||
# Run nvidia-smi command and capture output
|
||||
@@ -1431,25 +1408,6 @@ def mark_end(name):
|
||||
time_infos[name].pretty_print()
|
||||
|
||||
|
||||
def calculate_time(show=False, min_cost_ms=0.0):
|
||||
def wrapper(func):
|
||||
def inner_func(*args, **kwargs):
|
||||
torch.cuda.synchronize()
|
||||
if show:
|
||||
start_time = time.perf_counter()
|
||||
result = func(*args, **kwargs)
|
||||
torch.cuda.synchronize()
|
||||
if show:
|
||||
cost_time = (time.perf_counter() - start_time) * 1000
|
||||
if cost_time > min_cost_ms:
|
||||
print(f"Function {func.__name__} took {cost_time} ms to run.")
|
||||
return result
|
||||
|
||||
return inner_func
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class LayerFn(Protocol):
|
||||
def __call__(self, idx: int, prefix: str) -> torch.nn.Module: ...
|
||||
|
||||
@@ -1498,24 +1456,6 @@ def make_layers(
|
||||
return modules, start_layer, end_layer
|
||||
|
||||
|
||||
def make_layers_non_pp(
|
||||
num_hidden_layers: int,
|
||||
layer_fn: LayerFn,
|
||||
prefix: str = "",
|
||||
) -> torch.nn.ModuleList:
|
||||
from sglang.srt.utils.offloader import get_offloader
|
||||
|
||||
layers = torch.nn.ModuleList(
|
||||
get_offloader().wrap_modules(
|
||||
(
|
||||
layer_fn(idx=idx, prefix=add_prefix(idx, prefix))
|
||||
for idx in range(num_hidden_layers)
|
||||
)
|
||||
)
|
||||
)
|
||||
return layers
|
||||
|
||||
|
||||
def set_random_seed(seed: int) -> None:
|
||||
"""Set the random seed for all libraries."""
|
||||
random.seed(seed)
|
||||
@@ -2818,11 +2758,6 @@ def init_custom_process_group(
|
||||
return pg
|
||||
|
||||
|
||||
def crash_on_warnings():
|
||||
# Crash on warning if we are running CI tests
|
||||
return get_bool_env_var("SGLANG_IS_IN_CI")
|
||||
|
||||
|
||||
@functools.lru_cache(None)
|
||||
def print_warning_once(msg: str) -> None:
|
||||
# Set the stacklevel to 2 to print the caller's line info
|
||||
@@ -2956,26 +2891,6 @@ def set_gpu_proc_affinity(
|
||||
logger.info(f"Process {pid} gpu_id {gpu_id} is running on CPUs: {p.cpu_affinity()}")
|
||||
|
||||
|
||||
def permute_weight(x: torch.Tensor) -> torch.Tensor:
|
||||
b_ = x.shape[0]
|
||||
n_ = x.shape[1]
|
||||
k_ = x.shape[2]
|
||||
|
||||
x_ = x
|
||||
if x.dtype == torch.bfloat16 or x.dtype == torch.float16:
|
||||
x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 32), 4, 8)
|
||||
elif x.dtype == torch.float8_e4m3fnuz or x.dtype == torch.int8:
|
||||
x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 64), 4, 16)
|
||||
else:
|
||||
# return x_
|
||||
x_ = x_.view(int(b_), int(n_ / 16), 16, int(k_ / 8), 2, 4)
|
||||
|
||||
x_ = x_.permute(0, 1, 3, 4, 2, 5)
|
||||
x_ = x_.contiguous()
|
||||
x_ = x_.view(*x.shape)
|
||||
return x_
|
||||
|
||||
|
||||
class MultiprocessingSerializer:
|
||||
@staticmethod
|
||||
def serialize(obj, output_str: bool = False):
|
||||
@@ -3141,30 +3056,6 @@ def safe_pickle_loads(data):
|
||||
return SafeUnpickler(io.BytesIO(buf)).load()
|
||||
|
||||
|
||||
def debug_timing(func):
|
||||
# todo: replace with a more organized instrumentation
|
||||
def wrapper(*args, **kwargs):
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
tic = torch.cuda.Event(enable_timing=True)
|
||||
toc = torch.cuda.Event(enable_timing=True)
|
||||
tic.record()
|
||||
result = func(*args, **kwargs)
|
||||
toc.record()
|
||||
toc.synchronize() # Wait for the function to complete without synchronizing all ops on the GPU
|
||||
elapsed = tic.elapsed_time(toc)
|
||||
indices = kwargs.get("indices", args[1] if len(args) > 1 else None)
|
||||
num_tokens = len(indices) if indices is not None else 0
|
||||
throughput = num_tokens / elapsed * 1000 if elapsed > 0 else 0
|
||||
logger.debug(
|
||||
f"Transfer time: {elapsed} ms, throughput: {throughput} tokens/s"
|
||||
)
|
||||
return result
|
||||
else:
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def nullable_str(val: str):
|
||||
if not val or val == "None":
|
||||
return None
|
||||
@@ -3705,35 +3596,6 @@ def is_no_spec_infer_or_topk_one(cfg):
|
||||
)
|
||||
|
||||
|
||||
def is_fa3_default_architecture(hf_config):
|
||||
architectures = getattr(hf_config, "architectures", None)
|
||||
if not isinstance(architectures, list) or not architectures:
|
||||
return False
|
||||
default_archs = {
|
||||
"Llama4ForConditionalGeneration",
|
||||
"LlamaForCausalLM",
|
||||
"Olmo2ForCausalLM",
|
||||
"Gemma2ForCausalLM",
|
||||
"Gemma3ForConditionalGeneration",
|
||||
"MixtralForCausalLM",
|
||||
"Qwen2ForCausalLM",
|
||||
"Qwen3ForCausalLM",
|
||||
"Qwen3MoeForCausalLM",
|
||||
"Qwen3VLForConditionalGeneration",
|
||||
"Qwen3VLMoeForConditionalGeneration",
|
||||
"Glm4MoeForCausalLM",
|
||||
"Glm4vForConditionalGeneration",
|
||||
"Glm4vMoeForConditionalGeneration",
|
||||
"GlmOcrForConditionalGeneration",
|
||||
"Step3VLForConditionalGeneration",
|
||||
"StepVLForConditionalGeneration",
|
||||
"Step3p7ForConditionalGeneration",
|
||||
"MiMoV2ForCausalLM",
|
||||
"MiMoV2FlashForCausalLM",
|
||||
}
|
||||
return architectures[0] in default_archs
|
||||
|
||||
|
||||
# Can be more general if it is used in multiple places (keep it simple and thus not general now)
|
||||
class BumpAllocator:
|
||||
def __init__(self, buffer_size: int, dtype, device):
|
||||
|
||||
@@ -1,595 +0,0 @@
|
||||
"""Resolution reads its own decisions, not the record's fields.
|
||||
|
||||
`declare_resolution` records a decision in the declaration stash and writes
|
||||
nothing. The fields keep what the caller passed, so a resolver that reads a
|
||||
field another resolver may have decided reads the raw input -- silently, and
|
||||
only on the configurations where that other resolver fires. The whole pipeline
|
||||
therefore reads through `resolving_view` (or `resolved_view`, which is
|
||||
the same view after resolution has finished), and this pins that there is
|
||||
nothing left reading a field directly.
|
||||
|
||||
Subjects: every function in `arg_groups/` that takes a config, every
|
||||
`ServerArgs` handler the dispatcher reaches, and every member of `ServerArgs` /
|
||||
`PortArgs` -- the members are reached from the hooks and from business code,
|
||||
which the handler walk cannot see, and a member that recomputes from a raw field
|
||||
decides from what was typed. All three
|
||||
are derived -- a new hook file, a new handler or a new member is covered the
|
||||
moment it is written. Readers *outside* those
|
||||
two -- the platform defaults, `ModelConfig`, the spec-algo hook -- are reached by
|
||||
resolution too and have moved to the view as well, but enumerating them needs
|
||||
the call-graph derivation `test_resolution_reads_no_bag` owns; this file pins
|
||||
the two scopes it can derive exactly.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import pathlib
|
||||
import re
|
||||
|
||||
import msgspec
|
||||
import msgspec.structs
|
||||
|
||||
import sglang
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=45, suite="base-a-test-cpu")
|
||||
|
||||
_SRT = pathlib.Path(sglang.__file__).resolve().parent / "srt"
|
||||
_FIELDS = frozenset(field.name for field in msgspec.structs.fields(ServerArgs))
|
||||
|
||||
# Names a config travels under. `args` is included because the platform hooks
|
||||
# use it; a false positive would be a function taking an argparse Namespace and
|
||||
# reading an attribute that happens to be a ServerArgs field name, which the
|
||||
# allowlist below would then have to carry.
|
||||
_HOLDER_NAMES = frozenset({"server_args", "sa", "args"})
|
||||
|
||||
|
||||
def _holders(fn):
|
||||
names = {
|
||||
arg.arg
|
||||
for arg in list(fn.args.posonlyargs)
|
||||
+ list(fn.args.args)
|
||||
+ list(fn.args.kwonlyargs)
|
||||
if arg.arg in _HOLDER_NAMES
|
||||
}
|
||||
for arg in (
|
||||
list(fn.args.posonlyargs) + list(fn.args.args) + list(fn.args.kwonlyargs)
|
||||
):
|
||||
annotation = arg.annotation
|
||||
text = (
|
||||
annotation.value
|
||||
if isinstance(annotation, ast.Constant)
|
||||
else (
|
||||
annotation.id
|
||||
if isinstance(annotation, ast.Name)
|
||||
else annotation.attr
|
||||
if isinstance(annotation, ast.Attribute)
|
||||
else None
|
||||
)
|
||||
)
|
||||
if text == "ServerArgs":
|
||||
names.add(arg.arg)
|
||||
return names
|
||||
|
||||
|
||||
def _field_reads(fn, holders):
|
||||
for node in ast.walk(fn):
|
||||
if (
|
||||
isinstance(node, ast.Attribute)
|
||||
and node.attr in _FIELDS
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id in holders
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
):
|
||||
yield node.lineno, node.attr
|
||||
|
||||
|
||||
_DECLARERS = frozenset(
|
||||
{
|
||||
"declare_resolution",
|
||||
"record_foreign_defaults",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _declared_fields():
|
||||
"""The fields resolution decides, read off every shape that reaches the stash.
|
||||
|
||||
A keyword on a `declare_*` call is only one shape: the model-override and
|
||||
post-process passes build a mapping instead (`MODEL_OVERRIDES` literals,
|
||||
`overrides["dtype"] = ...`, a returned dict), and late resolution splats a
|
||||
variable-keyed one. Deriving from keywords alone leaves nineteen fields
|
||||
outside the subject set, `dtype` and `reasoning_parser` among them.
|
||||
"""
|
||||
fields = set()
|
||||
# The declaration calls live wherever a resolver does; the mapping channels
|
||||
# only exist where the override providers and post-process passes are.
|
||||
keyword_sources = [_SRT / "server_args.py"]
|
||||
for sub in ("arg_groups", "hardware_backend", "parser"):
|
||||
keyword_sources += sorted((_SRT / sub).rglob("*.py"))
|
||||
mapping_sources = {_SRT / "server_args.py", *(_SRT / "arg_groups").rglob("*.py")}
|
||||
for path in keyword_sources:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
for node in ast.walk(tree):
|
||||
# 1. `declare_resolution(sa, src, page_size=64)` and its siblings
|
||||
if isinstance(node, ast.Call):
|
||||
name = (
|
||||
node.func.id
|
||||
if isinstance(node.func, ast.Name)
|
||||
else getattr(node.func, "attr", None)
|
||||
)
|
||||
if name in _DECLARERS:
|
||||
for keyword in node.keywords:
|
||||
if keyword.arg:
|
||||
fields.add(keyword.arg)
|
||||
elif isinstance(keyword.value, ast.Dict):
|
||||
fields.update(_string_keys(keyword.value))
|
||||
# 2. every mapping literal in the files that declare through one:
|
||||
# the MODEL_OVERRIDES tables, the dicts the override providers
|
||||
# return, the ones the post-process passes build. Scanning
|
||||
# unrelated files here would collect a plain kwarg dict
|
||||
# (`tokenizer_config={"trust_remote_code": ...}`) and turn a
|
||||
# passthrough read into a violation.
|
||||
if isinstance(node, ast.Dict) and path in mapping_sources:
|
||||
fields.update(_string_keys(node))
|
||||
# 3. `overrides["field"] = ...`
|
||||
if (
|
||||
path in mapping_sources
|
||||
and isinstance(node, ast.Assign)
|
||||
and isinstance(node.targets[0], ast.Subscript)
|
||||
and isinstance(node.targets[0].slice, ast.Constant)
|
||||
and isinstance(node.targets[0].slice.value, str)
|
||||
):
|
||||
fields.add(node.targets[0].slice.value)
|
||||
return frozenset(fields & _FIELDS)
|
||||
|
||||
|
||||
def _string_keys(node: ast.Dict) -> set:
|
||||
return {
|
||||
key.value
|
||||
for key in node.keys
|
||||
if isinstance(key, ast.Constant) and isinstance(key.value, str)
|
||||
}
|
||||
|
||||
|
||||
def _record_members():
|
||||
"""Every member of `ServerArgs` / `PortArgs`, by class and name."""
|
||||
tree = ast.parse((_SRT / "server_args.py").read_text(encoding="utf-8-sig"))
|
||||
members = {}
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ClassDef) and node.name in ("ServerArgs", "PortArgs"):
|
||||
for member in node.body:
|
||||
if isinstance(member, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
members[f"{node.name}.{member.name}"] = member
|
||||
return members
|
||||
|
||||
|
||||
def _config_reading_helpers():
|
||||
"""Module functions that load a decided field off the config they are handed.
|
||||
|
||||
A member that hands them `self`, or a call site that hands them a record,
|
||||
reads the raw input through the callee -- the shape neither an attribute
|
||||
scan nor a `getattr` scan can see, because the field name is spelled in the
|
||||
helper and the record is spelled at the call site.
|
||||
"""
|
||||
decided = _declared_fields()
|
||||
helpers = {}
|
||||
sources = [_SRT / "server_args.py"] + sorted((_SRT / "arg_groups").rglob("*.py"))
|
||||
for path in sources:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
for fn in ast.walk(tree):
|
||||
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
params = {
|
||||
arg.arg
|
||||
for arg in list(fn.args.posonlyargs)
|
||||
+ list(fn.args.args)
|
||||
+ list(fn.args.kwonlyargs)
|
||||
} - {"self", "cls"}
|
||||
if not params:
|
||||
continue
|
||||
reads = {
|
||||
node.attr
|
||||
for node in ast.walk(fn)
|
||||
if isinstance(node, ast.Attribute)
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id in params
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
and node.attr in decided
|
||||
}
|
||||
if reads:
|
||||
helpers[fn.name] = sorted(reads)
|
||||
return helpers
|
||||
|
||||
|
||||
# The accessors that hand back the process-global record itself. A helper that
|
||||
# is handed one of these reads the raw input exactly as a bare `self` would.
|
||||
_RECORD_ACCESSORS = frozenset({"get_server_args", "global_server_args"})
|
||||
|
||||
# `self._server_args` is the same record under a private name; the scan has to
|
||||
# see it or a reader inside the context object escapes every shape above.
|
||||
_RECORD_ATTR = re.compile(r"^_*(server_args|sa)$")
|
||||
|
||||
|
||||
def _record_arguments(node, aliases=frozenset()):
|
||||
"""The bare-record arguments of a call.
|
||||
|
||||
Four spellings reach a helper with a record: the bare name (`self`, `sa`),
|
||||
an attribute (`runner.server_args`), the process-global accessor called
|
||||
inline (`get_server_args()`), and a local bound to either of the last two
|
||||
earlier in the same function.
|
||||
"""
|
||||
out = []
|
||||
for arg in node.args:
|
||||
if isinstance(arg, ast.Name) and arg.id in ("self", "server_args", "sa"):
|
||||
out.append(arg.id)
|
||||
elif isinstance(arg, ast.Attribute) and _RECORD_ATTR.match(arg.attr or ""):
|
||||
out.append(ast.unparse(arg))
|
||||
elif (
|
||||
isinstance(arg, ast.Call)
|
||||
and isinstance(arg.func, ast.Name)
|
||||
and arg.func.id in _RECORD_ACCESSORS
|
||||
):
|
||||
out.append(ast.unparse(arg))
|
||||
elif isinstance(arg, ast.Name) and arg.id in aliases:
|
||||
out.append(arg.id)
|
||||
return out
|
||||
|
||||
|
||||
def _record_aliases(function):
|
||||
"""Locals bound to the record under a name of their own.
|
||||
|
||||
`_sa = getattr(runner, "server_args", None)`, `cfg = get_server_args()` and
|
||||
`engine_args = ServerArgs.from_cli_args(args)` all put the record behind a
|
||||
name the argument scan does not recognise, so a later
|
||||
`getattr(_sa, "<decided leaf>")` reads what the operator typed.
|
||||
"""
|
||||
aliases = set()
|
||||
for node in ast.walk(function):
|
||||
if not isinstance(node, ast.Assign) or len(node.targets) != 1:
|
||||
continue
|
||||
target = node.targets[0]
|
||||
if not isinstance(target, ast.Name):
|
||||
continue
|
||||
value = node.value
|
||||
if isinstance(value, ast.Attribute):
|
||||
if _RECORD_ATTR.match(value.attr or ""):
|
||||
aliases.add(target.id)
|
||||
continue
|
||||
if not isinstance(value, ast.Call):
|
||||
continue
|
||||
func = value.func
|
||||
if isinstance(func, ast.Name):
|
||||
if func.id in _RECORD_ACCESSORS or func.id == "ServerArgs":
|
||||
aliases.add(target.id)
|
||||
elif (
|
||||
func.id == "getattr"
|
||||
and len(value.args) >= 2
|
||||
and isinstance(value.args[1], ast.Constant)
|
||||
and _RECORD_ATTR.match(str(value.args[1].value))
|
||||
):
|
||||
aliases.add(target.id)
|
||||
elif (
|
||||
isinstance(func, ast.Attribute)
|
||||
and func.attr == "from_cli_args"
|
||||
and isinstance(func.value, ast.Name)
|
||||
and func.value.id == "ServerArgs"
|
||||
):
|
||||
aliases.add(target.id)
|
||||
return aliases
|
||||
|
||||
|
||||
# The one reader for which the raw field is the right answer. The gateway sizes
|
||||
# its worker pool from the operator's requested replica count; `--dwdp-size`
|
||||
# makes resolution declare a `dp_size` describing one multi-rank server's
|
||||
# internal topology, so reading the decision there would spawn dp_size
|
||||
# single-rank children and ask for dp_size^2 GPUs. A new entry here needs that
|
||||
# kind of reason next to it.
|
||||
_NO_RESOLVED_SURFACE = frozenset(
|
||||
{
|
||||
(
|
||||
"sgl-model-gateway/bindings/python/src/sglang_router/launch_server.py",
|
||||
"server_args.dp_size",
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_record_base(node, aliases):
|
||||
"""Is this expression the record itself?
|
||||
|
||||
A local bound to one, a parameter that carries one (`server_args`, `sa`,
|
||||
`engine_args`), or an attribute holding one (`self._server_args`).
|
||||
"""
|
||||
if isinstance(node, ast.Name):
|
||||
return node.id in aliases
|
||||
if isinstance(node, ast.Attribute):
|
||||
return bool(_RECORD_ATTR.match(node.attr or ""))
|
||||
return False
|
||||
|
||||
|
||||
def _record_handoff_offenders(rel, tree, helpers, decided, is_record=False):
|
||||
"""Every way a decided leaf is reached through a record in one module.
|
||||
|
||||
Two shapes, both scanned under the record aliases the function binds:
|
||||
handing the record to a helper that loads a decided field, and loading one
|
||||
off the alias directly (`alias.<leaf>` or `getattr(alias, "<leaf>")`). The
|
||||
second is what the MiniMax backend spelled, and an argument scan cannot see
|
||||
it -- the leaf never appears at a call site.
|
||||
"""
|
||||
offenders, seen = [], set()
|
||||
|
||||
def record(lineno, text):
|
||||
if (lineno, text) in seen:
|
||||
return
|
||||
seen.add((lineno, text))
|
||||
offenders.append(f"{rel}:{lineno} {text}")
|
||||
|
||||
scopes = [(tree, frozenset())] + [
|
||||
(fn, _record_aliases(fn))
|
||||
for fn in ast.walk(tree)
|
||||
if isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
]
|
||||
for scope, aliases in scopes:
|
||||
for node in ast.walk(scope):
|
||||
if isinstance(node, ast.Call):
|
||||
name = (
|
||||
node.func.id
|
||||
if isinstance(node.func, ast.Name)
|
||||
else getattr(node.func, "attr", None)
|
||||
)
|
||||
if name in helpers:
|
||||
for arg in _record_arguments(node, aliases):
|
||||
if arg == "self" and not is_record:
|
||||
continue
|
||||
record(
|
||||
node.lineno,
|
||||
f"{name}({arg}) reads {', '.join(helpers[name])}",
|
||||
)
|
||||
if (
|
||||
isinstance(node.func, ast.Name)
|
||||
and node.func.id == "getattr"
|
||||
and len(node.args) >= 2
|
||||
and isinstance(node.args[0], ast.Name)
|
||||
and node.args[0].id in aliases
|
||||
and isinstance(node.args[1], ast.Constant)
|
||||
and node.args[1].value in decided
|
||||
):
|
||||
record(
|
||||
node.lineno,
|
||||
f'getattr({node.args[0].id}, "{node.args[1].value}")',
|
||||
)
|
||||
elif (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
and node.attr in decided
|
||||
and _is_record_base(node.value, aliases)
|
||||
):
|
||||
record(node.lineno, f"{ast.unparse(node.value)}.{node.attr}")
|
||||
return offenders
|
||||
|
||||
|
||||
# Source the scanner must read the same way whether or not the tree happens to
|
||||
# contain these shapes today. The first four are the spellings that reached
|
||||
# production and were converted; the last two are the legal forms next to them,
|
||||
# which have to stay quiet or the guard is unusable.
|
||||
_SPELLINGS = """
|
||||
def hands_the_alias_to_a_helper(runner):
|
||||
_sa = getattr(runner, "server_args", None)
|
||||
return m3_fp8_attn_gemm_enabled(_sa)
|
||||
|
||||
|
||||
def loads_a_leaf_off_the_alias(runner):
|
||||
_sa = getattr(runner, "server_args", None)
|
||||
return getattr(_sa, "speculative_num_draft_tokens", None)
|
||||
|
||||
|
||||
def reads_a_leaf_through_the_alias(runner):
|
||||
sa_local = runner.server_args
|
||||
return sa_local.attention_backend
|
||||
|
||||
|
||||
def hands_the_accessor_to_a_helper():
|
||||
return attention_backends_of(get_server_args())
|
||||
|
||||
|
||||
def reads_the_view(runner):
|
||||
cfg = resolving_view(runner.server_args)
|
||||
return cfg.attention_backend
|
||||
|
||||
|
||||
def reads_an_undecided_leaf(runner):
|
||||
_sa = runner.server_args
|
||||
return _sa.tp_size
|
||||
|
||||
|
||||
def reads_a_leaf_off_a_private_attribute(self):
|
||||
return self._server_args.attention_backend
|
||||
|
||||
|
||||
def reads_a_leaf_off_a_constructed_record(cli):
|
||||
engine_args = ServerArgs.from_cli_args(cli)
|
||||
engine_args.resolve_once()
|
||||
return engine_args.attention_backend
|
||||
"""
|
||||
|
||||
|
||||
class TestResolutionReadsTheDeclarations(CustomTestCase):
|
||||
def test_no_hook_reads_a_field_off_the_record(self):
|
||||
offenders = []
|
||||
files = sorted((_SRT / "arg_groups").glob("*.py"))
|
||||
self.assertGreater(len(files), 5, "the hook scan found almost nothing")
|
||||
for path in files:
|
||||
rel = f"arg_groups/{path.name}"
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
for fn in ast.walk(tree):
|
||||
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
holders = _holders(fn)
|
||||
if not holders:
|
||||
continue
|
||||
for lineno, field in _field_reads(fn, holders):
|
||||
offenders.append(f"{rel}:{lineno} {fn.name} reads .{field}")
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
"a resolution hook reads a field off the record; the field holds the "
|
||||
"raw input, so this decides from what was typed rather than from "
|
||||
"what resolution decided. Read `resolving_view(server_args)`:\n "
|
||||
+ "\n ".join(offenders),
|
||||
)
|
||||
|
||||
def test_the_record_hosts_no_resolution_handler(self):
|
||||
"""The pipeline and every step it runs live under `arg_groups/`.
|
||||
|
||||
While a step was a method, it could read a raw field off `self` and
|
||||
`test_no_handler_reads_a_field_off_self` had to say it could not. There
|
||||
is no such method left, so the invariant is now the stronger one: the
|
||||
record hosts none of them. What the steps read is checked on the
|
||||
package side, by `test_no_hook_reads_a_field_off_the_record`.
|
||||
"""
|
||||
handlers = sorted(
|
||||
name
|
||||
for name, node in _record_members().items()
|
||||
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and (
|
||||
name.split(".")[-1].startswith(("_handle_", "_validate_"))
|
||||
or "resolution_pipeline" in name
|
||||
)
|
||||
)
|
||||
self.assertEqual(
|
||||
handlers,
|
||||
[],
|
||||
"a resolution handler is back on the record; it belongs in an "
|
||||
"`arg_groups` family, where the package-side guards can see it:\n "
|
||||
+ "\n ".join(handlers),
|
||||
)
|
||||
|
||||
def test_no_member_recomputes_from_a_raw_field(self):
|
||||
decided = _declared_fields()
|
||||
self.assertGreater(
|
||||
len(decided), 100, f"the declaration set derived only {len(decided)} fields"
|
||||
)
|
||||
members = _record_members()
|
||||
# The floor is here to catch the scan collapsing, not to pin the
|
||||
# class's size -- it drops as derived members move to their namespaces
|
||||
# and become declarations rather than methods on the record.
|
||||
self.assertGreater(len(members), 10, f"only {len(members)} members were found")
|
||||
offenders = []
|
||||
for name, fn in sorted(members.items()):
|
||||
holders = _holders(fn) | {"self"}
|
||||
for lineno, field in _field_reads(fn, holders):
|
||||
if field in decided:
|
||||
offenders.append(f"server_args.py:{lineno} {name} reads .{field}")
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
"a record member recomputes from a field resolution decides; the "
|
||||
"field holds the raw input, so the member answers for what was "
|
||||
"typed. Bind `cfg = resolving_view(self)` and read that:\n "
|
||||
+ "\n ".join(offenders),
|
||||
)
|
||||
|
||||
def test_no_reader_hands_the_record_to_a_config_helper(self):
|
||||
helpers = _config_reading_helpers()
|
||||
self.assertGreater(
|
||||
len(helpers), 5, f"the helper derivation found only {len(helpers)}"
|
||||
)
|
||||
decided = _declared_fields()
|
||||
offenders = []
|
||||
# `scripts/`, `examples/` and the gateway binding are outside the
|
||||
# package but hold records they resolve themselves, and every reader
|
||||
# this scan found in them was reading a field resolution fills in.
|
||||
_REPO = _SRT.parent.parent.parent
|
||||
roots = (
|
||||
[_SRT]
|
||||
+ [_SRT.parent / d for d in ("benchmark", "lang")]
|
||||
+ [
|
||||
_REPO / d
|
||||
for d in (
|
||||
"scripts",
|
||||
"examples",
|
||||
"sgl-model-gateway/bindings/python/src",
|
||||
)
|
||||
]
|
||||
)
|
||||
for root in roots:
|
||||
if not root.exists():
|
||||
continue
|
||||
for path in sorted(root.rglob("*.py")):
|
||||
# Package files keep their `srt/...` spelling (the skips below
|
||||
# key on it); the repo-level roots are named from the repo.
|
||||
try:
|
||||
rel = path.relative_to(_SRT.parent).as_posix()
|
||||
except ValueError:
|
||||
rel = path.relative_to(_REPO).as_posix()
|
||||
if rel.startswith(("srt/arg_groups/", "multimodal_gen/")):
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
except SyntaxError:
|
||||
continue
|
||||
offenders += _record_handoff_offenders(
|
||||
rel, tree, helpers, decided, is_record=rel == "srt/server_args.py"
|
||||
)
|
||||
offenders = [
|
||||
line
|
||||
for line in offenders
|
||||
if (line.split(":", 1)[0], line.split(" ", 1)[1])
|
||||
not in _NO_RESOLVED_SURFACE
|
||||
]
|
||||
self.assertEqual(
|
||||
offenders,
|
||||
[],
|
||||
"a caller hands the record to a helper that loads a field "
|
||||
"resolution decides; the helper then reads the raw input. Hand it "
|
||||
"`resolving_view(record)` (or the published bag):\n "
|
||||
+ "\n ".join(offenders),
|
||||
)
|
||||
|
||||
def test_the_scan_sees_every_spelling_that_reached_production(self):
|
||||
"""Every spelling that reached production, pinned next to the scanner.
|
||||
|
||||
A shape the scan stops seeing is a silent hole, so each one is listed
|
||||
here with the legal forms beside it and the flagged set compared
|
||||
exactly.
|
||||
|
||||
What it does not reach: a record that arrives as a *parameter* and was
|
||||
resolved by the caller (`scripts/playground/bench_speculative.py` hands
|
||||
`main(args, server_args)` one). Binding that would need the call graph,
|
||||
and naming a parameter `server_args` is also how the resolution-time
|
||||
readers spell a view.
|
||||
"""
|
||||
helpers = _config_reading_helpers()
|
||||
decided = _declared_fields()
|
||||
for name in ("m3_fp8_attn_gemm_enabled", "attention_backends_of"):
|
||||
self.assertIn(name, helpers, f"the helper derivation lost {name}")
|
||||
for field in ("speculative_num_draft_tokens", "attention_backend"):
|
||||
self.assertIn(field, decided, f"the declared set lost {field}")
|
||||
|
||||
offenders = _record_handoff_offenders(
|
||||
"sample.py", ast.parse(_SPELLINGS), helpers, decided
|
||||
)
|
||||
flagged = {line.split(" ", 1)[1] for line in offenders}
|
||||
self.assertEqual(
|
||||
flagged,
|
||||
{
|
||||
"m3_fp8_attn_gemm_enabled(_sa)"
|
||||
" reads " + ", ".join(helpers["m3_fp8_attn_gemm_enabled"]),
|
||||
'getattr(_sa, "speculative_num_draft_tokens")',
|
||||
"sa_local.attention_backend",
|
||||
"self._server_args.attention_backend",
|
||||
"engine_args.attention_backend",
|
||||
"attention_backends_of(get_server_args())"
|
||||
" reads " + ", ".join(helpers["attention_backends_of"]),
|
||||
},
|
||||
"the scan lost a spelling, or started flagging a legal one:\n "
|
||||
+ "\n ".join(sorted(flagged)),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import unittest
|
||||
|
||||
unittest.main()
|
||||
@@ -1,583 +0,0 @@
|
||||
"""Nobody reaches the startup record through another object for a resolved value.
|
||||
|
||||
The supplied-instance census counts three spellings, all of which start from a
|
||||
`server_args` parameter -- the caller chose the object, which is the contract
|
||||
that makes those reads defensible. This pins the fourth: `model_runner.
|
||||
server_args.field`, `self.scheduler.server_args.field`, `tokenizer_manager.
|
||||
server_args.field`. A reference lifted off whatever object happened to hold the
|
||||
record carries no contract at all, and it was invisible to every census, which
|
||||
is how it grew to 105 reads across 36 files unnoticed.
|
||||
|
||||
They are gone, and this is what keeps them gone. Only fields resolution writes
|
||||
are pinned: reading `model_runner.server_args.host` off the record answers with
|
||||
what the caller asked for, which is what the record is for. The written set is
|
||||
derived from the declaration sites rather than listed, so a field that stops
|
||||
being resolution-written drops out on its own -- and a field that stops being
|
||||
*declared* cannot slip out that way, because bare assignment during resolution
|
||||
is refused by `server_args/test_resolution_declarations.py`.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import pathlib
|
||||
import unittest
|
||||
|
||||
import sglang
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=28, suite="base-a-test-cpu")
|
||||
|
||||
_PACKAGE = pathlib.Path(sglang.__file__).resolve().parent
|
||||
_SRT = _PACKAGE / "srt"
|
||||
|
||||
# The pipeline and its extension points: reading the in-flight record is their
|
||||
# job, and they run before anything is published.
|
||||
_OWNERS = ("server_args.py", "runtime_context.py", "arg_groups/")
|
||||
|
||||
# Where the writers are. Resolution lives in `srt` -- nothing outside it
|
||||
# declares -- so the written-field derivations scan `srt` while the *reads* are
|
||||
# counted across the whole shipped package: a borrowed read answers with the
|
||||
# startup default wherever it is written, and `benchmark/` ships too.
|
||||
_READS_SCANNED = _PACKAGE
|
||||
|
||||
_DECLARERS = ("declare_resolution",)
|
||||
|
||||
|
||||
def _declared_by_keyword():
|
||||
"""Fields named as a keyword at a declaration site."""
|
||||
written = set()
|
||||
for path in sorted(_SRT.rglob("*.py")):
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if not any(declarer in source for declarer in _DECLARERS):
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
raise AssertionError(f"unparsable module in the census: {path}")
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
if isinstance(node.func, ast.Attribute):
|
||||
name = node.func.attr
|
||||
elif isinstance(node.func, ast.Name):
|
||||
name = node.func.id
|
||||
else:
|
||||
continue
|
||||
if name in _DECLARERS:
|
||||
written |= {kw.arg for kw in node.keywords if kw.arg}
|
||||
return written
|
||||
|
||||
|
||||
def _returned_field_names(function):
|
||||
"""Field names a provider/pass writes: the keys of the mapping it returns.
|
||||
|
||||
Only the returned mapping counts -- walking every `ast.Dict` in the body
|
||||
also collects a dict-valued field's *nested* keys and any unrelated local
|
||||
mapping, and those stray names would reject valid borrowed reads of fields
|
||||
resolution never writes. The mapping is traced through four spellings: a
|
||||
returned literal, assignments (annotated or not) to a returned name, a
|
||||
literal-key subscript write on it, and `.update(field=...)` on it. A
|
||||
spelling this cannot see raises instead of skipping.
|
||||
|
||||
``overrides[name]`` is also accepted when ``name`` comes from
|
||||
``for name in ("a", "b", ...)`` -- the keys stay statically enumerable.
|
||||
"""
|
||||
names = set()
|
||||
returned = set()
|
||||
# for x in ("a", "b"): ... -> {"x": {"a", "b"}}
|
||||
loop_keys = {
|
||||
node.target.id: {elt.value for elt in node.iter.elts}
|
||||
for node in ast.walk(function)
|
||||
if isinstance(node, ast.For)
|
||||
and isinstance(node.target, ast.Name)
|
||||
and isinstance(node.iter, (ast.Tuple, ast.List))
|
||||
and node.iter.elts
|
||||
and all(
|
||||
isinstance(elt, ast.Constant) and isinstance(elt.value, str)
|
||||
for elt in node.iter.elts
|
||||
)
|
||||
}
|
||||
|
||||
def top_level_keys(mapping):
|
||||
for key in mapping.keys:
|
||||
if not (isinstance(key, ast.Constant) and isinstance(key.value, str)):
|
||||
raise AssertionError(f"non-literal key in {function.name}")
|
||||
names.add(key.value)
|
||||
|
||||
def add_subscript_key(key):
|
||||
if isinstance(key, ast.Constant):
|
||||
names.add(key.value)
|
||||
elif isinstance(key, ast.Name) and key.id in loop_keys:
|
||||
names.update(loop_keys[key.id])
|
||||
else:
|
||||
raise AssertionError(f"non-literal key in {function.name}")
|
||||
|
||||
for node in ast.walk(function):
|
||||
if isinstance(node, ast.Return) and node.value is not None:
|
||||
value = node.value
|
||||
if isinstance(value, ast.Dict):
|
||||
top_level_keys(value)
|
||||
elif isinstance(value, ast.Name):
|
||||
returned.add(value.id)
|
||||
elif isinstance(value, ast.Constant) and value.value is None:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"opaque return in {function.name}: {ast.unparse(value)}"
|
||||
)
|
||||
for node in ast.walk(function):
|
||||
if isinstance(node, (ast.Assign, ast.AnnAssign)):
|
||||
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
|
||||
for target in targets:
|
||||
if (
|
||||
isinstance(target, ast.Name)
|
||||
and target.id in returned
|
||||
and isinstance(node.value, ast.Dict)
|
||||
):
|
||||
top_level_keys(node.value)
|
||||
if isinstance(target, ast.Subscript) and (
|
||||
isinstance(target.value, ast.Name) and target.value.id in returned
|
||||
):
|
||||
add_subscript_key(target.slice)
|
||||
if (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Attribute)
|
||||
and node.func.attr == "update"
|
||||
and isinstance(node.func.value, ast.Name)
|
||||
and node.func.value.id in returned
|
||||
):
|
||||
names |= {kw.arg for kw in node.keywords if kw.arg}
|
||||
# A positional dict literal has to be read here. The Dict walk above
|
||||
# only reaches literals that are *returned* or assigned to a returned
|
||||
# name, so `d.update({"field": value})` was being type-checked and
|
||||
# then dropped -- silently, under a comment claiming otherwise.
|
||||
for arg in node.args:
|
||||
if isinstance(arg, ast.Dict):
|
||||
top_level_keys(arg)
|
||||
else:
|
||||
raise AssertionError(f"opaque update() argument in {function.name}")
|
||||
if any(kw.arg is None for kw in node.keywords):
|
||||
raise AssertionError(f"**kwargs update() in {function.name}")
|
||||
return names
|
||||
|
||||
|
||||
def _declared_by_registry_and_passes():
|
||||
"""Fields the model-override registry and the post-process passes write.
|
||||
|
||||
These field names are *data* -- dict keys, not keywords -- so a keyword
|
||||
scan misses every one of them. The callables are collected from the live
|
||||
registries rather than by matching decorator names: 26 of the 27 providers
|
||||
register through a `_register_for(...)` helper, so a scan for
|
||||
`@register_model_override*` sees exactly one of them and reports a healthy
|
||||
census over a channel it cannot see.
|
||||
"""
|
||||
import sys
|
||||
|
||||
from sglang.srt.arg_groups import overrides
|
||||
|
||||
# Resolve each callable's body in the file it actually lives in. The
|
||||
# declarations are spread over `arg_groups/model_overrides/`, one module per
|
||||
# model family, and a scan hard-coded to `overrides.py` would find none of
|
||||
# them -- and, worse, would keep reporting a healthy census while doing it.
|
||||
bodies_by_module = {}
|
||||
|
||||
def _bodies(module_name):
|
||||
if module_name not in bodies_by_module:
|
||||
path = getattr(sys.modules[module_name], "__file__", None)
|
||||
assert path, f"{module_name} has no source file"
|
||||
module_tree = ast.parse(pathlib.Path(path).read_text(encoding="utf-8-sig"))
|
||||
bodies_by_module[module_name] = {
|
||||
node.name: node
|
||||
for node in ast.walk(module_tree)
|
||||
if isinstance(node, ast.FunctionDef)
|
||||
}
|
||||
return bodies_by_module[module_name]
|
||||
|
||||
callables = {fn for fns in overrides._MODEL_OVERRIDE_FNS.values() for fn in fns}
|
||||
callables |= {
|
||||
fn for _predicate, fn in getattr(overrides, "_PREDICATE_OVERRIDE_FNS", ())
|
||||
}
|
||||
callables |= set(overrides.POST_PROCESS_PASSES)
|
||||
|
||||
fields = set()
|
||||
for fn in callables:
|
||||
name = getattr(fn, "__name__", "")
|
||||
body = _bodies(fn.__module__).get(name)
|
||||
# Loud, not silent: a body this scan cannot find is a field census it
|
||||
# is not taking, and a narrower census makes every check downstream of
|
||||
# it quietly vacuous.
|
||||
assert body is not None, f"{fn.__module__}.{name} has no body to scan"
|
||||
fields |= _returned_field_names(body)
|
||||
|
||||
# The literal arch -> {field: value} table, which has no callable at all.
|
||||
# It lives with the rest of the registry, in `model_override_base`.
|
||||
from sglang.srt.arg_groups import model_override_base
|
||||
|
||||
table_tree = ast.parse(
|
||||
pathlib.Path(model_override_base.__file__).read_text(encoding="utf-8-sig")
|
||||
)
|
||||
seen_table = False
|
||||
for node in table_tree.body:
|
||||
target = None
|
||||
if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
|
||||
target = node.targets[0].id
|
||||
elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
|
||||
target = node.target.id
|
||||
if target != "MODEL_OVERRIDES" or node.value is None:
|
||||
continue
|
||||
for inner in ast.walk(node.value):
|
||||
if not isinstance(inner, ast.Dict):
|
||||
continue
|
||||
for key, value in zip(inner.keys, inner.values):
|
||||
if isinstance(value, ast.Dict):
|
||||
continue
|
||||
if not isinstance(key, ast.Constant):
|
||||
raise AssertionError("non-literal override key")
|
||||
fields.add(key.value)
|
||||
seen_table = True
|
||||
assert seen_table, "MODEL_OVERRIDES is not where this scan looks for it"
|
||||
return fields
|
||||
|
||||
|
||||
def _written_after_publish():
|
||||
"""Fields the runtime overrides once the bags exist.
|
||||
|
||||
Imported from the supplied-instance ratchet rather than re-derived: it
|
||||
already enumerates `get_context().override(...)` and its named wrapper, and
|
||||
a second derivation of the same channel is what drifts narrower. A borrowed
|
||||
read of one of these answers with the startup value the same way a
|
||||
resolution-written one does -- the write just lands later.
|
||||
"""
|
||||
import importlib.util
|
||||
|
||||
companion = (
|
||||
pathlib.Path(__file__).resolve().parent
|
||||
/ "test_supplied_instance_exposure_ratchet.py"
|
||||
)
|
||||
spec = importlib.util.spec_from_file_location("_exposure_for_ratchet", companion)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return set(module.TestSuppliedInstanceExposure._override_written_fields())
|
||||
|
||||
|
||||
def _resolution_written():
|
||||
"""Every field the startup record answers wrong once it stays raw."""
|
||||
return (
|
||||
_declared_by_keyword()
|
||||
| _declared_by_registry_and_passes()
|
||||
| _written_after_publish()
|
||||
)
|
||||
|
||||
|
||||
def _borrowed_parking_spans(tree):
|
||||
"""[(span, parked attribute names)] for classes that park a borrowed record.
|
||||
|
||||
The supplied-instance census counts `self.server_args.<field>` only where
|
||||
the class was handed the record as a parameter -- `self.server_args =
|
||||
server_args` inside a method that takes one. A class that borrows it off
|
||||
another object instead (`self.server_args = scheduler.server_args`) is
|
||||
covered by neither census, and a read through that attribute is exactly the
|
||||
borrowed-record chain read this file is about. The parked name is whatever
|
||||
the class chose -- `self.args = scheduler.server_args` hides the same read,
|
||||
so the assignment target is recorded, not assumed.
|
||||
"""
|
||||
spans = []
|
||||
for cls in (n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)):
|
||||
parked_names = set()
|
||||
for node in ast.walk(cls):
|
||||
# Both assignment spellings -- `self.args = x.server_args` and the
|
||||
# annotated `self.args: ServerArgs = x.server_args`.
|
||||
if isinstance(node, ast.Assign):
|
||||
targets = node.targets
|
||||
elif isinstance(node, ast.AnnAssign) and node.value is not None:
|
||||
targets = [node.target]
|
||||
else:
|
||||
continue
|
||||
# A bare name on the right is the parameter the companion census
|
||||
# follows; an attribute chain ending in `.server_args` is a record
|
||||
# taken off another object.
|
||||
if not (
|
||||
isinstance(node.value, ast.Attribute)
|
||||
and node.value.attr == "server_args"
|
||||
):
|
||||
continue
|
||||
for target in targets:
|
||||
if (
|
||||
isinstance(target, ast.Attribute)
|
||||
and isinstance(target.value, ast.Name)
|
||||
and target.value.id == "self"
|
||||
):
|
||||
parked_names.add(target.attr)
|
||||
if parked_names:
|
||||
spans.append(((cls.lineno, cls.end_lineno), parked_names))
|
||||
return spans
|
||||
|
||||
|
||||
def _subtrees_with_their_own_record():
|
||||
"""Top-level package directories that define a second `ServerArgs`.
|
||||
|
||||
`multimodal_gen` ships one, so `x.server_args.<field>` inside it names a
|
||||
field of *that* class -- and `model_path` is a field of both. Keying on the
|
||||
spelling alone once put a resolution call into a diffusion entry point.
|
||||
Derived from the class definitions rather than named here, so a third
|
||||
record would be excluded the same way instead of silently counting.
|
||||
"""
|
||||
roots = set()
|
||||
for path in _PACKAGE.rglob("*.py"):
|
||||
rel = path.relative_to(_PACKAGE).as_posix()
|
||||
if rel.startswith("srt/") or "/" not in rel:
|
||||
continue
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "ServerArgs" not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
continue
|
||||
if any(
|
||||
isinstance(node, ast.ClassDef) and node.name == "ServerArgs"
|
||||
for node in ast.walk(tree)
|
||||
):
|
||||
roots.add(rel.split("/")[0])
|
||||
return roots
|
||||
|
||||
|
||||
def _reads_the_startup_record(tree):
|
||||
"""True when this module's `server_args` is the one `srt` resolves.
|
||||
|
||||
Inside a subtree that owns another record, only a module that imports the
|
||||
startup record is talking about it.
|
||||
"""
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.ImportFrom) and node.module:
|
||||
if node.module.startswith("sglang.srt"):
|
||||
# The *original* names: `ServerArgs as SrtServerArgs` is still
|
||||
# the startup record, whatever this module calls it.
|
||||
names = {alias.name for alias in node.names}
|
||||
if names & {"ServerArgs", "server_args", "prepare_server_args"}:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _chain_reads(written):
|
||||
"""`<expression>.server_args.<field>` where the field is resolution-written."""
|
||||
found = []
|
||||
other_records = _subtrees_with_their_own_record()
|
||||
for path in sorted(_READS_SCANNED.rglob("*.py")):
|
||||
rel = path.relative_to(_READS_SCANNED).as_posix()
|
||||
in_srt = rel.startswith("srt/")
|
||||
if in_srt:
|
||||
under_srt = rel[len("srt/") :]
|
||||
if path.name in _OWNERS or under_srt.startswith(_OWNERS[-1]):
|
||||
continue
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "server_args" not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
raise AssertionError(f"unparsable module in the census: {rel}")
|
||||
if (
|
||||
not in_srt
|
||||
and rel.split("/")[0] in other_records
|
||||
and not _reads_the_startup_record(tree)
|
||||
):
|
||||
continue
|
||||
parked = _borrowed_parking_spans(tree)
|
||||
|
||||
def parked_alias(lineno, name):
|
||||
return any(
|
||||
start <= lineno <= end and name in names
|
||||
for (start, end), names in parked
|
||||
)
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if not (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
and node.attr in written
|
||||
):
|
||||
continue
|
||||
base = node.value
|
||||
if not isinstance(base, ast.Attribute):
|
||||
continue
|
||||
through_self = isinstance(base.value, ast.Name) and base.value.id == "self"
|
||||
if base.attr == "server_args":
|
||||
# `self.server_args.field` is the parked spelling the
|
||||
# supplied-instance census counts -- but only where the record
|
||||
# arrived as a parameter.
|
||||
if through_self and not parked_alias(node.lineno, base.attr):
|
||||
continue
|
||||
elif not (through_self and parked_alias(node.lineno, base.attr)):
|
||||
# Any other attribute counts only as a recorded parked alias
|
||||
# (`self.args = scheduler.server_args` and later `self.args.x`).
|
||||
continue
|
||||
suffix = " (parked borrowed record)" if through_self else ""
|
||||
found.append(f"{rel}:{node.lineno} {ast.unparse(base)}.{node.attr}{suffix}")
|
||||
# `getattr(model_runner.server_args, "field", default)` reads the same
|
||||
# borrowed record through a `Call`, with the same stale default.
|
||||
for node in ast.walk(tree):
|
||||
if not (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "getattr"
|
||||
and len(node.args) >= 2
|
||||
and isinstance(node.args[1], ast.Constant)
|
||||
and node.args[1].value in written
|
||||
):
|
||||
continue
|
||||
base = node.args[0]
|
||||
if not isinstance(base, ast.Attribute):
|
||||
continue
|
||||
through_self = isinstance(base.value, ast.Name) and base.value.id == "self"
|
||||
if base.attr == "server_args":
|
||||
if through_self and not parked_alias(node.lineno, base.attr):
|
||||
continue
|
||||
elif not (through_self and parked_alias(node.lineno, base.attr)):
|
||||
continue
|
||||
found.append(
|
||||
f"{rel}:{node.lineno} getattr({ast.unparse(base)}, "
|
||||
f"{node.args[1].value!r})"
|
||||
)
|
||||
return sorted(found)
|
||||
|
||||
|
||||
def _passes_named_at_call_sites() -> set:
|
||||
"""Names passed to ``run_post_process_pass(sa, fn)`` anywhere in the tree.
|
||||
|
||||
A call whose pass is not a bare name is a hard failure, not a skip: this
|
||||
scan is the ground truth every registry-driven check below is derived from,
|
||||
so `run_post_process_pass(self, overrides._new_pass)` (an `ast.Attribute`)
|
||||
or `run_post_process_pass(self, fn=_new_pass)` (a keyword) would otherwise
|
||||
walk past all of them silently. Keeping the call shape uniform is the
|
||||
price of the scan being complete.
|
||||
"""
|
||||
names = set()
|
||||
for path in sorted(pathlib.Path(next(iter(sglang.__path__))).rglob("*.py")):
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "run_post_process_pass" not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
continue
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
func = node.func
|
||||
if isinstance(func, ast.Name):
|
||||
called = func.id
|
||||
elif isinstance(func, ast.Attribute):
|
||||
called = func.attr
|
||||
else:
|
||||
called = None
|
||||
if called != "run_post_process_pass":
|
||||
continue
|
||||
if (
|
||||
len(node.args) != 2
|
||||
or node.keywords
|
||||
or not isinstance(node.args[1], ast.Name)
|
||||
):
|
||||
raise AssertionError(
|
||||
f"{path}:{node.lineno}: run_post_process_pass takes the pass "
|
||||
"as a bare name in its second positional argument; "
|
||||
f"{ast.unparse(node)!r} is invisible to this scan and to "
|
||||
"every registry-driven check derived from it"
|
||||
)
|
||||
names.add(node.args[1].id)
|
||||
return names
|
||||
|
||||
|
||||
class TestEveryInvokedPassIsRegistered(CustomTestCase):
|
||||
"""The registry is what the scans above enumerate, so a pass missing from it
|
||||
is a pass nothing checks.
|
||||
|
||||
Being invoked and being registered are two edits, and `_a2a_fusion_adjustments`
|
||||
shipped with only the first: it ran in production while the registry-driven
|
||||
scans walked past it. The call sites are the ground truth here -- the registry
|
||||
is derived from a decorator someone has to remember.
|
||||
"""
|
||||
|
||||
def test_the_registry_covers_every_call_site(self):
|
||||
from sglang.srt.arg_groups import overrides
|
||||
|
||||
invoked = _passes_named_at_call_sites()
|
||||
self.assertGreater(
|
||||
len(invoked),
|
||||
20,
|
||||
f"only {len(invoked)} call sites found; the scan is broken, not the tree",
|
||||
)
|
||||
registered = {fn.__name__ for fn in overrides.POST_PROCESS_PASSES}
|
||||
self.assertEqual(
|
||||
set(),
|
||||
invoked - registered,
|
||||
"these passes are invoked but carry no @register_post_process, so "
|
||||
"every check that walks POST_PROCESS_PASSES skips them",
|
||||
)
|
||||
self.assertEqual(
|
||||
set(),
|
||||
registered - invoked,
|
||||
"these passes carry @register_post_process but no slot invokes "
|
||||
"them; deleting a call site and leaving the decorator behind "
|
||||
"leaves a pass that only the scans can see",
|
||||
)
|
||||
|
||||
|
||||
class TestNoChainReadsOfResolvedConfig(CustomTestCase):
|
||||
def test_the_census_has_something_to_count(self):
|
||||
"""A written set that collapsed would make the pin vacuous.
|
||||
|
||||
Each mechanism is checked on its own, because they fail
|
||||
independently: the keyword scan cannot see a field name that is data,
|
||||
and a scan for `@register_model_override*` sees one provider out of
|
||||
twenty-seven because the rest register through a helper. A hand-written
|
||||
expectation of the resulting field names is what hid that -- it stayed
|
||||
green while a whole channel went unscanned -- so each mechanism is
|
||||
pinned by a floor derived from the live registry instead.
|
||||
"""
|
||||
from sglang.srt.arg_groups import overrides
|
||||
|
||||
by_keyword = _declared_by_keyword()
|
||||
by_data = _declared_by_registry_and_passes()
|
||||
|
||||
self.assertGreater(
|
||||
len(by_keyword),
|
||||
100,
|
||||
f"only {len(by_keyword)} fields are declared by keyword; the scan broke",
|
||||
)
|
||||
providers = {fn for fns in overrides._MODEL_OVERRIDE_FNS.values() for fn in fns}
|
||||
providers |= {
|
||||
fn for _predicate, fn in getattr(overrides, "_PREDICATE_OVERRIDE_FNS", ())
|
||||
}
|
||||
self.assertGreater(
|
||||
len(providers) + len(overrides.POST_PROCESS_PASSES),
|
||||
50,
|
||||
"the registry and pass tables collapsed; the data-channel scan is "
|
||||
"reading an empty registry",
|
||||
)
|
||||
self.assertGreater(
|
||||
len(by_data),
|
||||
25,
|
||||
f"only {len(by_data)} fields come from the registry and the passes, "
|
||||
f"across {len(providers)} providers and "
|
||||
f"{len(overrides.POST_PROCESS_PASSES)} passes; the scan of the "
|
||||
"dict-key channel broke",
|
||||
)
|
||||
# The data channel is not the keyword scan's subset: if it became one,
|
||||
# that scan would be doing all the work and a regression here would be
|
||||
# invisible.
|
||||
self.assertTrue(by_data - by_keyword, "the data channel adds nothing")
|
||||
|
||||
def test_nothing_reads_a_resolved_field_off_a_borrowed_record(self):
|
||||
found = _chain_reads(_resolution_written())
|
||||
self.assertEqual(
|
||||
found,
|
||||
[],
|
||||
"these reach the startup record through another object for a value "
|
||||
"resolution decides, so they answer with the CLI default once the "
|
||||
"record stays raw; read the config bag instead:\n " + "\n ".join(found),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,874 +0,0 @@
|
||||
"""The supplied-instance surface is measured on two axes, and may only shrink.
|
||||
|
||||
A callee that takes ``server_args`` keeps the supplied-instance contract: the
|
||||
caller chose the object, so no global-read ratchet counts it. Step 12 changes
|
||||
what that object *carries* — the instance stays at the user's raw input — so a
|
||||
callee reading a field **resolution fills in** would start seeing the CLI default
|
||||
instead of the effective value.
|
||||
|
||||
This pins that intersection. Each entry is one (file, field) pair where a
|
||||
parameter named ``server_args`` is read for a field resolution writes; the plan
|
||||
doc carries the proposed disposition per field
|
||||
(``global_context/12-raw-input-config.md``, "the supplied-instance conversion
|
||||
list"). New pairs fail: a new one is new step-12 work, and the moment to decide
|
||||
where the value should come from is when the read is written, not during the
|
||||
flip. Pairs that disappear also fail, with the entry to delete — the list is the
|
||||
measurement, not a memory of one.
|
||||
|
||||
The written-field set is derived here rather than hardcoded: the
|
||||
representative configs in ``_MATRIX`` (one per resolution family it exercises)
|
||||
are resolved and compared against the dataclass defaults, the same matrix the
|
||||
context repo's audit tool uses. Ambient environment is normalized per entry --
|
||||
resolution branches on CI detection and leaves sticky process state, so each
|
||||
entry resolves from the pristine snapshot, and the CI shape is an explicit
|
||||
entry rather than an accident of the runner. The read scan mirrors that
|
||||
tool's three shapes — a parameter attribute, ``getattr(server_args, "literal")``,
|
||||
and the parameter parked on ``self`` — because two implementations of one census
|
||||
that disagree are worse than either alone.
|
||||
|
||||
The second axis is **already wrong today**, not after a flip. Some config is
|
||||
decided *after* publish and recorded with ``get_context().override(...)`` —
|
||||
elastic-EP resizing `ep_size`, a weight update rewriting `model_path` /
|
||||
`load_format`, HiCache attach naming a storage backend, adaptive speculative
|
||||
decoding moving `speculative_num_steps`. That write reaches the bags and never
|
||||
the record, so a supplied-instance read of one of those fields answers with the
|
||||
startup value from the moment the override lands. Whether that is a defect
|
||||
depends on ordering — a value copied at construction, before any override, is
|
||||
fine — so this axis is pinned as a measurement with the same growth guard rather
|
||||
than as a list of bugs. One of them *was* a defect and is fixed at the base of
|
||||
this stack: the linear-attn dispatch table rebuilt itself from the record after
|
||||
the SM100 GDN prefill decision had been recorded in the bag, so a second runner's
|
||||
rebuild dropped it. That choice is a per-runner stamp now and is not recorded
|
||||
process-wide at all, so neither the read nor the field is on this axis.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import msgspec
|
||||
import msgspec.structs
|
||||
|
||||
import sglang
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=23, suite="base-a-test-cpu")
|
||||
# Also on a CUDA runner: the written set is derived by resolving on the running
|
||||
# host, and `is_cuda()` / capability gates only open on real hardware. The pin
|
||||
# is split by host so both registrations stay exact: `_EXPOSED` is asserted
|
||||
# everywhere, and a pair whose write only happens on CUDA belongs in
|
||||
# `_EXPOSED_CUDA_ONLY` -- pinned on the CUDA runner, invisible to the CPU
|
||||
# assertion. Without the split, one shared exact list could not hold such a
|
||||
# pair at all: pinning it fails the CPU run as "gone", omitting it fails the
|
||||
# CUDA run as "new". (No AMD registration: an `is_hip()`-gated write would
|
||||
# shift the exact sets in ways none of the pinning hosts can verify; the ROCm
|
||||
# resolution surface is covered by `test_resolution_is_reproducible.py`
|
||||
# instead, whose assertion is device-agnostic.)
|
||||
register_cuda_ci(est_time=16, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
_PACKAGE_ROOT = Path(next(iter(sglang.__path__))) / "srt"
|
||||
|
||||
# The config the resolution pipeline owns; reading the in-flight record is their
|
||||
# job, not a supplied-instance read.
|
||||
_OWNERS = ("server_args.py", "runtime_context.py", "arg_groups/")
|
||||
|
||||
_MINI_CONFIG = {
|
||||
"architectures": ["LlamaForCausalLM"],
|
||||
"model_type": "llama",
|
||||
"hidden_size": 16,
|
||||
"intermediate_size": 32,
|
||||
"num_attention_heads": 2,
|
||||
"num_key_value_heads": 2,
|
||||
"num_hidden_layers": 2,
|
||||
"vocab_size": 128,
|
||||
"max_position_embeddings": 2048,
|
||||
}
|
||||
|
||||
# One config resolves only its own decisions, so the written set is a union.
|
||||
_MATRIX = (
|
||||
{},
|
||||
{
|
||||
"speculative_algorithm": "EAGLE",
|
||||
"speculative_num_steps": 3,
|
||||
"speculative_eagle_topk": 1,
|
||||
"speculative_num_draft_tokens": 4,
|
||||
},
|
||||
{"dp_size": 2, "tp_size": 2, "enable_dp_attention": True},
|
||||
# DWDP resolves dp_size and enable_dp_attention *itself* -- the plain DP
|
||||
# entry above passes them in, and passed-in fields are excluded from the
|
||||
# written set, so without this entry the dp_size readers would never pin.
|
||||
{"tp_size": 2, "dwdp_size": 2},
|
||||
{"enable_hierarchical_cache": True, "hicache_ratio": 2.0},
|
||||
{"disaggregation_mode": "prefill"},
|
||||
{"tp_size": 2, "attn_cp_size": 2},
|
||||
{"enable_lora": True, "max_lora_rank": 16},
|
||||
{"kv_cache_dtype": "fp8_e4m3", "page_size": 64},
|
||||
# MIS resolves disable_radix_cache (and friends) itself; the backend is
|
||||
# passed in because the handler asserts flashinfer rather than switching.
|
||||
{"enable_mis": True, "attention_backend": "flashinfer"},
|
||||
)
|
||||
|
||||
# `declare_resolution` call sites whose keyword expansion is built
|
||||
# dynamically; the written fields are spelled out here and drift-guarded.
|
||||
_LATE_RESOLUTION_DYNAMIC_SITES = {
|
||||
"parser/template_detection.py": frozenset({"reasoning_parser", "tool_call_parser"}),
|
||||
}
|
||||
|
||||
# `get_context().override(...)` declares through the same seam, but the fields
|
||||
# are the caller's -- a test names them one call at a time. There is no static
|
||||
# set to collect, and nothing resolution decides: whatever a caller overrides
|
||||
# there is exposure only through that caller's own reads.
|
||||
_CALLER_SUPPLIED_LATE_SITES = frozenset({"runtime_context.py"})
|
||||
|
||||
# Resolution also branches on ambient environment; those shapes are explicit
|
||||
# entries so the written set is the same on every host. `SGLANG_IS_IN_CI`
|
||||
# makes resolution fill `soft_watchdog_timeout`.
|
||||
_ENV_MATRIX = (({}, {"SGLANG_IS_IN_CI": "true"}),)
|
||||
|
||||
# Only true constructor inputs: `tokenizer_path` / `served_model_name` are
|
||||
# resolution-written (filled from `model_path` when unset), so their readers
|
||||
# are step-12 exposure like any other pair.
|
||||
_PASSED = frozenset({"model_path", "device", "random_seed"})
|
||||
|
||||
# Empty. A pair belongs here when a reader has no bag to read -- it runs before
|
||||
# its process publishes -- and cannot use `resolving_view` either. The launcher's
|
||||
# pre-publish reads (`_set_envs_and_config`, the auto-parser gate) and the
|
||||
# late-resolution detection it calls all read the declarations now, so nothing
|
||||
# qualifies. A new entry needs that kind of reason next to it.
|
||||
_EXPOSED: frozenset = frozenset()
|
||||
|
||||
# Pairs whose resolution write only happens on a CUDA host (capability or
|
||||
# `is_cuda()` gated): asserted on the CUDA registration, invisible to the CPU
|
||||
# one. Empty today -- the current written sets coincide across the two hosts --
|
||||
# but this is where a GPU-only write's readers get pinned without breaking the
|
||||
# CPU-exact assertion.
|
||||
_EXPOSED_CUDA_ONLY: frozenset = frozenset()
|
||||
|
||||
|
||||
# Axis two: (file, field) pairs where a supplied-instance read names a field that
|
||||
# some code overrides post-publish. Each needs an ordering judgment, not a blanket
|
||||
# conversion; the list exists so a new one is a decision made when it is written.
|
||||
_OVERRIDDEN_AND_READ: frozenset = frozenset()
|
||||
|
||||
|
||||
def _expanded_override_keys(rel, tree, call, kw) -> set:
|
||||
"""The statically visible keys behind an ``override(..., **expr)``.
|
||||
|
||||
Handles a dict literal, a conditional between dict literals, and a name
|
||||
bound to a dict literal in the enclosing function (plus constant-subscript
|
||||
stores onto it -- the HiCache attach shape). One expansion is unresolvable
|
||||
by design and exempted by name: ``update_server_args`` forwards
|
||||
operator-chosen fields, so its key set is the API's, not this file's.
|
||||
Anything else unresolvable fails -- a silently skipped expansion would
|
||||
shrink the written set.
|
||||
"""
|
||||
for a in call.args:
|
||||
if isinstance(a, ast.Constant) and a.value == "update_server_args":
|
||||
return set()
|
||||
for k in call.keywords:
|
||||
if (
|
||||
k.arg == "source"
|
||||
and isinstance(k.value, ast.Constant)
|
||||
and k.value.value == "update_server_args"
|
||||
):
|
||||
return set()
|
||||
|
||||
def loop_variable_values(name: str) -> set:
|
||||
"""The values a `for name, ... in (<literal tuples>)` loop binds.
|
||||
|
||||
A handler that records one field per loop iteration spells the field
|
||||
names in the loop's own literal, so they are still static.
|
||||
"""
|
||||
values = set()
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.For):
|
||||
continue
|
||||
target = node.target
|
||||
names = (
|
||||
[target]
|
||||
if isinstance(target, ast.Name)
|
||||
else list(getattr(target, "elts", []))
|
||||
)
|
||||
if not names or not isinstance(names[0], ast.Name) or names[0].id != name:
|
||||
continue
|
||||
if not (node.lineno <= call.lineno <= (node.end_lineno or node.lineno)):
|
||||
continue
|
||||
for item in getattr(node.iter, "elts", []):
|
||||
first = (
|
||||
item.elts[0] if isinstance(item, ast.Tuple) and item.elts else item
|
||||
)
|
||||
if isinstance(first, ast.Constant) and isinstance(first.value, str):
|
||||
values.add(first.value)
|
||||
return values
|
||||
|
||||
def dict_keys(node) -> set:
|
||||
assert isinstance(node, ast.Dict), (
|
||||
f"non-literal dict in override expansion at {rel}:{call.lineno}"
|
||||
)
|
||||
keys = set()
|
||||
for key in node.keys:
|
||||
if isinstance(key, ast.Constant):
|
||||
keys.add(key.value)
|
||||
continue
|
||||
assert isinstance(key, ast.Name), (
|
||||
f"non-literal dict key in override expansion at {rel}:{call.lineno}"
|
||||
)
|
||||
bound = loop_variable_values(key.id)
|
||||
assert bound, (
|
||||
f"dict key {key.id!r} at {rel}:{call.lineno} is not bound by a "
|
||||
"literal loop; extend the resolver"
|
||||
)
|
||||
keys |= bound
|
||||
return keys
|
||||
|
||||
if isinstance(kw.value, ast.Dict):
|
||||
return dict_keys(kw.value)
|
||||
if isinstance(kw.value, ast.IfExp):
|
||||
keys = set()
|
||||
for branch in (kw.value.body, kw.value.orelse):
|
||||
if isinstance(branch, ast.Dict) and branch.keys:
|
||||
keys |= dict_keys(branch)
|
||||
elif isinstance(branch, ast.Dict):
|
||||
pass
|
||||
else:
|
||||
raise AssertionError(
|
||||
f"unresolvable override expansion at {rel}:{call.lineno}"
|
||||
)
|
||||
return keys
|
||||
assert isinstance(kw.value, ast.Name), (
|
||||
f"unresolvable override expansion at {rel}:{call.lineno}"
|
||||
)
|
||||
name = kw.value.id
|
||||
enclosing = None
|
||||
for fn in ast.walk(tree):
|
||||
if isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
if (
|
||||
fn.lineno
|
||||
<= call.lineno
|
||||
<= max(getattr(fn, "end_lineno", fn.lineno), fn.lineno)
|
||||
):
|
||||
if enclosing is None or fn.lineno > enclosing.lineno:
|
||||
enclosing = fn
|
||||
assert enclosing is not None, (
|
||||
f"override expansion outside any function at {rel}:{call.lineno}"
|
||||
)
|
||||
keys = set()
|
||||
found = False
|
||||
for node in ast.walk(enclosing):
|
||||
if (
|
||||
isinstance(node, ast.Assign)
|
||||
and len(node.targets) == 1
|
||||
and isinstance(node.targets[0], ast.Name)
|
||||
and node.targets[0].id == name
|
||||
and isinstance(node.value, ast.Dict)
|
||||
):
|
||||
found = True
|
||||
keys |= dict_keys(node.value)
|
||||
elif (
|
||||
isinstance(node, ast.Assign)
|
||||
and len(node.targets) == 1
|
||||
and isinstance(node.targets[0], ast.Subscript)
|
||||
and isinstance(node.targets[0].value, ast.Name)
|
||||
and node.targets[0].value.id == name
|
||||
and isinstance(node.targets[0].slice, ast.Constant)
|
||||
):
|
||||
keys.add(node.targets[0].slice.value)
|
||||
assert found, (
|
||||
f"override expansion '{name}' at {rel}:{call.lineno} has no "
|
||||
"dict-literal assignment in its function; extend the resolver"
|
||||
)
|
||||
return keys
|
||||
|
||||
|
||||
class TestSuppliedInstanceExposure(CustomTestCase):
|
||||
def _callTestMethod(self, method):
|
||||
# No CI retry: a failed first attempt has already resolved the matrix
|
||||
# and mutated process state; a retry against that contamination could
|
||||
# pass on a drifted written set or mask a real drift.
|
||||
return unittest.TestCase._callTestMethod(self, method)
|
||||
|
||||
def setUp(self):
|
||||
# Resolving the matrix writes process state on the way through (the
|
||||
# multimodal transport handler sets SGLANG_USE_CUDA_IPC_TRANSPORT, and
|
||||
# `EnvField.set()` flips a descriptor flag `os.environ` does not carry).
|
||||
# Leaking it makes *later* files in the same worker fail, which is how
|
||||
# this was found -- so the case restores what it touched.
|
||||
super().setUp()
|
||||
state = (dict(os.environ), self._env_field_flags())
|
||||
self.addCleanup(self._restore_process_state, state)
|
||||
|
||||
@staticmethod
|
||||
def _env_field_flags() -> dict:
|
||||
from sglang.srt.environ import EnvField, envs
|
||||
|
||||
flags = {}
|
||||
for klass in reversed(type(envs).__mro__):
|
||||
for name, field in vars(klass).items():
|
||||
if isinstance(field, EnvField):
|
||||
flags[name] = field._set_to_none
|
||||
return flags
|
||||
|
||||
@staticmethod
|
||||
def _restore_process_state(state) -> None:
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
saved_environ, saved_flags = state
|
||||
os.environ.clear()
|
||||
os.environ.update(saved_environ)
|
||||
for name, was_none in saved_flags.items():
|
||||
getattr(type(envs), name)._set_to_none = was_none
|
||||
|
||||
def _config_dir(self) -> str:
|
||||
config_dir = tempfile.mkdtemp(prefix="supplied_instance_")
|
||||
self.addCleanup(shutil.rmtree, config_dir, ignore_errors=True)
|
||||
with open(os.path.join(config_dir, "config.json"), "w") as handle:
|
||||
json.dump(_MINI_CONFIG, handle)
|
||||
return config_dir
|
||||
|
||||
def _resolution_written_fields(self) -> set:
|
||||
"""The union of what resolution fills in across the matrix.
|
||||
|
||||
Every entry must resolve. A silently skipped one would shrink this set,
|
||||
which makes pinned pairs look like they disappeared -- the list would
|
||||
then drift by environment rather than by code, and the failure would
|
||||
point at the wrong thing. Each entry resolves from the pristine
|
||||
process snapshot (resolution writes env and EnvField flags on the way
|
||||
through, and DWDP flips `SGLANG_SCHEDULER_SKIP_ALL_GATHER`), so the
|
||||
union does not depend on matrix order; and the ambient CI marker is
|
||||
cleared, so a runner's identity cannot leak into the measurement --
|
||||
the CI-conditioned writes come from `_ENV_MATRIX`'s explicit entry.
|
||||
Declarers outside `arg_groups/` count too: the parser auto-detection
|
||||
runs at launcher stage and the NPU helper is called by the pipeline, so
|
||||
their target fields are collected statically from the call sites --
|
||||
resolution writes by definition, just not reached by the matrix.
|
||||
"""
|
||||
pristine = (dict(os.environ), self._env_field_flags())
|
||||
written = set()
|
||||
|
||||
def resolve_one(extra, env):
|
||||
self._restore_process_state(pristine)
|
||||
os.environ.pop("SGLANG_IS_IN_CI", None)
|
||||
os.environ.update(env)
|
||||
model_path = self._config_dir()
|
||||
try:
|
||||
resolved = ServerArgs(
|
||||
model_path=model_path, device="cuda", random_seed=42, **extra
|
||||
)
|
||||
resolved.resolve_once()
|
||||
except Exception as exc:
|
||||
self.fail(
|
||||
f"the matrix entry {extra} (env={env}) did not resolve in "
|
||||
f"this environment ({type(exc).__name__}: {exc}); the "
|
||||
"written-field union would be short and the pinned list "
|
||||
"would drift"
|
||||
)
|
||||
defaults = {}
|
||||
for field in msgspec.structs.fields(resolved):
|
||||
if field.default is not msgspec.NODEFAULT:
|
||||
defaults[field.name] = field.default
|
||||
elif field.default_factory is not msgspec.NODEFAULT:
|
||||
defaults[field.name] = field.default_factory()
|
||||
for field_name, default in defaults.items():
|
||||
if field_name in _PASSED or field_name in extra:
|
||||
continue
|
||||
if getattr(resolved, field_name) != default:
|
||||
written.add(field_name)
|
||||
|
||||
for extra in _MATRIX:
|
||||
resolve_one(extra, {})
|
||||
for extra, env in _ENV_MATRIX:
|
||||
resolve_one(extra, env)
|
||||
self._restore_process_state(pristine)
|
||||
written |= self._declared_outside_the_pipeline()
|
||||
written |= self._hook_assignment_targets()
|
||||
written |= self._record_method_assignment_targets()
|
||||
written |= self._declarative_override_fields()
|
||||
return written
|
||||
|
||||
def _hook_assignment_targets(self) -> set:
|
||||
"""Fields any resolution hook can write, collected statically.
|
||||
|
||||
The matrix can only enumerate families someone thought to add -- the
|
||||
DFLASH hole (its hook is the sole writer of
|
||||
`speculative_draft_attention_backend`, and no entry ran it) showed
|
||||
that a family nobody listed leaves its readers unpinned. The hook
|
||||
modules under `arg_groups/` are the resolution pipeline's extension
|
||||
points -- along with the NPU default helper, which the pipeline calls
|
||||
the same way -- and their write surface is the may-write set,
|
||||
family-blind by
|
||||
construction. A hook writes two ways: `server_args.field = ...`, and
|
||||
`declare_resolution(server_args, source, field=...)`, which records
|
||||
the write in the declaration stash on its way to the field. Counting
|
||||
only the assignment would read a hook's conversion to a declaration as
|
||||
the field having stopped being written. Collected like the
|
||||
late-resolution keywords: statically, failing loudly on an
|
||||
unparsable module. Underscore-prefixed targets are pipeline
|
||||
bookkeeping, not config leaves.
|
||||
"""
|
||||
targets = set()
|
||||
modules = sorted((_PACKAGE_ROOT / "arg_groups").glob("*.py"))
|
||||
modules.append(_PACKAGE_ROOT / "hardware_backend/npu/utils.py")
|
||||
for path in modules:
|
||||
try:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
except SyntaxError:
|
||||
self.fail(f"unparsable hook module in the census: {path.name}")
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Assign):
|
||||
tgts = node.targets
|
||||
elif isinstance(node, (ast.AnnAssign, ast.AugAssign)):
|
||||
tgts = [node.target]
|
||||
elif (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "declare_resolution"
|
||||
):
|
||||
targets |= {
|
||||
kw.arg
|
||||
for kw in node.keywords
|
||||
if kw.arg and not kw.arg.startswith("_")
|
||||
}
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
for tgt in tgts:
|
||||
if (
|
||||
isinstance(tgt, ast.Attribute)
|
||||
and isinstance(tgt.value, ast.Name)
|
||||
and tgt.value.id == "server_args"
|
||||
and not tgt.attr.startswith("_")
|
||||
):
|
||||
targets.add(tgt.attr)
|
||||
return targets
|
||||
|
||||
def _record_method_assignment_targets(self) -> set:
|
||||
"""Fields ``ServerArgs``'s own methods can write, collected statically.
|
||||
|
||||
The record's handlers are as family-conditional as the hooks -- the
|
||||
mooncake/layer_first layout rewrite, the deepseek-EP mode defaults,
|
||||
the seed fill that only runs when the caller did *not* supply one (so
|
||||
construct-and-diff can never see it: measuring requires supplying).
|
||||
A write site that can never fire is a dead branch to delete upstream,
|
||||
not a census exemption. Only names that are declared dataclass fields
|
||||
count; underscore bookkeeping does not.
|
||||
|
||||
Two spellings write: an assignment, and ``self._declare(source,
|
||||
field=value)``, which records the write in the declaration stash on
|
||||
its way to the field. Counting only assignments would read a handler's
|
||||
conversion to a declaration as the field having stopped being written,
|
||||
which would quietly retire every pinned pair that reads it.
|
||||
"""
|
||||
tree = ast.parse(
|
||||
(_PACKAGE_ROOT / "server_args.py").read_text(encoding="utf-8-sig")
|
||||
)
|
||||
sa_class = next(
|
||||
node
|
||||
for node in tree.body
|
||||
if isinstance(node, ast.ClassDef) and node.name == "ServerArgs"
|
||||
)
|
||||
declared = {
|
||||
node.target.id
|
||||
for node in sa_class.body
|
||||
if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name)
|
||||
}
|
||||
targets = set()
|
||||
for node in ast.walk(sa_class):
|
||||
if isinstance(node, ast.Assign):
|
||||
tgts = node.targets
|
||||
elif isinstance(node, (ast.AnnAssign, ast.AugAssign)):
|
||||
tgts = [node.target]
|
||||
elif (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "declare_resolution"
|
||||
):
|
||||
targets |= {
|
||||
kw.arg
|
||||
for kw in node.keywords
|
||||
if kw.arg in declared and not kw.arg.startswith("_")
|
||||
}
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
for tgt in tgts:
|
||||
if (
|
||||
isinstance(tgt, ast.Attribute)
|
||||
and isinstance(tgt.value, ast.Name)
|
||||
and tgt.value.id == "self"
|
||||
and tgt.attr in declared
|
||||
and not tgt.attr.startswith("_")
|
||||
):
|
||||
targets.add(tgt.attr)
|
||||
# The deprecated-alias normalization declares through `**renamed`, so
|
||||
# the keyword scan sees no names; its field set is pinned here.
|
||||
alias_fields = {
|
||||
"attention_backend",
|
||||
"decode_attention_backend",
|
||||
"prefill_attention_backend",
|
||||
"speculative_draft_attention_backend",
|
||||
}
|
||||
|
||||
# The handler lives in `arg_groups/serving_hook.py`, reached either as a
|
||||
# record method or as a bare-name call, so look the loop up by both.
|
||||
def _deprecated_alias_handler():
|
||||
for node in ast.walk(sa_class):
|
||||
if (
|
||||
isinstance(node, ast.FunctionDef)
|
||||
and node.name == "_handle_deprecated_args"
|
||||
and any(isinstance(n, ast.For) for n in ast.walk(node))
|
||||
):
|
||||
return node
|
||||
for path in sorted((_PACKAGE_ROOT / "arg_groups").glob("*.py")):
|
||||
tree = ast.parse(path.read_text(encoding="utf-8-sig"))
|
||||
for node in tree.body:
|
||||
if (
|
||||
isinstance(node, ast.FunctionDef)
|
||||
and node.name == "handle_deprecated_args"
|
||||
):
|
||||
return node
|
||||
raise AssertionError("the deprecated-alias handler was not found")
|
||||
|
||||
deprecated = _deprecated_alias_handler()
|
||||
found_tuples = [
|
||||
{elt.value for elt in node.iter.elts if isinstance(elt, ast.Constant)}
|
||||
for node in ast.walk(deprecated)
|
||||
if isinstance(node, ast.For) and isinstance(node.iter, ast.Tuple)
|
||||
]
|
||||
self.assertIn(
|
||||
alias_fields,
|
||||
found_tuples,
|
||||
"the deprecated-alias normalization loop moved or changed its "
|
||||
"field tuple; update alias_fields to match",
|
||||
)
|
||||
return targets | alias_fields
|
||||
|
||||
def _declarative_override_fields(self) -> set:
|
||||
"""Fields the declarative override registry can write.
|
||||
|
||||
``MODEL_OVERRIDES`` maps arch -> {field: value}, and the
|
||||
``@register_model_override``(-``_predicate``) providers return (or
|
||||
build by subscript) {field: value} dicts, which go straight into the
|
||||
declaration stash, so no assignment scan sees these writes and a
|
||||
llama-only matrix never triggers them. Keys must be
|
||||
string literals; anything else fails loudly.
|
||||
"""
|
||||
tree = ast.parse(
|
||||
(_PACKAGE_ROOT / "arg_groups" / "overrides.py").read_text(
|
||||
encoding="utf-8-sig"
|
||||
)
|
||||
)
|
||||
fields = set()
|
||||
for node in tree.body:
|
||||
target = None
|
||||
if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
|
||||
target = node.targets[0].id
|
||||
elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
|
||||
target = node.target.id
|
||||
if target != "MODEL_OVERRIDES" or node.value is None:
|
||||
continue
|
||||
for inner in ast.walk(node.value):
|
||||
if not isinstance(inner, ast.Dict):
|
||||
continue
|
||||
for key, value in zip(inner.keys, inner.values):
|
||||
if isinstance(value, ast.Dict):
|
||||
continue # arch -> {…} outer layer
|
||||
self.assertIsInstance(key, ast.Constant, "non-literal override key")
|
||||
fields.add(key.value)
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.FunctionDef):
|
||||
continue
|
||||
if not any(
|
||||
isinstance(dec, ast.Call)
|
||||
and isinstance(dec.func, ast.Name)
|
||||
and dec.func.id.startswith("register_model_override")
|
||||
for dec in node.decorator_list
|
||||
):
|
||||
continue
|
||||
for inner in ast.walk(node):
|
||||
if isinstance(inner, ast.Assign) and isinstance(
|
||||
inner.targets[0], ast.Subscript
|
||||
):
|
||||
key = inner.targets[0].slice
|
||||
self.assertIsInstance(
|
||||
key, ast.Constant, f"non-literal override key in {node.name}"
|
||||
)
|
||||
fields.add(key.value)
|
||||
if isinstance(inner, ast.Dict):
|
||||
for key in inner.keys:
|
||||
self.assertIsInstance(
|
||||
key,
|
||||
ast.Constant,
|
||||
f"non-literal override key in {node.name}",
|
||||
)
|
||||
fields.add(key.value)
|
||||
return fields
|
||||
|
||||
def _declared_outside_the_pipeline(self) -> set:
|
||||
"""Fields declared by a `declare_resolution` caller outside
|
||||
`arg_groups/`, collected statically.
|
||||
|
||||
Resolution's launcher-stage writes live here -- the auto-detected
|
||||
parsers need a tokenizer or chat-template load, so they cannot run in
|
||||
`__post_init__` -- alongside the NPU default helper and the expert-pack
|
||||
loader, which the pipeline calls the same way. The construct-and-diff
|
||||
pass above never sees any of them.
|
||||
|
||||
`arg_groups/` is deliberately excluded: `_hook_assignment_targets`
|
||||
covers it exactly, and it resolves the pipeline's own computed
|
||||
expansions (`record_foreign_defaults` declares a `**` dict this
|
||||
collector's resolver cannot read). The keywords at the call sites are
|
||||
the written fields; an expansion this cannot resolve fails loudly like
|
||||
the override collector's, except the named dynamic sites below, whose
|
||||
field sets are spelled out and drift-guarded (each name must still
|
||||
appear as a constant in the file)."""
|
||||
written = set()
|
||||
root = _PACKAGE_ROOT
|
||||
for path in sorted(root.rglob("*.py")):
|
||||
rel = path.relative_to(root).as_posix()
|
||||
if rel.startswith("arg_groups/"):
|
||||
continue
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "declare_resolution" not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
self.fail(f"unparsable module in the census: {rel}")
|
||||
for node in ast.walk(tree):
|
||||
if not (
|
||||
isinstance(node, ast.Call)
|
||||
and (
|
||||
(
|
||||
isinstance(node.func, ast.Name)
|
||||
and node.func.id == "declare_resolution"
|
||||
)
|
||||
or (
|
||||
isinstance(node.func, ast.Attribute)
|
||||
and node.func.attr
|
||||
in ("declare_resolution", "_declare_resolution")
|
||||
)
|
||||
)
|
||||
):
|
||||
continue
|
||||
if all(kw.arg is None for kw in node.keywords) and any(
|
||||
isinstance(kw.value, ast.Name) and kw.value.id == "fields"
|
||||
for kw in node.keywords
|
||||
):
|
||||
# The forwarding shim (`ServerArgs._late_resolution` /
|
||||
# the helper's own body) re-expands its caller's kwargs;
|
||||
# the write sites are the callers.
|
||||
continue
|
||||
for kw in node.keywords:
|
||||
if kw.arg and kw.arg != "source":
|
||||
written.add(kw.arg)
|
||||
elif kw.arg is None:
|
||||
if rel in _CALLER_SUPPLIED_LATE_SITES:
|
||||
continue
|
||||
dynamic = _LATE_RESOLUTION_DYNAMIC_SITES.get(rel)
|
||||
if dynamic is not None:
|
||||
constants = {
|
||||
c.value
|
||||
for c in ast.walk(tree)
|
||||
if isinstance(c, ast.Constant)
|
||||
}
|
||||
missing = dynamic - constants
|
||||
self.assertFalse(
|
||||
missing,
|
||||
f"{rel}: the declared dynamic field set drifted "
|
||||
f"from the file ({sorted(missing)} not found)",
|
||||
)
|
||||
written |= dynamic
|
||||
else:
|
||||
written |= _expanded_override_keys(rel, tree, node, kw)
|
||||
return written
|
||||
|
||||
_READS_CACHE = None
|
||||
|
||||
def _supplied_instance_reads(self) -> set:
|
||||
"""Three spellings of the same read: ``server_args.field`` off the
|
||||
parameter, ``getattr(server_args, "field", default)`` with a literal
|
||||
name, and the *parked* form -- ``self.x = server_args`` in a method
|
||||
that takes the parameter, read as ``self.x.field`` anywhere in the
|
||||
class. Parking under a different object, a container, or a computed
|
||||
name stays invisible, like in every census of this family. The
|
||||
loudest boundary *was* the *chain* spelling,
|
||||
``model_runner.server_args.field`` off some other object, which this
|
||||
census still does not count -- but those reads are gone for every
|
||||
resolution-written field and ``test_chain_read_ratchet.py`` holds them
|
||||
at zero, so the gap is no longer where the risk is."""
|
||||
if TestSuppliedInstanceExposure._READS_CACHE is not None:
|
||||
return TestSuppliedInstanceExposure._READS_CACHE
|
||||
pairs = set()
|
||||
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
|
||||
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
|
||||
if rel.startswith(_OWNERS):
|
||||
continue
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "server_args" not in source:
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
# A silently dropped module shrinks `found` and reads as
|
||||
# intentional surface shrinkage under the bidirectional pin.
|
||||
self.fail(f"unparsable module in the census: {rel}")
|
||||
for fn in ast.walk(tree):
|
||||
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
params = {a.arg for a in list(fn.args.args) + list(fn.args.kwonlyargs)}
|
||||
if "server_args" not in params:
|
||||
continue
|
||||
for node in ast.walk(fn):
|
||||
if (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id == "server_args"
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
):
|
||||
pairs.add((rel, node.attr))
|
||||
elif (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Name)
|
||||
and node.func.id == "getattr"
|
||||
and len(node.args) >= 2
|
||||
and isinstance(node.args[0], ast.Name)
|
||||
and node.args[0].id == "server_args"
|
||||
and isinstance(node.args[1], ast.Constant)
|
||||
and isinstance(node.args[1].value, str)
|
||||
):
|
||||
# The same read in optional clothing. Only a literal
|
||||
# name is censusable; a computed one is not.
|
||||
pairs.add((rel, node.args[1].value))
|
||||
for cls in ast.walk(tree):
|
||||
if not isinstance(cls, ast.ClassDef):
|
||||
continue
|
||||
# Parked: `self.x = server_args` in a method that takes the
|
||||
# parameter, read as `self.x.field` anywhere in the class.
|
||||
parked = set()
|
||||
for fn in cls.body:
|
||||
if not isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
continue
|
||||
if "server_args" not in {
|
||||
a.arg for a in list(fn.args.args) + list(fn.args.kwonlyargs)
|
||||
}:
|
||||
continue
|
||||
for node in ast.walk(fn):
|
||||
if (
|
||||
isinstance(node, ast.Assign)
|
||||
and len(node.targets) == 1
|
||||
and isinstance(node.targets[0], ast.Attribute)
|
||||
and isinstance(node.targets[0].value, ast.Name)
|
||||
and node.targets[0].value.id == "self"
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id == "server_args"
|
||||
):
|
||||
parked.add(node.targets[0].attr)
|
||||
for node in ast.walk(cls):
|
||||
if (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.ctx, ast.Load)
|
||||
and isinstance(node.value, ast.Attribute)
|
||||
and node.value.attr in parked
|
||||
and isinstance(node.value.value, ast.Name)
|
||||
and node.value.value.id == "self"
|
||||
):
|
||||
pairs.add((rel, node.attr))
|
||||
TestSuppliedInstanceExposure._READS_CACHE = pairs
|
||||
return pairs
|
||||
|
||||
@staticmethod
|
||||
def _override_written_fields() -> set:
|
||||
"""Fields written post-publish through ``get_context().override(...)``."""
|
||||
written = set()
|
||||
for path in sorted(_PACKAGE_ROOT.rglob("*.py")):
|
||||
rel = path.relative_to(_PACKAGE_ROOT).as_posix()
|
||||
source = path.read_text(encoding="utf-8-sig")
|
||||
if "record_config_updates" not in source and not (
|
||||
"get_context" in source and "override" in source
|
||||
):
|
||||
continue
|
||||
try:
|
||||
tree = ast.parse(source)
|
||||
except SyntaxError:
|
||||
raise AssertionError(f"unparsable module in the census: {rel}")
|
||||
for node in ast.walk(tree):
|
||||
if not (
|
||||
isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)
|
||||
):
|
||||
continue
|
||||
base = node.func.value
|
||||
is_override = node.func.attr == "override" and (
|
||||
isinstance(base, ast.Call)
|
||||
and isinstance(base.func, ast.Name)
|
||||
and base.func.id == "get_context"
|
||||
)
|
||||
# `record_config_updates` is a named wrapper over override, so
|
||||
# its call sites are override sites. Its body forwards **kwargs
|
||||
# and names no field, so skip the forwarding call itself.
|
||||
is_wrapper = node.func.attr == "record_config_updates"
|
||||
if not (is_override or is_wrapper):
|
||||
continue
|
||||
inside_wrapper = any(
|
||||
isinstance(fn, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and fn.name == "record_config_updates"
|
||||
and fn.lineno <= node.lineno <= (fn.end_lineno or fn.lineno)
|
||||
for fn in ast.walk(tree)
|
||||
)
|
||||
if inside_wrapper:
|
||||
continue
|
||||
for kw in node.keywords:
|
||||
if kw.arg == "source":
|
||||
# Override metadata, not a config field.
|
||||
continue
|
||||
if kw.arg:
|
||||
written.add(kw.arg)
|
||||
else:
|
||||
written |= _expanded_override_keys(rel, tree, node, kw)
|
||||
return written
|
||||
|
||||
def test_the_post_publish_override_surface_matches_the_pinned_list(self):
|
||||
written = self._override_written_fields()
|
||||
self.assertGreater(
|
||||
len(written), 5, "found almost no override targets; the scan broke"
|
||||
)
|
||||
found = {pair for pair in self._supplied_instance_reads() if pair[1] in written}
|
||||
new = sorted(found - _OVERRIDDEN_AND_READ)
|
||||
gone = sorted(_OVERRIDDEN_AND_READ - found)
|
||||
self.assertEqual(
|
||||
([], []),
|
||||
(new, gone),
|
||||
"the post-publish override surface drifted. A read here answers with "
|
||||
"the startup value once the override lands, so a new pair needs an "
|
||||
"ordering judgment: copied before any override (fine), or read after "
|
||||
"one (then it must come from the bags).\n"
|
||||
f" new: {new}\n"
|
||||
f" gone (delete from _OVERRIDDEN_AND_READ): {gone}",
|
||||
)
|
||||
|
||||
def test_the_exposed_set_matches_the_pinned_list(self):
|
||||
import torch
|
||||
|
||||
written = self._resolution_written_fields()
|
||||
found = {pair for pair in self._supplied_instance_reads() if pair[1] in written}
|
||||
expected = set(_EXPOSED)
|
||||
if torch.cuda.is_available():
|
||||
expected |= _EXPOSED_CUDA_ONLY
|
||||
new = sorted(found - expected)
|
||||
gone = sorted(expected - found)
|
||||
self.assertEqual(
|
||||
([], []),
|
||||
(new, gone),
|
||||
"the supplied-instance step-12 surface drifted.\n"
|
||||
f" new (decide where the resolved value comes from): {new}\n"
|
||||
f" gone (delete from _EXPOSED / _EXPOSED_CUDA_ONLY): {gone}",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user