Cancel local task processes when a worker is stopped by a command or a SIG[INT|TERM|TSTP] signal

Kobzol · Kobzol · commit f431ae326db0 · 2023-04-19T09:58:22.000Z
diff --git a/crates/tako/src/internal/worker/rpc.rs b/crates/tako/src/internal/worker/rpc.rs
@@ -2,7 +2,7 @@ use std::future::Future;
 use std::net::{Ipv4Addr, SocketAddr};
 use std::rc::Rc;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use bytes::{Bytes, BytesMut};
 use futures::{SinkExt, Stream, StreamExt};
@@ -221,7 +221,8 @@ pub async fn run_worker(
             _ = overview_fut => { unreachable!() }
         };
 
-        match result {
+        // Handle sending stop info to the server and finishing running tasks gracefully.
+        let result = match result {
             Ok(Some(msg)) => {
                 // Worker wants to end gracefully, send message to the server
                 {
@@ -239,11 +240,42 @@ pub async fn run_worker(
                 // Server has disconnected
                 tokio::select! {
                     _ = &mut try_start_tasks => { unreachable!() }
-                    r = finish_tasks_on_server_lost(state) => r
+                    r = finish_tasks_on_server_lost(state.clone()) => r
                 }
                 Err(e)
             }
+        };
+
+        // At this point, there can still be some tasks that are running.
+        // We cancel them here to make sure that we do not leak their spawned processes, if possible.
+        // The futures of the tasks are scheduled onto the current tokio Runtime using spawn_local,
+        // therefore we do not need to await any specific future to drive them forward.
+        // try_start_tasks is not being polled, therefore no new tasks should be started.
+        {
+            let mut state_mut = state.get_mut();
+            for task in state_mut.running_tasks.clone() {
+                state_mut.cancel_task(task);
+            }
         }
+
+        let start = Instant::now();
+        loop {
+            if state.get().running_tasks.is_empty() {
+                break;
+            }
+
+            // Do not wait for the tasks forever
+            if start.elapsed() > Duration::from_secs(5) {
+                break;
+            }
+
+            log::info!("Waiting for tasks to be shut down...");
+            // The await will drive the event loop forward, giving the task futures a chance
+            // to remove themselves from state.running_tasks
+            tokio::time::sleep(Duration::from_secs(1)).await;
+        }
+
+        result
     };
 
     // Provide a local task set for spawning futures
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -64,6 +64,7 @@ def is_process_alive():
                     ]
                     return False
             raise Exception(f"Process with pid {process.pid} not found")
+
         wait_until(lambda: not is_process_alive())
 
     def check_running_processes(self):
@@ -90,13 +91,13 @@ def get_processes_by_name(
             if n == name:
                 yield (i, p)
 
-    def kill_process(self, name: str):
+    def kill_process(self, name: str, signal: int = signal.SIGTERM) -> subprocess.Popen:
         for i, p in self.get_processes_by_name(name):
             del self.processes[i]
             # Kill the whole group since the process may spawn a child
             if p.returncode is None and not p.poll():
-                os.killpg(os.getpgid(p.pid), signal.SIGTERM)
-            return
+                os.killpg(os.getpgid(p.pid), signal)
+            return p
         else:
             raise Exception("Process not found")
 
@@ -220,14 +221,16 @@ def stop_server(self):
     def kill_server(self):
         self.kill_process("server")
 
-    def kill_worker(self, worker_id: int):
+    def kill_worker(self, worker_id: int, signal: int = signal.SIGTERM, wait=True):
         table = self.command(["worker", "info", str(worker_id)], as_table=True)
         pid = table.get_row_value("Process pid")
         process = self.find_process_by_pid(int(pid))
         if process is None:
             raise Exception(f"Worker {worker_id} not found")
 
-        self.kill_process(process[0])
+        process = self.kill_process(process[0], signal=signal)
+        if wait:
+            wait_until(lambda: process.poll() is not None)
 
     def find_process_by_pid(self, pid: int) -> Optional[Tuple[str, subprocess.Popen]]:
         for name, process in self.processes:
diff --git a/tests/pyapi/__init__.py b/tests/pyapi/__init__.py
@@ -1,11 +1,10 @@
 import os.path
 from typing import List, Tuple
 
+from hyperqueue import LocalCluster
 from hyperqueue.client import Client
 from hyperqueue.job import Job
 
-from hyperqueue import LocalCluster
-
 from ..conftest import HqEnv
 from ..utils.mock import ProgramMock
 
diff --git a/tests/pyapi/test_cluster.py b/tests/pyapi/test_cluster.py
@@ -1,4 +1,5 @@
 import pytest
+
 from hyperqueue.cluster import LocalCluster, WorkerConfig
 from hyperqueue.job import Job
 
diff --git a/tests/pyapi/test_function.py b/tests/pyapi/test_function.py
@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import pytest
+
 from hyperqueue.client import Client, FailedJobsException, PythonEnv
 from hyperqueue.ffi.protocol import ResourceRequest
 from hyperqueue.job import Job
diff --git a/tests/pyapi/test_job.py b/tests/pyapi/test_job.py
@@ -4,6 +4,7 @@
 
 import iso8601
 import pytest
+
 from hyperqueue.client import FailedJobsException
 from hyperqueue.ffi.protocol import ResourceRequest
 from hyperqueue.job import Job
diff --git a/tests/test_job.py b/tests/test_job.py
@@ -1,8 +1,11 @@
 import os
+import signal
+import subprocess
 import time
 from datetime import datetime
 from os.path import isdir, isfile, join
 from pathlib import Path
+from typing import Callable, Optional
 
 import pytest
 
@@ -12,7 +15,7 @@
 from .utils.io import check_file_contents, read_file
 from .utils.job import default_task_output, list_jobs
 from .utils.table import parse_multiline_cell
-from .utils.wait import wait_for_pid_exit, wait_until
+from .utils.wait import wait_for_pid_exit, wait_for_worker_state, wait_until
 
 
 def test_job_submit(hq_env: HqEnv):
@@ -468,38 +471,11 @@ def test_cancel_all(hq_env: HqEnv):
 
 
 def test_cancel_terminate_process_children(hq_env: HqEnv):
-    hq_env.start_server()
-    hq_env.start_worker()
+    def cancel(worker_process):
+        hq_env.command(["job", "cancel", "1"])
+        wait_for_job_state(hq_env, 1, "CANCELED")
 
-    hq_env.command(
-        [
-            "submit",
-            "--",
-            *python(
-                """
-import os
-import sys
-import time
-print(os.getpid(), flush=True)
-pid = os.fork()
-if pid > 0:
-    print(pid, flush=True)
-time.sleep(3600)
-"""
-            ),
-        ]
-    )
-    wait_for_job_state(hq_env, 1, "RUNNING")
-    wait_until(lambda: len(read_file(default_task_output()).splitlines()) == 2)
-
-    hq_env.command(["job", "cancel", "1"])
-    wait_for_job_state(hq_env, 1, "CANCELED")
-
-    pids = [int(pid) for pid in read_file(default_task_output()).splitlines()]
-
-    parent, child = pids
-    wait_for_pid_exit(parent)
-    wait_for_pid_exit(child)
+    check_child_process_exited(hq_env, cancel)
 
 
 def test_cancel_send_sigint(hq_env: HqEnv):
@@ -582,8 +558,8 @@ def signal_handler(sig, frame):
 def test_reporting_state_after_worker_lost(hq_env: HqEnv):
     hq_env.start_server()
     hq_env.start_workers(2, cpus=1)
-    hq_env.command(["submit", "sleep", "1"])
-    hq_env.command(["submit", "sleep", "1"])
+    hq_env.command(["submit", "sleep", "2"])
+    hq_env.command(["submit", "sleep", "2"])
 
     wait_for_job_state(hq_env, [1, 2], "RUNNING")
 
@@ -592,15 +568,17 @@ def test_reporting_state_after_worker_lost(hq_env: HqEnv):
     table.check_column_value("State", 1, "RUNNING")
     hq_env.kill_worker(1)
 
-    time.sleep(0.25)
+    def task_is_waiting():
+        table = list_jobs(hq_env)
+        if table.get_column_value("State")[0] == "WAITING":
+            return 0, 1
+        elif table.get_column_value("State")[1] == "WAITING":
+            return 1, 0
+        else:
+            return None
 
+    idx, other = wait_until(task_is_waiting)
     table = list_jobs(hq_env)
-    if table.get_column_value("State")[0] == "WAITING":
-        idx, other = 0, 1
-    elif table.get_column_value("State")[1] == "WAITING":
-        idx, other = 1, 0
-    else:
-        assert 0
     assert table.get_column_value("State")[other] == "RUNNING"
 
     wait_for_job_state(hq_env, other + 1, "FINISHED")
@@ -1422,20 +1400,22 @@ def test_zero_custom_error_message(hq_env: HqEnv):
 
 
 @pytest.mark.parametrize("count", [None, 1, 7])
-def test_crashing_job_status_default(count, hq_env: HqEnv):
+def test_crashing_job_status_default(count: Optional[int], hq_env: HqEnv):
     hq_env.start_server()
 
-    if count:
-        hq_env.command(["submit", f"--crash-limit={count}", "sleep", "10"])
-    else:
-        # Crashing tasks threshold is 5 by default
-        hq_env.command(["submit", "sleep", "10"])
-        count = 5
+    count = count if count is not None else 5
+
+    hq_env.command(["submit", f"--crash-limit={count}", "sleep", "10"])
 
     for i in range(count):
         hq_env.start_worker()
         wait_for_job_state(hq_env, 1, "RUNNING")
         hq_env.kill_worker(i + 1)
+        if i < count - 1:
+            wait_for_job_state(hq_env, 1, "WAITING")
+
+    wait_for_job_state(hq_env, 1, "FAILED")
+
     table = list_jobs(hq_env)
     table.check_column_value("State", 0, "FAILED")
 
@@ -1494,11 +1474,39 @@ def get_pid():
     wait_for_pid_exit(pid)
 
 
-# TODO: fix this somehow
+def test_kill_task_subprocess_when_worker_is_interrupted(hq_env: HqEnv):
+    def interrupt_worker(worker_process):
+        hq_env.kill_worker(1, signal=signal.SIGINT)
+
+    check_child_process_exited(hq_env, interrupt_worker)
+
+
 @pytest.mark.xfail
-def test_kill_task_subprocess_when_worker_dies(hq_env: HqEnv):
+def test_kill_task_subprocess_when_worker_is_terminated(hq_env: HqEnv):
+    def terminate_worker(worker_process):
+        hq_env.kill_worker(1, signal=signal.SIGTERM)
+
+    check_child_process_exited(hq_env, terminate_worker)
+
+
+def test_kill_task_subprocess_when_worker_is_stopped(hq_env: HqEnv):
+    def stop_worker(worker_process):
+        hq_env.command(["worker", "stop", "1"])
+        wait_for_worker_state(hq_env, 1, "STOPPED")
+        hq_env.check_process_exited(worker_process)
+
+    check_child_process_exited(hq_env, stop_worker)
+
+
+def check_child_process_exited(
+    hq_env: HqEnv, stop_fn: Callable[[subprocess.Popen], None]
+):
+    """
+    Creates a task that spawns a child, and then calls `stop_fn`, which should kill either the task
+    or the worker. The function then checks that both the task process and its child have been killed.
+    """
     hq_env.start_server()
-    hq_env.start_worker()
+    worker_process = hq_env.start_worker()
 
     hq_env.command(
         [
@@ -1507,28 +1515,23 @@ def test_kill_task_subprocess_when_worker_dies(hq_env: HqEnv):
             *python(
                 """
 import os
+import sys
 import time
-
-child_pid = os.fork()
-if child_pid == 0:
-    time.sleep(3600)
-else:
-    print(child_pid, flush=True)
+print(os.getpid(), flush=True)
+pid = os.fork()
+if pid > 0:
+    print(pid, flush=True)
 time.sleep(3600)
 """
             ),
         ]
     )
     wait_for_job_state(hq_env, 1, "RUNNING")
+    wait_until(lambda: len(read_file(default_task_output()).splitlines()) == 2)
+    pids = [int(pid) for pid in read_file(default_task_output()).splitlines()]
 
-    def get_pid():
-        pid = read_file(default_task_output()).strip()
-        if not pid:
-            return None
-        return int(pid)
-
-    pid = wait_until(get_pid)
-
-    hq_env.kill_worker(1)
+    stop_fn(worker_process)
 
-    wait_for_pid_exit(pid)
+    parent, child = pids
+    wait_for_pid_exit(parent)
+    wait_for_pid_exit(child)
diff --git a/tests/test_jobfile.py b/tests/test_jobfile.py
@@ -105,7 +105,7 @@ def test_job_file_resource_variants1(hq_env: HqEnv, tmp_path):
         """
 [[task]]
 id = 0
-command = ["sleep", "0"]
+command = ["sleep", "1"]
 
 [[task.request]]
 resources = { "cpus" = "8" }
@@ -116,6 +116,7 @@ def test_job_file_resource_variants1(hq_env: HqEnv, tmp_path):
     )
     hq_env.command(["job", "submit-file", "job.toml"])
 
+    wait_for_job_state(hq_env, 1, "RUNNING")
     table = hq_env.command(["task", "info", "1", "0"], as_table=True)
     table.check_row_value(
         "Resources",
diff --git a/tests/test_stream.py b/tests/test_stream.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`import pytest`
	`2`	`+`
`2`	`3`	`from hyperqueue.cluster import LocalCluster, WorkerConfig`
`3`	`4`	`from hyperqueue.job import Job`
`4`	`5`