Merge pull request #540 from minrk/launcher-cleanup

Improve launcher cleanup

Author: minrk

ipyparallel/cluster/cluster.py

@@ -16,6 +16,7 @@
 import string
 import sys
 import time
+import traceback
 from functools import partial
 from multiprocessing import cpu_count
 from weakref import WeakSet
@@ -56,7 +57,11 @@ def _atexit_cleanup_clusters(*args):
             continue
         if cluster.controller or cluster.engines:
             print(f"Stopping cluster {cluster}", file=sys.stderr)
-            cluster.stop_cluster_sync()
+            try:
+                cluster.stop_cluster_sync()
+            except Exception:
+                print(f"Error stopping cluster {cluster}", file=sys.stderr)
+                traceback.print_exception(*sys.exc_info())
 
 
 _atexit_cleanup_clusters.registered = False
@@ -354,22 +359,37 @@ def from_dict(cls, d, **kwargs):
             if attr in d:
                 setattr(self, attr, d[attr])
 
+        cluster_key = ClusterManager._cluster_key(self)
+
         if d.get("controller"):
             controller_info = d["controller"]
             cls = self.controller_launcher_class = import_item(controller_info["class"])
             if controller_info["state"]:
-                self.controller = cls.from_dict(controller_info["state"], parent=self)
+                try:
+                    self.controller = cls.from_dict(
+                        controller_info["state"], parent=self
+                    )
+                except launcher.NotRunning as e:
+                    self.log.error(f"Controller for {cluster_key} not running: {e}")
 
         engine_info = d.get("engines")
         if engine_info:
             cls = self.engine_launcher_class = import_item(engine_info["class"])
             for engine_set_id, engine_state in engine_info.get("sets", {}).items():
-                self.engines[engine_set_id] = cls.from_dict(
-                    engine_state,
-                    engine_set_id=engine_set_id,
-                    parent=self,
-                )
-
+                try:
+                    self.engines[engine_set_id] = cls.from_dict(
+                        engine_state,
+                        engine_set_id=engine_set_id,
+                        parent=self,
+                    )
+                except launcher.NotRunning as e:
+                    self.log.error(
+                        f"Engine set {cluster_key}{engine_set_id} not running: {e}"
+                    )
+        # check if state changed
+        if self.to_dict() != d:
+            # if so, update our cluster file
+            self.update_cluster_file()
         return self
 
     @classmethod
@@ -698,7 +718,8 @@ class ClusterManager(LoggingConfigurable):
 
     _clusters = Dict(help="My cluster objects")
 
-    def _cluster_key(self, cluster):
+    @staticmethod
+    def _cluster_key(cluster):
         """Return a unique cluster key for a cluster
 
         Default is {profile}:{cluster_id}

ipyparallel/cluster/launcher.py

@@ -82,6 +82,12 @@ class UnknownStatus(LauncherError):
     pass
 
 
+class NotRunning(LauncherError):
+    """Raised when a launcher is no longer running"""
+
+    pass
+
+
 class BaseLauncher(LoggingConfigurable):
     """An abstraction for starting, stopping and signaling a process."""
 
@@ -408,7 +414,10 @@ def from_dict(cls, d, **kwargs):
     def _reconstruct_process(self, d):
         """Reconstruct our process"""
         if 'pid' in d and d['pid'] > 0:
-            self.process = psutil.Process(d['pid'])
+            try:
+                self.process = psutil.Process(d['pid'])
+            except psutil.NoSuchProcess as e:
+                raise NotRunning(f"Process {d['pid']}")
             self._start_waiting()
 
     def _wait(self):
@@ -465,10 +474,18 @@ def start(self):
     async def join(self, timeout=None):
         """Wait for the process to exit"""
         with ThreadPoolExecutor(1) as pool:
+            wait = partial(self.process.wait, timeout)
             try:
-                await asyncio.wrap_future(
-                    pool.submit(partial(self.process.wait, timeout))
-                )
+                try:
+                    future = pool.submit(wait)
+                except RuntimeError:
+                    # e.g. called during process shutdown,
+                    # which raises
+                    # RuntimeError: cannot schedule new futures after interpreter shutdown
+                    # Instead, do the blocking call
+                    wait()
+                else:
+                    await asyncio.wrap_future(future)
             except psutil.TimeoutExpired:
                 raise TimeoutError(
                     f"Process {self.pid} did not complete in {timeout} seconds."
@@ -638,8 +655,20 @@ def to_dict(self):
     @classmethod
     def from_dict(cls, d, **kwargs):
         self = super().from_dict(d, **kwargs)
+        n = 0
         for i, engine_dict in d['engines'].items():
-            self.launchers[i] = self.launcher_class.from_dict(engine_dict, parent=self)
+            try:
+                self.launchers[i] = self.launcher_class.from_dict(
+                    engine_dict, parent=self
+                )
+            except NotRunning as e:
+                self.log.error(f"Engine {i} not running: {e}")
+            else:
+                n += 1
+        if n == 0:
+            raise NotRunning("No engines left")
+        else:
+            self.n = n
         return self
 
     def start(self, n):
@@ -1184,9 +1213,17 @@ def wait_one(self, timeout):
 
     async def join(self, timeout=None):
         with ThreadPoolExecutor(1) as pool:
-            await asyncio.wrap_future(
-                pool.submit(partial(self.wait_one, timeout=timeout))
-            )
+            wait = partial(self.wait_one, timeout=timeout)
+            try:
+                future = pool.submit(wait)
+            except RuntimeError:
+                # e.g. called during process shutdown,
+                # which raises
+                # RuntimeError: cannot schedule new futures after interpreter shutdown
+                # Instead, do the blocking call
+                wait()
+            else:
+                await asyncio.wrap_future(future)
 
     def signal(self, sig):
         if self.state == 'running':