Merge pull request #655 from minrk/chunksize-progress

Use chunksize in progress

Author: minrk

ipyparallel/client/asyncresult.py

@@ -72,6 +72,7 @@ class AsyncResult(Future):
     owner = False
     _last_display_prefix = ""
     _stream_trailing_newline = True
+    _chunk_sizes = None
 
     def __init__(
         self,
@@ -81,6 +82,7 @@ def __init__(
         targets=None,
         owner=False,
         return_exceptions=False,
+        chunk_sizes=None,
     ):
         super().__init__()
         if not isinstance(children, list):
@@ -90,6 +92,7 @@ def __init__(
             self._single_result = False
 
         self._return_exceptions = return_exceptions
+        self._chunk_sizes = chunk_sizes or {}
 
         if isinstance(children[0], str):
             self.msg_ids = children
@@ -748,8 +751,14 @@ def __iter__(self):
             # already done
             yield from rlist
 
+    @lru_cache()
     def __len__(self):
-        return len(self.msg_ids)
+        return self._count_chunks(*self.msg_ids)
+
+    @lru_cache()
+    def _count_chunks(self, *msg_ids):
+        """Count the granular tasks"""
+        return sum(self._chunk_sizes.setdefault(msg_id, 1) for msg_id in msg_ids)
 
     # -------------------------------------
     # Sugar methods and attributes
@@ -795,7 +804,9 @@ def progress(self):
         Fractional progress would be given by 1.0 * ar.progress / len(ar)
         """
         self.wait(0)
-        return len(self) - len(set(self.msg_ids).intersection(self._client.outstanding))
+        finished_msg_ids = set(self.msg_ids).intersection(self._client.outstanding)
+        finished_count = self._count_chunks(*finished_msg_ids)
+        return len(self) - finished_count
 
     @property
     def elapsed(self):
@@ -1069,6 +1080,7 @@ def __init__(
         fname='',
         ordered=True,
         return_exceptions=False,
+        chunk_sizes=None,
     ):
         self._mapObject = mapObject
         self.ordered = ordered
@@ -1078,6 +1090,7 @@ def __init__(
             children,
             fname=fname,
             return_exceptions=return_exceptions,
+            chunk_sizes=chunk_sizes,
         )
         self._single_result = False
 

ipyparallel/client/remotefunction.py

@@ -271,6 +271,9 @@ def __call__(self, *sequences, **kwargs):
 
         pf = PrePickled(self.func)
 
+        chunk_sizes = {}
+        chunk_size = 1
+
         for index, t in enumerate(targets):
             args = []
             for seq in sequences:
@@ -279,6 +282,10 @@ def __call__(self, *sequences, **kwargs):
 
             if sum(len(arg) for arg in args) == 0:
                 continue
+
+            if _mapping:
+                chunk_size = min(len(arg) for arg in args)
+
             args = [PrePickled(arg) for arg in args]
 
             if _mapping:
@@ -292,6 +299,8 @@ def __call__(self, *sequences, **kwargs):
                 ar = view.apply(f, *args)
                 ar.owner = False
 
+            msg_id = ar.msg_ids[0]
+            chunk_sizes[msg_id] = chunk_size
             futures.extend(ar._children)
 
         r = AsyncMapResult(
@@ -301,6 +310,7 @@ def __call__(self, *sequences, **kwargs):
             fname=getname(self.func),
             ordered=self.ordered,
             return_exceptions=self.return_exceptions,
+            chunk_sizes=chunk_sizes,
         )
 
         if self.block:

ipyparallel/tests/test_asyncresult.py

@@ -480,3 +480,11 @@ def fail(i):
         done, pending = amr.wait(timeout=0, return_when=ipp.FIRST_EXCEPTION)
         assert pending == set()
         assert len(done) == len(amr)
+
+    def test_progress(self):
+        dv = self.client[:]
+        amr = dv.map_async(time.sleep, [0.2] * 2 * len(dv))
+        assert len(amr) == len(dv) * 2
+        assert amr.progress == 0
+        amr.wait_interactive()
+        assert amr.progress == len(amr)

ipyparallel/tests/test_lbview.py

@@ -29,16 +29,19 @@ def f(x):
             return x ** 2
 
         data = list(range(16))
-        r = self.view.map_sync(f, data)
-        self.assertEqual(r, list(map(f, data)))
+        ar = self.view.map_async(f, data)
+        assert len(ar) == len(data)
+        r = ar.get()
+        assert r == list(map(f, data))
 
     def test_map_generator(self):
         def f(x):
             return x ** 2
 
         data = list(range(16))
-        r = self.view.map_sync(f, iter(data))
-        self.assertEqual(r, list(map(f, iter(data))))
+        ar = self.view.map_async(f, iter(data))
+        r = ar.get()
+        assert r == list(map(f, iter(data)))
 
     def test_map_short_first(self):
         def f(x, y):
@@ -51,8 +54,10 @@ def f(x, y):
         data = list(range(10))
         data2 = list(range(4))
 
-        r = self.view.map_sync(f, data, data2)
-        self.assertEqual(r, list(map(f, data, data2)))
+        ar = self.view.map_async(f, data, data2)
+        assert len(ar) == len(data2)
+        r = ar.get()
+        assert r == list(map(f, data, data2))
 
     def test_map_short_last(self):
         def f(x, y):
@@ -65,8 +70,10 @@ def f(x, y):
         data = list(range(4))
         data2 = list(range(10))
 
-        r = self.view.map_sync(f, data, data2)
-        self.assertEqual(r, list(map(f, data, data2)))
+        ar = self.view.map_async(f, data, data2)
+        assert len(ar) == len(data)
+        r = ar.get()
+        assert r == list(map(f, data, data2))
 
     def test_map_unordered(self):
         def f(x):

ipyparallel/tests/test_view.py

@@ -334,7 +334,9 @@ def f(x):
             return x ** 2
 
         data = list(range(16))
-        r = view.map_sync(f, data)
+        ar = view.map_async(f, data)
+        assert len(ar) == len(data)
+        r = ar.get()
         self.assertEqual(r, list(map(f, data)))
 
     def test_map_empty_sequence(self):
@@ -361,15 +363,17 @@ def test_map_numpy(self):
         view = self.client[:]
         # 101 is prime, so it won't be evenly distributed
         arr = numpy.arange(101)
-        r = view.map_sync(lambda x: x, arr)
+        ar = view.map_async(lambda x: x, arr)
+        assert len(ar) == len(arr)
+        r = ar.get()
         assert_array_equal(r, arr)
 
     def test_scatter_gather_nonblocking(self):
         data = list(range(16))
         view = self.client[:]
         view.scatter('a', data, block=False)
         ar = view.gather('a', block=False)
-        self.assertEqual(ar.get(), data)
+        assert ar.get() == data
 
     @skip_without('numpy')
     def test_scatter_gather_numpy_nonblocking(self):