Only print dde partial fx graph for export

bobrenjc93 · bobrenjc93 · commit ecb672b47bb0 · 2025-03-23T15:33:41.000-07:00
Lazos correctly pointed out this doesn't make sense for compile since we graph break in compile. This results in tons of unwanted user log spew. We do want this in export though since it's drastiaclly reduced the support load for DDEs. This PR does the refactor to keep it in export but remove it from compile ghstack-source-id: db17302 Pull Request resolved: #149831
diff --git a/torch/_dynamo/symbolic_convert.py b/torch/_dynamo/symbolic_convert.py
@@ -1342,15 +1342,13 @@ def run(self):
                 raise
             except RuntimeError as e:
                 if hasattr(e, "msg") and "Data-dependent" in e.msg:
-                    print(
-                        "\n"
-                        + torch.fx.GraphModule(
-                            self.output.nn_modules, self.output.graph
-                        ).print_readable(
-                            print_output=False, include_stride=True, include_device=True
-                        ),
-                        file=sys.stderr,
+                    readable_graph = torch.fx.GraphModule(
+                        self.output.nn_modules, self.output.graph
+                    ).print_readable(
+                        print_output=False, include_stride=True, include_device=True
                     )
+                    e.partial_fx_graph = readable_graph  # type: ignore[attr-defined]
+                    raise
 
                 raise
             except Exception as e:
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
@@ -5,6 +5,7 @@
 import inspect
 import logging
 import re
+import sys
 import time
 import warnings
 from contextlib import contextmanager, nullcontext
@@ -1096,6 +1097,13 @@ def wrapper(*args, **kwargs):
                     message=str(e),
                     flags=_EXPORT_FLAGS,
                 )
+
+            if hasattr(e, "partial_fx_graph"):
+                print(
+                    e.partial_fx_graph,
+                    file=sys.stderr,
+                )
+
             raise e
         finally:
             _EXPORT_FLAGS = None
diff --git a/torch/fx/_symbolic_trace.py b/torch/fx/_symbolic_trace.py
@@ -7,7 +7,6 @@
 import inspect
 import math
 import os
-import sys
 import warnings
 from itertools import chain
 from types import CodeType, FunctionType, ModuleType
@@ -843,14 +842,12 @@ def forward(*args, **kwargs):
             self.submodule_paths = None
         except RuntimeError as e:
             if isinstance(e.args[0], str) and "data-dependent" in e.args[0]:
-                print(
-                    "\n"
-                    + self.graph.python_code(
-                        root_module="self",
-                        verbose=True,
-                    ).src,
-                    file=sys.stderr,
-                )
+                partial_fx_graph = self.graph.python_code(
+                    root_module="self",
+                    verbose=True,
+                ).src
+                e.partial_fx_graph = partial_fx_graph  # type: ignore[attr-defined]
+                raise
 
             raise
         finally:
