Client.get_dataset to always create Futures attached to itself

.github/workflows/ci-windows.yaml

@@ -23,6 +23,13 @@ jobs:
           activate-environment: testenv
           auto-activate-base: false
 
+      - name: Install contextvars
+        shell: bash -l {0}
+        run: |
+          if [[ "${{ matrix.python-version }}" = "3.6" ]]; then
+            conda install -c conda-forge contextvars
+          fi
+
       - name: Install tornado
         shell: bash -l {0}
         run: |

continuous_integration/travis/install.sh

@@ -55,6 +55,10 @@ conda install -c conda-forge -q \
     zstandard \
     $PACKAGES
 
+if [[ $PYTHON == 3.6 ]]; then
+  conda install -c conda-forge -c defaults contextvars
+fi
+
 # stacktrace is not currently avaiable for Python 3.8.
 # Remove the version check block below when it is avaiable.
 if [[ $PYTHON != 3.8 ]]; then

distributed/client.py

@@ -5,6 +5,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures._base import DoneAndNotDoneFutures
 from contextlib import contextmanager
+from contextvars import ContextVar
 import copy
 import errno
 from functools import partial
@@ -90,6 +91,7 @@
 _global_clients = weakref.WeakValueDictionary()
 _global_client_index = [0]
 
+_deserialize_client = ContextVar("_deserialize_client", default=None)
 
 DEFAULT_EXTENSIONS = [PubSubClientExtension]
 
@@ -163,7 +165,7 @@ def __init__(self, key, client=None, inform=True, state=None):
         self.key = key
         self._cleared = False
         tkey = tokey(key)
-        self.client = client or _get_global_client()
+        self.client = client or _deserialize_client.get() or _get_global_client()
         self.client._inc_ref(tkey)
         self._generation = self.client.generation
 
@@ -354,11 +356,11 @@ def release(self, _in_destructor=False):
                 pass  # Shutting down, add_callback may be None
 
     def __getstate__(self):
-        return (self.key, self.client.scheduler.address)
+        return self.key, self.client.scheduler.address
 
     def __setstate__(self, state):
         key, address = state
-        c = get_client(address)
+        c = _deserialize_client.get() or get_client(address)
         Future.__init__(self, key, c)
         c._send_to_scheduler(
             {
@@ -2175,8 +2177,7 @@ def retry(self, futures, asynchronous=None):
         """
         return self.sync(self._retry, futures, asynchronous=asynchronous)
 
-    @gen.coroutine
-    def _publish_dataset(self, *args, name=None, **kwargs):
+    async def _publish_dataset(self, *args, name=None, **kwargs):
         with log_errors():
             coroutines = []
 
@@ -2202,7 +2203,7 @@ def add_coro(name, data):
             for name, data in kwargs.items():
                 add_coro(name, data)
 
-            yield coroutines
+            await asyncio.gather(*coroutines)
 
     def publish_dataset(self, *args, **kwargs):
         """
@@ -2282,13 +2283,30 @@ def list_datasets(self, **kwargs):
         return self.sync(self.scheduler.publish_list, **kwargs)
 
     async def _get_dataset(self, name):
-        out = await self.scheduler.publish_get(name=name, client=self.id)
-        if out is None:
-            raise KeyError("Dataset '%s' not found" % name)
+        if sys.version_info >= (3, 7):
+            # Insulate contextvars change with a task
+            async def _():
+                _deserialize_client.set(self)
+                return await self.scheduler.publish_get(name=name, client=self.id)
+
+            out = await asyncio.create_task(_())
+        else:
+            # Python 3.6; creating a task doesn't copy the context.
+            # We can still detect a race condition though.
+            if _deserialize_client.get() not in (self, None):
+                raise RuntimeError(  # pragma: nocover
+                    "Detected race condition where get_dataset() is invoked in "
+                    "parallel by multiple clients. Please upgrade to Python 3.7+."
+                )
+            tok = _deserialize_client.set(self)
+            try:
+                out = await self.scheduler.publish_get(name=name, client=self.id)
+            finally:
+                _deserialize_client.reset(tok)
 
-        with temp_default_client(self):
-            data = out["data"]
-        return data
+        if out is None:
+            raise KeyError(f"Dataset '{name}' not found")
+        return out["data"]
 
     def get_dataset(self, name, **kwargs):
         """

distributed/publish.py

@@ -6,10 +6,10 @@
 class PublishExtension:
     """ An extension for the scheduler to manage collections
 
-    *  publish-list
-    *  publish-put
-    *  publish-get
-    *  publish-delete
+    *  publish_list
+    *  publish_put
+    *  publish_get
+    *  publish_delete
     """
 
     def __init__(self, scheduler):
@@ -59,21 +59,60 @@ class Datasets(MutableMapping):
 
     """
 
+    __slots__ = ("_client",)
+
     def __init__(self, client):
-        self.__client = client
+        self._client = client
 
     def __getitem__(self, key):
-        return self.__client.get_dataset(key)
+        # When client is asynchronous, it returns a coroutine
+        return self._client.get_dataset(key)
 
     def __setitem__(self, key, value):
-        self.__client.publish_dataset(value, name=key)
+        if self._client.asynchronous:
+            # 'await obj[key] = value' is not supported by Python as of 3.8
+            raise TypeError(
+                "Can't use 'client.datasets[name] = value' when client is "
+                "asynchronous; please use 'client.publish_dataset(name=value)' instead"
+            )
+        self._client.publish_dataset(value, name=key)
 
     def __delitem__(self, key):
-        self.__client.unpublish_dataset(key)
+        if self._client.asynchronous:
+            # 'await del obj[key]' is not supported by Python as of 3.8
+            raise TypeError(
+                "Can't use 'del client.datasets[name]' when client is asynchronous; "
+                "please use 'client.unpublish_dataset(name)' instead"
+            )
+        return self._client.unpublish_dataset(key)
 
     def __iter__(self):
-        for key in self.__client.list_datasets():
+        if self._client.asynchronous:
+            raise TypeError(
+                "Can't invoke iter() or 'for' on client.datasets when client is "
+                "asynchronous; use 'async for' instead"
+            )
+        for key in self._client.list_datasets():
             yield key
 
+    def __aiter__(self):
+        if not self._client.asynchronous:
+            raise TypeError(
+                "Can't invoke 'async for' on client.datasets when client is "
+                "synchronous; use iter() or 'for' instead"
+            )
+
+        async def _():
+            for key in await self._client.list_datasets():
+                yield key
+
+        return _()
+
     def __len__(self):
-        return len(self.__client.list_datasets())
+        if self._client.asynchronous:
+            # 'await len(obj)' is not supported by Python as of 3.8
+            raise TypeError(
+                "Can't use 'len(client.datasets)' when client is asynchronous; "
+                "please use 'len(await client.list_datasets())' instead"
+            )
+        return len(self._client.list_datasets())

distributed/tests/test_publish.py

@@ -213,6 +213,23 @@ def test_datasets_iter(client):
     client.publish_dataset(**{str(key): key for key in keys})
     for n, key in enumerate(client.datasets):
         assert key == str(n)
+    with pytest.raises(TypeError):
+        client.datasets.__aiter__()
+
+
+@gen_cluster(client=True)
+async def test_datasets_async(c, s, a, b):
+    await c.publish_dataset(foo=1, bar=2)
+    assert await c.datasets["foo"] == 1
+    assert {k async for k in c.datasets} == {"foo", "bar"}
+    with pytest.raises(TypeError):
+        c.datasets["baz"] = 3
+    with pytest.raises(TypeError):
+        del c.datasets["foo"]
+    with pytest.raises(TypeError):
+        next(iter(c.datasets))
+    with pytest.raises(TypeError):
+        len(c.datasets)
 
 
 @gen_cluster(client=True)
@@ -229,3 +246,34 @@ async def test_pickle_safe(c, s, a, b):
 
         with pytest.raises(TypeError):
             await c2.get_dataset("z")
+
+
+@gen_cluster(client=True)
+async def test_deserialize_client(c, s, a, b):
+    """Test that the client attached to Futures returned by Client.get_dataset is always
+    the instance of the client that invoked the method.
+    Specifically:
+
+    - when the client is defined by hostname, test that it is not accidentally
+      reinitialised by IP;
+    - when multiple clients are connected to the same scheduler, test that they don't
+      interfere with each other.
+
+    See: https://github.com/dask/distributed/issues/3227
+    """
+    future = await c.scatter("123")
+    await c.publish_dataset(foo=future)
+    future = await c.get_dataset("foo")
+    assert future.client is c
+
+    for addr in (s.address, "localhost:" + s.address.split(":")[-1]):
+        async with Client(addr, asynchronous=True) as c2:
+            future = await c.get_dataset("foo")
+            assert future.client is c
+            future = await c2.get_dataset("foo")
+            assert future.client is c2
+
+    # Ensure cleanup
+    from distributed.client import _deserialize_client
+
+    assert _deserialize_client.get() is None

distributed/tests/test_steal.py

@@ -106,6 +106,7 @@ async def test_worksteal_many_thieves(c, s, *workers):
     assert sum(map(len, s.has_what.values())) < 150
 
 
+@pytest.mark.xfail(reason="GH#3574")
 @gen_cluster(client=True, nthreads=[("127.0.0.1", 1)] * 2)
 async def test_dont_steal_unknown_functions(c, s, a, b):
     futures = c.map(inc, range(100), workers=a.address, allow_other_workers=True)

requirements.txt

@@ -1,5 +1,6 @@
 click >= 6.6
 cloudpickle >= 0.2.2
+contextvars;python_version<'3.7'
 dask >= 2.9.0
 msgpack >= 0.6.0
 psutil >= 5.0