feat(profiling): real torch flame graphs

.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/internal/settings/_supported_configurations.py

@@ -375,6 +375,7 @@
         "DD_PROFILING_OUTPUT_PPROF",
         "DD_PROFILING_PYTORCH_ENABLED",
         "DD_PROFILING_PYTORCH_EVENTS_LIMIT",
+        "DD_PROFILING_PYTORCH_MAX_FRAMES",
         "DD_PROFILING_SAMPLE_POOL_CAPACITY",
         "DD_PROFILING_SAMPLE_SIZE",
         "DD_PROFILING_SAMPLING_INTERVAL",

ddtrace/internal/settings/profiling.py

@@ -540,6 +540,15 @@ class ProfilingConfigPytorch(DDConfig):
         help="How many events the PyTorch profiler records each collection",
     )
 
+    max_frames = DDConfig.v(
+        int,
+        "max_frames",
+        default=128,
+        validator=validators.range(1, t.cast(int, float("inf"))),
+        help_type="Integer",
+        help="Maximum number of frames to capture",
+    )
+
 
 class ProfilingConfigException(DDConfig):
     __item__ = __prefix__ = "exception"

ddtrace/internal/settings/profiling.pyi

@@ -55,6 +55,7 @@ class ProfilingConfigHeap(DDConfig):
 class ProfilingConfigPytorch(DDConfig):
     enabled: bool
     events_limit: int
+    max_frames: int
 
 class ProfilingConfigException(DDConfig):
     enabled: bool

ddtrace/profiling/collector/pytorch.py

@@ -19,6 +19,15 @@
 _NANOS_PER_MICROSECOND = 1e3
 
 
+# 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__(
         self,
@@ -112,10 +121,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 +158,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 < config.pytorch.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

releasenotes/notes/profiling-improve-pytorch-flame-graphs-6e0d84ac06f4f797.yaml

@@ -0,0 +1,3 @@
+features:
+  - |
+    profiling: The PyTorch profiler now properly shows nested frames in flame graphs.

supported-configurations.json

@@ -2836,6 +2836,13 @@
         "default": "1000000"
       }
     ],
+    "DD_PROFILING_PYTORCH_MAX_FRAMES": [
+      {
+        "implementation": "A",
+        "type": "int",
+        "default": "128"
+      }
+    ],
     "DD_PROFILING_SAMPLE_POOL_CAPACITY": [
       {
         "implementation": "A",

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,8 +7,61 @@
 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):
+    from ddtrace.profiling.collector.pytorch import _DEVICE_FRAME_FILE_NAME
+    from ddtrace.profiling.collector.pytorch import _FILE_PLACEHOLDER
+
     filename = str(tmp_path / "pprof")
     monkeypatch.setenv("DD_PROFILING_OUTPUT_PPROF", filename)
     monkeypatch.setenv("DD_PROFILING_ENABLED", "1")
@@ -28,12 +81,12 @@ def test_call_script_pytorch_gpu(tmp_path, monkeypatch):
         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=_DEVICE_FRAME_FILE_NAME,
                 line_no=0,
             ),
         ],

tests/telemetry/test_writer.py

@@ -348,6 +348,7 @@ def test_app_started_event_configuration_override(test_agent_session, run_python
         {"name": "DD_PROFILING_OUTPUT_PPROF", "origin": "default", "value": None},
         {"name": "DD_PROFILING_PYTORCH_ENABLED", "origin": "default", "value": False},
         {"name": "DD_PROFILING_PYTORCH_EVENTS_LIMIT", "origin": "default", "value": 1000000},
+        {"name": "DD_PROFILING_PYTORCH_MAX_FRAMES", "origin": "default", "value": 128},
         {"name": "DD_PROFILING_SAMPLE_POOL_CAPACITY", "origin": "default", "value": 4},
         {"name": "DD_PROFILING_STACK_ENABLED", "origin": "env_var", "value": False},
         {"name": "DD_PROFILING_STACK_NATIVE_FRAMES", "origin": "default", "value": True},