Repository: skypilot-org/skypilot
Pull request: Fix sky cancel for on-prem mode
~~`sky cancel` fails to kill the user processes when on-prem mode is used as `ps --forest` does not list processes with different users.~~
~~Replace it with `pstree` to find all descendent processes.~~


~~Also, replace `kill -s 0` with `ps -p` to check the liveness of process as `kill -s 0` can [return 1](https://stackoverflow.com/questions/3043978/how-to-check-if-a-process-id-pid-exists) even if the process is running.~~

-----
`sky cancel` fails to kill the user processes when on-prem mode is used as `ps --forest` does not list processes with different users.
This PR implements a python based subprocess_daemon with [psutil](https://psutil.readthedocs.io/en/latest/#) for 1) readability and 2) alternative command `pstree` is not available on GCP.  
However, the problem for on-prem is not fully resolved because `sky cancel` still fails to kill the parent process with `ray stop` for some reason. might need @michaelzhiluo to debug this.

-----

This PR also fixed https://github.com/sky-proj/sky/issues/835 by checking if the target process exists or not to avoid generating infinite processes in sky spot case.

TODOs
- [x] Add cancel tests for three clouds
- [x] Pass cancel tests for three clouds 
- [ ] Add a cancel test for on-prem (this is a bit complicated as we don't have an automatic script for setting up on-prem yet, cc @michaelzhiluo )

Problem statement:
potential performance issue with sky spot controller
There seems to be a performance issue of sky spot controller.
I've been using it for 3 days and submitted 20+ jobs. In the beginning it was pretty smooth but now it becomes super slow.
Now it's stuck for 5mins+ when I tried `sky spot status` or `sky spot logs`.
```
(sky) weichiang@blaze:~/repos/sky$ sky spot status
Fetching managed spot job statuses...
```
I login to the machine and cpu loading seems to be higher than #core. It might be related to some daemon job leaks.
Still investigating.

Hints from issue discussion:
Looks like the performance issue is critical. now the recovery mechanism is not functioning well.
What process takes most of the CPUs? Could you have a screenshot of the process list sorted by the CPU utilization?

Commit message: Fix sky spot issue

Patch to review:
diff --git a/examples/resnet_app.yaml b/examples/resnet_app.yaml
index b21e17791f8..88e8a153992 100644
--- a/examples/resnet_app.yaml
+++ b/examples/resnet_app.yaml
@@ -2,9 +2,7 @@ name: resnet-app
 
 resources:
   cloud: aws
-  instance_type: p3.2xlarge
-  # cloud: azure
-  # accelerators: V100
+  accelerators: V100
 
 inputs: {
   gs://cloud-tpu-test-dataset/fake_imagenet: 70,
diff --git a/examples/resnet_distributed_tf_app.py b/examples/resnet_distributed_tf_app.py
index ea16c875f35..db1979b8565 100644
--- a/examples/resnet_distributed_tf_app.py
+++ b/examples/resnet_distributed_tf_app.py
@@ -13,20 +13,13 @@
 cluster = f'test-distributed-tf-{_user_and_host}'
 
 with sky.Dag() as dag:
-    # The working directory contains all code and will be synced to remote.
-    workdir = '~/Downloads/tpu'
-    subprocess.run(
-        'cd ~/Downloads; '
-        '(git clone https://github.com/concretevitamin/tpu || true); '
-        f'cd {workdir} && git checkout 9459fee',
-        shell=True,
-        check=True)
-
     # Total Nodes, INCLUDING Head Node
     num_nodes = 2
 
     # The setup command.  Will be run under the working directory.
     setup = """
+            git clone https://github.com/concretevitamin/tpu || true
+            cd tpu && git checkout 9459fee
             conda create -n resnet python=3.7 -y
             conda activate resnet
             conda install cudatoolkit=11.0 -y
@@ -54,6 +47,7 @@ def run_fn(node_rank: int, ip_list: List[str]) -> Optional[str]:
         str_tf_config = json.dumps(tf_config)
         print(f'{str_tf_config!r}')
         run = f"""
+            cd tpu
             conda activate resnet
             rm -rf resnet_model-dir
             export TF_CONFIG={str_tf_config!r}
@@ -68,7 +62,6 @@ def run_fn(node_rank: int, ip_list: List[str]) -> Optional[str]:
 
     train = sky.Task(
         'train',
-        workdir=workdir,
         setup=setup,
         num_nodes=num_nodes,
         run=run_fn,
diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
index af4aea6115b..d6937efdb74 100644
--- a/sky/setup_files/setup.py
+++ b/sky/setup_files/setup.py
@@ -45,6 +45,7 @@
     # polling strategies`
     'grpcio<=1.43.0',
     'pulp',
+    'psutil',
 ]
 
 extras_require = {
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 0e87ecbce85..0606a4c5d93 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -145,13 +145,14 @@ def run_with_log(
         parent_pid = os.getpid()
         daemon_script = os.path.join(
             os.path.dirname(os.path.abspath(job_lib.__file__)),
-            'subprocess_daemon.sh')
+            'subprocess_daemon.py')
         daemon_cmd = [
-            '/bin/bash',
+            'python',
             daemon_script,
+            '--parent-pid',
             str(parent_pid),
+            '--proc-pid',
             str(proc.pid),
-            str(int(with_ray)),
         ]
         subprocess.Popen(
             daemon_cmd,
diff --git a/sky/skylet/subprocess_daemon.py b/sky/skylet/subprocess_daemon.py
new file mode 100644
index 00000000000..f00cb3c6a95
--- /dev/null
+++ b/sky/skylet/subprocess_daemon.py
@@ -0,0 +1,49 @@
+"""Sky subprocess daemon.
+
+Wait for parent_pid to exit, then SIGTERM (or SIGKILL if needed) the child
+processes of proc_pid.
+"""
+
+import psutil
+import argparse
+import time
+import sys
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--parent-pid', type=int, required=True)
+    parser.add_argument('--proc-pid', type=int, required=True)
+    args = parser.parse_args()
+
+    parent_process = psutil.Process(args.parent_pid)
+    process = psutil.Process(args.proc_pid)
+
+    if process is None or parent_process is None:
+        sys.exit()
+
+    # Fetch children first to avoid target process termination.
+    children = process.children(recursive=True)
+    children.append(process)
+
+    # Wait for either parent or target process to exit.
+    while True:
+        time.sleep(1)
+        if not process.is_running() or not parent_process.is_running():
+            break
+
+    for pid in children:
+        try:
+            pid.terminate()
+        except psutil.NoSuchProcess:
+            pass
+
+    # Wait 30s for the processes to exit gracefully.
+    time.sleep(30)
+
+    # SIGKILL if they're still running.
+    for pid in children:
+        try:
+            pid.kill()
+        except psutil.NoSuchProcess:
+            pass
diff --git a/sky/skylet/subprocess_daemon.sh b/sky/skylet/subprocess_daemon.sh
deleted file mode 100644
index 87ae8b3bca2..00000000000
--- a/sky/skylet/subprocess_daemon.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/bin/bash
-# Busy-waits for parent_pid to exit, then SIGTERM the child processes of proc_pid and SIGKILL proc_pid.
-
-parent_pid=$1
-proc_pid=$2
-remote=${3:-0}
-
-while kill -s 0 ${parent_pid}; do sleep 1; done
-
-if [ ${remote} -eq 0 ]; then
-    # This is to avoid the PIPE outputing to the console after being killed in next line.
-    pkill -PIPE -P ${proc_pid}
-fi
-if [ $OSTYPE == 'darwin'* ]; then
-    # MacOS do not have the ps --forest option, it is enough to do pkill here. (pkill
-    # only kills the first level descendant of the process, not recursively find all
-    # of the descendants.)
-    pkill -TERM -P ${proc_pid}
-    sleep 5
-    kill -9 ${proc_pid}
-else
-    # Uses pgid instead of sid, since sid can be changed by other processes.
-    pgid=`ps -o pgid= -p ${proc_pid}`
-    # Recursively gracefully kill (SIGTERM) all child processes of proc_pid.
-    # We should not run this command if pgid is empty, i.e. the process is already dead.
-    [[ -z "${pgid}" ]] || (ps --forest -o pid -g $pgid | tail -n +2 | xargs kill -15)
-    # Wait 30s for the processes to exit gracefully.
-    sleep 30
-    [[ -z "${pgid}" ]] || (ps --forest -o pid -g $pgid | tail -n +2 | xargs kill -9)
-fi
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 7b81daa9538..90f45de70f4 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -248,6 +248,7 @@ def test_tpu():
             f'sky launch -y -c {name} examples/tpu_app.yaml | grep "TPU .* already exists"',  # Ensure sky launch won't create another TPU.
         ],
         f'sky down -y {name}',
+        timeout=30 * 60,  # can take >20 mins
     )
     run_one_test(test)
 
@@ -346,24 +347,44 @@ def test_autostop():
     run_one_test(test)
 
 
-# ---------- Testing `sky cancel` ----------
-def test_cancel():
-    name = _get_cluster_name()
+def _get_cancel_task_with_cloud(name, cloud):
     test = Test(
-        'cancel',
+        f'{cloud}-cancel-task',
         [
-            f'sky launch -c {name} examples/resnet_app.yaml -y -d',
+            f'sky launch -c {name} examples/resnet_app.yaml --cloud {cloud} -y -d',
             # Wait the GPU process to start.
             'sleep 60',
             f'sky exec {name} "nvidia-smi | grep python"',
             f'sky logs {name} 2 --status',  # Ensure the job succeeded.
             f'sky cancel {name} 1',
             'sleep 60',
-            f'sky exec {name} "nvidia-smi | grep \'No running process\'"',
+            # check if the python job is gone.
+            f'sky exec {name} "! nvidia-smi | grep python"',
             f'sky logs {name} 3 --status',  # Ensure the job succeeded.
         ],
         f'sky down -y {name}',
     )
+    return test
+
+
+# ---------- Testing `sky cancel` on AWS ----------
+def test_cancel_aws():
+    name = _get_cluster_name()
+    test = _get_cancel_task_with_cloud(name, 'aws')
+    run_one_test(test)
+
+
+# ---------- Testing `sky cancel` on Azure ----------
+def test_cancel_azure():
+    name = _get_cluster_name()
+    test = _get_cancel_task_with_cloud(name, 'azure')
+    run_one_test(test)
+
+
+# ---------- Testing `sky cancel` on GCP ----------
+def test_cancel_gcp():
+    name = _get_cluster_name()
+    test = _get_cancel_task_with_cloud(name, 'gcp')
     run_one_test(test)
 
 


Code review comments:
diff_hunk: @@ -45,6 +45,7 @@
     # polling strategies`
     'grpcio<=1.43.0',
     'pulp',
+    'psutil',
path: sky/setup_files/setup.py
text: @user1:
nit: order before pulp