diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index f291343fd..6c7d3f6e5 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -152,10 +152,8 @@ 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`. + # Fully-qualified Python path "pkg.subpkg.module.fn_name" + # None -> use the default identity-by-rank fallback in _Grafter._default_transform. grafter_transform_path: Optional[str] = None @classmethod @@ -831,27 +829,34 @@ class GraftTransformInput: class _Grafter: - """1+1 cross-system tensor grafter. + """Cross-system tensor transplant. Triggered silently from dumper.dump. - 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_()`. + 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 ("baseline" | "target"), which + determines whether a name match means send or recv on this side. + + Graft global rank layout: baseline occupies ranks 0..baseline_world-1; + target occupies ranks baseline_world..baseline_world+target_world-1. Each + side derives its own rank from its local default PG via dist.get_rank(). + + Please refer to TestGrafterE2eExample in tests for an example. """ - def __init__(self, *, config: DumperConfig) -> None: + def __init__(self, *, config: DumperConfig): self._config = config - self._pg: Optional[dist.ProcessGroup] = None + self._pg = None + + @property + def enabled(self) -> bool: + return self._config.grafter_enable def maybe_intercept( - self, - *, - value, - tags: dict, - extras: Optional[dict] = None, + self, *, value: Any, tags: dict, extras: Optional[dict] = None ) -> None: + """Intercept a dumper.dump call. `extras` is per-call auxiliary data + (e.g., shard layout, dtype hint) that the sender attaches and the + recv side's transform receives as `received_extras_list`.""" cfg = self._config if not cfg.grafter_enable: return @@ -864,9 +869,9 @@ class _Grafter: _log( 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 " - f"narrow the filter or wrap the value in a tensor." + f"skipping graft. Common cause: dumper.dump called with a non-tensor " + f"value (dict, list, ...) on this name. Either narrow the filter or " + f"wrap the value in a tensor." ) return @@ -876,7 +881,8 @@ class _Grafter: # all-gather over the graft world; sender ranks contribute (value, # extras) tuples, recv ranks contribute None (their local target is - # private and shouldn't leak). + # private and shouldn't leak). all_gather_object is pickle-routed, + # so tensor shapes may differ across sender ranks. total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size my_contribution = (value, extras) if is_send else None gathered: list = [None] * total_world @@ -890,8 +896,8 @@ class _Grafter: return sender_contribs = self._sender_slice(direction=direction, gathered=gathered) - # Pickled CUDA tensors restore to their original-device name; that - # may not match this process's local device, so normalize. + # Pickled CUDA tensors are restored on their original-device name; + # that may not match this process's local device, so normalize. sender_tensors = [ (c[0].to(value.device) if isinstance(c[0], torch.Tensor) else c[0]) for c in sender_contribs @@ -928,66 +934,13 @@ class _Grafter: f"{traceback.format_exc()}" ) - def _sender_slice(self, *, direction: "_GraftDirection", gathered: list) -> list: - cfg = self._config - if direction == _GraftDirection.B2T: - return gathered[: cfg.grafter_baseline_world_size] - return gathered[cfg.grafter_baseline_world_size :] - - def _apply_transform( - self, - *, - tags: dict, - received_list: list, - received_extras_list: list, - target: torch.Tensor, - ) -> torch.Tensor: - graft_input = GraftTransformInput( - tags=tags, - received_list=received_list, - received_extras_list=received_extras_list, - target=target, - ) - path = self._config.grafter_transform_path - fn = self._default_transform if path is None else _load_function(path) - return fn(graft_input) - - @staticmethod - def _default_transform(graft_input: GraftTransformInput) -> torch.Tensor: - """Identity-by-rank fallback. Requires #senders == #recvs and - shape(received_list[my_recv_rank]) == shape(target). Otherwise raises - and asks the user for a transform.""" - received_list = graft_input.received_list - target = graft_input.target - my_recv_rank = dist.get_rank() - recv_world_size = dist.get_world_size() - if len(received_list) != recv_world_size: - raise RuntimeError( - f"[Grafter] no grafter_transform_path set; default " - f"identity-by-rank requires #senders == #recvs but got " - f"#senders={len(received_list)} vs #recvs={recv_world_size}. " - f"Provide a transform via " - f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." - ) - candidate = received_list[my_recv_rank] - 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[{my_recv_rank}].shape={tuple(candidate.shape)} " - f"!= target.shape={tuple(target.shape)}. Provide a transform " - f"via DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." - ) - return candidate - 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" + f"[Grafter] tags={tags} matched BOTH grafter_b2t_filter and grafter_t2b_filter" ) if match_b2t: return _GraftDirection.B2T @@ -1000,6 +953,12 @@ class _Grafter: # baseline is the sender for B2T names; target is the sender for T2B. return (role == _GraftRole.BASELINE) == (direction == _GraftDirection.B2T) + def _sender_slice(self, *, direction: "_GraftDirection", gathered: list) -> list: + cfg = self._config + if direction == _GraftDirection.B2T: + return gathered[: cfg.grafter_baseline_world_size] + return gathered[cfg.grafter_baseline_world_size :] + @staticmethod def _match(expr: Optional[str], tags: dict) -> bool: if expr is None: @@ -1050,6 +1009,62 @@ class _Grafter: timeout_seconds=cfg.grafter_timeout, ) + def _apply_transform( + self, + *, + tags: dict, + received_list: list, + received_extras_list: list, + target: torch.Tensor, + ) -> torch.Tensor: + # TODO: integrate with dump_comparator unsharder annotations once + # full inverse (sharded -> global -> sharded) transforms exist. + graft_input = GraftTransformInput( + tags=tags, + received_list=received_list, + received_extras_list=received_extras_list, + target=target, + ) + path = self._config.grafter_transform_path + fn = self._default_transform if path is None else _load_function(path) + return fn(graft_input) + + @staticmethod + def _default_transform(graft_input: GraftTransformInput) -> torch.Tensor: + """Identity-by-rank fallback. Requires #senders == #recvs and + shape(received_list[my_recv_rank]) == shape(target). Otherwise raises + and asks the user for a transform.""" + received_list = graft_input.received_list + target = graft_input.target + my_recv_rank = dist.get_rank() + recv_world_size = dist.get_world_size() + if len(received_list) != recv_world_size: + raise RuntimeError( + _Grafter._default_transform_error( + f"requires #senders == #recvs but got " + f"#senders={len(received_list)} vs #recvs={recv_world_size}" + ) + ) + candidate = received_list[my_recv_rank] + if candidate.shape != target.shape: + raise RuntimeError( + _Grafter._default_transform_error( + f"requires matching shapes but " + f"received_list[{my_recv_rank}].shape={tuple(candidate.shape)} " + f"!= target.shape={tuple(target.shape)}" + ) + ) + return candidate + + @staticmethod + def _default_transform_error(detail: str) -> str: + return ( + f"[Grafter] no grafter_transform_path set; default identity-by-rank " + f"{detail}. Provide a transform via " + f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol defining " + f"`transform(graft_input: GraftTransformInput) -> Tensor`." + ) + # -------------------------------------- util fn ------------------------------------------ @@ -1186,7 +1201,7 @@ def _compare_tensors_quick(a: "torch.Tensor", b: "torch.Tensor") -> str: sglang.srt.debug_utils.dump_comparator._compute_and_print_diff; intentionally inlined here to keep dumper.py free of cross-file imports. - Different dtypes are fine -- we unify by casting both to fp32, which is + Different dtypes are fine — we unify by casting both to fp32, which is enough for the order-of-magnitude diff summary we log.""" if a.shape != b.shape: return f"shape mismatch (a={tuple(a.shape)} vs b={tuple(b.shape)})" @@ -1510,12 +1525,11 @@ 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 + miles.utils.misc.load_function — kept inline so dumper.py has no cross-package dependency. """ import importlib diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 208e5aa1f..1cf9224c4 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -1,6 +1,7 @@ import io import multiprocessing import os +import re import sys import threading import time @@ -2629,76 +2630,90 @@ class TestRecomputeStatus: class TestGrafterConfig: + def test_from_env_parses_filters(self): + with temp_set_env( + DUMPER_GRAFTER_B2T_FILTER="name == 'x'", + DUMPER_GRAFTER_T2B_FILTER="name == 'y'", + ): + cfg = DumperConfig.from_env() + assert cfg.grafter_b2t_filter == "name == 'x'" + assert cfg.grafter_t2b_filter == "name == 'y'" + + def test_from_env_parses_int_fields(self): + with temp_set_env( + DUMPER_GRAFTER_BASELINE_WORLD_SIZE="8", + DUMPER_GRAFTER_TARGET_WORLD_SIZE="8", + DUMPER_GRAFTER_MASTER_PORT="29999", + DUMPER_GRAFTER_TIMEOUT="120", + ): + cfg = DumperConfig.from_env() + assert cfg.grafter_baseline_world_size == 8 + assert type(cfg.grafter_baseline_world_size) is int + assert cfg.grafter_target_world_size == 8 + assert cfg.grafter_master_port == 29999 + assert cfg.grafter_timeout == 120 + def test_from_env_role(self): + with temp_set_env(DUMPER_GRAFTER_ROLE="baseline"): + assert DumperConfig.from_env().grafter_role == "baseline" + + def test_from_env_enable_flag(self): + # enable=True requires all of role, master_address/port, world sizes, + # and at least one filter per DumperConfig.__post_init__. with temp_set_env( DUMPER_GRAFTER_ENABLE="1", DUMPER_GRAFTER_ROLE="baseline", DUMPER_GRAFTER_MASTER_ADDRESS="127.0.0.1", - DUMPER_GRAFTER_MASTER_PORT="29500", + DUMPER_GRAFTER_MASTER_PORT="29999", DUMPER_GRAFTER_BASELINE_WORLD_SIZE="1", DUMPER_GRAFTER_TARGET_WORLD_SIZE="1", DUMPER_GRAFTER_B2T_FILTER="name == 'x'", ): - cfg = DumperConfig.from_env() - assert cfg.grafter_enable is True - assert cfg.grafter_role == "baseline" - assert cfg.grafter_b2t_filter == "name == 'x'" - assert cfg.grafter_master_port == 29500 + assert DumperConfig.from_env().grafter_enable is True + with temp_set_env(DUMPER_GRAFTER_ENABLE="false"): + assert DumperConfig.from_env().grafter_enable is False def test_enable_without_required_fields_raises(self): - # missing role - with pytest.raises(AssertionError, match="grafter_role"): + with pytest.raises(AssertionError, match=r"grafter_role"): DumperConfig(grafter_enable=True) - # missing master_address - with pytest.raises(AssertionError, match="grafter_master_address"): + with pytest.raises(AssertionError, match=r"grafter_master_address"): DumperConfig(grafter_enable=True, grafter_role="baseline") - # non-positive port - with pytest.raises(AssertionError, match="grafter_master_port"): + with pytest.raises(AssertionError, match=r"grafter_master_port"): DumperConfig( grafter_enable=True, grafter_role="baseline", grafter_master_address="127.0.0.1", ) - # missing baseline_world_size - with pytest.raises(AssertionError, match="grafter_baseline_world_size"): + with pytest.raises(AssertionError, match=r"grafter_baseline_world_size"): DumperConfig( grafter_enable=True, grafter_role="baseline", grafter_master_address="127.0.0.1", - grafter_master_port=29500, + grafter_master_port=12345, ) - # missing target_world_size - with pytest.raises(AssertionError, match="grafter_target_world_size"): + with pytest.raises(AssertionError, match=r"neither grafter_b2t_filter nor"): DumperConfig( grafter_enable=True, grafter_role="baseline", grafter_master_address="127.0.0.1", - grafter_master_port=29500, - grafter_baseline_world_size=1, - ) - # no filter set - with pytest.raises(AssertionError, match="grafter_b2t_filter"): - DumperConfig( - grafter_enable=True, - grafter_role="baseline", - grafter_master_address="127.0.0.1", - grafter_master_port=29500, + grafter_master_port=12345, grafter_baseline_world_size=1, grafter_target_world_size=1, ) - def test_disabled_does_not_validate_other_fields(self): - # All grafter_* fields can be left at their absurd defaults when - # grafter_enable is False. - cfg = DumperConfig(grafter_enable=False) - assert cfg.grafter_enable is False + def test_env_name_for_grafter_field(self): + assert ( + DumperConfig._env_name("grafter_b2t_filter") == "DUMPER_GRAFTER_B2T_FILTER" + ) def _unit_grafter_config(**overrides) -> DumperConfig: - """Build a fully-valid DumperConfig for unit tests of `_Grafter` filter - matching, without spinning up a process group. Dummy values for required - fields are never reached because these tests short-circuit before - `_ensure_group` runs. + """Build a fully-valid DumperConfig for unit-test use. + + All grafter_* required fields default to dummy values; overrides patch + individual fields (e.g., grafter_enable=False or filter strings). + Dummy values are never reached because these unit tests short-circuit + before _ensure_group runs. """ base = dict( grafter_enable=True, @@ -2713,27 +2728,88 @@ def _unit_grafter_config(**overrides) -> DumperConfig: return DumperConfig(**base) +class TestLog: + def test_log_format(self): + with _capture_stdout() as captured: + _log("hello") + out = captured.getvalue() + assert "hello" in out, out + assert "[Dumper, rank=" in out, out + assert ", t=" in out, out + + +class TestCompareTensorsQuick: + def test_identical(self): + a = torch.tensor([1.0, 2.0, 3.0]) + s = _compare_tensors_quick(a, a.clone()) + assert "rel_diff=0" in s, s + assert "max_abs=0" in s, s + + def test_diverged(self): + a = torch.tensor([1.0, 2.0, 3.0]) + b = torch.tensor([1.0, 2.0, 4.0]) # last element differs by 1 + s = _compare_tensors_quick(a, b) + # rel_diff > 0 implies divergence; max_abs should equal 1.0 + assert "max_abs=1" in s, s + assert "rel_diff=" in s, s + + def test_shape_mismatch(self): + s = _compare_tensors_quick(torch.zeros(3), torch.zeros(4)) + assert "shape mismatch" in s, s + + def test_dtype_unified(self): + # Different dtypes should NOT error — both are cast to fp32 internally. + s = _compare_tensors_quick( + torch.zeros(3, dtype=torch.float32), + torch.zeros(3, dtype=torch.float64), + ) + assert "rel_diff=" in s, s + assert "max_abs=" in s, s + + def test_empty(self): + s = _compare_tensors_quick(torch.zeros(0), torch.zeros(0)) + assert s == "empty" + + class TestGrafterFilterMatching: - """Unit tests for the filter-matching short-circuit logic.""" + """Unit tests for the filter-matching short-circuit logic. + + These don't initialize a process group, so the network-related fields + are dummy values via _unit_grafter_config. + """ def test_disabled_returns_silently(self): grafter = _Grafter(config=_unit_grafter_config(grafter_enable=False)) grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"}) assert grafter._pg is None # never initialized - def test_unmatched_name_returns_silently(self): + def test_unmatched_non_tensor_silent(self): + """Non-tensor + unmatched name → silent skip, no print.""" grafter = _Grafter(config=_unit_grafter_config()) - grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "z"}) + 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_matched_non_tensor_prints_and_skips(self): - """Non-tensor that matches a filter -> log explanation, then skip.""" + """Non-tensor that matches a filter → print explanation, then skip. + + This catches misconfigured filters (e.g. matching a name that maps to a + dict/list at some call sites) without silently masking the issue.""" grafter = _Grafter(config=_unit_grafter_config()) with _capture_stdout() as captured: grafter.maybe_intercept(value={"not": "a tensor"}, tags={"name": "x"}) - out = captured.getvalue() + output = captured.getvalue() + assert grafter._pg is None # still no PG init + assert "value is not a torch.Tensor" in output, output + assert "type=dict" in output, output + + def test_unmatched_name_returns_silently(self): + grafter = _Grafter( + config=_unit_grafter_config(grafter_t2b_filter="name == 'y'") + ) + grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "z"}) assert grafter._pg is None - assert "value is not a torch.Tensor" in out, out def test_overlap_filters_raise(self): grafter = _Grafter( @@ -2748,14 +2824,6 @@ class TestGrafterFilterMatching: ): 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( @@ -2764,16 +2832,38 @@ class TestGrafterFilterMatching: grafter_t2b_filter="name == 'x' and layer_id < 3", ) ) - # layer_id=1 -> both filters match -> overlap raise (proves filter saw layer_id). + # 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. + # 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_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 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).""" @@ -2783,7 +2873,7 @@ class TestGrafterFilterMatching: grafter_t2b_filter="layer_id < 3", ) ) - # layer_id absent -> resolves to None; `None < 3` raises TypeError. + # layer_id absent → resolves to None; `None < 3` raises TypeError in py3. with pytest.raises(TypeError): grafter.maybe_intercept(value=torch.zeros(2), tags={"name": "x"}) @@ -2796,7 +2886,7 @@ class TestGrafterFilterMatching: grafter_t2b_filter="layer_id is None and name == 'x'", ) ) - # No `layer_id` in tags -> resolves to None -> filter matches -> tries + # 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"): @@ -2816,7 +2906,7 @@ class TestGrafterFilterMatching: 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 + (`'NoneType' object is not callable`) — loud enough to surface the misconfiguration.""" grafter = _Grafter( config=_unit_grafter_config( @@ -2834,43 +2924,25 @@ class TestGrafterFilterMatching: grafter_t2b_filter=None, ) ) - # name='attn_input' matches /attn.*/ -> tries to init group (hits + # 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. + # name='other' does not match → silent skip. 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). - Limited to 1+1 because the CI fleet has only 2 GPUs. Each process - initializes its OWN default PG (nccl, world_size=1) from the start, - mirroring production where baseline and target are independently launched. + Limited to 1+1 because CI machines we can rely on have only 2 GPUs. + Each process initializes its OWN default PG (nccl, world_size=1) from + the start, mirroring production where baseline and target are + independently launched. + + For asymmetric / multi-rank coverage that doesn't need GPU, see + `_run_graft_test_cpu_multi` below. """ import torch.multiprocessing as mp @@ -2897,6 +2969,8 @@ def _run_graft_test(worker_func, **kwargs): def _graft_worker_entry(rank, role_port, worker_func, result_queue, kwargs): + import traceback + torch.cuda.set_device(rank) dist.init_process_group( backend="nccl", @@ -2913,19 +2987,196 @@ def _graft_worker_entry(rank, role_port, worker_func, result_queue, kwargs): dist.destroy_process_group() +def _run_graft_test_split(worker_baseline, worker_target, **kwargs) -> dict: + """Like `_run_graft_test`, but each role runs its OWN dedicated worker + function (no `if rank == 0:` branching) and stdout is captured per role. + + Returns ``{"baseline": stdout_str, "target": stdout_str}`` so tests can + snapshot/assert on the per-role logs. Used by the E2E example for + educational clarity (each role's logic reads top-to-bottom) and to assert + the user-visible log output matches expectations. + """ + import torch.multiprocessing as mp + + role_ports = { + "baseline": find_available_port(29700), + "target": find_available_port(29800), + } + + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + processes = [] + for global_rank, (role, worker) in enumerate( + [("baseline", worker_baseline), ("target", worker_target)] + ): + p = ctx.Process( + target=_graft_split_worker_entry, + args=(global_rank, role, role_ports[role], worker, result_queue, kwargs), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + outputs: dict = {} + errors: list = [] + for _ in range(2): + role, error, captured = result_queue.get() + outputs[role] = captured + if error: + errors.append(f"role={role}: {error}") + if errors: + raise AssertionError( + "\n".join(errors) + + "\nCaptured outputs:\n" + + f"--- baseline ---\n{outputs.get('baseline', '')}\n" + + f"--- target ---\n{outputs.get('target', '')}" + ) + return outputs + + +def _graft_split_worker_entry( + global_rank, role, role_port, worker_func, result_queue, kwargs +): + import io + import traceback + + captured = io.StringIO() + old_stdout = sys.stdout + sys.stdout = captured + error = None + try: + # Set per-role env BEFORE we (re)build the module-level `dumper`. The + # parent left DUMPER_GRAFTER_ENABLE/ROLE unset because they vary per + # child; we set them here, then rebuild the global so that worker + # code can simply call `from sglang.srt.debug_utils.dumper import dumper` + # and get a properly-configured Grafter — exactly mirroring how + # production code uses the global. + os.environ["DUMPER_GRAFTER_ENABLE"] = "1" + os.environ["DUMPER_GRAFTER_ROLE"] = role + import sglang.srt.debug_utils.dumper as _dumper_module + + _dumper_module.dumper = _dumper_module._Dumper( + config=_dumper_module.DumperConfig.from_env() + ) + + torch.cuda.set_device(global_rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{role_port}", + world_size=1, + rank=0, + ) + try: + worker_func(**kwargs) + except Exception as e: + error = f"{e}\n{traceback.format_exc()}" + finally: + try: + dist.destroy_process_group() + except Exception: + pass + finally: + sys.stdout = old_stdout + result_queue.put((role, error, captured.getvalue())) + + +def _run_graft_test_cpu_multi( + worker_func, *, baseline_world: int, target_world: int, **kwargs +): + """Spawn (baseline_world + target_world) CPU-only processes (gloo backend). + + Used to exercise asymmetric multi-rank cases (e.g. 4 baseline ranks and + 2 target ranks) that we can't run on the 2-GPU CI fleet. Each role gets + its OWN default PG (gloo, world=role_world); the graft cross-system PG + spans all ranks. + + The worker function receives (role, local_rank, **kwargs). + """ + import torch.multiprocessing as mp + + # One default-PG port per role (baseline-side ranks share one PG, target + # ranks share another). Allocated up-front to avoid child races. + role_ports = { + "baseline": find_available_port(29800), + "target": find_available_port(29900), + } + + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + processes = [] + total = baseline_world + target_world + for global_rank in range(total): + if global_rank < baseline_world: + role = "baseline" + local_rank = global_rank + local_world = baseline_world + else: + role = "target" + local_rank = global_rank - baseline_world + local_world = target_world + p = ctx.Process( + target=_graft_cpu_worker_entry, + args=( + role, + local_rank, + local_world, + role_ports[role], + worker_func, + result_queue, + kwargs, + ), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + errors = [result_queue.get() for _ in range(total)] + errors = [e for e in errors if e] + if errors: + raise AssertionError("\n".join(errors)) + + +def _graft_cpu_worker_entry( + role, local_rank, local_world, port, worker_func, result_queue, kwargs +): + import traceback + + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + world_size=local_world, + rank=local_rank, + ) + try: + worker_func(role=role, local_rank=local_rank, **kwargs) + result_queue.put(None) + except Exception as e: + result_queue.put( + f"role={role} local_rank={local_rank}: {e}\n{traceback.format_exc()}" + ) + finally: + dist.destroy_process_group() + + def _make_grafter_test_config( *, rank: int, graft_port: int, group_name: str, timeout: int = 30, + transform_path: Optional[str] = None, 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 - default PG. + """Helper for distributed grafter tests. + + Same b2t/t2b filters on both sides; only `grafter_role` differs (rank 0 = + baseline, rank 1 = target). Both sides are world_size=1 within their own + role's default PG. """ role = "baseline" if rank == 0 else "target" return DumperConfig( @@ -2939,7 +3190,9 @@ def _make_grafter_test_config( grafter_target_world_size=1, grafter_group_name=group_name, grafter_timeout=timeout, - grafter_transform_path=transform_path, + # Loading the user transform on the recv side; for b2t the recv is + # the target side (rank 1). + grafter_transform_path=transform_path if rank == 1 else None, ) @@ -2964,8 +3217,11 @@ class TestGrafterDistributed: grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: target = torch.zeros(3, device="cuda:1") - grafter.maybe_intercept(value=target, tags={"name": "x"}) + with _capture_stdout() as captured: + grafter.maybe_intercept(value=target, tags={"name": "x"}) assert target.tolist() == [1.0, 2.0, 3.0], f"got {target.tolist()}" + # Success log must include the pre/new diff summary. + assert "diff_pre_vs_new=" in captured.getvalue(), captured.getvalue() finally: if grafter._pg is not None: dist.destroy_process_group(grafter._pg) @@ -3000,105 +3256,9 @@ class TestGrafterDistributed: 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( - self._test_unmatched_func, - graft_port=graft_port, - group_name="grafter_unmatched", - ) - - @staticmethod - def _test_unmatched_func(rank, graft_port, group_name): - grafter = _Grafter( - config=_make_grafter_test_config( - rank=rank, graft_port=graft_port, group_name=group_name - ) - ) - try: - target = torch.tensor([7.0, 7.0, 7.0], device=f"cuda:{rank}") - grafter.maybe_intercept(value=target, tags={"name": "other"}) - assert target.tolist() == [7.0, 7.0, 7.0], "tensor must not be modified" - assert grafter._pg is None, "group must not init for unmatched name" - finally: - if grafter._pg is not None: - dist.destroy_process_group(grafter._pg) - - def test_init_timeout_warns(self): - graft_port = find_available_port(29630) - _run_graft_test( - self._test_init_timeout_func, - graft_port=graft_port, - group_name="grafter_timeout", - ) - - @staticmethod - def _test_init_timeout_func(rank, graft_port, group_name): - grafter = _Grafter( - config=_make_grafter_test_config( - rank=rank, graft_port=graft_port, group_name=group_name, timeout=2 - ) - ) - try: - with _capture_stdout() as captured: - if rank == 1: - time.sleep(4) - tensor = torch.tensor([1.0, 2.0, 3.0], device=f"cuda:{rank}") - if rank == 0: - grafter.maybe_intercept(value=tensor, tags={"name": "x"}) - else: - target = torch.zeros(3, device=f"cuda:{rank}") - grafter.maybe_intercept(value=target, tags={"name": "x"}) - output = captured.getvalue() - if rank == 0: - assert "WARNING" in output, output - assert "has not completed after 2s" in output, output - finally: - if grafter._pg is not None: - dist.destroy_process_group(grafter._pg) - - def test_group_init_is_cached_across_calls(self): - """`_ensure_group` runs lazily on first matched dump and caches `_pg`; - subsequent matched dumps must reuse the same object.""" - graft_port = find_available_port(29660) - _run_graft_test( - self._test_group_cache_func, - graft_port=graft_port, - group_name="grafter_cache", - ) - - @staticmethod - def _test_group_cache_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: - t1 = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") - t2 = torch.tensor([4.0, 5.0, 6.0], device="cuda:0") - grafter.maybe_intercept(value=t1, tags={"name": "x"}) - pg_after_first = grafter._pg - assert pg_after_first is not None - grafter.maybe_intercept(value=t2, tags={"name": "x"}) - assert grafter._pg is pg_after_first, "_pg must be cached" - else: - t1 = torch.zeros(3, device="cuda:1") - t2 = torch.zeros(3, device="cuda:1") - grafter.maybe_intercept(value=t1, tags={"name": "x"}) - pg_after_first = grafter._pg - assert pg_after_first is not None - grafter.maybe_intercept(value=t2, tags={"name": "x"}) - assert grafter._pg is pg_after_first - assert t1.tolist() == [1.0, 2.0, 3.0] - assert t2.tolist() == [4.0, 5.0, 6.0] - finally: - 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_.""" + # Write a tiny module that defines `transform(graft_input)`. The + # worker prepends tmp_path to sys.path so import_module sees it. module_name = "_xform_user_basic" (tmp_path / f"{module_name}.py").write_text( "def transform(graft_input):\n" @@ -3138,6 +3298,30 @@ class TestGrafterDistributed: 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( + self._test_unmatched_func, + graft_port=graft_port, + group_name="grafter_unmatched", + ) + + @staticmethod + def _test_unmatched_func(rank, graft_port, group_name): + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, graft_port=graft_port, group_name=group_name + ) + ) + try: + target = torch.tensor([7.0, 7.0, 7.0], device=f"cuda:{rank}") + grafter.maybe_intercept(value=target, tags={"name": "other"}) + assert target.tolist() == [7.0, 7.0, 7.0], "tensor must not be modified" + assert grafter._pg is None, "group must not init for unmatched name" + 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.""" @@ -3157,10 +3341,13 @@ class TestGrafterDistributed: ) try: if rank == 0: + # Baseline sends shape=(3,) tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: + # Target's local target has shape=(4,) — mismatch with sender. target = torch.tensor([7.0, 7.0, 7.0, 7.0], device="cuda:1") + # No exception should propagate; tensor must stay unchanged. grafter.maybe_intercept(value=target, tags={"name": "x"}) assert target.tolist() == [ 7.0, @@ -3219,64 +3406,13 @@ class TestGrafterDistributed: output = captured.getvalue() assert "transform/copy_ raised RuntimeError" in output, output assert "intentional test error" in output, output + # Full traceback must be included so the bug is debuggable. 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(graft_input):\n" - " return torch.zeros(99, device=graft_input.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) - def test_extras_flow_to_recv_transform(self, tmp_path: Path): """Sender attaches per-call grafter_extras; recv transform reads them and uses them to compute the override value.""" @@ -3309,6 +3445,7 @@ class TestGrafterDistributed: ) try: if rank == 0: + # Baseline (sender) attaches an extras dict. tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") grafter.maybe_intercept( value=tensor, @@ -3318,14 +3455,53 @@ class TestGrafterDistributed: else: target = torch.zeros(3, device="cuda:1") grafter.maybe_intercept(value=target, tags={"name": "x"}) - assert target.tolist() == [42.0, 42.0, 42.0], target.tolist() + assert target.tolist() == [ + 42.0, + 42.0, + 42.0, + ], f"target should be filled from sender extras, got {target.tolist()}" + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_init_timeout_warns(self): + graft_port = find_available_port(29630) + _run_graft_test( + self._test_init_timeout_func, + graft_port=graft_port, + group_name="grafter_timeout", + ) + + @staticmethod + def _test_init_timeout_func(rank, graft_port, group_name): + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, graft_port=graft_port, group_name=group_name, timeout=2 + ) + ) + try: + with _capture_stdout() as captured: + if rank == 1: + time.sleep(4) + tensor = torch.tensor([1.0, 2.0, 3.0], device=f"cuda:{rank}") + if rank == 0: + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.zeros(3, device=f"cuda:{rank}") + grafter.maybe_intercept(value=target, tags={"name": "x"}) + output = captured.getvalue() + if rank == 0: + assert ( + "WARNING" in output + ), f"expected WARNING in rank 0 output: {output}" + assert "has not completed after 2s" in output, output finally: if grafter._pg is not None: dist.destroy_process_group(grafter._pg) def test_extras_default_none_flow(self): """When the sender omits `grafter_extras`, the recv transform sees a - list of Nones.""" + list of Nones — but len(received_extras_list) still matches n_senders.""" graft_port = find_available_port(29650) _run_graft_test( self._test_extras_none_func, @@ -3343,11 +3519,14 @@ class TestGrafterDistributed: try: if rank == 0: tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + # Note: extras kwarg omitted entirely → None on the wire. grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: target = torch.zeros(3, device="cuda:1") with _capture_stdout() as captured: grafter.maybe_intercept(value=target, tags={"name": "x"}) + # Default identity transform copies tensor through; recv log + # must reflect that received_extras_list == [None]. output = captured.getvalue() assert "sender_extras=[None]" in output, output assert target.tolist() == [1.0, 2.0, 3.0], target.tolist() @@ -3355,107 +3534,103 @@ class TestGrafterDistributed: if grafter._pg is not None: dist.destroy_process_group(grafter._pg) - -def _run_graft_test_cpu_multi( - worker_func, *, baseline_world: int, target_world: int, **kwargs -): - """Spawn (baseline_world + target_world) CPU-only processes (gloo backend). - - Used to exercise asymmetric multi-rank cases (e.g. 4 baseline ranks and - 2 target ranks) that we can't run on the 2-GPU CI fleet. Each role gets - its OWN default PG (gloo, world=role_world); the graft cross-system PG - spans all ranks. - - The worker function receives (role, local_rank, **kwargs). - """ - import torch.multiprocessing as mp - - role_ports = { - "baseline": find_available_port(29800), - "target": find_available_port(29900), - } - - ctx = mp.get_context("spawn") - result_queue = ctx.Queue() - processes = [] - total = baseline_world + target_world - for global_rank in range(total): - if global_rank < baseline_world: - role = "baseline" - local_rank = global_rank - local_world = baseline_world - else: - role = "target" - local_rank = global_rank - baseline_world - local_world = target_world - p = ctx.Process( - target=_graft_cpu_worker_entry, - args=( - role, - local_rank, - local_world, - role_ports[role], - worker_func, - result_queue, - kwargs, - ), + def test_group_init_is_cached_across_calls(self): + """The graft process group is initialized lazily on the first + matched dump() and cached afterwards — subsequent dumps must reuse + the same `_pg` object, not re-init.""" + graft_port = find_available_port(29660) + _run_graft_test( + self._test_group_cache_func, + graft_port=graft_port, + group_name="grafter_cache", ) - p.start() - processes.append(p) - for p in processes: - p.join() - - errors = [result_queue.get() for _ in range(total)] - errors = [e for e in errors if e] - if errors: - raise AssertionError("\n".join(errors)) - - -def _graft_cpu_worker_entry( - role, local_rank, local_world, port, worker_func, result_queue, kwargs -): - dist.init_process_group( - backend="gloo", - init_method=f"tcp://127.0.0.1:{port}", - world_size=local_world, - rank=local_rank, - ) - try: - worker_func(role=role, local_rank=local_rank, **kwargs) - result_queue.put(None) - except Exception as e: - result_queue.put( - f"role={role} local_rank={local_rank}: {e}\n{traceback.format_exc()}" + @staticmethod + def _test_group_cache_func(rank, graft_port, group_name): + grafter = _Grafter( + config=_make_grafter_test_config( + rank=rank, graft_port=graft_port, group_name=group_name + ) ) - finally: - dist.destroy_process_group() + try: + if rank == 0: + t1 = torch.tensor([1.0, 2.0, 3.0], device="cuda:0") + t2 = torch.tensor([4.0, 5.0, 6.0], device="cuda:0") + grafter.maybe_intercept(value=t1, tags={"name": "x"}) + pg_after_first = grafter._pg + assert pg_after_first is not None + grafter.maybe_intercept(value=t2, tags={"name": "x"}) + assert ( + grafter._pg is pg_after_first + ), "_pg must be cached across calls, not re-initialized" + else: + target1 = torch.zeros(3, device="cuda:1") + target2 = torch.zeros(3, device="cuda:1") + grafter.maybe_intercept(value=target1, tags={"name": "x"}) + pg_after_first = grafter._pg + assert pg_after_first is not None + grafter.maybe_intercept(value=target2, tags={"name": "x"}) + assert grafter._pg is pg_after_first + assert target1.tolist() == [1.0, 2.0, 3.0] + assert target2.tolist() == [4.0, 5.0, 6.0] + 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 + (same robustness contract as transform-throws).""" + module_name = "_xform_returns_wrong_shape" + (tmp_path / f"{module_name}.py").write_text( + "import torch\n" + "def transform(graft_input):\n" + " # Deliberately return a shape that copy_ will reject.\n" + " return torch.zeros(99, device=graft_input.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", + ) -def _make_multi_rank_config( - *, - role: str, - graft_port: int, - group_name: str, - baseline_world: int, - target_world: int, - transform_path: Optional[str], - direction: str, -) -> DumperConfig: - return DumperConfig( - grafter_enable=True, - grafter_role=role, - grafter_b2t_filter="name == 'x'" if direction == "b2t" else None, - grafter_t2b_filter="name == 'x'" if direction == "t2b" else None, - grafter_master_address="127.0.0.1", - grafter_master_port=graft_port, - grafter_baseline_world_size=baseline_world, - grafter_target_world_size=target_world, - grafter_backend="gloo", - grafter_group_name=group_name, - grafter_timeout=30, - grafter_transform_path=transform_path, - ) + @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"}) + # target must be unchanged; error must be logged with traceback. + 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) class TestGrafterMultiRankCpu: @@ -3463,7 +3638,9 @@ class TestGrafterMultiRankCpu: only 2 GPUs, which is too few for these cases).""" def test_4_baseline_2_target_b2t_with_user_transform(self, tmp_path: Path): - """4 baseline senders -> 2 target receivers via b2t graft.""" + """4 baseline senders -> 2 target receivers via b2t graft. + The user transform asserts received_list has length 4 with each + sender's tensor matching its rank, then returns a marker tensor.""" module_name = "_xform_assert_4_senders" (tmp_path / f"{module_name}.py").write_text( "import torch\n" @@ -3503,9 +3680,11 @@ class TestGrafterMultiRankCpu: grafter = _Grafter(config=cfg) try: if role == "baseline": + # rank-i baseline contributes [i, i, i]. tensor = torch.full((3,), float(local_rank)) grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: + # Target's local tensor (will be overwritten with 999s by transform). target = torch.full((3,), 99.0) grafter.maybe_intercept(value=target, tags={"name": "x"}) assert target.tolist() == [999.0, 999.0, 999.0], target.tolist() @@ -3514,7 +3693,9 @@ class TestGrafterMultiRankCpu: dist.destroy_process_group(grafter._pg) def test_2_target_4_baseline_t2b_with_user_transform(self, tmp_path: Path): - """Mirror image: 2 target senders -> 4 baseline receivers.""" + """Mirror image of the b2t case: 2 target senders -> 4 baseline + receivers via t2b graft. Confirms the (role, direction) algebra and + sender_slice work correctly when target is the SENDER side.""" module_name = "_xform_assert_2_senders_t2b" (tmp_path / f"{module_name}.py").write_text( "import torch\n" @@ -3523,7 +3704,9 @@ class TestGrafterMultiRankCpu: " assert len(rl) == 2, f'expected 2 senders, got {len(rl)}'\n" " for i, t in enumerate(rl):\n" " v = float(t.flatten()[0].item())\n" - " assert v == float(i + 100), f'rl[{i}][0]={v}'\n" + " assert v == float(i + 100), (\n" + " f'rl[{i}][0]={v}, want {float(i + 100)}'\n" + " )\n" " return torch.full_like(graft_input.target, 7.0)\n" ) graft_port = find_available_port(29670) @@ -3554,6 +3737,7 @@ class TestGrafterMultiRankCpu: grafter = _Grafter(config=cfg) try: if role == "target": + # rank-i target contributes [i+100, i+100, i+100]. tensor = torch.full((3,), float(local_rank + 100)) grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: @@ -3565,9 +3749,10 @@ class TestGrafterMultiRankCpu: dist.destroy_process_group(grafter._pg) def test_default_transform_with_asymmetric_world_logs_and_skips(self): - """Default identity-by-rank requires #senders == #recvs; with 4 vs 2 - and no user transform, recv catches RuntimeError + leaves target - unchanged.""" + """The default identity-by-rank fallback requires #senders == #recvs. + With baseline=4 and target=2 and no user transform, the recv side + must catch the RuntimeError, log it with traceback, and leave the + target unchanged.""" graft_port = find_available_port(29675) _run_graft_test_cpu_multi( self._test_default_asym_func, @@ -3585,7 +3770,7 @@ class TestGrafterMultiRankCpu: group_name=group_name, baseline_world=4, target_world=2, - transform_path=None, + transform_path=None, # default identity-by-rank fallback direction="b2t", ) grafter = _Grafter(config=cfg) @@ -3597,9 +3782,13 @@ class TestGrafterMultiRankCpu: target = torch.full((3,), 42.0) with _capture_stdout() as captured: grafter.maybe_intercept(value=target, tags={"name": "x"}) - assert target.tolist() == [42.0, 42.0, 42.0], target.tolist() + assert target.tolist() == [42.0, 42.0, 42.0], ( + f"target must be unchanged when default transform raises, " + f"got {target.tolist()}" + ) output = captured.getvalue() assert "transform/copy_ raised RuntimeError" in output, output + # The error message must explain WHY the default fell through. assert "#senders=4" in output and "#recvs=2" in output, output assert "Traceback (most recent call last)" in output, output finally: @@ -3608,16 +3797,23 @@ class TestGrafterMultiRankCpu: def test_mixed_shape_senders_via_user_transform(self, tmp_path: Path): """`all_gather_object` is pickle-routed, so sender ranks may - contribute tensors with DIFFERENT shapes. The user transform handles - the dispatch.""" + contribute tensors with DIFFERENT shapes. The user transform sees + the full list and is responsible for picking/reducing. Asserts that + rank-i baseline's tensor has shape (i+1,) and the transform + concatenates them on the recv side.""" module_name = "_xform_concat_mixed_shape" (tmp_path / f"{module_name}.py").write_text( "import torch\n" "def transform(graft_input):\n" - " expected_shapes = [(i + 1,) for i in range(len(graft_input.received_list))]\n" - " actual_shapes = [tuple(t.shape) for t in graft_input.received_list]\n" - " assert actual_shapes == expected_shapes, actual_shapes\n" - " return torch.cat(graft_input.received_list)\n" + " rl = graft_input.received_list\n" + " # Each baseline sent shape=(rank+1,) tensors filled with rank.\n" + " expected_shapes = [(i + 1,) for i in range(len(rl))]\n" + " actual_shapes = [tuple(t.shape) for t in rl]\n" + " assert actual_shapes == expected_shapes, (\n" + " f'shape mismatch: expected {expected_shapes}, got {actual_shapes}'\n" + " )\n" + " # Concat to length 1+2+3+4 = 10 == target's length.\n" + " return torch.cat(rl)\n" ) graft_port = find_available_port(29680) _run_graft_test_cpu_multi( @@ -3647,9 +3843,11 @@ class TestGrafterMultiRankCpu: grafter = _Grafter(config=cfg) try: if role == "baseline": + # rank-i baseline contributes shape=(i+1,) filled with i. tensor = torch.full((local_rank + 1,), float(local_rank)) grafter.maybe_intercept(value=tensor, tags={"name": "x"}) else: + # 1 + 2 + 3 + 4 = 10 elements after concat. target = torch.zeros(10) grafter.maybe_intercept(value=target, tags={"name": "x"}) expected = [0.0] + [1.0] * 2 + [2.0] * 3 + [3.0] * 4 @@ -3659,5 +3857,259 @@ class TestGrafterMultiRankCpu: dist.destroy_process_group(grafter._pg) +def _make_multi_rank_config( + *, + role: str, + graft_port: int, + group_name: str, + baseline_world: int, + target_world: int, + transform_path: Optional[str], + direction: str, +) -> DumperConfig: + return DumperConfig( + grafter_enable=True, + grafter_role=role, + grafter_b2t_filter="name == 'x'" if direction == "b2t" else None, + grafter_t2b_filter="name == 'x'" if direction == "t2b" else None, + grafter_master_address="127.0.0.1", + grafter_master_port=graft_port, + grafter_baseline_world_size=baseline_world, + grafter_target_world_size=target_world, + grafter_backend="gloo", + grafter_group_name=group_name, + grafter_timeout=30, + grafter_transform_path=transform_path, + ) + + +def _e2e_transform(graft_input): + """User transform used by the E2E example test. Demonstrates the two + customization hooks reviewers should learn from: + + 1. The transform receives a `GraftTransformInput` and returns the + tensor that the recv side will `.copy_()` into its local target. + 2. `graft_input.received_extras_list` carries whatever the sender + passed via `grafter_extras={...}` — useful for any per-call + metadata the recv side needs (layer ids, calibration knobs, ...). + + Here we keep the example minimal: the sender attaches a single dummy + key/value so the recv side has something concrete to assert on, then + the transform is just identity. Real workflows would compute a + non-trivial override (scale, reshape, decode, ...) using the extras. + """ + assert ( + graft_input.received_extras_list[0]["my_extra_key"] == "my_extra_value" + ), graft_input.received_extras_list + return graft_input.received_list[0] + + +class TestGrafterE2eExample: + """End-to-end example: target has a (suspected) buggy attention kernel. + + Story: target's attention kernel produces wrong outputs and we want to + test "if we replace target's attention with baseline's, does the rest of + the model converge?". The full graft wiring is: + + - At the attention call site, target sends its inputs (q/k/v) to + baseline → baseline's local inputs are overwritten by target's, so + baseline runs its (known-good) attention against the same inputs. + This is a t->b graft on `attn_input`. + - Both sides run the kernel. + - Baseline sends its outputs back to target → target's outputs are + overwritten by baseline's, so target's downstream sees baseline's + attention result. This is a b->t graft on `attn_output`. + + Net effect: target's attention is semantically replaced by baseline's, + without modifying target's source beyond inserting `dumper.dump` at the + input/output sites. This test additionally demonstrates two recv-side + customization hooks via `_e2e_transform`: + + * `grafter_extras={...}` per dump call — arbitrary per-call metadata + the recv side can consume. + * `DUMPER_GRAFTER_TRANSFORM_PATH` — a user-supplied function that + decides what value the recv side actually copy_'s in (defaults to + identity-by-rank when unset). + + The remaining call-site code is exactly: + + dumper.dump("attn_input", q, grafter_extras={"layer_id": 7}) # t -> b + out = target_attention_kernel(q, ...) + dumper.dump("attn_output", out, grafter_extras={"scale": 0.5}) # b -> t + """ + + def test_e2e_buggy_attn_replaced_by_baseline(self): + graft_port = find_available_port(29640) + # All non-role env is shared by both sides; we set it in the parent + # so the spawned subprocesses inherit it. DUMPER_GRAFTER_ENABLE and + # DUMPER_GRAFTER_ROLE are deliberately *not* set here — they are set + # by `_run_graft_test_split` per-rank, after which the global + # `dumper` is rebuilt (so workers can use the global directly). + with temp_set_env( + DUMPER_ENABLE="1", + DUMPER_ENABLE_OUTPUT_FILE="false", # skip disk I/O for the test + DUMPER_ENABLE_OUTPUT_CONSOLE="false", + # Pin exp_name so the dumper doesn't auto-pick + log "Choose + # exp_name=..." into the captured snapshot. + DUMPER_EXP_NAME="grafter_e2e_test", + DUMPER_GRAFTER_MASTER_ADDRESS="127.0.0.1", + DUMPER_GRAFTER_MASTER_PORT=str(graft_port), + DUMPER_GRAFTER_BASELINE_WORLD_SIZE="1", + DUMPER_GRAFTER_TARGET_WORLD_SIZE="1", + DUMPER_GRAFTER_B2T_FILTER="name == 'attn_output'", + DUMPER_GRAFTER_T2B_FILTER="name == 'attn_input'", + DUMPER_GRAFTER_GROUP_NAME="grafter_e2e", + DUMPER_GRAFTER_TIMEOUT="30", + DUMPER_GRAFTER_TRANSFORM_PATH=f"{__name__}._e2e_transform", + ): + outputs = _run_graft_test_split(self._worker_baseline, self._worker_target) + + self._assert_e2e_snapshot(outputs) + + @staticmethod + def _assert_e2e_snapshot(outputs: dict) -> None: + """Snapshot of the FULL per-role log timeline. + + Volatile fields (timestamps, ports, float diff values, tensor + min/max/mean/samples, struct addresses) are masked with ad-hoc regex + placeholders so the snapshot stays stable while still pinning + everything else. The snapshot doubles as documentation of the logs a + reader will see when running this E2E setup. + + Captured logs are unconditionally printed before asserting so a + snapshot failure doesn't require a re-run. + """ + baseline_log = outputs["baseline"] + target_log = outputs["target"] + + print("\n=========== captured baseline log ===========") + print(baseline_log) + print("=========== captured target log ===========") + print(target_log) + print("===========================================") + + # Convenience tokens for verbose volatile substrings. + prefix = r"\[Dumper, rank=\d+, t=\d+\.\d+\] " + # `get_tensor_info(t)` for our tensors expands to a long line; we + # match the leading struct fields verbatim and let the trailing + # min/max/mean/sample fields wildcard out. + tinfo_f32_4 = ( + r"type= shape=torch\.Size\(\[4\]\) " + r"dtype=torch\.float32 device=cuda:\d stride=\(1,\) " + r"req_grad=False .*" + ) + diff = r"rel_diff=[-\d.eE+]+ max_abs=[-\d.eE+]+ mean_abs=[-\d.eE+]+" + + # `_dump_inner` automatically annotates tags with `recompute_status` + # (always present, value depends on whether autograd recompute is + # active — "disabled" in this test env). + attn_input_tags = r"\{'name': 'attn_input', 'recompute_status': 'disabled'\}" + attn_output_tags = r"\{'name': 'attn_output', 'recompute_status': 'disabled'\}" + + # Same dummy extras dict travels in both directions. + extras_lit = r"\{'my_extra_key': 'my_extra_value'\}" + + baseline_pattern = ( + r"\A" + f"{prefix}\\[Grafter\\] init group: role=baseline " + r"baseline_world=1 target_world=1 rank=0 " + r"init_method=tcp://127\.0\.0\.1:\d+ backend=nccl " + r"name=grafter_e2e\n" + f"{prefix}\\[Grafter\\] recv role=baseline dir=t2b " + f"tags={attn_input_tags} n_senders=1 " + f"sender_extras=\\[{extras_lit}\\] " + f"before_overridden={tinfo_f32_4} " + f"to_override={tinfo_f32_4} " + f"diff_pre_vs_new={diff}\n" + f"{prefix}\\[Grafter\\] send role=baseline dir=b2t " + f"tags={attn_output_tags} extras={extras_lit} " + f"local={tinfo_f32_4}\n" + r"\Z" + ) + target_pattern = ( + r"\A" + f"{prefix}\\[Grafter\\] init group: role=target " + r"baseline_world=1 target_world=1 rank=1 " + r"init_method=tcp://127\.0\.0\.1:\d+ backend=nccl " + r"name=grafter_e2e\n" + f"{prefix}\\[Grafter\\] send role=target dir=t2b " + f"tags={attn_input_tags} extras={extras_lit} " + f"local={tinfo_f32_4}\n" + f"{prefix}\\[Grafter\\] recv role=target dir=b2t " + f"tags={attn_output_tags} n_senders=1 " + f"sender_extras=\\[{extras_lit}\\] " + f"before_overridden={tinfo_f32_4} " + f"to_override={tinfo_f32_4} " + f"diff_pre_vs_new={diff}\n" + r"\Z" + ) + + assert re.fullmatch(baseline_pattern, baseline_log, flags=re.DOTALL), ( + f"baseline log did not match snapshot.\n" + f"--- pattern ---\n{baseline_pattern}\n" + f"--- actual ---\n{baseline_log}" + ) + assert re.fullmatch(target_pattern, target_log, flags=re.DOTALL), ( + f"target log did not match snapshot.\n" + f"--- pattern ---\n{target_pattern}\n" + f"--- actual ---\n{target_log}" + ) + + @staticmethod + def _worker_baseline(): + # In production code, callers just `from sglang.srt.debug_utils.dumper + # import dumper` and call `dumper.dump(name, value)` — the env + # configures the global Grafter for them. We do the same here. + from sglang.srt.debug_utils.dumper import dumper + + # Step 1: graft input. target sends its q to baseline; baseline's + # `_e2e_transform` runs on the recv side, asserts the dummy extras + # made it across, then returns target's q so baseline's local + # placeholder is overwritten via .copy_(). + q = torch.tensor([99.0, 99.0, 99.0, 99.0], device="cuda:0") + dumper.dump("attn_input", q) + assert q.tolist() == [1.0, 2.0, 3.0, 4.0], ( + f"baseline's q should be overwritten by target's via the t->b graft, " + f"got {q.tolist()}" + ) + + # Step 2: baseline runs the known-good attention kernel. + attn_out = q * 10.0 # → [10, 20, 30, 40] + + # Step 3: graft output. baseline sends attn_out to target with a + # dummy extras key the recv-side transform will assert on. + dumper.dump( + "attn_output", + attn_out, + grafter_extras={"my_extra_key": "my_extra_value"}, + ) + + @staticmethod + def _worker_target(): + from sglang.srt.debug_utils.dumper import dumper + + # Step 1: graft input. target sends its real q to baseline along + # with a dummy extras key the recv-side transform will assert on. + q = torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda:1") + dumper.dump( + "attn_input", + q, + grafter_extras={"my_extra_key": "my_extra_value"}, + ) + + # Step 2: target runs the (suspected buggy) attention kernel — + # here it returns all zeros to mimic a broken implementation. + attn_out = torch.zeros_like(q) + + # Step 3: graft output. baseline sends its (good) attn_out to + # target; target's recv-side transform identity-passes it, so + # target's local attn_out ends up = baseline's [10, 20, 30, 40]. + dumper.dump("attn_output", attn_out) + assert attn_out.tolist() == [10.0, 20.0, 30.0, 40.0], ( + f"target's attn_out should be overwritten by baseline's via " + f"the b->t graft, got {attn_out.tolist()}" + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))