Proposal: Add a High-Level Frame Hook API to CPython

Ideas
1

Proposal: Add a High-Level Frame Hook API to CPython

Hi everyone,

I would like to start a discussion to propose a higher-level alternative API to PEP 523 for intercepting and modifying python frames prior to evaluation.

This proposal was developed in collaboration with William Wen, PyTorch core dev.

Disclaimer: Due to posting limits on new accounts, I’ve collected all links in the following gist:

Motivation

This proposal is motivated by the needs of PyTorch’s symbolic interpreter TorchDynamo (Dynamo). Dynamo is a bytecode-to-bytecode transpiler that does not execute frames itself - it inspects a frame’s bytecode, performs optimizations, and generates new bytecode that is then executed by the interpreter.

Dynamo currently relies on the eval frame API (PEP 523) to intercept and capture frame execution. Although PEP 523 is sufficient for implementing Dynamo, it exposes many low-level CPython implementation details that we would not like to excessively depend on. Supporting new Python versions has historically required substantial engineering effort, with changes to the internal frame evaluation mechanism being a significant source of difficulty.

We documented the steps needed to upgrade Dynamo for each new Python version from 3.11 to 3.13 here:

  • Supporting Dynamo in Python 3.11 - NULL
  • Supporting Dynamo in Python 3.11 - CPython Frame Evaluation
  • Supporting Dynamo in Python 3.12
  • Torch.compile support for Python 3.13 completed

In particular, there were some APIs that became private, such as frame creation/deletion and fastlocal buffer conversion to dict. However, these APIs are critical for our use case. To circumvent this, we had to copy-paste segments of CPython’s source code into our codebase, which is not ideal.

Some PEP 523-related issues we encountered included:

  • Formerly public APIs going private or being removed, for example, frame creation/deletion and fastlocal buffer conversion to dict
  • The eval frame hook was given the responsibility to free the frame argument

Part of our workarounds included copy-pasting parts of the CPython codebase into Dynamo, which is obviously not ideal.

To better address these issues, we propose adding a new high-level API to CPython that allows users to register frame hooks without needing to manage the low-level details of frame creation, management, and cleanup. This API would abstract away the complexities of CPython’s internal frame evaluation and provide a more stable interface for tools that need to intercept and customize frame execution.

Relationship to PEP 523

PEP 523 is intended to replace the interpreter’s frame evaluation function. This transfers to the new function the full responsibility for the frame lifecycle. In practice, this requires detailed knowledge of CPython internals and access to internal APIs that are hidden to the public.

The proposed API, on the other hand, aims to provide a structured way to register hooks that can modify the frame before it is executed. In this new model, a tool like torch.compile would only need to implement the hook function using public APIs, and the runtime would take care of calling the hooks and managing the frame lifecycle.

Proposal Overview

Registering a Frame Hook

This proposal introduces a new API that allows extensions to register a frame hook in a similar way to PEP 523. The frame hook function receives the frame about to be executed and may return a replacement code object:

typedef PyCodeObject* (*_PyFrameHookFunction)(struct _PyInterpreterFrame *);
PyUnstable_AddFrameHook(PyInterpreterState *interp, _PyFrameHookFunction hook_fn);
PyUnstable_RemoveFrameHook(PyInterpreterState *interp, _PyFrameHookFunction hook_fn);

This API allows users to register hooks that can modify a frame before it is executed without managing the frame lifecycle directly. The hooks would be called in the order they were registered, and each hook would have the opportunity to modify the frame before it is executed. This allows different tools to compose their behavior by registering independent hooks.

Example Usage

A simplified example of the frame hook function is shown below:

def frame_hook(frame: FrameType) -> types.CodeType | None:
    if not should_optimize(frame):
        return None

    if has_torch_tensor_in_frame(frame):
        # Replace the frame's code with optimized bytecode
        optimized_code: types.CodeType = dynamo.convert_frame(frame)
        return optimized_code

    return None

Interpreter Integration (_PyEval_*)

Internally, the interpreter checks for registered frame hooks before evaluating a frame. A conceptual sketch of the proposed API (_PyEval_FrameHook) might look as follows.

def _PyEval_FrameHook(tstate, frame, throw_flag):
    for hook in tstate.frame_hooks:
        # NOTE: this disabled field may or may not be needed by torch.compile
        if hook.disabled:
            continue
        # Hook is a function that takes a frame and returns a new code object
        code = hook(frame)        old_code = frame.f_code
        if code is not None and code != old_code:
            free_frame(frame)
            frame = new_frame_with_code(code)
        # else, hook did not modify the frame
    # 3.12+: the eval_frame function is responsible for freeing frame
    if tstate->interp->eval_frame == NULL:
        return _PyEval_EvalFrameDefault(tstate, frame, throwflag)
    return tstate->interp->eval_frame(tstate, frame, throw_flag)

This code is written in Python for the sake of simplicity, but the actual implementation is in C.

The hook interface described here is intentionally minimal. A hook receives a frame about to be executed and may return a replacement code object. For the purposes of this initial proposal, we assume that the original and new code objects have compatible localsplus layouts. In other words, the number and ordering of local, cell, and free variable slots must match between the two code objects. This simplifies the implementation and avoids the need to remap frame locals. With future discussions, we can extend the proposal to relax this restriction.

The _PyEval_EvalFrame function can be modified to check for a frame hook before evaluating the frame:

static inline PyObject*
_PyEval_EvalFrame(PyThreadState *tstate, _PyInterpreterFrame *frame, int throwflag)
{
    EVAL_CALL_STAT_INC(EVAL_CALL_TOTAL);
+
+   if (_PyInterpreterState_HasFrameHooks(tstate->interp) &&            
+           tstate->interp->enable_frame_hooks) {
+       return _PyEval_FrameHook(tstate, frame, throwflag);
+   }
+
    if (tstate->interp->eval_frame == NULL) {
        return _PyEval_EvalFrameDefault(tstate, frame, throwflag);
    }
    return tstate->interp->eval_frame(tstate, frame, throwflag);
}

Implementation

We have a prototype for the frame hook API can be found in the following links:

  • CPython
  • PyTorch
  • Frame Hook Example

Please, see the links in the disclaimer above.

Conclusion

The goal of this post is to gather feedback from the community on the proposed idea, its design, and its potential impact on CPython. We are open to suggestions and improvements to the proposed API and look forward to a constructive discussion.

4 Likes