refactor(python): share one task lifecycle across both dispatch paths#438
refactor(python): share one task lifecycle across both dispatch paths#438pratyush618 wants to merge 13 commits into
Conversation
The blocking wrapper and the async executor each implement the lifecycle a task runs through, and they have drifted: the native copy silently lost the queue hooks, the saga compensation context, soft_timeout and per-item batch handling. Patching the gaps into the executor would make a second copy that drifts again. Move the body to taskito/task_lifecycle.py and have the wrapper drive it. Pure move — the callers keep only what genuinely differs, which is where arguments come from and where the result goes. current_job needs no special handling: its context lookup already reads the contextvar before the thread-local. soft_timeout moves into a _task_soft_timeouts dict, mirroring the eleven other per-task options. It was the only one reachable solely through the wrapper's closure, which is why the executor could never honour it. _clear_context stays with the wrapper: it is edge state, and in the shared body it would null out the thread-local of whichever thread the executor loop runs on. No behaviour change — 1217 tests unchanged.
Drives each lifecycle guarantee through a real worker, parametrised over def and async def. The async cases route sync while native dispatch is dormant; they are the net for activating it.
`except Exception` misses a BaseException, so the cleanup saw no error and emitted JOB_COMPLETED for a task whose body never returned. Latent today; asyncio.CancelledError is a BaseException and is routine once async tasks dispatch natively.
The executor reimplemented the lifecycle and had drifted, silently missing the queue hooks, saga context, soft_timeout and per-item batch handling. It now awaits run_lifecycle and keeps only what differs: decoding the payload and answering the result sender.
The lifecycle held the exception in a plain local and passed it as a log argument, so both the frame and any retained LogRecord kept it — and through its traceback, every frame that raised. Latent on the blocking path; on native dispatch those frames own the result sender the worker's drain loop waits on.
The lifecycle gaps the old note cites are closed; shutdown is what is left.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR centralizes task lifecycle execution in ChangesShared Task Lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant AsyncTaskExecutor
participant run_lifecycle
participant TaskFunction
Worker->>AsyncTaskExecutor: Deserialize payload and start job
AsyncTaskExecutor->>run_lifecycle: Await lifecycle execution
run_lifecycle->>TaskFunction: Execute task and await awaitable result
TaskFunction-->>run_lifecycle: Return result or raise exception
run_lifecycle-->>AsyncTaskExecutor: Return result or re-raise exception
Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
sdks/python/tests/worker/test_native_async.py (1)
24-52: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake
_fake_queuereject undeclared lifecycle attributes.A bare
MagicMockstill fabricates newly accessed attributes, recreating the silent wrong-branch behavior this helper is intended to prevent. Usespec_setor a concrete stub so new lifecycle dependencies fail until explicitly configured.🤖 Prompt for 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. In `@sdks/python/tests/worker/test_native_async.py` around lines 24 - 52, Update _fake_queue to construct the queue mock with spec_set or a concrete stub that declares every lifecycle attribute configured and accessed by the helper. Preserve the existing defaults and overrides, while ensuring any newly accessed undeclared attribute raises immediately instead of being auto-created.
🤖 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 `@sdks/python/taskito/mixins/decorators.py`:
- Around line 419-421: The task registration logic around _wrap_task must remove
any existing _task_soft_timeouts entry when soft_timeout is None, while
continuing to store a provided timeout. Ensure replacing a task by task_name
cannot inherit a previous timeout.
In `@sdks/python/taskito/task_lifecycle.py`:
- Around line 269-314: Wrap the teardown operations in the finally
block—including compensation reset, resource release, proxy cleanup, hooks,
middleware, and event emission—in a nested try/finally, and move the existing
`error = None` into that nested finally. Ensure exception clearing runs even
when any teardown callback or cleanup operation raises, while preserving the
current cleanup order and event behavior.
- Around line 117-187: Move the cleanup guard to begin before proxy
reconstruction and resource acquisition in the task execution flow. Enclose
proxy cleanup, resource release callbacks, middleware setup, and before_task
hooks in an outer try/finally so exceptions during reconstruction, acquisition,
middleware lookup, or hooks still run all registered cleanup; preserve the
existing task execution error handling inside that guard.
- Around line 188-191: The run_lifecycle execution path must await awaitable
results returned by synchronous tasks instead of calling run_maybe_async inside
its active event loop. Update task_lifecycle.py lines 188-191 around
run_lifecycle to detect and await returned awaitables while preserving direct
handling for non-awaitable values; retain run_maybe_async in
mixins/decorators.py lines 157-166 as the blocking entrypoint, requiring no
direct change there.
- Around line 240-256: Handle non-cancellation BaseException cases in
async_support/executor.py around the async executor flow by emitting a terminal
non-retryable failure or equivalent disposition before re-raising
KeyboardInterrupt, SystemExit, GeneratorExit, and similar exceptions. In
task_lifecycle.py around the existing BaseException handler, preserve the
teardown-only behavior and do not add failure hooks or disposition reporting
there; the executor must own terminal reporting before propagation.
In `@sdks/python/tests/worker/test_lifecycle_parity.py`:
- Around line 124-150: Update the failure capture in on_fail and the
corresponding hook handlers at
sdks/python/tests/worker/test_lifecycle_parity.py:124-150, 320-349, and 366-404
to store only each exception’s type and message (or equivalent derived values),
never the exception object or traceback; adjust the assertions to validate those
stored values while preserving the existing failure checks before _stop_worker
runs.
- Around line 65-68: Update _stop_worker to verify that the worker thread has
terminated after thread.join(timeout=10), and fail teardown if thread.is_alive()
remains true. Preserve the existing shutdown request and timeout while ensuring
a still-running daemon thread cannot allow the test to pass.
---
Nitpick comments:
In `@sdks/python/tests/worker/test_native_async.py`:
- Around line 24-52: Update _fake_queue to construct the queue mock with
spec_set or a concrete stub that declares every lifecycle attribute configured
and accessed by the helper. Preserve the existing defaults and overrides, while
ensuring any newly accessed undeclared attribute raises immediately instead of
being auto-created.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0a52a6f4-2f29-4aa4-8f2f-0cb3cd52f4c1
📒 Files selected for processing (7)
sdks/python/taskito/app.pysdks/python/taskito/async_support/executor.pysdks/python/taskito/mixins/decorators.pysdks/python/taskito/task_lifecycle.pysdks/python/tests/observability/test_task_lifecycle_logs.pysdks/python/tests/worker/test_lifecycle_parity.pysdks/python/tests/worker/test_native_async.py
A plain `def` handing back a coroutine worked before the lifecycle became one: run_maybe_async now nests inside the lifecycle's own loop and raises. Deciding on the returned value covers both kinds and drops the is_async parameter.
Resources and proxies are acquired before the body is entered, so a raising before_task hook stranded them. The teardown could also skip clearing the exception if a hook raised, which was the point of clearing it.
Registering a name again replaces the task, but left the old timeout in the map for the new one to inherit.
A BaseException that is not a cancellation slipped past every clause, so the job reported nothing and sat Running until the reaper called it a timeout. The queue stub is now concrete: a bare MagicMock invents whatever the lifecycle asks for, which is how a missing attribute takes the wrong branch in silence.
The stop helper returned on a live worker, which is how the shutdown stall in this series nearly escaped: assertions had already passed, so only the clock showed it. The hooks also held live exceptions — the stall's own trigger.
|
Thanks — all seven addressed, and one of them was a genuine regression I'd introduced. Details, since a couple needed digging: Awaitable returns ( Fixed as suggested — await what is awaitable — which also made the Cleanup guard before acquiring ( Guaranteed exception clearing ( Non-cancellation
Tests retaining exceptions. Accepted — they now record type and message. Retention is the documented stall condition, so holding live exceptions in the very tests meant to be the net for activation would have made each one slow while proving nothing extra. Nitpick — Verification: full suite green, |
There was a problem hiding this comment.
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 `@sdks/python/taskito/task_lifecycle.py`:
- Around line 178-210: Update the setup-failure handling around the lifecycle
orchestration so every middleware recorded in completed_mw receives its after()
teardown when a later before_task hook or saga setup step raises. Invoke
teardown for completed middleware before releasing resources, while preserving
compensation-context reset, cleanup, and exception propagation.
In `@sdks/python/tests/worker/test_native_async.py`:
- Around line 576-600: Update test_base_exception_reports_failure to verify
shutdown after executor.stop(), using the existing shutdown assertion helper if
available or asserting that the executor thread is no longer alive. Keep the
existing failure-reporting and KeyboardInterrupt behavior unchanged.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4c0d22d4-546c-4ef2-ac38-dc5b44db0e6f
📒 Files selected for processing (5)
sdks/python/taskito/async_support/executor.pysdks/python/taskito/mixins/decorators.pysdks/python/taskito/task_lifecycle.pysdks/python/tests/worker/test_lifecycle_parity.pysdks/python/tests/worker/test_native_async.py
🚧 Files skipped from review as they are similar to previous changes (1)
- sdks/python/taskito/mixins/decorators.py
`completed_mw` exists so every before() gets an after(); the setup unwind released resources but skipped it, leaving middleware-owned state behind.
`stop()` only logs when its join times out, so the test could pass on a leaked thread — the same way the worker stop helper hid a stall.
|
Both accepted — each one a gap in my own fix from the last round. Middleware teardown on setup failure. Right, and I'd drawn the line in the wrong place. I reasoned that a task which never started shouldn't fire The hooks stay as they were, deliberately: there's no Asserting the executor thread stopped. Fair, and the same trap as Full suite green, |
Why
Native async dispatch has never run: the pool reads
_taskito_is_asyncoff the registry entry, butthe decorator sets it on the
TaskWrapperit returns, so every async task takes the blocking branch.Activating it was a two-line change gated behind one problem —
_wrap_taskandAsyncTaskExecutor._run_jobwere two independent implementations of one lifecycle, and the nativecopy had silently lost the queue hooks, the saga compensation context,
soft_timeoutand per-itembatch handling.
This PR removes that gate by making the two paths share one body, and fixes the bugs that surfaced
while doing it. It does not activate native dispatch — see below.
What changed
One lifecycle. New
taskito/task_lifecycle.py::run_lifecycleholds everything between "we havethe arguments" and "we have a result".
_wrap_taskdrives it withrun_maybe_async; the executorawaits it. Each caller keeps only what genuinely differs — how arguments arrive, and how the result
is delivered.
executor.pyloses 101 lines and no longer owns a single lifecycle concern, so a newlifecycle rule can't reach one path and miss the other.
Real bugs, each fixed with a regression test watched failing first:
except Exceptiondoesn't catch aBaseException,so the cleanup saw no error and emitted
JOB_COMPLETED— and handederror=Nonetoafter_task/middleware — for a task whose body never returned. Latent on the blocking path;asyncio.CancelledErroris aBaseExceptionand is routine once dispatch goes native.auto-clears an
except ... asname for exactly this reason;erroris not one) and passed as alogging argument, which a
LogRecordkeeps. Either one retains the exception, its traceback, andevery frame that raised.
masteranddead-lettered here: the lifecycle is itself a coroutine now, so the inner
run_maybe_asyncnestedinside its own event loop and raised. Awaiting whatever is awaitable covers both kinds and made
the
is_asyncparameter redundant.entered, but only the body's teardown gave them back — so a raising
before_taskhook leakedthem. Setup now unwinds explicitly.
BaseExceptionthat is not a cancellation slippedpast every clause in the executor, so the job reported nothing and sat Running until the reaper
called it a timeout.
Tests: 22 new, parametrised over
def/async defwhere the guarantee is shared and driven through a real@queue.task()anda real worker — hooks, saga context,
soft_timeout(which had no coverage at all), per-item batch,logging, middleware-on-cancel. Each was watched failing with its guarantee removed. The async cases
route sync while dispatch is dormant; they are the net for activating it.
tests/worker/test_native_async.pygrows a_fake_queue()helper, replacing ten hand-listed mocksetups. That mattered immediately: the lifecycle branches on
is not Noneand truthiness, so aMagicMock's auto-made
_task_batch_configsreads as a per-item batch config and every task fails thereturn-type contract. One list, not ten.
Why activation isn't here
Flipping both flags works — coroutines run on
taskito-async-executorand all 18 parity tests passon the native path. Then measurement showed a failed async task stalls worker shutdown for the full
30s drain timeout: 30.13s with nothing more exotic than an
on_failurehook that collects errors.That is past a typical container's grace period, so in-flight jobs get SIGKILLed.
The worker's drain loop ends only when the result channel disconnects, which needs the executor's
PyResultSenderto drop — and a failed task's traceback pins the frame that owns it. The tworetention sources above are fixed, but a user hook that keeps the exception can't be fixed from
Python: every frame in that traceback can reach the sender. The drain has to stop depending on Python
object lifetime, which is a design change worth its own PR.
ROADMAP.mdrecords the findings for it,including why
Scheduler.in_flightis a tempting but unsafe drain signal.So the flags stay off and the note explaining why is updated to name the real blocker.
Verification
1239 passed, 4 skipped— the 1217 baseline plus these 22, no regressions.ruff check/ruff format/mypyclean ontaskito/andtests/.cargo check --workspace, and with--features native-async(no Rust changed).than inferred (0.17s vs 30.13s).
Summary by CodeRabbit