fix: localmode subprocess parent process not sending SIGTERM to child

Author: ranzvi

setup.py

@@ -52,7 +52,6 @@ def read_version():
         "urllib3>=1.21.1,!=1.25,!=1.25.1",
         "docker-compose>=1.25.2",
         "PyYAML>=5.3, <6",  # PyYAML version has to match docker-compose requirements
-        "psutil",
     ],
     "scipy": ["scipy>=0.19.0"],
 }

src/sagemaker/local/image.py

@@ -31,10 +31,7 @@
 import tempfile
 
 from distutils.spawn import find_executable
-from signal import SIGTERM
 from threading import Thread
-
-import psutil
 from six.moves.urllib.parse import urlparse
 
 import sagemaker
@@ -843,8 +840,8 @@ def run(self):
 
     def down(self):
         """Placeholder docstring"""
-        for process in psutil.Process(self.process.pid).children():
-            process.send_signal(SIGTERM)
+        if os.name != 'nt':
+            sagemaker.local.utils.kill_child_processes(self.process.pid)
         self.process.terminate()
 
 

src/sagemaker/local/utils.py

@@ -15,6 +15,7 @@
 
 import os
 import shutil
+import subprocess
 
 from distutils.dir_util import copy_tree
 from six.moves.urllib.parse import urlparse
@@ -88,3 +89,24 @@ def recursive_copy(source, destination):
     """
     if os.path.isdir(source):
         copy_tree(source, destination)
+
+
+def kill_child_processes(pid):
+    child_pids = get_child_process_ids(pid)
+    for pid in child_pids:
+        os.kill(pid, 15)
+
+
+def get_child_process_ids(pid):
+    cmd = f"pgrep -P {pid}".split()
+    output, err = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    ).communicate()
+    if err:
+        return []
+    pids = [int(pid) for pid in output.decode('utf-8').split()]
+    if pids:
+        for pid in pids:
+            return pids + get_child_process_ids(pid)
+    else:
+        return []