Support arbitrary filtering in dumper (#19593)

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(
+        if (f := self._config.filter) is not None and not _evaluate_filter(f, tags):
-            f, _format_tags(tags)
-        ) is None:
             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()

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]"])
+        cfg = DumperConfig.from_kv_pairs(
-        assert cfg.filter == "layer_id=[0-3]"
+            ["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"])
+        cfg = DumperConfig.from_kv_pairs(["filter=name == 'foo'"])
-        assert cfg.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):
+    def test_filter_expr_range(self, tmp_path):
-        d = _make_test_dumper(tmp_path, filter=r"layer_id=[0-2]")
+        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]