Log compilation_metrics to TORCH_TRACE (#122638)

Summary: It's not technically needed as you can get it from Scuba too, but it's more convenient for tlparse to get at it this way. Signed-off-by: Edward Z. Yang <[email protected]> X-link: pytorch/pytorch#122638 Approved by: https://github.com/albanD Reviewed By: atalman Differential Revision: D55397237 Pulled By: ezyang fbshipit-source-id: 5508d97e207049eca63acb1f5955924d1e142c91
pytorch · Mar 27, 2024 · 697752c · 697752c
1 parent fd2228f
commit 697752c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/userbenchmark/dynamo/dynamobench/_dynamo/utils.py b/userbenchmark/dynamo/dynamobench/_dynamo/utils.py
@@ -673,6 +673,13 @@ class CompilationMetrics:
 def record_compilation_metrics(compilation_metrics: CompilationMetrics):
     global _compilation_metrics
     _compilation_metrics.append(compilation_metrics)
+    torch._logging.trace_structured(
+        "compilation_metrics",
+        lambda: {
+            k: list(v) if isinstance(v, set) else v
+            for k, v in dataclasses.asdict(compilation_metrics).items()
+        },
+    )
     if config.log_compilation_metrics:
         log_compilation_event(compilation_metrics)