Revert "Refactor stack_trace preservation for node meta preservation (pytorch#90803)"

This reverts commit 0f1302e.

Reverted pytorch#90803 on behalf of https://github.com/DanilBaibak due to Break internal build

Author: pytorchmergebot

test/test_functionalization.py

@@ -178,7 +178,7 @@ def g(x):
             from torch._functorch.aot_autograd import setup_stacktrace_preservation_hooks
             import torch.fx.traceback as fx_traceback
             setup_stacktrace_preservation_hooks([loss.grad_fn])
-            with fx_traceback.preserve_node_meta():
+            with fx_traceback.override_stack_trace():
                 loss.backward()
             return x.grad
 

torch/_dynamo/eval_frame.py

@@ -668,7 +668,7 @@ def run_node(self, n):
     if aten_graph:
         # Running graph with interpreter is needed for propagating the stack_trace
         def graph_with_interpreter(*args):
-            with torch.fx.traceback.preserve_node_meta():
+            with torch.fx.traceback.override_stack_trace():
                 return torch.fx.Interpreter(graph).run(*args)
 
         graph = make_fx(

torch/_functorch/aot_autograd.py

@@ -790,7 +790,7 @@ def joint_forward_backward(
         backward_out = []
         # Call the backwards pass
         if grad_primals:
-            with fx_traceback.preserve_node_meta():
+            with fx_traceback.override_stack_trace():
                 backward_out = torch.autograd.grad(
                     needed_outs,
                     grad_primals,
@@ -2319,7 +2319,7 @@ def functional_call(*args, **kwargs):
             mod, pytree.tree_unflatten(args[:params_len], params_spec)
         ):
             if isinstance(mod, torch.fx.GraphModule):
-                with fx_traceback.preserve_node_meta(), warnings.catch_warnings():
+                with fx_traceback.override_stack_trace(), warnings.catch_warnings():
                     warnings.filterwarnings(
                         "ignore", "Anomaly Detection has been enabled."
                     )

torch/fx/interpreter.py

@@ -153,7 +153,7 @@ def run(self, *args, initial_env : Optional[Dict[Node, Any]] = None, enable_io_p
 
     @contextmanager
     def _set_current_node(self, node):
-        with fx_traceback.set_current_meta(node.meta):
+        with fx_traceback.append_stack_trace(node.stack_trace), fx_traceback.set_current_meta(node.meta):
             yield
 
     @compatibility(is_backward_compatible=True)
@@ -477,7 +477,7 @@ def transform(self) -> GraphModule:
         Transform ``self.module`` and return the transformed
         ``GraphModule``.
         """
-        with fx_traceback.preserve_node_meta():
+        with fx_traceback.override_stack_trace():
             result = super().run(enable_io_processing=False)
         if result is not None:
             def strip_proxy(a : Union[Argument, Proxy]) -> Any:

torch/fx/proxy.py

@@ -76,19 +76,10 @@ def create_proxy(self, kind: str, target: Target, args: Tuple[Any, ...], kwargs:
             proxy = proxy_factory_fn(node)
 
         # Optionally set stack trace on the created Node for debugging purposes
-        if fx_traceback.has_preserved_node_meta():
-            current_meta: Dict[str, Any] = fx_traceback.get_current_meta()
-
-            # Explicitly set the stack_trace and nn_module_stack on the node.meta
-            # If other meta fields are needed, they can be added here
-            stack_trace = current_meta.get("stack_trace")
-            if stack_trace:
-                proxy.node.stack_trace = stack_trace
-
-            nn_module_stack = current_meta.get("nn_module_stack")
-            if nn_module_stack:
-                proxy.node.meta["nn_module_stack"] = nn_module_stack
-
+        if fx_traceback.is_stack_trace_overridden():
+            proxy.node.meta = fx_traceback.get_current_meta()
+            stacks = fx_traceback.format_stack()
+            proxy.node.stack_trace = '\n'.join(reversed(stacks))
         elif self.record_stack_traces:
             user_frame = self._find_user_frame()
             if user_frame:

torch/fx/traceback.py

@@ -1,63 +1,80 @@
 import traceback
 from contextlib import contextmanager
-from typing import List, Any, Dict
+from typing import Optional, List, Any, Dict
 from ._compatibility import compatibility
 
-__all__ = ['preserve_node_meta', 'has_preserved_node_meta',
-           'set_stack_trace', 'format_stack',
-           'set_current_meta', 'get_current_meta']
+__all__ = ['override_stack_trace', 'set_stack_trace', 'append_stack_trace', 'format_stack',
+           'is_stack_trace_overridden', 'get_current_meta', 'set_current_meta']
 
+
+current_stack: List[str] = []
 current_meta: Dict[str, Any] = {}
-should_preserve_node_meta = False
+is_overridden = False
 
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
-def preserve_node_meta():
-    global should_preserve_node_meta
+def override_stack_trace():
+    global is_overridden
 
-    saved_should_preserve_node_meta = should_preserve_node_meta
+    saved_is_overridden = is_overridden
     try:
-        should_preserve_node_meta = True
+        is_overridden = True
         yield
     finally:
-        should_preserve_node_meta = saved_should_preserve_node_meta
-
+        is_overridden = saved_is_overridden
 
 @compatibility(is_backward_compatible=False)
 def set_stack_trace(stack : List[str]):
-    global current_meta
+    global current_stack
+
+    if is_overridden and stack:
+        current_stack = stack
+
+@compatibility(is_backward_compatible=False)
+@contextmanager
+def append_stack_trace(stack : Optional[str]):
+    """
+    The content of stack here is an entire stacktraces as a string
+    """
+    global current_stack
 
-    if should_preserve_node_meta and stack:
-        current_meta["stack_trace"] = "".join(stack)
+    if is_overridden and stack:
+        try:
+            current_stack.append(stack)
+            yield
+        finally:
+            current_stack.pop()
+    else:
+        yield
 
 
 @compatibility(is_backward_compatible=False)
 def format_stack() -> List[str]:
-    if should_preserve_node_meta:
-        return [current_meta.get("stack_trace", "")]
+    if is_overridden:
+        return current_stack.copy()
     else:
         # fallback to traceback.format_stack()
         return traceback.format_list(traceback.extract_stack()[:-1])
 
 
 @compatibility(is_backward_compatible=False)
-def has_preserved_node_meta() -> bool:
-    return should_preserve_node_meta
+def is_stack_trace_overridden() -> bool:
+    return is_overridden
 
 
 @compatibility(is_backward_compatible=False)
 @contextmanager
 def set_current_meta(meta : Dict[str, Any]):
     global current_meta
 
-    if should_preserve_node_meta and meta:
-        saved_meta = current_meta
+    old_meta = current_meta
+    if is_overridden and meta:
         try:
             current_meta = meta
             yield
         finally:
-            current_meta = saved_meta
+            current_meta = old_meta
     else:
         yield