Pick up dask's Client.loop when possible; fixes partition(timeout)

Originally, `Stream().partition(2, timeout=1.).scatter()` would result
in an error about two different event loops running. This is because the
`partition` command would require streamz to allocate an IOLoop, and
then later get the `dask.distributed.Client` event loop as well.

This patch introduces a guarded import in core to see if there is a dask
client running, and if so, uses that loop instead of starting a new one.

Also adds a regression test for async and sync behavior.

Author: wwoods

streamz/core.py

@@ -21,6 +21,11 @@
 except ImportError:
     PollIOLoop = None  # dropped in tornado 6.0
 
+try:
+    from distributed.client import default_client as _dask_default_client
+except ImportError:
+    _dask_default_client = None
+
 from collections.abc import Iterable
 
 from .compatibility import get_thread_identity
@@ -42,6 +47,15 @@ def get_io_loop(asynchronous=None):
     if asynchronous:
         return IOLoop.current()
 
+    if _dask_default_client is not None:
+        try:
+            client = _dask_default_client()
+        except ValueError:
+            # No dask client found; continue
+            pass
+        else:
+            return client.loop
+
     if not _io_loops:
         loop = IOLoop()
         thread = threading.Thread(target=loop.start)

streamz/dask.py

@@ -38,10 +38,6 @@ class DaskStream(Stream):
     --------
     dask.distributed.Client
     """
-    def __init__(self, *args, **kwargs):
-        if 'loop' not in kwargs:
-            kwargs['loop'] = default_client().loop
-        super(DaskStream, self).__init__(*args, **kwargs)
 
 
 @DaskStream.register_api()

streamz/tests/test_dask.py

@@ -7,7 +7,7 @@
 from tornado import gen
 
 from streamz.dask import scatter
-from streamz import Stream
+from streamz import RefCounter, Stream
 
 from distributed import Future, Client
 from distributed.utils import sync
@@ -50,6 +50,47 @@ def add_to_dict(d):
         assert item["i"] == i
 
 
+@gen_cluster(client=True)
+def test_partition_then_scatter_async(c, s, a, b):
+    # Ensure partition w/ timeout before scatter works correctly for
+    # asynchronous
+    start = time.monotonic()
+    source = Stream(asynchronous=True)
+
+    L = source.partition(2, timeout=.1).scatter().map(
+            lambda x: [xx+1 for xx in x]).buffer(2).gather().flatten().sink_to_list()
+
+    rc = RefCounter(loop=source.loop)
+    for i in range(3):
+        yield source.emit(i, metadata=[{'ref': rc}])
+
+    while rc.count != 0 and time.monotonic() - start < 1.:
+        yield gen.sleep(1e-2)
+
+    assert L == [1, 2, 3]
+
+
+@pytest.mark.slow
+def test_partition_then_scatter_sync(loop):
+    # Ensure partition w/ timeout before scatter works correctly for synchronous
+    with cluster() as (s, [a, b]):
+        with Client(s['address'], loop=loop) as client:  # noqa: F841
+            start = time.monotonic()
+            source = Stream()
+            L = source.partition(2, timeout=.1).scatter().map(
+                    lambda x: [xx+1 for xx in x]).gather().flatten().sink_to_list()
+            assert source.loop is client.loop
+
+            rc = RefCounter()
+            for i in range(3):
+                source.emit(i, metadata=[{'ref': rc}])
+
+            while rc.count != 0 and time.monotonic() - start < 2.:
+                time.sleep(1e-2)
+
+            assert L == [1, 2, 3]
+
+
 @gen_cluster(client=True)
 def test_scan(c, s, a, b):
     source = Stream(asynchronous=True)