[SOT][3.13] fix get frame.f_locals segmentation fault (#69738)

gouzil · web-flow · commit 4d99ac377bb7 · 2024-11-27T10:09:45.000+08:00
diff --git a/paddle/fluid/pybind/sot/cpython_internals.c b/paddle/fluid/pybind/sot/cpython_internals.c
@@ -167,127 +167,59 @@ void Internal_PyEvalFrameClearAndPop(PyThreadState *tstate,
 }
 
 #if PY_3_13_PLUS
-// Returns borrowed reference or NULL
-static PyObject *framelocalsproxy_getval(_PyInterpreterFrame *frame,
-                                         PyCodeObject *co,
-                                         int i) {
-  PyObject **fast = _PyFrame_GetLocalsArray(frame);
-  _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
-
-  PyObject *value = fast[i];
-  PyObject *cell = NULL;
-
-  if (value == NULL) {
-    return NULL;
+// This function is used to get the locals mapping of the frame.
+void update_framelocals_mapping(PyObject *mapping,
+                                PyCodeObject *code,
+                                int i,
+                                PyObject *value) {
+  _PyLocals_Kind kind = _PyLocals_GetKind(code->co_localspluskinds, i);
+
+  if (kind & CO_FAST_FREE && !(code->co_flags & CO_OPTIMIZED)) {
+    return;
   }
 
-  if (kind == CO_FAST_FREE || kind & CO_FAST_CELL) {
-    // The cell was set when the frame was created from
-    // the function's closure.
-    assert(PyCell_Check(value));
-    cell = value;
+  if (kind & CO_FAST_HIDDEN) {
+    return;
   }
 
-  if (cell != NULL) {
-    value = PyCell_GET(cell);
+  if (kind & CO_FAST_FREE) {
+    assert(value != NULL && PyCell_Check(value));
+    value = PyCell_GET(value);
   }
 
-  if (value == NULL) {
-    return NULL;
+  if (value != NULL) {
+    PyDict_SetItem(
+        mapping, PyTuple_GET_ITEM(code->co_localsplusnames, i), value);
   }
-
-  return value;
 }
 
-bool Internal_PyFrame_HasHiddenLocals(_PyInterpreterFrame *frame) {
-  /*
-   * This function returns if there are hidden locals introduced by PEP 709,
-   * which are the isolated fast locals for inline comprehensions
-   */
-  PyCodeObject *co = _PyFrame_GetCode(frame);
-
-  for (int i = 0; i < co->co_nlocalsplus; i++) {
-    _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
-
-    if (kind & CO_FAST_HIDDEN) {
-      PyObject *value = framelocalsproxy_getval(frame, co, i);
+// simplified version `frame_get_var`, `frame_init_get_vars` and
+// `PyFrame_GetLocals`
+PyObject *get_framelocals_mapping(_PyInterpreterFrame *frame) {
+  PyObject *mapping = PyDict_New();
 
-      if (value != NULL) {
-        return true;
-      }
-    }
+  // If the frame is not yet executed, return an empty mapping, see
+  // `frame_get_var` function
+  if (!frame->stacktop) {
+    return mapping;
   }
 
-  return false;
-}
-static PyObject *Internal_framelocalsproxy_new(PyTypeObject *type,
-                                               PyObject *args,
-                                               PyObject *kwds) {
-  if (PyTuple_GET_SIZE(args) != 1) {
-    PyErr_Format(PyExc_TypeError,
-                 "FrameLocalsProxy expected 1 argument, got %zd",
-                 PyTuple_GET_SIZE(args));
-    return NULL;
-  }
-  PyObject *item = PyTuple_GET_ITEM(args, 0);
+  PyCodeObject *co = PyFrame_GET_CODE(frame);
 
-  if (!PyFrame_Check(item)) {
-    PyErr_Format(PyExc_TypeError, "expect frame, not %T", item);
-    return NULL;
+  // Get local variables, see `frame_get_var` function
+  int offset = co->co_nlocalsplus - co->co_nfreevars;
+  for (int i = 0; i < offset; i++) {
+    update_framelocals_mapping(mapping, co, i, frame->localsplus[i]);
   }
-  PyFrameObject *frame = (PyFrameObject *)item;
 
-  if (kwds != NULL && PyDict_Size(kwds) != 0) {
-    PyErr_SetString(PyExc_TypeError,
-                    "FrameLocalsProxy takes no keyword arguments");
-    return 0;
-  }
-
-  PyFrameLocalsProxyObject *self =
-      (PyFrameLocalsProxyObject *)type->tp_alloc(type, 0);
-  if (self == NULL) {
-    return NULL;
-  }
-
-  ((PyFrameLocalsProxyObject *)self)->frame = (PyFrameObject *)Py_NewRef(frame);
-
-  return (PyObject *)self;
-}
-
-PyObject *Internal_PyFrameLocalsProxy_New(PyFrameObject *frame) {
-  PyObject *args = PyTuple_Pack(1, frame);
-  if (args == NULL) {
-    return NULL;
-  }
-
-  PyObject *proxy = (PyObject *)Internal_framelocalsproxy_new(
-      &PyFrameLocalsProxy_Type, args, NULL);
-  Py_DECREF(args);
-  return proxy;
-}
-
-PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame) {
-  // We should try to avoid creating the FrameObject if possible.
-  // So we check if the frame is a module or class level scope
-  PyCodeObject *co = _PyFrame_GetCode(frame);
-
-  if (!(co->co_flags & CO_OPTIMIZED) &&
-      !Internal_PyFrame_HasHiddenLocals(frame)) {
-    if (frame->f_locals == NULL) {
-      // We found cases when f_locals is NULL for non-optimized code.
-      // We fill the f_locals with an empty dict to avoid crash until
-      // we find the root cause.
-      frame->f_locals = PyDict_New();
-      if (frame->f_locals == NULL) {
-        return NULL;
-      }
-    }
-    return Py_NewRef(frame->f_locals);
+  // Get closure variables, see `frame_init_get_vars` function
+  PyObject *closure = ((PyFunctionObject *)frame->f_funcobj)->func_closure;
+  for (int i = 0; i < co->co_nfreevars; ++i) {
+    update_framelocals_mapping(
+        mapping, co, offset + i, PyTuple_GET_ITEM(closure, i));
   }
 
-  PyFrameObject *f = Internal_PyFrame_GetFrameObject(frame);
-
-  return Internal_PyFrameLocalsProxy_New(f);
+  return mapping;
 }
 
 #else
diff --git a/paddle/fluid/pybind/sot/cpython_internals.h b/paddle/fluid/pybind/sot/cpython_internals.h
@@ -40,7 +40,7 @@ static int Internal_PyFrame_OpAlreadyRan(_PyInterpreterFrame *frame,
                                          int opcode,
                                          int oparg);
 #if PY_3_13_PLUS
-PyObject *Internal_PyFrame_GetLocals(_PyInterpreterFrame *frame);
+PyObject *get_framelocals_mapping(_PyInterpreterFrame *frame);
 #else
 int Internal_PyFrame_FastToLocalsWithError(_PyInterpreterFrame *frame);
 #endif
diff --git a/paddle/fluid/pybind/sot/eval_frame.c b/paddle/fluid/pybind/sot/eval_frame.c
@@ -354,7 +354,7 @@ static PyObject *_custom_eval_frame(PyThreadState *tstate,
   // _PyFrame_FastToLocalsWithError directly. But this is an internal API, so we
   // copy many code from CPython project into our project.
 #if PY_3_13_PLUS
-  PyObject *f_locals = Internal_PyFrame_GetLocals(frame);
+  PyObject *f_locals = get_framelocals_mapping(frame);
   if (f_locals == NULL) {
 #else
   if (Internal_PyFrame_FastToLocalsWithError(frame) < 0) {
