Repository: skypilot-org/skypilot
Pull request #4912: [API server] cleanup executor processes on shutdown
Author: login: aylei
Created: 2025-03-07T06:39:02Z | Merged: 2025-03-19T00:47:33Z

<!-- Describe the changes in this PR -->

close https://github.com/skypilot-org/skypilot/issues/4894

SIGTERM handling is relevant not only because `sky cancel`, but also because Kubernetes sends SIGTERM to restart failed server when liveness probe fails. Besides, some frameworks like openkruise on k8s also send SIGTERM to restart the Pod in-place when updating the container image.

This also closes https://github.com/skypilot-org/skypilot/issues/4856 via properly interrupting executors.

Future works:

1. unify the cleanup in `sky api stop` to server-side, then `sky api stop` only sends SIGTERM for consistency
2. handle SIGKILL, typically sent by OOM killer. This requires setting up some local state I think, either a monitor process or a filelock that hints the restarted server there is a uncleaned state and some cleanup work should be taken

@cg505 @Michaelvll PTAL if this makes sense

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Manually tested SIGTERM a server with and without `--foreground`
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`



Diff:
diff --git a/sky/server/requests/executor.py b/sky/server/requests/executor.py
index 5aa977a9139..a16c1c5f9ee 100644
--- a/sky/server/requests/executor.py
+++ b/sky/server/requests/executor.py
@@ -49,7 +49,6 @@
 from sky.utils import common_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
-from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     import types
@@ -221,6 +220,10 @@ def _restore_output(original_stdout: int, original_stderr: int) -> None:
     os.close(original_stderr)
 
 
+def _sigterm_handler(signum: int, frame: Optional['types.FrameType']) -> None:
+    raise KeyboardInterrupt
+
+
 def _request_execution_wrapper(request_id: str,
                                ignore_return_value: bool) -> None:
     """Wrapper for a request execution.
@@ -232,12 +235,8 @@ def _request_execution_wrapper(request_id: str,
     3. Redirect the stdout and stderr of the execution to log file;
     4. Handle the SIGTERM signal to abort the request gracefully.
     """
-
-    def sigterm_handler(signum: int,
-                        frame: Optional['types.FrameType']) -> None:
-        raise KeyboardInterrupt
-
-    signal.signal(signal.SIGTERM, sigterm_handler)
+    # Handle the SIGTERM signal to abort the request processing gracefully.
+    signal.signal(signal.SIGTERM, _sigterm_handler)
 
     pid = multiprocessing.current_process().pid
     logger.info(f'Running request {request_id} with pid {pid}')
@@ -355,6 +354,8 @@ def request_worker(worker: RequestWorker, max_parallel_size: int) -> None:
     Args:
         max_parallel_size: Maximum number of parallel jobs this worker can run.
     """
+    # Handle the SIGTERM signal to abort the executor process gracefully.
+    signal.signal(signal.SIGTERM, _sigterm_handler)
     proc_group = f'{worker.schedule_type.value}-{worker.id}'
     setproctitle.setproctitle(f'SkyPilot:worker:{proc_group}')
     queue = _get_queue(worker.schedule_type)
@@ -388,19 +389,11 @@ def process_request(executor: concurrent.futures.ProcessPoolExecutor):
                 logger.info(f'[{worker}] Finished request: {request_id}')
             else:
                 logger.info(f'[{worker}] Submitted request: {request_id}')
-        except KeyboardInterrupt:
-            # Interrupt the worker process will stop request execution, but
-            # the SIGTERM request should be respected anyway since it might
-            # be explicitly sent by user.
-            # TODO(aylei): crash the API server or recreate the worker process
-            # to avoid broken state.
-            logger.error(f'[{worker}] Worker process interrupted')
-            with ux_utils.print_exception_no_traceback():
-                raise
         except (Exception, SystemExit) as e:  # pylint: disable=broad-except
             # Catch any other exceptions to avoid crashing the worker process.
             logger.error(
-                f'[{worker}] Error processing request {request_id}: '
+                f'[{worker}] Error processing request: '
+                f'{request_id if "request_id" in locals() else ""} '
                 f'{common_utils.format_exception(e, use_bracket=True)}')
 
     # Use concurrent.futures.ProcessPoolExecutor instead of multiprocessing.Pool
@@ -409,12 +402,33 @@ def process_request(executor: concurrent.futures.ProcessPoolExecutor):
     # We use executor instead of individual multiprocessing.Process to avoid
     # the overhead of forking a new process for each request, which can be about
     # 1s delay.
-    with concurrent.futures.ProcessPoolExecutor(
+    try:
+        executor = concurrent.futures.ProcessPoolExecutor(
             max_workers=max_parallel_size,
             initializer=executor_initializer,
-            initargs=(proc_group,)) as executor:
+            initargs=(proc_group,))
         while True:
             process_request(executor)
+    # TODO(aylei): better to distinct between KeyboardInterrupt and SIGTERM.
+    except KeyboardInterrupt:
+        pass
+    finally:
+        # In most cases, here we receive either ctrl-c in foreground execution
+        # or SIGTERM on server exiting. Gracefully exit the worker process and
+        # the executor.
+        # TODO(aylei): worker may also be killed by system daemons like OOM
+        # killer, crash the API server or recreate the worker process to avoid
+        # broken state in such cases.
+        logger.info(f'[{worker}] Worker process interrupted')
+        executor_processes = list(executor._processes.values())  # pylint: disable=protected-access,line-too-long
+        # Shutdown the executor so that executor process can exit once the
+        # running task is finished or interrupted.
+        executor.shutdown(wait=False)
+        # Proactively interrupt the running task to avoid indefinite waiting.
+        subprocess_utils.run_in_parallel(
+            subprocess_utils.kill_process_with_grace_period,
+            executor_processes,
+            num_threads=len(executor_processes))
 
 
 def start(deploy: bool) -> List[multiprocessing.Process]:
diff --git a/sky/server/server.py b/sky/server/server.py
index fa81b3d3384..cc1bec24a1a 100644
--- a/sky/server/server.py
+++ b/sky/server/server.py
@@ -1140,6 +1140,9 @@ def cleanup(proc: multiprocessing.Process) -> None:
                 # The process may not be started yet, close it anyway.
                 proc.close()
 
+        # Terminate processes in reverse order in case dependency, especially
+        # queue server. Terminate queue server first does not affect the
+        # correctness of cleanup but introduce redundant error messages.
         subprocess_utils.run_in_parallel(cleanup,
-                                         sub_procs,
+                                         list(reversed(sub_procs)),
                                          num_threads=len(sub_procs))
diff --git a/sky/utils/subprocess_utils.py b/sky/utils/subprocess_utils.py
index b8eae506733..f0bafd8d3af 100644
--- a/sky/utils/subprocess_utils.py
+++ b/sky/utils/subprocess_utils.py
@@ -1,4 +1,5 @@
 """Utility functions for subprocesses."""
+import multiprocessing
 from multiprocessing import pool
 import os
 import random
@@ -181,29 +182,6 @@ def kill_children_processes(parent_pids: Optional[Union[
     if isinstance(parent_pids, int):
         parent_pids = [parent_pids]
 
-    def kill(proc: psutil.Process):
-        if not proc.is_running():
-            # Skip if the process is not running.
-            return
-        logger.debug(f'Killing process {proc.pid}')
-        try:
-            if force:
-                proc.kill()
-            else:
-                proc.terminate()
-            proc.wait(timeout=10)
-        except psutil.NoSuchProcess:
-            # The child process may have already been terminated.
-            pass
-        except psutil.TimeoutExpired:
-            logger.debug(
-                f'Process {proc.pid} did not terminate after 10 seconds')
-            # Attempt to force kill if the normal termination fails
-            if not force:
-                logger.debug(f'Force killing process {proc.pid}')
-                proc.kill()
-                proc.wait(timeout=5)  # Shorter timeout after force kill
-
     parent_processes = []
     if parent_pids is None:
         parent_processes = [psutil.Process()]
@@ -218,10 +196,54 @@ def kill(proc: psutil.Process):
     for parent_process in parent_processes:
         child_processes = parent_process.children(recursive=True)
         if parent_pids is not None:
-            kill(parent_process)
+            kill_process_with_grace_period(parent_process, force=force)
         logger.debug(f'Killing child processes: {child_processes}')
         for child in child_processes:
-            kill(child)
+            kill_process_with_grace_period(child, force=force)
+
+
+def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
+                                               psutil.Process],
+                                   force: bool = False,
+                                   grace_period: int = 10) -> None:
+    """Kill a process with SIGTERM and wait for it to exit.
+
+    Args:
+        proc: The process to kill, either a multiprocessing.Process or a
+            psutil.Process.
+        force: Whether to force kill the process.
+        grace_period: The grace period seconds to wait for the process to exit.
+    """
+    if isinstance(proc, psutil.Process):
+        alive = proc.is_running
+        wait = proc.wait
+    else:
+        alive = proc.is_alive
+        wait = proc.join
+    if not alive():
+        # Skip if the process is not running.
+        return
+    logger.debug(f'Killing process {proc.pid}')
+    try:
+        if force:
+            proc.kill()
+        else:
+            proc.terminate()
+        wait(timeout=grace_period)
+    except (psutil.NoSuchProcess, ValueError):
+        # The child process may have already been terminated.
+        return
+    except psutil.TimeoutExpired:
+        # Pass to finally to force kill the process.
+        pass
+    finally:
+        logger.debug(f'Process {proc.pid} did not terminate after '
+                     f'{grace_period} seconds')
+        # Attempt to force kill if the normal termination fails
+        if not force:
+            logger.debug(f'Force killing process {proc.pid}')
+            # Shorter timeout after force kill
+            kill_process_with_grace_period(proc, force=True, grace_period=5)
 
 
 def run_with_retries(
diff --git a/tests/unit_tests/sky/utils/test_subprocess_utils.py b/tests/unit_tests/sky/utils/test_subprocess_utils.py
index 710ad2636dc..d9bc88c50f3 100644
--- a/tests/unit_tests/sky/utils/test_subprocess_utils.py
+++ b/tests/unit_tests/sky/utils/test_subprocess_utils.py
@@ -1,8 +1,12 @@
 """Unit tests for subprocess_utils.py."""
 import logging
+import multiprocessing
+import signal
 import time
+import unittest
 from unittest import mock
 
+import psutil
 import pytest
 
 from sky.utils import subprocess_utils
@@ -158,3 +162,129 @@ def mock_sleep_side_effect(delay):
     # Start 1 at a time
     expected_started_counts = [1, 2, 3, 4]
     assert batch_sizes == expected_started_counts
+
+
+def _ignore_sigterm():
+    """Process function that ignores SIGTERM signal."""
+    signal.signal(signal.SIGTERM, signal.SIG_IGN)
+    time.sleep(30)
+
+
+class TestKillProcessWithGrace(unittest.TestCase):
+    """Test cases for kill_process_with_grace function."""
+
+    def setUp(self):
+        """Set up test cases."""
+        self.processes = []
+
+    def tearDown(self):
+        """Clean up any remaining processes."""
+        for process in self.processes:
+            try:
+                if isinstance(process, psutil.Process):
+                    if process.is_running():
+                        process.kill()
+                else:
+                    if process.is_alive():
+                        process.kill()
+                # Wait for the process to be fully terminated
+                time.sleep(0.1)
+            except (psutil.NoSuchProcess, ValueError):
+                pass
+
+    def _create_dummy_process(self):
+        """Create a dummy process that sleeps."""
+        process = multiprocessing.Process(target=time.sleep, args=(30,))
+        self.processes.append(process)
+        return process
+
+    def _wait_for_process_start(self, process, timeout=1):
+        """Wait for process to start and stabilize."""
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            try:
+                if isinstance(process, psutil.Process):
+                    if process.is_running():
+                        return True
+                else:
+                    if process.is_alive():
+                        return True
+            except (psutil.NoSuchProcess, ValueError):
+                pass
+            time.sleep(0.1)
+        return False
+
+    def test_normal_termination(self):
+        """Test normal process termination with SIGTERM."""
+        process = self._create_dummy_process()
+        process.start()
+        self.assertTrue(self._wait_for_process_start(process))
+
+        # Kill the process with grace period
+        subprocess_utils.kill_process_with_grace_period(process,
+                                                        force=False,
+                                                        grace_period=2)
+
+        # Process should be terminated
+        self.assertFalse(process.is_alive())
+
+    def test_force_kill(self):
+        """Test force kill with SIGKILL."""
+        process = self._create_dummy_process()
+        process.start()
+        self.assertTrue(self._wait_for_process_start(process))
+
+        # Force kill the process
+        subprocess_utils.kill_process_with_grace_period(process,
+                                                        force=True,
+                                                        grace_period=2)
+
+        # Process should be terminated
+        self.assertFalse(process.is_alive())
+
+    def test_psutil_process(self):
+        """Test killing process using psutil.Process."""
+        process = self._create_dummy_process()
+        process.start()
+        self.assertTrue(self._wait_for_process_start(process))
+
+        psutil_proc = psutil.Process(process.pid)
+        self.processes.append(psutil_proc)
+        self.assertTrue(self._wait_for_process_start(psutil_proc))
+
+        # Kill the process with grace period
+        subprocess_utils.kill_process_with_grace_period(psutil_proc,
+                                                        force=False,
+                                                        grace_period=2)
+
+        self.assertFalse(psutil_proc.is_running())
+
+    def test_already_dead(self):
+        """Test killing an already dead process."""
+        process = self._create_dummy_process()
+        process.start()
+        self.assertTrue(self._wait_for_process_start(process))
+        process.terminate()
+        time.sleep(0.1)  # Give some time for the process to terminate
+
+        # This should not raise any exception
+        subprocess_utils.kill_process_with_grace_period(process,
+                                                        force=False,
+                                                        grace_period=1)
+        self.assertFalse(process.is_alive())
+
+    def test_timeout(self):
+        """Test process that ignores SIGTERM and requires SIGKILL."""
+        process = multiprocessing.Process(target=_ignore_sigterm)
+        self.processes.append(process)
+        process.start()
+        self.assertTrue(self._wait_for_process_start(process))
+
+        # Try to kill with SIGTERM first
+        subprocess_utils.kill_process_with_grace_period(process,
+                                                        force=False,
+                                                        grace_period=1)
+
+        # Process should be terminated by SIGKILL after grace period
+        time.sleep(0.1)  # Give some time for the process to be fully terminated
+        self.assertFalse(process.is_alive())


Reviews:
user: login: cg505
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:19:15Z
user: login: cg505
body: Thanks for the change, I think this is good. @Michaelvll should review as well.

> 2. handle SIGKILL, typically sent by OOM killer. This requires setting up some local state I think, either a monitor process or a filelock that hints the restarted server there is a uncleaned state and some cleanup work should be taken

Another option could be that the worker itself periodically checks if the api server PID is still alive, and exits if not. This could avoid needing a separate monitor process. Just an idea, not sure what's the best option.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-10T18:26:10Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T09:03:44Z
user: login: Michaelvll
body: Thanks @aylei! LGTM.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-13T00:06:17Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-13T00:52:47Z

Review comments:
diff_hunk: @@ -385,12 +380,32 @@ def process_request(executor: concurrent.futures.ProcessPoolExecutor):
     # We use executor instead of individual multiprocessing.Process to avoid
     # the overhead of forking a new process for each request, which can be about
     # 1s delay.
-    with concurrent.futures.ProcessPoolExecutor(
+    try:
+        executor = concurrent.futures.ProcessPoolExecutor(
             max_workers=max_parallel_size,
             initializer=executor_initializer,
-            initargs=(proc_group,)) as executor:
+            initargs=(proc_group,))
         while True:
             process_request(executor)
+    except KeyboardInterrupt:
path: sky/server/requests/executor.py
user: login: cg505
body: Seems that ideally we shouldn't use KeyboardInterrupt for the worker SIGTERM. It makes sense for request execution since all the core code is designed to work with KeyboardInterrupt, but here we can fully control the exception type we throw and catch. We could use a custom exception.
Super low priority, probably better to just merge this. Could add a TODO comment.
created_at: 2025-03-10T18:19:15Z
author_association: COLLABORATOR
diff_hunk: @@ -385,12 +380,32 @@ def process_request(executor: concurrent.futures.ProcessPoolExecutor):
     # We use executor instead of individual multiprocessing.Process to avoid
     # the overhead of forking a new process for each request, which can be about
     # 1s delay.
-    with concurrent.futures.ProcessPoolExecutor(
+    try:
+        executor = concurrent.futures.ProcessPoolExecutor(
             max_workers=max_parallel_size,
             initializer=executor_initializer,
-            initargs=(proc_group,)) as executor:
+            initargs=(proc_group,))
         while True:
             process_request(executor)
+    except KeyboardInterrupt:
path: sky/server/requests/executor.py
user: login: aylei
body: Sure! My original intent is to handle Ctrl-C (in `--foreground` mode) and SIGTERM from parent process in a uniform way, but agree it is better to use a distinct exception for SIGTERM
created_at: 2025-03-11T09:03:44Z
author_association: COLLABORATOR
diff_hunk: @@ -216,10 +194,54 @@ def kill(proc: psutil.Process):
     for parent_process in parent_processes:
         child_processes = parent_process.children(recursive=True)
         if parent_pids is not None:
-            kill(parent_process)
+            kill_process_with_grace(parent_process, force=force)
         logger.debug(f'Killing child processes: {child_processes}')
         for child in child_processes:
-            kill(child)
+            kill_process_with_grace(child, force=force)
+
+
+def kill_process_with_grace(proc: Union[multiprocessing.Process,
path: sky/utils/subprocess_utils.py
user: login: Michaelvll
body: ```suggestion
def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
```
created_at: 2025-03-12T18:41:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,139 @@
+"""Unit tests for subprocess_utils.py."""
path: tests/unit_tests/sky/test_subprocess_utils.py
user: login: Michaelvll
body: Should this be under `tests/unit_tests/sky/utils/`?
created_at: 2025-03-12T18:43:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,139 @@
+"""Unit tests for subprocess_utils.py."""
path: tests/unit_tests/sky/test_subprocess_utils.py
user: login: aylei
body: Ah, good catch
created_at: 2025-03-13T00:52:47Z
author_association: COLLABORATOR

PR comments:
user: login: Michaelvll
created_at: 2025-03-13T00:07:53Z
body: /smoke-test -k test_multi_echo
/smoke-test -k test_cancel_pytorch
/smoke-test -k test_cancel_aws --aws
/smoke-test -k test_job_queue
author_association: COLLABORATOR
user: login: aylei
created_at: 2025-03-13T01:08:55Z
body: /smoke-test -k test_multi_echo
/smoke-test -k test_cancel_pytorch
/smoke-test -k test_cancel_aws --aws
/smoke-test -k test_job_queue
author_association: COLLABORATOR
user: login: aylei
created_at: 2025-03-13T02:25:14Z
body: /smoke-test -k test_job_queue
author_association: COLLABORATOR
user: login: aylei
created_at: 2025-03-18T16:39:46Z
body: /smoke-test -k test_multi_echo
/smoke-test -k test_cancel_pytorch
/smoke-test -k test_cancel_aws --aws
/smoke-test -k test_job_queue
author_association: COLLABORATOR