From 1dc5f30f93353b43c149dcf475375a0396101996 Mon Sep 17 00:00:00 2001
From: Hongsheng Zeng <zenghongsheng@baidu.com>
Date: Fri, 5 Jun 2020 17:34:55 +0800
Subject: [PATCH] provide kill process function (#283)

Co-authored-by: TomorrowIsAnOtherDay <2466956298@qq.com>
---
 parl/remote/scripts.py | 25 ++++---------------------
 parl/utils/utils.py    | 25 ++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/parl/remote/scripts.py b/parl/remote/scripts.py
index 417e848..08a06c2 100644
--- a/parl/remote/scripts.py
+++ b/parl/remote/scripts.py
@@ -27,7 +27,7 @@ import tempfile
 import warnings
 import zmq
 from multiprocessing import Process
-from parl.utils import get_ip_address, to_str, _IS_WINDOWS
+from parl.utils import get_ip_address, to_str, _IS_WINDOWS, kill_process
 from parl.remote.remote_constants import STATUS_TAG
 
 # A flag to mark if parl is started from a command line
@@ -244,26 +244,9 @@ def start_worker(address, cpu_num):
 
 @click.command("stop", help="Exit the cluster.")
 def stop():
-    if _IS_WINDOWS:
-        command = r'''for /F "skip=2 tokens=2 delims=," %a in ('wmic process where "commandline like '%remote\\job.py%'" get processid^,status /format:csv') do taskkill /F /T /pid %a'''
-        os.popen(command).read()
-
-        command = r'''for /F "skip=2 tokens=2 delims=," %a in ('wmic process where "commandline like '%remote\\start.py%'" get processid^,status /format:csv') do taskkill /F /pid %a'''
-        os.popen(command).read()
-
-        command = r'''for /F "skip=2 tokens=2 delims=," %a in ('wmic process where "commandline like '%remote\\monitor.py%'" get processid^,status /format:csv') do taskkill /F /pid %a'''
-        os.popen(command).read()
-    else:
-        command = (
-            "ps aux | grep remote/start.py | awk '{print $2}' | xargs kill -9")
-        subprocess.call([command], shell=True)
-        command = (
-            "ps aux | grep remote/job.py | awk '{print $2}' | xargs kill -9")
-        subprocess.call([command], shell=True)
-        command = (
-            "ps aux | grep remote/monitor.py | awk '{print $2}' | xargs kill -9"
-        )
-        subprocess.call([command], shell=True)
+    kill_process('remote/start.py')
+    kill_process('remote/job.py')
+    kill_process('remote/monitor.py')
 
 
 @click.command("status")
diff --git a/parl/utils/utils.py b/parl/utils/utils.py
index 71bf3bc..02c8ef2 100644
--- a/parl/utils/utils.py
+++ b/parl/utils/utils.py
@@ -13,11 +13,14 @@
 # limitations under the License.
 
 import sys
+import os
+import subprocess
 import numpy as np
 
 __all__ = [
     'has_func', 'action_mapping', 'to_str', 'to_byte', 'is_PY2', 'is_PY3',
-    'MAX_INT32', '_HAS_FLUID', '_HAS_TORCH', '_IS_WINDOWS', '_IS_MAC'
+    'MAX_INT32', '_HAS_FLUID', '_HAS_TORCH', '_IS_WINDOWS', '_IS_MAC',
+    'kill_process'
 ]
 
 
@@ -99,3 +102,23 @@ except ImportError:
 
 _IS_WINDOWS = (sys.platform == 'win32')
 _IS_MAC = (sys.platform == 'darwin')
+
+
+def kill_process(regex_pattern):
+    """kill process whose execution commnad is matched by regex pattern
+
+    Args:
+        regex_pattern(string): regex pattern used to filter the process to be killed
+    
+    NOTE:
+        In windows, we will replace sep `/` with `\\\\`
+    """
+    if _IS_WINDOWS:
+        regex_pattern = regex_pattern.replace('/', '\\\\')
+        command = r'''for /F "skip=2 tokens=2 delims=," %a in ('wmic process where "commandline like '%{}%'" get processid^,status /format:csv') do taskkill /F /T /pid %a'''.format(
+            regex_pattern)
+        os.popen(command).read()
+    else:
+        command = "ps aux | grep {} | awk '{{print $2}}' | xargs kill -9".format(
+            regex_pattern)
+        subprocess.call([command], shell=True)
-- 
GitLab