Process worker: drain in-flight tasks before exiting

[wolever-gl]

Summary

This PR creates a GrowthLoop vendor build of Prefect 3.6.27+growthloop and adds an opt-in graceful-drain mode for process workers during SIGTERM shutdown.

The production use case is Kubernetes rollout/termination behavior. When a process-worker pod receives SIGTERM, the worker should stop accepting new flow runs while allowing already-started flow subprocesses to finish, instead of immediately forwarding an interrupt that can kill in-flight work.

Behavior

When PREFECT_WORKER_DRAIN_ON_SIGTERM=true is set for a process worker:

• The worker CLI registers the SIGTERM handler after the worker instance exists.
• On Unix, the first SIGTERM requests drain instead of forwarding SIGINT immediately.
• Drain mode stops the process worker polling/sync service loops so the worker stops accepting new scheduled work.
• The process worker waits for real in-flight run infrastructure to exit before stopping.
• A second SIGTERM escalates by sending SIGINT to the worker process.

When the env var is unset or not true, worker SIGTERM behavior should remain upstream behavior. Non-process workers are not wired into this vendor drain path.

Deploy / Install Instructions

The current wheel has been deployed to the public GrowthLoop bucket:

• GCS: gs://growthloop-wolever-public/prefect/prefect-3.6.27+growthloop-py3-none-any.whl
• Versioned GCS: gs://growthloop-wolever-public/prefect/3.6.27+growthloop/prefect-3.6.27+growthloop-py3-none-any.whl
• Public URL: https://storage.googleapis.com/growthloop-wolever-public/prefect/prefect-3.6.27%2Bgrowthloop-py3-none-any.whl
• SHA256: ec6635650c11294f85f5aa2a10d898d560c6266570b3cebc67b422226b5e6fc1
• Size: 1820051 bytes

Install or upgrade an environment directly from the wheel:

python -m pip install --force-reinstall \
  'https://storage.googleapis.com/growthloop-wolever-public/prefect/prefect-3.6.27%2Bgrowthloop-py3-none-any.whl'

With uv:

uv pip install --reinstall-package prefect \
  'https://storage.googleapis.com/growthloop-wolever-public/prefect/prefect-3.6.27%2Bgrowthloop-py3-none-any.whl'

Verify the installed version:

python -c 'import prefect; print(prefect.__version__)'
# 3.6.27+growthloop

Enable drain behavior for process workers by setting the env var on the worker container/process:

export PREFECT_WORKER_DRAIN_ON_SIGTERM=true
prefect worker start --pool <work-pool-name> --type process

For Kubernetes deployments, add PREFECT_WORKER_DRAIN_ON_SIGTERM=true to the worker pod env and make sure terminationGracePeriodSeconds is long enough for the longest in-flight flow run you intend to let finish.

Implementation Details

Core changes in this PR:

• pyproject.toml pins the vendor package version to 3.6.27+growthloop.
• README.md documents that this repository is a GrowthLoop-maintained vendor fork and includes wheel build instructions.
• src/prefect/cli/worker.py defers worker signal handler setup until after worker construction and scopes drain wiring to ProcessWorker instances.
• src/prefect/utilities/processutils/__init__.py adds drain-state helpers and SIGTERM drain handling support. open_process() should skip terminating a child only on a clean drain exit, but should still close process resources and should terminate on cancellation/escalation.
• src/prefect/workers/base.py adds worker drain state and keeps cleanup of submitting flow-run ids in finally paths.
• src/prefect/workers/process.py handles process-worker drain shutdown and tracks in-flight infrastructure lifetime past subprocess PID handoff.
• tools/write_build_info.py is used as a Hatch custom build hook so clean static-version wheel builds still include prefect._build_info.

PR Comments / Review Follow-up

Cursor Bugbot reported two issues:

1. _active_flow_run_ids was cleared after subprocess start, not subprocess completion.
2. open_process() skipped process.aclose() during drain.

The follow-up fixes for both issues, the clean-wheel _build_info failure, non-process worker drain scoping, and the in-flight submit-loop race have been committed and pushed to this PR branch. The relevant Bugbot threads are resolved.

Current Status

Resolved locally and included in the deployed wheel:

• Clean checkout wheel build/install/import works with prefect.__version__ == "3.6.27+growthloop".
• open_process() closes resources during drain and terminates on cancellation/escalation.
• Drain wiring is scoped to process workers.
• Process-worker drain waits on in-flight infrastructure lifetime, including unlimited workers.

Still open:

• The PR branch is still merge-conflicted with main. A direct in-place merge was conflict-free after resolution, but it staged a large upstream delta, so it was intentionally unwound. This needs a deliberate rebase/merge strategy before final merge.
• A full CLI/OS-signal test against a local server remains recommended before relying on this in production. The process-worker lifecycle test now exercises an actual flow subprocess held past drain, confirms a post-drain heartbeat, and verifies the worker exits only after release.

Automated Test Plan

Implemented or prepared locally:

• Build-info unit test for static vendor version generation.
• Clean archive wheel smoke: build from a git archive without src/prefect/_build_info.py, install into a fresh venv, and import prefect.
• Processutils tests for clean drain completion and cancellation/escalation cleanup.
• CLI tests for default/unset env behavior, process-worker drain wiring, and non-process worker opt-out.
• Process-worker tests proving drain waits after drain request and in-flight tracking survives subprocess PID handoff for limit=None and limit=1.
• Real process-worker lifecycle test that holds a flow subprocess during drain, observes a post-drain heartbeat, and verifies the worker exits only after the flow completes.

Recommended additional test before production rollout:

• End-to-end CLI worker signal test: start a local Prefect server and process worker, submit a long-running flow, send SIGTERM to the worker, assert the flow completes and the worker stops only after completion.
• Escalation test: repeat with a second SIGTERM/SIGINT while the flow is still running and assert no orphan child process remains.
• Kubernetes smoke: run the worker in a pod with terminationGracePeriodSeconds, launch a sleep flow, delete the pod, and assert the flow completes before the pod exits; repeat with a short grace period or second signal to verify escalation.

Validation Run Locally

uv run pytest tests/test_build_info.py \
  tests/utilities/test_processutils.py \
  tests/cli/test_worker.py::test_start_worker_does_not_wire_sigterm_drain_by_default \
  tests/cli/test_worker.py::test_start_worker_wires_sigterm_drain_for_process_worker \
  tests/cli/test_worker.py::test_start_worker_does_not_wire_sigterm_drain_for_non_process_worker \
  tests/workers/test_process_worker.py::test_process_worker_start_waits_after_drain_request \
  tests/workers/test_process_worker.py::test_process_worker_tracks_in_flight_runs_until_infrastructure_exits -q

Result: the focused follow-up suite passed 26 passed, 3 skipped; the real process-worker drain lifecycle module passed 19 passed; and the base-worker suite passed 104 passed, 1 skipped.

uv run ruff format --check tools/write_build_info.py tests/test_build_info.py \
  src/prefect/utilities/processutils/__init__.py tests/utilities/test_processutils.py \
  src/prefect/workers/process.py src/prefect/cli/worker.py \
  tests/cli/test_worker.py tests/workers/test_process_worker.py

uv run ruff check pyproject.toml tools/write_build_info.py tests/test_build_info.py \
  src/prefect/utilities/processutils/__init__.py tests/utilities/test_processutils.py \
  src/prefect/workers/process.py src/prefect/cli/worker.py \
  tests/cli/test_worker.py tests/workers/test_process_worker.py

Result: format check and ruff check passed.

---

Note

Medium Risk
Changes worker shutdown, subprocess lifecycle, and packaging for a forked runtime—important for K8s rollouts but scoped behind an env flag and process workers only.

Overview
This PR vendors Prefect as 3.6.27+growthloop (README fork notes, pinned pyproject.toml version, Hatch write_build_info hook for wheels without git) and adds opt-in graceful shutdown for process workers via PREFECT_WORKER_DRAIN_ON_SIGTERM=true.

When enabled, the worker CLI registers SIGTERM after the worker exists and passes ProcessWorker.request_drain into signal setup. On Unix, the first SIGTERM enters drain mode (stops polling, keeps heartbeats/healthchecks healthy) instead of immediately forwarding SIGINT; a second SIGTERM escalates to SIGINT. BaseWorker stops scheduling new runs while draining; ProcessWorker tracks in-flight runs until subprocess infrastructure exits. open_process skips terminate() only on a clean drain exit but still **aclose()**s; cancellation still terminates.

Default SIGTERM behavior is unchanged when the env var is unset; non-process workers are not wired to drain.

^{Reviewed by Cursor Bugbot for commit 3103838. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 3103838. Configure here.}

[cursor]

Drain escalation cancels nothing

High Severity

A second SIGTERM is supposed to escalate out of drain, but ProcessWorker.start swallows asyncio.CancelledError whenever _draining is set. That turns escalation into a normal context exit, so teardown waits on _runs_task_group instead of cancelling it, and in-flight runs keep draining. The SIGTERM handler also only ever sends SIGINT, never SIGKILL, so a stuck drain may never force-stop.

Additional Locations (1)

• src/prefect/utilities/processutils/__init__.py#L592-L598

 

^{Reviewed by Cursor Bugbot for commit 3103838. Configure here.}

[wolever-gl]

This fixes a bug in BaseWorker.start (link) which assumes that self._limiter can be used to track the number of in-flight run executions, and it's safe to exit when bool(self._limiter) == False. When a subprocess limit is set, this is safe. But when self._limiter = None - the default - the main process will terminate even if there are in-flight runs.