Merge pull request #401 from wwoods/master2streamz

Test showing synchronous partition(timeout) + dask scatter is broken

Author: martindurant

setup.cfg

@@ -39,3 +39,8 @@ max-line-length = 120
 
 [bdist_wheel]
 universal=1
+
+[tool:pytest]
+markers:
+  network: Test requires an internet connection
+  slow: Skipped unless --runslow passed

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:  # pragma: no cover
+    _dask_default_client = None
+
 from collections.abc import Iterable
 
 from threading import get_ident as 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)
@@ -163,6 +177,7 @@ class Stream(object):
 
     def __init__(self, upstream=None, upstreams=None, stream_name=None,
                  loop=None, asynchronous=None, ensure_io_loop=False):
+        self.name = stream_name
         self.downstreams = OrderedWeakrefSet()
         if upstreams is not None:
             self.upstreams = list(upstreams)
@@ -182,8 +197,6 @@ def __init__(self, upstream=None, upstreams=None, stream_name=None,
             if upstream:
                 upstream.downstreams.add(self)
 
-        self.name = stream_name
-
     def _set_loop(self, loop):
         self.loop = None
         if loop is not None:
@@ -343,8 +356,6 @@ def __str__(self):
                     s = str(at)
                 elif hasattr(at, '__name__'):
                     s = getattr(self, m).__name__
-                elif hasattr(at.__class__, '__name__'):
-                    s = getattr(self, m).__class__.__name__
                 else:
                     s = None
             if s:
@@ -994,7 +1005,8 @@ def __init__(self, upstream, n, timeout=None, key=None, **kwargs):
         self._buffer = defaultdict(lambda: [])
         self._metadata_buffer = defaultdict(lambda: [])
         self._callbacks = {}
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
 
     def _get_key(self, x):
         if self._key is None:
@@ -1206,7 +1218,8 @@ def __init__(self, upstream, interval, **kwargs):
         self.metadata_buffer = []
         self.last = gen.moment
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
 
         self.loop.add_callback(self.cb)
 
@@ -1308,7 +1321,8 @@ def __init__(
         self._buffer = {}
         self._metadata_buffer = {}
         self.last = gen.moment
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
         self.loop.add_callback(self.cb)
 
     def _get_key(self, x):
@@ -1355,7 +1369,8 @@ def __init__(self, upstream, interval, **kwargs):
         self.interval = convert_interval(interval)
         self.queue = Queue()
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream,**kwargs)
 
         self.loop.add_callback(self.cb)
 
@@ -1393,7 +1408,8 @@ def __init__(self, upstream, interval, **kwargs):
         self.interval = convert_interval(interval)
         self.next = 0
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
 
     @gen.coroutine
     def update(self, x, who=None, metadata=None):
@@ -1418,7 +1434,8 @@ class buffer(Stream):
     def __init__(self, upstream, n, **kwargs):
         self.queue = Queue(maxsize=n)
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
 
         self.loop.add_callback(self.cb)
 
@@ -1862,7 +1879,8 @@ def __init__(self, upstream, **kwargs):
         self.next = []
         self.next_metadata = None
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
 
         self.loop.add_callback(self.cb)
 
@@ -1918,7 +1936,8 @@ def __init__(self, upstream, topic, producer_config, **kwargs):
         self.topic = topic
         self.producer = ck.Producer(producer_config)
 
-        Stream.__init__(self, upstream, ensure_io_loop=True, **kwargs)
+        kwargs["ensure_io_loop"] = True
+        Stream.__init__(self, upstream, **kwargs)
         self.stopped = False
         self.polltime = 0.2
         self.loop.add_callback(self.poll)

streamz/dask.py

@@ -38,9 +38,8 @@ 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)
+        kwargs["ensure_io_loop"] = True
+        super().__init__(*args, **kwargs)
 
 
 @DaskStream.register_api()

streamz/dataframe/tests/test_dataframes.py

@@ -11,7 +11,7 @@
 from tornado import gen
 
 from streamz import Stream
-from streamz.utils_test import gen_test
+from streamz.utils_test import gen_test, wait_for
 from streamz.dataframe import DataFrame, Series, DataFrames, Aggregation
 import streamz.dataframe as sd
 from streamz.dask import DaskStream
@@ -231,6 +231,8 @@ def test_index(stream):
     a.emit(df)
     a.emit(df)
 
+    wait_for(lambda: len(L) > 1, timeout=2, period=0.05)
+
     assert_eq(L[0], df.index + 5)
     assert_eq(L[1], df.index + 5)
 

streamz/tests/test_core.py

@@ -403,10 +403,11 @@ def test_timed_window():
 
 @gen_test()
 def test_timed_window_ref_counts():
-    source = Stream()
+    source = Stream(asynchronous=True)
     _ = source.timed_window(0.01)
 
     ref1 = RefCounter()
+    assert str(ref1) == "<RefCounter count=0>"
     source.emit(1, metadata=[{'ref': ref1}])
     assert ref1.count == 1
     yield gen.sleep(0.05)
@@ -417,6 +418,13 @@ def test_timed_window_ref_counts():
     assert ref2.count == 1
 
 
+def test_mixed_async():
+    s1 = Stream(asynchronous=False)
+    with pytest.raises(ValueError):
+        Stream(asynchronous=True, upstream=s1)
+
+
+
 @gen_test()
 def test_timed_window_metadata():
     source = Stream()

streamz/tests/test_dask.py

@@ -8,7 +8,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
@@ -51,6 +51,46 @@ 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]
+
+
+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_non_unique_emit(c, s, a, b):
     """Regression for https://github.com/python-streamz/streams/issues/397
@@ -123,7 +163,6 @@ def test_accumulate(c, s, a, b):
     assert L[-1][1] == 3
 
 
-@pytest.mark.slow
 def test_sync(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as client:  # noqa: F841
@@ -140,7 +179,6 @@ def f():
             assert L == list(map(inc, range(10)))
 
 
-@pytest.mark.slow
 def test_sync_2(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop):  # noqa: F841
@@ -154,8 +192,7 @@ def test_sync_2(loop):  # noqa: F811
             assert L == list(map(inc, range(10)))
 
 
-@pytest.mark.slow
-@gen_cluster(client=True, ncores=[('127.0.0.1', 1)] * 2)
+@gen_cluster(client=True, nthreads=[('127.0.0.1', 1)] * 2)
 def test_buffer(c, s, a, b):
     source = Stream(asynchronous=True)
     L = source.scatter().map(slowinc, delay=0.5).buffer(5).gather().sink_to_list()
@@ -181,7 +218,6 @@ def test_buffer(c, s, a, b):
     assert source.loop == c.loop
 
 
-@pytest.mark.slow
 def test_buffer_sync(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as c:  # noqa: F841
@@ -206,7 +242,6 @@ def test_buffer_sync(loop):  # noqa: F811
 
 
 @pytest.mark.xfail(reason='')
-@pytest.mark.slow
 def test_stream_shares_client_loop(loop):  # noqa: F811
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as client:  # noqa: F841