Skip to content

Commit 9847c98

Browse files
committed
Add CUPTI monitor profiler support
## Human Note ## Agent note Add an opt-in CUPTI monitor configuration to the shared profiler wrapper, including PM metric sampling, environment counters, CUDA Graph dependency recording, event-node attribution, and metric discovery. Constructing the profiler now prepares the backend immediately, which lets graph recorders arm before capture while preserving existing context-manager usage. Monitor traces use PyTorch's native Perfetto exporter and normalize its unconditional `.gz` suffix, so callers receive the requested `.pftrace` or JSON path instead of having to discover a differently named artifact. The optional dependency keeps cupti-python out of the default installation. ## Test Plan ```bash ~/.venvs/nightly/bin/python -m pytest test/test_profiler.py test/test_perfetto.py -q uvx ruff check transformer_nuggets/utils/benchmark.py transformer_nuggets/utils/__init__.py test/test_profiler.py uvx ruff format --check transformer_nuggets/utils/benchmark.py transformer_nuggets/utils/__init__.py test/test_profiler.py prek ```
1 parent 833b8a1 commit 9847c98

4 files changed

Lines changed: 266 additions & 55 deletions

File tree

pyproject.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,10 @@ dev = [
4242
"docstring-parser",
4343
]
4444

45+
cupti-monitor = [
46+
"cupti-python>=13.3",
47+
]
48+
4549
llama = [
4650
"sentencepiece==0.1.99",
4751
"datasets==2.15.0",

test/test_profiler.py

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import gzip
2+
from pathlib import Path
3+
4+
import pytest
5+
6+
from transformer_nuggets.utils import benchmark
7+
from transformer_nuggets.utils.benchmark import (
8+
CuptiMonitorConfig,
9+
_write_cupti_monitor_trace,
10+
)
11+
12+
13+
def test_cupti_monitor_config_builds_pytorch_options():
14+
config = CuptiMonitorConfig(
15+
environment_counters=True,
16+
pm_metrics=("sm__cycles_active.avg.pct_of_peak_sustained_elapsed",),
17+
graph_dependencies=True,
18+
event_node_ids=True,
19+
cuda_sync_events=True,
20+
pftrace_compression_level=4,
21+
)
22+
23+
assert config.custom_profiler_config() == {
24+
"backend": "cupti_monitor",
25+
"enable_environment_counters": True,
26+
"enable_pm_sampling": True,
27+
"pm_metrics": ["sm__cycles_active.avg.pct_of_peak_sustained_elapsed"],
28+
"enable_graph_dependencies": True,
29+
"enable_event_node_ids": True,
30+
"enable_cuda_sync_events": True,
31+
"pftrace_compression_level": 4,
32+
}
33+
34+
35+
@pytest.mark.parametrize("level", [-1, 10])
36+
def test_cupti_monitor_config_rejects_invalid_compression(level):
37+
with pytest.raises(ValueError, match="between 0 and 9"):
38+
CuptiMonitorConfig(pftrace_compression_level=level)
39+
40+
41+
class GzipExportProfiler:
42+
def __init__(self, payload: bytes) -> None:
43+
self.payload = payload
44+
45+
def export_chrome_trace(self, path: str) -> None:
46+
with gzip.open(f"{path}.gz", "wb") as output:
47+
output.write(self.payload)
48+
49+
50+
def test_cupti_monitor_pftrace_export_hides_gzip_suffix(tmp_path):
51+
path = tmp_path / "trace.pftrace"
52+
profiler = GzipExportProfiler(b"compressed native trace")
53+
54+
_write_cupti_monitor_trace(profiler, path, trace_format="track_event")
55+
56+
assert path.exists()
57+
assert not Path(f"{path}.gz").exists()
58+
with gzip.open(path, "rb") as trace:
59+
assert trace.read() == b"compressed native trace"
60+
61+
62+
def test_cupti_monitor_json_export_hides_gzip_suffix(tmp_path):
63+
path = tmp_path / "trace.json"
64+
profiler = GzipExportProfiler(b'{"traceEvents": []}')
65+
66+
_write_cupti_monitor_trace(profiler, path, trace_format="chrome_json")
67+
68+
assert path.read_bytes() == b'{"traceEvents": []}'
69+
assert not Path(f"{path}.gz").exists()
70+
71+
72+
def test_cupti_monitor_profiler_is_constructed_before_context_entry(monkeypatch, tmp_path):
73+
experimental_config = object()
74+
constructed = []
75+
76+
class FakeProfiler:
77+
def start(self):
78+
pass
79+
80+
def stop(self):
81+
pass
82+
83+
def fake_profile(**kwargs):
84+
constructed.append(kwargs)
85+
return FakeProfiler()
86+
87+
monkeypatch.setattr(
88+
benchmark,
89+
"_cupti_monitor_experimental_config",
90+
lambda config: experimental_config,
91+
)
92+
monkeypatch.setattr(benchmark.torch.profiler, "profile", fake_profile)
93+
94+
context = benchmark.profiler(
95+
tmp_path / "trace",
96+
cupti_monitor=CuptiMonitorConfig(graph_dependencies=True),
97+
)
98+
99+
assert constructed[0]["experimental_config"] is experimental_config
100+
with context:
101+
pass
Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,23 @@
11
from transformer_nuggets.utils.benchmark import (
2-
benchmark_torch_function_in_microseconds,
2+
DEFAULT_CUPTI_MONITOR_PM_METRICS,
3+
CudaBenchmarkStats,
4+
CuptiMonitorConfig,
5+
ProfileConfig,
6+
attach_oom_observer,
37
benchmark_cuda_function_in_microseconds,
48
benchmark_cuda_function_in_microseconds_triton,
59
benchmark_cuda_function_stats,
10+
benchmark_torch_function_in_microseconds,
11+
cuda_memory_usage,
612
locked_clocks,
713
max_memory_usage,
8-
cuda_memory_usage,
914
profile_function,
10-
ProfileConfig,
11-
CudaBenchmarkStats,
12-
save_memory_snapshot,
1315
profiler,
14-
attach_oom_observer,
16+
save_memory_snapshot,
17+
supported_cupti_monitor_metrics,
1518
)
16-
from transformer_nuggets.utils.tracing import LoggingMode, NanInfDetect
17-
from transformer_nuggets.utils.triton import print_sass
18-
from transformer_nuggets.utils.merge_traces import merge_traces
1919
from transformer_nuggets.utils.memory_viz import generate_memory_comparison_html
20+
from transformer_nuggets.utils.merge_traces import merge_traces
2021
from transformer_nuggets.utils.perfetto import (
2122
TraceFormat,
2223
default_trace_path,
@@ -28,5 +29,7 @@
2829
write_trace,
2930
write_track_event_trace,
3031
)
32+
from transformer_nuggets.utils.tracing import LoggingMode, NanInfDetect
3133
from transformer_nuggets.utils.track_event import chrome_trace_to_track_event_trace
34+
from transformer_nuggets.utils.triton import print_sass
3235
# from transformer_nuggets.utils.model_extraction import extract_attention_data

0 commit comments

Comments
 (0)