diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 8ec1011b9..0249b6854 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -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() diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 9937370f0..643250012 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -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]