Repository: PrefectHQ/prefect

Issue:
Handle SIGTERM received by prefect-agent gracefully
This issue is about allowing flow runs some time to finish when prefect-agent receives a SIGTERM signal from user/docker/kubernetes, specifically deployed in combination with a horizontal pod autoscaler.

The use case of a k8s user wanting to (auto)scale up their data pipelines using prefect:
- easiest would be to run `prefect-agent` Deployment behind a HPA (like is currently already possible in the [orion server helm chart](https://github.com/PrefectHQ/prefect-helm/blob/2022.12.22/charts/prefect-orion/values.yaml#L41-L53))
  - different worker Deployments for different queues, depending on average mem usage, cpu-bound/io-bound, etc
- set the `terminationGracePeriodSeconds` of the agent to e.g. ~2x the average duration of a flow run in its queue
  - when a Pod is scaled down in kubernetes, k8s sends a SIGTERM to the container(s) main PID(s). then it waits `terminationGracePeriodSeconds` to send a SIGKILL to get rid of the Pod. this guy [explains it better](https://cloud.google.com/blog/products/containers-kubernetes/kubernetes-best-practices-terminating-with-grace)
- when the main agent PID receives SIGTERM, allow the flow runs to finish gracefully
- if somehow a flow run is not done after `terminationGracePeriodSeconds`, the main gets a SIGKILL (like when k8s kills the main PID due to OOM).
  - in the case of for instance the AMQP protocol, an orphaned/zombie task would be marked as open again in the queue and another worker would pick it up.
  - marking as crashed is the other alternative, and the user would add a limited amount of retries in case of undefined crash (lost connection or whatever). if there was an actual python Exception, that's a different case and the user might not want to retry.

e.g. FastAPI works like this: it will allow ongoing requests to finish and will gracefully close the loop when a SIGTERM is received. k8s counts on this behaviour in the scale-down mechanic of the Deployment/StatefulSet/ReplicaSet.

there is also still the issue of k8s SIGKILLing subprocesses to avoid the container from going OOM, which I think is relevant for this particular issue (and the CRASH detection) as well: 
> on another note (agent pod, not orion server), we just watched an agent's subprocess (a flow run) get OOM SIGKILLed by kubernetes. interestingly, it only kills the subprocess ([ref](https://discuss.kubernetes.io/t/subprocess-killed-with-a-137-error/22058/4)) and the main process hangs (configured to sequentially handle only 1 flow run at a time) ~indefinitely and~ (nvm, it takes the next task) happily. 
>
> probably doesn't warrant an issue or so, but making sure you're aware of this behaviour of k8s
> 
> ```
> <frozen runpy>:128: RuntimeWarning: 'prefect.engine' found in sys.modules after import of package 'prefect', but prior to execution of 'prefect.engine'; this may result in unpredictable behaviour
> 19:44:53.133 | INFO | Flow run 'phi2-kelis-m' - Created task run 'parse_crawl-259eac6e-0' for task 'parse_crawl'
> 19:44:53.136 | INFO | Flow run 'phi2-kelis-m' - Executing 'parse_crawl-259eac6e-0' immediately...
> 19:45:15.900 | ERROR | prefect.infrastructure.process - Process 'process-0-2-0rc4' exited with status code: -9; This indicates that the process exited due to a SIGKILL signal. Typically, this is caused by high memory usage causing the operating system to terminate the process.
> 19:45:15.900 | INFO | prefect.agent - Reporting flow run 'fd0e2945-6994-4eee-87c4-57ed659f2bf4' as crashed due to non-zero status code.
> ```
https://github.com/PrefectHQ/prefect/pull/7948#issuecomment-1380932813

_Originally posted by @ddelange in https://github.com/PrefectHQ/prefect/issues/7800#issuecomment-1386911707_
            


Discussion:

Resolving patch:
diff --git a/src/prefect/cli/agent.py b/src/prefect/cli/agent.py
--- a/src/prefect/cli/agent.py
+++ b/src/prefect/cli/agent.py
@@ -1,6 +1,7 @@
 """
 Command line interface for working with agent services
 """
+import os
 from functools import partial
 from typing import List
 from uuid import UUID
@@ -20,6 +21,7 @@
     PREFECT_AGENT_QUERY_INTERVAL,
     PREFECT_API_URL,
 )
+from prefect.utilities.processutils import setup_signal_handlers_agent
 from prefect.utilities.services import critical_service_loop
 
 agent_app = PrefectTyper(
@@ -161,6 +163,11 @@ async def start(
                 f"Starting v{prefect.__version__} agent with ephemeral API..."
             )
 
+    agent_process_id = os.getpid()
+    setup_signal_handlers_agent(
+        agent_process_id, "the Prefect agent", app.console.print
+    )
+
     async with PrefectAgent(
         work_queues=work_queues,
         work_queue_prefix=work_queue_prefix,
diff --git a/src/prefect/utilities/processutils.py b/src/prefect/utilities/processutils.py
--- a/src/prefect/utilities/processutils.py
+++ b/src/prefect/utilities/processutils.py
@@ -320,12 +320,20 @@ def forward_signal_handler(
     """Forward subsequent signum events (e.g. interrupts) to respective signums."""
     current_signal, future_signals = signums[0], signums[1:]
 
+    # avoid RecursionError when setting up a direct signal forward to the same signal for the main pid
+    avoid_infinite_recursion = signum == current_signal and pid == os.getpid()
+    if avoid_infinite_recursion:
+        # store the vanilla handler so it can be temporarily restored below
+        original_handler = signal.getsignal(current_signal)
+
     def handler(*args):
         print_fn(
             f"Received {getattr(signum, 'name', signum)}. "
             f"Sending {getattr(current_signal, 'name', current_signal)} to"
             f" {process_name} (PID {pid})..."
         )
+        if avoid_infinite_recursion:
+            signal.signal(current_signal, original_handler)
         os.kill(pid, current_signal)
         if future_signals:
             forward_signal_handler(
@@ -345,12 +353,31 @@ def setup_signal_handlers_server(pid: int, process_name: str, print_fn: Callable
     setup_handler = partial(
         forward_signal_handler, pid, process_name=process_name, print_fn=print_fn
     )
-    # on Windows, use CTRL_BREAK_EVENT as SIGTERM is useless:
-    # https://bugs.python.org/issue26350
+    # when server receives a signal, it needs to be propagated to the uvicorn subprocess
     if sys.platform == "win32":
+        # on Windows, use CTRL_BREAK_EVENT as SIGTERM is useless:
+        # https://bugs.python.org/issue26350
         setup_handler(signal.SIGINT, signal.CTRL_BREAK_EVENT)
     else:
         # first interrupt: SIGTERM, second interrupt: SIGKILL
         setup_handler(signal.SIGINT, signal.SIGTERM, signal.SIGKILL)
         # forward first SIGTERM directly, send SIGKILL on subsequent SIGTERM
         setup_handler(signal.SIGTERM, signal.SIGTERM, signal.SIGKILL)
+
+
+def setup_signal_handlers_agent(pid: int, process_name: str, print_fn: Callable):
+    """Handle interrupts of the agent gracefully."""
+    setup_handler = partial(
+        forward_signal_handler, pid, process_name=process_name, print_fn=print_fn
+    )
+    # when agent receives SIGINT, it stops dequeueing new FlowRuns, and runs until the subprocesses finish
+    # the signal is not forwarded to subprocesses, so they can continue to run and hopefully still complete
+    if sys.platform == "win32":
+        # on Windows, use CTRL_BREAK_EVENT as SIGTERM is useless:
+        # https://bugs.python.org/issue26350
+        setup_handler(signal.SIGINT, signal.CTRL_BREAK_EVENT)
+    else:
+        # forward first SIGINT directly, send SIGKILL on subsequent interrupt
+        setup_handler(signal.SIGINT, signal.SIGINT, signal.SIGKILL)
+        # first SIGTERM: send SIGINT, send SIGKILL on subsequent SIGTERM
+        setup_handler(signal.SIGTERM, signal.SIGINT, signal.SIGKILL)


Test patch:
diff --git a/tests/cli/test_start_agent.py b/tests/cli/test_start_agent.py
new file mode 100644
--- /dev/null
+++ b/tests/cli/test_start_agent.py
@@ -0,0 +1,170 @@
+import os
+import signal
+import sys
+import tempfile
+
+import anyio
+import pytest
+
+from prefect.settings import get_current_settings
+from prefect.utilities.processutils import open_process
+
+POLL_INTERVAL = 0.5
+STARTUP_TIMEOUT = 20
+SHUTDOWN_TIMEOUT = 5
+
+
+async def safe_shutdown(process):
+    try:
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await process.wait()
+    except TimeoutError:
+        # try twice in case process.wait() hangs
+        with anyio.fail_after(SHUTDOWN_TIMEOUT):
+            await process.wait()
+
+
+@pytest.fixture(scope="function")
+async def agent_process(use_hosted_api_server):
+    """
+    Runs an agent listening to all queues.
+    Yields:
+        The anyio.Process.
+    """
+    out = tempfile.TemporaryFile()  # capture output for test assertions
+
+    # Will connect to the same database as normal test clients
+    async with open_process(
+        command=[
+            "prefect",
+            "agent",
+            "start",
+            "--match=nonexist",
+        ],
+        stdout=out,
+        stderr=out,
+        env={**os.environ, **get_current_settings().to_environment_variables()},
+    ) as process:
+        process.out = out
+
+        for _ in range(int(STARTUP_TIMEOUT / POLL_INTERVAL)):
+            await anyio.sleep(POLL_INTERVAL)
+            if out.tell() > 400:
+                await anyio.sleep(2)
+                break
+
+        assert out.tell() > 400, "The agent did not start up in time"
+        assert process.returncode is None, "The agent failed to start up"
+
+        # Yield to the consuming tests
+        yield process
+
+        # Then shutdown the process
+        try:
+            process.terminate()
+        except ProcessLookupError:
+            pass
+        out.close()
+
+
+class TestAgentSignalForwarding:
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigint_sends_sigterm(self, agent_process):
+        agent_process.send_signal(signal.SIGINT)
+        await safe_shutdown(agent_process)
+        agent_process.out.seek(0)
+        out = agent_process.out.read().decode()
+
+        assert "Sending SIGINT" in out, (
+            "When sending a SIGINT, the main process should receive a SIGINT."
+            f" Output:\n{out}"
+        )
+        assert "Agent stopped!" in out, (
+            "When sending a SIGINT, the main process should shutdown gracefully."
+            f" Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigterm_sends_sigterm_directly(self, agent_process):
+        agent_process.send_signal(signal.SIGTERM)
+        await safe_shutdown(agent_process)
+        agent_process.out.seek(0)
+        out = agent_process.out.read().decode()
+
+        assert "Sending SIGINT" in out, (
+            "When sending a SIGTERM, the main process should receive a SIGINT."
+            f" Output:\n{out}"
+        )
+        assert "Agent stopped!" in out, (
+            "When sending a SIGTERM, the main process should shutdown gracefully."
+            f" Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigint_sends_sigterm_then_sigkill(self, agent_process):
+        agent_process.send_signal(signal.SIGINT)
+        await anyio.sleep(0.01)  # some time needed for the recursive signal handler
+        agent_process.send_signal(signal.SIGINT)
+        await safe_shutdown(agent_process)
+        agent_process.out.seek(0)
+        out = agent_process.out.read().decode()
+
+        assert (
+            # either the main PID is still waiting for shutdown, so forwards the SIGKILL
+            "Sending SIGKILL" in out
+            # or SIGKILL came too late, and the main PID is already closing
+            or "KeyboardInterrupt" in out
+            or "Agent stopped!" in out
+            or "Aborted." in out
+        ), (
+            "When sending two SIGINT shortly after each other, the main process should"
+            f" first receive a SIGINT and then a SIGKILL. Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform == "win32",
+        reason="SIGTERM is only used in non-Windows environments",
+    )
+    async def test_sigterm_sends_sigterm_then_sigkill(self, agent_process):
+        agent_process.send_signal(signal.SIGTERM)
+        await anyio.sleep(0.01)  # some time needed for the recursive signal handler
+        agent_process.send_signal(signal.SIGTERM)
+        await safe_shutdown(agent_process)
+        agent_process.out.seek(0)
+        out = agent_process.out.read().decode()
+
+        assert (
+            # either the main PID is still waiting for shutdown, so forwards the SIGKILL
+            "Sending SIGKILL" in out
+            # or SIGKILL came too late, and the main PID is already closing
+            or "KeyboardInterrupt" in out
+            or "Agent stopped!" in out
+            or "Aborted." in out
+        ), (
+            "When sending two SIGTERM shortly after each other, the main process should"
+            f" first receive a SIGINT and then a SIGKILL. Output:\n{out}"
+        )
+
+    @pytest.mark.skipif(
+        sys.platform != "win32",
+        reason="CTRL_BREAK_EVENT is only defined in Windows",
+    )
+    async def test_sends_ctrl_break_win32(self, agent_process):
+        agent_process.send_signal(signal.SIGINT)
+        await safe_shutdown(agent_process)
+        agent_process.out.seek(0)
+        out = agent_process.out.read().decode()
+
+        assert "Sending CTRL_BREAK_EVENT" in out, (
+            "When sending a SIGINT, the main process should send a CTRL_BREAK_EVENT to"
+            f" the uvicorn subprocess. Output:\n{out}"
+        )
diff --git a/tests/cli/test_start_server.py b/tests/cli/test_start_server.py
--- a/tests/cli/test_start_server.py
+++ b/tests/cli/test_start_server.py
@@ -11,6 +11,7 @@
 from prefect.testing.fixtures import is_port_in_use
 from prefect.utilities.processutils import open_process
 
+POLL_INTERVAL = 0.5
 STARTUP_TIMEOUT = 20
 SHUTDOWN_TIMEOUT = 20
 
@@ -70,7 +71,7 @@ async def server_process():
                         if response.status_code == 200:
                             await anyio.sleep(0.5)  # extra sleep for less flakiness
                             break
-                    await anyio.sleep(0.1)
+                    await anyio.sleep(POLL_INTERVAL)
             if response:
                 response.raise_for_status()
             if not response: