Only print dde partial fx graph for export (pytorch#153218)

* Only print dde partial fx graph for export

Get pytorch#149831 into 2.7.1

* [dynamo] Add test to ensure we don't print fx graph upon data dependent graph break

This adds a regression test for pytorch#149831, also as part of getting it
cherry-picked into 2.7.1.

ghstack-source-id: fedc9ea
Pull Request resolved: pytorch#153416

Author: StrongerXi

test/dynamo/test_repros.py

@@ -15,6 +15,7 @@
 import itertools
 import os
 import random
+import sys
 import types
 import typing
 import unittest
@@ -6870,6 +6871,31 @@ def fn(x):
         x = torch.randn(4)
         self.assertEqual(fn(x), opt_fn(x))
 
+    def test_data_dependent_error_log_no_print(self):
+        # This is a regression test case for
+        # https://github.com/pytorch/pytorch/pull/149831
+        from io import StringIO
+
+        capturedOutput = StringIO()
+        sys.stderr = capturedOutput
+
+        @torch.compile(fullgraph=True)
+        def func(a):
+            if a.sum() > 0:
+                return a + 1
+            return a + 2
+
+        a = torch.rand(10, 10)
+        try:
+            func(a)
+        except Exception:
+            pass
+        sys.stderr = sys.__stderr__
+
+        # Make sure we don't _print_ out the graph module.
+        output = capturedOutput.getvalue()
+        self.assertNotIn("class GraphModule", output)
+
 
 instantiate_parametrized_tests(ReproTests)
 

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:

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
@@ -1089,6 +1090,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

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: