Support user-supplied recv-side transform in dumper grafter (#24509)

python/sglang/srt/debug_utils/dumper.py

@@ -7,6 +7,7 @@ import re
 import socket
 import threading
 import time
+import traceback
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -151,6 +152,11 @@ class DumperConfig(_BaseConfig):
     grafter_backend: str = "nccl"
     grafter_group_name: str = "graft"
     grafter_timeout: int = 300
+    # Fully-qualified Python path "pkg.subpkg.module.fn_name". When set, the
+    # recv side calls this function with (received_list, target) and copies
+    # the result into target. None -> use the default identity-by-rank
+    # fallback in `_Grafter._default_transform`.
+    grafter_transform_path: Optional[str] = None
 
     @classmethod
     def _env_prefix(cls) -> str:
@@ -837,21 +843,65 @@ class _Grafter:
         is_send = self._is_sender(role=role, direction=direction)
 
         # 1+1 broadcast: sender side ships the tensor as a pickled object;
-        # recv side calls `value.copy_()` with the received tensor.
+        # recv side feeds it through the user transform (default: identity)
+        # and `value.copy_()` the result.
         sender_rank = 0 if direction == _GraftDirection.B2T else 1
         obj_list: list = [None]
         if is_send:
             obj_list = [value]
             _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:
+        if is_send:
-            received = obj_list[0]
+            return
-            if isinstance(received, torch.Tensor):
+
-                # Pickled CUDA tensors restore to their original-device name;
+        received = obj_list[0]
-                # that may not match this process's local device, so normalize.
+        if isinstance(received, torch.Tensor):
-                received = received.to(value.device)
+            # 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)
+        # Transform + copy_ are wrapped: a buggy user transform must NOT
+        # crash the whole training/inference run. On error we log the full
+        # traceback and skip this graft point; downstream sees the recv
+        # side's original tensor unchanged.
+        try:
+            value_to_override = self._apply_transform([received], target=value)
             _log(f"[Grafter] recv role={role.value} dir={direction.value} tags={tags}")
-            value.copy_(received)
+            value.copy_(value_to_override)
+        except Exception as e:
+            _log(
+                f"[Grafter] recv role={role.value} dir={direction.value} "
+                f"tags={tags} transform/copy_ raised {type(e).__name__}: {e}; "
+                f"skipping graft for this call (target tensor unchanged)\n"
+                f"{traceback.format_exc()}"
+            )
+
+    def _apply_transform(
+        self,
+        received_list: list,
+        *,
+        target: torch.Tensor,
+    ) -> torch.Tensor:
+        path = self._config.grafter_transform_path
+        if path is None:
+            return self._default_transform(received_list, target=target)
+        return _load_function(path)(received_list, target)
+
+    @staticmethod
+    def _default_transform(
+        received_list: list, *, target: torch.Tensor
+    ) -> torch.Tensor:
+        """Identity-by-rank fallback. For the 1+1 setup currently supported,
+        just returns the single received tensor; requires shape match."""
+        candidate = received_list[0]
+        if candidate.shape != target.shape:
+            raise RuntimeError(
+                f"[Grafter] no grafter_transform_path set; default "
+                f"identity-by-rank requires matching shapes but "
+                f"received_list[0].shape={tuple(candidate.shape)} != "
+                f"target.shape={tuple(target.shape)}. Provide a transform via "
+                f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol."
+            )
+        return candidate
 
     def _classify_direction(self, tags: dict) -> Optional["_GraftDirection"]:
         cfg = self._config
@@ -1339,6 +1389,26 @@ def _get_local_ip_by_remote() -> Optional[str]:
     return None
 
 
+@functools.lru_cache(maxsize=None)
+def _load_function(path: str) -> Callable:
+    """Resolve a fully-qualified Python path 'pkg.module.symbol' to its object.
+
+    Copied (verbatim, minus the function-registry branch) from
+    miles.utils.misc.load_function -- kept inline so dumper.py has no
+    cross-package dependency.
+    """
+    import importlib
+
+    module_path, _, attr = path.rpartition(".")
+    if not module_path:
+        raise ValueError(
+            f"_load_function expects 'pkg.module.symbol', got {path!r} "
+            f"(missing dotted prefix)"
+        )
+    module = importlib.import_module(module_path)
+    return getattr(module, attr)
+
+
 def _init_custom_process_group(
     *,
     backend: str,

test/registered/debug_utils/test_dumper.py

@@ -23,6 +23,7 @@ from sglang.srt.debug_utils.dumper import (
     _format_tags,
     _get_default_exp_name,
     _Grafter,
+    _load_function,
     _log,
     _map_tensor,
     _materialize_value,
@@ -2809,6 +2810,28 @@ class TestGrafterFilterMatching:
         grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "other"})
         assert grafter._pg is None
 
+    def test_load_function_bad_module(self):
+        with pytest.raises(ModuleNotFoundError):
+            _load_function("no_such_pkg.no_such_module.transform")
+
+    def test_load_function_missing_attr(self):
+        # `os.path` exists but has no `definitely_no_such_attr`.
+        with pytest.raises(AttributeError):
+            _load_function("os.path.definitely_no_such_attr")
+
+    def test_load_function_no_dotted_prefix(self):
+        with pytest.raises(ValueError, match=r"missing dotted prefix"):
+            _load_function("only_one_segment")
+
+    def test_load_function_non_callable_resolves_but_call_fails(self):
+        """`_load_function` itself only does attribute lookup -- it doesn't
+        verify the result is callable. A non-callable target manifests at
+        call time as TypeError; we still want the failure to be debuggable."""
+        sep = _load_function("os.path.sep")  # str, not a callable
+        assert isinstance(sep, str)
+        with pytest.raises(TypeError):
+            sep()
+
 
 def _run_graft_test(worker_func, **kwargs):
     """Spawn one GPU-using process per role (rank 0 = baseline, rank 1 = target).
@@ -2866,6 +2889,7 @@ def _make_grafter_test_config(
     timeout: int = 30,
     b2t_filter: Optional[str] = "name == 'x'",
     t2b_filter: Optional[str] = None,
+    transform_path: 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
@@ -2883,6 +2907,7 @@ def _make_grafter_test_config(
         grafter_target_world_size=1,
         grafter_group_name=group_name,
         grafter_timeout=timeout,
+        grafter_transform_path=transform_path,
     )
 
 
@@ -3040,6 +3065,186 @@ class TestGrafterDistributed:
             if grafter._pg is not None:
                 dist.destroy_process_group(grafter._pg)
 
+    def test_recv_with_user_transform(self, tmp_path: Path):
+        """User transform doubles the received tensor before copy_."""
+        module_name = "_xform_user_basic"
+        (tmp_path / f"{module_name}.py").write_text(
+            "def transform(received_list, target):\n"
+            "    return received_list[0] * 2\n"
+        )
+        graft_port = find_available_port(29610)
+        _run_graft_test(
+            self._test_user_transform_func,
+            graft_port=graft_port,
+            group_name="grafter_transform",
+            transform_dir=str(tmp_path),
+            transform_path=f"{module_name}.transform",
+        )
+
+    @staticmethod
+    def _test_user_transform_func(
+        rank, graft_port, group_name, transform_dir, transform_path
+    ):
+        sys.path.insert(0, transform_dir)
+        grafter = _Grafter(
+            config=_make_grafter_test_config(
+                rank=rank,
+                graft_port=graft_port,
+                group_name=group_name,
+                transform_path=transform_path,
+            )
+        )
+        try:
+            if rank == 0:
+                tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
+                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
+            else:
+                target = torch.zeros(3, device="cuda:1")
+                grafter.maybe_intercept(value=target, tags={"name": "x"})
+                assert target.tolist() == [2.0, 4.0, 6.0], f"got {target.tolist()}"
+        finally:
+            if grafter._pg is not None:
+                dist.destroy_process_group(grafter._pg)
+
+    def test_default_fallback_shape_mismatch_does_not_crash(self):
+        """When sender shape != target shape, default identity fallback raises;
+        the grafter must catch it, log, and leave target unchanged."""
+        graft_port = find_available_port(29615)
+        _run_graft_test(
+            self._test_shape_mismatch_func,
+            graft_port=graft_port,
+            group_name="grafter_shape_mismatch",
+        )
+
+    @staticmethod
+    def _test_shape_mismatch_func(rank, graft_port, group_name):
+        grafter = _Grafter(
+            config=_make_grafter_test_config(
+                rank=rank, graft_port=graft_port, group_name=group_name
+            )
+        )
+        try:
+            if rank == 0:
+                tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
+                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
+            else:
+                target = torch.tensor([7.0, 7.0, 7.0, 7.0], device="cuda:1")
+                grafter.maybe_intercept(value=target, tags={"name": "x"})
+                assert target.tolist() == [
+                    7.0,
+                    7.0,
+                    7.0,
+                    7.0,
+                ], f"target should be unchanged after shape-mismatch graft, got {target.tolist()}"
+        finally:
+            if grafter._pg is not None:
+                dist.destroy_process_group(grafter._pg)
+
+    def test_user_transform_exception_does_not_crash(self, tmp_path: Path):
+        """A user transform that raises must NOT bring down the system; the
+        grafter logs and skips the copy_, leaving target unchanged."""
+        module_name = "_xform_throws"
+        (tmp_path / f"{module_name}.py").write_text(
+            "def transform(received_list, target):\n"
+            "    raise RuntimeError('intentional test error from user transform')\n"
+        )
+        graft_port = find_available_port(29635)
+        _run_graft_test(
+            self._test_transform_throws_func,
+            graft_port=graft_port,
+            group_name="grafter_throws",
+            transform_dir=str(tmp_path),
+            transform_path=f"{module_name}.transform",
+            module_name=module_name,
+        )
+
+    @staticmethod
+    def _test_transform_throws_func(
+        rank, graft_port, group_name, transform_dir, transform_path, module_name
+    ):
+        sys.path.insert(0, transform_dir)
+        grafter = _Grafter(
+            config=_make_grafter_test_config(
+                rank=rank,
+                graft_port=graft_port,
+                group_name=group_name,
+                transform_path=transform_path,
+            )
+        )
+        try:
+            if rank == 0:
+                tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
+                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
+            else:
+                target = torch.tensor([9.0, 9.0, 9.0], device="cuda:1")
+                with _capture_stdout() as captured:
+                    grafter.maybe_intercept(value=target, tags={"name": "x"})
+                assert target.tolist() == [
+                    9.0,
+                    9.0,
+                    9.0,
+                ], f"target must be unchanged when transform throws, got {target.tolist()}"
+                output = captured.getvalue()
+                assert "transform/copy_ raised RuntimeError" in output, output
+                assert "intentional test error" in output, output
+                assert "Traceback (most recent call last)" in output, output
+                assert f"{module_name}.py" in output, output
+        finally:
+            if grafter._pg is not None:
+                dist.destroy_process_group(grafter._pg)
+
+    def test_copy_failure_does_not_crash(self, tmp_path: Path):
+        """If the user transform returns a tensor whose shape doesn't match
+        target, `value.copy_(value_to_override)` raises -- and that error
+        must be caught, logged with traceback, and target left unchanged."""
+        module_name = "_xform_returns_wrong_shape"
+        (tmp_path / f"{module_name}.py").write_text(
+            "import torch\n"
+            "def transform(received_list, target):\n"
+            "    return torch.zeros(99, device=target.device)\n"
+        )
+        graft_port = find_available_port(29665)
+        _run_graft_test(
+            self._test_copy_failure_func,
+            graft_port=graft_port,
+            group_name="grafter_copy_fail",
+            transform_dir=str(tmp_path),
+            transform_path=f"{module_name}.transform",
+        )
+
+    @staticmethod
+    def _test_copy_failure_func(
+        rank, graft_port, group_name, transform_dir, transform_path
+    ):
+        sys.path.insert(0, transform_dir)
+        grafter = _Grafter(
+            config=_make_grafter_test_config(
+                rank=rank,
+                graft_port=graft_port,
+                group_name=group_name,
+                transform_path=transform_path,
+            )
+        )
+        try:
+            if rank == 0:
+                tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
+                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
+            else:
+                target = torch.tensor([7.0, 7.0, 7.0], device="cuda:1")
+                with _capture_stdout() as captured:
+                    grafter.maybe_intercept(value=target, tags={"name": "x"})
+                assert target.tolist() == [
+                    7.0,
+                    7.0,
+                    7.0,
+                ], f"target must be unchanged on copy_ failure, got {target.tolist()}"
+                output = captured.getvalue()
+                assert "transform/copy_ raised" in output, output
+                assert "Traceback (most recent call last)" in output, output
+        finally:
+            if grafter._pg is not None:
+                dist.destroy_process_group(grafter._pg)
+
 
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))