diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index a3fb617b6..f94de7c01 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/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, ( - "grafter_enable=True but grafter_b2t_filter is not set; " - "nothing would ever be grafted" + assert ( + self.grafter_b2t_filter is not None + 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 - baseline -> target). The only per-side difference is `grafter_role`, - which tells the side whether it's the sender (baseline) or the - receiver (target). Receiver overwrites its local target tensor with - the sender's via `value.copy_()`. + Both sides set the SAME `grafter_b2t_filter` (names that flow baseline -> + target) and `grafter_t2b_filter` (names that flow target -> baseline). + The only per-side difference is `grafter_role`, which tells the side + whether it's the sender or the receiver for the matched direction. + 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 - # (graft rank 1). Use broadcast_object_list so receiver can have an - # arbitrarily shaped placeholder; sender ships a pickled tensor. + # 1+1 broadcast: sender side ships the tensor as a pickled object; + # recv side calls `value.copy_()` with the received 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}") - dist.broadcast_object_list(obj_list, src=0, group=self._pg) - if role == _GraftRole.TARGET: + _log(f"[Grafter] send role={role.value} dir={direction.value} tags={tags}") + dist.broadcast_object_list(obj_list, src=sender_rank, group=self._pg) + 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: diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index bb293f3db..8a26ff625 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/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(