How does torch dynamo set custom eval frame callback?

huijjj: <p>Hi, I am a machine learning engineer new to torch dynamo, just started studying.</p>
<p>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.</p>
<p>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?</p>
<p>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.</p>
<p>Simple C code to setup the callback.</p>
<pre data-code-wrap="C"><code class="lang-C">// my_eval.c
#define PY_SSIZE_T_CLEAN
#include &lt;Python.h&gt;
#include &lt;frameobject.h&gt;

// 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", &amp;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(&amp;moduledef);
}

</code></pre>
<p>Simple setup script to install the C bindings.</p>
<pre data-code-wrap="Python"><code class="lang-Python">from setuptools import setup, Extension

setup(
    name="myeval",
    ext_modules=[
        Extension(
            "myeval", 
            ["my_eval.c"],
        ),
    ],
)
</code></pre>
<p>Main python code</p>
<pre data-code-wrap="Python"><code class="lang-Python">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()
</code></pre>
<p>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.</p>
<p>Can anyone help me out with this?</p>

huijjj: <p>It seems like torch dynamo custom eval frame function only evaluates the target code (ex, function decorated with <span class="mention">@torch.compile</span>) and the custom eval frame function itself is evaluated by cpython default callback, is my understanding correct?</p>