feat(celery): finalize in-flight task on worker eviction so the chord proceeds

[ocervell]

Problem

On a K8s pod eviction (SIGTERM → grace → SIGKILL) the worker running a task dies. With task_acks_late on the Redis broker, its message is only redelivered after the broker visibility timeout — which we derive above each pool's deadline (hours, on the long-task pool). So the surrounding chord/workflow stalls for that entire window.

This is the one OOM/eviction case the other layers don't cover:

Scenario	Recovery
Task timeout / memory soft-limit	secator monitor self-stops → partial results (existing)
Fast-spike OOM (child killed, master survives)	task_reject_on_worker_lost requeues immediately (~18s), capped by #1199
Whole-pod eviction (master dies too)	❌ nothing requeues fast on Redis → visibility-timeout stall

A whole-pod eviction has no surviving master, so reject_on_worker_lost can't help and Redis falls back to the visibility timeout.

Fix

Catch the worker_shutting_down signal and raise a flag the running task's monitor thread polls. On shutdown it stops the task early via the existing stop_process(exit_ok=True) path — the same self-stop already used for the timeout and memory-limit cases — so the task returns its partial results through the normal completion path and the chord proceeds in seconds instead of waiting for the visibility timeout.

• celery_signals.py — SHUTDOWN_FLAG + worker_shutting_down_handler (wired unconditionally); stale flag cleared on worker boot. File-based so it crosses the prefork master→child boundary.
• command.py — the monitor checks the flag, and its poll interval is capped (MONITOR_POLL_SECONDS=5) so an eviction is caught well within the pod's terminationGracePeriodSeconds (60s).

Goes through Celery's normal completion path, so the chord is satisfied the ordinary way — no pool change (stays prefork), no Celery upgrade, no revoke (which yields REVOKED/TaskRevokedError, the wrong tool here — that's for cancel, e.g. stop_runner).

Tests

• test_shutdown_flag_lifecycle — flag set/clear.
• test_monitor_stops_running_command_on_shutdown — a real sleep 30 command stops early (well under 30s) and emits the eviction Warning.

Validated end-to-end locally: evicting a worker mid-chord, the chord callback fired in ~5s (vs. the ~hours visibility timeout) with partial results from the evicted member.

Part of the OOM/eviction robustness set alongside #1199 (worker-loss retry cap) and #1202 (on_build runner identity).

🤖 Generated with Claude Code

https://claude.ai/code/session_01P5vSjfkBuGAAHdKxHS3ySm

Summary by CodeRabbit

• New Features

  • Implemented graceful worker shutdown that terminates running commands and automatically saves partial results, preventing indefinite task hangs during system shutdown
  • Enhanced command monitoring to responsively detect shutdown signals and exit long-running tasks quickly

• Tests

  • Added comprehensive unit tests validating worker shutdown eviction and command termination behavior

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 769d31df-c3ea-4935-90d6-11086a1bb6c3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

Walkthrough

Adds a file-based worker shutdown flag (SHUTDOWN_FLAG) in celery_signals.py with helpers to set, check, and clear it. The flag is written when Celery fires worker_shutting_down. The command monitor loop in command.py polls this flag and, when set, sends SIGTERM to the running subprocess and exits early. A new unit test module validates the full eviction lifecycle.

Changes

Worker Eviction Self-Finalize

Layer / File(s)	Summary
Shutdown flag API and Celery signal wiring secator/celery_signals.py	Defines SHUTDOWN_FLAG path, is_worker_shutting_down(), clear_shutdown_flag(), and worker_shutting_down_handler() that writes the marker on eviction. setup_handlers() is updated to clear stale flags on boot and register the handler unconditionally with signals.worker_shutting_down.
Command monitor shutdown detection secator/runners/command.py	Adds MONITOR_POLL_SECONDS = 5, imports is_worker_shutting_down inside _monitor_process, inserts a shutdown branch that enqueues a Warning and calls stop_process(exit_ok=True) with SIGTERM, and caps the monitor sleep to min(stat_update_frequency, MONITOR_POLL_SECONDS) for faster reaction.
Eviction unit tests tests/unit/test_eviction.py	TestEvictionSelfFinalize verifies flag set/clear lifecycle, confirms a sleep 30 command terminates promptly when the shutdown flag is raised mid-execution (with MONITOR_POLL_SECONDS patched to speed up polling), and checks that an eviction warning appears in results.

Sequence Diagram(s)

sequenceDiagram
    participant Celery as Celery Worker
    participant SignalHandler as worker_shutting_down_handler
    participant FlagFile as SHUTDOWN_FLAG (filesystem)
    participant Monitor as _monitor_process loop
    participant Subprocess as Running subprocess

    Celery->>SignalHandler: signals.worker_shutting_down fires
    SignalHandler->>FlagFile: touch(SHUTDOWN_FLAG)
    loop every min(stat_update_frequency, 5s)
        Monitor->>FlagFile: is_worker_shutting_down()
        FlagFile-->>Monitor: True
        Monitor->>Monitor: enqueue Warning("worker shutting down")
        Monitor->>Subprocess: SIGTERM via stop_process(exit_ok=True)
        Monitor->>Monitor: break — save incomplete results
    end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• freelabz/secator#727: Introduces the _monitor_process subsystem in secator/runners/command.py that this PR extends with the is_worker_shutting_down() early-exit branch.

Poem

🐇 Hop hop, the worker must go,
A flag file appears, soft and low.
The monitor peeks every five,
"SIGTERM!" — no process alive.
Incomplete tasks saved with grace,
The chord moves on to its rightful place! 🏁

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'feat(celery): finalize in-flight task on worker eviction so the chord proceeds' directly summarizes the main change: gracefully finalizing in-flight tasks during worker eviction to allow chord workflows to proceed quickly.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/worker-eviction-finalize

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@secator/celery_signals.py`:
- Around line 36-41: The worker_shutting_down_handler function silently swallows
all exceptions when writing to SHUTDOWN_FLAG, which can mask real issues like
missing or unwritable state directories and cause tasks to stall. Modify the
function to first ensure the parent directory of SHUTDOWN_FLAG exists by
creating it if necessary, then replace the broad exception handler with one that
specifically catches and logs at least OSError exceptions with an appropriate
error message, while still allowing the function to complete gracefully without
raising.

In `@tests/unit/test_eviction.py`:
- Around line 19-23: The setUp and tearDown methods are operating on the real
global shutdown flag file, which causes test interference when running tests in
parallel or with shared STATE_DIR. Modify the setUp method to patch
secator.celery_signals.SHUTDOWN_FLAG to point to a temporary test-specific path
before calling clear_shutdown_flag(), and modify the tearDown method to clean up
that patch after calling clear_shutdown_flag() to ensure each test uses an
isolated shutdown flag path.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e0e1675a-0d72-4b1b-a024-c404d0ceaabe

📥 Commits

Reviewing files that changed from the base of the PR and between d39d440 and 4fe6fb1.

📒 Files selected for processing (3)

• secator/celery_signals.py
• secator/runners/command.py
• tests/unit/test_eviction.py

[ocervell]

Superseded — this work is already merged into ai-resiliency and ships to canary via #1241 (the branch is 0 commits ahead of ai-resiliency). Closing to keep the PR list clear; no work is lost, the commits live in ai-resiliency. Branch deleted.