diff --git a/pyworkflow/celery/app.py b/pyworkflow/celery/app.py index d35ef50..bb49b19 100644 --- a/pyworkflow/celery/app.py +++ b/pyworkflow/celery/app.py @@ -12,12 +12,19 @@ garbage collector and Celery's saferepr module. It does not affect functionality. """ +import contextlib import os from datetime import timedelta from typing import Any from celery import Celery -from celery.signals import worker_init, worker_process_init, worker_shutdown +from celery.signals import ( + task_postrun, + task_prerun, + worker_init, + worker_process_init, + worker_shutdown, +) from kombu import Exchange, Queue from pyworkflow.observability.logging import configure_logging @@ -281,6 +288,9 @@ def create_celery_app( PYWORKFLOW_CELERY_SENTINEL_MASTER: Sentinel master name (used if sentinel_master_name param not provided) PYWORKFLOW_WORKER_MAX_MEMORY: Max memory per worker child (KB) (used if worker_max_memory_per_child param not provided) PYWORKFLOW_WORKER_MAX_TASKS: Max tasks per worker child (used if worker_max_tasks_per_child param not provided) + PYWORKFLOW_RESCHEDULE_ON_SIGTERM: Re-enqueue in-flight tasks immediately when a worker is shut down + (SIGTERM / spot reclaim), instead of waiting out the broker visibility_timeout. Default on; set to + "0"/"false"/"no" to disable. Implemented as a consumer bootstep; see pyworkflow/celery/reschedule.py. Examples: # Default configuration (uses env vars if set, otherwise localhost Redis) @@ -471,6 +481,14 @@ def create_celery_app( # to ensure proper initialization AFTER process forking. # See on_worker_init() and on_worker_process_init() below. + # Re-enqueue in-flight tasks on worker shutdown (SIGTERM / spot reclaim) so a + # live worker resumes them within seconds instead of waiting out the broker + # visibility_timeout. Hooked as a consumer bootstep so it fires exactly as the + # consumer is torn down (event-based). See pyworkflow/celery/reschedule.py. + from pyworkflow.celery.reschedule import RescheduleConsumerStep + + app.steps["consumer"].add(RescheduleConsumerStep) + # Auto-discover workflows from environment variable or configured modules discover_workflows() @@ -555,6 +573,29 @@ def on_worker_shutdown(**kwargs): close_worker_loop() +@task_prerun.connect +def on_task_prerun(task_id=None, task=None, args=None, kwargs=None, **_): + """ + Track the task now executing in this child process so it can be re-enqueued + if the worker is shut down (SIGTERM) before the task finishes. + + Wrapped defensively: tracking must never interfere with task execution. + """ + from pyworkflow.celery.reschedule import track_task_start + + with contextlib.suppress(Exception): + track_task_start(task, task_id, args, kwargs) + + +@task_postrun.connect +def on_task_postrun(task_id=None, **_): + """Clear the in-flight task once it has finished (clean completion).""" + from pyworkflow.celery.reschedule import track_task_end + + with contextlib.suppress(Exception): + track_task_end(task_id) + + def get_celery_app() -> Celery: """ Get the global Celery application instance. diff --git a/pyworkflow/celery/reschedule.py b/pyworkflow/celery/reschedule.py new file mode 100644 index 0000000..15c8edc --- /dev/null +++ b/pyworkflow/celery/reschedule.py @@ -0,0 +1,316 @@ +""" +Reschedule in-flight tasks when a worker is shut down (SIGTERM). + +When a Celery worker pod is reclaimed (e.g. an AWS spot node drain), it receives +SIGTERM and warm-shuts-down. Any step task that was executing is *not* promptly +redelivered: with ``task_acks_late=True`` the broker only re-queues it after the +Redis ``visibility_timeout`` (an hour by default), and the parent workflow stays +SUSPENDED that whole time -- so the user-facing flow silently hangs with no output. + +This module makes the worker re-enqueue its in-flight task(s) *immediately* on +shutdown so a live worker resumes them within seconds instead of ~1h. + +Why a shared Redis registry (not process-local state) +----------------------------------------------------- +Tasks execute in forked **child** processes, but the shutdown hook runs in the +**main** process. The main process can't see a child's in-memory task args, so +children publish what they are running to a per-pod Redis hash +(``pyworkflow:inflight:``) on ``task_prerun`` and remove it on +``task_postrun``. On shutdown the main process reads that hash and re-enqueues +every entry. + +Why a consumer bootstep (event-based, no sleep) +----------------------------------------------- +The re-enqueue must happen *after* this worker's broker consumer is cancelled -- +otherwise the worker's own still-active consumer grabs the message straight back +into ``unacked`` (stranded until the visibility_timeout). There is no Celery +*signal* for "consumer cancelled" (``worker_shutting_down`` fires before it; +``worker_shutdown`` fires after the pool has drained the long task -- far too +late). So we hook the consumer blueprint directly with a bootstep that +``requires`` the ``Tasks`` step: its ``stop()`` runs as the consumer is torn +down (early in shutdown, before the pool drain). There we cancel the task +consumer ourselves -- guaranteeing this worker stops fetching -- and only then +re-enqueue, so the message lands in the *ready* queue and another worker's +blocking BRPOP picks it up immediately. Event-based, deterministic, no sleep. + +Safety against double execution +------------------------------- +``execute_step_task`` short-circuits on an existing ``STEP_COMPLETED`` event (and +on a terminal run status), and ``_record_step_completion_and_resume`` re-checks +before recording, so at most one completion is ever recorded no matter how many +copies run. Only ``SingletonWorkflowTask`` tasks that declare ``unique_on`` are +rescheduled -- those are idempotent by run_id/step_id. Tasks without a uniqueness +key are left alone. + +Gate: ``PYWORKFLOW_RESCHEDULE_ON_SIGTERM`` (default on; set to ``0``/``false``/``no`` +to disable and fall back to the old visibility_timeout redelivery behaviour). +""" + +import json +import os +import socket +import threading +from typing import Any + +from celery import bootsteps +from loguru import logger + +# Registry hash self-expires so a hard-killed pod (no postrun) cannot leak entries. +_REGISTRY_TTL_SECONDS = 6 * 3600 +_REGISTRY_KEY_PREFIX = "pyworkflow:inflight:" + +_backend_lock = threading.Lock() +_backend: Any = None +_backend_resolved = False + + +def _enabled() -> bool: + return os.getenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "1").strip().lower() not in ( + "0", + "false", + "no", + ) + + +def _pod_id() -> str: + """Stable per-pod id shared by the main process and its forked children.""" + return os.environ.get("HOSTNAME") or socket.gethostname() + + +def _registry_key() -> str: + return f"{_REGISTRY_KEY_PREFIX}{_pod_id()}" + + +def _get_backend() -> Any: + """ + Cached sentinel-aware Redis backend, reusing the singleton lock configuration + (same broker/sentinel the task locks already use). Returns None when no Redis + backend is configured, in which case rescheduling is a no-op. + """ + global _backend, _backend_resolved + if _backend_resolved: + return _backend + with _backend_lock: + if _backend_resolved: + return _backend + try: + from pyworkflow.celery.app import celery_app + from pyworkflow.celery.singleton import RedisLockBackend, SingletonConfig + + cfg = SingletonConfig(celery_app) + url = cfg.backend_url + if url: + _backend = RedisLockBackend( + url, + is_sentinel=cfg.is_sentinel, + sentinel_master=cfg.sentinel_master, + ) + except Exception as exc: + logger.warning(f"reschedule: could not initialise redis backend: {exc}") + _backend = None + _backend_resolved = True + return _backend + + +def _resolve_task(name: str) -> Any: + """Look up a registered task instance by name (used to re-enqueue).""" + from pyworkflow.celery.app import celery_app + + return celery_app.tasks.get(name) + + +def _resolve_queue(task: Any) -> str | None: + """Best-effort queue the task was delivered on, so the re-enqueue lands there.""" + try: + delivery_info = getattr(task.request, "delivery_info", None) or {} + routing_key = delivery_info.get("routing_key") + if routing_key: + return routing_key + except Exception: + pass + return getattr(task, "queue", None) + + +def track_task_start( + task: Any, + task_id: str | None, + args: Any, + kwargs: Any, +) -> None: + """Publish a task as in-flight (per pod) so it can be re-enqueued on SIGTERM.""" + if not _enabled() or task is None or task_id is None: + return + + # Import here to avoid import cycles at module load. + from pyworkflow.celery.singleton import SingletonWorkflowTask + + # Only re-enqueue tasks keyed by a uniqueness field (run_id/step_id). These are + # idempotent; blindly re-enqueuing a non-unique task could start duplicate runs. + if not isinstance(task, SingletonWorkflowTask) or not task.unique_on: + return + + backend = _get_backend() + if backend is None: + return + + descriptor = json.dumps( + { + "name": task.name, + "args": list(args or ()), + "kwargs": dict(kwargs or {}), + "queue": _resolve_queue(task), + }, + default=str, + ) + key = _registry_key() + try: + backend._execute_with_refresh(lambda: backend.redis.hset(key, task_id, descriptor)) + backend._execute_with_refresh(lambda: backend.redis.expire(key, _REGISTRY_TTL_SECONDS)) + except Exception as exc: + logger.warning(f"reschedule: failed to record in-flight task {task_id}: {exc}") + + +def track_task_end(task_id: str | None) -> None: + """Remove a task from the in-flight registry once it has finished.""" + if not _enabled() or task_id is None: + return + backend = _get_backend() + if backend is None: + return + try: + backend._execute_with_refresh(lambda: backend.redis.hdel(_registry_key(), task_id)) + except Exception as exc: + logger.warning(f"reschedule: failed to clear in-flight task {task_id}: {exc}") + + +def reschedule_inflight_on_shutdown() -> None: + """ + Re-enqueue every task still recorded as in-flight for this pod. Assumes the + caller has already cancelled this worker's task consumer (see + :class:`RescheduleConsumerStep`) so the fresh copy is not re-grabbed. Never + raises -- shutdown must not be blocked. + """ + if not _enabled(): + return + backend = _get_backend() + if backend is None: + return + + key = _registry_key() + try: + entries = backend._execute_with_refresh(lambda: backend.redis.hgetall(key)) or {} + except Exception as exc: + logger.warning(f"reschedule: failed to read in-flight registry: {exc}") + return + + if not entries: + return + + for task_id, raw in entries.items(): + try: + descriptor = json.loads(raw) + task = _resolve_task(descriptor["name"]) + if task is None: + logger.warning( + f"reschedule: unknown task {descriptor.get('name')!r}, " + f"cannot re-enqueue {task_id}" + ) + continue + + args = descriptor.get("args") or [] + kwargs = descriptor.get("kwargs") or {} + queue = descriptor.get("queue") + + # Release the singleton lock held by the dying task so the re-enqueued + # copy is not rejected as a duplicate. release_lock() regenerates the + # exact lock key from the same args/kwargs the running task holds. + try: + task.release_lock(task_args=list(args), task_kwargs=kwargs) + except Exception as exc: + logger.warning( + f"reschedule: failed to release lock for {task.name} " + f"(old_task_id={task_id}): {exc}" + ) + + options = {"queue": queue} if queue else {} + task.apply_async(args=list(args), kwargs=kwargs, **options) + logger.warning( + f"reschedule: re-enqueued in-flight task on shutdown " + f"task={task.name} old_task_id={task_id} queue={queue}" + ) + except Exception as exc: + logger.opt(exception=True).error( + f"reschedule: failed to re-enqueue task {task_id}: {exc}" + ) + + # Drop the registry for this pod; a fresh worker starts with a clean slate. + try: + backend._execute_with_refresh(lambda: backend.redis.delete(key)) + except Exception as exc: + logger.warning(f"reschedule: failed to clear in-flight registry: {exc}") + + +def _is_worker_stopping() -> bool: + """ + True only when the worker is actually shutting down (SIGTERM warm or cold), + not on a transient consumer restart (e.g. broker reconnect), which also tears + the consumer down. Celery's SIGTERM handler sets these flags before the + consumer is stopped. + """ + try: + from celery.worker import state + + # Identity checks (matching celery.worker.state): the flags are set to the + # exit code, which can be 0 (EX_OK) -- `0 == False`, so membership/equality + # against False would wrongly treat a clean exitcode-0 shutdown as "not + # stopping" and skip the reschedule. + def _set(v: Any) -> bool: + return v is not None and v is not False + + return _set(state.should_stop) or _set(state.should_terminate) + except Exception: + return False + + +def on_consumer_stop(consumer: Any) -> None: + """ + Invoked from the consumer bootstep as the consumer is torn down. Cancels this + worker's task consumer (so it stops fetching) and re-enqueues in-flight work. + """ + if not _enabled() or not _is_worker_stopping(): + return + + # Cancel our own task consumer first: once we stop BRPOPing, the re-enqueued + # message stays in the ready list for another worker instead of being grabbed + # back here into `unacked`. Tasks.stop() will cancel again later (idempotent). + task_consumer = getattr(consumer, "task_consumer", None) + if task_consumer is not None: + try: + task_consumer.cancel() + except Exception as exc: + logger.warning(f"reschedule: failed to cancel task consumer: {exc}") + + reschedule_inflight_on_shutdown() + + +class RescheduleConsumerStep(bootsteps.StartStopStep): + """ + Consumer bootstep that re-enqueues this worker's in-flight tasks the moment + the consumer is torn down on shutdown -- event-based, no polling/sleep. + + ``requires`` the task consumer (``Tasks``) so that (a) ``c.task_consumer`` + exists when our ``stop()`` runs and (b) we run *before* ``Tasks.stop()``, + letting us cancel the consumer ourselves and re-enqueue while it's guaranteed + not fetching. + """ + + requires = ("celery.worker.consumer:Tasks",) + + def start(self, c: Any) -> None: # noqa: D401 - nothing to start + pass + + def stop(self, c: Any) -> None: + try: + on_consumer_stop(c) + except Exception as exc: # never block consumer teardown + logger.warning(f"reschedule: consumer-stop hook failed: {exc}") diff --git a/tests/integration/_reschedule_e2e_app.py b/tests/integration/_reschedule_e2e_app.py new file mode 100644 index 0000000..4158152 --- /dev/null +++ b/tests/integration/_reschedule_e2e_app.py @@ -0,0 +1,41 @@ +""" +Standalone workflow module for the SIGTERM-reschedule end-to-end test. + +Imported by the worker subprocess (via PYWORKFLOW_MODULE) and by the CLI that +starts the workflow. The single step blocks until a `release` marker file +appears, so the test can deterministically catch it mid-execution, SIGTERM the +worker, and prove a fresh worker re-runs it and the workflow completes. + +Marker files live under $E2E_MARKER_DIR: +- writes `started` as soon as the step runs (so the test knows it's in-flight) +- waits for `release` before returning +""" + +import asyncio +import os +from pathlib import Path + +from pyworkflow import step, workflow + + +def _marker_dir() -> Path: + return Path(os.environ["E2E_MARKER_DIR"]) + + +@step(max_retries=1) +async def blocking_step() -> str: + d = _marker_dir() + d.mkdir(parents=True, exist_ok=True) + # Signal that the step is executing (the test waits for this before SIGTERM). + (d / "started").write_text("1") + # Block until released, so the step is reliably in-flight when the worker dies. + for _ in range(240): # up to ~120s + if (d / "release").exists(): + return "released" + await asyncio.sleep(0.5) + raise TimeoutError("release marker never appeared") + + +@workflow(durable=True) +async def e2e_blocking_workflow() -> str: + return await blocking_step() diff --git a/tests/integration/test_reschedule.py b/tests/integration/test_reschedule.py new file mode 100644 index 0000000..14ad68b --- /dev/null +++ b/tests/integration/test_reschedule.py @@ -0,0 +1,133 @@ +""" +Integration tests for SIGTERM task rescheduling against a REAL Redis. + +These exercise the actual pieces the unit tests fake out: +- the per-pod in-flight registry hash in real Redis, +- real singleton-lock release + re-acquire, +- a real broker publish (apply_async lands a message on the queue's Redis list). + +Requires a Redis at localhost:6379. Skipped otherwise. Uses a dedicated db (14) +which is flushed around each test, so it never touches app data on db 0. +""" + +import pytest + +from pyworkflow.celery import reschedule + +try: + import redis as _redis + + _client = _redis.from_url("redis://localhost:6379/0") + _client.ping() + REDIS_AVAILABLE = True +except Exception: + REDIS_AVAILABLE = False + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif(not REDIS_AVAILABLE, reason="Redis not available"), +] + +TEST_DB_URL = "redis://localhost:6379/14" +STEPS_QUEUE = "pyworkflow.steps" + + +@pytest.fixture +def env(monkeypatch): + """Real Redis-backed reschedule wiring on an isolated db, fixed pod id.""" + from pyworkflow.celery.app import create_celery_app + from pyworkflow.celery.singleton import RedisLockBackend, SingletonWorkflowTask + + raw = _redis.from_url(TEST_DB_URL, decode_responses=True) + raw.flushdb() + + monkeypatch.delenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", raising=False) + monkeypatch.setenv("HOSTNAME", "itest-pod") + + app = create_celery_app(broker_url=TEST_DB_URL, result_backend=TEST_DB_URL) + + @app.task( + base=SingletonWorkflowTask, + name="test.reschedule_probe", + queue=STEPS_QUEUE, + unique_on=["run_id"], + ) + def probe(run_id, payload=None): # pragma: no cover - never executed (no worker) + return None + + backend = RedisLockBackend(TEST_DB_URL) + monkeypatch.setattr(reschedule, "_get_backend", lambda: backend) + monkeypatch.setattr(reschedule, "_resolve_task", lambda name: app.tasks.get(name)) + + registry_key = "pyworkflow:inflight:itest-pod" + + yield { + "raw": raw, + "app": app, + "probe": probe, + "backend": backend, + "registry_key": registry_key, + } + + raw.flushdb() + + +def test_track_writes_and_clears_real_hash(env): + probe, raw, key = env["probe"], env["raw"], env["registry_key"] + + reschedule.track_task_start(probe, "tid-1", (), {"run_id": "rX"}) + entries = raw.hgetall(key) + assert "tid-1" in entries + assert "test.reschedule_probe" in entries["tid-1"] # descriptor carries task name + + reschedule.track_task_end("tid-1") + assert raw.hgetall(key) == {} + + +def test_reschedule_releases_lock_and_reenqueues(env): + probe, raw, key = env["probe"], env["raw"], env["registry_key"] + + # Simulate the running task holding its singleton lock. + lock_key = probe.generate_lock("test.reschedule_probe", [], {"run_id": "rX"}) + assert probe.acquire_lock(lock_key, "tid-1") is True + assert raw.get(lock_key) == "tid-1" + + reschedule.track_task_start(probe, "tid-1", (), {"run_id": "rX"}) + llen_before = raw.llen(STEPS_QUEUE) + + reschedule.reschedule_inflight_on_shutdown() + + # A fresh copy was published to the real broker queue... + assert raw.llen(STEPS_QUEUE) == llen_before + 1 + # ...and it re-acquired the singleton lock under a NEW task id (so the dying + # task's lock no longer blocks it). + holder = raw.get(lock_key) + assert holder is not None and holder != "tid-1" + # Registry drained. + assert raw.hgetall(key) == {} + + +def test_direct_apply_async_is_deduped_without_release(env): + """Documents WHY reschedule must release the lock first: a plain re-enqueue + while the lock is held is silently deduped and never reaches the broker.""" + probe, raw = env["probe"], env["raw"] + + lock_key = probe.generate_lock("test.reschedule_probe", [], {"run_id": "rX"}) + assert probe.acquire_lock(lock_key, "tid-1") is True + llen_before = raw.llen(STEPS_QUEUE) + + probe.apply_async(kwargs={"run_id": "rX"}) # no release -> singleton dedup + + assert raw.llen(STEPS_QUEUE) == llen_before # nothing enqueued + assert raw.get(lock_key) == "tid-1" # lock untouched + + +def test_disabled_writes_nothing(env, monkeypatch): + monkeypatch.setenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "0") + probe, raw, key = env["probe"], env["raw"], env["registry_key"] + + reschedule.track_task_start(probe, "tid-1", (), {"run_id": "rX"}) + reschedule.reschedule_inflight_on_shutdown() + + assert raw.hgetall(key) == {} + assert raw.llen(STEPS_QUEUE) == 0 diff --git a/tests/integration/test_reschedule_e2e.py b/tests/integration/test_reschedule_e2e.py new file mode 100644 index 0000000..28f73c7 --- /dev/null +++ b/tests/integration/test_reschedule_e2e.py @@ -0,0 +1,200 @@ +""" +End-to-end test for SIGTERM task rescheduling with REAL worker subprocesses. + +Proves the production scenario: a prefork Celery worker is running a step when it +receives SIGTERM (spot reclaim); its ``worker_shutting_down`` handler re-enqueues +the in-flight step; a fresh worker picks it up and the workflow COMPLETES within +seconds -- instead of hanging until the broker visibility_timeout (~1h). + +Sequence (deterministic): +1. Worker A runs the step, which blocks until a `release` marker appears. +2. SIGTERM A; wait until A *logs* the re-enqueue (its worker_shutting_down handler). +3. SIGKILL A's whole process group, so A cannot finish the step itself -- only a + fresh worker can (mirrors k8s SIGKILL after the grace period). +4. Worker B starts, the step is released, and the workflow completes. + +Requires Redis at localhost:6379. Uses a dedicated broker db (13) and a temp +FileStorageBackend shared across processes. Marked slow (boots real workers). +""" + +import asyncio +import contextlib +import os +import re +import signal +import subprocess +import sys +import time +from pathlib import Path + +import pytest + +try: + import redis as _redis + + _redis.from_url("redis://localhost:6379/0").ping() + REDIS_AVAILABLE = True +except Exception: + REDIS_AVAILABLE = False + +pytestmark = [ + pytest.mark.integration, + pytest.mark.slow, + pytest.mark.skipif(not REDIS_AVAILABLE, reason="Redis not available"), +] + +BROKER = "redis://localhost:6379/13" +APP_MODULE = "_reschedule_e2e_app" +WORKFLOW = "e2e_blocking_workflow" +RESCHEDULE_LOG = "re-enqueued in-flight task on shutdown" +_APP_DIR = os.path.dirname(os.path.abspath(__file__)) +_RUN_ID_RE = re.compile(r"run_[0-9a-f]{16}") + + +def _base_env(storage_path: str, marker_dir: str) -> dict: + env = os.environ.copy() + env.update( + { + "PYWORKFLOW_CELERY_BROKER": BROKER, + "PYWORKFLOW_CELERY_RESULT_BACKEND": BROKER, + "PYWORKFLOW_RUNTIME": "celery", + "PYWORKFLOW_STORAGE_TYPE": "file", + "PYWORKFLOW_STORAGE_PATH": storage_path, + "PYWORKFLOW_MODULE": APP_MODULE, + "E2E_MARKER_DIR": marker_dir, + "PYWORKFLOW_RESCHEDULE_ON_SIGTERM": "1", + "PYTHONPATH": _APP_DIR + os.pathsep + env.get("PYTHONPATH", ""), + } + ) + return env + + +def _spawn_worker(env: dict, hostname: str, logfile) -> subprocess.Popen: + # start_new_session so we can signal the whole group (main + prefork children). + return subprocess.Popen( + [ + sys.executable, + "-m", + "pyworkflow.cli", + "--module", + APP_MODULE, + "worker", + "run", + "--loglevel", + "warning", + ], + env={**env, "HOSTNAME": hostname}, + stdout=logfile, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + +def _killpg(proc: subprocess.Popen | None) -> None: + if proc and proc.poll() is None: + with contextlib.suppress(Exception): + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + with contextlib.suppress(Exception): + proc.wait(timeout=10) + + +def _run_status(storage_path: str, run_id: str) -> str: + from pyworkflow.storage.file import FileStorageBackend + + storage = FileStorageBackend(base_path=storage_path) + + async def _get(): + run = await storage.get_run(run_id) + return run.status.value if run else "NONE" + + return asyncio.run(_get()) + + +def _wait_for(predicate, timeout: float, interval: float = 0.5) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(interval) + return False + + +def test_sigterm_reschedule_resumes_workflow(tmp_path): + storage_path = str(tmp_path / "storage") + marker_dir = tmp_path / "markers" + marker_dir.mkdir() + env = _base_env(storage_path, str(marker_dir)) + log_a = tmp_path / "worker_a.log" + log_b = tmp_path / "worker_b.log" + + raw = _redis.from_url(BROKER, decode_responses=True) + raw.flushdb() + + worker_a = worker_b = None + with open(log_a, "w") as fa, open(log_b, "w") as fb: + try: + # 1. Worker A (prefork) -- will run the step, then we SIGTERM it. + worker_a = _spawn_worker(env, "e2e-pod-a", fa) + + # 2. Start the workflow (separate CLI process, same broker/storage). + started = subprocess.run( + [ + sys.executable, + "-m", + "pyworkflow.cli", + "--module", + APP_MODULE, + "workflows", + "run", + WORKFLOW, + "--no-wait", + ], + env=env, + capture_output=True, + text=True, + timeout=60, + ) + m = _RUN_ID_RE.search(started.stdout + started.stderr) + if m: + run_id = m.group(0) + else: + runs = list((Path(storage_path) / "runs").glob("run_*.json")) + assert runs, f"no run created:\n{started.stdout}\n{started.stderr}" + run_id = runs[0].stem + + # 3. Wait until the step is actually executing on worker A. + assert _wait_for(lambda: (marker_dir / "started").exists(), timeout=60), ( + f"step never started on worker A\n{log_a.read_text()}" + ) + + # 4. SIGTERM A and wait until it LOGS the re-enqueue (handler fired). + worker_a.send_signal(signal.SIGTERM) + assert _wait_for(lambda: RESCHEDULE_LOG in log_a.read_text(), timeout=30), ( + f"worker A never logged a reschedule on SIGTERM\n{log_a.read_text()}" + ) + + # 5. Hard-kill A's whole group so it cannot complete the step itself; + # only the re-enqueued copy on a fresh worker can finish it. + _killpg(worker_a) + + # 6. Fresh worker B picks up the re-enqueued step. + worker_b = _spawn_worker(env, "e2e-pod-b", fb) + + # 7. Release the (re-run) step so the workflow can finish. + (marker_dir / "release").write_text("go") + + # 8. Must COMPLETE within seconds (visibility_timeout is 3600s, so a + # completion now proves the reschedule -- not broker redelivery). + _wait_for( + lambda: _run_status(storage_path, run_id) in ("completed", "failed"), + timeout=90, + ) + final = _run_status(storage_path, run_id) + assert final == "completed", ( + f"run {run_id} ended {final!r}, expected completed\n" + f"--- A ---\n{log_a.read_text()}\n--- B ---\n{log_b.read_text()}" + ) + finally: + _killpg(worker_a) + _killpg(worker_b) + raw.flushdb() diff --git a/tests/unit/test_reschedule.py b/tests/unit/test_reschedule.py new file mode 100644 index 0000000..355d7ce --- /dev/null +++ b/tests/unit/test_reschedule.py @@ -0,0 +1,301 @@ +""" +Unit tests for re-enqueueing in-flight tasks on worker shutdown (SIGTERM). + +See pyworkflow/celery/reschedule.py. The in-flight registry lives in Redis (a +per-pod hash) so the main process can read what the child processes are running; +here we substitute a tiny in-memory fake for that Redis. +""" + +from unittest.mock import MagicMock + +import pytest + +from pyworkflow.celery import reschedule +from pyworkflow.celery.singleton import SingletonWorkflowTask + + +class _FakeRedis: + """Minimal hash-only fake supporting the ops reschedule.py uses.""" + + def __init__(self): + self.hashes: dict[str, dict[str, str]] = {} + + def hset(self, key, field, value): + self.hashes.setdefault(key, {})[field] = value + return 1 + + def hdel(self, key, field): + return self.hashes.get(key, {}).pop(field, None) is not None + + def hgetall(self, key): + return dict(self.hashes.get(key, {})) + + def expire(self, key, ttl): + return True + + def delete(self, key): + self.hashes.pop(key, None) + return 1 + + +class _FakeBackend: + """Stands in for RedisLockBackend (passthrough _execute_with_refresh + .redis).""" + + def __init__(self): + self.redis = _FakeRedis() + + def _execute_with_refresh(self, fn, *args, **kwargs): + return fn(*args, **kwargs) + + +def _make_singleton_task( + name: str = "pyworkflow.execute_step", + unique_on=("run_id", "step_id"), + queue: str | None = "pyworkflow.steps", + routing_key: str | None = "pyworkflow.steps", +): + """A MagicMock that passes ``isinstance(_, SingletonWorkflowTask)``.""" + task = MagicMock(spec=SingletonWorkflowTask) + task.name = name + task.unique_on = list(unique_on) if unique_on else unique_on + task.queue = queue + task.request.delivery_info = {"routing_key": routing_key} if routing_key else {} + return task + + +@pytest.fixture(autouse=True) +def fake_backend(monkeypatch): + """Force the feature on, give it a fake Redis, and a name->task resolver.""" + monkeypatch.delenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", raising=False) + monkeypatch.setenv("HOSTNAME", "test-pod-0") + + backend = _FakeBackend() + monkeypatch.setattr(reschedule, "_get_backend", lambda: backend) + + # Re-enqueue resolves tasks by name; map names to the mocks the test created. + registry: dict[str, MagicMock] = {} + monkeypatch.setattr(reschedule, "_resolve_task", lambda name: registry.get(name)) + + backend.registry = registry # let tests register name -> task + return backend + + +def _track(backend, task, task_id, kwargs): + """Register the task for name-resolution and record it as in-flight.""" + backend.registry[task.name] = task + reschedule.track_task_start(task, task_id, (), kwargs) + + +class TestTracking: + def test_singleton_task_is_tracked(self, fake_backend): + task = _make_singleton_task() + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + assert fake_backend.redis.hashes["pyworkflow:inflight:test-pod-0"] + + def test_non_singleton_task_is_ignored(self, fake_backend): + task = MagicMock() # not a SingletonWorkflowTask + reschedule.track_task_start(task, "task-1", (), {"run_id": "r1"}) + assert fake_backend.redis.hashes == {} + + def test_singleton_without_unique_on_is_ignored(self, fake_backend): + task = _make_singleton_task(unique_on=None) + reschedule.track_task_start(task, "task-1", (), {"run_id": "r1"}) + assert fake_backend.redis.hashes == {} + + def test_track_end_clears_entry(self, fake_backend): + task = _make_singleton_task() + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + reschedule.track_task_end("task-1") + assert fake_backend.redis.hgetall("pyworkflow:inflight:test-pod-0") == {} + + def test_none_task_id_is_ignored(self, fake_backend): + task = _make_singleton_task() + reschedule.track_task_start(task, None, (), {"run_id": "r1", "step_id": "s1"}) + assert fake_backend.redis.hashes == {} + + +class TestReschedule: + def test_reenqueues_inflight_task(self, fake_backend): + task = _make_singleton_task() + kwargs = {"run_id": "r1", "step_id": "s1", "args_json": "[]"} + _track(fake_backend, task, "task-1", kwargs) + + reschedule.reschedule_inflight_on_shutdown() + + # Lock released with the captured args/kwargs so the fresh copy isn't deduped. + task.release_lock.assert_called_once_with(task_args=[], task_kwargs=kwargs) + # Re-enqueued exactly once, back onto the original queue. + task.apply_async.assert_called_once() + _, call_kwargs = task.apply_async.call_args + assert call_kwargs["queue"] == "pyworkflow.steps" + assert call_kwargs["kwargs"] == kwargs + # Registry drained so a later shutdown call is a no-op. + assert fake_backend.redis.hgetall("pyworkflow:inflight:test-pod-0") == {} + + def test_clean_completion_is_not_reenqueued(self, fake_backend): + task = _make_singleton_task() + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + reschedule.track_task_end("task-1") # task finished normally + + reschedule.reschedule_inflight_on_shutdown() + + task.apply_async.assert_not_called() + + def test_queue_falls_back_to_task_queue(self, fake_backend): + task = _make_singleton_task(routing_key=None, queue="pyworkflow.schedules") + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + + reschedule.reschedule_inflight_on_shutdown() + + _, call_kwargs = task.apply_async.call_args + assert call_kwargs["queue"] == "pyworkflow.schedules" + + def test_handler_swallows_apply_async_errors(self, fake_backend): + task = _make_singleton_task() + task.apply_async.side_effect = RuntimeError("broker down") + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + + # Must never raise out of the shutdown path. + reschedule.reschedule_inflight_on_shutdown() + + def test_release_lock_failure_does_not_block_reenqueue(self, fake_backend): + task = _make_singleton_task() + task.release_lock.side_effect = RuntimeError("redis down") + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + + reschedule.reschedule_inflight_on_shutdown() + + task.apply_async.assert_called_once() + + def test_unknown_task_name_is_skipped(self, fake_backend): + task = _make_singleton_task(name="pyworkflow.execute_step") + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + # Simulate the task no longer being registered at shutdown time. + fake_backend.registry.clear() + + # Should not raise even though the name can't be resolved. + reschedule.reschedule_inflight_on_shutdown() + task.apply_async.assert_not_called() + + def test_multiple_inflight_tasks_all_reenqueued(self, fake_backend): + t1 = _make_singleton_task(name="pyworkflow.execute_step") + t2 = _make_singleton_task( + name="pyworkflow.resume_workflow", + queue="pyworkflow.schedules", + routing_key="pyworkflow.schedules", + ) + _track(fake_backend, t1, "task-1", {"run_id": "r1", "step_id": "s1"}) + _track(fake_backend, t2, "task-2", {"run_id": "r2", "step_id": "s2"}) + + reschedule.reschedule_inflight_on_shutdown() + + t1.apply_async.assert_called_once() + t2.apply_async.assert_called_once() + + +class TestNoBackend: + def test_no_redis_backend_is_noop(self, monkeypatch): + monkeypatch.setattr(reschedule, "_get_backend", lambda: None) + task = _make_singleton_task() + # Tracking and reschedule both degrade to no-ops without raising. + reschedule.track_task_start(task, "task-1", (), {"run_id": "r1", "step_id": "s1"}) + reschedule.reschedule_inflight_on_shutdown() + task.apply_async.assert_not_called() + + +class TestDisabled: + def test_disabled_skips_tracking(self, fake_backend, monkeypatch): + monkeypatch.setenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "0") + task = _make_singleton_task() + reschedule.track_task_start(task, "task-1", (), {"run_id": "r1", "step_id": "s1"}) + assert fake_backend.redis.hashes == {} + + @pytest.mark.parametrize("value", ["false", "no", "0", "FALSE", "No"]) + def test_disabled_values(self, monkeypatch, value): + monkeypatch.setenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", value) + assert reschedule._enabled() is False + + @pytest.mark.parametrize("value", ["1", "true", "yes", ""]) + def test_enabled_values(self, monkeypatch, value): + monkeypatch.setenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", value) + assert reschedule._enabled() is True + + def test_disabled_reschedule_is_noop(self, fake_backend, monkeypatch): + task = _make_singleton_task() + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + monkeypatch.setenv("PYWORKFLOW_RESCHEDULE_ON_SIGTERM", "0") + + reschedule.reschedule_inflight_on_shutdown() + + task.apply_async.assert_not_called() + + +@pytest.fixture +def _stopping(monkeypatch): + """Simulate the worker being in SIGTERM shutdown (state.should_stop set).""" + from celery.worker import state + + monkeypatch.setattr(state, "should_stop", 0, raising=False) + monkeypatch.setattr(state, "should_terminate", None, raising=False) + + +class TestConsumerStopHook: + """ + Drives the prerun tracking signal + the consumer bootstep teardown hook + (on_consumer_stop), which is the real event-based path on shutdown. + """ + + def test_prerun_then_consumer_stop_roundtrip(self, fake_backend, _stopping): + from pyworkflow.celery import app as app_module + + task = _make_singleton_task() + fake_backend.registry[task.name] = task + kwargs = {"run_id": "r1", "step_id": "s1"} + + app_module.on_task_prerun(task_id="task-1", task=task, args=(), kwargs=kwargs) + assert fake_backend.redis.hgetall("pyworkflow:inflight:test-pod-0") + + consumer = MagicMock() # has .task_consumer.cancel() + reschedule.on_consumer_stop(consumer) + + consumer.task_consumer.cancel.assert_called_once() # stop fetching first + task.apply_async.assert_called_once() # then re-enqueue + + def test_consumer_stop_is_noop_when_not_shutting_down(self, fake_backend, monkeypatch): + # Transient consumer restart (e.g. broker reconnect): flags are clear. + from celery.worker import state + + monkeypatch.setattr(state, "should_stop", None, raising=False) + monkeypatch.setattr(state, "should_terminate", None, raising=False) + + task = _make_singleton_task() + _track(fake_backend, task, "task-1", {"run_id": "r1", "step_id": "s1"}) + + consumer = MagicMock() + reschedule.on_consumer_stop(consumer) + + consumer.task_consumer.cancel.assert_not_called() + task.apply_async.assert_not_called() + + def test_postrun_clears_before_consumer_stop(self, fake_backend, _stopping): + from pyworkflow.celery import app as app_module + + task = _make_singleton_task() + fake_backend.registry[task.name] = task + app_module.on_task_prerun( + task_id="task-1", task=task, args=(), kwargs={"run_id": "r1", "step_id": "s1"} + ) + app_module.on_task_postrun(task_id="task-1") + reschedule.on_consumer_stop(MagicMock()) + task.apply_async.assert_not_called() + + def test_bootstep_requires_tasks_and_delegates(self, fake_backend, _stopping, monkeypatch): + # The bootstep must run after the Tasks consumer step exists. + assert "celery.worker.consumer:Tasks" in reschedule.RescheduleConsumerStep.requires + # And its stop() delegates to on_consumer_stop. + called = {} + monkeypatch.setattr(reschedule, "on_consumer_stop", lambda c: called.setdefault("c", c)) + step = object.__new__(reschedule.RescheduleConsumerStep) # skip Step.__init__ + sentinel = MagicMock() + step.stop(sentinel) + assert called.get("c") is sentinel diff --git a/tests/unit/test_singleton.py b/tests/unit/test_singleton.py index 4041c14..c19fc34 100644 --- a/tests/unit/test_singleton.py +++ b/tests/unit/test_singleton.py @@ -2,7 +2,7 @@ Unit tests for SingletonWorkflowTask and related functionality. """ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -950,3 +950,60 @@ def test_sentinel_master_none(self): config = SingletonConfig(mock_app) assert config.sentinel_master is None + + +class TestReleaseLockKeyReproduction: + """ + The SIGTERM reschedule path (pyworkflow/celery/reschedule.py) releases a + running task's singleton lock by calling release_lock() with the task's + captured args/kwargs, then re-enqueues. This verifies release_lock() unlocks + exactly the key generate_lock() produced for a keyword-arg dispatch (how + execute_step_task is actually called). + """ + + @pytest.fixture + def mock_celery_app(self): + app = MagicMock() + app.conf.get.side_effect = lambda key, default=None: { + "singleton_key_prefix": "pyworkflow:lock:", + "singleton_lock_expiry": 3600, + "singleton_raise_on_duplicate": False, + }.get(key, default) + app.conf.broker_url = "redis://localhost:6379/0" + return app + + def test_release_lock_unlocks_generated_key_for_kwargs_dispatch(self, mock_celery_app): + class StepTask(SingletonWorkflowTask): + name = "pyworkflow.execute_step" + unique_on = ["run_id", "step_id"] + + def run(self, run_id, step_id, args_json=None, storage_config=None): + return None + + task = StepTask() + task.bind(mock_celery_app) + + backend = MagicMock() + with patch.object( + SingletonWorkflowTask, + "singleton_backend", + new_callable=PropertyMock, + return_value=backend, + ): + # Mirrors the real dispatch: keyword-only args. + task.release_lock( + task_args=[], + task_kwargs={ + "run_id": "run_123", + "step_id": "step_456", + "args_json": "[]", + }, + ) + + expected_key = generate_lock_key( + "pyworkflow.execute_step", + ["run_123", "step_456"], + {}, + key_prefix="pyworkflow:lock:", + ) + backend.unlock.assert_called_once_with(expected_key)