[SOT] Attach need breakgraph info to some exceptions (#70904)

SigureMo · web-flow · commit 76f9a12bb7c3 · 2025-01-21T09:00:46.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -52,7 +52,7 @@
     map_if,
     switch_symbol_registry,
 )
-from ...utils.exceptions import BreakGraphError
+from ...utils.exceptions import BreakGraphError, SotExtraInfo
 from ..instruction_utils import get_instructions
 from .guard import Guard, StringifiedExpression, make_guard
 from .mutable_data import MutationDel, MutationNew, MutationSet
@@ -703,6 +703,13 @@ def try_infer_meta_fn(args, kwargs) -> Any:
                 kwmetas = convert_to_meta(kwargs)
                 return args, kwargs, infer_meta_fn(func, *metas, **kwmetas)
 
+            except Exception as e:
+                if SotExtraInfo.from_exception(e).need_breakgraph:
+                    raise BreakGraphError(
+                        f"API {func} encountered a need break graph error {e}"
+                    )
+                raise e
+
         if ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
             args, kwargs, out_metas = try_infer_meta_fn(args, kwargs)
         else:
diff --git a/python/paddle/jit/sot/utils/exceptions.py b/python/paddle/jit/sot/utils/exceptions.py
@@ -69,3 +69,28 @@ def impl(*args, **kwargs):
 
 class ExportError(SotErrorBase):
     pass
+
+
+class SotExtraInfo:
+    SOT_EXTRA_INFO_ATTR_NAME = "__SOT_EXTRA_INFO__"
+
+    def __init__(self, *, need_breakgraph: bool = False):
+        self.need_breakgraph = need_breakgraph
+
+    def set_need_breakgraph(self, need_breakgraph: bool):
+        self.need_breakgraph = need_breakgraph
+
+    def attach(self, err: BaseException):
+        setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, self)
+
+    @staticmethod
+    def default() -> SotExtraInfo:
+        return SotExtraInfo()
+
+    @staticmethod
+    def from_exception(err: BaseException) -> SotExtraInfo:
+        info = getattr(
+            err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, SotExtraInfo.default()
+        )
+        setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, info)
+        return info
diff --git a/python/paddle/tensor/creation.py b/python/paddle/tensor/creation.py
@@ -2780,10 +2780,15 @@ def convert_scalar(x):
         value_name = "values"
         values = input.ravel().tolist()
         if input.size > 1024 * 1024:
-            raise ValueError(
+            from paddle.jit.sot.utils.exceptions import SotExtraInfo
+
+            sot_extra_info = SotExtraInfo(need_breakgraph=True)
+            err = ValueError(
                 "The size of input is too big. Please consider "
                 "saving it to file and 'load_op' to load it"
             )
+            sot_extra_info.attach(err)
+            raise err
         if in_dynamic_or_pir_mode():
             if output is None:
                 output = zeros(list(input.shape), dtype)
diff --git a/test/sot/test_numpy.py b/test/sot/test_numpy.py
@@ -31,6 +31,10 @@ def tensor_add_numpy(x, y):
     return ret
 
 
+def large_numpy_array_to_tensor(x):
+    return paddle.to_tensor(x)
+
+
 class TestNumpy(TestCaseBase):
     @strict_mode_guard(False)
     def test_numpy_add(self):
@@ -51,6 +55,12 @@ def test_tensor_add_numpy_array(self):
         self.assert_results(tensor_add_numpy, x, y)
         self.assert_results(tensor_add_numpy, y, x)
 
+    def test_large_numpy_array_to_tensor(self):
+        # size should be larger than 1024*1024, because we throw an exception
+        # when the size is larger than 1024*1024 in assign API (to_tensor static branch)
+        x = np.random.rand(1024, 1024, 2).astype(np.float32)
+        self.assert_results(large_numpy_array_to_tensor, x)
+
 
 if __name__ == "__main__":
     unittest.main()
