Preserve contextvars during comm offload

Helps with setting the current client in worker while deserializing. Implementation referenced from python/cpython#9688

Author: gjoseph92

distributed/tests/test_utils.py

@@ -1,5 +1,6 @@
 import array
 import asyncio
+import contextvars
 import functools
 import io
 import os
@@ -554,6 +555,22 @@ async def test_offload():
     assert (await offload(lambda x, y: x + y, 1, y=2)) == 3
 
 
+@pytest.mark.asyncio
+async def test_offload_preserves_contextvars():
+    var = contextvars.ContextVar("var", default="foo")
+
+    def change_var():
+        var.set("bar")
+        return var.get()
+
+    o1 = offload(var.get)
+    o2 = offload(change_var)
+
+    r1, r2 = await asyncio.gather(o1, o2)
+    assert (r1, r2) == ("foo", "bar")
+    assert var.get() == "foo"
+
+
 def test_serialize_for_cli_deprecated():
     with pytest.warns(FutureWarning, match="serialize_for_cli is deprecated"):
         from distributed.utils import serialize_for_cli

distributed/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
 import functools
 import importlib
 import inspect
@@ -1322,7 +1323,11 @@ def import_term(name: str):
 
 async def offload(fn, *args, **kwargs):
     loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(_offload_executor, lambda: fn(*args, **kwargs))
+    # Retain context vars while deserializing; see https://bugs.python.org/issue34014
+    context = contextvars.copy_context()
+    return await loop.run_in_executor(
+        _offload_executor, lambda: context.run(fn, *args, **kwargs)
+    )
 
 
 class EmptyContext: