feat(profiling): real torch flame graphs

Author: KowalskiThomas

.github/workflows/pytorch_gpu_tests.yml

@@ -77,6 +77,7 @@ jobs:
           cp "$REPO_ROOT/tests/profiling/__init__.py" "$TEST_ROOT/profiling/"
           cp "$REPO_ROOT/tests/profiling/test_pytorch.py" "$TEST_ROOT/profiling/"
           cp "$REPO_ROOT/tests/profiling/simple_program_pytorch_gpu.py" "$TEST_ROOT/profiling/"
+          cp "$REPO_ROOT/tests/profiling/simple_program_pytorch_cpu.py" "$TEST_ROOT/profiling/"
           cp -R "$REPO_ROOT/tests/profiling/collector" "$TEST_ROOT/profiling/"
 
           cd "$ISOLATED_DIR"

ddtrace/profiling/collector/pytorch.py

@@ -18,6 +18,18 @@
 
 _NANOS_PER_MICROSECOND = 1e3
 
+# Safety bound on how far we walk up the cpu_parent chain when reconstructing
+# the operator call tree, to guard against pathological depths.
+_MAX_FRAMES = 128
+
+# Frames require a file name, but GPU frames are not from a Python file.
+# We use the following as a placeholder.
+_FILE_PLACEHOLDER = "<native>"
+
+# We use a frame to group GPU frames under a device.
+# This frame also needs a file name, none can really make sense.
+_DEVICE_FRAME_FILE_NAME = "<torch>"
+
 
 class _WrappedTorchProfiler(wrapt.ObjectProxy):
     def __init__(
@@ -112,10 +124,23 @@ def _handle_torch_trace(prof: Any) -> None:
         LOG.debug("Dropped events. events_limit %d, len(events): %d", events_limit, num_events)
         events = random.sample(events, events_limit)  # nosec: used for sampling, not security
 
-    # Determine which attributes to use once (avoid per-event getattr checks)
+    # Determine which attributes to use once (avoid per-event getattr checks).
+    # For CPU operators we use the "self" (exclusive of children) variants so that,
+    # once we reconstruct the operator call tree below, parent frames aggregate to
+    # their inclusive totals without double counting their children. CUDA device
+    # events are leaves and report 0 for the "self" device metrics (they are async),
+    # so for those we use the event's own totals instead.
     sample_event = events[0]
-    use_device_time = hasattr(sample_event, "device_time")
-    use_device_memory = hasattr(sample_event, "device_memory_usage")
+    self_gpu_time_attr = (
+        "self_device_time_total" if hasattr(sample_event, "self_device_time_total") else "self_cuda_time_total"
+    )
+    total_gpu_time_attr = "device_time_total" if hasattr(sample_event, "device_time_total") else "cuda_time_total"
+    self_gpu_memory_attr = (
+        "self_device_memory_usage" if hasattr(sample_event, "self_device_memory_usage") else "self_cuda_memory_usage"
+    )
+    total_gpu_memory_attr = (
+        "device_memory_usage" if hasattr(sample_event, "device_memory_usage") else "cuda_memory_usage"
+    )
 
     # Earlier PyTorch versions use microseconds, later versions use nanoseconds
     kineto_results = prof.profiler.kineto_results
@@ -136,44 +161,62 @@ def _handle_torch_trace(prof: Any) -> None:
         handle = ddup.SampleHandle()
         data_added = False
 
-        # cpu time sample
-        cpu_time: int = e.cpu_time
-        if cpu_time > 0:
+        # Number of times this event aggregates. The time metrics below are
+        # totals across those occurrences, so we divide by count to recover the
+        # per-occurrence value and let ddup re-multiply by count.
+        count: int = e.count or 1
+
+        # Cache str(device_type) since we use it multiple times. CPU operators get
+        # the call tree reconstructed and use exclusive ("self") metrics; CUDA
+        # device events are leaves and use their full totals.
+        device_type_str = str(e.device_type)
+        is_cpu = device_type_str.startswith("DeviceType.CPU")
+
+        # cpu time sample (exclusive of children)
+        self_cpu_time: int = e.self_cpu_time_total
+        if self_cpu_time > 0:
             data_added = True
-            handle.push_cputime(int(cpu_time * _NANOS_PER_MICROSECOND), e.count)
+            handle.push_cputime(int(self_cpu_time / count * _NANOS_PER_MICROSECOND), count)
 
-        # gpu time sample
-        gpu_time: int = e.device_time if use_device_time else e.cuda_time
+        # gpu time sample: exclusive for CPU operators, full device time for leaves
+        gpu_time: int = getattr(e, self_gpu_time_attr) if is_cpu else getattr(e, total_gpu_time_attr)
         if gpu_time > 0:
             data_added = True
-            handle.push_gpu_gputime(int(gpu_time * _NANOS_PER_MICROSECOND), e.count)
+            handle.push_gpu_gputime(int(gpu_time / count * _NANOS_PER_MICROSECOND), count)
 
         # gpu flops sample
         flops: int = e.flops
         if flops is not None and flops > 0:
             data_added = True
-            handle.push_gpu_flops(flops, e.count)
+            handle.push_gpu_flops(flops, count)
 
-        # GPU memory usage
-        gpu_memory: int = e.device_memory_usage if use_device_memory else e.cuda_memory_usage
+        # GPU memory usage: exclusive for CPU operators, full usage for leaves
+        gpu_memory: int = getattr(e, self_gpu_memory_attr) if is_cpu else getattr(e, total_gpu_memory_attr)
         if gpu_memory is not None and gpu_memory > 0:
             data_added = True
-            handle.push_gpu_memory(gpu_memory, e.count)
+            handle.push_gpu_memory(gpu_memory, count)
 
         if not data_added:
             if empty_events_count % 1000 == 0:
                 LOG.debug("%d events with no data to record: %s", empty_events_count, e)
             empty_events_count += 1
             continue
 
-        # Cache str(device_type) since we use it multiple times
-        device_type_str = str(e.device_type)
-
-        handle.push_frame(e.name, "unknown-file", 0, 0)
+        # Reconstruct the operator call tree by walking up the cpu_parent chain.
+        # Frames are pushed leaf-first (this event), then each ancestor, so that
+        # the flame graph nests children under their parents instead of rendering
+        # every operator as a flat leaf. Stacks go root last.
+        handle.push_frame(e.name, _FILE_PLACEHOLDER, 0, 0)
+        parent = getattr(e, "cpu_parent", None)
+        depth = 0
+        while parent is not None and depth < _MAX_FRAMES:
+            handle.push_frame(parent.name, _FILE_PLACEHOLDER, 0, 0)
+            parent = getattr(parent, "cpu_parent", None)
+            depth += 1
         # Pushing pseudoframes for the device name ("device.CPU" or "device.CUDA")
         # onto the stack allows differentiation of pytorch frames from other profiling frames
         # in the flame graph. Note that stacks go root last, so this goes at the end.
-        handle.push_frame(f"PYTORCH_{device_type_str}", "unknown-file", 0, 0)
+        handle.push_frame(f"PYTORCH_{device_type_str}", _DEVICE_FRAME_FILE_NAME, 0, 0)
         handle.push_gpu_device_name(f"cuda {e.device_index}")
 
         # Get thread info from cache or compute and cache it

tests/profiling/simple_program_pytorch_cpu.py

@@ -0,0 +1,38 @@
+import torch
+import torch.nn
+from torch.profiler import ProfilerActivity
+
+
+class TinyModel(torch.nn.Module):
+    """A tiny model whose forward pass produces a nested CPU operator tree.
+
+    A ``torch.nn.Linear`` forward records a parent ``aten::linear`` operator
+    containing child operators (``aten::addmm``/``aten::matmul``/``aten::t``),
+    which is what lets us assert that the profiler reconstructs nesting.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.net = torch.nn.Sequential(
+            torch.nn.Linear(64, 64),
+            torch.nn.ReLU(),
+            torch.nn.Linear(64, 16),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+def run() -> None:
+    model = TinyModel()
+    inputs = torch.randn(32, 64)
+
+    with torch.profiler.profile(
+        activities=[ProfilerActivity.CPU],
+    ):
+        for _ in range(5):
+            model(inputs).sum().backward()
+
+
+if __name__ == "__main__":
+    run()

tests/profiling/test_pytorch.py

@@ -7,6 +7,56 @@
 from tests.utils import call_program
 
 
+try:
+    import torch  # noqa: F401
+
+    _HAS_TORCH = True
+except ImportError:
+    _HAS_TORCH = False
+
+
+@pytest.mark.skipif(not _HAS_TORCH, reason="torch is not installed")
+def test_call_script_pytorch_cpu(tmp_path, monkeypatch):
+    """The torch profiler integration should reconstruct the operator call tree.
+
+    Each torch event's stack is built by walking the ``cpu_parent`` chain, so a
+    nested operator (e.g. ``aten::addmm`` under ``aten::linear``) must show up as
+    a multi-frame stack rooted at the ``PYTORCH_DeviceType.CPU`` pseudo-frame,
+    rather than as a flat two-frame stack.
+    """
+    filename = str(tmp_path / "pprof")
+    monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
+    monkeypatch.setenv("DD_PROFILING_ENABLED", "1")
+    monkeypatch.setenv("DD_PROFILING_PYTORCH_ENABLED", "1")
+    _, stderr, exitcode, _ = call_program(
+        "ddtrace-run", sys.executable, os.path.join(os.path.dirname(__file__), "simple_program_pytorch_cpu.py")
+    )
+    assert exitcode == 0, f"Profiler exited with code {exitcode}. Stderr: {stderr}"
+
+    profile = pprof_utils.parse_newest_profile(filename)
+    samples = pprof_utils.get_samples_with_value_type(profile, "cpu-time")
+    assert len(samples) > 0, "Expected at least one cpu-time sample"
+
+    # Find at least one sample with a reconstructed (nested) stack: more than the
+    # old flat two frames ([op_name, PYTORCH_DeviceType.CPU]) and rooted at the
+    # CPU device pseudo-frame.
+    nested_found = False
+    for sample in samples:
+        locations = [pprof_utils.get_location_from_id(profile, loc_id) for loc_id in sample.location_id]
+        if len(locations) < 3:
+            continue
+        if locations[-1].function_name != "PYTORCH_DeviceType.CPU":
+            continue
+        # The frames between the leaf and the device root are the operator
+        # ancestor chain; require at least one genuine parent operator.
+        op_frames = [loc.function_name for loc in locations[:-1]]
+        if len(op_frames) >= 2:
+            nested_found = True
+            break
+
+    assert nested_found, "Expected at least one cpu-time sample with a nested operator stack"
+
+
 @pytest.mark.skipif(not os.getenv("DD_PROFILING_PYTORCH_ENABLED", False), reason="Not testing pytorch GPU")
 def test_call_script_pytorch_gpu(tmp_path, monkeypatch):
     filename = str(tmp_path / "pprof")
@@ -24,16 +74,18 @@ def test_call_script_pytorch_gpu(tmp_path, monkeypatch):
     print("number of gpu time samples: ", len(samples))
     print("first sample: ", samples[0])
 
+    from ddtrace.profiling.collector.pytorch import _FILE_PLACEHOLDER
+
     expected_sample = pprof_utils.StackEvent(
         locations=[
             pprof_utils.StackLocation(
                 function_name="Memset (Device)",
-                filename="unknown-file",
+                filename=_FILE_PLACEHOLDER,
                 line_no=0,
             ),
             pprof_utils.StackLocation(
                 function_name="PYTORCH_DeviceType.CUDA",
-                filename="unknown-file",
+                filename=_FILE_PLACEHOLDER,
                 line_no=0,
             ),
         ],