chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Support arbitrary filtering in dumper (#19593)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:31:21 +08:00
committed by GitHub
parent 8240a87306
commit f5a10e04cd
2 changed files with 88 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
python/sglang/srt/debug_utils/dumper.py
View File

@@ -424,9 +424,7 @@ class _Dumper:
**self._state.global_ctx,
)
if (f := self._config.filter) is not None and re.search(
f, _format_tags(tags)
) is None:
if (f := self._config.filter) is not None and not _evaluate_filter(f, tags):
return
if not (enable_value or enable_curr_grad or enable_future_grad):
@@ -897,6 +895,27 @@ def _format_tags(kwargs: dict) -> str:
return "___".join(f"{k}={v}" for k, v in kwargs.items())
class _DefaultNoneDict(dict):
"""dict subclass that returns None for missing keys, for filter expression eval."""
def __missing__(self, key: str):
return None
_FILTER_BUILTINS: dict[str, Any] = {"search": re.search, "match": re.match}
def _evaluate_filter(filter_expr: str, tags: dict[str, Any]) -> bool:
"""Evaluate a Python filter expression against the tags dict.
Unknown tag keys resolve to None, so `layer_id is None` works when layer_id is absent.
`re.search` and `re.match` are available as `search()` and `match()`.
"""
namespace = _DefaultNoneDict(tags)
namespace.update(_FILTER_BUILTINS)
return bool(eval(filter_expr, {"__builtins__": {}}, namespace))
def _deepcopy_or_clone(x):
if isinstance(x, torch.Tensor):
return x.clone()

89
test/registered/debug_utils/test_dumper.py
View File

@@ -193,8 +193,10 @@ class TestKvPairsParsing:
assert type(cfg.server_port) is str
def test_from_kv_pairs_optional_str_field(self):
cfg = DumperConfig.from_kv_pairs(["filter=layer_id=[0-3]"])
assert cfg.filter == "layer_id=[0-3]"
cfg = DumperConfig.from_kv_pairs(
["filter=layer_id is not None and layer_id < 3"]
)
assert cfg.filter == "layer_id is not None and layer_id < 3"
def test_from_kv_pairs_optional_str_exp_name(self):
cfg = DumperConfig.from_kv_pairs(["exp_name=my_experiment"])
@@ -205,14 +207,14 @@ class TestKvPairsParsing:
[
"enable=true",
"dir=/my/dir",
"filter=name=foo",
"filter=name == 'foo'",
"collective_timeout=30",
"enable_grad=1",
]
)
assert cfg.enable is True
assert cfg.dir == "/my/dir"
assert cfg.filter == "name=foo"
assert cfg.filter == "name == 'foo'"
assert cfg.collective_timeout == 30
assert cfg.enable_grad is True
@@ -237,8 +239,8 @@ class TestKvPairsParsing:
assert DumperConfig._kv_pairs_to_dict([]) == {}
def test_from_kv_pairs_value_with_equals_in_value(self):
cfg = DumperConfig.from_kv_pairs(["filter=name=foo"])
assert cfg.filter == "name=foo"
cfg = DumperConfig.from_kv_pairs(["filter=name == 'foo'"])
assert cfg.filter == "name == 'foo'"
def test_from_kv_pairs_type_validation_still_works(self):
with pytest.raises(TypeError, match="collective_timeout.*expected int"):
@@ -427,7 +429,7 @@ class TestDumperDistributed:
dumper.set_ctx(ctx_arg=None)
dumper.step()
dumper.configure(filter=r"^$")
dumper.configure(filter="False")
dumper.dump("tensor_skip", tensor)
dumper.configure(filter=None)
dumper.step()
@@ -497,7 +499,7 @@ class TestDumperFileWriteControl:
with temp_set_env(
DUMPER_ENABLE="1",
DUMPER_DIR=str(tmp_path),
DUMPER_FILTER="name=keep",
DUMPER_FILTER="name.startswith('keep')",
):
run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
@@ -618,7 +620,7 @@ class TestOutputControl:
pass
def test_capture_output_respects_filter(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="name=keep")
d = _make_test_dumper(tmp_path, filter="'keep' in name")
with d.capture_output() as captured:
d.dump("keep_this", torch.randn(3, 3))
@@ -866,7 +868,7 @@ class TestKvFilter:
assert _format_tags({}) == ""
def test_filter_matches_extra_kwargs(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="layer_id=0")
d = _make_test_dumper(tmp_path, filter="layer_id == 0")
d.dump("tensor_a", torch.randn(3), layer_id=0)
d.dump("tensor_b", torch.randn(3), layer_id=1)
@@ -874,7 +876,7 @@ class TestKvFilter:
_assert_files(filenames, exist=["tensor_a"], not_exist=["tensor_b"])
def test_filter_matches_global_ctx(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="ctx_arg=200")
d = _make_test_dumper(tmp_path, filter="ctx_arg == 200")
d.set_ctx(ctx_arg=200)
d.dump("tensor_a", torch.randn(3))
d.set_ctx(ctx_arg=None)
@@ -884,15 +886,15 @@ class TestKvFilter:
_assert_files(filenames, exist=["tensor_a"], not_exist=["tensor_b"])
def test_filter_matches_name(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="name=keep")
d = _make_test_dumper(tmp_path, filter="'keep' in name")
d.dump("keep_this", torch.randn(3))
d.dump("skip_this", torch.randn(3))
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["keep_this"], not_exist=["skip_this"])
def test_filter_regex(self, tmp_path):
d = _make_test_dumper(tmp_path, filter=r"layer_id=[0-2]")
def test_filter_expr_range(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="layer_id is not None and layer_id < 3")
d.dump("t0", torch.randn(3), layer_id=0)
d.dump("t1", torch.randn(3), layer_id=1)
d.dump("t5", torch.randn(3), layer_id=5)
@@ -900,6 +902,37 @@ class TestKvFilter:
filenames = _get_filenames(tmp_path)
_assert_files(filenames, exist=["name=t0", "name=t1"], not_exist=["name=t5"])
def test_filter_expr_with_none(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="layer_id is None or layer_id < 3")
d.dump("no_layer", torch.randn(3))
d.dump("layer0", torch.randn(3), layer_id=0)
d.dump("layer5", torch.randn(3), layer_id=5)
filenames = _get_filenames(tmp_path)
_assert_files(
filenames,
exist=["no_layer", "layer0"],
not_exist=["layer5"],
)
def test_filter_expr_with_re_search(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="search(r'attn|mlp', name)")
d.dump("self_attn", torch.randn(3))
d.dump("mlp_proj", torch.randn(3))
d.dump("layernorm", torch.randn(3))
filenames = _get_filenames(tmp_path)
_assert_files(
filenames,
exist=["self_attn", "mlp_proj"],
not_exist=["layernorm"],
)
def test_filter_expr_syntax_error(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="layer_id ===")
with pytest.raises(SyntaxError):
d.dump("tensor", torch.randn(3))
def test_no_filter_dumps_all(self, tmp_path):
d = _make_test_dumper(tmp_path)
d.dump("a", torch.randn(3))
@@ -952,7 +985,10 @@ class TestDumpModel:
def test_filter(self, tmp_path):
d = _make_test_dumper(
tmp_path, enable_model_value=True, enable_model_grad=True, filter="weight"
tmp_path,
enable_model_value=True,
enable_model_grad=True,
filter="'weight' in name",
)
model = torch.nn.Linear(4, 2)
x = torch.randn(3, 4)
@@ -1355,16 +1391,16 @@ class TestDumperHttp:
dumper_http_url,
"configure",
enable=True,
filter="layer_id=0",
filter="layer_id == 0",
dir="/tmp/test_http",
)
states = self._post(dumper_http_url, "get_state")
self._assert_all_ranks(states, "config.enable", True)
self._assert_all_ranks(states, "config.filter", "layer_id=0")
self._assert_all_ranks(states, "config.filter", "layer_id == 0")
self._assert_all_ranks(states, "config.dir", "/tmp/test_http")
def test_configure_clear_optional(self, dumper_http_url: str):
self._post(dumper_http_url, "configure", filter="layer_id=0")
self._post(dumper_http_url, "configure", filter="layer_id == 0")
self._post(dumper_http_url, "configure", filter=None)
states = self._post(dumper_http_url, "get_state")
self._assert_all_ranks(states, "config.filter", None)
@@ -1377,10 +1413,17 @@ class TestDumperHttp:
self._assert_all_ranks(states, "step", 0)
def test_get_state(self, dumper_http_url: str):
self._post(dumper_http_url, "configure", enable=True, filter="layer_id=[0-3]")
self._post(
dumper_http_url,
"configure",
enable=True,
filter="layer_id is not None and layer_id < 3",
)
states = self._post(dumper_http_url, "get_state")
self._assert_all_ranks(states, "config.enable", True)
self._assert_all_ranks(states, "config.filter", "layer_id=[0-3]")
self._assert_all_ranks(
states, "config.filter", "layer_id is not None and layer_id < 3"
)
for state in states:
assert "dump_index" in state
assert "step" in state
@@ -1713,7 +1756,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase):
return self.relu(self.linear(x))
captured, x, output = self._run(
tmp_path, Inner, filter="name=non_intrusive__model.linear.output"
tmp_path, Inner, filter="name == 'non_intrusive__model.linear.output'"
)
assert "non_intrusive__model.linear.output" in captured
@@ -1936,7 +1979,7 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase):
assert meta["layer_id"] == 1
def test_filter_by_layer_id(self, tmp_path):
"""filter='layer_id=0' keeps only layer 0 dumps."""
"""filter='layer_id == 0' keeps only layer 0 dumps."""
class Inner(torch.nn.Module):
def __init__(self):
@@ -1950,7 +1993,7 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase):
x = layer(x)
return x
captured, x, output = self._run(tmp_path, Inner, filter="layer_id=0")
captured, x, output = self._run(tmp_path, Inner, filter="layer_id == 0")
layer0_keys = [k for k in captured if "layers.0" in k]
layer1_keys = [k for k in captured if "layers.1" in k]