Repository: mars-project/mars

Issue:
[BUG] When stopping worker processes, subprocesses also need to be stopped
Mars executor launches subprocesses for TensorFlow and PyTorch, and these processes need to be terminated when worker processes are terminated.


Discussion:

Resolving patch:
diff --git a/mars/deploy/local/core.py b/mars/deploy/local/core.py
--- a/mars/deploy/local/core.py
+++ b/mars/deploy/local/core.py
@@ -31,7 +31,7 @@
 from ...resource import cpu_count
 from ...scheduler.service import SchedulerService
 from ...session import new_session
-from ...utils import get_next_port
+from ...utils import get_next_port, kill_process_tree
 from ...worker.service import WorkerService
 from .distributor import gen_distributor
 
@@ -266,20 +266,7 @@ def _ensure_process_finish(proc):
         if proc is None or not proc.is_alive():
             return
         proc.join(3)
-
-        # in case the process does not finish
-        if proc.is_alive():  # pragma: no cover
-            try:
-                import psutil
-                for subproc in psutil.Process(proc.pid).children(recursive=True):
-                    try:
-                        subproc.kill()
-                    except psutil.NoSuchProcess:  # pragma: no cover
-                        pass
-            except ImportError:
-                pass
-            finally:
-                proc.terminate()
+        kill_process_tree(proc.pid)
 
     def stop(self):
         try:
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -23,6 +23,7 @@
 import os
 import pkgutil
 import random
+import shutil
 import socket
 import struct
 import sys
@@ -702,3 +703,29 @@ def has_unknown_shape(tiled):
     if any(np.isnan(s) for s in itertools.chain(*tiled.nsplits)):
         return True
     return False
+
+
+def kill_process_tree(pid, include_parent=True):
+    try:
+        import psutil
+    except ImportError:  # pragma: no cover
+        return
+    try:
+        proc = psutil.Process(pid)
+    except psutil.NoSuchProcess:
+        return
+
+    plasma_sock_dir = None
+    children = proc.children(recursive=True)
+    if include_parent:
+        children.append(proc)
+    for p in children:
+        try:
+            if 'plasma' in p.name():
+                plasma_sock_dir = next((conn.laddr for conn in p.connections('unix')
+                                        if 'plasma' in conn.laddr), None)
+            p.kill()
+        except psutil.NoSuchProcess:  # pragma: no cover
+            pass
+    if plasma_sock_dir:
+        shutil.rmtree(plasma_sock_dir, ignore_errors=True)
diff --git a/mars/worker/daemon.py b/mars/worker/daemon.py
--- a/mars/worker/daemon.py
+++ b/mars/worker/daemon.py
@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import signal
 from collections import defaultdict
 
 import psutil
 
 from .utils import WorkerActor
+from ..utils import kill_process_tree
 
 
 class WorkerDaemonActor(WorkerActor):
@@ -85,7 +84,7 @@ def kill_actor_process(self, actor_ref):
         try:
             pid = self._proc_pids[proc_idx]
             self._killed_pids.add(pid)
-            os.kill(pid, signal.SIGKILL)
+            kill_process_tree(pid)
         except (KeyError, OSError):
             pass
 


Test patch:
diff --git a/mars/scheduler/tests/integrated/base.py b/mars/scheduler/tests/integrated/base.py
--- a/mars/scheduler/tests/integrated/base.py
+++ b/mars/scheduler/tests/integrated/base.py
@@ -14,7 +14,6 @@
 
 import logging
 import os
-import shutil
 import signal
 import subprocess
 import sys
@@ -27,7 +26,7 @@
 from mars import resource
 from mars.config import options
 from mars.tests.core import EtcdProcessHelper
-from mars.utils import get_next_port
+from mars.utils import get_next_port, kill_process_tree
 from mars.actors.core import new_client
 from mars.scheduler import SessionManagerActor, ResourceActor
 from mars.scheduler.graph import GraphState
@@ -106,23 +105,7 @@ def tearDown(self):
     def kill_process_tree(self, proc, intentional=True):
         if intentional:
             self.intentional_death_pids.add(proc.pid)
-
-        import psutil
-        proc = psutil.Process(proc.pid)
-        plasma_sock_dir = None
-        for p in proc.children(recursive=True):
-            try:
-                if 'plasma' in p.name():
-                    socks = [conn.laddr for conn in p.connections('unix')
-                             if 'plasma' in conn.laddr]
-                    if socks:
-                        plasma_sock_dir = os.path.dirname(socks[0])
-                p.kill()
-            except psutil.NoSuchProcess:
-                continue
-        proc.kill()
-        if plasma_sock_dir:
-            shutil.rmtree(plasma_sock_dir, ignore_errors=True)
+        kill_process_tree(proc.pid)
 
     def add_state_file(self, environ):
         fn = os.environ[environ] = os.path.join(