From c8b69727cd89cea7ef7590d695ffe426d4f2b700 Mon Sep 17 00:00:00 2001
From: Simon Mo <xmo@berkeley.edu>
Date: Sun, 24 Nov 2019 16:32:07 -0800
Subject: [PATCH] `ray stop` only kills process with `ray` keyword (#6257)

* Use psutil to kill processes

* Psutil as core requirement

* Revert "Psutil as core requirement"

This reverts commit d3235ce3d994d2bb7db39e3ad4a46049703898bb.

* Revert "Use psutil to kill processes"

This reverts commit de0ed874fed673f5e98715950688f418bbcc415c.

* Revert back to subproc

* Add comments, grep for ray as well

* SIGTERM
---
 python/ray/scripts/scripts.py | 37 +++++++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/python/ray/scripts/scripts.py b/python/ray/scripts/scripts.py
index 1d5c2fb36..9144bfed4 100644
--- a/python/ray/scripts/scripts.py
+++ b/python/ray/scripts/scripts.py
@@ -415,7 +415,12 @@ def start(node_ip_address, redis_address, address, redis_port,
 
 
 @cli.command()
-def stop():
+@click.option(
+    "-f",
+    "--force",
+    is_flag=True,
+    help="If set, ray will send SIGKILL instead of SIGTERM.")
+def stop(force):
     # Note that raylet needs to exit before object store, otherwise
     # it cannot exit gracefully.
     processes_to_kill = [
@@ -426,33 +431,45 @@ def stop():
         # See STANDARD FORMAT SPECIFIERS section of
         # http://man7.org/linux/man-pages/man1/ps.1.html
         # about comm and args. This can help avoid killing non-ray processes.
+
+        # Format:
+        # Keyword to filter, filter by command (True)/filter by args (False)
         ["raylet", True],
         ["plasma_store", True],
         ["raylet_monitor", True],
         ["monitor.py", False],
         ["redis-server", True],
         ["default_worker.py", False],  # Python worker.
-        [" ray_", True],  # Python worker.
+        ["ray_", True],  # Python worker.
         ["org.ray.runtime.runner.worker.DefaultWorker", False],  # Java worker.
         ["log_monitor.py", False],
         ["reporter.py", False],
         ["dashboard.py", False],
     ]
 
+    signal_name = "TERM"
+    if force:
+        signal_name = "KILL"
+
     for process in processes_to_kill:
-        filter = process[0]
-        if process[1]:
-            format = "pid,comm"
+        keyword, filter_by_cmd = process
+        if filter_by_cmd:
+            ps_format = "pid,comm"
             # According to https://superuser.com/questions/567648/ps-comm-format-always-cuts-the-process-name,  # noqa: E501
             # comm only prints the first 15 characters of the executable name.
-            if len(filter) > 15:
+            if len(keyword) > 15:
                 raise ValueError("The filter string should not be more than" +
                                  " 15 characters. Actual length: " +
-                                 str(len(filter)) + ". Filter: " + filter)
+                                 str(len(keyword)) + ". Filter: " + keyword)
         else:
-            format = "pid,args"
-        command = ("kill -9 $(ps ax -o " + format + " | grep '" + filter +
-                   "' | grep -v grep | " + "awk '{ print $1 }') 2> /dev/null")
+            ps_format = "pid,args"
+        command = (
+            "kill -s {} $(ps ax -o {} | grep {} | grep -v grep | grep ray | "
+            "awk '{{ print $1 }}') 2> /dev/null".format(
+                # ^^ This is how you escape braces in python format string.
+                signal_name,
+                ps_format,
+                keyword))
         subprocess.call([command], shell=True)