[Unblock internal] Fix log capture issue on internal tests (#1076)

Author: yf225

helion/_testing.py

@@ -5,6 +5,8 @@
 import functools
 import importlib
 import inspect
+import io
+import logging
 import operator
 import os
 from pathlib import Path
@@ -39,6 +41,39 @@
 EXAMPLES_DIR: Path = PROJECT_ROOT / "examples"
 
 
+class _LogCapture(logging.Handler):
+    """Simple logging handler to capture log records."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.records: list[logging.LogRecord] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        self.records.append(record)
+
+    def clear(self) -> None:
+        self.records.clear()
+
+
+class _OutputCapture:
+    """Simple output capture class for stdout/stderr."""
+
+    def __init__(self) -> None:
+        self.stdout = io.StringIO()
+        self.stderr = io.StringIO()
+
+    def readouterr(self) -> tuple[str, str]:
+        """Read and clear captured output, returning (stdout, stderr) tuple."""
+        stdout_val = self.stdout.getvalue()
+        stderr_val = self.stderr.getvalue()
+        # Clear the buffers
+        self.stdout.seek(0)
+        self.stdout.truncate()
+        self.stderr.seek(0)
+        self.stderr.truncate()
+        return (stdout_val, stderr_val)
+
+
 def is_cuda() -> bool:
     """Return True if running on CUDA (NVIDIA GPU)."""
     return (
@@ -940,3 +975,26 @@ def assertExpectedJournal(self, value: str) -> None:
             expected,
             msg="To accept the new output, re-run test with env EXPECTTEST_ACCEPT=1",
         )
+
+    @contextlib.contextmanager
+    def capture_logs(self) -> Generator[_LogCapture, None, None]:
+        """Context manager to capture logs."""
+        handler = _LogCapture()
+        handler.setLevel(logging.DEBUG)
+        logger = logging.getLogger()
+        logger.addHandler(handler)
+        try:
+            yield handler
+        finally:
+            logger.removeHandler(handler)
+
+    @contextlib.contextmanager
+    def capture_output(self) -> Generator[_OutputCapture, None, None]:
+        """Context manager to capture stdout/stderr."""
+        capture = _OutputCapture()
+        old_stdout, old_stderr = sys.stdout, sys.stderr
+        sys.stdout, sys.stderr = capture.stdout, capture.stderr
+        try:
+            yield capture
+        finally:
+            sys.stdout, sys.stderr = old_stdout, old_stderr

test/test_debug_utils.py

@@ -50,13 +50,6 @@ def _with_print_repro_enabled(self):
             else:
                 os.environ["HELION_PRINT_REPRO"] = original
 
-    def _clear_captures(self):
-        """Clear pytest capture fixtures if available."""
-        if hasattr(self, "_capfd"):
-            self._capfd.readouterr()
-        if hasattr(self, "_caplog"):
-            self._caplog.clear()
-
     def _create_kernel(self, **kwargs):
         """Create a simple 1D kernel for testing.
 
@@ -85,21 +78,19 @@ def test_print_repro_env_var(self):
             torch.manual_seed(0)
             x = torch.randn([128], dtype=torch.float32, device=DEVICE)
 
-            self._clear_captures()
-
-            result = kernel(x)
-            torch.testing.assert_close(result, x + 1)
+            with self.capture_logs() as log_capture:
+                result = kernel(x)
+                torch.testing.assert_close(result, x + 1)
 
-            # Extract repro script from logs (use records to get the raw message without formatting)
-            assert hasattr(self, "_caplog"), "caplog fixture not available"
-            repro_script = None
-            for record in self._caplog.records:
-                if "# === HELION KERNEL REPRO ===" in record.message:
-                    repro_script = record.message
-                    break
+                # Extract repro script from logs (use records to get the raw message without formatting)
+                repro_script = None
+                for record in log_capture.records:
+                    if "# === HELION KERNEL REPRO ===" in record.message:
+                        repro_script = record.message
+                        break
 
-            if repro_script is None:
-                self.fail("No repro script found in logs")
+                if repro_script is None:
+                    self.fail("No repro script found in logs")
 
             # Normalize range_warp_specializes=[None] to [] for comparison
             normalized_script = repro_script.replace(
@@ -149,8 +140,6 @@ def test_print_repro_on_autotune_error(self):
             torch.manual_seed(0)
             x = torch.randn([128], dtype=torch.float32, device=DEVICE)
 
-            self._clear_captures()
-
             # Mock do_bench to fail on the second config with PTXASError (warn level)
             from torch._inductor.runtime.triton_compat import PTXASError
             from triton.testing import do_bench as original_do_bench
@@ -163,13 +152,13 @@ def mock_do_bench(*args, **kwargs):
                     raise PTXASError("Mocked PTXAS error")
                 return original_do_bench(*args, **kwargs)
 
-            with mock.patch("helion.autotuner.base_search.do_bench", mock_do_bench):
-                # Autotune will try both configs, second one will fail and print repro
-                kernel.autotune([x], force=False)
+            with self.capture_output() as output_capture:
+                with mock.patch("helion.autotuner.base_search.do_bench", mock_do_bench):
+                    # Autotune will try both configs, second one will fail and print repro
+                    kernel.autotune([x], force=False)
 
-            # Extract repro script from stderr
-            assert hasattr(self, "_capfd"), "capfd fixture not available"
-            captured = "".join(self._capfd.readouterr())
+                # Extract repro script from stderr
+                captured = "".join(output_capture.readouterr())
 
             # Verify that a repro script was printed for the failing config
             self.assertIn("# === HELION KERNEL REPRO ===", captured)