PaddlePaddle
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/exception_stack.py
Copy file name to clipboard
+31-21Lines changed: 31 additions & 21 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/exception_stack.py
Copy file name to clipboard
+31-21Lines changed: 31 additions & 21 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+22-18Lines changed: 22 additions & 18 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
Copy file name to clipboard
+22-18Lines changed: 22 additions & 18 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
Copy file name to clipboard
+6-5Lines changed: 6 additions & 5 deletions b/‎python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
Copy file name to clipboard
+6-5Lines changed: 6 additions & 5 deletions
@@ -15,18 +15,21 @@
 from __future__ import annotations
 
 import dataclasses
+from typing import TYPE_CHECKING
 
 from ...utils import InnerError
 from .variables import ConstantVariable, ExceptionVariable
 
+if TYPE_CHECKING:
+    from .function_graph import FunctionGraph
+
 
 @dataclasses.dataclass
 class ExceptionStack:
-
     # This data structure manages exceptions as in CPython, primarily handling
     # the __context__ attribute of SotCapturedException.
 
-    _exception_stack: list[ExceptionVariable | None] = dataclasses.field(
+    _exception_stack: list[ExceptionVariable] = dataclasses.field(
         default_factory=list
     )
     _current_exception: ExceptionVariable | None = dataclasses.field(
@@ -36,7 +39,9 @@ class ExceptionStack:
     def clear_current_exception(self):
         self._current_exception = None
 
-    def set_current_exception(self, val: ExceptionVariable, graph):
+    def set_current_exception(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
         self._set_context_and_break_context_reference_cycle(val, graph)
         self._current_exception = val
 
@@ -49,11 +54,18 @@ def get_current_exception(self):
             raise InnerError("Current exception should not be None")
         return self._current_exception
 
-    def _set_context_recursive(self, val: ExceptionVariable, prev_idx):
+    def _set_context_and_break_context_reference_cycle(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
+        # set Exception.__context__
+        self._set_context_recursive(val, len(self._exception_stack) - 1)
+        self._break_context_reference_cycle(val, graph)
+
+    def _set_context_recursive(self, val: ExceptionVariable, prev_idx: int):
         # Recursively sets the __context__ attribute for ExceptionVariable objects
         # in self._exception_stack. Ensures that __context__ is properly linked
         # to the previous exception in the stack.
-        if (ctx := val.__context__) and type(ctx) is not ConstantVariable:
+        if (ctx := val.__context__) and not isinstance(ctx, ConstantVariable):
             return val
         if (
             len(self._exception_stack) + prev_idx > 0
@@ -63,16 +75,18 @@ def _set_context_recursive(self, val: ExceptionVariable, prev_idx):
             val.setattr("__context__", prev)
         return val
 
-    def _break_context_reference_cycle(self, val: ExceptionVariable, graph):
+    def _break_context_reference_cycle(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
         # Detects and breaks cycles in exception __context__ chains using Floyd's algorithm,
         # following CPython's implementation.
 
         fast = slow = val
         slow_update_toggle = False
         while True:
             context = fast.__context__
-            if (
-                type(context) is ConstantVariable
+            if isinstance(
+                context, ConstantVariable
             ):  # End of the chain; no context set
                 break
 
@@ -92,28 +106,24 @@ def _break_context_reference_cycle(self, val: ExceptionVariable, graph):
                 slow = slow.__context__
             slow_update_toggle = not slow_update_toggle
 
-    def _set_context_and_break_context_reference_cycle(
-        self, val: ExceptionVariable, graph
-    ):
-        # set Exception.__context__
-        self._set_context_recursive(val, len(self._exception_stack) - 1)
-        self._break_context_reference_cycle(val, graph)
-
-    def pop(self):
+    def pop(self) -> ExceptionVariable:
         return self._exception_stack.pop()
 
-    def push(self, val):
+    def push(self, val: ExceptionVariable) -> None:
         self._exception_stack.append(val)
 
+    def empty(self) -> bool:
+        return len(self._exception_stack) == 0
+
     def __len__(self):
         return len(self._exception_stack)
 
-    def __str__(self):
-        return f"{self._exception_stack}"
+    def __repr__(self):
+        return f"ExceptionStack({self._exception_stack})"
 
-    def __getitem__(self, idx):
+    def __getitem__(self, idx: int) -> ExceptionVariable:
         return self._exception_stack[idx]
 
-    def cleanup(self):
+    def cleanup(self) -> None:
         self._exception_stack[:] = []
         self._current_exception = None
@@ -60,6 +60,7 @@
 from ..instruction_utils.opcode_info import (
     NEED_TO_BOOL,
     RETURN,
+    ExceptionHandler,
     JumpDirection,
     PopJumpCond,
 )
@@ -138,10 +139,7 @@
 
 # In Python 3.13, the method layout is changed, and a NULL will be pushed after the value.
 CALL_METHOD_LAYOUT_NULL_AFTER_VALUE = sys.version_info >= (3, 13)
-ALREADY_SUPPORTED_EXCEPTION = sys.version_info >= (
-    3,
-    9,
-) and sys.version_info < (
+ALREADY_SUPPORTED_EXCEPTION = sys.version_info < (
     3,
     11,
 )
@@ -356,7 +354,7 @@ def inner(*args, **kwargs):
 
 def fallback_if_python_version_unsupported(fn: Callable):
     def inner(*args, **kwargs):
-        if sys.version_info >= (3, 11):
+        if not ALREADY_SUPPORTED_EXCEPTION:
             raise FallbackError(
                 "SOT currently only partially supports exception handling (Python 3.10 and below). "
                 "Unsupported exception bytecode will fall back to dynamic graph mode."
@@ -679,20 +677,24 @@ def step(self, instr: Instruction):
 
     def handle_exception(self, e: SotCapturedException):
         # TODO(DrRyanHuang): The newly created ExceptionVariable might differ from the previous one
-        e_var = VariableFactory.from_value(e, self._graph, DummyTracker([]))
+        e_var = VariableFactory.from_value(
+            e,
+            self._graph,
+            DummyTracker(e.tracked_args),
+        )
 
         # The exception is not raised by `raise Exception`
         if (
-            len(self.exception_stack) == 0
+            self.exception_stack.empty()
             or self.exception_stack.get_current_exception() != e_var
         ):
             self.exception_stack.set_current_exception(e_var, self._graph)
 
-        if len(self.vframe.block_stack):
+        if self.vframe.block_stack:
             # The implementation is referenced from the exception_unwind section
             # of CPython's main_loop.
             block_stack_entry = self.vframe.block_stack.pop()
-            while block_stack_entry.inst.opname == "EXCEPT_HANDLER":
+            while block_stack_entry.inst.opname == ExceptionHandler.opname:
                 # Remove previous EXCEPT_HANDLER entries, which indicate that the
                 # exception has already been handled. Continue until a SETUP_FINALLY
                 # block is encountered, which signifies an active exception handler.
@@ -703,7 +705,7 @@ def handle_exception(self, e: SotCapturedException):
                     # frame, so the exception is propagated to the outer function for handling.
                     self.stack.pop_n(
                         len(self.stack)
-                    )  # clear stack to prevent memory leaks
+                    )  # Just like CPython; no memory leaks in our simulation
                     raise e
                 block_stack_entry = self.vframe.block_stack.pop()
 
@@ -718,7 +720,7 @@ def handle_exception(self, e: SotCapturedException):
             # Push a dummy EXCEPT_HANDLER block onto the stack to indicate that exception
             # handling has begun and to record the current stack level.
             EXCEPT_HANDLER_INSTRUCTION = Instruction(
-                257, "EXCEPT_HANDLER", None, 0
+                ExceptionHandler.opcode, ExceptionHandler.opname, None, 0
             )
             self.vframe.block_stack.append(
                 BlockStackItem(
@@ -1988,7 +1990,7 @@ def LOAD_ASSERTION_ERROR(self, instr: Instruction):
     def POP_EXCEPT(self, instr: Instruction):
         assert len(self.vframe.block_stack) > 0
 
-        if self.vframe.block_stack[-1].inst.opname != "EXCEPT_HANDLER":
+        if self.vframe.block_stack[-1].inst.opname != ExceptionHandler.opname:
             raise FallbackError(
                 "Bug in SOT tracing of exception handling."
                 "Top of the block stack is not EXCEPT_HANDLER."
@@ -1997,7 +1999,7 @@ def POP_EXCEPT(self, instr: Instruction):
         self.vframe.block_stack.pop()
         self.stack.pop_n(3)
 
-        assert len(self.exception_stack)
+        assert self.exception_stack
         self.exception_stack.pop()
 
     @staticmethod
@@ -2028,7 +2030,7 @@ def _raise_exception_instance(
     @fallback_if_python_version_unsupported
     def RAISE_VARARGS(self, instr: Instruction):
         if instr.arg == 0:
-            if not len(self.exception_stack):
+            if self.exception_stack.empty():
                 msg = ConstantVariable.wrap_literal(
                     "No active exception to reraise", self._graph
                 )
@@ -2085,7 +2087,7 @@ def raise_sot_captured_exception(
             exc_type, self._graph, DummyTracker(list(args))
         ).call_function(*args, **kwargs)
         self.exception_stack.set_current_exception(exc, self._graph)
-        raise SotCapturedExceptionFactory.get(exc_type)
+        raise SotCapturedExceptionFactory.create(exc.get_py_value())
 
 
 class OpcodeExecutor(OpcodeExecutorBase):
@@ -2244,13 +2246,13 @@ def get_compute_fn_and_update_changed_vars(
                 compile_graph_result, store_vars, store_var_info
             )
 
-    def fallback_when_block_stack_is_empty(self):
+    def fallback_when_block_stack_is_not_empty(self):
         """
         SOT currently doesn't support a non-empty block stack (related to exception handling),
         triggering a fallback.
         """
 
-        if len(self.vframe.block_stack):
+        if self.vframe.block_stack:
             raise FallbackError(
                 'SOT currently does not support a non-empty block stack, '
                 'triggering a fallback\n'
@@ -2266,7 +2268,7 @@ def _break_graph_when_if(self, result: TensorVariable, instr: Instruction):
             instr: The jump instruction.
 
         """
-        self.fallback_when_block_stack_is_empty()
+        self.fallback_when_block_stack_is_not_empty()
         self._graph.add_global_guarded_variable(result)
 
         # 1. analyse info
@@ -2426,6 +2428,7 @@ def _break_graph_when_call(
             push_n: The number of elements to be pushed onto the stack.
 
         """
+        self.fallback_when_block_stack_is_not_empty()
         self.stack = origin_stack
 
         # 1. collect infomations
@@ -2527,6 +2530,7 @@ def create_resume_fn(null_indices):
     def _break_graph_when_for_loop(
         self, iterator: VariableBase, for_iter: Instruction
     ):
+        self.fallback_when_block_stack_is_not_empty()
         # 1. find the range of loop body
         assert for_iter.jump_to is not None
         for_iter_idx = self.indexof(for_iter)
 
@@ -631,12 +631,12 @@ def dispatch_list_ne(lhs: ListVariable, rhs: ListVariable):
 )
 
 
-def register_exception(exc):
-    @Dispatcher.register_decorator(exc)
+def register_exception(exc_type: type[Exception]):
+    @Dispatcher.register_decorator(exc_type)
     def builtin_exception_dispatcher(*args) -> int:
+        exc = exc_type(*args)
         return ExceptionVariable(
             exc,
-            *args,
             graph=Dispatcher.graph,
             tracker=DummyTracker([]),
         )
@@ -1524,7 +1524,7 @@ def numpy_binary_dispatcher(var: NumPyVariable, other: NumPyVariable):
 
 
 # `operator.eq` of `ExceptionVariable` dispatch
-def exception_variable_equal(left, right):
+def exception_variable_equal(left: ExceptionVariable, right: ExceptionVariable):
     result = (left is right) or (left.get_py_value() == right.get_py_value())
     return VariableFactory.from_value(
         result,
@@ -1623,10 +1623,11 @@ def dispatch_all(var: ContainerVariable | IterVariable):
 Dispatcher.register(
     operator_exception_match,
     ("BuiltinVariable | ExceptionVariable", "BuiltinVariable | TupleVariable"),
-    lambda exc_instance, expected_exc_types: ConstantVariable.wrap_literal(
+    lambda exc_instance, expected_exc_types: ConstantVariable(
         ExceptionVariable.check_if_exception_matches(
             exc_instance, expected_exc_types
         ),
         exc_instance.graph,
+        DummyTracker([exc_instance, expected_exc_types]),
     ),
 )