Hi, I am a machine learning engineer new to torch dynamo, just started studying.
After watching the deep dive tutorial youtube video, I figured out some facts that torch dynamo internally utilizes cpython api (_PyInterpreterState_SetEvalFrameFunc) to hijack the normal python execution flow and modify bytecodes before actually running them. I also found out that the callback itself is written in python.
I got confused and came up with following questions: Python code run by custom interpreter which is also written in python? Than how the python codes for the custom interpreter is interpreted?
To clearly understand what is going on inside the hood, especially with the custom cpython eval frame part, I wrote some codes to mimic the behavior of torch dynamo.
Simple C code to setup the callback.
// my_eval.c
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <frameobject.h>
// This is CPython-internal; must match the interpreter layout
typedef PyObject *(*_PyFrameEvalFunction)(PyThreadState *, PyFrameObject *, int);
extern void _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState *interp, _PyFrameEvalFunction eval_frame);
extern PyInterpreterState *PyInterpreterState_Get(void);
static PyObject* g_custom_eval_func = NULL;
static PyObject* my_eval_frame(PyThreadState *tstate, PyFrameObject *frame, int throwflag) {
if (!g_custom_eval_func) return NULL;
PyObject *args = Py_BuildValue("(Oi)", frame, throwflag);
if (!args) return NULL;
PyObject *result = PyObject_CallObject(g_custom_eval_func, args);
Py_DECREF(args);
return result;
}
static PyObject* set_eval_frame(PyObject *self, PyObject *args) {
PyObject *py_callback;
if (!PyArg_ParseTuple(args, "O:set_eval_frame", &py_callback))
return NULL;
if (!PyCallable_Check(py_callback)) {
PyErr_SetString(PyExc_TypeError, "Expected a callable");
return NULL;
}
Py_XINCREF(py_callback);
Py_XDECREF(g_custom_eval_func);
g_custom_eval_func = py_callback;
PyInterpreterState *interp = PyInterpreterState_Get();
_PyInterpreterState_SetEvalFrameFunc(interp, my_eval_frame);
Py_RETURN_NONE;
}
static PyMethodDef methods[] = {
{"set_eval_frame", set_eval_frame, METH_VARARGS, "Set the custom eval frame callback"},
{NULL, NULL, 0, NULL}
};
static struct PyModuleDef moduledef = {
PyModuleDef_HEAD_INIT,
"myeval",
NULL,
-1,
methods
};
PyMODINIT_FUNC PyInit_myeval(void) {
return PyModule_Create(&moduledef);
}
Simple setup script to install the C bindings.
from setuptools import setup, Extension
setup(
name="myeval",
ext_modules=[
Extension(
"myeval",
["my_eval.c"],
),
],
)
Main python code
import myeval
def custom_eval(frame, throw_flag):
print("Intercepted frame:", frame.f_code.co_name)
return None
def test():
x = 1 + 2
return x
myeval.set_eval_frame(custom_eval)
test()
Unfortunately, I couldn’t quite get this running. my_eval_frame gets called repeatedly without actually running any code until the code crashes with segmentation fault.
Can anyone help me out with this?