sglang/test/registered/debug_utils/test_dumper.py

import io
import multiprocessing
import os
import re
import sys
import threading
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Optional

import pytest
import requests
import torch
import torch.distributed as dist

from sglang.srt.debug_utils.dumper import (
    DumperConfig,
    _collect_parallel_rank_tags,
    _collective_with_timeout,
    _compare_tensors_quick,
    _deepcopy_or_clone,
    _detect_recompute_status,
    _Dumper,
    _format_tags,
    _get_default_exp_name,
    _Grafter,
    _load_function,
    _log,
    _map_tensor,
    _materialize_value,
    _MegatronPlugin,
    _obj_to_dict,
    _RecomputeStatus,
    _register_forward_hook_or_replace_fn,
    _SGLangPlugin,
    _torch_save,
    dumper,
    get_tensor_info,
    get_truncated_value,
)
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.common import temp_set_env
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    find_available_port,
    popen_launch_server,
    run_distributed_test,
)

register_cuda_ci(est_time=30, suite="nightly-2-gpu", nightly=True)
register_amd_ci(est_time=60, suite="nightly-amd", nightly=True)


@contextmanager
def _capture_stdout():
    captured = io.StringIO()
    old_stdout = sys.stdout
    sys.stdout = captured
    try:
        yield captured
    finally:
        sys.stdout = old_stdout


class TestDumperConfig:
    def test_from_env_defaults_match_dataclass_defaults(self):
        assert DumperConfig.from_env() == DumperConfig()

    def test_from_env_bool(self):
        with temp_set_env(DUMPER_ENABLE="1"):
            assert DumperConfig.from_env().enable is True
        with temp_set_env(DUMPER_ENABLE="false"):
            assert DumperConfig.from_env().enable is False

    def test_from_env_str(self):
        with temp_set_env(DUMPER_FILTER="layer_id=0"):
            assert DumperConfig.from_env().filter == "layer_id=0"

    def test_from_env_dir(self):
        with temp_set_env(DUMPER_DIR="/my/dir"):
            assert DumperConfig.from_env().dir == "/my/dir"

    def test_from_env_int(self):
        with temp_set_env(DUMPER_COLLECTIVE_TIMEOUT="120"):
            assert DumperConfig.from_env().collective_timeout == 120

    def test_configure_overrides(self):
        d = _make_test_dumper("/tmp")
        d.configure(enable=False)
        assert d._config.enable is False
        d.configure(enable=True)
        assert d._config.enable is True

    def test_type_validation(self):
        with pytest.raises(TypeError, match="enable.*expected bool.*got str"):
            DumperConfig(enable="yes")
        with pytest.raises(
            TypeError, match="collective_timeout.*expected int.*got str"
        ):
            DumperConfig(collective_timeout="abc")
        with pytest.raises(TypeError, match="filter.*expected str.*got int"):
            DumperConfig(filter=123)

    def test_configure_default_skips_when_env_set(self):
        with temp_set_env(DUMPER_FILTER="from_env"):
            d = _Dumper(config=DumperConfig.from_env())
            d.configure_default(filter="from_code")
            assert d._config.filter == "from_env"

    def test_configure_default_applies_when_no_env(self):
        d = _Dumper(config=DumperConfig.from_env())
        d.configure_default(filter="from_code")
        assert d._config.filter == "from_code"

    def test_from_env_whitespace_treated_as_unset(self):
        with temp_set_env(DUMPER_FILTER="   "):
            assert DumperConfig.from_env().filter is None

    def test_may_enable_default_false(self):
        d = _Dumper(config=DumperConfig())
        assert d.may_enable is False

    def test_may_enable_true_when_enabled(self):
        d = _Dumper(config=DumperConfig(enable=True))
        assert d.may_enable is True

    def test_may_enable_true_when_server_port_set(self):
        d = _Dumper(config=DumperConfig(server_port="40000"))
        assert d.may_enable is True

        d2 = _Dumper(config=DumperConfig(server_port="reuse"))
        assert d2.may_enable is True


class TestServerPortParsed:
    def test_negative_returns_none(self):
        assert DumperConfig(server_port="-1").server_port_parsed is None

    def test_zero_returns_none(self):
        assert DumperConfig(server_port="0").server_port_parsed is None

    def test_positive_returns_int(self):
        result = DumperConfig(server_port="40000").server_port_parsed
        assert result == 40000
        assert isinstance(result, int)

    def test_reuse_returns_string(self):
        assert DumperConfig(server_port="reuse").server_port_parsed == "reuse"


class TestDefaultExpName:
    def test_starts_with_prefix(self):
        name = _get_default_exp_name(timeout_seconds=5)
        assert name.startswith("dump_")

    def test_suffix_format(self):
        name = _get_default_exp_name(timeout_seconds=5)
        suffix = name[len("dump_") :]
        assert len(suffix) == 22
        assert suffix[8] == "_"


class TestKvPairsParsing:
    def test_from_kv_pairs_none_returns_defaults(self):
        assert DumperConfig.from_kv_pairs(None) == DumperConfig()

    def test_from_kv_pairs_empty_returns_defaults(self):
        assert DumperConfig.from_kv_pairs([]) == DumperConfig()

    def test_from_kv_pairs_bool_field(self):
        cfg = DumperConfig.from_kv_pairs(["enable=true"])
        assert cfg.enable is True
        assert cfg.dir == "/tmp/dumper"

    def test_from_kv_pairs_bool_numeric(self):
        assert DumperConfig.from_kv_pairs(["enable=1"]).enable is True
        assert DumperConfig.from_kv_pairs(["enable=0"]).enable is False

    def test_from_kv_pairs_int_field(self):
        cfg = DumperConfig.from_kv_pairs(["collective_timeout=120"])
        assert cfg.collective_timeout == 120
        assert type(cfg.collective_timeout) is int

    def test_from_kv_pairs_int_field_zero_stays_int(self):
        cfg = DumperConfig.from_kv_pairs(["collective_timeout=0"])
        assert cfg.collective_timeout == 0
        assert type(cfg.collective_timeout) is int

    def test_from_kv_pairs_str_field_not_coerced(self):
        cfg = DumperConfig.from_kv_pairs(["server_port=0"])
        assert cfg.server_port == "0"
        assert type(cfg.server_port) is str

    def test_from_kv_pairs_str_field_one_stays_str(self):
        cfg = DumperConfig.from_kv_pairs(["server_port=1"])
        assert cfg.server_port == "1"
        assert type(cfg.server_port) is str

    def test_from_kv_pairs_optional_str_field(self):
        cfg = DumperConfig.from_kv_pairs(
            ["filter=layer_id is not None and layer_id < 3"]
        )
        assert cfg.filter == "layer_id is not None and layer_id < 3"

    def test_from_kv_pairs_optional_str_exp_name(self):
        cfg = DumperConfig.from_kv_pairs(["exp_name=my_experiment"])
        assert cfg.exp_name == "my_experiment"

    def test_from_kv_pairs_multiple_fields(self):
        cfg = DumperConfig.from_kv_pairs(
            [
                "enable=true",
                "dir=/my/dir",
                "filter=name == 'foo'",
                "collective_timeout=30",
                "enable_grad=1",
            ]
        )
        assert cfg.enable is True
        assert cfg.dir == "/my/dir"
        assert cfg.filter == "name == 'foo'"
        assert cfg.collective_timeout == 30
        assert cfg.enable_grad is True

    def test_from_kv_pairs_missing_equals_raises(self):
        with pytest.raises(ValueError, match="missing '='"):
            DumperConfig.from_kv_pairs(["enable"])

    def test_from_kv_pairs_unknown_key_raises(self):
        with pytest.raises(ValueError, match="Unknown config key"):
            DumperConfig.from_kv_pairs(["nonexistent=true"])

    def test_kv_pairs_to_dict_returns_only_explicit(self):
        d = DumperConfig._kv_pairs_to_dict(["enable=true", "dir=/x"])
        assert d == {"enable": True, "dir": "/x"}
        assert "filter" not in d
        assert "collective_timeout" not in d

    def test_kv_pairs_to_dict_none_returns_empty(self):
        assert DumperConfig._kv_pairs_to_dict(None) == {}

    def test_kv_pairs_to_dict_empty_returns_empty(self):
        assert DumperConfig._kv_pairs_to_dict([]) == {}

    def test_from_kv_pairs_value_with_equals_in_value(self):
        cfg = DumperConfig.from_kv_pairs(["filter=name == 'foo'"])
        assert cfg.filter == "name == 'foo'"

    def test_from_kv_pairs_type_validation_still_works(self):
        with pytest.raises(TypeError, match="collective_timeout.*expected int"):
            DumperConfig.from_kv_pairs(["collective_timeout=not_a_number"])


class TestDumperPureFunctions:
    def test_get_truncated_value(self):
        assert get_truncated_value(None) is None
        assert get_truncated_value(42) == 42
        assert len(get_truncated_value((torch.randn(10), torch.randn(20)))) == 2
        assert get_truncated_value(torch.randn(10, 10)).shape == (10, 10)
        assert get_truncated_value(torch.randn(100, 100)).shape == (5, 5)

    def test_obj_to_dict(self):
        assert _obj_to_dict({"a": 1}) == {"a": 1}

        class Obj:
            x, y = 10, 20

            def method(self):
                pass

        result = _obj_to_dict(Obj())
        assert result["x"] == 10
        assert "method" not in result

    def test_deepcopy_or_clone_tensor(self):
        original = torch.randn(3, 3)
        cloned = _deepcopy_or_clone(original)
        assert torch.equal(cloned, original)
        original.fill_(999.0)
        assert not torch.equal(cloned, original)

    def test_deepcopy_or_clone_non_tensor(self):
        original = {"a": [1, 2, 3]}
        cloned = _deepcopy_or_clone(original)
        assert cloned == original
        assert cloned is not original
        original["a"].append(4)
        assert len(cloned["a"]) == 3

    def test_get_tensor_info(self):
        info = get_tensor_info(torch.randn(10, 10))
        for key in ["shape=", "dtype=", "min=", "max=", "mean="]:
            assert key in info

        assert "value=42" in get_tensor_info(42)
        assert "min=None" in get_tensor_info(torch.tensor([]))


class TestMapTensor:
    def test_bare_tensor(self):
        t = torch.randn(4)
        result = _map_tensor(t, lambda x: x * 2)
        assert torch.equal(result, t * 2)

    def test_bare_tensor_no_change(self):
        t = torch.randn(4)
        result = _map_tensor(t, lambda x: x)
        assert result is t

    def test_dict_with_tensor_values(self):
        t1 = torch.randn(3)
        t2 = torch.randn(5)
        value = {"a": t1, "b": t2, "meta": "not a tensor"}
        result = _map_tensor(value, lambda x: x.clone())
        assert torch.equal(result["a"], t1)
        assert torch.equal(result["b"], t2)
        assert result["a"] is not t1
        assert result["b"] is not t2
        assert result["meta"] == "not a tensor"

    def test_dict_no_tensors(self):
        value = {"a": 1, "b": "hello"}
        result = _map_tensor(value, lambda x: x.clone())
        assert result == value

    def test_nested_dict(self):
        inner_t = torch.randn(3)
        value = {"outer": {"inner": inner_t, "label": "ok"}, "top": torch.randn(2)}
        result = _map_tensor(value, lambda x: x.clone())
        assert torch.equal(result["outer"]["inner"], inner_t)
        assert result["outer"]["inner"] is not inner_t
        assert result["outer"]["label"] == "ok"
        assert result is not value
        assert result["outer"] is not value["outer"]

    def test_non_tensor_non_dict(self):
        result = _map_tensor(42, lambda x: x.clone())
        assert result == 42


class TestTorchSave:
    def test_normal(self, tmp_path):
        path = str(tmp_path / "a.pt")
        tensor = torch.randn(3, 3)

        _torch_save(tensor, path)

        assert torch.equal(torch.load(path, weights_only=True), tensor)

    def test_parameter_fallback(self, tmp_path):
        class BadParam(torch.nn.Parameter):
            def __reduce_ex__(self, protocol):
                raise RuntimeError("not pickleable")

        path = str(tmp_path / "b.pt")
        param = BadParam(torch.randn(4))

        _torch_save(param, path)

        assert torch.equal(torch.load(path, weights_only=True), param.data)

    def test_shared_storage_not_bloated(self, tmp_path):
        big = torch.randn(1000, 1000)
        view = big[0]
        path = str(tmp_path / "view.pt")

        _torch_save({"value": view, "meta": {}}, path)

        file_size = Path(path).stat().st_size
        expected_max = view.nelement() * view.element_size() * 10
        assert file_size < expected_max, (
            f"File {file_size} bytes but view is only "
            f"{view.nelement() * view.element_size()} bytes — "
            f"torch.save likely serialized the full "
            f"{big.nelement() * big.element_size()} byte storage"
        )

    def test_silent_skip(self, tmp_path, capsys):
        path = str(tmp_path / "c.pt")

        _torch_save({"fn": lambda: None}, path)

        captured = capsys.readouterr()
        assert "[Dumper, rank=" in captured.out
        assert "Observe error=" in captured.out
        assert "skip the tensor" in captured.out


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)
        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):
        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 TestCollectiveTimeout:
    def test_watchdog_fires_on_timeout(self):
        block_event = threading.Event()
        output = ""

        def run_with_timeout():
            nonlocal output
            with _capture_stdout() as captured:
                _collective_with_timeout(
                    lambda: block_event.wait(),
                    operation_name="test_blocked_op",
                    timeout_seconds=2,
                )
            output = captured.getvalue()

        worker = threading.Thread(target=run_with_timeout)
        worker.start()

        time.sleep(4)
        block_event.set()
        worker.join(timeout=5)

        print(f"Captured output: {output!r}")
        assert "WARNING" in output
        assert "test_blocked_op" in output
        assert "2s" in output


class TestDumperDistributed:
    def test_basic(self, tmp_path):
        with temp_set_env(
            DUMPER_ENABLE="1",
            DUMPER_DIR=str(tmp_path),
        ):
            run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path))

    @staticmethod
    def _test_basic_func(rank, tmpdir):
        tensor = torch.randn(10, 10, device=f"cuda:{rank}")

        dumper.dump("tensor_a", tensor, arg=100)
        dumper.step()

        dumper.set_ctx(ctx_arg=200)
        dumper.dump("tensor_b", tensor)
        dumper.set_ctx(ctx_arg=None)
        dumper.step()

        dumper.configure(filter="False")
        dumper.dump("tensor_skip", tensor)
        dumper.configure(filter=None)
        dumper.step()

        dumper.dump_dict("obj", {"a": torch.randn(3, device=f"cuda:{rank}"), "b": 42})
        dumper.step()

        dist.barrier()
        filenames = _get_filenames(tmpdir)
        _assert_files(
            filenames,
            exist=["tensor_a", "tensor_b", "arg=100", "ctx_arg=200", "obj_a", "obj_b"],
            not_exist=["tensor_skip"],
        )

    def test_collective_timeout(self):
        with temp_set_env(DUMPER_ENABLE="1"):
            run_distributed_test(self._test_collective_timeout_func)

    @staticmethod
    def _test_collective_timeout_func(rank):
        dumper = _Dumper(
            config=DumperConfig(
                enable=True,
                collective_timeout=3,
            ),
        )

        with _capture_stdout() as captured:
            if rank != 0:
                time.sleep(6)
            dumper.step()

        output = captured.getvalue()
        print(f"Rank {rank} captured output: {output!r}")

        if rank == 0:
            assert "WARNING" in output, f"Expected WARNING in rank 0 output: {output}"
            assert "has not completed after 3s" in output

    def test_file_content_correctness(self, tmp_path):
        with temp_set_env(
            DUMPER_ENABLE="1",
            DUMPER_DIR=str(tmp_path),
        ):
            run_distributed_test(self._test_file_content_func, tmpdir=str(tmp_path))

    @staticmethod
    def _test_file_content_func(rank, tmpdir):
        tensor = torch.arange(12, device=f"cuda:{rank}").reshape(3, 4).float()

        dumper.dump("content_check", tensor)
        dumper.step()

        dist.barrier()
        path = _find_dump_file(tmpdir, rank=rank, name="content_check")
        raw = _load_dump(path)
        assert isinstance(raw, dict), f"Expected dict, got {type(raw)}"
        assert "value" in raw and "meta" in raw
        assert torch.equal(raw["value"], tensor.cpu())
        assert raw["meta"]["name"] == "content_check"
        assert raw["meta"]["rank"] == rank


class TestDumperFileWriteControl:
    def test_filter(self, tmp_path):
        with temp_set_env(
            DUMPER_ENABLE="1",
            DUMPER_DIR=str(tmp_path),
            DUMPER_FILTER="name.startswith('keep')",
        ):
            run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))

    @staticmethod
    def _test_filter_func(rank, tmpdir):
        dumper.dump("keep_this", torch.randn(5, device=f"cuda:{rank}"))
        dumper.dump("skip_this", torch.randn(5, device=f"cuda:{rank}"))
        dumper.dump("not_keep_this", torch.randn(5, device=f"cuda:{rank}"))
        dumper.step()

        dist.barrier()
        filenames = _get_filenames(tmpdir)
        _assert_files(
            filenames,
            exist=["keep_this"],
            not_exist=["skip_this", "not_keep_this"],
        )

    def test_save_false(self, tmp_path):
        with temp_set_env(
            DUMPER_ENABLE="1",
            DUMPER_DIR=str(tmp_path),
        ):
            run_distributed_test(self._test_save_false_func, tmpdir=str(tmp_path))

    @staticmethod
    def _test_save_false_func(rank, tmpdir):
        dumper.dump("no_save_tensor", torch.randn(5, device=f"cuda:{rank}"), save=False)
        dumper.step()

        dist.barrier()
        assert len(_get_filenames(tmpdir)) == 0


class TestDumpEnableFlags:
    def test_all_enables_false_no_output(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_value=False, enable_grad=False)
        d.dump("should_skip", torch.randn(3, 3))
        assert len(_get_filenames(tmp_path)) == 0


class TestOutputControl:
    def test_file_enabled_by_default(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        d.dump("file_on", torch.randn(3, 3))

        _assert_files(_get_filenames(tmp_path), exist=["file_on"])

    def test_file_disabled(self, tmp_path, capsys):
        d = _make_test_dumper(tmp_path, enable_output_file=False)
        d.dump("file_off", torch.randn(3, 3))

        assert len(_get_filenames(tmp_path)) == 0
        assert "file_off" in capsys.readouterr().out

    def test_console_enabled_by_default(self, tmp_path, capsys):
        d = _make_test_dumper(tmp_path)
        d.dump("console_on", torch.randn(3, 3))

        captured = capsys.readouterr()
        assert "[Dumper.Value]" in captured.out
        assert "console_on" in captured.out

    def test_console_disabled(self, tmp_path, capsys):
        d = _make_test_dumper(tmp_path, enable_output_console=False)
        d.dump("console_off", torch.randn(3, 3))

        assert "console_off" not in capsys.readouterr().out
        _assert_files(_get_filenames(tmp_path), exist=["console_off"])

    def test_capture_output_basic(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(4, 4)

        with d.capture_output() as captured:
            d.dump("cap_basic", tensor)

        assert "cap_basic" in captured
        assert set(captured["cap_basic"].keys()) == {"value", "meta"}
        assert torch.equal(captured["cap_basic"]["value"], tensor)
        assert captured["cap_basic"]["meta"]["name"] == "cap_basic"

    def test_capture_output_no_file(self, tmp_path):
        d = _make_test_dumper(tmp_path)

        with d.capture_output() as captured:
            d.dump("cap_no_file", torch.randn(3, 3))

        assert "cap_no_file" in captured
        assert len(_get_filenames(tmp_path)) == 0

    def test_capture_output_multiple(self, tmp_path):
        d = _make_test_dumper(tmp_path)

        with d.capture_output() as captured:
            d.dump("first", torch.randn(2, 2))
            d.dump("second", torch.randn(3, 3))

        assert set(captured.keys()) == {"first", "second"}
        assert captured["first"]["value"].shape == (2, 2)
        assert captured["second"]["value"].shape == (3, 3)

    def test_capture_output_value_cloned(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        tensor = torch.zeros(3, 3)

        with d.capture_output() as captured:
            d.dump("clone_check", tensor)

        tensor.fill_(999.0)
        assert torch.equal(captured["clone_check"]["value"], torch.zeros(3, 3))

    def test_capture_output_nested_raises(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        with d.capture_output():
            with pytest.raises(AssertionError):
                with d.capture_output():
                    pass

    def test_capture_output_respects_filter(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="'keep' in name")

        with d.capture_output() as captured:
            d.dump("keep_this", torch.randn(3, 3))
            d.dump("skip_this", torch.randn(3, 3))

        assert "keep_this" in captured
        assert "skip_this" not in captured


class TestDumpDictFormat:
    """Verify that dump files use the dict output format: {"value": ..., "meta": {...}}."""

    def test_dict_format_structure(self, tmp_path):
        dumper = _make_test_dumper(tmp_path)
        tensor = torch.randn(4, 4)
        dumper.dump("fmt_test", tensor, custom_key="hello")

        path = _find_dump_file(str(tmp_path), rank=0, name="fmt_test")
        raw = _load_dump(path)

        assert isinstance(raw, dict)
        assert set(raw.keys()) == {"value", "meta"}
        assert torch.equal(raw["value"], tensor)

        meta = raw["meta"]
        assert meta["name"] == "fmt_test"
        assert meta["custom_key"] == "hello"
        assert "step" in meta
        assert "rank" in meta
        assert "dump_index" in meta

    def test_dict_format_with_context(self, tmp_path):
        dumper = _make_test_dumper(tmp_path)
        dumper.set_ctx(ctx_val=42)
        tensor = torch.randn(2, 2)
        dumper.dump("ctx_fmt", tensor)

        path = _find_dump_file(str(tmp_path), rank=0, name="ctx_fmt")
        raw = _load_dump(path)

        assert raw["meta"]["ctx_val"] == 42
        assert torch.equal(raw["value"], tensor)


def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
    """Create a _Dumper for CPU testing without distributed."""
    defaults = dict(
        enable=True,
        dir=str(tmp_path),
        exp_name="test",
    )
    defaults.update(overrides)
    config = DumperConfig(**defaults)
    return _Dumper(config=config)


def _get_filenames(tmpdir):
    return {f.name for f in Path(tmpdir).glob("*/*.pt")}


def _assert_files(filenames, *, exist=(), not_exist=()):
    for p in exist:
        assert any(p in f for f in filenames), f"{p} not found in {filenames}"
    for p in not_exist:
        assert not any(
            p in f for f in filenames
        ), f"{p} should not exist in {filenames}"


def _load_dump(path: Path) -> dict:
    """Load a dump file and return the raw dict (with 'value' and 'meta' keys)."""
    return torch.load(path, map_location="cpu", weights_only=False)


def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path:
    matches = [
        f
        for f in Path(tmpdir).glob("*/*.pt")
        if f"rank={rank}" in f.name and name in f.name
    ]
    assert (
        len(matches) == 1
    ), f"Expected 1 file matching rank={rank} name={name}, got {matches}"
    return matches[0]


class TestMaterializeValue:
    def test_materialize_value_callable(self):
        tensor = torch.randn(3, 3)
        result = _materialize_value(lambda: tensor)
        assert torch.equal(result, tensor)

    def test_materialize_value_passthrough(self):
        tensor = torch.randn(3, 3)
        result = _materialize_value(tensor)
        assert result is tensor

    def test_dump_with_callable_value(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(4, 4)
        d.dump("lazy_tensor", lambda: tensor)

        _assert_files(_get_filenames(tmp_path), exist=["name=lazy_tensor"])

        path = _find_dump_file(tmp_path, rank=0, name="lazy_tensor")
        assert torch.equal(_load_dump(path)["value"], tensor)


class TestSaveValue:
    def test_dump_output_format(self, tmp_path):
        dumper = _make_test_dumper(tmp_path)
        tensor = torch.randn(4, 4)

        dumper.dump("dict_test", tensor)

        path = _find_dump_file(tmp_path, rank=0, name="dict_test")
        loaded = _load_dump(path)
        assert torch.equal(loaded["value"], tensor)
        assert loaded["meta"]["name"] == "dict_test"
        assert loaded["meta"]["rank"] == 0


class TestStaticMetadata:
    def test_static_meta_contains_world_info(self):
        dumper = _make_test_dumper("/tmp")
        meta = dumper._static_meta
        assert "world_rank" in meta
        assert "world_size" in meta
        assert meta["world_rank"] == 0
        assert meta["world_size"] == 1

    def test_static_meta_caching(self):
        dumper = _make_test_dumper("/tmp")
        meta1 = dumper._static_meta
        meta2 = dumper._static_meta
        assert meta1 is meta2

    def test_parallel_info_graceful_fallback(self):
        sglang_info = _SGLangPlugin().collect_parallel_info()
        assert isinstance(sglang_info, dict)

        megatron_info = _MegatronPlugin().collect_parallel_info()
        assert isinstance(megatron_info, dict)

    def test_dump_includes_static_meta(self, tmp_path):
        dumper = _make_test_dumper(tmp_path)
        tensor = torch.randn(2, 2)

        dumper.dump("meta_test", tensor)

        path = _find_dump_file(tmp_path, rank=0, name="meta_test")
        loaded = _load_dump(path)
        meta = loaded["meta"]
        assert "world_rank" in meta
        assert "world_size" in meta


class TestDumpGrad:
    def test_dump_grad_basic(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=True)
        x = torch.randn(3, 3, requires_grad=True)
        y = (x * 2).sum()

        d.dump("test_tensor", x)
        y.backward()

        filenames = _get_filenames(tmp_path)
        assert any("name=test_tensor" in f and "grad__" not in f for f in filenames)
        _assert_files(filenames, exist=["grad__test_tensor"])

    def test_dump_grad_non_tensor_skipped(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=True)
        d.dump("not_tensor", 42)

        _assert_files(_get_filenames(tmp_path), not_exist=["grad__"])

    def test_dump_grad_no_requires_grad_skipped(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=True)
        x = torch.randn(3, 3, requires_grad=False)
        d.dump("no_grad_tensor", x)

        _assert_files(
            _get_filenames(tmp_path),
            exist=["name=no_grad_tensor"],
            not_exist=["grad__"],
        )

    def test_dump_grad_captures_step(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=True)
        d._state.step = 42
        x = torch.randn(3, 3, requires_grad=True)
        y = (x * 2).sum()

        d.dump("id_test", x)
        d._state.step = 999
        y.backward()

        grad_file = _find_dump_file(tmp_path, name="grad__id_test")
        assert "step=42" in grad_file.name

    def test_dump_grad_file_content(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=True)
        x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
        y = (x * 3).sum()

        d.dump("content_check", x)
        y.backward()

        grad_path = _find_dump_file(tmp_path, name="grad__content_check")
        expected_grad = torch.full((2, 2), 3.0)
        assert torch.equal(_load_dump(grad_path)["value"], expected_grad)

    def test_disable_value(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_value=False, enable_grad=True)
        x = torch.randn(3, 3, requires_grad=True)
        y = (x * 2).sum()

        d.dump("fwd_disabled", x)
        y.backward()

        filenames = _get_filenames(tmp_path)
        assert not any(
            "name=fwd_disabled" in f and "grad__" not in f for f in filenames
        )
        _assert_files(filenames, exist=["grad__fwd_disabled"])

    def test_disable_grad(self, tmp_path):
        d = _make_test_dumper(tmp_path, enable_grad=False)
        x = torch.randn(3, 3, requires_grad=True)
        y = (x * 2).sum()

        d.dump("grad_disabled", x)
        y.backward()

        _assert_files(
            _get_filenames(tmp_path),
            exist=["name=grad_disabled"],
            not_exist=["grad__"],
        )


class TestKvFilter:
    def test_format_tags(self):
        assert _format_tags({"a": 1, "b": "hello"}) == "a=1___b=hello"
        assert _format_tags({}) == ""

    def test_filter_matches_extra_kwargs(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="layer_id == 0")
        d.dump("tensor_a", torch.randn(3), layer_id=0)
        d.dump("tensor_b", torch.randn(3), layer_id=1)

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["tensor_a"], not_exist=["tensor_b"])

    def test_filter_matches_global_ctx(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="ctx_arg == 200")
        d.set_ctx(ctx_arg=200)
        d.dump("tensor_a", torch.randn(3))
        d.set_ctx(ctx_arg=None)
        d.dump("tensor_b", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["tensor_a"], not_exist=["tensor_b"])

    def test_filter_matches_name(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="'keep' in name")
        d.dump("keep_this", torch.randn(3))
        d.dump("skip_this", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["keep_this"], not_exist=["skip_this"])

    def test_filter_expr_range(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="layer_id is not None and layer_id < 3")
        d.dump("t0", torch.randn(3), layer_id=0)
        d.dump("t1", torch.randn(3), layer_id=1)
        d.dump("t5", torch.randn(3), layer_id=5)

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["name=t0", "name=t1"], not_exist=["name=t5"])

    def test_filter_expr_with_none(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="layer_id is None or layer_id < 3")
        d.dump("no_layer", torch.randn(3))
        d.dump("layer0", torch.randn(3), layer_id=0)
        d.dump("layer5", torch.randn(3), layer_id=5)

        filenames = _get_filenames(tmp_path)
        _assert_files(
            filenames,
            exist=["no_layer", "layer0"],
            not_exist=["layer5"],
        )

    def test_filter_expr_with_re_search(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="search(r'attn|mlp', name)")
        d.dump("self_attn", torch.randn(3))
        d.dump("mlp_proj", torch.randn(3))
        d.dump("layernorm", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(
            filenames,
            exist=["self_attn", "mlp_proj"],
            not_exist=["layernorm"],
        )

    def test_filter_expr_syntax_error(self, tmp_path):
        d = _make_test_dumper(tmp_path, filter="layer_id ===")
        with pytest.raises(SyntaxError):
            d.dump("tensor", torch.randn(3))

    def test_no_filter_dumps_all(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        d.dump("a", torch.randn(3))
        d.dump("b", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["name=a", "name=b"])


class TestDumpModel:
    def test_grad_basic(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2)
        x = torch.randn(3, 4)
        y = model(x).sum()
        y.backward()

        d.dump_model(model, name_prefix="model")

        _assert_files(
            _get_filenames(tmp_path),
            exist=["grad__model__weight", "grad__model__bias"],
        )

    def test_value_basic(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )
        model = torch.nn.Linear(4, 2, bias=False)

        d.dump_model(model, name_prefix="model")

        _assert_files(
            _get_filenames(tmp_path),
            exist=["model__weight"],
        )

    def test_no_grad_skipped(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2)

        d.dump_model(model, name_prefix="model")

        filenames = _get_filenames(tmp_path)
        assert len(filenames) == 0

    def test_filter(self, tmp_path):
        d = _make_test_dumper(
            tmp_path,
            enable_model_value=True,
            enable_model_grad=True,
            filter="'weight' in name",
        )
        model = torch.nn.Linear(4, 2)
        x = torch.randn(3, 4)
        y = model(x).sum()
        y.backward()

        d.dump_model(model, name_prefix="model")

        _assert_files(
            _get_filenames(tmp_path),
            exist=["model__weight", "grad__model__weight"],
            not_exist=["model__bias", "grad__model__bias"],
        )

    def test_grad_file_content(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        x = torch.ones(1, 4)
        y = model(x).sum()
        y.backward()

        d.dump_model(model, name_prefix="p")

        path = _find_dump_file(tmp_path, name="grad__p__weight")
        assert torch.equal(_load_dump(path)["value"], model.weight.grad)

    def test_disable_model_grad(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )
        model = torch.nn.Linear(4, 2)
        x = torch.randn(3, 4)
        y = model(x).sum()
        y.backward()

        d.dump_model(model, name_prefix="model")

        filenames = _get_filenames(tmp_path)
        assert all("grad" not in f for f in filenames)

    def test_parameter_saved_as_parameter(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )
        model = torch.nn.Linear(4, 2, bias=False)

        d.dump_model(model, name_prefix="p")

        path = _find_dump_file(tmp_path, name="p__weight")
        loaded = _load_dump(path)
        assert isinstance(loaded["value"], torch.nn.Parameter)
        assert torch.equal(loaded["value"], model.weight)

    def test_unpicklable_parameter_falls_back_to_data(self, tmp_path):
        class BadParam(torch.nn.Parameter):
            def __reduce_ex__(self, protocol):
                raise RuntimeError("not pickleable")

        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        model.weight = BadParam(model.weight.data)

        d.dump_model(model, name_prefix="p")

        path = _find_dump_file(tmp_path, name="p__weight")
        loaded = _load_dump(path)
        assert isinstance(loaded["value"], torch.Tensor)
        assert not isinstance(loaded["value"], torch.nn.Parameter)
        assert torch.equal(loaded["value"], model.weight.data)

    def test_disable_model_value(self, tmp_path):
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        x = torch.ones(1, 4)
        y = model(x).sum()
        y.backward()

        d.dump_model(model, name_prefix="model")

        filenames = _get_filenames(tmp_path)
        assert all("grad" in f for f in filenames)


class TestDumpModelGradInjection:
    def test_get_grad_overrides_param_grad(self, tmp_path):
        """dump_model dumps the tensor returned by get_grad instead of param.grad."""
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        y = model(torch.ones(1, 4)).sum()
        y.backward()
        injected = torch.full_like(model.weight, 7.0)

        d.dump_model(model, name_prefix="p", get_grad=lambda param: injected)

        path = _find_dump_file(tmp_path, name="grad__p__weight")
        assert torch.equal(_load_dump(path)["value"], injected)

    def test_get_grad_reads_custom_grad_storage(self, tmp_path):
        """get_grad lets callers surface grads living outside param.grad (e.g. main_grad)."""
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        assert model.weight.grad is None
        model.weight.main_grad = torch.full_like(model.weight, 3.0)

        d.dump_model(
            model,
            name_prefix="p",
            get_grad=lambda param: getattr(param, "main_grad", None),
        )

        path = _find_dump_file(tmp_path, name="grad__p__weight")
        assert torch.equal(_load_dump(path)["value"], model.weight.main_grad)

    def test_get_grad_returning_none_skips_grad_dump(self, tmp_path):
        """A get_grad returning None suppresses the grad dump for that param."""
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        y = model(torch.ones(1, 4)).sum()
        y.backward()

        d.dump_model(model, name_prefix="p", get_grad=lambda param: None)

        assert len(_get_filenames(tmp_path)) == 0

    def test_default_uses_param_grad_without_main_grad_fallback(self, tmp_path):
        """Without get_grad, only param.grad is dumped; a main_grad attribute is ignored."""
        d = _make_test_dumper(
            tmp_path, enable_model_grad=True, enable_model_value=False
        )
        model = torch.nn.Linear(4, 2, bias=False)
        assert model.weight.grad is None
        model.weight.main_grad = torch.full_like(model.weight, 3.0)

        d.dump_model(model, name_prefix="p")

        assert len(_get_filenames(tmp_path)) == 0


class TestDumpModelStepOverride:
    def test_step_override_in_filenames(self, tmp_path):
        """An explicit step overrides the dumper's ambient step for value and grad files."""
        d = _make_test_dumper(tmp_path, enable_model_value=True, enable_model_grad=True)
        d._state.step = 3
        model = torch.nn.Linear(4, 2, bias=False)
        y = model(torch.ones(1, 4)).sum()
        y.backward()

        d.dump_model(model, name_prefix="p", step=7)

        filenames = _get_filenames(tmp_path)
        assert len(filenames) == 2
        assert all("step=7" in f for f in filenames)

    def test_step_none_uses_ambient_step(self, tmp_path):
        """Without an explicit step, dump_model records the dumper's current step."""
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )
        d._state.step = 3
        model = torch.nn.Linear(4, 2, bias=False)

        d.dump_model(model, name_prefix="p")

        filenames = _get_filenames(tmp_path)
        assert filenames and all("step=3" in f for f in filenames)


class TestParallelRankInFilename:
    def test_config_default_false(self):
        """include_parallel_rank_in_filename defaults to False."""
        assert DumperConfig().include_parallel_rank_in_filename is False

    def test_config_from_kv_pairs(self):
        """include_parallel_rank_in_filename is parsed as a bool from kv pairs."""
        cfg = DumperConfig.from_kv_pairs(["include_parallel_rank_in_filename=true"])
        assert cfg.include_parallel_rank_in_filename is True

    def test_collect_tags_merges_keys_across_plugins(self, monkeypatch):
        """_collect_parallel_rank_tags keeps only rank keys, merging across plugins."""
        plugin_a = type(
            "PluginA",
            (),
            {"collect_parallel_info": lambda self: {"pp_rank": 1, "ignored": 9}},
        )()
        plugin_b = type(
            "PluginB",
            (),
            {"collect_parallel_info": lambda self: {"tp_rank": 2, "cp_rank": 3}},
        )()
        monkeypatch.setattr(
            "sglang.srt.debug_utils.dumper._plugins", [plugin_a, plugin_b]
        )

        tags = _collect_parallel_rank_tags()
        assert tags == {"pp_rank": 1, "tp_rank": 2, "cp_rank": 3}

    def test_collect_tags_first_plugin_wins_on_conflict(self, monkeypatch):
        """When two plugins report the same rank key, the first plugin wins."""
        plugin_a = type(
            "PluginA", (), {"collect_parallel_info": lambda self: {"pp_rank": 1}}
        )()
        plugin_b = type(
            "PluginB", (), {"collect_parallel_info": lambda self: {"pp_rank": 7}}
        )()
        monkeypatch.setattr(
            "sglang.srt.debug_utils.dumper._plugins", [plugin_a, plugin_b]
        )

        assert _collect_parallel_rank_tags() == {"pp_rank": 1}

    def test_collect_tags_skips_empty_plugin_info(self, monkeypatch):
        """Plugins that report no parallel info are skipped without error."""
        plugin_empty = type(
            "PluginEmpty", (), {"collect_parallel_info": lambda self: {}}
        )()
        plugin_real = type(
            "PluginReal", (), {"collect_parallel_info": lambda self: {"tp_rank": 4}}
        )()
        monkeypatch.setattr(
            "sglang.srt.debug_utils.dumper._plugins", [plugin_empty, plugin_real]
        )

        assert _collect_parallel_rank_tags() == {"tp_rank": 4}

    def test_disabled_filename_has_no_rank_tags(self, tmp_path, monkeypatch):
        """When disabled, dump filenames do not include parallel-rank tags."""
        monkeypatch.setattr(
            _SGLangPlugin,
            "collect_parallel_info",
            lambda self: {"pp_rank": 2, "tp_rank": 3},
        )

        d = _make_test_dumper(tmp_path, include_parallel_rank_in_filename=False)
        d.dump("hidden", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        assert filenames
        assert all("pp_rank=" not in f and "tp_rank=" not in f for f in filenames)

    def test_enabled_filename_includes_rank_tags(self, tmp_path, monkeypatch):
        """When enabled, dump filenames include the collected parallel-rank tags."""
        monkeypatch.setattr(
            _SGLangPlugin,
            "collect_parallel_info",
            lambda self: {"pp_rank": 2, "tp_rank": 3},
        )

        d = _make_test_dumper(tmp_path, include_parallel_rank_in_filename=True)
        d.dump("hidden", torch.randn(3))

        filenames = _get_filenames(tmp_path)
        assert any("pp_rank=2" in f and "tp_rank=3" in f for f in filenames)


class TestTransformModelParamName:
    def test_base_plugin_returns_none(self):
        """The default plugin hook keeps the original name (returns None)."""
        plugin = _SGLangPlugin()
        assert (
            plugin.transform_model_param_name(torch.nn.Linear(2, 2), "layers.0.weight")
            is None
        )

    def test_dump_model_keeps_name_without_transform(self, tmp_path):
        """With no plugin rewriting names, dump_model uses the original param name."""
        model = torch.nn.Module()
        model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)])
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )

        d.dump_model(model, name_prefix="m")

        _assert_files(_get_filenames(tmp_path), exist=["m__layers.0.weight"])

    def test_dump_model_applies_plugin_transform(self, tmp_path, monkeypatch):
        """dump_model rewrites param names through the plugin transform hook."""

        def _shift_layers(self, model, param_name):
            return re.sub(
                r"layers\.(\d+)", lambda m: f"layers.{int(m.group(1)) + 4}", param_name
            )

        monkeypatch.setattr(_SGLangPlugin, "transform_model_param_name", _shift_layers)

        model = torch.nn.Module()
        model.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2, bias=False)])
        d = _make_test_dumper(
            tmp_path, enable_model_value=True, enable_model_grad=False
        )

        d.dump_model(model, name_prefix="m")

        _assert_files(
            _get_filenames(tmp_path),
            exist=["m__layers.4.weight"],
            not_exist=["m__layers.0.weight"],
        )

    def test_get_model_config_unwraps_module_chain(self):
        """_get_model_config peels nested .module wrappers to find .config."""
        config = object()
        leaf = type("Leaf", (), {"config": config})()
        wrapped = type("W", (), {"module": type("W2", (), {"module": leaf})()})()
        assert _MegatronPlugin._get_model_config(wrapped) is config

    def test_get_model_config_returns_none_when_absent(self):
        """_get_model_config returns None when no .config is reachable."""
        assert _MegatronPlugin._get_model_config(object()) is None


class _FakeMpu:
    _pp_size = 2

    @classmethod
    def get_pipeline_model_parallel_world_size(cls):
        return cls._pp_size


class TestMegatronTransformModelParamName:
    @pytest.fixture(autouse=True)
    def _patch_megatron(self, monkeypatch):
        monkeypatch.setattr(_MegatronPlugin, "_available", True)
        monkeypatch.setattr(_MegatronPlugin, "_mpu", _FakeMpu, raising=False)
        monkeypatch.setattr(
            _MegatronPlugin,
            "_get_model_config",
            staticmethod(lambda model: object()),
        )
        monkeypatch.setattr(
            _MegatronPlugin,
            "_get_transformer_layer_offset",
            staticmethod(lambda config: 4),
        )

    def test_remaps_local_layer_index_to_global(self):
        """layers.N is shifted by the PP-stage offset; other tokens untouched."""
        plugin = _MegatronPlugin()
        result = plugin.transform_model_param_name(object(), "decoder.layers.0.weight")
        assert result == "decoder.layers.4.weight"

    def test_remaps_all_layer_occurrences(self):
        """Every ``layers.N`` occurrence in the name is shifted."""
        plugin = _MegatronPlugin()
        result = plugin.transform_model_param_name(object(), "layers.1.x.layers.2.y")
        assert result == "layers.5.x.layers.6.y"

    def test_returns_none_when_not_available(self, monkeypatch):
        """No transform when megatron is unavailable."""
        monkeypatch.setattr(_MegatronPlugin, "_available", False)
        plugin = _MegatronPlugin()
        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None

    def test_returns_none_when_pp_size_one(self, monkeypatch):
        """No transform without pipeline parallelism (pp_size == 1)."""
        monkeypatch.setattr(_FakeMpu, "_pp_size", 1)
        plugin = _MegatronPlugin()
        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None
        monkeypatch.setattr(_FakeMpu, "_pp_size", 2)

    def test_returns_none_when_offset_zero(self, monkeypatch):
        """A zero offset (e.g. first PP stage) leaves the name unchanged (None)."""
        monkeypatch.setattr(
            _MegatronPlugin,
            "_get_transformer_layer_offset",
            staticmethod(lambda config: 0),
        )
        plugin = _MegatronPlugin()
        assert plugin.transform_model_param_name(object(), "layers.0.weight") is None


class TestCleanup:
    def test_cleanup_removes_old_dumps(self, tmp_path):
        old_dir = tmp_path / "dump_old"
        old_dir.mkdir()
        (old_dir / "dummy.pt").touch()

        dumper = _make_test_dumper(tmp_path, cleanup_previous=True)
        dumper.dump("new_tensor", torch.randn(3, 3))

        assert not old_dir.exists()
        _assert_files(_get_filenames(tmp_path), exist=["new_tensor"])

    def test_cleanup_removes_exp_name_dir(self, tmp_path):
        exp_name = "my_custom_exp"
        old_exp_dir = tmp_path / exp_name
        old_exp_dir.mkdir()
        (old_exp_dir / "old_data.pt").touch()

        dumper = _make_test_dumper(tmp_path, exp_name=exp_name, cleanup_previous=True)
        dumper.dump("new_tensor", torch.randn(3, 3))

        assert not (tmp_path / exp_name / "old_data.pt").exists()
        _assert_files(_get_filenames(tmp_path), exist=["new_tensor"])

    def test_cleanup_removes_both_dump_prefix_and_exp_name(self, tmp_path):
        old_dump = tmp_path / "dump_old"
        old_dump.mkdir()
        (old_dump / "dummy.pt").touch()

        exp_name = "custom_run"
        old_exp = tmp_path / exp_name
        old_exp.mkdir()
        (old_exp / "stale.pt").touch()

        dumper = _make_test_dumper(tmp_path, exp_name=exp_name, cleanup_previous=True)
        dumper.dump("new_tensor", torch.randn(3, 3))

        assert not old_dump.exists()
        assert not (tmp_path / exp_name / "stale.pt").exists()
        _assert_files(_get_filenames(tmp_path), exist=["new_tensor"])

    def test_no_cleanup_by_default(self, tmp_path):
        old_dir = tmp_path / "dump_old"
        old_dir.mkdir()
        (old_dir / "dummy.pt").touch()

        dumper = _make_test_dumper(tmp_path)
        dumper.dump("new_tensor", torch.randn(3, 3))

        assert old_dir.exists()
        _assert_files(_get_filenames(tmp_path), exist=["new_tensor"])


class TestReset:
    def test_reset_clears_state(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        d.set_ctx(layer_id=1)
        d.dump("before_reset", torch.randn(3, 3))

        d.reset()

        assert d._state.dump_index == 0
        assert d._state.step == 0
        assert d._state.global_ctx == {}

    def test_dump_works_after_reset(self, tmp_path):
        d = _make_test_dumper(tmp_path)
        d.dump("pre", torch.randn(3, 3))

        d.reset()
        d.dump("post", torch.randn(3, 3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["pre", "post"])
        post_file = _find_dump_file(tmp_path, name="post")
        assert "dump_index=1" in post_file.name

    def test_cleanup_previous_re_triggers_after_reset(self, tmp_path):
        """Miles pattern: reset() + configure(cleanup_previous=True) should re-clean."""
        exp_alpha = "exp_alpha"
        exp_beta = "exp_beta"

        (tmp_path / exp_alpha).mkdir()
        (tmp_path / exp_alpha / "stale.pt").touch()
        (tmp_path / exp_beta).mkdir()
        (tmp_path / exp_beta / "stale.pt").touch()

        d = _make_test_dumper(tmp_path, exp_name=exp_alpha, cleanup_previous=True)
        d.dump("phase1", torch.randn(2, 2))

        d.reset()
        d.configure(exp_name=exp_beta, cleanup_previous=True)
        d.dump("phase2", torch.randn(2, 2))

        assert not (tmp_path / exp_alpha / "stale.pt").exists()
        assert not (tmp_path / exp_beta / "stale.pt").exists()
        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["phase1", "phase2"])

    def test_no_cleanup_when_config_false(self, tmp_path):
        """cleanup_previous=False: handled stays False but no cleanup runs."""
        old_dir = tmp_path / "dump_old"
        old_dir.mkdir()
        (old_dir / "dummy.pt").touch()

        d = _make_test_dumper(tmp_path, cleanup_previous=False)
        d.dump("tensor", torch.randn(2, 2))

        assert old_dir.exists()
        assert d._state.cleanup_previous_handled is False

    def test_multi_phase_switch(self, tmp_path):
        """Simulate Miles multi-phase: configure → dump → reset → configure new phase → dump."""
        d = _make_test_dumper(tmp_path, cleanup_previous=True)

        d.configure(exp_name="fwd_only")
        d.dump("weight", torch.randn(2, 2))
        d.step()
        d.configure(enable=False)

        d.reset()
        d.configure(exp_name="fwd_bwd", enable=True, cleanup_previous=True)
        d.dump("weight", torch.randn(2, 2))
        d.step()

        fwd_only_files = list(Path(tmp_path).glob("fwd_only/*.pt"))
        fwd_bwd_files = list(Path(tmp_path).glob("fwd_bwd/*.pt"))
        assert len(fwd_only_files) > 0
        assert len(fwd_bwd_files) > 0
        assert d._state.step == 1
        assert d._state.dump_index == 1

    def test_reset_removes_non_intrusive_hooks(self, tmp_path):
        model = torch.nn.Sequential(
            torch.nn.Linear(4, 4),
            torch.nn.ReLU(),
            torch.nn.Linear(4, 4),
        )
        d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
        d.register_non_intrusive_dumper(model)

        x = torch.randn(2, 4)
        with d.capture_output() as captured:
            model(x)
        assert len(captured) > 0

        d.reset()
        d.configure(enable=True, dir=str(tmp_path), non_intrusive_mode="all")

        with d.capture_output() as captured_after:
            model(x)
        assert len(captured_after) == 0

    def test_reset_removes_non_intrusive_hooks_multiple_models(self, tmp_path):
        model_a = torch.nn.Sequential(
            torch.nn.Linear(4, 4),
            torch.nn.ReLU(),
        )
        model_b = torch.nn.Sequential(
            torch.nn.Linear(4, 4),
            torch.nn.ReLU(),
        )
        d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
        d.register_non_intrusive_dumper(model_a)
        d.register_non_intrusive_dumper(model_b)

        x = torch.randn(2, 4)
        with d.capture_output() as captured:
            model_a(x)
            model_b(x)
        assert len(captured) > 0

        d.reset()
        d.configure(enable=True, dir=str(tmp_path), non_intrusive_mode="all")

        with d.capture_output() as captured_a:
            model_a(x)
        assert len(captured_a) == 0

        with d.capture_output() as captured_b:
            model_b(x)
        assert len(captured_b) == 0


def _dumper_worker(rank, http_port: int, stop_event):
    """Minimal distributed dumper worker: configure, step (triggers ZMQ setup), then wait."""
    dumper.configure(enable=False, server_port=str(http_port))
    dumper.step()
    stop_event.wait()


def _wait_for_dumper_http(url: str, timeout: float = 30) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            requests.post(f"{url}/dumper/configure", json={}, timeout=2)
            return
        except requests.ConnectionError:
            time.sleep(0.5)
    raise TimeoutError(f"Dumper HTTP server not reachable at {url}")


class TestZmqPortIsolation:
    """Multiple independent dumper instances (each with 2 ranks) must not conflict on ZMQ ports."""

    NUM_INSTANCES = 3

    def test_concurrent_instances_no_port_conflict(self):
        ports = [
            find_available_port(40000 + i * 1000) for i in range(self.NUM_INSTANCES)
        ]
        stop_events = []
        threads = []
        ctx = multiprocessing.get_context("spawn")

        for port in ports:
            stop_event = ctx.Event()
            stop_events.append(stop_event)
            thread = threading.Thread(
                target=run_distributed_test,
                args=(_dumper_worker,),
                kwargs={"http_port": port, "stop_event": stop_event},
            )
            thread.start()
            threads.append(thread)

        try:
            for port in ports:
                _wait_for_dumper_http(f"http://127.0.0.1:{port}")

            for i, port in enumerate(ports):
                resp = requests.post(
                    f"http://127.0.0.1:{port}/dumper/get_state", json={}
                )
                resp.raise_for_status()
                states = resp.json()
                assert (
                    len(states) == 2
                ), f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}"
        finally:
            for event in stop_events:
                event.set()
            for thread in threads:
                thread.join(timeout=10)


class TestDumperHttp:
    """Test /dumper/* HTTP control — parametrized over standalone vs sglang server."""

    @pytest.fixture(scope="class", params=["standalone", "sglang"])
    def dumper_http_url(self, request):
        if request.param == "standalone":
            http_port = find_available_port(40000)
            base_url = f"http://127.0.0.1:{http_port}"
            stop_event = multiprocessing.get_context("spawn").Event()
            thread = threading.Thread(
                target=run_distributed_test,
                args=(_dumper_worker,),
                kwargs={"http_port": http_port, "stop_event": stop_event},
            )
            thread.start()
            try:
                _wait_for_dumper_http(base_url)
                yield base_url
            finally:
                stop_event.set()
                thread.join(timeout=10)
        else:
            base_url = DEFAULT_URL_FOR_TEST
            env = {**os.environ, "DUMPER_SERVER_PORT": "reuse"}
            proc = popen_launch_server(
                "Qwen/Qwen3-0.6B",
                base_url,
                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
                other_args=["--max-total-tokens", "128"],
                env=env,
            )
            try:
                yield base_url
            finally:
                kill_process_tree(proc.pid)

    @staticmethod
    def _post(base_url: str, method: str, **kwargs) -> list[dict]:
        resp = requests.post(f"{base_url}/dumper/{method}", json=kwargs or None)
        resp.raise_for_status()
        states = resp.json()
        assert isinstance(states, list) and len(states) >= 1
        return states

    @staticmethod
    def _assert_all_ranks(states: list[dict], path: str, expected):
        """Assert that ``state[path]`` equals ``expected`` on every rank."""
        keys = path.split(".")
        for rank, state in enumerate(states):
            val = state
            for k in keys:
                val = val[k]
            assert (
                val == expected
            ), f"rank {rank}: {path}={val!r}, expected {expected!r}"

    def test_configure_enable_toggle(self, dumper_http_url: str):
        for enable in [True, False]:
            self._post(dumper_http_url, "configure", enable=enable)
            states = self._post(dumper_http_url, "get_state")
            self._assert_all_ranks(states, "config.enable", enable)

    def test_configure_multi_field(self, dumper_http_url: str):
        self._post(
            dumper_http_url,
            "configure",
            enable=True,
            filter="layer_id == 0",
            dir="/tmp/test_http",
        )
        states = self._post(dumper_http_url, "get_state")
        self._assert_all_ranks(states, "config.enable", True)
        self._assert_all_ranks(states, "config.filter", "layer_id == 0")
        self._assert_all_ranks(states, "config.dir", "/tmp/test_http")

    def test_configure_clear_optional(self, dumper_http_url: str):
        self._post(dumper_http_url, "configure", filter="layer_id == 0")
        self._post(dumper_http_url, "configure", filter=None)
        states = self._post(dumper_http_url, "get_state")
        self._assert_all_ranks(states, "config.filter", None)

    def test_reset(self, dumper_http_url: str):
        self._post(dumper_http_url, "configure", enable=True)
        self._post(dumper_http_url, "reset")
        states = self._post(dumper_http_url, "get_state")
        self._assert_all_ranks(states, "dump_index", 0)
        self._assert_all_ranks(states, "step", 0)

    def test_get_state(self, dumper_http_url: str):
        self._post(
            dumper_http_url,
            "configure",
            enable=True,
            filter="layer_id is not None and layer_id < 3",
        )
        states = self._post(dumper_http_url, "get_state")
        self._assert_all_ranks(states, "config.enable", True)
        self._assert_all_ranks(
            states, "config.filter", "layer_id is not None and layer_id < 3"
        )
        for state in states:
            assert "dump_index" in state
            assert "step" in state

    def test_all_ranks_consistent(self, dumper_http_url: str):
        self._post(dumper_http_url, "configure", enable=True, dir="/tmp/multi")
        states = self._post(dumper_http_url, "get_state")
        configs = [s["config"] for s in states]
        for rank_config in configs[1:]:
            assert rank_config == configs[0], f"rank configs diverged: {configs}"

    def test_error_unknown_field(self, dumper_http_url: str):
        resp = requests.post(
            f"{dumper_http_url}/dumper/configure",
            json={"nonexistent_field": 123},
        )
        assert resp.status_code == 400

    def test_error_unknown_method(self, dumper_http_url: str):
        resp = requests.post(
            f"{dumper_http_url}/dumper/nonexistent",
            json={},
        )
        assert resp.status_code == 400

    def test_error_wrong_type(self, dumper_http_url: str):
        resp = requests.post(
            f"{dumper_http_url}/dumper/configure",
            json={"enable": "not_a_bool"},
        )
        assert resp.status_code == 400


class TestRegisterForwardHookOrReplaceFn:
    def test_unknown_mode_raises(self):
        module = torch.nn.Linear(4, 4)
        with pytest.raises(ValueError, match="Unknown mode"):
            _register_forward_hook_or_replace_fn(
                module,
                pre_hook=lambda _mod, _input: None,
                hook=lambda _mod, _input, _output: None,
                mode="bad",
            )


class _NonIntrusiveTestBase:
    _PREFIX = "non_intrusive__"

    @staticmethod
    def _assert_captured_contains(
        captured: dict, expected: list[str], prefix: str = "non_intrusive__"
    ) -> None:
        for suffix in expected:
            key = f"{prefix}{suffix}"
            assert key in captured, f"missing {key}"

    @staticmethod
    def _wrap_as_outer(inner_cls: type) -> torch.nn.Module:
        """Wrap an inner module class as OuterModel.model, mimicking typical model nesting."""

        class OuterModel(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.model = inner_cls()

            def forward(self, *args, **kwargs):
                return self.model(*args, **kwargs)

        return OuterModel()

    @staticmethod
    def _make_dumper(tmp_path, **overrides) -> "_Dumper":
        return _make_test_dumper(tmp_path, non_intrusive_mode="all", **overrides)

    def _run(self, tmp_path, inner_cls, **dumper_overrides):
        d = self._make_dumper(tmp_path, **dumper_overrides)
        model = self._wrap_as_outer(inner_cls)
        d.register_non_intrusive_dumper(model)
        x = torch.randn(2, 4)
        with d.capture_output() as captured:
            output = model(x)
        return captured, x, output


class TestNonIntrusiveDumper(_NonIntrusiveTestBase):
    """Tests for mode='all' — hooks on every module, non_intrusive__ prefix."""

    def test_basic_inputs_and_outputs(self, tmp_path):
        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                return self.relu(self.linear(x))

        captured, x, output = self._run(tmp_path, Inner)

        self._assert_captured_contains(
            captured,
            [
                "output",
                "inputs.0",
                "model.output",
                "model.inputs.0",
                "model.linear.output",
                "model.linear.inputs.0",
                "model.relu.output",
                "model.relu.inputs.0",
            ],
        )
        P = self._PREFIX
        assert torch.allclose(captured[f"{P}output"]["value"], output)

    def test_inputs_dumped_before_forward(self, tmp_path):
        """Inputs are captured *before* forward(); in-place mutation must not affect them."""

        class Mutator(torch.nn.Module):
            def forward(self, x):
                x.fill_(999.0)
                return x

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mutator = Mutator()

            def forward(self, x):
                return self.mutator(x)

        d = self._make_dumper(tmp_path)
        model = self._wrap_as_outer(Inner)
        d.register_non_intrusive_dumper(model)

        x = torch.randn(2, 4)
        original_x = x.clone()
        with d.capture_output() as captured:
            model(x)

        P = self._PREFIX
        dumped_input = captured[f"{P}model.mutator.inputs.0"]["value"]
        assert torch.allclose(dumped_input, original_x), (
            f"pre-hook should capture inputs before forward mutates them; "
            f"got {dumped_input} but expected {original_x}"
        )

        dumped_output = captured[f"{P}model.mutator.output"]["value"]
        assert (
            dumped_output == 999.0
        ).all(), "post-hook should capture outputs after forward"

    def test_hooks_all_module_levels(self, tmp_path):
        class Attention(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.qkv_proj = torch.nn.Linear(4, 12)
                self.o_proj = torch.nn.Linear(4, 4)

            def forward(self, x):
                _qkv = self.qkv_proj(x)
                return self.o_proj(x)

        class Layer(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.self_attn = Attention()
                self.mlp = torch.nn.Linear(4, 4)

            def forward(self, x):
                x = self.self_attn(x)
                return self.mlp(x)

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = torch.nn.ModuleList([Layer()])

            def forward(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        self._assert_captured_contains(
            captured,
            [
                "output",
                "model.output",
                "model.layers.0.output",
                "model.layers.0.self_attn.output",
                "model.layers.0.self_attn.qkv_proj.output",
                "model.layers.0.self_attn.o_proj.output",
                "model.layers.0.mlp.output",
                "model.layers.0.self_attn.qkv_proj.inputs.0",
                "model.layers.0.self_attn.o_proj.inputs.0",
                "model.layers.0.mlp.inputs.0",
            ],
        )
        P = self._PREFIX
        assert f"{P}model.layers.output" not in captured

    def test_multi_tensor_tuple_output(self, tmp_path):
        class TupleModule(torch.nn.Module):
            def forward(self, x):
                return x, x * 2

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.split = TupleModule()
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, x):
                a, b = self.split(x)
                return self.linear(a + b)

        captured, x, output = self._run(tmp_path, Inner)

        assert "non_intrusive__model.split.output.0" in captured
        assert "non_intrusive__model.split.output.1" in captured
        assert torch.allclose(
            captured["non_intrusive__model.split.output.0"]["value"], x
        )

    def test_single_tensor_tuple_collapses(self, tmp_path):
        class SingleTupleModule(torch.nn.Module):
            def forward(self, x):
                return (x * 3,)

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.wrap = SingleTupleModule()

            def forward(self, x):
                return self.wrap(x)[0]

        captured, x, output = self._run(tmp_path, Inner)

        assert "non_intrusive__model.wrap.output" in captured
        assert "non_intrusive__model.wrap.output.0" not in captured

    def test_multiple_forward_inputs(self, tmp_path):
        class TwoInputModule(torch.nn.Module):
            def forward(self, x, mask):
                return x * mask

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.mul = TwoInputModule()

            def forward(self, x):
                mask = torch.ones_like(x)
                return self.mul(x, mask)

        captured, x, output = self._run(tmp_path, Inner)

        assert "non_intrusive__model.mul.inputs.0" in captured
        assert "non_intrusive__model.mul.inputs.1" in captured

    def test_none_output_only_dumps_inputs(self, tmp_path):
        class NoneModule(torch.nn.Module):
            def forward(self, x):
                return None

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.sink = NoneModule()

            def forward(self, x):
                self.sink(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        assert "non_intrusive__model.sink.inputs.0" in captured
        assert not any(
            k.startswith("non_intrusive__model.sink.output") for k in captured
        )

    def test_non_tensor_value_silently_skipped(self, tmp_path):
        class IntModule(torch.nn.Module):
            def forward(self, x):
                return 42

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.const = IntModule()

            def forward(self, x):
                self.const(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        assert "non_intrusive__model.const.inputs.0" in captured
        assert not any(
            k.startswith("non_intrusive__model.const.output") for k in captured
        )

    def test_root_module_name_no_malformed_dots(self, tmp_path):
        d = self._make_dumper(tmp_path)
        model = torch.nn.Linear(4, 4)
        d.register_non_intrusive_dumper(model)

        x = torch.randn(2, 4)
        with d.capture_output() as captured:
            model(x)

        for key in captured:
            assert not key.startswith("non_intrusive__."), f"malformed key: {key}"
            assert ".." not in key, f"double dot in key: {key}"

        assert "non_intrusive__output" in captured
        assert "non_intrusive__inputs.0" in captured

    def test_respects_dumper_filter(self, tmp_path):
        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)
                self.relu = torch.nn.ReLU()

            def forward(self, x):
                return self.relu(self.linear(x))

        captured, x, output = self._run(
            tmp_path, Inner, filter="name == 'non_intrusive__model.linear.output'"
        )

        assert "non_intrusive__model.linear.output" in captured
        assert "non_intrusive__model.relu.output" not in captured
        assert "non_intrusive__model.linear.inputs.0" not in captured

    def test_disabled_dumper_no_output(self, tmp_path):
        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, x):
                return self.linear(x)

        d = self._make_dumper(tmp_path)
        d.configure(enable=False)
        model = self._wrap_as_outer(Inner)
        d.register_non_intrusive_dumper(model)

        x = torch.randn(2, 4)
        with d.capture_output() as captured:
            model(x)

        assert len(captured) == 0


def _make_forward_batch():
    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode

    return ForwardBatch(
        forward_mode=ForwardMode.DECODE,
        batch_size=2,
        input_ids=torch.tensor([10, 20]),
        req_pool_indices=torch.zeros(2, dtype=torch.long),
        seq_lens=torch.tensor([5, 6]),
        out_cache_loc=torch.zeros(2, dtype=torch.long),
        seq_lens_sum=11,
        positions=torch.tensor([0, 1]),
    )


class TestNonIntrusiveDumperConfigMode(_NonIntrusiveTestBase):
    @staticmethod
    def _build_model() -> torch.nn.Module:
        class SubLayer(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, forward_batch):
                return self.linear(
                    forward_batch.input_ids.float().unsqueeze(-1).expand(-1, 4)
                )

        class Root(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layer = SubLayer()

            def forward(self, forward_batch):
                return self.layer(forward_batch)

        return Root()

    def _run(self, tmp_path, mode: str) -> tuple:
        d = _make_test_dumper(tmp_path, non_intrusive_mode=mode)
        model = self._build_model()
        d.register_non_intrusive_dumper(model)
        forward_batch = _make_forward_batch()
        with d.capture_output() as captured:
            model(forward_batch)
        return captured, forward_batch

    def test_off_mode(self, tmp_path):
        captured, _ = self._run(tmp_path, "off")
        assert len(captured) == 0

    def test_core_mode(self, tmp_path):
        captured, fb = self._run(tmp_path, "core")

        # core fields dumped with clean names
        assert "input_ids" in captured
        assert "positions" in captured
        assert "seq_lens" in captured
        assert torch.equal(captured["input_ids"]["value"], fb.input_ids)
        assert torch.equal(captured["positions"]["value"], fb.positions)
        assert torch.equal(captured["seq_lens"]["value"], fb.seq_lens)

        # nothing with non_intrusive__ prefix
        assert not any(k.startswith("non_intrusive__") for k in captured)

    def test_all_mode(self, tmp_path):
        captured, fb = self._run(tmp_path, "all")

        # core fields dumped with clean names
        assert "input_ids" in captured
        assert "positions" in captured
        assert "seq_lens" in captured
        assert torch.equal(captured["input_ids"]["value"], fb.input_ids)
        assert torch.equal(captured["positions"]["value"], fb.positions)
        assert torch.equal(captured["seq_lens"]["value"], fb.seq_lens)

        # core fields NOT duplicated with prefix
        for field in ("input_ids", "positions", "seq_lens"):
            assert not any(
                k.startswith("non_intrusive__") and k.endswith(field) for k in captured
            )

        # ForwardBatch skipped on sub-modules (no duplication)
        assert not any(
            k.startswith("non_intrusive__layer.inputs.") and "seq_lens" in k
            for k in captured
        ), f"ForwardBatch skipped on sub-module, got: {list(captured.keys())}"

        # regular tensor outputs on sub-modules still dumped
        assert "non_intrusive__layer.linear.output" in captured
        assert "non_intrusive__layer.output" in captured


class _LayerWithNumber(torch.nn.Module):
    """Test helper: module with a ``layer_number`` attribute (Megatron style)."""

    def __init__(self, layer_number: int):
        super().__init__()
        self.layer_number = layer_number
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x):
        return self.linear(x)


class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase):
    """Tests for automatic layer_id context injection via set_ctx."""

    def test_layer_id_from_layer_number(self, tmp_path):
        """Megatron PP: layer_number (1-based global) -> layer_id = layer_number - 1."""

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = torch.nn.ModuleList(
                    [_LayerWithNumber(10), _LayerWithNumber(11)]
                )

            def forward(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        layer0_key = "non_intrusive__model.layers.0.linear.output"
        layer1_key = "non_intrusive__model.layers.1.linear.output"
        assert layer0_key in captured
        assert layer1_key in captured
        assert captured[layer0_key]["meta"]["layer_id"] == 9
        assert captured[layer1_key]["meta"]["layer_id"] == 10

        root_key = "non_intrusive__output"
        assert root_key in captured
        assert "layer_id" not in captured[root_key]["meta"]

    def test_layer_id_from_layer_id_attr(self, tmp_path):
        """SGLang style: module has layer_id attribute directly."""

        class Layer(torch.nn.Module):
            def __init__(self, layer_id: int):
                super().__init__()
                self.layer_id = layer_id
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, x):
                return self.linear(x)

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = torch.nn.ModuleList([Layer(5)])

            def forward(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        layer_key = "non_intrusive__model.layers.0.linear.output"
        assert layer_key in captured
        assert captured[layer_key]["meta"]["layer_id"] == 5

    def test_layer_id_fallback_from_module_name(self, tmp_path):
        """layers.N modules without layer_number/layer_id -> layer_id from module name."""

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = torch.nn.ModuleList(
                    [torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)]
                )

            def forward(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        captured, x, output = self._run(tmp_path, Inner)

        assert len(captured) > 0
        input_keys: list[str] = [
            k for k in captured if "model.layers." in k and "inputs" in k
        ]
        assert len(input_keys) > 0
        for key in input_keys:
            meta = captured[key]["meta"]
            assert "layer_id" in meta, f"{key} missing layer_id"
            if "layers.0" in key:
                assert meta["layer_id"] == 0
            elif "layers.1" in key:
                assert meta["layer_id"] == 1

    def test_filter_by_layer_id(self, tmp_path):
        """filter='layer_id == 0' keeps only layer 0 dumps."""

        class Inner(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.layers = torch.nn.ModuleList(
                    [_LayerWithNumber(1), _LayerWithNumber(2)]
                )

            def forward(self, x):
                for layer in self.layers:
                    x = layer(x)
                return x

        captured, x, output = self._run(tmp_path, Inner, filter="layer_id == 0")

        layer0_keys = [k for k in captured if "layers.0" in k]
        layer1_keys = [k for k in captured if "layers.1" in k]
        assert len(layer0_keys) > 0, "layer 0 dumps should be kept"
        assert len(layer1_keys) == 0, f"layer 1 dumps should be filtered: {layer1_keys}"


class TestDumperE2E:
    def test_step_and_non_intrusive_hooks(self, tmp_path):
        base_url = DEFAULT_URL_FOR_TEST
        dump_dir = str(tmp_path)
        env = {
            **os.environ,
            "DUMPER_SERVER_PORT": "reuse",
        }
        proc = popen_launch_server(
            "Qwen/Qwen3-0.6B",
            base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            other_args=["--tp", "2", "--max-total-tokens", "128"],
            env=env,
        )
        try:
            states = requests.post(f"{base_url}/dumper/get_state", json={}).json()
            assert len(states) == 2, f"Expected 2 ranks (tp=2), got {len(states)}"
            for state in states:
                assert state["config"]["enable"] is False
                assert state["step"] == 0

            requests.post(
                f"{base_url}/dumper/configure",
                json={"enable": True, "dir": dump_dir},
            ).raise_for_status()

            states = requests.post(f"{base_url}/dumper/get_state", json={}).json()
            assert len(states) == 2
            for rank, state in enumerate(states):
                assert (
                    state["config"]["enable"] is True
                ), f"rank {rank}: enable should be True after configure"
                assert state["config"]["dir"] == dump_dir

            resp = requests.post(
                f"{base_url}/generate",
                json={"text": "Hello", "sampling_params": {"max_new_tokens": 8}},
            )
            assert resp.status_code == 200, f"Generate failed: {resp.text}"

            states = requests.post(f"{base_url}/dumper/get_state", json={}).json()
            assert len(states) == 2
            steps = [s["step"] for s in states]
            for rank, step in enumerate(steps):
                assert step > 0, f"rank {rank}: step should be > 0, got {step}"
            assert steps[0] == steps[1], f"step mismatch across ranks: {steps}"

            dump_files = list(Path(dump_dir).glob("dump_*/*.pt"))
            assert len(dump_files) > 0, f"No dump files in {dump_dir}"
            filenames = {f.name for f in dump_files}

            for field in ("input_ids", "positions", "rids"):
                assert any(f"name={field}" in f for f in filenames), (
                    f"Missing {field} dump from non-intrusive hooks, "
                    f"got: {sorted(filenames)[:10]}"
                )

            for rank in range(2):
                assert any(
                    f"rank={rank}" in f for f in filenames
                ), f"No dump files for rank {rank}"

            sample_file = dump_files[0]
            loaded = torch.load(sample_file, map_location="cpu", weights_only=False)
            assert isinstance(loaded, dict), f"Expected dict, got {type(loaded)}"
            assert (
                "value" in loaded and "meta" in loaded
            ), f"Missing value/meta keys: {loaded.keys()}"
            assert "name" in loaded["meta"]
            assert "rank" in loaded["meta"]
            assert "step" in loaded["meta"]

            par = loaded["meta"].get("sglang_parallel_info", {})
            expected_keys = [
                "tp_rank",
                "tp_size",
                "pp_rank",
                "pp_size",
                "moe_ep_rank",
                "moe_ep_size",
                "moe_tp_rank",
                "moe_tp_size",
                "moe_dp_rank",
                "moe_dp_size",
                "enable_dp_attention",
                "attn_tp_rank",
                "attn_tp_size",
                "attn_dp_rank",
                "attn_dp_size",
                "attn_cp_rank",
                "attn_cp_size",
            ]
            for key in expected_keys:
                assert (
                    key in par
                ), f"Missing {key} in sglang_parallel_info, got: {sorted(par)}"

            rids_files = [f for f in dump_files if "name=rids" in f.name]
            rids_loaded = torch.load(
                rids_files[0], map_location="cpu", weights_only=False
            )
            rids_value = rids_loaded["value"]
            assert isinstance(
                rids_value, list
            ), f"rids should be a list, got {type(rids_value)}"
            assert len(rids_value) > 0, "rids should be non-empty"
            assert all(
                isinstance(r, str) for r in rids_value
            ), f"each rid should be a str, got {[type(r) for r in rids_value]}"
        finally:
            kill_process_tree(proc.pid)


class TestRegisterForwardHook:
    @pytest.mark.parametrize("mode", ["hook", "replace_fn"])
    def test_handles_removable(self, mode):
        call_log: list[str] = []

        def pre_hook(_module, _args, _kwargs):
            call_log.append("pre")

        def hook(_module, _input, _output):
            call_log.append("post")

        module = torch.nn.Linear(4, 4)
        handles = _register_forward_hook_or_replace_fn(
            module,
            pre_hook=pre_hook,
            hook=hook,
            mode=mode,
        )

        x = torch.randn(2, 4)
        if mode == "hook":
            module(x)
        else:
            module.forward(x)
        assert call_log == ["pre", "post"]

        call_log.clear()
        for h in handles:
            h.remove()

        if mode == "hook":
            module(x)
        else:
            module.forward(x)
        assert call_log == []

    @pytest.mark.parametrize("mode", ["hook", "replace_fn"])
    def test_kwargs_passed_to_pre_hook(self, mode):
        received: list[tuple] = []

        class KwargsModule(torch.nn.Module):
            def forward(self, x, *, scale=1.0):
                return x * scale

        def pre_hook(_module, _args, _kwargs):
            received.append((_args, _kwargs))

        def hook(_module, _input, _output):
            pass

        module = KwargsModule()
        _register_forward_hook_or_replace_fn(
            module,
            pre_hook=pre_hook,
            hook=hook,
            mode=mode,
        )

        x = torch.randn(2, 4)
        if mode == "hook":
            module(x, scale=2.0)
        else:
            module.forward(x, scale=2.0)

        assert len(received) == 1
        args, kwargs = received[0]
        assert len(args) == 1
        assert torch.equal(args[0], x)
        assert kwargs == {"scale": 2.0}

    def test_replace_fn_remove_asserts_on_rewrap(self):
        module = torch.nn.Linear(4, 4)
        handles = _register_forward_hook_or_replace_fn(
            module,
            pre_hook=lambda _m, _a, _kw: None,
            hook=lambda _m, _i, _o: None,
            mode="replace_fn",
        )

        module.forward = lambda *a, **kw: None

        with pytest.raises(AssertionError):
            handles[0].remove()


class TestPluginCoreFields:
    def test_sglang_core_fields(self):
        plugin = _SGLangPlugin()
        assert plugin.core_fields() == frozenset(
            {"input_ids", "positions", "seq_lens", "req_pool_indices", "rids"}
        )

    def test_megatron_core_fields(self):
        plugin = _MegatronPlugin()
        assert plugin.core_fields() == frozenset(
            {"input_ids", "position_ids", "cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}
        )


class TestMegatronConvertValue:
    @pytest.fixture(autouse=True)
    def _patch_megatron(self, monkeypatch):
        class FakePackedSeqParams:
            def __init__(self, **kwargs):
                for k, v in kwargs.items():
                    setattr(self, k, v)

        monkeypatch.setattr(_MegatronPlugin, "_available", True)
        monkeypatch.setattr(
            _MegatronPlugin, "PackedSeqParams", FakePackedSeqParams, raising=False
        )
        self._FakePackedSeqParams = FakePackedSeqParams

    def test_extracts_packed_seq_params(self):
        plugin = _MegatronPlugin()
        cu_q = torch.tensor([0, 3, 7])
        cu_kv = torch.tensor([0, 3, 7])
        value = self._FakePackedSeqParams(
            cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, qkv_format="thd"
        )

        result = plugin.convert_value(value, skip_forward_batch=False)
        assert set(result.keys()) == {"cu_seqlens_q", "cu_seqlens_kv", "qkv_format"}
        assert torch.equal(result["cu_seqlens_q"], cu_q)
        assert torch.equal(result["cu_seqlens_kv"], cu_kv)
        assert result["qkv_format"] == "thd"

    def test_non_packed_returns_none(self):
        plugin = _MegatronPlugin()
        assert plugin.convert_value(torch.randn(4), skip_forward_batch=False) is None
        assert plugin.convert_value("hello", skip_forward_batch=False) is None


class TestNonIntrusiveKwargsModel(_NonIntrusiveTestBase):
    def test_kwargs_core_fields(self, tmp_path):
        class KwargsModel(torch.nn.Module):
            def forward(self, *, input_ids, position_ids):
                return input_ids + position_ids

        model = KwargsModel()
        d = _make_test_dumper(tmp_path, non_intrusive_mode="core")
        d.register_non_intrusive_dumper(model)

        ids = torch.randn(4)
        pos = torch.randn(4)
        with d.capture_output() as captured:
            model(input_ids=ids, position_ids=pos)

        assert "input_ids" in captured
        assert "position_ids" in captured
        assert torch.equal(captured["input_ids"]["value"], ids)
        assert torch.equal(captured["position_ids"]["value"], pos)

    def test_kwargs_all_mode(self, tmp_path):
        class KwargsModel(torch.nn.Module):
            def forward(self, *, input_ids, position_ids, custom_value):
                return input_ids + position_ids + custom_value

        model = KwargsModel()
        d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
        d.register_non_intrusive_dumper(model)

        ids = torch.randn(4)
        pos = torch.randn(4)
        custom = torch.randn(4)
        with d.capture_output() as captured:
            model(input_ids=ids, position_ids=pos, custom_value=custom)

        assert "input_ids" in captured
        assert "position_ids" in captured

        P = self._PREFIX
        assert f"{P}inputs.custom_value" in captured

    def test_mixed_args_and_kwargs(self, tmp_path):
        class MixedModel(torch.nn.Module):
            def forward(self, x, *, input_ids):
                return x + input_ids

        model = MixedModel()
        d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
        d.register_non_intrusive_dumper(model)

        x = torch.randn(4)
        ids = torch.randn(4)
        with d.capture_output() as captured:
            model(x, input_ids=ids)

        assert "input_ids" in captured

        P = self._PREFIX
        assert f"{P}inputs.0" in captured

    def test_packed_seq_params_core_fields(self, tmp_path, monkeypatch):
        class FakePackedSeqParams:
            def __init__(self, **kwargs):
                for k, v in kwargs.items():
                    setattr(self, k, v)

        monkeypatch.setattr(_MegatronPlugin, "_available", True)
        monkeypatch.setattr(
            _MegatronPlugin, "PackedSeqParams", FakePackedSeqParams, raising=False
        )

        class MegatronLikeModel(torch.nn.Module):
            def forward(self, *, input_ids, packed_seq_params):
                return input_ids

        model = MegatronLikeModel()
        d = _make_test_dumper(tmp_path, non_intrusive_mode="core")
        d.register_non_intrusive_dumper(model)

        ids = torch.randn(4)
        cu_q = torch.tensor([0, 3, 7])
        cu_kv = torch.tensor([0, 3, 7])
        psp = FakePackedSeqParams(
            cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, qkv_format="thd"
        )
        with d.capture_output() as captured:
            model(input_ids=ids, packed_seq_params=psp)

        assert "input_ids" in captured
        assert torch.equal(captured["input_ids"]["value"], ids)
        assert "cu_seqlens_q" in captured
        assert torch.equal(captured["cu_seqlens_q"]["value"], cu_q)
        assert "cu_seqlens_kv" in captured
        assert torch.equal(captured["cu_seqlens_kv"]["value"], cu_kv)
        assert "qkv_format" in captured
        assert captured["qkv_format"]["value"] == "thd"


class TestDumperDims:
    def test_dims_in_meta_not_filename(self, tmp_path) -> None:
        dumper = _make_test_dumper(tmp_path)
        tensor = torch.randn(4, 8)
        dumper.dump("hidden", tensor, dims="b h(tp)")
        dumper.step()

        exp_dir = tmp_path / dumper._config.exp_name
        pt_files = list(exp_dir.glob("*.pt"))
        assert len(pt_files) == 1

        assert "dims" not in pt_files[0].stem

        data = torch.load(pt_files[0], weights_only=False)
        assert "dims" in data["meta"]
        assert data["meta"]["dims"] == "b h(tp)"

    def test_dims_grad_override(self, tmp_path) -> None:
        dumper = _Dumper(
            config=DumperConfig(
                enable=True,
                dir=str(tmp_path),
                enable_grad=True,
            )
        )

        tensor = torch.randn(4, 8, requires_grad=True)
        dumper.dump("hidden", tensor, dims="b h(tp)", dims_grad="b h(tp:partial)")
        dumper.step()

        tensor.backward(torch.ones_like(tensor))

        exp_dir = tmp_path / dumper._config.exp_name
        pt_files = sorted(exp_dir.glob("*.pt"))
        assert len(pt_files) == 2

        value_file = [f for f in pt_files if "grad__" not in f.stem][0]
        grad_file = [f for f in pt_files if "grad__" in f.stem][0]

        value_data = torch.load(value_file, weights_only=False)
        assert value_data["meta"]["dims"] == "b h(tp)"
        assert value_data["meta"]["dims_grad"] == "b h(tp:partial)"

        grad_data = torch.load(grad_file, weights_only=False)
        assert grad_data["meta"]["dims"] == "b h(tp:partial)"

    def test_dims_grad_inherits(self, tmp_path) -> None:
        dumper = _Dumper(
            config=DumperConfig(
                enable=True,
                dir=str(tmp_path),
                enable_grad=True,
            )
        )

        tensor = torch.randn(4, 8, requires_grad=True)
        dumper.dump("hidden", tensor, dims="b h(tp)")
        dumper.step()

        tensor.backward(torch.ones_like(tensor))

        exp_dir = tmp_path / dumper._config.exp_name
        grad_file = [f for f in exp_dir.glob("*.pt") if "grad__" in f.stem][0]
        grad_data = torch.load(grad_file, weights_only=False)
        assert grad_data["meta"]["dims"] == "b h(tp)"


class TestCtxDecorator:
    def test_ctx_dynamic_lambda(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)

        class FakeLayer:
            def __init__(self, layer_id: int) -> None:
                self.layer_id = layer_id

            @d.ctx(lambda self: dict(layer_id=self.layer_id))
            def forward(self, x: torch.Tensor) -> torch.Tensor:
                d.dump("hidden", x)
                return x

        layer = FakeLayer(layer_id=42)
        layer.forward(torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["layer_id=42"])

    def test_ctx_static_kwargs(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)

        @d.ctx(phase="decode")
        def decode_step(x: torch.Tensor) -> torch.Tensor:
            d.dump("step_out", x)
            return x

        decode_step(torch.randn(3))

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["phase=decode"])

    def test_ctx_clears_on_exception(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)

        @d.ctx(phase="train")
        def buggy_fn() -> None:
            raise RuntimeError("boom")

        with pytest.raises(RuntimeError, match="boom"):
            buggy_fn()

        assert d._state.global_ctx == {}

    def test_ctx_rejects_mixed_args(self) -> None:
        d = _make_test_dumper("/tmp")

        with pytest.raises(ValueError, match="cannot mix"):
            d.ctx(lambda self: dict(a=1), phase="x")

    def test_ctx_rejects_empty_args(self) -> None:
        d = _make_test_dumper("/tmp")

        with pytest.raises(ValueError, match="must provide"):
            d.ctx()


class TestRecomputeStatus:
    def test_disabled_by_default(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(3, 3)
        d.dump("test_tensor", tensor)

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["recompute_status=disabled"])

    def test_recompute_status_in_embedded_meta(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(3, 3)
        d.dump("test_tensor", tensor)

        path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
        raw = _load_dump(path)
        assert raw["meta"]["recompute_status"] == "disabled"

    def test_recompute_status_recompute(self, tmp_path: Path, monkeypatch) -> None:
        import sglang.srt.debug_utils.dumper as dumper_mod

        monkeypatch.setattr(
            dumper_mod, "_detect_recompute_status", lambda: _RecomputeStatus.RECOMPUTE
        )

        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(3, 3)
        d.dump("test_tensor", tensor)

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["recompute_status=recompute"])

        path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
        raw = _load_dump(path)
        assert raw["meta"]["recompute_status"] == "recompute"
        assert raw["meta"]["recompute_pseudo_rank"] == 1
        assert raw["meta"]["recompute_pseudo_size"] == 2

    def test_recompute_status_original(self, tmp_path: Path, monkeypatch) -> None:
        import sglang.srt.debug_utils.dumper as dumper_mod

        monkeypatch.setattr(
            dumper_mod,
            "_detect_recompute_status",
            lambda: _RecomputeStatus.ORIGINAL,
        )

        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(3, 3)
        d.dump("test_tensor", tensor)

        filenames = _get_filenames(tmp_path)
        _assert_files(filenames, exist=["recompute_status=original"])

        path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
        raw = _load_dump(path)
        assert raw["meta"]["recompute_status"] == "original"
        assert raw["meta"]["recompute_pseudo_rank"] == 0
        assert raw["meta"]["recompute_pseudo_size"] == 2

    def test_disabled_no_recompute_pseudo_fields(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path)
        tensor = torch.randn(3, 3)
        d.dump("test_tensor", tensor)

        path = _find_dump_file(tmp_path, rank=0, name="test_tensor")
        raw = _load_dump(path)
        assert "recompute_pseudo_rank" not in raw["meta"]
        assert "recompute_pseudo_size" not in raw["meta"]

    def test_grad_hook_has_no_recompute_status(self, tmp_path: Path) -> None:
        d = _make_test_dumper(tmp_path, enable_grad=True)
        x = torch.randn(3, 3, requires_grad=True)
        y = (x * 2).sum()

        d.dump("test_tensor", x)
        y.backward()

        grad_files = [f for f in _get_filenames(tmp_path) if "grad__test_tensor" in f]
        assert len(grad_files) == 1
        assert "recompute_status" not in grad_files[0]

    def test_non_intrusive_hooks_have_recompute_status(self, tmp_path: Path) -> None:
        class Simple(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(4, 4)

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return self.linear(x)

        model = Simple()
        d = _make_test_dumper(tmp_path, non_intrusive_mode="all")
        d.register_non_intrusive_dumper(model)

        with d.capture_output() as captured:
            model(torch.randn(2, 4))

        for key, data in captured.items():
            assert (
                "recompute_status" in data["meta"]
            ), f"missing recompute_status in {key}"
            assert data["meta"]["recompute_status"] == "disabled"

    def test_detect_recompute_status_default(self) -> None:
        assert _detect_recompute_status() == _RecomputeStatus.DISABLED


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="29999",
            DUMPER_GRAFTER_BASELINE_WORLD_SIZE="1",
            DUMPER_GRAFTER_TARGET_WORLD_SIZE="1",
            DUMPER_GRAFTER_B2T_FILTER="name == 'x'",
        ):
            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):
        with pytest.raises(AssertionError, match=r"grafter_role"):
            DumperConfig(grafter_enable=True)
        with pytest.raises(AssertionError, match=r"grafter_master_address"):
            DumperConfig(grafter_enable=True, grafter_role="baseline")
        with pytest.raises(AssertionError, match=r"grafter_master_port"):
            DumperConfig(
                grafter_enable=True,
                grafter_role="baseline",
                grafter_master_address="127.0.0.1",
            )
        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=12345,
            )
        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=12345,
                grafter_baseline_world_size=1,
                grafter_target_world_size=1,
            )

    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-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,
        grafter_role="baseline",
        grafter_master_address="127.0.0.1",
        grafter_master_port=12345,
        grafter_baseline_world_size=1,
        grafter_target_world_size=1,
        grafter_b2t_filter="name == 'x'",
    )
    base.update(overrides)
    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.

    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_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_matched_non_tensor_prints_and_skips(self):
        """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"})
        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

    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_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_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)."""
        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 in py3.
        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).

    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

    role_ports = [find_available_port(29700 + i * 100) for i in range(2)]

    ctx = mp.get_context("spawn")
    result_queue = ctx.Queue()
    processes = []
    for rank in range(2):
        p = ctx.Process(
            target=_graft_worker_entry,
            args=(rank, role_ports[rank], worker_func, result_queue, kwargs),
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    errors = [result_queue.get() for _ in range(2)]
    errors = [e for e in errors if e]
    if errors:
        raise AssertionError("\n".join(errors))


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",
        init_method=f"tcp://127.0.0.1:{role_port}",
        world_size=1,
        rank=0,
    )
    try:
        worker_func(rank=rank, **kwargs)
        result_queue.put(None)
    except Exception as e:
        result_queue.put(f"rank={rank}: {e}\n{traceback.format_exc()}")
    finally:
        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,
) -> DumperConfig:
    """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(
        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,
        grafter_target_world_size=1,
        grafter_group_name=group_name,
        grafter_timeout=timeout,
        # 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,
    )


class TestGrafterDistributed:
    def test_b2t_copy_roundtrip(self):
        """Baseline (rank 0) sends 'x' to target (rank 1), target.copy_'s it."""
        graft_port = find_available_port(29600)
        _run_graft_test(
            self._test_b2t_func, graft_port=graft_port, group_name="grafter_b2t"
        )

    @staticmethod
    def _test_b2t_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.zeros(3, device="cuda:1")
                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)

    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_recv_with_user_transform(self, tmp_path: Path):
        # 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"
            "    return graft_input.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_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."""
        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:
                # 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,
                    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(graft_input):\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
                # 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_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."""
        module_name = "_xform_uses_extras"
        (tmp_path / f"{module_name}.py").write_text(
            "import torch\n"
            "def transform(graft_input):\n"
            "    fill = graft_input.received_extras_list[0]['fill_value']\n"
            "    return torch.full_like(graft_input.target, fill)\n"
        )
        graft_port = find_available_port(29645)
        _run_graft_test(
            self._test_extras_func,
            graft_port=graft_port,
            group_name="grafter_extras",
            transform_dir=str(tmp_path),
            transform_path=f"{module_name}.transform",
        )

    @staticmethod
    def _test_extras_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:
                # Baseline (sender) attaches an extras dict.
                tensor = torch.tensor([1.0, 2.0, 3.0], device="cuda:0")
                grafter.maybe_intercept(
                    value=tensor,
                    tags={"name": "x"},
                    extras={"fill_value": 42.0},
                )
            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,
                ], 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 — but len(received_extras_list) still matches n_senders."""
        graft_port = find_available_port(29650)
        _run_graft_test(
            self._test_extras_none_func,
            graft_port=graft_port,
            group_name="grafter_extras_none",
        )

    @staticmethod
    def _test_extras_none_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")
                # 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()
        finally:
            if grafter._pg is not None:
                dist.destroy_process_group(grafter._pg)

    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",
        )

    @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 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",
        )

    @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:
    """Coverage of asymmetric multi-rank cases via CPU/gloo (CI fleet has
    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.
        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"
            "def transform(graft_input):\n"
            "    rl = graft_input.received_list\n"
            "    assert len(rl) == 4, f'expected 4 senders, got {len(rl)}'\n"
            "    for i, t in enumerate(rl):\n"
            "        v = float(t.flatten()[0].item())\n"
            "        assert v == float(i), f'rl[{i}][0]={v}, want {float(i)}'\n"
            "    return torch.full_like(graft_input.target, 999.0)\n"
        )
        graft_port = find_available_port(29655)
        _run_graft_test_cpu_multi(
            self._test_4b_2t_func,
            baseline_world=4,
            target_world=2,
            graft_port=graft_port,
            group_name="grafter_4b_2t",
            transform_dir=str(tmp_path),
            transform_path=f"{module_name}.transform",
        )

    @staticmethod
    def _test_4b_2t_func(
        role, local_rank, graft_port, group_name, transform_dir, transform_path
    ):
        sys.path.insert(0, transform_dir)
        cfg = _make_multi_rank_config(
            role=role,
            graft_port=graft_port,
            group_name=group_name,
            baseline_world=4,
            target_world=2,
            transform_path=transform_path,
            direction="b2t",
        )
        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()
        finally:
            if grafter._pg is not None:
                dist.destroy_process_group(grafter._pg)

    def test_2_target_4_baseline_t2b_with_user_transform(self, tmp_path: Path):
        """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"
            "def transform(graft_input):\n"
            "    rl = graft_input.received_list\n"
            "    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), (\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)
        _run_graft_test_cpu_multi(
            self._test_2t_4b_func,
            baseline_world=4,
            target_world=2,
            graft_port=graft_port,
            group_name="grafter_2t_4b",
            transform_dir=str(tmp_path),
            transform_path=f"{module_name}.transform",
        )

    @staticmethod
    def _test_2t_4b_func(
        role, local_rank, graft_port, group_name, transform_dir, transform_path
    ):
        sys.path.insert(0, transform_dir)
        cfg = _make_multi_rank_config(
            role=role,
            graft_port=graft_port,
            group_name=group_name,
            baseline_world=4,
            target_world=2,
            transform_path=transform_path,
            direction="t2b",
        )
        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:
                target = torch.full((3,), 99.0)
                grafter.maybe_intercept(value=target, tags={"name": "x"})
                assert target.tolist() == [7.0, 7.0, 7.0], target.tolist()
        finally:
            if grafter._pg is not None:
                dist.destroy_process_group(grafter._pg)

    def test_default_transform_with_asymmetric_world_logs_and_skips(self):
        """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,
            baseline_world=4,
            target_world=2,
            graft_port=graft_port,
            group_name="grafter_default_asym",
        )

    @staticmethod
    def _test_default_asym_func(role, local_rank, graft_port, group_name):
        cfg = _make_multi_rank_config(
            role=role,
            graft_port=graft_port,
            group_name=group_name,
            baseline_world=4,
            target_world=2,
            transform_path=None,  # default identity-by-rank fallback
            direction="b2t",
        )
        grafter = _Grafter(config=cfg)
        try:
            if role == "baseline":
                tensor = torch.full((3,), float(local_rank))
                grafter.maybe_intercept(value=tensor, tags={"name": "x"})
            else:
                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], (
                    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:
            if grafter._pg is not None:
                dist.destroy_process_group(grafter._pg)

    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 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"
            "    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(
            self._test_mixed_shape_func,
            baseline_world=4,
            target_world=2,
            graft_port=graft_port,
            group_name="grafter_mixed_shape",
            transform_dir=str(tmp_path),
            transform_path=f"{module_name}.transform",
        )

    @staticmethod
    def _test_mixed_shape_func(
        role, local_rank, graft_port, group_name, transform_dir, transform_path
    ):
        sys.path.insert(0, transform_dir)
        cfg = _make_multi_rank_config(
            role=role,
            graft_port=graft_port,
            group_name=group_name,
            baseline_world=4,
            target_world=2,
            transform_path=transform_path,
            direction="b2t",
        )
        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
                assert target.tolist() == expected, target.tolist()
        finally:
            if grafter._pg is not None:
                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=<class 'torch\.Tensor'> 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__]))