This is an automated email from the ASF dual-hosted git repository.
dheerajturaga pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9bf9408d07e Edge worker: bound drain with SIGTERM/SIGKILL escalation
(#65838)
9bf9408d07e is described below
commit 9bf9408d07e5c4eadf6bb96d242bb0f852e11606
Author: Dheeraj Turaga <[email protected]>
AuthorDate: Mon Apr 27 07:36:14 2026 -0500
Edge worker: bound drain with SIGTERM/SIGKILL escalation (#65838)
---
providers/edge3/provider.yaml | 19 ++++
.../src/airflow/providers/edge3/cli/worker.py | 123 +++++++++++++++++++--
.../airflow/providers/edge3/get_provider_info.py | 14 +++
.../edge3/tests/unit/edge3/cli/test_worker.py | 63 +++++++++++
4 files changed, 209 insertions(+), 10 deletions(-)
diff --git a/providers/edge3/provider.yaml b/providers/edge3/provider.yaml
index f296400fdad..c160b23b278 100644
--- a/providers/edge3/provider.yaml
+++ b/providers/edge3/provider.yaml
@@ -116,6 +116,25 @@ config:
type: integer
example: "10"
default: "30"
+ drain_timeout_sec:
+ description: |
+ Maximum seconds the worker waits for running jobs to finish after
entering drain
+ (shutdown requested via SIGINT / SIGTERM / SHUTDOWN_REQUEST /
version mismatch).
+ Once this timeout elapses the worker sends SIGTERM to each running
job, then SIGKILL
+ after ``drain_kill_grace_sec``, and exits regardless of remaining
jobs.
+ Set to 0 (the default) to wait indefinitely and preserve legacy
behavior.
+ version_added: 3.6.0
+ type: integer
+ example: "3600"
+ default: "0"
+ drain_kill_grace_sec:
+ description: |
+ Grace period in seconds after SIGTERM before the worker escalates to
SIGKILL and exits.
+ Only used when ``drain_timeout_sec`` is greater than 0.
+ version_added: 3.6.0
+ type: integer
+ example: "30"
+ default: "30"
worker_concurrency:
description: |
The concurrency defines the default max parallel running task
instances and can also be set during
diff --git a/providers/edge3/src/airflow/providers/edge3/cli/worker.py
b/providers/edge3/src/airflow/providers/edge3/cli/worker.py
index f5396127219..799b34aefd4 100644
--- a/providers/edge3/src/airflow/providers/edge3/cli/worker.py
+++ b/providers/edge3/src/airflow/providers/edge3/cli/worker.py
@@ -20,8 +20,9 @@ import logging
import os
import signal
import sys
+import time
import traceback
-from asyncio import Task, create_task, get_running_loop, sleep
+from asyncio import Task, create_task, gather, get_running_loop, sleep
from collections.abc import Awaitable, Callable
from contextlib import suppress
from datetime import datetime
@@ -105,6 +106,10 @@ class EdgeWorker:
"""List of jobs that the worker is running currently."""
drain: bool = False
"""Flag if job processing should be completed and no new jobs fetched for
a graceful stop/shutdown."""
+ drain_started_at: float | None = None
+ """``time.monotonic()`` timestamp of when drain was first requested."""
+ drain_timed_out: bool = False
+ drain_kill_deadline: float | None = None
maintenance_mode: bool = False
"""Flag if job processing should be completed and no new jobs fetched for
maintenance mode. """
maintenance_comments: str | None = None
@@ -142,6 +147,8 @@ class EdgeWorker:
self.job_poll_interval = self.conf.getint("edge", "job_poll_interval")
self.hb_interval = self.conf.getint("edge", "heartbeat_interval")
+ self.drain_timeout_sec = self.conf.getint("edge", "drain_timeout_sec")
+ self.drain_kill_grace_sec = self.conf.getint("edge",
"drain_kill_grace_sec")
self.base_log_folder: str = (
self.conf.get("logging", "base_log_folder", fallback="NOT
AVAILABLE") or ""
)
@@ -210,18 +217,112 @@ class EdgeWorker:
)
def signal_drain(self):
- self.drain = True
- logger.info("Request to shut down Edge Worker received, waiting for
jobs to complete.")
+ if self._start_draining():
+ logger.info(
+ "Request to shut down Edge Worker received, waiting for jobs
to complete. %s",
+ self._drain_policy_description(),
+ )
def shutdown_handler(self):
- if self.drain:
+ if not self._start_draining():
return
- self.drain = True
logger.info("SIGTERM received. Sending SIGTERM to all jobs and quit")
+ try:
+ loop = get_running_loop()
+ except RuntimeError:
+ pass
+ else:
+ task = loop.create_task(
+ self._push_drain_notice_to_all_jobs(
+ "Edge worker received external SIGTERM; terminating task
supervisor."
+ )
+ )
+ self.background_tasks.add(task)
+ task.add_done_callback(self.background_tasks.discard)
+ self._terminate_jobs(signal.SIGTERM)
+
+ def _start_draining(self) -> bool:
+ """Mark drain start. Returns ``True`` on the first call, ``False`` on
subsequent calls."""
+ if self.drain:
+ return False
+ self.drain = True
+ self.drain_started_at = time.monotonic()
+ return True
+
+ def _drain_policy_description(self) -> str:
+ """One-line description of the configured drain-timeout policy for
logging."""
+ if self.drain_timeout_sec <= 0:
+ return "Drain timeout disabled; will wait indefinitely."
+ return (
+ f"Drain timeout: {self.drain_timeout_sec}s (then SIGTERM), "
+ f"kill grace: {self.drain_kill_grace_sec}s (then SIGKILL)."
+ )
+
+ def _terminate_jobs(self, sig: int) -> None:
+ """Send ``sig`` to every running job process. Safe to call
repeatedly."""
for job in self.jobs:
if job.process.pid:
with suppress(ProcessLookupError, PermissionError):
- os.kill(job.process.pid, signal.SIGTERM)
+ os.kill(job.process.pid, sig)
+
+ async def _push_drain_notice_to_all_jobs(self, message: str) -> None:
+ """Best-effort push of ``message`` into each running job's task log
stream."""
+
+ async def push_one(job: Job) -> None:
+ try:
+ await logs_push(
+ task=job.edge_job.key,
+ log_chunk_time=timezone.utcnow(),
+ log_chunk_data=f"{message}\n",
+ )
+ except Exception:
+ logger.exception("Failed to push drain notice to task log for
%s", job.edge_job.identifier)
+
+ await gather(*(push_one(job) for job in list(self.jobs)))
+
+ async def _enforce_drain_timeout(self) -> bool:
+ """
+ Apply drain-timeout policy when configured.
+
+ Two-phase escalation: once ``drain_timeout_sec`` elapses, SIGTERM
remaining jobs;
+ after ``drain_kill_grace_sec`` more, SIGKILL and return ``True`` so
the loop exits.
+ Returns ``False`` otherwise (not configured, not draining, deadline
not hit, or waiting out grace).
+ """
+ if self.drain_timeout_sec <= 0:
+ return False
+ if not self.drain or not self.jobs or self.drain_started_at is None:
+ return False
+ now = time.monotonic()
+ if now - self.drain_started_at < self.drain_timeout_sec:
+ return False
+ if not self.drain_timed_out:
+ self.drain_timed_out = True
+ self.drain_kill_deadline = now + self.drain_kill_grace_sec
+ logger.warning(
+ "Drain timeout of %ds exceeded with %d job(s) still running.
Sending SIGTERM.",
+ self.drain_timeout_sec,
+ len(self.jobs),
+ )
+ await self._push_drain_notice_to_all_jobs(
+ f"Edge worker drain timeout of {self.drain_timeout_sec}s
expired; "
+ f"sending SIGTERM to task supervisor. "
+ f"Will escalate to SIGKILL after {self.drain_kill_grace_sec}s
grace."
+ )
+ self._terminate_jobs(signal.SIGTERM)
+ return False
+ if self.drain_kill_deadline is not None and now >=
self.drain_kill_deadline:
+ logger.warning(
+ "Drain kill grace of %ds exceeded with %d job(s) still
running. Sending SIGKILL and exiting.",
+ self.drain_kill_grace_sec,
+ len(self.jobs),
+ )
+ await self._push_drain_notice_to_all_jobs(
+ f"Edge worker drain kill-grace of {self.drain_kill_grace_sec}s
expired; "
+ f"sending SIGKILL and exiting worker."
+ )
+ self._terminate_jobs(signal.SIGKILL)
+ return True
+ return False
async def _get_sysinfo(self) -> dict[str, str | int | float | datetime]:
"""Produce the sysinfo from worker to post to central site."""
@@ -406,6 +507,9 @@ class EdgeWorker:
else:
logger.info("%i %s running", len(self.jobs), "job is" if
len(self.jobs) == 1 else "jobs are")
+ if await self._enforce_drain_timeout():
+ break
+
await self.interruptible_sleep()
async def fetch_and_run_job(self) -> None:
@@ -459,7 +563,7 @@ class EdgeWorker:
await logs_push(
task=job.edge_job.key,
log_chunk_time=timezone.utcnow(),
- log_chunk_data=f"Error starting job:\n{ex_txt}",
+ log_chunk_data=f"Error executing job:\n{ex_txt}",
)
await jobs_set_state(job.edge_job.key, TaskInstanceState.FAILED)
@@ -498,14 +602,13 @@ class EdgeWorker:
self.maintenance_comments = worker_info.maintenance_comments
else:
self.maintenance_comments = None
- if worker_info.state == EdgeWorkerState.SHUTDOWN_REQUEST:
+ if worker_info.state == EdgeWorkerState.SHUTDOWN_REQUEST and
self._start_draining():
logger.info("Shutdown requested!")
- self.drain = True
worker_state_changed = worker_info.state != state
except EdgeWorkerVersionException:
logger.info("Version mismatch of Edge worker and Core. Shutting
down worker.")
- self.drain = True
+ self._start_draining()
return worker_state_changed
async def interruptible_sleep(self):
diff --git a/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
b/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
index 441cee6b65a..7e0085982b2 100644
--- a/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
+++ b/providers/edge3/src/airflow/providers/edge3/get_provider_info.py
@@ -67,6 +67,20 @@ def get_provider_info():
"example": "10",
"default": "30",
},
+ "drain_timeout_sec": {
+ "description": "Maximum seconds the worker waits for
running jobs to finish after entering drain\n(shutdown requested via SIGINT /
SIGTERM / SHUTDOWN_REQUEST / version mismatch).\nOnce this timeout elapses the
worker sends SIGTERM to each running job, then SIGKILL\nafter
``drain_kill_grace_sec``, and exits regardless of remaining jobs.\nSet to 0
(the default) to wait indefinitely and preserve legacy behavior.\n",
+ "version_added": "3.6.0",
+ "type": "integer",
+ "example": "3600",
+ "default": "0",
+ },
+ "drain_kill_grace_sec": {
+ "description": "Grace period in seconds after SIGTERM
before the worker escalates to SIGKILL and exits.\nOnly used when
``drain_timeout_sec`` is greater than 0.\n",
+ "version_added": "3.6.0",
+ "type": "integer",
+ "example": "30",
+ "default": "30",
+ },
"worker_concurrency": {
"description": "The concurrency defines the default
max parallel running task instances and can also be set during\nstart of worker
with the ``airflow edge worker`` command parameter. The size of the
workers\nand the resources must support the nature of your tasks. The
parameter\nworks together with the concurrency_slots parameter of a task.\n",
"version_added": None,
diff --git a/providers/edge3/tests/unit/edge3/cli/test_worker.py
b/providers/edge3/tests/unit/edge3/cli/test_worker.py
index 62565cdf98c..be33ea68d36 100644
--- a/providers/edge3/tests/unit/edge3/cli/test_worker.py
+++ b/providers/edge3/tests/unit/edge3/cli/test_worker.py
@@ -464,6 +464,69 @@ class TestEdgeWorker:
await worker_with_job.heartbeat()
assert worker_with_job.drain
+ @patch("airflow.providers.edge3.cli.worker.os.kill")
+ @patch("airflow.providers.edge3.cli.worker.time.monotonic")
+ async def test_enforce_drain_timeout_disabled(
+ self, mock_monotonic, mock_kill, worker_with_job: EdgeWorker
+ ):
+ worker_with_job.drain_timeout_sec = 0
+ worker_with_job.drain = True
+ worker_with_job.drain_started_at = 0.0
+ mock_monotonic.return_value = 9999.0
+ assert await worker_with_job._enforce_drain_timeout() is False
+ mock_kill.assert_not_called()
+
+ @patch("airflow.providers.edge3.cli.worker.os.kill")
+ @patch("airflow.providers.edge3.cli.worker.time.monotonic")
+ async def test_enforce_drain_timeout_within_deadline(
+ self, mock_monotonic, mock_kill, worker_with_job: EdgeWorker
+ ):
+ worker_with_job.drain_timeout_sec = 60
+ worker_with_job.drain = True
+ worker_with_job.drain_started_at = 100.0
+ worker_with_job.jobs[0].process = mock.MagicMock(pid=4242)
+ mock_monotonic.return_value = 120.0 # 20s < 60s
+ assert await worker_with_job._enforce_drain_timeout() is False
+ mock_kill.assert_not_called()
+
+ @patch("airflow.providers.edge3.cli.worker.logs_push")
+ @patch("airflow.providers.edge3.cli.worker.os.kill")
+ @patch("airflow.providers.edge3.cli.worker.time.monotonic")
+ async def test_enforce_drain_timeout_sigterm_on_deadline(
+ self, mock_monotonic, mock_kill, mock_logs_push, worker_with_job:
EdgeWorker
+ ):
+ worker_with_job.drain_timeout_sec = 60
+ worker_with_job.drain_kill_grace_sec = 30
+ worker_with_job.drain = True
+ worker_with_job.drain_started_at = 100.0
+ worker_with_job.jobs[0].process = mock.MagicMock(pid=4242)
+ mock_monotonic.return_value = 170.0 # 70s >= 60s, first trigger
+ assert await worker_with_job._enforce_drain_timeout() is False
+ mock_kill.assert_called_once_with(4242, signal.SIGTERM)
+ assert worker_with_job.drain_timed_out is True
+ assert worker_with_job.drain_kill_deadline == 200.0
+ mock_logs_push.assert_called_once()
+ assert "SIGTERM" in mock_logs_push.call_args.kwargs["log_chunk_data"]
+
+ @patch("airflow.providers.edge3.cli.worker.logs_push")
+ @patch("airflow.providers.edge3.cli.worker.os.kill")
+ @patch("airflow.providers.edge3.cli.worker.time.monotonic")
+ async def test_enforce_drain_timeout_sigkill_after_grace(
+ self, mock_monotonic, mock_kill, mock_logs_push, worker_with_job:
EdgeWorker
+ ):
+ worker_with_job.drain_timeout_sec = 60
+ worker_with_job.drain_kill_grace_sec = 30
+ worker_with_job.drain = True
+ worker_with_job.drain_started_at = 100.0
+ worker_with_job.drain_timed_out = True
+ worker_with_job.drain_kill_deadline = 200.0
+ worker_with_job.jobs[0].process = mock.MagicMock(pid=4242)
+ mock_monotonic.return_value = 210.0 # past kill deadline
+ assert await worker_with_job._enforce_drain_timeout() is True
+ mock_kill.assert_called_once_with(4242, signal.SIGKILL)
+ mock_logs_push.assert_called_once()
+ assert "SIGKILL" in mock_logs_push.call_args.kwargs["log_chunk_data"]
+
@pytest.mark.parametrize(
"http_error",
[
