Support t2b direction and overlap protection in dumper grafter (#24508)

python/sglang/srt/debug_utils/dumper.py

@@ -143,6 +143,7 @@ class DumperConfig(_BaseConfig):
     grafter_enable: bool = False
     grafter_role: str = ""  # required if enabled: "baseline" or "target"
     grafter_b2t_filter: Optional[str] = None  # names flowing baseline -> target
+    grafter_t2b_filter: Optional[str] = None  # names flowing target -> baseline
     grafter_master_address: str = ""  # required if enabled
     grafter_master_port: int = -1  # required if enabled (positive port)
     grafter_baseline_world_size: int = -1  # required if enabled
@@ -178,9 +179,12 @@ class DumperConfig(_BaseConfig):
                 f"grafter_target_world_size must be > 0 when grafter_enable=True, "
                 f"got {self.grafter_target_world_size}"
             )
-            assert self.grafter_b2t_filter is not None, (
+            assert (
-                "grafter_enable=True but grafter_b2t_filter is not set; "
+                self.grafter_b2t_filter is not None
-                "nothing would ever be grafted"
+                or self.grafter_t2b_filter is not None
+            ), (
+                "grafter_enable=True but neither grafter_b2t_filter nor "
+                "grafter_t2b_filter is set; nothing would ever be grafted"
             )
 
     @property
@@ -789,14 +793,20 @@ class _GraftRole(enum.Enum):
     TARGET = "target"
 
 
+class _GraftDirection(enum.Enum):
+    B2T = "b2t"  # name flows baseline -> target
+    T2B = "t2b"  # name flows target -> baseline
+
+
 class _Grafter:
     """1+1 cross-system tensor grafter.
 
-    Both sides set the SAME `grafter_b2t_filter` (names that flow
+    Both sides set the SAME `grafter_b2t_filter` (names that flow baseline ->
-    baseline -> target). The only per-side difference is `grafter_role`,
+    target) and `grafter_t2b_filter` (names that flow target -> baseline).
-    which tells the side whether it's the sender (baseline) or the
+    The only per-side difference is `grafter_role`, which tells the side
-    receiver (target). Receiver overwrites its local target tensor with
+    whether it's the sender or the receiver for the matched direction.
-    the sender's via `value.copy_()`.
+    Receiver overwrites its local target tensor with the sender's via
+    `value.copy_()`.
     """
 
     def __init__(self, *, config: DumperConfig) -> None:
@@ -808,12 +818,13 @@ class _Grafter:
         if not cfg.grafter_enable:
             return
 
-        if not self._match(cfg.grafter_b2t_filter, tags):
+        direction = self._classify_direction(tags)
+        if direction is None:
             return
 
         if not isinstance(value, torch.Tensor):
             _log(
-                f"[Grafter] tags={tags} matched grafter_b2t_filter but "
+                f"[Grafter] tags={tags} matched grafter_{direction.value}_filter but "
                 f"value is not a torch.Tensor (got type={type(value).__name__}); "
                 f"skipping graft. Common cause: dumper.dump called with a "
                 f"non-tensor value (dict, list, ...) on this name. Either "
@@ -823,24 +834,45 @@ class _Grafter:
 
         self._ensure_group()
         role = _GraftRole(cfg.grafter_role)
+        is_send = self._is_sender(role=role, direction=direction)
 
-        # b2t with 1+1: baseline rank is sender (graft rank 0), target is recv
+        # 1+1 broadcast: sender side ships the tensor as a pickled object;
-        # (graft rank 1). Use broadcast_object_list so receiver can have an
+        # recv side calls `value.copy_()` with the received tensor.
-        # arbitrarily shaped placeholder; sender ships a pickled tensor.
+        sender_rank = 0 if direction == _GraftDirection.B2T else 1
         obj_list: list = [None]
-        if role == _GraftRole.BASELINE:
+        if is_send:
             obj_list = [value]
-            _log(f"[Grafter] send role=baseline tags={tags}")
+            _log(f"[Grafter] send role={role.value} dir={direction.value} tags={tags}")
-        dist.broadcast_object_list(obj_list, src=0, group=self._pg)
+        dist.broadcast_object_list(obj_list, src=sender_rank, group=self._pg)
-        if role == _GraftRole.TARGET:
+        if not is_send:
             received = obj_list[0]
             if isinstance(received, torch.Tensor):
                 # Pickled CUDA tensors restore to their original-device name;
                 # that may not match this process's local device, so normalize.
                 received = received.to(value.device)
-            _log(f"[Grafter] recv role=target tags={tags}")
+            _log(f"[Grafter] recv role={role.value} dir={direction.value} tags={tags}")
             value.copy_(received)
 
+    def _classify_direction(self, tags: dict) -> Optional["_GraftDirection"]:
+        cfg = self._config
+        match_b2t = self._match(cfg.grafter_b2t_filter, tags)
+        match_t2b = self._match(cfg.grafter_t2b_filter, tags)
+        if match_b2t and match_t2b:
+            raise RuntimeError(
+                f"[Grafter] tags={tags} matched BOTH grafter_b2t_filter "
+                f"and grafter_t2b_filter"
+            )
+        if match_b2t:
+            return _GraftDirection.B2T
+        if match_t2b:
+            return _GraftDirection.T2B
+        return None
+
+    @staticmethod
+    def _is_sender(*, role: "_GraftRole", direction: "_GraftDirection") -> bool:
+        # baseline is the sender for B2T names; target is the sender for T2B.
+        return (role == _GraftRole.BASELINE) == (direction == _GraftDirection.B2T)
+
     @staticmethod
     def _match(expr: Optional[str], tags: dict) -> bool:
         if expr is None:

test/registered/debug_utils/test_dumper.py

@@ -2702,6 +2702,113 @@ class TestGrafterFilterMatching:
         assert grafter._pg is None
         assert "value is not a torch.Tensor" in out, out
 
+    def test_overlap_filters_raise(self):
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter="name == 'x'",
+                grafter_t2b_filter="name == 'x'",
+            )
+        )
+        with pytest.raises(
+            RuntimeError,
+            match=r"matched BOTH grafter_b2t_filter and grafter_t2b_filter",
+        ):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"})
+
+    def test_unmatched_non_tensor_silent(self):
+        """Non-tensor + unmatched name -> silent skip, no print."""
+        grafter = _Grafter(config=_unit_grafter_config())
+        with _capture_stdout() as captured:
+            grafter.maybe_intercept(value=42, tags={"name": "other"})
+        assert grafter._pg is None
+        assert "[Grafter]" not in captured.getvalue(), captured.getvalue()
+
+    def test_filter_expression_uses_extra_tags(self):
+        """Filter expressions can reference any tag key, not just 'name'."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter="name == 'x' and layer_id < 3",
+                grafter_t2b_filter="name == 'x' and layer_id < 3",
+            )
+        )
+        # layer_id=1 -> both filters match -> overlap raise (proves filter saw layer_id).
+        with pytest.raises(RuntimeError, match=r"matched BOTH"):
+            grafter.maybe_intercept(
+                value=torch.zeros(2),
+                tags={"name": "x", "layer_id": 1},
+            )
+        # layer_id=5 -> neither filter matches -> silent skip.
+        grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x", "layer_id": 5})
+        assert grafter._pg is None
+
+    def test_filter_expression_only_uses_non_name_tag(self):
+        """A filter that doesn't reference `name` at all is still valid; it
+        should match purely on the other tag(s)."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter=None,
+                grafter_t2b_filter="layer_id < 3",
+            )
+        )
+        # layer_id absent -> resolves to None; `None < 3` raises TypeError.
+        with pytest.raises(TypeError):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"})
+
+    def test_filter_expression_unknown_tag_resolves_to_none(self):
+        """Unknown tag keys resolve to None inside filter expressions, so
+        `layer_id is None` works as an "absent" probe without raising."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter=None,
+                grafter_t2b_filter="layer_id is None and name == 'x'",
+            )
+        )
+        # No `layer_id` in tags -> resolves to None -> filter matches -> tries
+        # to init the recv group (which we can't actually do here without a
+        # real PG, so we expect the assertion failure from _ensure_group).
+        with pytest.raises(AssertionError, match="default torch.distributed"):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"})
+
+    def test_filter_expression_syntax_error_raises(self):
+        """A filter string that isn't valid Python should surface as a
+        SyntaxError so the misconfiguration is loud, not silent."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(grafter_b2t_filter="name == "),
+        )
+        with pytest.raises(SyntaxError):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"})
+
+    def test_filter_expression_undefined_helper_raises(self):
+        """Referencing an undefined helper inside a filter (e.g. a function
+        the user expected to be in scope) should NOT be silently treated as
+        False. The filter namespace is a `_DefaultNoneDict` (unknown keys
+        resolve to None), so calling an undefined helper raises TypeError
+        (`'NoneType' object is not callable`) -- loud enough to surface the
+        misconfiguration."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter="totally_undefined_helper(name)"
+            ),
+        )
+        with pytest.raises(TypeError, match=r"NoneType.* not callable"):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"})
+
+    def test_filter_can_use_re_search(self):
+        """`re.search` is exposed inside filter expressions as `search()`."""
+        grafter = _Grafter(
+            config=_unit_grafter_config(
+                grafter_b2t_filter="search(r'attn.*', name) is not None",
+                grafter_t2b_filter=None,
+            )
+        )
+        # name='attn_input' matches /attn.*/ -> tries to init group (hits
+        # the no-default-PG assertion, proving the regex matched).
+        with pytest.raises(AssertionError, match="default torch.distributed"):
+            grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "attn_input"})
+        # name='other' does not match -> silent skip.
+        grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "other"})
+        assert grafter._pg is None
+
 
 def _run_graft_test(worker_func, **kwargs):
     """Spawn one GPU-using process per role (rank 0 = baseline, rank 1 = target).
@@ -2758,6 +2865,7 @@ def _make_grafter_test_config(
     group_name: str,
     timeout: int = 30,
     b2t_filter: Optional[str] = "name == 'x'",
+    t2b_filter: Optional[str] = None,
 ) -> DumperConfig:
     """Build a DumperConfig for distributed grafter tests. rank 0 -> baseline,
     rank 1 -> target. Both sides are world_size=1 within their own role's
@@ -2768,6 +2876,7 @@ def _make_grafter_test_config(
         grafter_enable=True,
         grafter_role=role,
         grafter_b2t_filter=b2t_filter,
+        grafter_t2b_filter=t2b_filter,
         grafter_master_address="127.0.0.1",
         grafter_master_port=graft_port,
         grafter_baseline_world_size=1,
@@ -2804,6 +2913,36 @@ class TestGrafterDistributed:
             if grafter._pg is not None:
                 dist.destroy_process_group(grafter._pg)
 
+    def test_t2b_copy_roundtrip(self):
+        """Target (rank 1) sends 'x' to baseline (rank 0), baseline.copy_'s it."""
+        graft_port = find_available_port(29605)
+        _run_graft_test(
+            self._test_t2b_func, graft_port=graft_port, group_name="grafter_t2b"
+        )
+
+    @staticmethod
+    def _test_t2b_func(rank, graft_port, group_name):
+        grafter = _Grafter(
+            config=_make_grafter_test_config(
+                rank=rank,
+                graft_port=graft_port,
+                group_name=group_name,
+                b2t_filter=None,
+                t2b_filter="name == 'x'",
+            )
+        )
+        try:
+            if rank == 1:
+                tensor = torch.tensor([4.0, 5.0, 6.0], device="cuda:1")
+                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
+            else:
+                target = torch.zeros(3, device="cuda:0")
+                grafter.maybe_intercept(value=target, tags={"name": "x"})
+                assert target.tolist() == [4.0, 5.0, 6.0], f"got {target.tolist()}"
+        finally:
+            if grafter._pg is not None:
+                dist.destroy_process_group(grafter._pg)
+
     def test_unmatched_name_skipped(self):
         graft_port = find_available_port(29620)
         _run_graft_test(