[SOT] Handle in_sot_simulation_mode same as psdb functions (#73332)

DrRyanHuang · web-flow · commit 5024e244ca16 · 2025-06-15T05:25:50.000+08:00
diff --git a/python/paddle/base/dygraph/base.py b/python/paddle/base/dygraph/base.py
@@ -66,10 +66,17 @@ def in_to_static_mode() -> bool:
 
 def in_sot_simulation_mode() -> bool:
     """
-    Return a bool value that indicates whether running code under SOT simulation context.
+    Returns whether the code is running under the SOT simulation context.
 
+    NOTE: Always returns False because if this function is called directly from native Python,
+    it is not within the SOT simulation process. In that case, returning False is correct.
+    If the code is running within the SOT simulation process, the function will be represented
+    by UserDefinedFunctionVariable, which is specially handled in its `call_function` method
+    to return True when this function is called.
+
+    This design avoids introducing `global_var` into the guard logic.
     """
-    return global_var._in_sot_simulation_mode_
+    return False
 
 
 # TODO(Aurelius84): Need to remove this alias after clean usage in PaddleX
@@ -119,19 +126,6 @@ def to_static_mode_guard(
         global_var._in_to_static_mode_ = original_val
 
 
-@signature_safe_contextmanager
-def sot_simulation_mode_guard(
-    is_sot_simulation: bool = True,
-) -> Generator[None, None, None]:
-    global global_var
-    original_val = global_var._in_sot_simulation_mode_
-    global_var._in_sot_simulation_mode_ = is_sot_simulation
-    try:
-        yield
-    finally:
-        global_var._in_sot_simulation_mode_ = original_val
-
-
 @signature_safe_contextmanager
 def param_guard(
     parameters: OrderedDict[str, Tensor],
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -220,7 +220,6 @@ def __init__(self):
         """
         global _dygraph_tracer_
         self._in_to_static_mode_ = False
-        self._in_sot_simulation_mode_ = False
         self._functional_dygraph_context_manager = None
         self._dygraph_tracer_ = _dygraph_tracer_
         env_pir_enabled = os.environ.get("FLAGS_enable_pir_api")
@@ -242,9 +241,6 @@ def __init__(self):
     def __str__(self):
         strings = []
         strings.append("_in_to_static_mode_:" + str(self._in_to_static_mode_))
-        strings.append(
-            "_in_sot_simulation_mode_:" + str(self._in_sot_simulation_mode_)
-        )
         strings.append(
             "_functional_dygraph_context_manager:"
             + str(self._functional_dygraph_context_manager)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -19,7 +19,6 @@
 from typing import TYPE_CHECKING
 
 import paddle
-from paddle.base.dygraph.base import sot_simulation_mode_guard
 
 from ...profiler import EventGuard, event_register
 from ...psdb import NO_FALLBACK_CODES
@@ -363,11 +362,12 @@ def start_translate(
         simulator = OpcodeExecutor(vframe, graph)
         simulator.check_code_simulatable()
         InfoCollector().attach(CompileCountInfo, frame.f_code)
-        with sot_simulation_mode_guard(True):
-            new_custom_code, guard_fn = simulator.transform(frame)
-            if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
-                assert guard_fn(frame)
-                assert guard_fn.mirror_guard(frame)
+
+        new_custom_code, guard_fn = simulator.transform(frame)
+        if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+            assert guard_fn(frame)
+            assert guard_fn.mirror_guard(frame)
+
         if not simulator._graph.need_cache:
             return (
                 CustomCode(None, True),
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -30,7 +30,10 @@
 )
 
 import paddle
-from paddle.base.dygraph.base import _DecoratorContextManager
+from paddle.base.dygraph.base import (
+    _DecoratorContextManager,
+    in_sot_simulation_mode,
+)
 
 from .... import psdb
 from ....profiler import EventGuard
@@ -277,7 +280,7 @@ def handle_psdb_function(self, /, *args, **kwargs):
                 f"Fallback by psdb.fallback (recursive={recursive_var.get_py_value()})",
                 disable_eval_frame=recursive_var.get_py_value(),
             )
-        elif self.value is psdb.in_sot:
+        elif self.value in {psdb.in_sot, in_sot_simulation_mode}:
             return ConstantVariable.wrap_literal(True, self.graph)
         return None
 
