Adapting Python's `contextvars` library for use in distributed systems
Fear and Loathing in Python: Building a Distributed Context System for Wool
[](https://gist.github.com/conradbzura/885a542ff0ccd548aa16fd05525a7a71#fear-and-loathing-in-python-building-a-distributed-context-system-for-wool)
I would like to share one of the more interesting features I've implemented for Wool recently: a distributed context system with (near) complete parity to Python's native contextvars library. I undertook this endeavor thinking it'd be a neat little feature enabling Wool's users to get more creative with how they design their distributed applications. Well, it took me two attempts to get it right (I hope), and it was enough of an ordeal that I decided to put the experience to paper. Enjoy.
False start
[](https://gist.github.com/conradbzura/885a542ff0ccd548aa16fd05525a7a71#false-start)
When I first approached the challenge of extending Python's contextvars library to Wool, I'll admit I didn't have an adequate understanding of how it interfaced with asyncio—I was under the impression that newly constructed asyncio tasks always got a shallow copy of their parent context. This is evident in the fact that I chose to store worker proxies in a context var. While in the vast majority of situations my solution would have worked just fine, the real deficiency was the risk of silent corruption in the cases where it didn't.
You see, while Python provides sane defaults that keep context inheritance and mutations coherent, the language provides the capability for consenting adults to specify their own desired context when scheduling asynchronous tasks and callbacks:
import asyncio import contextvars
async def main(): ctx = contextvars.Context() loop = asyncio.get_running_loop()
A task or a callback can be handed an explicit context to
run in, instead of the default copy of the current one...
asyncio.create_task(worker(), context=ctx) loop.call_soon(callback, context=ctx)
Herein lies the problem. Imagine a developer decides to schedule a dozen concurrent coroutines using a task group—and they give each one the same context instance:
import asyncio import contextvars
async def main(): ctx = contextvars.Context()
Twelve coroutines, every one handed the same context instance...
async with asyncio.TaskGroup() as tg: for n in range(12): tg.create_task(handle(n), context=ctx)
In a vanilla, single-process event loop, the concurrent context var mutations would race and the context state of each task would reflect the latest mutation across the fan-out:
import asyncio import contextvars
var = contextvars.ContextVar("request_id", default="?")
async def handler(request_id, work): token = var.set(request_id) # claim the slot... await asyncio.sleep(work) # ...do some async work... print(f"set {request_id!r}, read back {var.get()!r}") var.reset(token) # ...release it
async def main(): ctx = contextvars.Context() # one context, shared by both await asyncio.gather( asyncio.create_task(handler("A", 0.1), context=ctx), asyncio.create_task(handler("B", 0.2), context=ctx), )
asyncio.run(main())
set 'A', read back 'B' <- B clobbered A's value mid-sleep
set 'B', read back '?' <- A's reset wiped the slot before B awoke
This is probably not a desired outcome, but this type of implementation, intentional or not, would silently corrupt Wool's operational state—crucially, the current proxy through which Wool routines are dispatched—leading to unexpected behavior without necessarily failing loudly. I considered this a pretty nasty latent bug and decided it needed fixing-by extending Python's context semantics to Wool in a way that would allow me to enforce certain invariants that are strictly necessary to make distributed execution work while maintaining functional parity with the stdlib.
What I settled on was a design that preserved parity (or so I thought) with a subset of Python's context semantics as they pertain to asyncio—namely, the default copy-on-fork context inheritance (via `asyncio.create_task` and the like) that guarantees isolated task-local state for concurrent tasks. Yield boundaries (via `await`), however, should not fork the context-we know that a caller that yields to another coroutine cannot progress until that coroutine completes, maintaining write-isolation and, thus, a coherent, deterministic task state.
Because Wool routines are solely dispatched upon `await`, Wool should treat an inter-process yield boundary like any other—the context should propagate forward and back across the wire between peers, just as it would between cooperative coroutines.
To resolve the latent state corruption bug, however, I decided that Wool must prevent users from executing coroutines concurrently within the same context-a rare exception to Wool's core "be unopinionated" principle.
This turned out to be more complicated than I thought, on a few levels, which I'll dive into now.
First, the identity of a stdlib `ContextVar` is the object instance itself (the string you pass to the constructor is purely informational)—there is no trivial way to preserve this identity perfectly across Wool's serialization boundary (a long story for another time). Here, I decided to promote a Wool var's name to be its identity, along with a package-derived namespace to prevent collisions. This was the easy part.
Second, Python doesn't expose any way to access the live context, at least not as of yet. In my first attempt at this, I decided to create a custom Wool context type that would live alongside the stdlib's. The Wool context would be forcibly forked on any new asyncio task, guaranteeing the desired task-local state isolation, and this was tracked with a global mapping of asyncio task to wool context (leaning heavily on `asyncio.current_task`), along with some trickery to determine which context serves as the parent when forking. This design worked well enough, until I realized that event loop callbacks are not associated with an asyncio task, even though they do inherit a copy of the scheduling task's context. This seemingly minor oversight turned out to have major implications, as the Wool context was effectively invisible within callbacks, and, because my design tied context identity to an asyncio task, I needed a complete redesign if I wanted parity. Woof.
Running it back
[](https://gist.github.com/conradbzura/885a542ff0ccd548aa16fd05525a7a71#running-it-back)
In thinking through the problem for a second time, I came to the realization that I didn't actually need a context type at all. What I needed was a way to reliably catch collisions, i.e., concurrent asyncio tasks sharing the same stdlib context. This check could even be narrowed down to only the contexts that have had a Wool context var set. To facilitate this, I needed to be able to identify a logical execution chain that was guaranteed to execute serially, even across workers. This was what the task-pinning in my previous attempt was for, but I knew now that that approach was not viable—I needed my chain concept to include any callbacks scheduled by the task. To this end, I settled on a task factory implementation that would mint a new chain on fork, and associate the chain with the current stdlib task and context via a context var. This has proven to be a much cleaner design, providing a few major advantages over the prior design (besides the obvious one—being correct):
- Users don't have an additional context type to consider—they can continue to work with the stdlib context system and everything just works as expected. The chain propagates on copy as it should (except when a task is created, in which case a new one is minted), and the guard only fires when concurrent tasks are executed in the same chain (via the same stdlib context):
import asyncio import contextvars import wool
foo = wool.ContextVar("foo")
async def worker(n): foo.set(n) await asyncio.sleep(0) return foo.get()
async def main():
Each create_task forks a fresh chain, so concurrent tasks stay
isolated - every worker reads back its own value...
async with asyncio.TaskGroup() as tg: tasks = [tg.create_task(worker(n)) for n in range(12)] assert sorted(t.result() for t in tasks) == list(range(12))
Force the old pathological case by scheduling concurrent tasks
in one shared context - Wool fails loudly instead of corrupting
silently...
ctx = contextvars.copy_context() a = asyncio.create_task(worker(1), context=ctx) b = asyncio.create_task(worker(2), context=ctx) await asyncio.gather(a, b) # Raises wool.ChainContention
- Wool context vars are also backed by stdlib vars, which means they behave identically, with the only difference being, of course, that the Wool variant propagates across worker boundaries and guards against distributed programming anti-patterns. Because of the stdlib backing, Wool context vars also inherit stdlib token-based reset semantics for free:
import wool
foo = wool.ContextVar("foo", default="default")
token = foo.set("override") assert foo.get() == "override"
foo.reset(token) # Restore the value from before the set assert foo.get() == "default"
As good as I felt about this approach, however, there remained a number of minor details that I had to get right. One nuanced area to support was async generators. Wool already fully supports distributed async generators with stdlib semantics, and I need to make sure adding contexts to the picture maintained parity. Here's the crux of the issue: when a distributed async generator is iterated, what really happens is that Wool dispatches the appropriate step function-i.e., `anext`, `asend`, or `athrow`-to the worker pool as a request frame. Of course, Wool cannot serialize the stdlib context in which the step is executed (and the context can differ between steps, because of course it can), so we rely on our chain concept instead.
The catch, however, is that each step executes as a distinct asyncio task on the worker, and the worker mints a fresh context for each task. This is at odds with the stdlib, because, typically, generators are iterated in the same context, and a user could decide to, for example, cache a token returned by a context var `set` operation to reuse later:
import wool
foo = wool.ContextVar("foo", default="outer")
@wool.routine async def gen(): token = foo.set("inner") # Set on one step, stash the token... yield foo.get() # "inner" foo.reset(token) # ...and reuse it on a later step yield foo.get() # "outer"
Our naive context-per-step approach would diverge from stdlib semantics, as tokens must be consumed within their spawning context, making this hypothetical impossible in Wool. The solution here was fairly simple—I opted to cache the context by chain ID in a weak-value dictionary, so that subsequent steps are executed in the same context as their predecessor.
As for running iteration steps in different contexts: it's moot (thankfully). The only idiomatic ways to schedule a coroutine with an arbitrary context all pass through the event loop's task factory, in which case the task receives a fresh Wool chain and everything works as you'd expect. Direct instantiations of `asyncio.Task` that go on to interact with a Wool context var will fail loudly via the owning-task gate. If you've dreamed up some other creative way to schedule coroutines-please, don't (at least, if you want to have a good experience with Wool). Iteration steps may still be executed in separate contexts, they'll just also be associated with distinct execution chains-perfectly valid.
Oh, and I should mention that Wool's context system is completely event-loop-agnostic, meaning third-party implementations like uvloop are fully compatible. Phew.
Where it landed
[](https://gist.github.com/conradbzura/885a542ff0ccd548aa16fd05525a7a71#where-it-landed)
After all that, here's the payoff: with one new type, `wool.ContextVar`, Wool users can utilize Python's context system across distributed workers with complete functional parity (at least, within the guardrails Wool imposes). If your current context vars satisfy the following constraints, `wool.ContextVar` will be a drop-in replacement:
- Context var values must be picklable. They have to travel across the wire, after all.
- Coroutine tasks should not be executed concurrently within the same context. I'd wager that this invariant probably holds in most asynchronous applications, as I don't see how breaking it would achieve any desirable effects.
- Context var names must be unique within their package. Remember, the package (i.e., namespace) and name together constitute a Wool var's identity, strictly necessary for deterministic round-trip serialization.
Here are some examples illustrating the behavior you can expect:
- Forward propagation, from caller to worker: A value set before dispatch is visible inside the routine, even though the routine body runs in another process.
- Back propagation, from worker to caller: A mutation made on the worker lands in the caller's context when the `await` returns.
- Propagation across generator steps: State set on one iteration survives into the next and propagates cooperatively between caller and worker.
- Propagation through nested dispatch: When a routine dispatches another routine, the chain carries the variable all the way through the call stack.
So, to reiterate, you get a distributed context emulating stdlib semantics via a narrow, declarative API surface that's a drop-in replacement for Python's `contextvars.ContextVar` type so long as the three aforementioned invariants are satisfied. This aligns with Wool's broader objective—to provide distributed execution semantics through Python's existing async primitives with minimal boilerplate and few opinions. If you're already comfortable with async Python and its context system, Wool will look and feel very familiar to you.
This little memoir only scratches the surface of the design rationale and implementation specifics of this feature—if you're interested in taking a closer look, I encourage you to check out the pull requests for both the first and second iterations on GitHub. Or, if you'd prefer, you can try it out for yourself—it's available in the latest pre-release (`pip install --pre wool`) and will go live in a week or so.
I'm actively developing Wool—fixing bugs, adding bugs, etc. To stay up-to-date, star or watch the repo on GitHub.