Suppress noisy Kineto benchmark logs

Author: drisspg

test/test_benchmark_stats.py

@@ -1,9 +1,10 @@
-import warnings
+import os
 
 import pytest
 
 from transformer_nuggets.utils.benchmark import (
     CudaBenchmarkStats,
+    benchmark_cuda_function_in_microseconds,
     benchmark_cuda_function_stats,
 )
 
@@ -61,24 +62,56 @@ def benchmark_gpu(self, fn, **kwargs):
     assert stats.median_ci_us == pytest.approx((10.0, 10.0))
 
 
-@pytest.mark.parametrize(
-    "message",
-    [
-        "CUDA warning: SyncActivityProfilerHandler::start failed to stop cleanly",
-        "Detected call of profiler_start while another profiler is active",
-        "Detected call of profiler_stop without a matching start",
-    ],
-)
-def test_benchmark_utils_suppresses_known_profiler_warnings(message):
-    with warnings.catch_warnings(record=True) as caught:
-        warnings.warn(message, UserWarning, stacklevel=1)
+def test_benchmark_cuda_function_sets_kineto_log_level_around_profiler_call(monkeypatch):
+    monkeypatch.delenv("KINETO_LOG_LEVEL", raising=False)
+
+    def fake_do_bench_using_profiling(fn, *, rep, is_vetted_benchmarking):
+        assert rep == 3
+        assert is_vetted_benchmarking is False
+        assert os.environ["KINETO_LOG_LEVEL"] == "6"
+        fn()
+        return 0.123
+
+    monkeypatch.setattr(
+        "transformer_nuggets.utils.benchmark.do_bench_using_profiling",
+        fake_do_bench_using_profiling,
+    )
+
+    latency_us = benchmark_cuda_function_in_microseconds(lambda: None, NUM_ITERS=3)
+
+    assert latency_us == pytest.approx(123.0)
+    assert "KINETO_LOG_LEVEL" not in os.environ
+
+
+def test_benchmark_cuda_function_restores_existing_kineto_log_level(monkeypatch):
+    monkeypatch.setenv("KINETO_LOG_LEVEL", "2")
+
+    def fake_do_bench_using_profiling(fn, *, rep, is_vetted_benchmarking):
+        assert os.environ["KINETO_LOG_LEVEL"] == "6"
+        fn()
+        return 0.123
+
+    monkeypatch.setattr(
+        "transformer_nuggets.utils.benchmark.do_bench_using_profiling",
+        fake_do_bench_using_profiling,
+    )
 
-    assert caught == []
+    benchmark_cuda_function_in_microseconds(lambda: None)
 
+    assert os.environ["KINETO_LOG_LEVEL"] == "2"
 
-def test_benchmark_utils_does_not_suppress_unrelated_warnings():
-    with warnings.catch_warnings(record=True) as caught:
-        warnings.warn("transformer-nuggets unrelated benchmark warning", UserWarning, stacklevel=1)
 
-    assert len(caught) == 1
-    assert str(caught[0].message) == "transformer-nuggets unrelated benchmark warning"
+def test_benchmark_cuda_function_stats_sets_kineto_log_level_around_profiler_call(monkeypatch):
+    class FakeBenchmarker:
+        def benchmark_gpu(self, fn, **kwargs):
+            assert os.environ["KINETO_LOG_LEVEL"] == "6"
+            fn()
+            return [0.010, 0.012, 0.011]
+
+    monkeypatch.delenv("KINETO_LOG_LEVEL", raising=False)
+    monkeypatch.setattr("torch._inductor.runtime.benchmarking.benchmarker", FakeBenchmarker())
+
+    stats = benchmark_cuda_function_stats(lambda: None)
+
+    assert stats.samples_us == pytest.approx((10.0, 12.0, 11.0))
+    assert "KINETO_LOG_LEVEL" not in os.environ

transformer_nuggets/utils/benchmark.py

@@ -5,7 +5,6 @@
 import os
 import random
 import statistics
-import warnings
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -20,13 +19,42 @@
 from torch.cuda._memory_viz import profile_plot  # type: ignore
 from torch.profiler import profile, ProfilerActivity, record_function, schedule
 
-warnings.filterwarnings("ignore", message=".*SyncActivityProfilerHandler.*")
-warnings.filterwarnings("ignore", message=".*profiler_start.*")
-warnings.filterwarnings("ignore", message=".*profiler_stop.*")
-
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
+_KINETO_LOG_LEVEL_ENV = "KINETO_LOG_LEVEL"
+_KINETO_SUPPRESS_ALL_LOGS_LEVEL = "6"
+_KEEP_KINETO_LOG_LEVEL_ENV = "TRANSFORMER_NUGGETS_KEEP_KINETO_LOG_LEVEL"
+
+
+@contextmanager
+def _suppress_sync_activity_profiler_logs():
+    """Suppress noisy Kineto profiler start/stop logs during benchmark timing.
+
+    Recent PyTorch nightlies can print lines like::
+
+        USDT:... SyncActivityProfilerHandler.cpp:52] profiler_start
+        USDT:... SyncActivityProfilerHandler.cpp:59] profiler_stop
+
+    These are emitted by Kineto native logging, not Python ``warnings``.
+    ``KINETO_LOG_LEVEL=6`` disables these logs and is read dynamically by
+    Kineto, so setting it only around the profiler-backed benchmark call avoids
+    subprocesses and fd-level stderr redirection.
+    """
+    if os.environ.get(_KEEP_KINETO_LOG_LEVEL_ENV):
+        yield
+        return
+
+    previous_level = os.environ.get(_KINETO_LOG_LEVEL_ENV)
+    os.environ[_KINETO_LOG_LEVEL_ENV] = _KINETO_SUPPRESS_ALL_LOGS_LEVEL
+    try:
+        yield
+    finally:
+        if previous_level is None:
+            os.environ.pop(_KINETO_LOG_LEVEL_ENV, None)
+        else:
+            os.environ[_KINETO_LOG_LEVEL_ENV] = previous_level
+
 
 def _nvml():
     library_path = ctypes.util.find_library("nvidia-ml") or "libnvidia-ml.so.1"
@@ -267,7 +295,8 @@ def _call_do_bench_using_profiling(
     call_kwargs = {"rep": rep}
     if "is_vetted_benchmarking" in params:
         call_kwargs["is_vetted_benchmarking"] = is_vetted_benchmarking
-    return do_bench_using_profiling(fn, **call_kwargs)
+    with _suppress_sync_activity_profiler_logs():
+        return do_bench_using_profiling(fn, **call_kwargs)
 
 
 def _benchmark_cuda_graph_replay_samples_us(
@@ -373,13 +402,14 @@ def benchmark_cuda_function_stats(func: Callable, *args, **kwargs) -> CudaBenchm
         no_args = lambda: func(*args, **kwargs)
         from torch._inductor.runtime.benchmarking import benchmarker
 
-        samples_ms = benchmarker.benchmark_gpu(
-            no_args,
-            benchmark_iters=num_iters,
-            memory_warmup_iters=memory_warmup_iters,
-            return_mode="all",
-            is_vetted_benchmarking=is_vetted_benchmarking,
-        )
+        with _suppress_sync_activity_profiler_logs():
+            samples_ms = benchmarker.benchmark_gpu(
+                no_args,
+                benchmark_iters=num_iters,
+                memory_warmup_iters=memory_warmup_iters,
+                return_mode="all",
+                is_vetted_benchmarking=is_vetted_benchmarking,
+            )
     return CudaBenchmarkStats.from_samples(
         (float(sample) * 1e3 for sample in samples_ms),
         confidence=confidence,