Use ContextVar for current worker #5517
Conversation
Progress towards dask#5485. This was surprisingly easy and actually seems to work. Not sure yet what it breaks. I was pleasantly surprised to find that Tornado `loop.add_callback`s, `PeriodicCallback`s, etc. all play correctly with asyncio's [built-in contextvar management](https://www.python.org/dev/peps/pep-0567/#asyncio). With that, just setting the contextvar during `__init__` and `start` probably catches almost all cases, because all the long-running callbacks/coroutines (including comms) will inherit the context that's set when they're created. Where else should we add this `as_current_worker` decorator? This gives me confidence we'll be able to use the same pattern for a single current client contextvar as mentioned in dask#5467 (comment).
| "current_worker" | ||
| ) | ||
|
|
||
| P = ParamSpec("P") |
There was a problem hiding this comment.
I was recently burned by ParamSpec in that not all ecosystem libraries supported typing_extensions~=3.10 (in my case it was tensorflow). That may have improved in the last few months, but something to consider
There was a problem hiding this comment.
Yeah, trying to get ParamSpec working probably took 4x longer than actually implementing the content of this PR. I'm just really excited about it.
Since mypy doesn't even support it properly, there's not much value in it here yet. I had thought we could leave it around and remove the # type: ignores in the future, but I probably should just take it out before merging this.
There was a problem hiding this comment.
I'm also excited about ParamSpec, and would like to use it for every decorator I see
Progress towards #5485.
This was surprisingly simple and actually seems to work. Not sure yet what it breaks.
I was pleasantly surprised to find that Tornado
loop.add_callbacks,PeriodicCallbacks, etc. all play correctly with asyncio's built-in contextvar management. With that, just setting the contextvar during__init__andstartprobably catches almost all cases, because all the long-running callbacks/coroutines (including comms) will inherit the context that's set when they're created. Where else should we add thisas_current_workerdecorator?This gives me confidence we'll be able to use the same pattern for a single current client contextvar as mentioned in #5467 (comment).
cc @crusaderky @fjetter
threading.local()with ContextVar #5485pre-commit run --all-files