Merge pull request #487 from minrk/tqdm-notebook

add some more control to interactive waits

Author: minrk

ipyparallel/client/asyncresult.py

@@ -10,25 +10,23 @@
 from contextlib import contextmanager
 from datetime import datetime
 from functools import partial
+from queue import Queue
 from threading import Event
 
-try:
-    from queue import Queue
-except ImportError:  # py2
-    from Queue import Queue
-
-from decorator import decorator
-import tqdm
 import zmq
-from zmq import MessageTracker
-
+from decorator import decorator
 from IPython import get_ipython
-from IPython.core.display import display, display_pretty
-from ipyparallel import error
-from ipyparallel.util import utcnow, compare_datetimes, _parse_date
+from IPython.core.display import display
+from IPython.core.display import display_pretty
 from ipython_genutils.py3compat import string_types
 
-from .futures import MessageFuture, multi_future
+from .futures import MessageFuture
+from .futures import multi_future
+from ipyparallel import error
+from ipyparallel.util import _parse_date
+from ipyparallel.util import compare_datetimes
+from ipyparallel.util import progress
+from ipyparallel.util import utcnow
 
 
 def _raw_text(s):
@@ -38,7 +36,7 @@ def _raw_text(s):
 _default = object()
 
 # global empty tracker that's always done:
-finished_tracker = MessageTracker()
+finished_tracker = zmq.MessageTracker()
 
 
 @decorator
@@ -75,7 +73,6 @@ def __init__(
         fname='unknown',
         targets=None,
         owner=False,
-        progress_bar=tqdm.tqdm,
     ):
         super(AsyncResult, self).__init__()
         if not isinstance(children, list):
@@ -95,7 +92,6 @@ def __init__(
         self._fname = fname
         self._targets = targets
         self.owner = owner
-        self.progress_bar = progress_bar
 
         self._ready = False
         self._ready_event = Event()
@@ -381,7 +377,7 @@ def _handle_sent(self, f):
         """Resolve sent Future, build MessageTracker"""
         trackers = f.result()
         trackers = [t for t in trackers if t is not None]
-        self._tracker = MessageTracker(*trackers)
+        self._tracker = zmq.MessageTracker(*trackers)
         self._sent_event.set()
 
     @property
@@ -569,13 +565,27 @@ def wall_time(self):
         """
         return self.timedelta(self.submitted, self.received)
 
-    def wait_interactive(self, interval=1.0, timeout=-1):
-        """interactive wait, printing progress at regular intervals."""
+    def wait_interactive(self, interval=0.1, timeout=-1, widget=None):
+        """interactive wait, printing progress at regular intervals.
+
+        Parameters
+        ----------
+        interval : float
+            Interval on which to update progress display.
+        timeout : float
+            Time (in seconds) to wait before raising a TimeoutError.
+            -1 (default) means no timeout.
+        widget : bool
+            default: True if in an IPython kernel (notebook), False otherwise.
+            Override default context-detection behavior for whether a widget-based progress bar
+            should be used.
+        """
         if timeout is None:
             timeout = -1
         N = len(self)
         tic = time.perf_counter()
-        progress_bar = self.progress_bar(total=N, unit='tasks', desc=self._fname)
+        progress_bar = progress(widget=widget, total=N, unit='tasks', desc=self._fname)
+
         n_prev = 0
         while not self.ready() and (
             timeout < 0 or time.perf_counter() - tic <= timeout

ipyparallel/client/client.py

@@ -1307,7 +1307,9 @@ def _futures_for_msgs(self, msg_ids):
                 futures.append(f)
         return futures
 
-    def wait_for_engines(self, n, *, timeout=-1, block=True):
+    def wait_for_engines(
+        self, n, *, timeout=-1, block=True, interactive=None, widget=None
+    ):
         """Wait for `n` engines to become available.
 
         Returns when `n` engines are available,
@@ -1322,6 +1324,14 @@ def wait_for_engines(self, n, *, timeout=-1, block=True):
             Time (in seconds) to wait before raising a TimeoutError
         block : bool
             if False, return Future instead of waiting
+        interactive : bool
+            default: True if in IPython, False otherwise.
+            if True, show a progress bar while waiting for engines
+        widget : bool
+            default: True if in an IPython kernel (notebook), False otherwise.
+            Only has an effect if `interactive` is True.
+            if True, forces use of widget progress bar.
+            If False, forces use of terminal tqdm.
 
         Returns
         ------
@@ -1347,12 +1357,25 @@ def wait_for_engines(self, n, *, timeout=-1, block=True):
             deadline = None
             seconds_remaining = 1000
 
+        if interactive is None:
+            interactive = get_ipython() is not None
+
+        if interactive:
+            progress_bar = util.progress(
+                widget=widget, initial=len(self.ids), total=n, unit='engine'
+            )
+
         future = Future()
 
         def notify(_):
             if future.done():
                 return
+            if interactive:
+                progress_bar.update(len(self.ids) - progress_bar.n)
             if len(self.ids) >= n:
+                # ensure we refresh when we finish
+                if interactive:
+                    progress_bar.close()
                 future.set_result(None)
 
         future.add_done_callback(lambda f: self._registration_callbacks.remove(notify))

ipyparallel/util.py

@@ -30,6 +30,7 @@ def lru_cache(maxsize=0):
     SIGKILL = None
 from types import FunctionType
 
+import tqdm
 from dateutil.parser import parse as dateutil_parse
 from dateutil.tz import tzlocal
 
@@ -636,3 +637,29 @@ def _patch_jupyter_client_dates():
 
 # FIXME: remove patch when we require jupyter_client 5.0
 _patch_jupyter_client_dates()
+
+
+def progress(*args, widget=None, **kwargs):
+    """Create a tqdm progress bar
+
+    If `widget` is None, autodetects if IPython widgets should be used,
+    otherwise use basic tqdm.
+    """
+    if widget is None:
+        # auto widget if in a kernel
+        ip = get_ipython()
+        if ip is not None and getattr(ip, 'kernel', None) is not None:
+            try:
+                import ipywidgets  # noqa
+            except ImportError:
+                widget = False
+            else:
+                widget = True
+        else:
+            widget = False
+    if widget:
+        f = tqdm.tqdm_notebook
+    else:
+        kwargs.setdefault("file", sys.stdout)
+        f = tqdm.tqdm
+    return f(*args, **kwargs)