async-task-api-worker

SDK to write workers for the AsyncTaskAPI platform.

This is a work in progress: the package ships the protocol models (inbound envelope decoder and the byte-exact outbound message encoder), the error types, the settings, the runtime engine with the Worker definition API, the RabbitMQ transport, the process lifecycle (graceful drain), the probes, the metrics and the JSON logs. The conformance kit and the extras land in later releases.

Declaring a worker

from pydantic import BaseModel, Field

from async_task_api_worker import TaskContext, Worker


class Input(BaseModel):
    text: str = Field(min_length=1, max_length=100_000)


class Output(BaseModel):
    words: int


worker = Worker("count-words")


@worker.task(timeout=30.0)
async def handle(task_input: Input, ctx: TaskContext) -> Output:
    await ctx.progress(0.5)
    return Output(words=len(task_input.text.split()))

• async def handlers run on the event loop; def handlers (with a SyncTaskContext) run in a thread pool. The context parameter must be annotated exactly TaskContext or SyncTaskContext (an identity check: no subclass or Annotated wrapper). Hooks and probes (on_startup, on_shutdown, readiness_probe) must be async def.
• The input and output models are read from the annotations. Exactly one task per worker.
• Raise FatalError(code, message) to fail a task with a code the client can match, RetryableError for a transient failure. Any other exception is reported as INTERNAL_ERROR without its text.
• ctx.progress(fraction, payload) reports the latest state of the task, not an event stream: updates are throttled (at most one per WORKER_PROGRESS_MIN_INTERVAL_S, default 1 s), a newer update replaces a held one, fraction is in [0, 1) and never goes back. The terminal success or failure message is the 100%.

Running a worker

export BROKER_URL=amqp://user:password@broker:5672/ IN_QUEUE_NAME=count-words_queue_in OUT_QUEUE_NAME=count-words_queue_out
async-task-api-worker run my_module:worker      # or: python -m async_task_api_worker run my_module:worker

my_module:worker is a Worker, or a zero-argument factory function returning one; modules in the current directory can be imported. A definition error in the module exits with code 2 and one line. From code, worker.run() blocks until the process is stopped and await worker.serve() serves from an existing event loop.

• The settings are read and validated first: an invalid configuration, or a task timeout that does not fit under the broker's consumer timeout (task timeout + WORKER_RETRY_BACKOFF_MAX_S + 60 s < WORKER_CONSUMER_TIMEOUT_S), exits with code 2 before any connection. Exit codes: 0 clean stop, 1 failure, 2 configuration or usage error.
• The queues are declared durable and with no arguments (the broker decides their type and policies), and consumed with manual acknowledgements: a task is acked only after its result is published.
• SIGTERM or SIGINT stops consuming, lets the in-flight tasks finish within WORKER_SHUTDOWN_GRACE_S (default 25 s), requeues the ones that did not, closes the connection and exits 0. A second signal skips the rest of the grace period. A def handler whose thread is still running after the drain, or an async def handler whose asyncio.to_thread thread is (given 2 s after the drain), is logged and the process exits at once with code 1 (only for run, never for serve()). The whole shutdown is bounded: the process exits at most max(WORKER_SHUTDOWN_GRACE_S, 5) + 17 s after the signal (5 s for the nacks of the cancelled tasks, 5 s to close the connection, and 5 s in total, WORKER_SHUTDOWN_HOOKS_TIMEOUT_S, for all the shutdown hooks together, each getting what is left). Set a pod's terminationGracePeriodSeconds to at least that. A signal during startup does not wait for it: the connection attempt and the startup hooks are cancelled and the process exits 0 (a hook that blocks the event loop cannot be interrupted).
• Retries and crash-loop protection need quorum queues (the broker's default queue type, as in the platform's chart). A classic queue gives no delivery counter: there, a RetryableError gets one retry, then concludes with its own code, and a warning says so.
• Refused broker credentials exit with code 2 right away instead of being retried.
• The attempt number of a delivery (ctx.attempt) counts every time the broker handed the message out. A RetryableError requeues until attempt WORKER_DELIVERY_LIMIT, then concludes with its own code. A delivery past that limit (a crash loop) is concluded with TASK_ABANDONED_AFTER_REDELIVERIES without running the handler.

Probes, metrics and logs

A small HTTP server runs on HEALTH_CHECK_HOST:HEALTH_CHECK_PORT (default 0.0.0.0:8083) in both modes, from the very start of the process until the end of the shutdown.

Errors use the platform's format, {"status": "error", "error": {"code", "message"}} (404 NOT_FOUND, 405 METHOD_NOT_ALLOWED, 500 PROBE_FAILED). broker and draining are reserved component names. /ready is also 503 in job mode once the quota is reached and the consumer cancelled.

The server is bounded: at most 32 connections are served at once (the others are closed at once), reading the request, writing the answer and closing each take at most 5 s, and concurrent /ready requests share a single evaluation of the probes.

Startup. /health answers while the startup hooks run, but only as long as the event loop is free: a hook that loads a model or a large file synchronously blocks it, /health stops answering and the kubelet restarts the pod in a loop. Run heavy loading in asyncio.to_thread, or give the container a startupProbe.

Metrics have no task, client or service label (the scrape adds the pod and service): worker_tasks_total{status}, worker_task_duration_seconds, worker_messages_published_total{message_type}, event_loop_lag_seconds{component}, worker_task_failures_total{code}, worker_task_retries_total, worker_poison_messages_total{reason}, worker_publish_failures_total{message_type}, worker_ack_failures_total, worker_tasks_in_flight, worker_shutdown_requeued_total, worker_progress_dropped_total{reason} and worker_sdk_info, plus the process, platform and GC collectors.

Logs are one JSON object per line on stdout (ts, level, logger, msg, and task_id, service, attempt when bound; ctx.logger carries them, and so do the records of any library logging inside a handler, on the SDK's handler). LOG_LEVEL sets the level. If the application configured the root logger before running the worker, its handlers are kept and the SDK adds none (no line is printed twice): the format is then the application's, and so is the level of the aio_pika.tools logger, whose two ERROR lines on a dropped connection the SDK otherwise logs as warnings (the connection then reconnects by itself). To still get the task fields on your own handler, chain a record factory:

import logging
from async_task_api_worker.observability.logging import TaskContextFilter

task_fields = TaskContextFilter()
previous_factory = logging.getLogRecordFactory()


def factory(*args, **kwargs):
    record = previous_factory(*args, **kwargs)
    task_fields.filter(record)
    return record


logging.setLogRecordFactory(factory)

Every delivery ends with exactly one INFO decision line, for instance Delivery settled: outcome=failure action=conclude code=NOT_ALLOWED attempt=1/5, with the fields task_id, attempt, delivery_limit, outcome (success, failure, poison, retry, publish_failed, shutdown_requeued, cancelled, engine_error), action (ack, conclude, nack-requeue) and code: grep <task_id> tells the whole story of a task. Payloads and secrets are never logged.

Deployment or Job

The same worker code and image run as a long-lived consumer (a Kubernetes Deployment) or as a one-shot process (a Job, for instance one started per message by KEDA's ScaledJob). The deployment picks the mode with WORKER_MODE; the Worker definition does not change.

In job mode the prefetch is forced to 1, whatever WORKER_CONCURRENCY says (it then only sizes the handler semaphore and the thread pool). The quota counts deliveries on arrival: any delivery counts, poison ones and redeliveries past the limit included, and consuming stops as soon as the last one is in hand. The exit code reports the pod, not the task: a task that succeeds, fails (FatalError, timeout, invalid input), or is requeued (RetryableError) all exit 0, and so does an idle exit, because a failed task is a business outcome already published to the client. Exit 1 (an infrastructure failure such as a broker unreachable for WORKER_JOB_CONNECT_TIMEOUT_S, or a def handler stuck after the drain) and 2 (invalid configuration) stay non-zero, so the Job's backoffLimit applies to them. SIGTERM and activeDeadlineSeconds drain exactly as in service mode.

Testing your worker

async_task_api_worker.testing runs your task in memory through the real engine (decoding, time limit, retries, error mapping, acknowledgements), with no broker and no Docker:

from pydantic import BaseModel

from async_task_api_worker import FatalError, TaskContext, Worker
from async_task_api_worker.testing import Outcome, ProgressRecorder, WorkerHarness, make_task_context


class Input(BaseModel):
    text: str


class Output(BaseModel):
    words: int


worker = Worker("count-words")


@worker.task(timeout=30.0)
async def handle(task_input: Input, ctx: TaskContext) -> Output:
    if not task_input.text.strip():
        raise FatalError("EMPTY_TEXT", "There is nothing to count.")
    await ctx.progress(0.5)
    return Output(words=len(task_input.text.split()))


async def test_counts_the_words() -> None:
    async with WorkerHarness(worker) as harness:
        run = await harness.submit({"text": "hello big world"})

    run.assert_succeeded()
    assert run.response == {"words": 3}
    assert run.message_types == ("started", "progress", "success")


async def test_an_empty_text_fails_with_its_code() -> None:
    async with WorkerHarness(worker) as harness:
        run = await harness.submit({"text": "  "})

    run.assert_failed(code="EMPTY_TEXT")
    assert run.outcome is Outcome.ACKED


async def test_the_handler_alone_with_a_context_you_build() -> None:
    recorder = ProgressRecorder()
    ctx = make_task_context(recorder=recorder)

    assert await handle(Input(text="a b"), ctx) == Output(words=2)
    assert [update.fraction for update in recorder.updates] == [0.5]

• harness.submit(body, client_id=..., input_refs=..., delivery_count=...) builds the envelope the API publishes; submit_raw(bytes) sends bytes as they are (a poison message). The result is a TaskRun: messages (decoded) and raw (the published bytes), message_types (all three are tuples), outcome (ACKED, NACKED_REQUEUE), response, error, and the two assertions. task_id is the envelope's, or None for bytes that do not carry one.
• delivery_count is the broker's x-delivery-count: pass 0, then delivery_count=n to see a RetryableError requeue until the last attempt (worker_delivery_limit), which concludes with the error's own code. Retry backoffs are a millisecond; every other setting is the production default (progress throttling included).
• settings_overrides={"worker_task_timeout_s": 0.1, ...} takes WorkerSettings field names; no environment variable is read, and the time budget is checked as at startup. harness.fail_next_publish("success") makes the next publish (of that message type, or of any) fail once, to test the requeue. Use async with (or await harness.aclose()).
• make_task_context and make_sync_task_context build the context of an async def and of a def handler, to call the handler directly; ProgressRecorder collects what it reports.

Development

From this directory:

uv sync --locked --all-extras
uv run ruff check . && uv run ruff format --check .
uv run ty check src
uv run pytest -m "not integration" -q --cov=async_task_api_worker --cov-branch --cov-fail-under=95

The integration tests need Docker (they start RabbitMQ with the platform's policy through testcontainers):

uv run pytest -m integration -q

Add dependencies with uv add from this directory, never by editing pyproject.toml.