use client.wait_for_engines in cluster context managers

Author: minrk

ipyparallel/cluster/cluster.py

@@ -5,7 +5,6 @@
 starts/stops/polls controllers, engines, etc.
 """
 import asyncio
-import atexit
 import inspect
 import logging
 import os
@@ -79,6 +78,15 @@ def _default_profile_dir(self):
              a shortcut for specifying profile_dir within $IPYTHONDIR.""",
     )
 
+    engine_timeout = Integer(
+        60,
+        help="""Timeout to use when waiting for engines to register
+
+        before giving up.
+        """,
+        config=True,
+    )
+
     controller_launcher_class = Launcher(
         default_value=launcher.LocalControllerLauncher,
         kind='Controller',
@@ -424,9 +432,11 @@ async def __aenter__(self):
         client = self._context_client = self.connect_client()
         if self.n:
             # wait for engine registration
-            # TODO: timeout
-            while len(client) < self.n:
-                await asyncio.sleep(0.1)
+            await asyncio.wrap_future(
+                client.wait_for_engines(
+                    self.n, block=False, timeout=self.engine_timeout
+                )
+            )
         return client
 
     async def __aexit__(self, *args):
@@ -442,8 +452,7 @@ def __enter__(self):
         client = self._context_client = self.connect_client()
         if self.n:
             # wait for engine registration
-            while len(client) < self.n:
-                time.sleep(0.1)
+            client.wait_for_engines(self.n, block=True, timeout=self.engine_timeout)
         return client
 
     def __exit__(self, *args):