more test coverage for Cluster API

reasonably complete now

fix various little things along the way

Author: minrk

ipyparallel/cluster/cluster.py

@@ -19,6 +19,7 @@
 import IPython
 import traitlets.log
 from IPython.core.profiledir import ProfileDir
+from IPython.core.profiledir import ProfileDirError
 from traitlets import Any
 from traitlets import default
 from traitlets import Dict
@@ -68,9 +69,13 @@ def _default_profile_dir(self):
             ip = IPython.get_ipython()
             if ip is not None:
                 return ip.profile_dir.location
-        return ProfileDir.find_profile_dir_by_name(
-            IPython.paths.get_ipython_dir(), name=self.profile or 'default'
-        ).location
+        ipython_dir = IPython.paths.get_ipython_dir()
+        profile_name = self.profile or 'default'
+        try:
+            pd = ProfileDir.find_profile_dir_by_name(ipython_dir, name=profile_name)
+        except ProfileDirError:
+            pd = ProfileDir.create_profile_dir_by_name(ipython_dir, name=profile_name)
+        return pd.location
 
     profile = Unicode(
         "",
@@ -215,16 +220,22 @@ def _default_log(self):
     _engine_sets = Dict()
 
     def __repr__(self):
-        profile_dir = self.profile_dir
-        home_dir = os.path.expanduser("~")
-        if profile_dir.startswith(home_dir + os.path.sep):
-            # truncate $HOME/. -> ~/...
-            profile_dir = "~" + profile_dir[len(home_dir) :]
 
         fields = {
             "cluster_id": repr(self.cluster_id),
-            "profile_dir": repr(profile_dir),
         }
+        profile_dir = self.profile_dir
+        profile_prefix = os.path.join(IPython.paths.get_ipython_dir(), "profile_")
+        if profile_dir.startswith(profile_prefix):
+            fields["profile"] = repr(profile_dir[len(profile_prefix) :])
+        else:
+            home_dir = os.path.expanduser("~")
+
+            if profile_dir.startswith(home_dir + os.path.sep):
+                # truncate $HOME/. -> ~/...
+                profile_dir = "~" + profile_dir[len(home_dir) :]
+            fields["profile_dir"] = repr(profile_dir)
+
         if self._controller:
             fields["controller"] = "<running>"
         if self._engine_sets:
@@ -473,7 +484,7 @@ class ClusterManager(LoggingConfigurable):
 
     _clusters = Dict(help="My cluster objects")
 
-    def load_clusters(self, serialized_state):
+    def from_dict(self, serialized_state):
         """Load serialized cluster state"""
         raise NotImplementedError("Serializing clusters not implemented")
 
@@ -484,12 +495,13 @@ def list_clusters(self):
         # just cluster ids for now
         return sorted(self._clusters)
 
-    def new_cluster(self, cluster_cls, **kwargs):
+    def new_cluster(self, **kwargs):
         """Create a new cluster"""
-        cluster = Cluster(parent=self)
+        cluster = Cluster(parent=self, **kwargs)
         if cluster.cluster_id in self._clusters:
             raise KeyError(f"Cluster {cluster.cluster_id} already exists!")
-        self._clusters[cluster]
+        self._clusters[cluster.cluster_id] = cluster
+        return cluster
 
     def get_cluster(self, cluster_id):
         """Get a Cluster object by id"""
@@ -499,19 +511,3 @@ def remove_cluster(self, cluster_id):
         """Delete a cluster by id"""
         # TODO: check running?
         del self._clusters[cluster_id]
-
-    def _cluster_method(self, method_name, cluster_id, *args, **kwargs):
-        """Wrapper around single-cluster methods
-
-        Defines ClusterManager.method(cluster_id, ...)
-
-        which returns ClusterManager.clusters[cluster_id].method(...)
-        """
-        cluster = self._clusters[cluster_id]
-        method = getattr(cluster, method_name)
-        return method(*args, **kwargs)
-
-    def __getattr__(self, key):
-        if key in Cluster.__dict__:
-            return partial(self._cluster_method, key)
-        return super().__getattr__(self, key)

ipyparallel/cluster/launcher.py

@@ -412,6 +412,7 @@ def __init__(self, work_dir=u'.', config=None, **kwargs):
 
     def start(self, n):
         """Start n engines by profile or profile_dir."""
+        self.n = n
         dlist = []
         for i in range(n):
             if i > 0:

ipyparallel/tests/conftest.py

@@ -1,5 +1,8 @@
 """pytest fixtures"""
 import inspect
+import os
+from tempfile import TemporaryDirectory
+from unittest import mock
 
 import pytest
 from IPython.terminal.interactiveshell import TerminalInteractiveShell
@@ -9,6 +12,13 @@
 from . import teardown
 
 
+@pytest.fixture(autouse=True, scope="session")
+def temp_ipython():
+    with TemporaryDirectory(suffix="dotipython") as td:
+        with mock.patch.dict(os.environ, {"IPYTHONDIR": td}):
+            yield
+
+
 def pytest_collection_modifyitems(items):
     """This function is automatically run by pytest passing all collected test
     functions.

ipyparallel/tests/test_cluster.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 import shutil
 import signal
 import sys
@@ -12,6 +13,10 @@
 from ipyparallel import cluster
 from ipyparallel.cluster.launcher import find_launcher_class
 
+_engine_launcher_classes = ["Local"]
+if shutil.which("mpiexec"):
+    _engine_launcher_classes.append("MPI")
+
 
 @pytest.fixture
 def Cluster(request):
@@ -52,6 +57,13 @@ async def test_cluster_id(Cluster):
     assert cluster.cluster_id == 'abc'
 
 
+async def test_ipython_log(ipython):
+    c = cluster.Cluster(parent=ipython)
+    assert c.log.name == f"{cluster.Cluster.__module__}.{c.cluster_id}"
+    assert len(c.log.handlers) == 1
+    assert c.log.handlers[0].stream is sys.stdout
+
+
 async def test_start_stop_controller(Cluster):
     cluster = Cluster()
     await cluster.start_controller()
@@ -72,7 +84,7 @@ async def test_start_stop_controller(Cluster):
     # TODO: test file cleanup
 
 
-@pytest.mark.parametrize("engine_launcher_class", ["Local", "MPI"])
+@pytest.mark.parametrize("engine_launcher_class", _engine_launcher_classes)
 async def test_start_stop_engines(Cluster, engine_launcher_class):
     cluster = Cluster(engine_launcher_class=engine_launcher_class)
     await cluster.start_controller()
@@ -89,7 +101,23 @@ async def test_start_stop_engines(Cluster, engine_launcher_class):
     await cluster.stop_controller()
 
 
-@pytest.mark.parametrize("engine_launcher_class", ["Local", "MPI"])
+@pytest.mark.parametrize("engine_launcher_class", _engine_launcher_classes)
+async def test_start_stop_cluster(Cluster, engine_launcher_class):
+    n = 2
+    cluster = Cluster(engine_launcher_class=engine_launcher_class, n=n)
+    await cluster.start_cluster()
+    controller = cluster._controller
+    assert controller is not None
+    assert len(cluster._engine_sets) == 1
+
+    rc = cluster.connect_client()
+    rc.wait_for_engines(n, timeout=10)
+    await cluster.stop_cluster()
+    assert cluster._controller is None
+    assert cluster._engine_sets == {}
+
+
+@pytest.mark.parametrize("engine_launcher_class", _engine_launcher_classes)
 async def test_signal_engines(Cluster, engine_launcher_class):
     cluster = Cluster(engine_launcher_class=engine_launcher_class)
     await cluster.start_controller()
@@ -117,6 +145,25 @@ async def test_signal_engines(Cluster, engine_launcher_class):
     await cluster.stop_controller()
 
 
+@pytest.mark.parametrize("engine_launcher_class", _engine_launcher_classes)
+async def test_restart_engines(Cluster, engine_launcher_class):
+    n = 3
+    async with Cluster(engine_launcher_class=engine_launcher_class, n=n) as rc:
+        cluster = rc.cluster
+        engine_set_id = next(iter(cluster._engine_sets))
+        engine_set = cluster._engine_sets[engine_set_id]
+        assert rc.ids == list(range(n))
+        before_pids = rc[:].apply_sync(os.getpid)
+        await cluster.restart_engines()
+        # wait for unregister
+        while any(eid in rc.ids for eid in range(n)):
+            await asyncio.sleep(0.1)
+        # wait for register
+        rc.wait_for_engines(n, timeout=10)
+        after_pids = rc[:].apply_sync(os.getpid)
+        assert set(after_pids).intersection(before_pids) == set()
+
+
 async def test_async_with(Cluster):
     async with Cluster(n=5) as rc:
         assert sorted(rc.ids) == list(range(5))
@@ -160,3 +207,22 @@ async def test_cluster_repr(Cluster):
         repr(c)
         == "<Cluster(cluster_id='test', profile_dir='/tmp', controller=<running>, engine_sets=['engineid'])>"
     )
+
+
+async def test_cluster_manager():
+    m = cluster.ClusterManager()
+    assert m.list_clusters() == []
+    c = m.new_cluster(profile_dir="/tmp")
+    assert c.profile_dir == "/tmp"
+    assert m.get_cluster(c.cluster_id) is c
+    with pytest.raises(KeyError):
+        m.get_cluster("nosuchcluster")
+
+    with pytest.raises(KeyError):
+        m.new_cluster(cluster_id=c.cluster_id)
+
+    assert m.list_clusters() == [c.cluster_id]
+    m.remove_cluster(c.cluster_id)
+    assert m.list_clusters() == []
+    with pytest.raises(KeyError):
+        m.remove_cluster("nosuchcluster")