Support async predictors (#2010)
* Permit async predictors
First steps towards allowing `async def predict` method signatures. This
commit adds support to Worker for starting an `asyncio` event loop if
the `predict` function returns an awaitable or an async generator.
For now, we drop support for output capture as well as cancelation.
* Async-compatible stream interception
In an async context, attempting to intercept stream writes at the file
descriptor layer is futile. We can do it, but we will have no way of
associating a write made from native code with a specific prediction --
and the only reason to intercept/swap out the STDOUT/STDERR file
descriptors is so that we can catch writes from native code.
This commit adds an altogether simpler implementation which can work for
async code with that restriction. All it does is patch `sys.stdout` and
`sys.stderr` with objects that can redirect (or tee) the output to a
callback function.
* Implement single-task cancelation for async predictors
This implements basic cancelation for async predictors. Whereas regular
predictors implement cancelation using a custom CancelationException,
asyncio already has a concept of task cancelation, so we use that.
When cancelation is requested, we send a `Cancel()` event down the
events pipe to the child. Regular predictors ignore these, but async
predictors cancel the currently-running task when they receive one.
In future, these `Cancel()` events will specify which running prediction
they are intended to cancel.
* Ensure that graceful shutdown works as expected for async predictors
When a `Shutdown()` event is sent, any running prediction should be
allowed to completed.
For now, we implement this by awaiting any task that is tracked when we
break out of the child worker's event loop.
* Update support-async-predictors branch for Pydantic v2
- Use renamed _ChildWorker type
- Set initial `__url__` to `None` prior to URL parsing potentially
throwing `ValueError`
- Declare `pid` field in `FakeChildWorker`
- Do not use nested redirectors
---------
Co-authored-by: Nick Stenning <nick@whiteink.com>
Co-authored-by: Dominic Baggott <dominic.baggott@gmail.com>
