Enhance hook mechanism in dumper (#19073)

test/registered/debug_utils/test_dump_comparator.py

@@ -108,7 +108,7 @@ class TestEndToEnd(CustomTestCase):
                 dumper.step()
                 dumper.dump("tensor_b", tensor * 2)
                 dumper.step()
-                dump_dirs.append(Path(d) / f"sglang_dump_{dumper._config.partial_name}")
+                dump_dirs.append(Path(d) / dumper._config.exp_name)
 
             args = Namespace(
                 baseline_path=str(dump_dirs[0]),

test/registered/debug_utils/test_dumper.py

@@ -440,7 +440,7 @@ def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
     config = _DumperConfig(
         enable=True,
         dir=str(tmp_path),
-        partial_name="test",
+        exp_name="test",
         enable_http_server=False,
         **overrides,
     )
@@ -448,7 +448,7 @@ def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
 
 
 def _get_filenames(tmpdir):
-    return {f.name for f in Path(tmpdir).glob("sglang_dump_*/*.pt")}
+    return {f.name for f in Path(tmpdir).glob("*/*.pt")}
 
 
 def _assert_files(filenames, *, exist=(), not_exist=()):
@@ -468,7 +468,7 @@ def _load_dump(path: Path) -> dict:
 def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path:
     matches = [
         f
-        for f in Path(tmpdir).glob("sglang_dump_*/*.pt")
+        for f in Path(tmpdir).glob("*/*.pt")
         if f"rank={rank}" in f.name and name in f.name
     ]
     assert (
@@ -770,7 +770,7 @@ class TestDumpModel:
 
 class TestCleanup:
     def test_cleanup_removes_old_dumps(self, tmp_path):
-        old_dir = tmp_path / "sglang_dump_old"
+        old_dir = tmp_path / "dump_old"
         old_dir.mkdir()
         (old_dir / "dummy.pt").touch()
 
@@ -781,7 +781,7 @@ class TestCleanup:
         _assert_files(_get_filenames(tmp_path), exist=["new_tensor"])
 
     def test_no_cleanup_by_default(self, tmp_path):
-        old_dir = tmp_path / "sglang_dump_old"
+        old_dir = tmp_path / "dump_old"
         old_dir.mkdir()
         (old_dir / "dummy.pt").touch()
 
@@ -1020,6 +1020,43 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase):
         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):