Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.13k stars 69 forks source link

in generators (and async generators / coroutines, too), exception handling during yield should be done by the generator #815

Open t-vi opened 1 month ago

t-vi commented 1 month ago

Currently in generators, we catch exceptions during the yield and raise them. This is not correct, we should be setting the error and let the frame do the error handling. In other words, the current implementation ignores try blocks when the error happens during yield.

https://github.com/Lightning-AI/lightning-thunder/blob/d15b64c71aca487397c47452ac6872d853d9acda/thunder/core/interpreter.py#L5888-L5894

This is also the only occasion when Python 3.12 reaches the CLEANUP_THROW opcode.

cc @apaz-cli

tfogal commented 1 month ago

triage review:

t-vi commented 1 month ago

So what needs to be done is in the above snippet to call into run_frame but instead of executing the next instruction to immediately jump to the exception handling. Probably just an if in _run_frame and calling _run_frame in the above snippet instead of raising.