diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 1ebf1f037..9475ff694 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -138,6 +138,7 @@ class DumperConfig(_BaseConfig): collective_timeout: int = 60 server_port: str = "-1" non_intrusive_mode: str = "core" + source_patcher_config: Optional[str] = None @classmethod def _env_prefix(cls) -> str: @@ -325,6 +326,25 @@ class _Dumper: return decorator + def apply_source_patches(self) -> None: + """Apply source patches from DUMPER_SOURCE_PATCHER_CONFIG if set. + + Automatically injects ``from sglang.srt.debug_utils.dumper import dumper`` + into every replacement block so users don't need to write it in YAML. + """ + config_path = self._config.source_patcher_config + if not config_path: + return + + from sglang.srt.debug_utils.source_patcher import apply_patches_from_config + + yaml_content: str = Path(config_path).read_text() + print(f"[source_patcher] loading config from {config_path}") + apply_patches_from_config( + yaml_content, + extra_imports=["from sglang.srt.debug_utils.dumper import dumper"], + ) + def register_non_intrusive_dumper( self, model: "torch.nn.Module", @@ -588,11 +608,13 @@ class _NonIntrusiveDumper: def _detect_module_ctx( cls, module_name: str, module: "torch.nn.Module" ) -> Optional[dict]: - if cls._LAYER_NAME_RE.fullmatch(module_name): + match = cls._LAYER_NAME_RE.fullmatch(module_name) + if match: for plugin in _plugins: layer_id = plugin.detect_layer_id(module) if layer_id is not None: return {"layer_id": layer_id} + return {"layer_id": int(match.group(1))} return None def _register_ctx_hooks(self, module: "torch.nn.Module", *, ctx: dict) -> None: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 98045d8bc..cee3dd2ea 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1061,6 +1061,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) if dumper.may_enable: + dumper.apply_source_patches() dumper.register_non_intrusive_dumper(self.model) # Pre-expand RoPE cache before CUDA Graph capture @@ -1763,9 +1764,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): init_new_workspace=init_new_workspace, ) - self.prefill_attention_backend_str, self.decode_attention_backend_str = ( - self.server_args.get_attention_backends() - ) + ( + self.prefill_attention_backend_str, + self.decode_attention_backend_str, + ) = self.server_args.get_attention_backends() if self.decode_attention_backend_str != self.prefill_attention_backend_str: from sglang.srt.layers.attention.hybrid_attn_backend import ( diff --git a/test/registered/debug_utils/source_patcher/test_dumper_integration.py b/test/registered/debug_utils/source_patcher/test_dumper_integration.py new file mode 100644 index 000000000..70a06eb45 --- /dev/null +++ b/test/registered/debug_utils/source_patcher/test_dumper_integration.py @@ -0,0 +1,54 @@ +"""Test dumper.apply_source_patches() integration with source_patcher.""" + +from pathlib import Path +from types import ModuleType + +import yaml + +from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper + +SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module" + + +class TestDumperApplySourcePatches: + def test_no_config_is_noop(self) -> None: + config = DumperConfig(source_patcher_config=None) + d = _Dumper(config=config) + d.apply_source_patches() + + def test_patches_applied_from_yaml( + self, sample_module: ModuleType, tmp_path: Path + ) -> None: + cls = sample_module.SampleClass + obj = cls() + assert obj.greet("world") == "hello world" + + original_code = cls.greet.__code__ + + patch_config = { + "patches": [ + { + "target": f"{SAMPLE_MODULE_NAME}.SampleClass.greet", + "edits": [ + { + "match": 'greeting = f"hello {name}"', + "replacement": 'greeting = f"dumper_patched {name}"', + } + ], + } + ] + } + + config_path = tmp_path / "patch_config.yaml" + config_path.write_text(yaml.dump(patch_config)) + + config = DumperConfig(source_patcher_config=str(config_path)) + d = _Dumper(config=config) + + try: + d.apply_source_patches() + assert obj.greet("world") == "dumper_patched world" + finally: + cls.greet.__code__ = original_code + + assert obj.greet("world") == "hello world" diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index ac5470a5a..42177bf97 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -1755,7 +1755,6 @@ def _make_forward_batch(): class TestNonIntrusiveDumperConfigMode(_NonIntrusiveTestBase): - @staticmethod def _build_model() -> torch.nn.Module: class SubLayer(torch.nn.Module): @@ -1904,8 +1903,8 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): assert layer_key in captured assert captured[layer_key]["meta"]["layer_id"] == 5 - def test_no_layer_id_when_no_attr(self, tmp_path): - """layers.N modules without layer_number/layer_id -> no layer_id injected.""" + 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): @@ -1922,8 +1921,17 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): captured, x, output = self._run(tmp_path, Inner) assert len(captured) > 0 - for key, entry in captured.items(): - assert "layer_id" not in entry["meta"], f"{key} has unexpected layer_id" + 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.""" @@ -2171,7 +2179,6 @@ class TestMegatronConvertValue: class TestNonIntrusiveKwargsModel(_NonIntrusiveTestBase): - def test_kwargs_core_fields(self, tmp_path): class KwargsModel(torch.nn.Module): def forward(self, *, input_ids, position_ids): diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py new file mode 100644 index 000000000..59477f910 --- /dev/null +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -0,0 +1,216 @@ +"""E2E test: source patcher + dumper + comparator on SGLang server. + +Patches Qwen3DecoderLayer.forward to insert dumper.dump() calls, +launches 1-GPU baseline and 2-GPU TP=2 target servers, runs inference, +verifies patched dump fields exist, then runs comparator to verify +numerical consistency. + +The dumper.apply_source_patches() auto-injects ``from ... import dumper`` +so the YAML only needs ``dumper.dump(...)`` calls. +""" + +import os +import subprocess +import tempfile +from pathlib import Path + +import pytest +import requests + +from sglang.srt.debug_utils.comparator.output_types import ( + AnyRecord, + SummaryRecord, + parse_record_json, +) +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + popen_launch_server, +) + +register_cuda_ci(est_time=120, suite="nightly-2-gpu", nightly=True) + +MODEL = "Qwen/Qwen3-0.6B" +EXP_NAME = "e2e_source_patcher" +DUMPER_FILTER = r"layer_id=[012]" + +PATCH_CONFIG_YAML: str = """\ +patches: + - target: sglang.srt.models.qwen3.Qwen3DecoderLayer.forward + edits: + - match: "hidden_states = self.mlp(hidden_states)" + prepend: "dumper.dump('patched_attn_output', hidden_states, dims='t h')" + - match: "return hidden_states, residual" + prepend: "dumper.dump('patched_mlp_output', hidden_states, dims='t h')" +""" + + +class TestSourcePatcherE2ESGLang: + """E2E: patch Qwen3 forward -> dump -> compare 1gpu vs 2gpu-tp2.""" + + @pytest.mark.timeout(300) + def test_patch_dump_and_compare(self, tmp_path: Path) -> None: + patched_fields: list[str] = ["patched_attn_output", "patched_mlp_output"] + base_url: str = DEFAULT_URL_FOR_TEST + + config_path: Path = tmp_path / "patch_config.yaml" + config_path.write_text(PATCH_CONFIG_YAML) + + # Run 1: baseline (1 GPU) + baseline_dir: Path = tmp_path / "baseline" + _run_server_and_generate( + dump_dir=baseline_dir, + config_path=config_path, + tp=1, + base_url=base_url, + ) + _verify_patched_fields(dump_dir=baseline_dir, field_names=patched_fields) + + # Run 2: target (2 GPU TP=2) + target_dir: Path = tmp_path / "target" + _run_server_and_generate( + dump_dir=target_dir, + config_path=config_path, + tp=2, + base_url=base_url, + ) + _verify_patched_fields(dump_dir=target_dir, field_names=patched_fields) + + # Compare baseline vs target + baseline_exp: Path = baseline_dir / EXP_NAME + target_exp: Path = target_dir / EXP_NAME + + result: subprocess.CompletedProcess[str] = subprocess.run( + [ + "python", + "-m", + "sglang.srt.debug_utils.comparator", + "--baseline-path", + str(baseline_exp), + "--target-path", + str(target_exp), + "--output-format", + "json", + "--grouping", + "logical", + ], + capture_output=True, + text=True, + ) + + debug_file: Path = _save_comparator_output( + stdout=result.stdout, stderr=result.stderr + ) + print(f"Comparator debug output: {debug_file}") + + assert result.returncode == 0, ( + f"Comparator failed (rc={result.returncode}). " + f"Debug output: {debug_file}" + ) + + records: list[AnyRecord] = [ + parse_record_json(line) + for line in result.stdout.strip().splitlines() + if line.strip() + ] + assert ( + len(records) > 0 + ), f"Comparator produced no output records. Debug: {debug_file}" + + summary: SummaryRecord = _find_summary(records=records, debug_file=debug_file) + assert ( + summary.passed > 0 + ), f"No comparisons passed (total={summary.total}). Debug: {debug_file}" + assert summary.failed == 0, ( + f"{summary.failed} comparisons failed " + f"(passed={summary.passed}, skipped={summary.skipped}). " + f"Debug: {debug_file}" + ) + + +# --------------------------------- helpers --------------------------------- + + +def _run_server_and_generate( + *, + dump_dir: Path, + config_path: Path, + tp: int, + base_url: str, +) -> None: + """Launch SGLang server with source patcher + dumper, send a generate request.""" + env: dict[str, str] = { + **os.environ, + "DUMPER_SOURCE_PATCHER_CONFIG": str(config_path), + "DUMPER_DIR": str(dump_dir), + "DUMPER_EXP_NAME": EXP_NAME, + "DUMPER_SERVER_PORT": "reuse", + } + + proc = popen_launch_server( + MODEL, + base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=["--tp", str(tp), "--max-total-tokens", "128"], + env=env, + ) + try: + requests.post( + f"{base_url}/dumper/configure", + json={ + "enable": True, + "filter": DUMPER_FILTER, + "cleanup_previous": True, + }, + ).raise_for_status() + + resp = requests.post( + f"{base_url}/generate", + json={ + "text": "The capital of France is", + "sampling_params": {"max_new_tokens": 8}, + }, + ) + assert resp.status_code == 200, f"Generate failed: {resp.text}" + finally: + kill_process_tree(proc.pid) + + +def _verify_patched_fields(*, dump_dir: Path, field_names: list[str]) -> None: + """Verify that patched dump fields exist as .pt files.""" + for field in field_names: + matches: list[Path] = list(dump_dir.rglob(f"*name={field}*.pt")) + assert len(matches) > 0, ( + f"Expected patched field '{field}' not found under {dump_dir}. " + f"Available files: {sorted(f.name for f in dump_dir.rglob('*.pt'))[:20]}" + ) + + +def _find_summary(*, records: list[AnyRecord], debug_file: Path) -> SummaryRecord: + """Extract the SummaryRecord from comparator output.""" + summaries: list[SummaryRecord] = [ + r for r in records if isinstance(r, SummaryRecord) + ] + assert len(summaries) == 1, ( + f"Expected 1 summary record, got {len(summaries)}. " + f"Record types: {[type(r).__name__ for r in records]}. " + f"Debug: {debug_file}" + ) + return summaries[0] + + +def _save_comparator_output(*, stdout: str, stderr: str) -> Path: + """Save comparator stdout+stderr to a temp file that persists for debugging.""" + fd, path_str = tempfile.mkstemp(prefix="comparator_e2e_", suffix=".log", dir="/tmp") + with os.fdopen(fd, "w") as f: + f.write("=== STDOUT ===\n") + f.write(stdout) + f.write("\n=== STDERR ===\n") + f.write(stderr) + return Path(path_str) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])