Pluggable JITs and PEP-523 (frame evaluation API)

[itamaro]

This was discussed in the PyCon US 2022 open space, the purpose of this thread is to capture what was discussed (disclaimer: just the PEP-523 aspects, not the entire discussion).

@markshannon had concerns with the current state of PEP-523, which requires checking for eval_frame inefficiently in the interpreter and in specialized code.

there was general agreement that for JIT use-cases, changing PEP-523 to allow hooking in initially (and taking total control over the interpreter) and not per-function would be sufficient, and will not interfere with further optimizations in the default frame evaluator and specializations. it was pointed out that this might not be the case for debuggers. if my understanding is correct, it should be sufficient also for TorchDynamo's use-case (confirmation from @jansel needed).

looking at the 3.11 code now, it appears that some (most?) of the problematic checks are actually missing, making the current situation pretty close to what's described above. so perhaps all that's needed is to update PEP-523 to make the assumptions about at what point it can be invoked explicit?

PEP-523 users that are interested in taking over full control over the interpreter will need to plug in "early enough" (to be defined). if such users delegate back to _PyEval_EvalFrameDefault, they might lose some or all control (or fight against the default frame eval) because eval_frame is not checked inside the default eval loop and in the specializer. to guarantee control, users will need to reimplement the eval loop (likely by duplicating it in their extension and modifying it accordingly).

[jansel]

The current plan for TorchDynamo would involve dynamically enabling/disabling eval frame hooks. We would install the eval_frame hook on entry to a torch.nn.Module, then remove the eval_frame hook when you leave the outermost torch.nn.Module.

It might be possible to emulate this behavior by installing a persistent hook that checks the count of torch.nn.Module you are in and calls custom_eval_frame() if the count is nonzero and _PyEval_EvalFrameDefault() if the count is zero.

[markshannon]

I worry that bouncing back and forth from the PyTorch JIT to the interpreter is going to be really bad for performance.
The assumption is that a JIT enabled once, so the cost of doing so is unimportant.

OOI, why not use a decorator like numba?

[jansel]

For most PyTorch models python overheads are tiny compared to the large tensor operations being done inside PyTorch. The goal is maximizing flexibility, where the user code may not be fully converted, and we mix execution of unmodified python bytecode with compiled PyTorch subgraphs.

[kmod]

I realized after our discussion that in the current "Pyston's JIT as an extension module" that we're working on, it's not quite safe to use _PyEval_EvalFrameDefault after switching over to our customized frame evaluation function since we use a different opcache format than the 3.8 interpreter does. This is probably less of an issue on >=3.11, but it would be nice to make sure that we don't end up in _PyEval_EvalFrameDefault after registering the JIT, whether by returning to a frame that exists at the time of transition, or by someone calling _PyEval_EvalFrameDefault directly (though thankfully that was hidden in 3.11).

A similar-but-different problematic case is a generator that is created before changing eval_frame and then continued afterwards -- I think this will end up using the new eval on a frame that was created and executed using the default one.

I think one approach to these issues would be to have some sort of hook that gets run before any Python code is executed, but I'm not aware of an existing one like this.

[markshannon]

We need to:

• Provide an alternative API for debuggers, so that they do not need to use PEP 523; like PEP 669.
• Provide an API that suits the needs of:
  • Cinder
  • Pyston
  • CPython itself
• Provide a regular API, not just a mishmash of the above requirements.

Both Cinder and Pyston want to attach early in the process, and take over control of the interpreter.

I think the best approach is to piggy back on the instrumentation mechanism of PEP 669, treating custom eval functions like profiling. Each call would be instrumented, so that it would call the custom eval function without additional checks.

[carljm]

have some sort of hook that gets run before any Python code is executed

We just talked about this briefly at the sprints and it does seem like this would make things a lot cleaner (compared to having to swap in at an arbitrary point after some Python code has already run.) A simple way to do it would be something like -X jit=modulename where modulename is an extension module, and there's a protocol for -X jit to find an eval-frame hook in that module.

[markshannon]

Would the -X jit=modulename option do anything other than import the specified module?
If not, then it isn't jit specific at all and we should maybe add a -X preload=modulename option. Might be useful for other tools that want complete visibility, like coverage.py.

[itamaro]

Would the -X jit=modulename option do anything other than import the specified module?

I think it's possible we will want to design a specialized protocol between the interpreter and the registered JIT, so it might be more than just importing (though I don't have a sense for what that protocol might include at this point).

the -X preload=modulename idea might not apply to the JIT use case, but I think it has merit on its own, and could potentially serve as a way to start experimenting with pluggable JITs (e.g. - let's see what we can do with just preload-and-import and use this to inform the more specific protocol if one is needed). while not directly related to this discussion, my mind went to the PEP 690 (lazy imports) discussion thread, where such a flag could serve as a way to guarantee that lazy imports are enabled "early enough".

[ericsnowcurrently]

to guarantee control, users will need to reimplement the eval loop (likely by duplicating it in their extension and modifying it accordingly).

FWIW, there has been some discussion about generating the eval loop at build time. If we do that right, the same tooling could be used to generate custom eval loops outside CPython.

[itamaro]

that would be amazing

[markshannon]

I'd like to make a few concrete proposals, to see whether they suit Pyston and/or Cinder.

1. Add a -X preload= command line option to pre-load one or modules before normal execution.
2. If there are several preloaded modules, they will executed in left to right order.
3. Guarantee that if _PyInterpreterState_SetEvalFrameFunc is called in a preload module, then only the provided frame eval function will be used in subsequent preloaded modules and in the main program
4. If _PyInterpreterState_SetEvalFrameFunc is called in normal code, then the VM is free to use either the provided frame eval function or _PyEval_EvalFrameDefault for subsequent execution.
5. If a preloaded module includes Python code, or it imports Python code, and it calls _PyInterpreterState_SetEvalFrameFunc, then there are no guarantees about what eval frame function will be used in the preload module or the imported module.

I think that covers it.

Basically, implement your preload module in C and call _PyInterpreterState_SetEvalFrameFunc in it, otherwise weird stuff might happen.

[carljm]

There may be pushback from core devs about perceived overlap between preload and the existing sitecustomize.py feature. I think it will be important to highlight that preload is really intended for things that need to hook in in C before any Python code runs, which distinguishes it from sitecustomize.py.

[mpage]

One thing I'd like to explore is if it would be possible to avoid tying alternate execution engines into the frame infrastructure that's used by CPython. I think this might be possible if in addition to what is already proposed, we provided the following:

1. An API to set the default vectorcall functions that are set on PyFunctionObjects when they are created. This would provide a mechanism to allocate an alternate frame and pass it to a different interpreter loop.
2. Restore the ability to set a callback that is invoked when PyEval_GetFrame is called. This would allow us to materialize the call stack of PyFrameObjects that correspond to the alternate stack being used.
3. Add a void* to PyThreadState that could be used to store a pointer to an alternate top-of-stack. This provides a convenient place to store the top-of-stack that is pushed/popped by the alternate interpreter/JIT.

[kmod]

I'm definitely in support of providing lower-level hooks like this; Pyston wouldn't make use of it initially but it would be a nice ability to have.

Re env var, I'm just a bit worried that users might decide that theoretical-safety is not as important as being able to use their existing workflows, and decide to continue using a .pth-file approach. Just a thought but maybe an alternative to a preloading hook is to try to make late-loading safe? If there's a small enough number of potential safety issues, maybe we could just handle them directly.

[markshannon]

decide to continue using a .pth-file approach

Do you mean that they might set the environment variable in their .pth file? If -Xpreload happens first, then that won't work, so people won't do it.

maybe an alternative to a preloading hook is to try to make late-loading safe?

I don't know how to make late-loading safe without runtime overhead, and there might be more corner cases than you think.

[kmod]

We will likely offer users a choice of either updating their workflows to set the command line arg / environment variable, or to pip install a special "preload" package that uses a .pth file to install the jit with the caveat of some theoretical instability.

I know I was originally worried about late-loading, but the Pyston team thinks we've figured out all of the issues we know about with this pth-file approach, at least with how things work in 3.8. I'm curious, what issues do you think there are?

[markshannon]

I think using an environment variable makes sense, given the issue with sub-processes.
If we are going to use an environment variable, would it make sense to use only the environment variable and drop the command line option? Or would the command line option could set the environment variable, so that sub-processes automatically inherit it?

[itamaro]

+1 for command line option setting the env variable

[kmod]

Personally I would expect command-line args to not be inherited by subprocesses; I believe this would be different than how all of the other python arguments work. I don't have a strong investment in the decision, but I would be -0.5 on having the arg set the env variable.

Also I wanted to point out that an environment variable is not a panacea; for example it will not work with pyperformance (which resets all the environment variables).

[markshannon]

As you're against using environment variables, what would method would you propose for communicating with sub-processes?

[kmod]

Sorry I guess I expressed myself more negatively than intended -- I'm not "against" environment variables, I just think it's more user-friendly if there's a mechanism that can be configured once for a given environment, especially if that one step lives within existing tooling such as by being a pip-installable package.

Mechanically we can already do this using the conventional-but-disliked .pth-file approach. The downside is that .pth files get loaded after Python code has already been executed, raising some of the safety concerns that we talked about at PyCon. But after PyCon I've spent a little bit of time working on this and I am not aware of any further concerns at the moment, so I was hoping to open back up the conversation about whether early-loading is necessary. (Same caveat as before, this is based on my experience with 3.8 not 3.11)

I know it's a bit scarier to load in a JIT after the default interpreter has already run, but my thinking is that 1) maybe it's not that much work to get that to be safe, and 2) maybe users would prefer that over an environment variable.

Also, this isn't exclusive with having an environment variable that guarantees safety.

[tekknolagi]

We have some use for exposing object offsets to JITed code. For example, sometimes the JITed code would like to reach into a Python object like a listiterator. Ideally we would not expose the listiterobject struct in Include, since it's otherwise situated very nicely in Objects/listobject.c. Maybe there is generic way to expose field offsets?

[tekknolagi]

I mean, if people rally and say, "no, actually, we should expose all the structs in some non-public C header for JIT consumption", that would be fine too, but it seems like a big change.

[ambv]

I was unaware of this discussion here earlier and I'm not sure if it's still relevant, but now that I'm looking at it, I'd like to point out that -Xpresite= is now a thing. Currently it's used for reliable gathering of stdlib test coverage, and it requires a Py_DEBUG build.

To me it looks like what -Xpreload= was meant to do, as proposed in this discussion, so if this is desired, we can promote -Xpresite= to work on non-debug builds, too.

[itamaro]

To me it looks like what -Xpreload= was meant to do, as proposed in this discussion, so if this is desired, we can promote -Xpresite= to work on non-debug builds, too.

yeah, I think -Xpresite= achieves similar goals. we were debating internally recently how to initialize the Cinder JIT as an extension module early enough in the lifecycle, and I believe this mechanism would be better than the alternative (e.g. in sitecustomize.py). @jbower-fb what do you think?

[carljm]

I think this would be useful for the Cinder JIT. @ambv do you see any downsides to enabling -Xpresite= and PYTHON_PRESITE on non-debug builds, too? If not, I'd say let's do it.

[ambv]

No downside. I originally only put it for debug builds because there were concerns at the core sprint of possible user misuse.

While I'm intuitively skeptical that this can be abused, I didn't spend time on analyzing PRESITE= as a possible security vulnerability vector.

So, I'm supportive in general but I'd like to have this conversation first.

[carljm]

Makes sense.

Considering that PYTHONHOME already exists, and allows anyone able to set it to completely replace the standard library (including many modules that are always imported on startup), I don't see how PYTHON_PRESITE introduces a new attack vector.

I assume -E (and thus also -I) prevent PYTHON_PRESITE from taking effect?