add activate arg to start_and_connect

creates a blocking DirectView on all engines, registering `%px` magics and friends

`wait_for_engines` can also derive `n` from parent cluster, if available

Author: minrk

ipyparallel/client/client.py

@@ -1455,7 +1455,7 @@ def _futures_for_msgs(self, msg_ids):
         return futures
 
     def wait_for_engines(
-        self, n, *, timeout=-1, block=True, interactive=None, widget=None
+        self, n=None, *, timeout=-1, block=True, interactive=None, widget=None
     ):
         """Wait for `n` engines to become available.
 
@@ -1490,6 +1490,18 @@ def wait_for_engines(
         ------
         TimeoutError : if timeout is reached.
         """
+        if n is None:
+            # get n from cluster, if not specified
+            if self.cluster is None:
+                raise TypeError("n engines to wait for must be specified")
+
+            if self.cluster.n:
+                n = self.cluster.n
+            else:
+                # compute n from engine sets,
+                # e.g. the default where n is calculated at runtime from `cpu_count()`
+                n = sum(engine_set.n for engine_set in self.cluster.engines.values())
+
         if len(self.ids) >= n:
             if block:
                 return

ipyparallel/cluster/cluster.py

@@ -711,25 +711,49 @@ def _engines_stopped(self, engine_set_id, stop_data=None):
         log(f"engine set stopped {engine_set_id}: {stop_data}")
         self.update_cluster_file()
 
-    async def start_and_connect(self, n=None):
+    async def start_and_connect(self, n=None, activate=False):
         """Single call to start a cluster and connect a client
 
+        If `activate` is given, a blocking DirectView on all engines will be created
+        and activated, registering `%px` magics for use in IPython
+
+        Example::
+
+            rc = await Cluster(engines="mpi").start_and_connect(n=8, activate=True)
+
+            %px print("hello, world!")
+
         Equivalent to::
 
             await self.start_cluster(n)
             client = await self.connect_client()
             await client.wait_for_engines(n, block=False)
 
-        .. versionadded: 7.1
+        .. versionadded:: 7.1
+
+        .. versionadded:: 8.1
+
+            activate argument.
         """
         if n is None:
             n = self.n
         await self.start_cluster(n=n)
         client = await self.connect_client()
+
+        if n is None:
+            # number of engines to wait for
+            # if not specified, derive current value from EngineSets
+            n = sum(engine_set.n for engine_set in self.engines.values())
+
         if n:
             await asyncio.wrap_future(
                 client.wait_for_engines(n, block=False, timeout=self.engine_timeout)
             )
+
+        if activate:
+            view = client[:]
+            view.block = True
+            view.activate()
         return client
 
     async def start_cluster(self, n=None):

ipyparallel/tests/test_cluster.py

@@ -359,3 +359,17 @@ async def test_wait_for_engines_crash(Cluster):
         rc = c.connect_client_sync()
         with pytest.raises(ipp.error.EngineError):
             rc.wait_for_engines(3, timeout=20)
+
+
+@pytest.mark.parametrize("activate", (True, False))
+def test_start_and_connect_activate(ipython, Cluster, activate):
+    rc = Cluster(n=2, log_level=10).start_and_connect_sync(activate=activate)
+    with rc:
+        if activate:
+            assert "px" in ipython.magics_manager.magics["cell"]
+            px = ipython.magics_manager.magics["cell"]["px"]
+            assert px.__self__.view.client is rc
+        else:
+            if "px" in ipython.magics_manager.magics["cell"]:
+                px = ipython.magics_manager.magics["cell"]["px"]
+                assert px.__self__.view.client is not rc