From 27d2a3e80c992bc324d659cf56a223fbd18075a8 Mon Sep 17 00:00:00 2001 From: David <20960328+totallynotdavid@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:23:52 +0000 Subject: [PATCH] Migrate compute service to rqueue Align queue ownership and database access so cancelled or crashed simulations can be recovered without stale worker writes. --- .env.example | 6 + ARCHITECTURE.md | 35 +- docker-compose.yml | 7 +- mise.toml | 2 +- packages/api/api/core/db.py | 208 +++-- packages/api/api/core/procrastinate_app.py | 13 - packages/api/api/core/queue.py | 40 + packages/api/api/core/repository.py | 552 +++++++++---- packages/api/api/core/schema.py | 10 +- packages/api/api/core/settings.py | 54 +- packages/api/api/core/tasks.py | 738 ++++++++++++++---- packages/api/api/main.py | 14 +- packages/api/api/migrate.py | 2 +- packages/api/api/queue_migrate.py | 192 ----- packages/api/api/routes.py | 70 +- packages/api/api/worker.py | 91 ++- packages/api/pyproject.toml | 11 +- packages/api/readme.md | 35 +- packages/api/tests/conftest.py | 29 +- packages/api/tests/test_api.py | 155 ++-- packages/api/tests/test_db.py | 217 ++++- packages/api/tests/test_jobs.py | 81 +- .../api/tests/test_migrations_integration.py | 35 +- .../api/tests/test_repository_integration.py | 559 +++++++++---- packages/api/tests/test_tasks.py | 720 ++++++++++------- packages/api/tests/test_worker_integration.py | 679 ++++++++++++++++ pyproject.toml | 1 + scripts/e2e/compute_stack_smoke.sh | 6 +- scripts/e2e/crash_recovery_e2e.sh | 38 +- scripts/integration.sh | 4 +- uv.lock | 95 +-- 31 files changed, 3446 insertions(+), 1253 deletions(-) delete mode 100644 packages/api/api/core/procrastinate_app.py create mode 100644 packages/api/api/core/queue.py delete mode 100644 packages/api/api/queue_migrate.py create mode 100644 packages/api/tests/test_worker_integration.py diff --git a/.env.example b/.env.example index bac4630..e580076 100644 --- a/.env.example +++ b/.env.example @@ -14,3 +14,9 @@ APP_DB_PASSWORD= # Browser-reachable MinIO endpoint for output download URLs. MINIO_PUBLIC_ENDPOINT=localhost:9000 + +# Task queue. The queue name and the schema rqueue owns; both reach the +# migration step and every process that reads or writes the queue, so a value +# set here applies everywhere or nowhere. +COMPUTE_QUEUE=simulations +COMPUTE_QUEUE_SCHEMA=task_queue diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 2c00558..f195b79 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -16,7 +16,7 @@ SvelteKit web app | v FastAPI compute service - |-- compute.jobs and the Procrastinate queue + |-- compute.jobs and the rqueue task_queue schema |-- worker -> tsdhn engine -> MinIO ``` @@ -48,7 +48,7 @@ job ID, a compute-service selector, or output storage keys. The compute service owns: - the API used by the web server; -- `compute.jobs` and the Procrastinate queue tables; +- `compute.jobs` and the `task_queue` schema rqueue owns; - the internal compute job ID; - job progress, retry state, and worker heartbeats; - simulation work directories and checkpoints; @@ -93,8 +93,8 @@ submission_error created_at ``` -The compute service creates and writes `compute.jobs` and the Procrastinate -queue tables. `compute.jobs.simulation_id` links a compute job to the web +The compute service creates and writes `compute.jobs` and the `task_queue` +schema rqueue owns. `compute.jobs.simulation_id` links a compute job to the web simulation. The value is unique because repeating a submission must return the same compute job. @@ -191,3 +191,30 @@ The engine records enough state to continue valid completed work after a retry. The compute service keeps the job in `compute.jobs` and reports the final failure when retries are exhausted. Deployment settings determine retry limits, worker recovery, storage, and cleanup. + +A worker asked to stop gracefully stops taking new work and gives what it is +already running a short grace period. A simulation runs far longer than that, so +it is cancelled and its claim is handed straight back for another worker, which +resumes from the checkpoints in the job's work directory. + +A worker that dies mid-run loses its lease, and the queue reclaims the job +without asking the dead worker anything. When that happens on the job's last +attempt the queue records the failure by itself, so no running code is left to +update `compute.jobs`. The worker process therefore reconciles: it periodically +finds jobs the queue has finished that `compute.jobs` still shows as running, +and marks them failed with an error saying the status was reconciled rather +than reported by the run. A job that reported its own outcome is never +overwritten. + +Because the simulation runs on a thread that the service cannot stop on demand, +a job records which attempt currently owns it, and every write an attempt makes +is accepted only if that attempt still owns the job and the job is not already +finished. A late write from an attempt that has been replaced, or from one whose +job has already been reconciled, is refused rather than applied. + +The same reasoning covers the job's working directory, which holds the +checkpoints a retry resumes from. An attempt holds an exclusive claim on that +directory for as long as its simulation is actually running, so a replacement +never reads checkpoints another attempt is still writing. If the worker process +dies the claim is released with it, and the replacement resumes normally; if the +previous attempt is still running, the replacement waits and retries instead. diff --git a/docker-compose.yml b/docker-compose.yml index 81d2f61..283a04c 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -44,13 +44,14 @@ services: condition: service_healthy environment: COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn + COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue} APP_DB_ROLE: ${APP_DB_ROLE:-tsdhn_app} APP_DB_PASSWORD: ${APP_DB_PASSWORD:?set APP_DB_PASSWORD in .env} command: [ "sh", "-lc", - "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev tsdhn-procrastinate-migrate", + "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev rqueue --database-url \"$$COMPUTE_DATABASE_URL\" --schema \"$$COMPUTE_QUEUE_SCHEMA\" migrate", ] restart: "no" @@ -69,6 +70,8 @@ services: condition: service_healthy environment: COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn + COMPUTE_QUEUE: ${COMPUTE_QUEUE:-simulations} + COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue} MINIO_ENDPOINT: minio:9000 MINIO_PUBLIC_ENDPOINT: ${MINIO_PUBLIC_ENDPOINT:-localhost:9000} MINIO_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minioadmin} @@ -94,6 +97,8 @@ services: condition: service_healthy environment: COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn + COMPUTE_QUEUE: ${COMPUTE_QUEUE:-simulations} + COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue} MINIO_ENDPOINT: minio:9000 MINIO_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minioadmin} MINIO_SECRET_KEY: ${MINIO_SECRET_KEY:-minioadmin} diff --git a/mise.toml b/mise.toml index 71b1efb..ab047db 100644 --- a/mise.toml +++ b/mise.toml @@ -155,7 +155,7 @@ description = "Apply compute, queue, and web migrations to local PostgreSQL" depends = ["db:start"] run = [ "uv run tsdhn-compute-migrate", - "uv run tsdhn-procrastinate-migrate", + "uv run rqueue --database-url \"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" --schema \"${COMPUTE_QUEUE_SCHEMA:-task_queue}\" migrate", "DATABASE_URL=\"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" bun --filter web db:migrate", "uv run tsdhn-web-grants", ] diff --git a/packages/api/api/core/db.py b/packages/api/api/core/db.py index 072cbd0..43807a6 100644 --- a/packages/api/api/core/db.py +++ b/packages/api/api/core/db.py @@ -1,84 +1,188 @@ -"""Database connections for API reads, workers, and job notifications.""" +"""The process-wide asyncpg pool, shared by the API and the worker. + +Both processes open one pool in `open_pool(...)` and take a connection from it +per statement (or per transaction). Nothing holds a connection across a +simulation run: the only long-lived borrow in the system is the one +`rqueue.Worker` makes to LISTEN on its wake channel, which is why +`settings.worker_pool_size()` sizes the worker's pool the way it does. +""" -import threading import uuid -from collections.abc import Iterator -from contextlib import contextmanager +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from typing import Any -import psycopg -from psycopg.rows import dict_row -from psycopg_pool import ConnectionPool +import asyncpg from api.core.errors import TransientInfraError -from api.core.settings import ( - COMPUTE_DATABASE_URL, - DB_POOL_MAX_SIZE, - DB_POOL_MIN_SIZE, -) +from api.core.settings import COMPUTE_DATABASE_URL __all__ = [ "CONNECT_TIMEOUT", "JobRow", + "acquire", "close_pool", "connect", "get_pool", + "is_transient", "notify_channel", - "pooled", + "open_pool", + "transient_connection_errors", ] JobRow = dict[str, Any] CONNECT_TIMEOUT = 2 -_pool: ConnectionPool[psycopg.Connection[JobRow]] | None = None -_pool_lock = threading.Lock() - - -def _new_pool() -> ConnectionPool[psycopg.Connection[JobRow]]: - return ConnectionPool( - COMPUTE_DATABASE_URL, - min_size=DB_POOL_MIN_SIZE, - max_size=DB_POOL_MAX_SIZE, - kwargs={"row_factory": dict_row, "connect_timeout": CONNECT_TIMEOUT}, - open=False, - ) - - -def get_pool() -> ConnectionPool[psycopg.Connection[JobRow]]: - """Return the process-wide pool, creating it on first use.""" +_pool: asyncpg.Pool | None = None + +# Client-side failures: the server was never reached, or the socket died. +_CLIENT_ERRORS = (ConnectionError, OSError, TimeoutError) + +# Server-side failures are classified by SQLSTATE *class* rather than by +# exception type. An enumerated tuple of asyncpg classes is the wrong basis: +# it looked complete and was not (`TooManyConnectionsError` is an +# `InsufficientResourcesError`, not a `PostgresConnectionError`, so pool +# exhaustion read as permanent), and it grows by one entry per incident. +# +# 08 connection_exception the connection broke or was refused +# 53 insufficient_resources too many connections, out of memory, disk full +# 57 operator_intervention admin shutdown, crash shutdown, cannot connect now +# 40 transaction_rollback serialization failure, deadlock detected +# +# The asymmetry justifies the width. Classifying a transient failure as +# permanent throws away a simulation that would have succeeded -- tens of +# minutes of compute, and an operator retry to get it back. Classifying a +# permanent failure as transient costs two extra attempts and roughly 45s of +# backoff before it fails anyway with the same error. So the few permanent +# members inside these classes (57P04 database_dropped, say) are worth carrying +# to catch every transient one. +# +# Deliberately *not* included: 22 (data exception), 23 (integrity constraint), +# 42 (syntax and access rules). Those are application bugs, and a retry there +# re-runs a whole simulation to reach the identical failure. +_TRANSIENT_SQLSTATE_CLASSES = frozenset({"08", "40", "53", "57"}) + + +def is_transient(exc: BaseException) -> bool: + """Whether this failure could plausibly succeed on a later attempt.""" + if isinstance(exc, _CLIENT_ERRORS): + return True + sqlstate = getattr(exc, "sqlstate", None) + return isinstance(sqlstate, str) and sqlstate[:2] in _TRANSIENT_SQLSTATE_CLASSES + + +async def open_pool(*, min_size: int, max_size: int) -> asyncpg.Pool: + """Create the process-wide pool. Idempotent within one process.""" global _pool - with _pool_lock: - if _pool is None or _pool.closed: - _pool = _new_pool() - _pool.open(wait=False) - return _pool + if _pool is None: + _pool = await asyncpg.create_pool( + COMPUTE_DATABASE_URL, + min_size=min_size, + max_size=max_size, + timeout=CONNECT_TIMEOUT, + ) + return _pool -def close_pool() -> None: +async def close_pool() -> None: global _pool - with _pool_lock: - if _pool is not None and not _pool.closed: - _pool.close() + if _pool is not None: + await _pool.close() _pool = None -@contextmanager -def pooled() -> Iterator[psycopg.Connection[JobRow]]: - """Provide a short-lived connection for API queries.""" - with get_pool().connection() as conn: - yield conn +def get_pool() -> asyncpg.Pool: + """Return the open pool, or explain that startup did not open one.""" + if _pool is None: + raise RuntimeError("database pool is not open; call db.open_pool() first") + return _pool + + +@asynccontextmanager +async def transient_connection_errors( + connection: asyncpg.Connection, +) -> AsyncIterator[None]: + """Translate a connection lost mid-statement into `TransientInfraError`. + + `acquire()` covers only the borrow, which is not where a database restart + under a running simulation actually lands. asyncpg reports a connection + that dies around a statement two different ways, and only one of them was + already covered: + + - the backend dies with the statement in flight -> `ConnectionDoesNotExist`, + SQLSTATE 08003, which `is_transient` recognises; + - the backend is already gone when the statement is issued -> a bare + `InterfaceError("connection is closed")`, which is not a connection + error class at all. + + `InterfaceError` is also what asyncpg raises for a genuine programming + error ("the server expects 2 arguments for this query, 1 was passed"), so + it is translated only when the connection really is gone. That distinction + is load bearing: `TRANSIENT_RETRY` retries `TransientInfraError`, so + mistranslating a bug would spend a job's whole retry budget re-running + something that cannot succeed. A programming error leaves the connection + open and propagates unchanged. + """ + try: + yield + except asyncpg.InterfaceError as e: + if not _connection_is_gone(connection): + raise + raise TransientInfraError("database unavailable") from e + except Exception as e: + if not is_transient(e): + raise + raise TransientInfraError("database unavailable") from e + + +def _connection_is_gone(connection: asyncpg.Connection) -> bool: + """Whether this connection can still be used for anything at all. + + Not simply `is_closed()`. When a pooled connection's backend dies, the + pool terminates the connection and takes the proxy back, after which + *every* method on that proxy raises `InterfaceError("cannot call + Connection.is_closed(): connection has been released back to the pool")` -- + the question included. A proxy that cannot answer is certainly gone, + whereas a connection that merely rejected a malformed call answers + `False` and keeps working. + """ + try: + return bool(connection.is_closed()) + except asyncpg.InterfaceError: + return True + + +@asynccontextmanager +async def acquire() -> AsyncIterator[asyncpg.Connection]: + """Borrow a pooled connection for one statement or one transaction.""" + pool = get_pool() + try: + connection = await pool.acquire() + except Exception as e: + # Pool exhaustion arrives here as TooManyConnectionsError (53300), which + # is every bit as retryable as a refused socket. + if not is_transient(e): + raise + raise TransientInfraError("database unavailable") from e + try: + yield connection + finally: + await pool.release(connection) + +async def connect() -> asyncpg.Connection: + """Open a connection outside the pool. -def connect() -> psycopg.Connection[JobRow]: - """Open a dedicated worker connection.""" + The SSE endpoint holds one of these for the lifetime of a stream, which can + be half an hour. Taking that from the request pool would let a handful of + watching browsers starve every other route. + """ try: - return psycopg.connect( - COMPUTE_DATABASE_URL, - connect_timeout=CONNECT_TIMEOUT, - row_factory=dict_row, - ) - except psycopg.OperationalError as e: - raise TransientInfraError("database connection failed") from e + return await asyncpg.connect(COMPUTE_DATABASE_URL, timeout=CONNECT_TIMEOUT) + except Exception as e: + if not is_transient(e): + raise + raise TransientInfraError("database unavailable") from e def notify_channel(simulation_id: uuid.UUID) -> str: diff --git a/packages/api/api/core/procrastinate_app.py b/packages/api/api/core/procrastinate_app.py deleted file mode 100644 index 94171c3..0000000 --- a/packages/api/api/core/procrastinate_app.py +++ /dev/null @@ -1,13 +0,0 @@ -import procrastinate - -from api.core.settings import COMPUTE_DATABASE_URL, PROCRASTINATE_SEARCH_PATH - -__all__ = ["app"] - -app = procrastinate.App( - connector=procrastinate.PsycopgConnector( - conninfo=COMPUTE_DATABASE_URL, - kwargs={"options": f"-c search_path={PROCRASTINATE_SEARCH_PATH}"}, - ), - import_paths=("api.core.tasks",), -) diff --git a/packages/api/api/core/queue.py b/packages/api/api/core/queue.py new file mode 100644 index 0000000..fcafe2e --- /dev/null +++ b/packages/api/api/core/queue.py @@ -0,0 +1,40 @@ +"""The rqueue handle for the simulation queue. + +Three modules import rqueue, and no others: this one for `Queue`, `tasks.py` +for the task and retry types, and `worker.py` for `Worker`. `repository.py` in +particular does not -- it reaches the queue through the `defer` callback +`create_or_get_job` takes, and takes the queue's schema and name as plain +arguments where reconciliation needs them, so the business code stays +independent of which queue is behind it. +""" + +import asyncpg +from rqueue import Queue + +from api.core.settings import COMPUTE_QUEUE, COMPUTE_QUEUE_SCHEMA + +__all__ = ["build_queue", "get_queue", "set_queue"] + +_queue: Queue | None = None + + +def build_queue(pool: asyncpg.Pool) -> Queue: + """Bind the process-wide queue to `pool` and return it. + + Both processes call this: the worker to serve the queue, the API to + enqueue on the connection that also writes `compute.jobs`. + """ + queue = Queue(pool, name=COMPUTE_QUEUE, schema=COMPUTE_QUEUE_SCHEMA) + set_queue(queue) + return queue + + +def set_queue(queue: Queue | None) -> None: + global _queue + _queue = queue + + +def get_queue() -> Queue: + if _queue is None: + raise RuntimeError("simulation queue is not built; call build_queue() first") + return _queue diff --git a/packages/api/api/core/repository.py b/packages/api/api/core/repository.py index fd571f8..4b1100e 100644 --- a/packages/api/api/core/repository.py +++ b/packages/api/api/core/repository.py @@ -1,14 +1,22 @@ -"""Read and write the compute service's job state.""" +"""Read and write the compute service's job state. +Every JSON column is written as `$n::text::jsonb` and read back with an +explicit `::text` cast, decoded here with `json.loads`. asyncpg has no `Jsonb` +wrapper, and whether a given connection carries a jsonb codec depends on who +configured it; casting on both sides makes the behaviour identical either way. +This is the same convention rqueue's own storage layer uses. +""" + +import json import logging import uuid from datetime import datetime from typing import Any, cast -import psycopg -from psycopg.types.json import Jsonb +import anyio +import asyncpg -from api.core.db import JobRow, connect, notify_channel, pooled +from api.core.db import JobRow, acquire, notify_channel, transient_connection_errors from api.core.storage import iso, output_store from tsdhn.domain import EarthquakeInput, JobStatus from tsdhn.engine import SimulationResult @@ -25,6 +33,7 @@ "is_database_connected", "list_abandoned_work_dirs", "mark_started", + "reconcile_terminal_jobs", "record_failure", "record_progress", ] @@ -35,11 +44,38 @@ # These details are safe to show to researchers. FAILED_JOB_DETAILS = "Pipeline failed - check error logs" +RECONCILED_JOB_DETAILS = "Failed - reconciled from the task queue" + +# The queue's error_type is an exception class name ("LeaseExpired", +# "RuntimeError"), the same shape _public_error already exposes, so it is safe +# to show. The message says plainly that no run reported this outcome. +RECONCILED_ERROR = ( + "Simulation stopped without reporting a result; " + "status reconciled from the task queue (%s)" +) + +# Queue states that mean the queue gave up on the job. 'succeeded' is excluded +# deliberately: the queue only records it after run_simulation_task returned, +# which it cannot do before complete_job has committed compute.jobs, so a +# succeeded job that still looks unfinished here is not a state we can reach. +QUEUE_GAVE_UP = ("failed", "cancelled") + +# A job that reported its own outcome is never overwritten, by anyone. +TERMINAL_STATUSES = [JobStatus.COMPLETED.value, JobStatus.FAILED.value] + +# The canonical form `str(uuid.UUID(...))` produces, and the only form +# `enqueue_simulation` ever writes. Anything else is skipped by reconciliation +# rather than casting -- see `_reconcile_sql`. +CANONICAL_UUID_RE = "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$" + +# The jsonb columns come back as text and are decoded below. +_JSON_COLUMNS = ("input_params", "calculation", "travel_times", "outputs") _COLUMNS = """ - id, simulation_id, status, input_params, details, step, step_index, - total_steps, calculation, travel_times, outputs, error, created_at, - updated_at, started_at, finished_at + id, simulation_id, status, input_params::text AS input_params, details, step, + step_index, total_steps, calculation::text AS calculation, + travel_times::text AS travel_times, outputs::text AS outputs, error, + created_at, updated_at, started_at, finished_at """ @@ -48,14 +84,14 @@ def _sql(template: str) -> str: SELECT_BY_SIMULATION_ID = _sql( - "SELECT {columns} FROM compute.jobs WHERE simulation_id = %s" + "SELECT {columns} FROM compute.jobs WHERE simulation_id = $1" ) -SELECT_BY_ID = _sql("SELECT {columns} FROM compute.jobs WHERE id = %s") +SELECT_BY_ID = _sql("SELECT {columns} FROM compute.jobs WHERE id = $1") INSERT_JOB_SQL = _sql( """ INSERT INTO compute.jobs (id, simulation_id, status, input_params, details) - VALUES (%s, %s, %s, %s, %s) + VALUES ($1, $2, $3, $4::text::jsonb, $5) ON CONFLICT (simulation_id) DO NOTHING RETURNING {columns} """ @@ -73,9 +109,21 @@ def _model_dump(data: EarthquakeInput) -> dict[str, Any]: return cast(dict[str, Any], dumped) -def _notify(conn: psycopg.Connection[JobRow], simulation_id: uuid.UUID) -> None: +def _decode(record: asyncpg.Record | None) -> JobRow | None: + """Turn one selected row into a plain mapping with JSON already parsed.""" + if record is None: + return None + row: JobRow = dict(record) + for column in _JSON_COLUMNS: + encoded = row.get(column) + if isinstance(encoded, str): + row[column] = json.loads(encoded) + return row + + +async def _notify(conn: asyncpg.Connection, simulation_id: uuid.UUID) -> None: """PostgreSQL delivers the notification when the update commits.""" - conn.execute(f"NOTIFY {notify_channel(simulation_id)}") + await conn.execute(f"NOTIFY {notify_channel(simulation_id)}") def status_from_row(row: JobRow) -> dict[str, Any]: @@ -109,55 +157,56 @@ def _public_error(e: Exception, step: str | None) -> str: return f"Simulation failed ({type(e).__name__})" -def fetch_by_id(conn: psycopg.Connection[JobRow], job_id: uuid.UUID) -> JobRow | None: - return conn.execute(SELECT_BY_ID, [job_id]).fetchone() +async def fetch_by_id(conn: asyncpg.Connection, job_id: uuid.UUID) -> JobRow | None: + async with transient_connection_errors(conn): + return _decode(await conn.fetchrow(SELECT_BY_ID, job_id)) -def create_or_get_job( +async def create_or_get_job( *, data: EarthquakeInput, simulation_id: str, defer: Any ) -> dict[str, Any]: """Insert a job and enqueue it, atomically. - `defer` receives the open connection and new compute job id. Keeping it as - an argument leaves this module independent of the queue implementation. - The insert and enqueue commit together, so a job row always has a queue - entry. + `defer` is awaited with the open connection and the new compute job id. + Keeping it as an argument leaves this module independent of the queue + implementation. The insert and the enqueue commit together, so a job row + always has a queue entry. """ simulation_uuid = as_uuid(simulation_id) input_params = _model_dump(data) - with pooled() as conn, conn.transaction(): + async with acquire() as conn, conn.transaction(): compute_job_id = uuid.uuid4() - inserted = conn.execute( - INSERT_JOB_SQL, - [ + inserted = _decode( + await conn.fetchrow( + INSERT_JOB_SQL, compute_job_id, simulation_uuid, JobStatus.QUEUED.value, - Jsonb(input_params), + json.dumps(input_params), "Queued for simulation worker", - ], - ).fetchone() + ) + ) if inserted is None: - existing = conn.execute( - SELECT_BY_SIMULATION_ID, [simulation_uuid] - ).fetchone() + existing = _decode( + await conn.fetchrow(SELECT_BY_SIMULATION_ID, simulation_uuid) + ) if existing is None: raise RuntimeError("Compute job was not persisted") if existing["input_params"] != input_params: raise ValueError("Job id already exists with different input") return status_from_row(existing) - defer(conn, compute_job_id) + await defer(conn, compute_job_id) return status_from_row(inserted) -def get_job_status(simulation_id: str) -> dict[str, Any]: +async def get_job_status(simulation_id: str) -> dict[str, Any]: try: simulation_uuid = as_uuid(simulation_id) - with pooled() as conn: - row = conn.execute(SELECT_BY_SIMULATION_ID, [simulation_uuid]).fetchone() + async with acquire() as conn: + row = _decode(await conn.fetchrow(SELECT_BY_SIMULATION_ID, simulation_uuid)) except Exception as e: logger.error("Job lookup failed for %s: %s", sanitize_for_log(simulation_id), e) raise ValueError("Invalid or unknown job ID") from e @@ -167,166 +216,332 @@ def get_job_status(simulation_id: str) -> dict[str, Any]: return status_from_row(row) -def get_outputs(simulation_id: str) -> list[StoredOutput]: +async def get_outputs(simulation_id: str) -> list[StoredOutput]: """Return the outputs recorded with a completed job.""" simulation_uuid = as_uuid(simulation_id) - with pooled() as conn: - row = conn.execute( - "SELECT status, outputs FROM compute.jobs WHERE simulation_id = %s", - [simulation_uuid], - ).fetchone() + async with acquire() as conn: + row = await conn.fetchrow( + """ + SELECT status, outputs::text AS outputs FROM compute.jobs + WHERE simulation_id = $1 + """, + simulation_uuid, + ) if row is None: raise ValueError("Invalid or unknown job ID") if row["status"] != JobStatus.COMPLETED.value: return [] - return cast(list[StoredOutput], row["outputs"] or []) + # `outputs` is NOT NULL DEFAULT '[]' today, so this cannot be NULL. The + # check is here because this query bypasses _decode and reads the column + # as text: if the constraint were ever relaxed, json.loads(None) would + # raise rather than fall through to the `or []` that used to cover it. + encoded = row["outputs"] + if encoded is None: + return [] + return cast(list[StoredOutput], json.loads(encoded) or []) -def get_current_step( - conn: psycopg.Connection[JobRow], job_uuid: uuid.UUID -) -> str | None: - row = conn.execute( - "SELECT step FROM compute.jobs WHERE id = %s", [job_uuid] - ).fetchone() +async def get_current_step(conn: asyncpg.Connection, job_uuid: uuid.UUID) -> str | None: + async with transient_connection_errors(conn): + row = await conn.fetchrow( + "SELECT step FROM compute.jobs WHERE id = $1", job_uuid + ) return str(row["step"]) if row and row["step"] else None -def list_abandoned_work_dirs(cutoff: datetime) -> list[str]: +async def list_abandoned_work_dirs(cutoff: datetime) -> list[str]: """Return failed job workspaces older than the retention cutoff.""" - with pooled() as conn: - rows = conn.execute( + async with acquire() as conn: + rows = await conn.fetch( """ SELECT simulation_id FROM compute.jobs - WHERE status = %s AND finished_at IS NOT NULL AND finished_at < %s + WHERE status = $1 AND finished_at IS NOT NULL AND finished_at < $2 """, - [JobStatus.FAILED.value, cutoff], - ).fetchall() + JobStatus.FAILED.value, + cutoff, + ) return [str(row["simulation_id"]) for row in rows] -def is_database_connected() -> bool: +def _reconcile_sql(queue_schema: str) -> str: + """Build the reconciliation statement for one queue schema. + + `queue_schema` is interpolated rather than bound because PostgreSQL has no + bind parameter for an identifier. The only value ever passed is + `rqueue.Queue.schema`, which rqueue validated against + `[a-z_][a-z0-9_]*` when the queue was constructed, so the interpolated + name is byte-identical to the schema rqueue itself created. + + The `MATERIALIZED` CTE is load bearing, not stylistic. Casting + `payload->>'compute_job_id'` to uuid *aborts the whole statement* on one + malformed value -- `SELECT ('{{"compute_job_id":"x"}}'::jsonb->> + 'compute_job_id')::uuid` raises `invalid input syntax for type uuid` -- + so a single bad row would silently un-reconcile every other stuck job in + the pass. Filtering with a regex in the same `WHERE` as the cast does not + fix that: PostgreSQL is free to evaluate the qualifiers in either order. + Putting the cast in a materialized CTE's select list, behind the regex in + that CTE's `WHERE`, does: the projection is computed only for rows that + passed the filter, and `MATERIALIZED` stops the planner from folding the + two levels back together. A payload this codebase did not write is skipped + rather than fatal. + """ + return f""" + WITH gave_up AS MATERIALIZED ( + SELECT (payload->>'compute_job_id')::uuid AS compute_job_id, + state, error_type, finished_at + FROM {queue_schema}.jobs + WHERE queue = $4 + AND task = $5 + AND state = ANY($6::text[]) + AND finished_at < $7 + AND payload->>'compute_job_id' ~* $9 + ) + UPDATE compute.jobs AS j + SET status = $1, + details = $2, + error = format($3::text, COALESCE(q.error_type, q.state)), + finished_at = COALESCE(j.finished_at, q.finished_at, now()), + updated_at = now() + FROM gave_up AS q + WHERE j.id = q.compute_job_id + AND j.status <> ALL($8::text[]) + RETURNING j.id, j.simulation_id + """ # noqa: S608 + + +async def reconcile_terminal_jobs( + *, queue_schema: str, queue_name: str, task: str, cutoff: datetime +) -> list[uuid.UUID]: + """Fail compute jobs the queue gave up on without the run saying so. + + A simulation normally writes its own outcome: `record_failure` on the last + attempt, `complete_job` on success. Several paths end a job without that + write ever happening -- a worker whose lease expired with its attempt + budget spent, which rqueue fails by a pure SQL update that never calls the + handler; a failure inside `complete_job` after the kernel finished; a + second outage that defeats `record_failure` itself. In every one of them + the queue row is terminal while compute.jobs is still `running`, and + without this pass nothing would ever correct it. + + Only rows the queue gave up on and that compute.jobs has not already + finished are touched, so running this twice, or from two workers at once, + changes nothing the first pass did not: the guard is evaluated under the + row lock the UPDATE takes. + """ + async with acquire() as conn, conn.transaction(): + rows = await conn.fetch( + _reconcile_sql(queue_schema), + JobStatus.FAILED.value, + RECONCILED_JOB_DETAILS, + RECONCILED_ERROR, + queue_name, + task, + list(QUEUE_GAVE_UP), + cutoff, + TERMINAL_STATUSES, + CANONICAL_UUID_RE, + ) + # Same transaction as the update, so a watching SSE stream is only + # woken for a state that has committed. + for row in rows: + await _notify(conn, row["simulation_id"]) + return [row["id"] for row in rows] + + +async def is_database_connected() -> bool: try: - with pooled() as conn: - conn.execute("SELECT 1") + async with acquire() as conn: + await conn.execute("SELECT 1") return True except Exception: return False -def mark_started( - conn: psycopg.Connection[JobRow], job_uuid: uuid.UUID, simulation_id: uuid.UUID -) -> None: - conn.execute( - """ - UPDATE compute.jobs - SET status = %s, details = %s, - started_at = COALESCE(started_at, now()), updated_at = now() - WHERE id = %s - """, - [JobStatus.RUNNING.value, "Simulation worker started", job_uuid], - ) - _notify(conn, simulation_id) - conn.commit() +async def mark_started( + conn: asyncpg.Connection, + job_uuid: uuid.UUID, + simulation_id: uuid.UUID, + attempt: int, +) -> bool: + """Claim this row for `attempt`, and say whether the claim was won. + + Three things happen in one statement, and all three have to: taking + ownership, refusing to reopen a finished job, and marking the job running. + A claim is lost when a *newer* attempt already owns the row, or when the + job is already `completed` -- at-least-once delivery means a completed job + can come back, and re-running it would flip a finished row to `running` for + everyone watching. + + `completed` is refused; `failed` deliberately is not. `rqueue.Admin.retry_job` + is the documented way an operator restarts a terminally failed job, and it + works by putting the same row back to `pending` for an ordinary claim. If a + `failed` compute row could not be reclaimed, that retry would run a whole + simulation whose every write was rejected. + + A reclaimed row also has its previous attempt's `error` and `finished_at` + cleared. That is not cosmetic: `status_from_row` hands both to the status + endpoint and to SSE beside the new `running`, so leaving them makes a live + run indistinguishable from a stale failure to every client. + """ + # asyncpg commits each statement on its own, so the NOTIFY and the UPDATE + # it describes are wrapped together: a client must never be woken to read a + # state that has not committed yet. + async with transient_connection_errors(conn), conn.transaction(): + claimed = await conn.fetchval( + """ + UPDATE compute.jobs + SET status = $1, details = $2, owner_attempt = $3, + error = NULL, finished_at = NULL, + started_at = COALESCE(started_at, now()), updated_at = now() + WHERE id = $4 + AND (owner_attempt IS NULL OR owner_attempt <= $3) + AND status <> $5 + RETURNING id + """, + JobStatus.RUNNING.value, + "Simulation worker started", + attempt, + job_uuid, + JobStatus.COMPLETED.value, + ) + if claimed is None: + return False + await _notify(conn, simulation_id) + return True -def record_progress( - conn: psycopg.Connection[JobRow], +async def record_progress( + conn: asyncpg.Connection, job_uuid: uuid.UUID, simulation_id: uuid.UUID, message: str, details: dict[str, Any], -) -> None: - conn.execute( - """ - UPDATE compute.jobs - SET status = %s, - details = %s, - step = COALESCE(%s, step), - step_index = COALESCE(%s, step_index), - total_steps = COALESCE(%s, total_steps), - calculation = COALESCE(%s, calculation), - travel_times = COALESCE(%s, travel_times), - updated_at = now() - WHERE id = %s - """, - [ + attempt: int, +) -> bool: + """Write one progress update, fenced on owning the attempt. + + This is the only write reachable from the kernel *thread*, which outlives + the coroutine that started it (see `api.core.tasks`). Two predicates make + that safe, and both are part of the same statement as the write: + + - `owner_attempt = $9` refuses a write from an attempt that has been + superseded, so a thread abandoned by attempt 1 cannot scribble over what + attempt 2 is doing; + - `status <> ALL(TERMINAL_STATUSES)` refuses a write to a job that is + already finished, which is the case that matters most: reconciliation + marks an abandoned job `failed`, and a stale write that flipped it back + to `running` would silently undo exactly the repair reconciliation + exists to make. The abandoned attempt is still the *owner* there, so the + attempt fence alone would not catch it. + + Returns whether the write landed. A refused write is normal, not an error. + """ + calculation = details.get("calculation") + travel_times = details.get("travel_times") + async with transient_connection_errors(conn), conn.transaction(): + written = await conn.fetchval( + """ + UPDATE compute.jobs + SET status = $1, + details = $2, + step = COALESCE($3::text, step), + step_index = COALESCE($4::integer, step_index), + total_steps = COALESCE($5::integer, total_steps), + calculation = COALESCE($6::text::jsonb, calculation), + travel_times = COALESCE($7::text::jsonb, travel_times), + updated_at = now() + WHERE id = $8 + AND owner_attempt = $9 + AND status <> ALL($10::text[]) + RETURNING id + """, JobStatus.RUNNING.value, message, details.get("step"), details.get("step_index"), details.get("total_steps"), - Jsonb(details["calculation"]) if "calculation" in details else None, - Jsonb(details["travel_times"]) if "travel_times" in details else None, + json.dumps(calculation) if "calculation" in details else None, + json.dumps(travel_times) if "travel_times" in details else None, job_uuid, - ], - ) - _notify(conn, simulation_id) - conn.commit() + attempt, + TERMINAL_STATUSES, + ) + if written is None: + return False + await _notify(conn, simulation_id) + return True -def record_failure( - conn: psycopg.Connection[JobRow], +async def record_failure( + conn: asyncpg.Connection, job_uuid: uuid.UUID, simulation_id: uuid.UUID, exc: Exception, *, step: str | None, will_retry: bool, + attempt: int, ) -> None: """Record a failed run. When the exception is about to be retried this only updates `details` and leaves the status RUNNING, so an in-flight retry does not look terminal to anyone watching. + + Both statements refuse to touch a row that already reached a terminal + state. That is not defensive padding: `complete_job` can fail *after* its + UPDATE committed (a connection lost at commit is ambiguous by definition), + and the caller then reports a failure for a job that is, in fact, finished. + The guard is the same one `_fail_exhausted` carried on `bump`, and it is + paired with the same `owner_attempt` fence every other write in an + attempt's lifetime carries. """ logger.exception("Simulation failed for compute job %s", job_uuid) - if will_retry: - conn.execute( - "UPDATE compute.jobs SET details = %s, updated_at = now() WHERE id = %s", - [f"Retrying after transient error ({type(exc).__name__})", job_uuid], - ) - else: - conn.execute( - """ - UPDATE compute.jobs - SET status = %s, details = %s, error = %s, - finished_at = now(), updated_at = now() - WHERE id = %s - """, - [ + async with transient_connection_errors(conn), conn.transaction(): + if will_retry: + await conn.execute( + "UPDATE compute.jobs SET details = $1, updated_at = now() " + "WHERE id = $2 AND owner_attempt = $3 " + "AND status <> ALL($4::text[])", + f"Retrying after transient error ({type(exc).__name__})", + job_uuid, + attempt, + TERMINAL_STATUSES, + ) + else: + await conn.execute( + """ + UPDATE compute.jobs + SET status = $1, details = $2, error = $3, + finished_at = now(), updated_at = now() + WHERE id = $4 AND owner_attempt = $5 + AND status <> ALL($6::text[]) + """, JobStatus.FAILED.value, FAILED_JOB_DETAILS, _public_error(exc, step), job_uuid, - ], - ) - _notify(conn, simulation_id) - conn.commit() - - -def fail_job( - conn: psycopg.Connection[JobRow], - job_uuid: uuid.UUID, - simulation_id: uuid.UUID, - error: str, -) -> None: - conn.execute( - """ - UPDATE compute.jobs - SET status = %s, details = %s, error = %s, - finished_at = now(), updated_at = now() - WHERE id = %s - """, - [JobStatus.FAILED.value, FAILED_JOB_DETAILS, error, job_uuid], - ) - _notify(conn, simulation_id) - conn.commit() - - -def complete_job( - conn: psycopg.Connection[JobRow], row: JobRow, result: SimulationResult -) -> None: - """Upload the result and commit its manifest with the terminal state.""" + attempt, + TERMINAL_STATUSES, + ) + await _notify(conn, simulation_id) + + +async def complete_job( + conn: asyncpg.Connection, row: JobRow, result: SimulationResult, attempt: int +) -> bool: + """Upload the result and commit its manifest with the terminal state. + + Fenced on `owner_attempt` like every other write an attempt makes, but + without the terminal-status guard the others carry: a run that genuinely + produced a result should still be able to record it, and only the attempt + that owns the row can reach this at all. + + Returns whether the manifest landed. It matters here more than anywhere + else that the caller can tell: the upload has already happened by the time + the UPDATE runs, so a silently skipped write leaves objects in MinIO with + nothing in `compute.jobs` pointing at them, and no trace of why. + """ now = datetime.now().astimezone() simulation_id = str(row["simulation_id"]) compute_job_id = str(row["id"]) @@ -354,36 +569,49 @@ def complete_job( "travel_times": travel_times, "outputs": outputs, } - bucket, metadata_key = output_store.upload_simulation_result( - simulation_id=simulation_id, - compute_job_id=compute_job_id, - outputs=result.outputs, - metadata=metadata, + # MinIO's client is blocking; uploading on the event loop would stall every + # other coroutine in the worker, heartbeats included. + bucket, metadata_key = await anyio.to_thread.run_sync( + lambda: output_store.upload_simulation_result( + simulation_id=simulation_id, + compute_job_id=compute_job_id, + outputs=result.outputs, + metadata=metadata, + ) ) - conn.execute( - """ - UPDATE compute.jobs - SET status = %s, details = %s, calculation = %s, travel_times = %s, - outputs = %s, result_bucket = %s, result_key = %s, error = NULL, - finished_at = %s, updated_at = now() - WHERE id = %s - """, - [ + async with transient_connection_errors(conn), conn.transaction(): + written = await conn.fetchval( + """ + UPDATE compute.jobs + SET status = $1, details = $2, calculation = $3::text::jsonb, + travel_times = $4::text::jsonb, outputs = $5::text::jsonb, + result_bucket = $6, result_key = $7, error = NULL, + finished_at = $8, updated_at = now() + WHERE id = $9 AND owner_attempt = $10 + RETURNING id + """, JobStatus.COMPLETED.value, "Simulation completed successfully", - Jsonb(calculation), - Jsonb(travel_times), - Jsonb(outputs), + json.dumps(calculation), + json.dumps(travel_times), + json.dumps(outputs), bucket, metadata_key, now, row["id"], - ], - ) - _notify(conn, row["simulation_id"]) - conn.commit() - - -def open_worker_connection() -> psycopg.Connection[JobRow]: - return connect() + attempt, + ) + if written is None: + logger.error( + "Compute job %s no longer belongs to attempt %d; its result was " + "uploaded to %s/%s but not recorded, and nothing in compute.jobs " + "now points at it", + compute_job_id, + attempt, + bucket, + metadata_key, + ) + return False + await _notify(conn, row["simulation_id"]) + return True diff --git a/packages/api/api/core/schema.py b/packages/api/api/core/schema.py index f793576..7ebf10a 100644 --- a/packages/api/api/core/schema.py +++ b/packages/api/api/core/schema.py @@ -23,9 +23,17 @@ created_at timestamptz NOT NULL DEFAULT now(), updated_at timestamptz NOT NULL DEFAULT now(), started_at timestamptz, - finished_at timestamptz + finished_at timestamptz, + -- Which task-queue attempt currently owns this row. Every write an attempt + -- makes during its run carries its own number and matches on this column in + -- the same statement, so a write from an attempt that has been superseded + -- matches zero rows instead of racing the attempt that replaced it. NULL + -- until the first attempt claims the row. + owner_attempt integer ); +ALTER TABLE compute.jobs ADD COLUMN IF NOT EXISTS owner_attempt integer; + DO $$ BEGIN IF EXISTS ( diff --git a/packages/api/api/core/settings.py b/packages/api/api/core/settings.py index 8e2b991..37ab65e 100644 --- a/packages/api/api/core/settings.py +++ b/packages/api/api/core/settings.py @@ -5,6 +5,8 @@ "APP_DB_PASSWORD", "APP_DB_ROLE", "COMPUTE_DATABASE_URL", + "COMPUTE_QUEUE", + "COMPUTE_QUEUE_SCHEMA", "DB_POOL_MAX_SIZE", "DB_POOL_MIN_SIZE", "JOBS_DIR", @@ -17,10 +19,12 @@ "MINIO_SECURE", "NUMBA_THREADS", "OUTPUT_URL_TTL", - "PROCRASTINATE_QUEUE", - "PROCRASTINATE_SCHEMA", - "PROCRASTINATE_SEARCH_PATH", "SSE_MAX_DURATION", + "WORKER_CONCURRENCY", + "WORKER_ID", + "WORKER_LEASE_SECONDS", + "api_pool_size", + "worker_pool_size", ] COMPUTE_DATABASE_URL = os.environ.get( @@ -28,18 +32,52 @@ "postgresql://tsdhn:tsdhn@localhost:5432/tsdhn", ) -PROCRASTINATE_QUEUE = os.environ.get("PROCRASTINATE_QUEUE", "simulations") -PROCRASTINATE_SCHEMA = "compute" -PROCRASTINATE_SEARCH_PATH = f"{PROCRASTINATE_SCHEMA},public" +COMPUTE_QUEUE = os.environ.get("COMPUTE_QUEUE", "simulations") + +# rqueue owns its own schema and never shares one with application tables, so +# this must not be `compute`: that schema is compute.jobs' own. +COMPUTE_QUEUE_SCHEMA = os.environ.get("COMPUTE_QUEUE_SCHEMA", "task_queue") JOBS_DIR: Path = Path(os.environ.get("TSDHN_JOBS_DIR", "jobs")).resolve() LOG_LEVEL = os.environ.get("TSDHN_LOG_LEVEL", "INFO").upper() -# API reads use a pool; worker runs use dedicated connections. -DB_POOL_MIN_SIZE = int(os.environ.get("DB_POOL_MIN_SIZE", "1")) +# Zero keeps the pool lazy, so a process starts before PostgreSQL is +# reachable and reports the outage through /health instead of refusing to +# boot -- the behaviour the psycopg pool had with open(wait=False). +DB_POOL_MIN_SIZE = int(os.environ.get("DB_POOL_MIN_SIZE", "0")) DB_POOL_MAX_SIZE = int(os.environ.get("DB_POOL_MAX_SIZE", "10")) +# One simulation saturates the machine's cores, so the worker runs one job at a +# time unless a deployment says otherwise. +WORKER_CONCURRENCY = int(os.environ.get("TSDHN_WORKER_CONCURRENCY", "1")) + +# A lease this long tolerates a slow database round trip; rqueue heartbeats at +# a third of it while the simulation runs off the event loop. +WORKER_LEASE_SECONDS = float(os.environ.get("TSDHN_WORKER_LEASE_SECONDS", "60")) + +WORKER_ID = os.environ.get("TSDHN_WORKER_ID", "") + + +def api_pool_size() -> tuple[int, int]: + """Return the API process's pool bounds.""" + return DB_POOL_MIN_SIZE, max(DB_POOL_MIN_SIZE, DB_POOL_MAX_SIZE) + + +def worker_pool_size() -> tuple[int, int]: + """Return the worker process's pool bounds. + + This floor is load bearing, not cosmetic. `rqueue.Worker` holds one + connection for the whole run to LISTEN on its wake channel, borrows one per + poll to recover expired leases and claim, one per heartbeat, and one per + progress write from a running simulation. When the pool cannot spare a + connection for the listener, rqueue logs a warning and falls back to + polling only, which is a latency regression that is easy to misdiagnose. + """ + max_size = max(DB_POOL_MAX_SIZE, 2 * WORKER_CONCURRENCY + 4) + return min(DB_POOL_MIN_SIZE, max_size), max_size + + APP_DB_ROLE = os.environ.get("APP_DB_ROLE", "tsdhn_app") APP_DB_PASSWORD = os.environ.get("APP_DB_PASSWORD", "") diff --git a/packages/api/api/core/tasks.py b/packages/api/api/core/tasks.py index 8e6d11d..d44a2ef 100644 --- a/packages/api/api/core/tasks.py +++ b/packages/api/api/core/tasks.py @@ -1,194 +1,640 @@ -"""Queue tasks for simulation runs and worker maintenance.""" - +"""Queue tasks for simulation runs and worker maintenance. + +Abandoned attempts, and how their writes are fenced +--------------------------------------------------- +The simulation kernel is synchronous and runs on a thread via +`asyncio.to_thread`. rqueue cancels a handler's *coroutine* when the heartbeat +discovers the lease is gone (`Worker._heartbeat_loop` -> `handler_task.cancel()`), +but Python cannot kill the thread underneath it. That thread keeps running and +holds a live reference to the event loop through `on_progress`, so it can still +try to write `compute.jobs` for a job another worker has taken over -- or for a +job `reconcile_terminal_jobs` has already, correctly, marked failed. + +Procrastinate did not have this exposure: its task was synchronous, with no +background thread able to outlive a cancelled coroutine. It is new here. + +The fence is in the database, not in this process. `mark_started` claims the row +for `context.attempt` by writing `compute.jobs.owner_attempt`, and every write +an attempt makes afterwards carries its own attempt number in the *same* +statement as the write (`WHERE id = $1 AND owner_attempt = $n ...`). A write +from a superseded attempt therefore matches zero rows. `record_progress` adds +`AND status <> ALL(TERMINAL_STATUSES)`, because the abandoned attempt is often +still the row's owner -- the case that matters is a stale write flipping a +reconciled `failed` job back to `running`, undoing the repair reconciliation +just made, and only the status predicate catches that. + +A refused write is not an error; it is the fence working. `on_progress` treats +one as the signal to stop, raising `AbandonedAttempt` so the kernel unwinds and +the thread ends rather than burning a core on work nobody will read. + +That exception surfaces in two different places, and the difference matters: + +- **Refused write.** The coroutine is still live and awaiting the thread, so + `AbandonedAttempt` propagates into `run_simulation_task` as a real exception. + It is caught there and re-raised as `rqueue.CancelJob`, because standing down + is not a failure: letting it reach rqueue's generic exception path would + record a `rqueue.job.failed`, finalize the job as failed, and log two + tracebacks for a fence working correctly -- and an on-call engineer would + have no way to tell that from a simulation that actually broke. +- **Cancellation.** rqueue has already cancelled the coroutine, so there is + nothing left to propagate into and the thread just ends. The same `abandoned` + flag is set directly there, which usually stops the thread a step *earlier* + than a refused write would. + +The workspace is fenced separately, and has to be +------------------------------------------------- +The database fence protects `compute.jobs`. It does nothing for the *files*: +the workspace is keyed by simulation id, not by attempt, and an abandoned +thread mid-step keeps writing checkpoints into it. A replacement attempt that +resumed from a half-written checkpoint would produce a wrong scientific result +rather than an error, which is worse than any status confusion. + +`claim_workspace` therefore takes an exclusive `flock` on a lock file beside the +workspace. A second attempt is refused while a live owner holds it -- including +from another thread of the same process -- and gets `TransientInfraError`, so +rqueue's backoff waits for the abandoned thread to unwind rather than racing it. +When the holding *process* dies the kernel drops the lock, so an ordinary worker +crash still leaves the workspace resumable. + +The claim spans more than the kernel. `complete_job` reads the result files back +out of the same directory to upload them, after the kernel thread has returned, +so a claim that ended with the kernel would leave that upload unprotected -- +the identical corruption a few lines later. The claim therefore carries two +shares, one given back by the kernel thread and one by this coroutine after the +upload, and the descriptor closes only when both have. Neither holder can be +cancelled, and a cancelled coroutine ends while the thread is still writing, so +"whoever finishes last closes it" is the only rule that works. + +`remove_workspace` takes the lock before it removes anything, for the same +reason in reverse: `unlink` on a locked file drops the directory entry while the +holder keeps its lock on the orphaned inode, so the next `O_CREAT` at that path +creates a *new* inode that locks uncontended. It refuses rather than forces, and +the next sweep retries. + +Two limits are worth knowing. A thread that never reaches another progress +callback holds the workspace until it finishes on its own, and the replacement +can exhaust its attempts waiting -- a visible, bounded failure rather than +silent corruption, which is the trade being made. And `flock` is a +single-machine primitive: two workers on different hosts sharing one network +volume are not protected by it. Neither the compose deployment (one worker, a +local volume) nor any current configuration is exposed to the second. +""" + +import asyncio +import contextlib +import fcntl import logging +import os import shutil +import threading import uuid +from collections.abc import Awaitable, Callable from datetime import datetime, timedelta -from typing import Any, Literal +from pathlib import Path +from typing import Any -import psycopg -from procrastinate import JobContext, RetryStrategy -from procrastinate import exceptions as procrastinate_exceptions -from procrastinate.jobs import Job as ProcrastinateJob -from procrastinate.jobs import Status as ProcrastinateStatus +import asyncpg +from rqueue import CancelJob, Job, PermanentFailure, Queue, RetryPolicy, TaskContext -from api.core import repository -from api.core.db import JobRow, connect, pooled +from api.core import db, repository from api.core.errors import TransientInfraError -from api.core.procrastinate_app import app -from api.core.settings import JOBS_DIR, PROCRASTINATE_QUEUE +from api.core.queue import get_queue +from api.core.settings import JOBS_DIR from tsdhn.domain import EarthquakeInput, JobStatus from tsdhn.engine import run_simulation -from tsdhn.utils.file_utils import sanitize_for_log __all__ = [ + "MAX_ATTEMPTS", + "RUN_SIMULATION", + "TRANSIENT_RETRY", + "AbandonedAttempt", + "decode_payload", "enqueue_simulation", - "reap_action", - "reap_stalled_jobs_task", + "reconcile_terminal_jobs", + "register_tasks", + "run_periodic_reconcile", + "run_periodic_sweep", "run_simulation_task", - "sweep_abandoned_work_dirs_task", + "sweep_abandoned_work_dirs", ] -type ReapAction = Literal["retry", "exhausted"] - logger = logging.getLogger(__name__) +RUN_SIMULATION = "api.run_simulation" + + +class AbandonedAttempt(Exception): + """Raised into the kernel thread when its attempt no longer owns the job. + + Reaches `run_simulation_task` only on the refused-write path, where the + coroutine is still live and awaiting the thread. It is caught there and + turned into `rqueue.CancelJob`, so standing down never counts as a task + failure. On the cancellation path the coroutine is already gone, so nothing + propagates it and the thread simply ends. + """ + + # Retry only infrastructure failures. Domain and pipeline errors are terminal. +# `retry_on` is rqueue's exception allowlist, the direct counterpart of +# Procrastinate's `retry_exceptions`. MAX_ATTEMPTS = 3 -TRANSIENT_RETRY = RetryStrategy( +TRANSIENT_RETRY = RetryPolicy( max_attempts=MAX_ATTEMPTS, - exponential_wait=15, - retry_exceptions=(TransientInfraError,), + initial_backoff=15.0, + multiplier=2.0, + retry_on=(TransientInfraError,), ) -# This is a heartbeat timeout, not a simulation runtime limit. -STALLED_HEARTBEAT_SECONDS = 90 +# A simulation runs for tens of minutes and has no meaningful upper bound, so +# it is left untimed; rqueue's lease recovery, not a timeout, is what reclaims +# a run whose worker died. +RUN_TIMEOUT_SECONDS: float | None = None -# Delay requeued work so several jobs do not start at once after a worker crash. -CRASH_REQUEUE_DELAY_SECONDS = 30 - -CRASH_BUDGET_EXHAUSTED_ERROR = ( - "Simulation worker stopped responding mid-run (retry budget exhausted)" -) +# The workspace ownership lock, kept beside the workspace rather than inside +# it: `prepare_simulation_workspace` removes the whole directory when a run +# starts without resuming, which would take the lock with it. +WORKSPACE_LOCK_SUFFIX = ".lock" # Keep failed workspaces for local inspection and manual recovery. WORK_DIR_TTL = timedelta(hours=24) +SWEEP_INTERVAL_SECONDS = 3600.0 + +# How long a queue job must have been terminal before reconciliation claims it. +# The queue row being terminal already implies no worker holds its lease, so in +# principle nothing else can be writing compute.jobs by then. The grace covers +# the one case where that reasoning is not airtight: a worker wedged long +# enough to lose its lease, whose recovery -- and whose own write of the +# outcome -- can still land after rqueue has failed the row underneath it. +# Five minutes is comfortably longer than any such write, and the delay costs +# nothing: this only ever runs on jobs that are already over. +RECONCILE_GRACE = timedelta(minutes=5) + +# A minute, against the sweep's hour, because the two clean up different +# things. The sweep reclaims disk from jobs that are already reported failed, +# where an hour of delay is invisible. This pass is what ends a *user-visible* +# stuck status: until it runs, a researcher watching /events sees "running" +# for a simulation that no longer exists anywhere. The whole pass is one +# indexed statement over jobs the queue has already finished, so polling it +# this often is cheap; the stuck window is then bounded by the lease duration +# plus the grace, not by the interval. +RECONCILE_INTERVAL_SECONDS = 60.0 + + +def decode_payload(payload: Any) -> uuid.UUID: + """Turn a queued payload into the compute job id it names. + + A payload that cannot be decoded fails the job durably without spending + the retry budget: the same bytes will not decode on a later attempt. + """ + if not isinstance(payload, dict): + raise ValueError("simulation payload must be a JSON object") + return repository.as_uuid(str(payload["compute_job_id"])) + + +async def enqueue_simulation( + connection: asyncpg.Connection, compute_job_id: uuid.UUID +) -> Job: + """Queue the task on `connection` so it commits with the job row. + + `dedupe_key` replaces Procrastinate's `queueing_lock` (one queued job per + compute job) and `concurrency_key` replaces its `lock` (one simulation + running at a time for a compute job). rqueue keeps those two as separate + primitives, so both survive the swap unchanged in meaning. + """ + return await get_queue().enqueue( + connection, + task=RUN_SIMULATION, + payload={"compute_job_id": str(compute_job_id)}, + dedupe_key=f"simulation:{compute_job_id}", + on_conflict="return_existing", + concurrency_key=f"compute-job:{compute_job_id}", + ) -def enqueue_simulation( - conn: psycopg.Connection[JobRow], compute_job_id: uuid.UUID -) -> None: - """Queue the task on `conn` so it commits with the job row.""" - run_simulation_task.configure( - connection=conn, - queue=PROCRASTINATE_QUEUE, - queueing_lock=f"simulation:{compute_job_id}", - lock=f"compute-job:{compute_job_id}", - ).defer(compute_job_id=str(compute_job_id)) - - -@app.task( - name="api.run_simulation", - queue=PROCRASTINATE_QUEUE, - pass_context=True, - retry=TRANSIENT_RETRY, -) -def run_simulation_task(context: JobContext, compute_job_id: str) -> None: - """Run one simulation end to end, streaming progress into compute.jobs.""" - job_uuid = repository.as_uuid(compute_job_id) - - # A simulation keeps its connection for the duration of the run. - with connect() as conn: - row = repository.fetch_by_id(conn, job_uuid) - if row is None: - raise RuntimeError(f"Unknown compute job {compute_job_id}") - - simulation_id: uuid.UUID = row["simulation_id"] - work_dir = JOBS_DIR / str(simulation_id) - data = EarthquakeInput(**row["input_params"]) +class WorkspaceClaim: + """One exclusive `flock` on a workspace, released by whoever finishes last. + + Two things write to, or read out of, the workspace during an attempt, and + neither can be cancelled once started: the kernel thread, and the upload + `complete_job` makes on another thread. They do not overlap, but they can + *end* in either order -- a cancelled coroutine ends while the kernel thread + is still going -- so a single owner cannot close the descriptor at the right + moment. Each side takes a share and gives it back at its own end; the + descriptor closes, and the lock with it, only when both have. + """ + + def __init__(self, fd: int, shares: int = 2) -> None: + self._fd = fd + self._shares = shares + self._guard = threading.Lock() + + def release(self) -> None: + """Give back one share. Closing is the last releaser's job.""" + with self._guard: + self._shares -= 1 + if self._shares > 0: + return + fd, self._fd = self._fd, -1 + if fd >= 0: + os.close(fd) + + +def _lock_path(work_dir: Path) -> Path: + return work_dir.with_name(work_dir.name + WORKSPACE_LOCK_SUFFIX) + + +def claim_workspace(work_dir: Path, attempt: int) -> tuple[WorkspaceClaim, bool]: + """Take this attempt's exclusive claim on `work_dir`, and say whether to resume. + + The database fence stops a superseded attempt writing to `compute.jobs`. It + does nothing about the *filesystem*: the workspace is keyed by simulation + id, not by attempt, and a kernel thread abandoned mid-step keeps writing + into it. A replacement attempt resuming from those checkpoints could read a + half-written one -- which is worse than a confusing status, because a + corrupted checkpoint yields a wrong scientific result rather than an error. + + `flock` is the right primitive because of how it is released. It is held + against the open file description, so a second attempt is refused even from + another thread of the same process (which is exactly the zombie case, and + what a concurrency above one makes possible) -- and the kernel drops it when + the holding **process** dies, SIGKILL included. So an ordinary worker crash + still leaves the workspace resumable, which is the behaviour + `crash_recovery_e2e.sh` scenario 1 depends on, while a live zombie does not. + + The returned claim carries two shares, because the claim has to outlast the + kernel: `complete_job` reads the result files out of this same directory + after the kernel thread has returned, and a lease lost in between would + otherwise let a redelivered attempt start writing there mid-upload. + """ + lock_path = _lock_path(work_dir) + lock_path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as e: + os.close(fd) + # Transient on purpose: the previous attempt's thread unwinds at its + # next progress write, which the database fence refuses, so the + # workspace frees itself. rqueue's backoff is the wait. + raise TransientInfraError( + f"simulation workspace {work_dir.name} is still held by an earlier attempt" + ) from e + os.ftruncate(fd, 0) + os.write(fd, f"attempt {attempt}\n".encode()) + # Evaluated under the lock: without it, "are there checkpoints to resume + # from" is a question about a directory someone else may be halfway + # through writing. + return WorkspaceClaim(fd), work_dir.exists() + + +def remove_workspace(work_dir: Path) -> bool: + """Remove a workspace and its lock, if nothing still holds the lock. + + Taking the lock before unlinking it is what makes removal safe, and the + ordering is not fussiness. `unlink` on a locked file is the classic `flock` + footgun: it removes the directory entry while the holder keeps its lock on + the now-orphaned inode, so the next `O_CREAT` at that path makes a *new* + inode and locks it uncontended -- leaving two threads each believing they + own the workspace exclusively. The sweep can reach a job whose row is + terminal while its kernel thread is still alive, so this is reachable. + + Returns whether the workspace was removed. A refusal is not an error: the + next sweep tries again, and until then the directory is exactly where its + owner expects it. + """ + lock_path = _lock_path(work_dir) + try: + fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) + except OSError: + # No lock file and none creatable -- nothing has ever claimed this. + shutil.rmtree(work_dir, ignore_errors=True) + return True + try: + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + logger.info( + "Workspace %s is still held; leaving it for the next sweep", + work_dir.name, + ) + return False + shutil.rmtree(work_dir, ignore_errors=True) + with contextlib.suppress(OSError): + lock_path.unlink() + finally: + os.close(fd) + return True - repository.mark_started(conn, job_uuid, simulation_id) - def on_progress(message: str, details: dict[str, Any]) -> None: - repository.record_progress(conn, job_uuid, simulation_id, message, details) +async def run_simulation_task(compute_job_id: uuid.UUID, context: TaskContext) -> None: + """Run one simulation end to end, streaming progress into compute.jobs.""" + async with db.acquire() as conn: + row = await repository.fetch_by_id(conn, compute_job_id) + if row is None: + # There is no attempt at which this job could succeed. + raise PermanentFailure(f"Unknown compute job {compute_job_id}") + + simulation_id: uuid.UUID = row["simulation_id"] + work_dir = JOBS_DIR / str(simulation_id) + data = EarthquakeInput(**row["input_params"]) + + # The claim is the decision. Reading `row["status"]` above and branching on + # it would be the same check-then-write gap the owner_attempt fence exists + # to close: an attempt completing between that SELECT and this UPDATE could + # still have its `completed` overwritten with `running`. + async with db.acquire() as conn: + claimed = await repository.mark_started( + conn, compute_job_id, simulation_id, context.attempt + ) + if not claimed: + await _stand_down(compute_job_id, work_dir, context.attempt) + return - try: - result = run_simulation( - data, - work_dir, - # Keep outputs and checkpoints when a retry has a work directory. - resume=work_dir.exists(), - on_progress=on_progress, + # Captured before the hop into the thread: the callback below runs on a + # worker thread and has to get back to the loop that owns the pool. + loop = asyncio.get_running_loop() + # Set when this attempt's coroutine is cancelled out from under the thread + # still running the kernel. See "Abandoned attempts" in the module docstring. + abandoned = threading.Event() + + def on_progress(message: str, details: dict[str, Any]) -> None: + if abandoned.is_set(): + # Raising, rather than returning quietly, is the point: it unwinds + # `run_simulation` and ends the thread at this step boundary + # instead of letting it burn a core on work nobody will read. + raise AbandonedAttempt( + f"attempt {context.attempt} of compute job {compute_job_id} " + "no longer owns this job" ) - repository.complete_job(conn, row, result) - except Exception as e: - will_retry = ( - isinstance(e, TransientInfraError) - and context.job.attempts < MAX_ATTEMPTS + # Blocking on the result keeps the old synchronous contract: progress + # is durable before the simulation moves on, and a write that fails + # surfaces in the simulation rather than in a dropped background task. + written = asyncio.run_coroutine_threadsafe( + _write_progress( + compute_job_id, simulation_id, message, details, context.attempt + ), + loop, + ).result() + if not written: + # The database refused it: this attempt no longer owns the row, or + # the row is already finished. Stop the same way as above -- the + # write was already fenced, so nothing was corrupted; there is just + # no reason to keep running. + abandoned.set() + raise AbandonedAttempt( + f"attempt {context.attempt} of compute job {compute_job_id} " + "no longer owns this job" ) - repository.record_failure( - conn, - job_uuid, - simulation_id, - e, - step=repository.get_current_step(conn, job_uuid), - will_retry=will_retry, + + def run_kernel(claim: WorkspaceClaim, resume: bool) -> Any: + try: + # Keep outputs and checkpoints when a retry has a work directory. + return run_simulation( + data, work_dir, resume=resume, on_progress=on_progress ) - raise - else: - shutil.rmtree(work_dir, ignore_errors=True) + finally: + # On this thread, so the claim is given back when the *thread* + # stops writing -- which is not when the coroutine stops waiting. + claim.release() + + claim: WorkspaceClaim | None = None + try: + # rqueue installs a ThreadPoolExecutor sized from the worker's + # concurrency as the loop's default executor, so these hops are capacity + # bounded without building an executor here. + held, resume = await asyncio.to_thread( + claim_workspace, work_dir, context.attempt + ) + claim = held + result = await asyncio.to_thread(run_kernel, held, resume) + async with db.acquire() as conn: + # Still under this attempt's claim: complete_job reads the result + # files back out of the workspace to upload them, and a lease lost + # in the moment between the kernel returning and the upload + # finishing would otherwise let a redelivered attempt take the + # freed lock and start writing into the directory being read. + recorded = await repository.complete_job(conn, row, result, context.attempt) + except asyncio.CancelledError: + # rqueue cancels this coroutine when the heartbeat finds the lease gone. + # The coroutine ends here; the OS thread underneath it does not, because + # Python cannot kill a thread. Tell it to stop. + abandoned.set() + raise + except AbandonedAttempt as e: + # A refused write, not a broken run. This is the one path where the + # coroutine is still live when the guard fires, so the exception really + # does arrive here -- and letting it fall through to `except Exception` + # would spend a `rqueue.job.failed`, a `_fail_terminal`, and two + # exception tracebacks on the fence doing exactly its job. rqueue's + # failure accounting has to mean genuine failures, or it means nothing. + # + # `CancelJob` is the honest signal: stop, do not retry. Reaching here + # means a newer attempt owns the row or the job is already terminal, + # both of which imply this attempt's lease is gone, so rqueue's + # finalization will hit `LeaseLost` and log that instead. If the lease + # somehow survived, the job becomes `cancelled` -- which reconciliation + # already treats as the queue giving up, so `compute.jobs` is repaired + # rather than left running forever. + logger.warning("Standing down attempt %d: %s", context.attempt, e) + raise CancelJob(str(e)) from e + except Exception as e: + # Finalization is as retryable as the run: a MinIO outage raises + # TransientInfraError from complete_job, and without this the retry + # happened but nobody watching compute.jobs was ever told. + # record_failure refuses to touch a row that is already terminal, so an + # UPDATE that committed before the connection dropped is not reported + # as a failure. + await _record_failure(compute_job_id, simulation_id, e, context.attempt) + raise + finally: + if claim is not None: + claim.release() + + if not recorded: + # Superseded between the last progress write and here. `complete_job` + # has already logged the orphaned upload. Stand down the same way a + # refused progress write does -- and in particular do not remove the + # work directory, which the attempt that took over is resuming from. + raise CancelJob( + f"attempt {context.attempt} of compute job {compute_job_id} " + "no longer owns this job" + ) + # The claim is given back above, and remove_workspace re-takes it itself: + # unlinking a lock file somebody still holds is what creates two owners. + await asyncio.to_thread(remove_workspace, work_dir) + + +async def _stand_down(compute_job_id: uuid.UUID, work_dir: Path, attempt: int) -> None: + """Handle a claim this attempt did not win. + + The claim already told us we lost; this only decides which of the two ways + to lose it was, so a stale read here is harmless. + """ + async with db.acquire() as conn: + current = await repository.fetch_by_id(conn, compute_job_id) + + if current is not None and current["status"] == JobStatus.COMPLETED.value: + # rqueue is at-least-once by design (REQUIREMENTS.md ss3): a worker that + # died between complete_job's commit and rqueue's own finalization gets + # this job delivered again. The result is already durable, so let this + # delivery succeed, and finish the one step the dead attempt may not + # have reached. + logger.info( + "Compute job %s is already completed; skipping redelivered attempt %d", + compute_job_id, + attempt, + ) + await asyncio.to_thread(remove_workspace, work_dir) + return -def reap_action(job: ProcrastinateJob) -> ReapAction: - return "retry" if job.attempts < MAX_ATTEMPTS else "exhausted" + logger.warning( + "Compute job %s is owned by a newer attempt than %d; standing down", + compute_job_id, + attempt, + ) + raise CancelJob( + f"attempt {attempt} of compute job {compute_job_id} no longer owns this job" + ) -@app.periodic(cron="*/2 * * * *") -@app.task(name="api.reap_stalled_jobs", queue=PROCRASTINATE_QUEUE) -async def reap_stalled_jobs_task(timestamp: int) -> None: - stalled = list( - await app.job_manager.get_stalled_jobs( - seconds_since_heartbeat=STALLED_HEARTBEAT_SECONDS +async def _write_progress( + compute_job_id: uuid.UUID, + simulation_id: uuid.UUID, + message: str, + details: dict[str, Any], + attempt: int, +) -> bool: + async with db.acquire() as conn: + return await repository.record_progress( + conn, compute_job_id, simulation_id, message, details, attempt ) - ) - if not stalled: - return - retry_at = datetime.now().astimezone() + timedelta( - seconds=CRASH_REQUEUE_DELAY_SECONDS - ) - for job in stalled: - if job.id is None: # pragma: no cover - persisted jobs always have an id - continue - compute_job_id = str(job.task_kwargs.get("compute_job_id", "")) - if not compute_job_id: - continue - - if reap_action(job) == "retry": - logger.warning( - "Requeuing stalled job %s (compute job %s): heartbeat went stale", - job.id, - sanitize_for_log(compute_job_id), +async def _record_failure( + compute_job_id: uuid.UUID, + simulation_id: uuid.UUID, + exc: Exception, + attempt: int, +) -> None: + """Persist a failed run, matching the decision rqueue is about to make. + + `should_retry` is the same predicate `rqueue.Worker` consults, so the + `details` a watching client sees ("retrying" versus terminal) cannot drift + from what the queue actually does with the job. + """ + try: + async with db.acquire() as conn: + await repository.record_failure( + conn, + compute_job_id, + simulation_id, + exc, + step=await repository.get_current_step(conn, compute_job_id), + will_retry=TRANSIENT_RETRY.should_retry(exc, attempt=attempt), + attempt=attempt, ) - try: - await app.job_manager.retry_job_by_id_async( - job_id=job.id, retry_at=retry_at - ) - except procrastinate_exceptions.ConnectorException: - logger.info("Stalled job %s resolved before requeue", job.id) - continue + except Exception: + # A database outage here must not replace the failure it is reporting. + logger.exception("Could not record the failure of job %s", compute_job_id) + +async def sweep_abandoned_work_dirs() -> None: + """Delete the workspaces of jobs that failed longer than WORK_DIR_TTL ago.""" + cutoff = datetime.now().astimezone() - WORK_DIR_TTL + for simulation_id in await repository.list_abandoned_work_dirs(cutoff): + await asyncio.to_thread(remove_workspace, JOBS_DIR / simulation_id) + + +async def reconcile_terminal_jobs(*, grace: timedelta = RECONCILE_GRACE) -> None: + """Sync compute.jobs to jobs the queue finished without the run reporting. + + Procrastinate's reaper used to close this gap for one of its causes: a + stalled job whose retries ran out. rqueue's fenced leases replace the + queue-state half of that reaper outright and do it better, but the half + that wrote `compute.jobs` had no replacement, and the same hole is reachable + from more than stalling alone. This is that replacement, written against + the condition rather than against any one of the paths that reaches it. + """ + queue = get_queue() + reconciled = await repository.reconcile_terminal_jobs( + queue_schema=queue.schema, + queue_name=queue.name, + task=RUN_SIMULATION, + cutoff=datetime.now().astimezone() - grace, + ) + for compute_job_id in reconciled: + # Never routine: every one of these is a run that ended without being + # able to say so, so it is worth a line in the worker log. logger.warning( - "Retry budget exhausted for compute job %s: marking FAILED", - sanitize_for_log(compute_job_id), + "compute job %s was terminal in the queue but still unfinished; " + "status reconciled to failed", + compute_job_id, ) - _fail_exhausted(compute_job_id) + + +async def _run_periodically( + description: str, + run_pass: Callable[[], Awaitable[None]], + stop: asyncio.Event, + interval: float, +) -> None: + """Run one maintenance pass on an interval until `stop` is set. + + Plain asyncio tasks in the worker process, not `rqueue.Scheduler` jobs. + Both passes are idempotent and harmless to run twice, so neither needs the + scheduler's occurrence-key machinery, and a scheduler process would be dead + weight beside them. A failed pass is logged and the loop continues: a + database outage must not silently end maintenance for the process lifetime. + """ + while not stop.is_set(): try: - await app.job_manager.finish_job_by_id_async( - job_id=job.id, status=ProcrastinateStatus.FAILED, delete_job=False - ) - except procrastinate_exceptions.ConnectorException: - logger.info("Stalled job %s was already resolved", job.id) - - -def _fail_exhausted(compute_job_id: str) -> None: - job_uuid = repository.as_uuid(compute_job_id) - with pooled() as conn: - row = repository.fetch_by_id(conn, job_uuid) - if row is None or row["status"] in { - JobStatus.COMPLETED.value, - JobStatus.FAILED.value, - }: - return - repository.fail_job( - conn, job_uuid, row["simulation_id"], CRASH_BUDGET_EXHAUSTED_ERROR - ) + await run_pass() + except Exception: + logger.exception("%s failed", description) + with contextlib.suppress(TimeoutError): + async with asyncio.timeout(interval): + await stop.wait() -@app.periodic(cron="0 * * * *") -@app.task(name="api.sweep_abandoned_work_dirs", queue=PROCRASTINATE_QUEUE) -def sweep_abandoned_work_dirs_task(timestamp: int) -> None: - cutoff = datetime.now().astimezone() - WORK_DIR_TTL - for simulation_id in repository.list_abandoned_work_dirs(cutoff): - shutil.rmtree(JOBS_DIR / simulation_id, ignore_errors=True) +async def run_periodic_sweep( + stop: asyncio.Event, *, interval: float = SWEEP_INTERVAL_SECONDS +) -> None: + """Sweep abandoned workspaces until `stop` is set.""" + await _run_periodically( + "Sweep of abandoned work directories", + lambda: sweep_abandoned_work_dirs(), + stop, + interval, + ) + + +async def run_periodic_reconcile( + stop: asyncio.Event, *, interval: float = RECONCILE_INTERVAL_SECONDS +) -> None: + """Reconcile queue-terminal jobs until `stop` is set.""" + await _run_periodically( + "Reconciliation of queue-terminal jobs", + lambda: reconcile_terminal_jobs(), + stop, + interval, + ) + + +def register_tasks(queue: Queue) -> Queue: + """Register every task this deployment runs, and return the queue. + + Both processes register. The producer needs it too: `Queue.build_insert` + reads the registration to stamp the task's retry budget and timeout onto + the row it writes, so an unregistered producer would quietly enqueue jobs + with the queue's defaults instead of TRANSIENT_RETRY's. + """ + if RUN_SIMULATION not in queue.tasks: + queue.register( + name=RUN_SIMULATION, + handler=run_simulation_task, + decoder=decode_payload, + retry=TRANSIENT_RETRY, + timeout=RUN_TIMEOUT_SECONDS, + ) + return queue diff --git a/packages/api/api/main.py b/packages/api/api/main.py index ba431de..d9c837a 100644 --- a/packages/api/api/main.py +++ b/packages/api/api/main.py @@ -8,8 +8,10 @@ from fastapi.middleware.cors import CORSMiddleware from api import __version__ -from api.core.db import close_pool, get_pool -from api.core.settings import LOG_LEVEL +from api.core import db +from api.core.queue import build_queue +from api.core.settings import LOG_LEVEL, api_pool_size +from api.core.tasks import register_tasks from api.routes import get_calculator, ops_router, router # Send logs to stdout so container runtimes can collect them. @@ -23,13 +25,15 @@ @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: get_calculator() - # Open the pool without waiting for Postgres. - get_pool() + min_size, max_size = api_pool_size() + # A zero floor keeps this from waiting on Postgres to start. + pool = await db.open_pool(min_size=min_size, max_size=max_size) + register_tasks(build_queue(pool)) logger.info("TSDHN API ready") try: yield finally: - close_pool() + await db.close_pool() def create_app() -> FastAPI: diff --git a/packages/api/api/migrate.py b/packages/api/api/migrate.py index 3315e97..f256ed1 100644 --- a/packages/api/api/migrate.py +++ b/packages/api/api/migrate.py @@ -1,6 +1,6 @@ """Create the compute schema and provision the web database role. -Run this before applying the Procrastinate and web schemas. The command uses +Run this before applying the queue and web schemas. The command uses the database owner for migrations; the web role is a runtime-only role. """ diff --git a/packages/api/api/queue_migrate.py b/packages/api/api/queue_migrate.py deleted file mode 100644 index d7f3527..0000000 --- a/packages/api/api/queue_migrate.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Apply the vendor queue schema once and make the operation repeatable.""" - -from __future__ import annotations - -import logging - -import procrastinate -import psycopg -from psycopg import sql - -from api.core.settings import ( - COMPUTE_DATABASE_URL, - PROCRASTINATE_SCHEMA, - PROCRASTINATE_SEARCH_PATH, -) - -logger = logging.getLogger(__name__) - - -def queue_schema_state( - conn: psycopg.Connection[tuple[str, ...]], schema: str = PROCRASTINATE_SCHEMA -) -> tuple[bool, ...]: - """Return whether the required Procrastinate types and tables exist.""" - result = conn.execute( - """ - SELECT - to_regtype(%s) IS NOT NULL, - to_regtype(%s) IS NOT NULL, - to_regtype(%s) IS NOT NULL, - to_regclass(%s) IS NOT NULL, - to_regclass(%s) IS NOT NULL, - to_regclass(%s) IS NOT NULL, - to_regclass(%s) IS NOT NULL - """, - [ - f"{schema}.procrastinate_job_status", - f"{schema}.procrastinate_job_event_type", - f"{schema}.procrastinate_job_to_defer_v1", - f"{schema}.procrastinate_workers", - f"{schema}.procrastinate_jobs", - f"{schema}.procrastinate_periodic_defers", - f"{schema}.procrastinate_events", - ], - ).fetchone() - if result is None: - raise RuntimeError("failed to inspect the Procrastinate schema") - return tuple(bool(value) for value in result) - - -def move_legacy_schema(conn: psycopg.Connection[tuple[str, ...]]) -> None: - """Move a pre-compute-schema Procrastinate install into `compute`.""" - for type_name in ( - "procrastinate_job_status", - "procrastinate_job_event_type", - "procrastinate_job_to_defer_v1", - ): - conn.execute( - sql.SQL("ALTER TYPE public.{name} SET SCHEMA compute").format( - name=sql.Identifier(type_name) - ) - ) - - objects = conn.execute( - """ - SELECT c.relname, c.relkind - FROM pg_class AS c - JOIN pg_namespace AS n ON n.oid = c.relnamespace - WHERE n.nspname = 'public' - AND c.relname LIKE 'procrastinate_%' - AND c.relkind IN ('r', 'p') - """ - ).fetchall() - for name, _kind in objects: - conn.execute( - sql.SQL("ALTER TABLE public.{name} SET SCHEMA compute").format( - name=sql.Identifier(name) - ) - ) - - sequences = conn.execute( - """ - SELECT sequence.relname, table_.relname, column_.attname - FROM pg_class AS sequence - JOIN pg_namespace AS sequence_schema - ON sequence_schema.oid = sequence.relnamespace - JOIN pg_depend AS dependency - ON dependency.classid = 'pg_class'::regclass - AND dependency.objid = sequence.oid - AND dependency.deptype = 'a' - JOIN pg_class AS table_ ON table_.oid = dependency.refobjid - JOIN pg_namespace AS table_schema - ON table_schema.oid = table_.relnamespace - JOIN pg_attribute AS column_ - ON column_.attrelid = table_.oid - AND column_.attnum = dependency.refobjsubid - WHERE sequence_schema.nspname = 'public' - AND sequence.relkind = 'S' - AND table_schema.nspname = 'compute' - """ - ).fetchall() - for sequence_name, table_name, column_name in sequences: - sequence_identifier = sql.Identifier(sequence_name) - conn.execute( - sql.SQL("ALTER SEQUENCE public.{name} OWNED BY NONE").format( - name=sequence_identifier - ) - ) - conn.execute( - sql.SQL("ALTER SEQUENCE public.{name} SET SCHEMA compute").format( - name=sequence_identifier - ) - ) - conn.execute( - sql.SQL( - "ALTER SEQUENCE compute.{sequence_name} " - "OWNED BY compute.{table_name}.{column_name}" - ).format( - sequence_name=sequence_identifier, - table_name=sql.Identifier(table_name), - column_name=sql.Identifier(column_name), - ) - ) - - conn.execute( - """ - DO $$ - DECLARE function_record record; - BEGIN - FOR function_record IN - SELECT p.proname, pg_get_function_identity_arguments(p.oid) AS arguments - FROM pg_proc AS p - JOIN pg_namespace AS n ON n.oid = p.pronamespace - WHERE n.nspname = 'public' AND p.proname LIKE 'procrastinate_%' - LOOP - EXECUTE format( - 'ALTER FUNCTION public.%I(%s) SET SCHEMA compute', - function_record.proname, - function_record.arguments - ); - END LOOP; - END $$; - """ - ) - - -def apply_schema(conninfo: str) -> None: - """Apply Procrastinate's schema in `compute`, rejecting partial installs.""" - with psycopg.connect(conninfo) as conn: - state = queue_schema_state(conn) - legacy_state = queue_schema_state(conn, "public") - - if all(state): - logger.info("Procrastinate schema already applied") - return - if all(legacy_state): - with psycopg.connect(conninfo) as conn: - move_legacy_schema(conn) - if not all(queue_schema_state(conn)): - raise RuntimeError("failed to move the Procrastinate schema to compute") - conn.commit() - logger.info("Procrastinate schema moved from public to compute") - return - if any(state): - raise RuntimeError( - "partial Procrastinate schema in compute; repair it before retrying" - ) - if any(legacy_state): - raise RuntimeError( - "partial Procrastinate schema in public; repair it before retrying" - ) - - connector = procrastinate.PsycopgConnector( - conninfo=conninfo, - kwargs={"options": f"-c search_path={PROCRASTINATE_SEARCH_PATH}"}, - ) - sync_connector = connector.get_sync_connector() - sync_connector.open() - try: - schema_manager = procrastinate.App(connector=connector).schema_manager - schema_manager.apply_schema() - finally: - sync_connector.close() - logger.info("Procrastinate schema applied") - - -def main() -> None: # pragma: no cover - logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s") - apply_schema(COMPUTE_DATABASE_URL) - - -if __name__ == "__main__": # pragma: no cover - main() diff --git a/packages/api/api/routes.py b/packages/api/api/routes.py index 9db99b5..09c330e 100644 --- a/packages/api/api/routes.py +++ b/packages/api/api/routes.py @@ -1,3 +1,5 @@ +import asyncio +import contextlib import json import logging import tempfile @@ -8,14 +10,13 @@ from typing import Any import anyio -import psycopg from fastapi import APIRouter, Depends, HTTPException, status from fastapi.responses import RedirectResponse, StreamingResponse from api import __version__ -from api.core import repository -from api.core.db import CONNECT_TIMEOUT, notify_channel -from api.core.settings import COMPUTE_DATABASE_URL, SSE_MAX_DURATION +from api.core import db, repository +from api.core.db import notify_channel +from api.core.settings import SSE_MAX_DURATION from api.core.storage import output_store from api.core.tasks import enqueue_simulation from api.schemas import ( @@ -56,9 +57,7 @@ def get_calculator() -> TsunamiCalculator: @ops_router.get("/health", response_model=HealthStatus) async def health() -> HealthStatus: - database_connected = await anyio.to_thread.run_sync( - repository.is_database_connected - ) + database_connected = await repository.is_database_connected() storage_connected = await anyio.to_thread.run_sync(output_store.is_connected) return HealthStatus( status="healthy" if database_connected and storage_connected else "degraded", @@ -95,13 +94,10 @@ def compute() -> CalculationPreview: async def create_job(req: JobRequest) -> JobCreated: simulation_id = str(req.simulation_id) try: - job_status = await anyio.to_thread.run_sync( - partial( - repository.create_or_get_job, - data=req.input, - simulation_id=simulation_id, - defer=enqueue_simulation, - ) + job_status = await repository.create_or_get_job( + data=req.input, + simulation_id=simulation_id, + defer=enqueue_simulation, ) except ValueError as e: raise HTTPException( @@ -129,7 +125,7 @@ async def get_job(simulation_id: str) -> JobStatusResponse: @router.get("/jobs/{simulation_id}/outputs", response_model=OutputList) async def list_outputs(simulation_id: str) -> OutputList: try: - outputs = await anyio.to_thread.run_sync(repository.get_outputs, simulation_id) + outputs = await repository.get_outputs(simulation_id) except ValueError as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e return OutputList( @@ -152,7 +148,7 @@ async def list_outputs(simulation_id: str) -> OutputList: ) async def get_output(simulation_id: str, name: str) -> RedirectResponse: try: - outputs = await anyio.to_thread.run_sync(repository.get_outputs, simulation_id) + outputs = await repository.get_outputs(simulation_id) except ValueError as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e @@ -193,16 +189,27 @@ async def stream() -> AsyncIterator[str]: return deadline = anyio.current_time() + SSE_MAX_DURATION - aconn = await psycopg.AsyncConnection.connect( - COMPUTE_DATABASE_URL, autocommit=True, connect_timeout=CONNECT_TIMEOUT - ) + # asyncpg delivers notifications to a callback, not to an async + # generator, so a one-slot queue bridges the callback into this loop. + # It only ever carries "something changed"; the status is re-read from + # the row, which is what the client is actually shown. + wakeups: asyncio.Queue[None] = asyncio.Queue(maxsize=1) + + def on_notify( + _connection: object, _pid: int, _channel: str, _payload: str + ) -> None: + with contextlib.suppress(asyncio.QueueFull): + wakeups.put_nowait(None) + + # Deliberately not a pooled connection: a stream can hold this for + # SSE_MAX_DURATION, and a handful of watchers would drain the pool + # every other route shares. + connection = await db.connect() try: - await aconn.execute(f"LISTEN {channel}") + await connection.add_listener(channel, on_notify) last = job_status while anyio.current_time() < deadline: - notified = False - async for _ in aconn.notifies(timeout=_KEEPALIVE_SECONDS, stop_after=1): - notified = True + notified = await _wait_for_notification(wakeups) # Read on every tick to cover notifications that race the wait. job_status = await _job_status(simulation_id) @@ -214,7 +221,7 @@ async def stream() -> AsyncIterator[str]: elif not notified: yield ": keepalive\n\n" finally: - await aconn.close() + await connection.close() return StreamingResponse( stream(), @@ -223,9 +230,22 @@ async def stream() -> AsyncIterator[str]: ) +async def _wait_for_notification(wakeups: asyncio.Queue[None]) -> bool: + """Wait for one notification, or report the keepalive timeout instead.""" + try: + async with asyncio.timeout(_KEEPALIVE_SECONDS): + await wakeups.get() + except TimeoutError: + return False + # Collapse a burst: one re-read covers every notification behind it. + while not wakeups.empty(): + wakeups.get_nowait() + return True + + async def _job_status(simulation_id: str) -> dict[str, Any]: try: - return await anyio.to_thread.run_sync(repository.get_job_status, simulation_id) + return await repository.get_job_status(simulation_id) except ValueError as e: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) from e diff --git a/packages/api/api/worker.py b/packages/api/api/worker.py index 7d82150..258dccc 100644 --- a/packages/api/api/worker.py +++ b/packages/api/api/worker.py @@ -1,10 +1,28 @@ +import asyncio import logging +import os +import signal +import socket import numba +from rqueue import Worker -from api.core.db import close_pool -from api.core.procrastinate_app import app -from api.core.settings import LOG_LEVEL, NUMBA_THREADS, PROCRASTINATE_QUEUE +from api.core import db +from api.core.queue import build_queue +from api.core.settings import ( + COMPUTE_QUEUE, + LOG_LEVEL, + NUMBA_THREADS, + WORKER_CONCURRENCY, + WORKER_ID, + WORKER_LEASE_SECONDS, + worker_pool_size, +) +from api.core.tasks import ( + register_tasks, + run_periodic_reconcile, + run_periodic_sweep, +) logging.basicConfig( level=LOG_LEVEL, @@ -13,6 +31,64 @@ logger = logging.getLogger(__name__) +def worker_id() -> str: + """Name this worker so its heartbeats and leases are attributable.""" + if WORKER_ID: + return WORKER_ID + # rqueue restricts a worker id to letters, digits, '_', '.', ':' and '-'. + host = "".join( + c if c.isalnum() or c in "_.-" else "-" for c in socket.gethostname() + ) + return f"tsdhn-worker-{host or 'unknown'}-{os.getpid()}"[:128] + + +async def run() -> None: + min_size, max_size = worker_pool_size() + pool = await db.open_pool(min_size=min_size, max_size=max_size) + try: + worker = Worker( + register_tasks(build_queue(pool)), + worker_id=worker_id(), + concurrency=WORKER_CONCURRENCY, + lease_duration=WORKER_LEASE_SECONDS, + ) + + loop = asyncio.get_running_loop() + for received in (signal.SIGTERM, signal.SIGINT): + # stop() stops claiming and gives in-flight work a bounded grace + # period (rqueue's shutdown_timeout, 30s by default). A simulation + # runs for tens of minutes, so in practice it does not finish: + # rqueue cancels it and hands the lease straight back as `pending` + # rather than leaving it to expire, and the next worker picks it up + # and resumes from the checkpoints in its work directory. + # + # The lease is handed back on time. The *process* is not: rqueue's + # bounded default executor shuts down with wait=True, so run() does + # not return until the cancelled kernel thread finishes on its own. + # That is an rqueue defect, filed there rather than worked around + # here; in practice the orchestrator's own SIGKILL bounds it. + loop.add_signal_handler(received, worker.stop) + + # The worker process owns maintenance because it is the process that + # owns recovery: reconciliation exists to repair the compute.jobs rows + # rqueue's own lease recovery leaves behind. + stop_maintenance = asyncio.Event() + maintenance = [ + asyncio.create_task(run_periodic_sweep(stop_maintenance)), + asyncio.create_task(run_periodic_reconcile(stop_maintenance)), + ] + logger.info("simulation worker serving queue %s", COMPUTE_QUEUE) + try: + await worker.run() + finally: + stop_maintenance.set() + for task in maintenance: + task.cancel() + await asyncio.gather(*maintenance, return_exceptions=True) + finally: + await db.close_pool() + + def main() -> None: # pragma: no cover if NUMBA_THREADS is not None: # Numba lacks type stubs; suppress type checking. @@ -22,14 +98,7 @@ def main() -> None: # pragma: no cover NUMBA_THREADS, ) - app.open() - try: - app.run_worker(queues=[PROCRASTINATE_QUEUE]) - finally: - app.close() - # The periodic tasks use the read pool; close it so its worker - # threads are joined before interpreter shutdown. - close_pool() + asyncio.run(run()) if __name__ == "__main__": # pragma: no cover diff --git a/packages/api/pyproject.toml b/packages/api/pyproject.toml index b0c9817..ff95ef6 100644 --- a/packages/api/pyproject.toml +++ b/packages/api/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "tsdhn-api" version = "0.0.1" -description = "TSDHN FastAPI service and Procrastinate worker" +description = "TSDHN FastAPI service and rqueue simulation worker" requires-python = ">=3.14" readme = "readme.md" license = { text = "MIT" } @@ -14,16 +14,16 @@ dependencies = [ "uvicorn==0.51.0", "pydantic==2.13.4", "anyio==4.14.1", - "psycopg[binary,pool]>=3.3.4", + "psycopg[binary]>=3.3.4", "minio>=7.2.20", - "procrastinate>=3.9.0", + "asyncpg>=0.31.0", + "rqueue", ] [project.scripts] tsdhn-api = "api.main:start_app" tsdhn-worker = "api.worker:main" tsdhn-compute-migrate = "api.migrate:main" -tsdhn-procrastinate-migrate = "api.queue_migrate:main" tsdhn-web-grants = "api.web_grants:main" [build-system] @@ -33,3 +33,6 @@ build-backend = "uv_build" [tool.uv.build-backend] module-name = "api" module-root = "" + +[tool.uv.sources] +rqueue = { git = "https://github.com/totallynotdavid/transactions.git", rev = "545d67aa341372c663972427b2b520fcd07faf80" } diff --git a/packages/api/readme.md b/packages/api/readme.md index 278840c..68deb00 100644 --- a/packages/api/readme.md +++ b/packages/api/readme.md @@ -1,7 +1,7 @@ # tsdhn-api `tsdhn-api` accepts requests from the web server, records job state in -PostgreSQL, queues work with Procrastinate, runs the shared `tsdhn` engine, +PostgreSQL, queues work with `rqueue`, runs the shared `tsdhn` engine, and stores output files in MinIO. The browser does not call this service. The web app calls it with @@ -18,13 +18,30 @@ uv run tsdhn-api uv run tsdhn-worker ``` -Create the compute and queue tables before starting the API or worker: +Create the compute and queue tables before starting the API or worker. +`tsdhn-compute-migrate` also provisions the web application's database role, so +it needs a password for that role; there is no default, because a default +password is worse than an error: ```sh +export APP_DB_PASSWORD="$(openssl rand -hex 32)" # or set it in .env + uv run tsdhn-compute-migrate -uv run tsdhn-procrastinate-migrate +uv run rqueue \ + --database-url "${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@localhost:5432/tsdhn}" \ + --schema "${COMPUTE_QUEUE_SCHEMA:-task_queue}" \ + migrate ``` +Everything else defaults to the values `api/core/settings.py` uses, so with +`APP_DB_PASSWORD` set the block works on a fresh clone with nothing else +exported. `rqueue`'s CLI reads `RQUEUE_DATABASE_URL`, not +`COMPUTE_DATABASE_URL`, which is why the URL is passed explicitly. + +`rqueue` owns the queue tables and their migrations; `--schema` is a global +flag, before the subcommand. They live in `COMPUTE_QUEUE_SCHEMA` (default +`task_queue`), never in `compute`. + For a complete deployment, including web migrations and grants, follow [`DEPLOY.md`](../../DEPLOY.md). For local PostgreSQL, `mise run db-migrate` applies every database change. `mise run test-integration` creates disposable @@ -49,12 +66,14 @@ download flows. - `api/schemas.py` defines public request and response models. - `api/security.py` checks `COMPUTE_API_TOKEN`. - `api/core/repository.py` reads and updates compute jobs. -- `api/core/tasks.py` runs queued simulations and records progress. +- `api/core/tasks.py` registers and runs the queued simulation task. - `api/core/storage.py` uploads output files and creates download URLs. -- `api/core/procrastinate_app.py` defines the queue and scheduled cleanup work. -- `api/migrate.py`, `api/queue_migrate.py`, and `api/web_grants.py` apply the - database changes described in `DEPLOY.md`. -- `api/worker.py` starts a worker for the configured queue. +- `api/core/queue.py` holds the `rqueue.Queue` the API and worker share. +- `api/core/db.py` owns the process-wide asyncpg pool. +- `api/migrate.py` and `api/web_grants.py` apply the database changes described + in `DEPLOY.md`; `rqueue migrate` applies the queue's own. +- `api/worker.py` starts an `rqueue.Worker` for the configured queue and the + periodic workspace sweep. ## Tests diff --git a/packages/api/tests/conftest.py b/packages/api/tests/conftest.py index afc1440..d91da34 100644 --- a/packages/api/tests/conftest.py +++ b/packages/api/tests/conftest.py @@ -3,12 +3,18 @@ from __future__ import annotations import uuid -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator import psycopg import pytest +import pytest_asyncio +from rqueue import Queue, migrations +from api.core import db +from api.core.queue import build_queue, set_queue from api.core.schema import COMPUTE_SCHEMA_SQL +from api.core.settings import COMPUTE_QUEUE_SCHEMA +from api.core.tasks import register_tasks from scripts.database import create_database, drop_database LOCAL_DATABASE_URL = "postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn" @@ -33,3 +39,24 @@ def isolated_database() -> Iterator[str]: yield target.database_url finally: drop_database(LOCAL_DATABASE_URL, database_name) + + +@pytest_asyncio.fixture +async def queue( + isolated_database: str, monkeypatch: pytest.MonkeyPatch +) -> AsyncIterator[Queue]: + """Point the process-wide pool and queue at a disposable database. + + Both schemas are applied the way a deployment applies them: the compute + schema by `isolated_database`, the queue schema by rqueue's own migration + runner. Neither is ever created implicitly at import or worker startup. + """ + monkeypatch.setattr(db, "COMPUTE_DATABASE_URL", isolated_database) + pool = await db.open_pool(min_size=1, max_size=6) + async with pool.acquire() as connection: + await migrations.migrate(connection, schema=COMPUTE_QUEUE_SCHEMA) + try: + yield register_tasks(build_queue(pool)) + finally: + set_queue(None) + await db.close_pool() diff --git a/packages/api/tests/test_api.py b/packages/api/tests/test_api.py index 895b4d9..5b9ab11 100644 --- a/packages/api/tests/test_api.py +++ b/packages/api/tests/test_api.py @@ -1,5 +1,6 @@ +import asyncio import json -from collections.abc import AsyncIterator, Iterator +from collections.abc import Callable, Iterator from pathlib import Path from typing import Any, cast @@ -8,7 +9,7 @@ from fastapi.testclient import TestClient from api import routes -from api.core import repository +from api.core import db, repository from api.core.storage import output_store from api.main import app @@ -33,7 +34,17 @@ def _service_token(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.fixture -def client() -> Iterator[TestClient]: +def client(monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + # The lifespan opens the process pool; these tests stub the repository + # instead, so it is replaced with one that never reaches PostgreSQL. + async def open_pool(**_kwargs: Any) -> object: + return object() + + async def close_pool() -> None: + return None + + monkeypatch.setattr(db, "open_pool", open_pool) + monkeypatch.setattr(db, "close_pool", close_pool) with TestClient(app) as test_client: yield test_client @@ -42,10 +53,19 @@ def _auth() -> dict[str, str]: return {"Authorization": f"Bearer {TOKEN}"} +def _async(value: Any) -> Callable[..., Any]: + """Return an async stub yielding `value`, for the now-async repository.""" + + async def stub(*_args: Any, **_kwargs: Any) -> Any: + return value + + return stub + + def test_health_is_unauthenticated( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(repository, "is_database_connected", lambda: True) + monkeypatch.setattr(repository, "is_database_connected", _async(True)) monkeypatch.setattr(output_store, "is_connected", lambda: True) response = client.get("/api/v1/health") @@ -62,7 +82,7 @@ def test_health_is_unauthenticated( def test_health_reports_degraded_when_a_dependency_is_down( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(repository, "is_database_connected", lambda: True) + monkeypatch.setattr(repository, "is_database_connected", _async(True)) monkeypatch.setattr(output_store, "is_connected", lambda: False) assert client.get("/api/v1/health").json()["status"] == "degraded" @@ -80,14 +100,16 @@ def test_get_job_returns_repository_status( monkeypatch.setattr( repository, "get_job_status", - lambda _simulation_id: { - "status": "running", - "details": "Processing tsunami", - "step": "tsunami", - "step_index": 3, - "total_steps": 8, - "outputs": [], - }, + _async( + { + "status": "running", + "details": "Processing tsunami", + "step": "tsunami", + "step_index": 3, + "total_steps": 8, + "outputs": [], + } + ), ) response = client.get(f"/api/v1/jobs/{SIMULATION_ID}", headers=_auth()) @@ -101,7 +123,7 @@ def test_get_job_returns_repository_status( def test_get_job_maps_unknown_job_to_not_found( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - def unknown(_simulation_id: str) -> dict[str, Any]: + async def unknown(_simulation_id: str) -> dict[str, Any]: raise ValueError("unknown") monkeypatch.setattr( @@ -143,7 +165,7 @@ def test_jobs_use_simulation_id_to_reuse_an_existing_job( ) -> None: calls: list[dict[str, Any]] = [] - def create_or_get_job(**kwargs: Any) -> dict[str, Any]: + async def create_or_get_job(**kwargs: Any) -> dict[str, Any]: calls.append(kwargs) return { "status": "queued", @@ -163,6 +185,8 @@ def create_or_get_job(**kwargs: Any) -> dict[str, Any]: "status": "queued", } assert calls[0]["simulation_id"] == SIMULATION_ID + # The route hands the repository the queue seam, not a queue import. + assert calls[0]["defer"] is routes_module.enqueue_simulation def test_outputs_require_a_token(client: TestClient) -> None: @@ -172,7 +196,7 @@ def test_outputs_require_a_token(client: TestClient) -> None: def test_outputs_list_is_empty_until_the_job_completes( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(repository, "get_outputs", lambda _id: []) + monkeypatch.setattr(repository, "get_outputs", _async([])) response = client.get(f"/api/v1/jobs/{SIMULATION_ID}/outputs", headers=_auth()) assert response.status_code == 200 @@ -185,14 +209,16 @@ def test_outputs_list_names_what_the_job_produced( monkeypatch.setattr( repository, "get_outputs", - lambda _id: [ - { - "name": "max_height_map", - "key": f"simulations/{SIMULATION_ID}/outputs/maxola.pdf", - "filename": "maxola.pdf", - "content_type": "application/pdf", - } - ], + _async( + [ + { + "name": "max_height_map", + "key": f"simulations/{SIMULATION_ID}/outputs/maxola.pdf", + "filename": "maxola.pdf", + "content_type": "application/pdf", + } + ] + ), ) response = client.get(f"/api/v1/jobs/{SIMULATION_ID}/outputs", headers=_auth()) @@ -214,14 +240,16 @@ def test_output_download_redirects_to_a_presigned_url( monkeypatch.setattr( repository, "get_outputs", - lambda _id: [ - { - "name": "max_height_map", - "key": f"simulations/{SIMULATION_ID}/outputs/maxola.pdf", - "filename": "maxola.pdf", - "content_type": "application/pdf", - } - ], + _async( + [ + { + "name": "max_height_map", + "key": f"simulations/{SIMULATION_ID}/outputs/maxola.pdf", + "filename": "maxola.pdf", + "content_type": "application/pdf", + } + ] + ), ) monkeypatch.setattr( output_store, @@ -241,7 +269,7 @@ def test_output_download_redirects_to_a_presigned_url( def test_unknown_output_name_is_404( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(repository, "get_outputs", lambda _id: []) + monkeypatch.setattr(repository, "get_outputs", _async([])) response = client.get( f"/api/v1/jobs/{SIMULATION_ID}/outputs/nope", @@ -254,7 +282,7 @@ def test_unknown_output_name_is_404( def test_output_download_maps_an_unknown_job_to_not_found( client: TestClient, monkeypatch: pytest.MonkeyPatch ) -> None: - def unknown(_simulation_id: str) -> list[dict[str, str]]: + async def unknown(_simulation_id: str) -> list[dict[str, str]]: raise ValueError("unknown") monkeypatch.setattr(repository, "get_outputs", unknown) @@ -277,20 +305,18 @@ def test_job_events_maps_an_invalid_job_id_to_not_found(client: TestClient) -> N assert response.status_code == 404 -class _FakeAsyncConnection: +class _FakeListenerConnection: + """Stands in for the dedicated asyncpg connection the SSE stream opens.""" + def __init__(self, *, notify: bool) -> None: self.notify = notify - self.executed: list[str] = [] + self.channels: list[str] = [] self.closed = False - async def execute(self, statement: str) -> None: - self.executed.append(statement) - - async def notifies( - self, *, timeout: float, stop_after: int - ) -> AsyncIterator[object]: + async def add_listener(self, channel: str, callback: Callable[..., None]) -> None: + self.channels.append(channel) if self.notify: - yield object() + callback(self, 0, channel, "") async def close(self) -> None: self.closed = True @@ -300,10 +326,17 @@ async def close(self) -> None: async def test_job_events_emits_a_terminal_snapshot_without_opening_a_listener( monkeypatch: pytest.MonkeyPatch, ) -> None: + opened: list[object] = [] + async def status(_simulation_id: str) -> dict[str, Any]: return {"status": "completed", "outputs": ["result"]} + async def connect() -> object: + opened.append(object()) + return opened[-1] + monkeypatch.setattr(routes, "_job_status", status) + monkeypatch.setattr(routes_module.db, "connect", connect) response = await routes.job_events(SIMULATION_ID) chunks = [chunk async for chunk in response.body_iterator] @@ -314,6 +347,7 @@ async def status(_simulation_id: str) -> dict[str, Any]: "status": "completed", "outputs": ["result"], } + assert opened == [] @pytest.mark.asyncio @@ -341,16 +375,16 @@ async def test_job_events_reads_a_changed_snapshot_and_closes_the_listener( {"status": "completed", "step": "tsunami"}, ] ) - connection = _FakeAsyncConnection(notify=True) + connection = _FakeListenerConnection(notify=True) async def status(_simulation_id: str) -> dict[str, Any]: return next(snapshots) - async def connect(*args: Any, **kwargs: Any) -> _FakeAsyncConnection: + async def connect() -> _FakeListenerConnection: return connection monkeypatch.setattr(routes, "_job_status", status) - monkeypatch.setattr(routes_module.psycopg.AsyncConnection, "connect", connect) + monkeypatch.setattr(routes_module.db, "connect", connect) monkeypatch.setattr(routes_module.anyio, "current_time", lambda: 0.0) response = await routes.job_events(SIMULATION_ID) @@ -358,7 +392,7 @@ async def connect(*args: Any, **kwargs: Any) -> _FakeAsyncConnection: assert len(chunks) == 2 assert '"status": "completed"' in chunks[1] - assert connection.executed == ["LISTEN tsdhn_job_4cfe522f7e7d46e096ca7b98743fb9f5"] + assert connection.channels == ["tsdhn_job_4cfe522f7e7d46e096ca7b98743fb9f5"] assert connection.closed @@ -367,19 +401,21 @@ async def test_job_events_sends_keepalive_when_state_does_not_change( monkeypatch: pytest.MonkeyPatch, ) -> None: snapshot = {"status": "running", "step": "tsunami"} - connection = _FakeAsyncConnection(notify=False) + connection = _FakeListenerConnection(notify=False) snapshots = iter([snapshot, snapshot]) clock = iter([0.0, 0.0, float("inf")]) async def status(_simulation_id: str) -> dict[str, Any]: return next(snapshots) - async def connect(*args: Any, **kwargs: Any) -> _FakeAsyncConnection: + async def connect() -> _FakeListenerConnection: return connection monkeypatch.setattr(routes, "_job_status", status) - monkeypatch.setattr(routes_module.psycopg.AsyncConnection, "connect", connect) + monkeypatch.setattr(routes_module.db, "connect", connect) monkeypatch.setattr(routes_module.anyio, "current_time", lambda: next(clock)) + # Without a notification the stream falls through on the keepalive timeout. + monkeypatch.setattr(routes_module, "_KEEPALIVE_SECONDS", 0.01) response = await routes.job_events(SIMULATION_ID) chunks = [chunk async for chunk in response.body_iterator] @@ -388,6 +424,25 @@ async def connect(*args: Any, **kwargs: Any) -> _FakeAsyncConnection: assert connection.closed +@pytest.mark.asyncio +async def test_the_notification_bridge_collapses_a_burst_into_one_reread() -> None: + wakeups: asyncio.Queue[None] = asyncio.Queue(maxsize=1) + wakeups.put_nowait(None) + + assert await routes_module._wait_for_notification(wakeups) is True + assert wakeups.empty() + + +@pytest.mark.asyncio +async def test_the_notification_bridge_reports_the_keepalive_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(routes_module, "_KEEPALIVE_SECONDS", 0.01) + wakeups: asyncio.Queue[None] = asyncio.Queue(maxsize=1) + + assert await routes_module._wait_for_notification(wakeups) is False + + def test_legacy_simulations_endpoint_is_removed(client: TestClient) -> None: response = client.post( "/api/v1/simulations", diff --git a/packages/api/tests/test_db.py b/packages/api/tests/test_db.py index 957fc54..f37617a 100644 --- a/packages/api/tests/test_db.py +++ b/packages/api/tests/test_db.py @@ -1,44 +1,219 @@ """Connection boundary behavior for the compute service.""" -from collections.abc import Iterator -from contextlib import contextmanager from typing import Any -import psycopg +import asyncpg import pytest -from api.core import db +from api.core import db, settings from api.core.errors import TransientInfraError db_module: Any = db -class _Pool: - def __init__(self) -> None: - self.connection_value = object() +@pytest.mark.asyncio +async def test_acquire_borrows_and_returns_a_pooled_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connection = object() + released: list[object] = [] + + class _Pool: + async def acquire(self) -> object: + return connection + + async def release(self, borrowed: object) -> None: + released.append(borrowed) + + monkeypatch.setattr(db_module, "get_pool", lambda: _Pool()) + + async with db.acquire() as borrowed: + assert borrowed is connection + assert released == [connection] + + +@pytest.mark.asyncio +async def test_acquire_classifies_a_database_outage_as_transient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Pool: + async def acquire(self) -> object: + raise ConnectionRefusedError("database unavailable") + + monkeypatch.setattr(db_module, "get_pool", lambda: _Pool()) + + with pytest.raises(TransientInfraError, match="database unavailable"): + async with db.acquire(): + pytest.fail("acquire must not yield when the database is unreachable") + + +@pytest.mark.asyncio +async def test_connect_classifies_a_database_outage_as_transient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def fail(*args: Any, **kwargs: Any) -> object: + raise OSError("database unavailable") + + monkeypatch.setattr(db_module.asyncpg, "connect", fail) + + with pytest.raises(TransientInfraError, match="database unavailable"): + await db.connect() + + +@pytest.mark.asyncio +async def test_a_postgres_connection_error_is_transient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # asyncpg raises this, not OSError, when the server answers and refuses. + async def fail(*args: Any, **kwargs: Any) -> object: + raise asyncpg.PostgresConnectionError("server closed the connection") + + monkeypatch.setattr(db_module.asyncpg, "connect", fail) + + with pytest.raises(TransientInfraError): + await db.connect() + + +def test_get_pool_says_which_call_was_missed(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(db_module, "_pool", None) - @contextmanager - def connection(self) -> Iterator[object]: - yield self.connection_value + with pytest.raises(RuntimeError, match=r"db\.open_pool"): + db.get_pool() -def test_pooled_yields_a_connection_from_the_process_pool( +def test_the_worker_pool_leaves_room_for_the_listener_connection( monkeypatch: pytest.MonkeyPatch, ) -> None: - pool = _Pool() - monkeypatch.setattr(db_module, "get_pool", lambda: pool) + # rqueue.Worker._listener() holds one connection for the whole run. A pool + # sized to concurrency alone would starve it and silently drop back to + # polling, so the floor has to exceed the worker's own concurrency. + monkeypatch.setattr(settings, "WORKER_CONCURRENCY", 4) + monkeypatch.setattr(settings, "DB_POOL_MAX_SIZE", 2) - with db.pooled() as connection: - assert connection is pool.connection_value + _, max_size = settings.worker_pool_size() + assert max_size > settings.WORKER_CONCURRENCY + 1 + assert max_size == 2 * 4 + 4 -def test_connect_classifies_a_database_outage_as_transient( + +def test_the_worker_pool_never_shrinks_below_the_configured_maximum( monkeypatch: pytest.MonkeyPatch, ) -> None: - def fail(*args: Any, **kwargs: Any) -> object: - raise psycopg.OperationalError("database unavailable") + monkeypatch.setattr(settings, "WORKER_CONCURRENCY", 1) + monkeypatch.setattr(settings, "DB_POOL_MAX_SIZE", 20) + + assert settings.worker_pool_size()[1] == 20 + + +def test_notify_channel_is_a_bare_identifier_per_job() -> None: + import uuid + + simulation_id = uuid.UUID("4cfe522f-7e7d-46e0-96ca-7b98743fb9f5") + channel = db.notify_channel(simulation_id) + + assert channel == "tsdhn_job_4cfe522f7e7d46e096ca7b98743fb9f5" + # No hyphens or quoting needed: it goes straight into LISTEN/NOTIFY. + assert channel.replace("_", "").isalnum() + assert db.notify_channel(uuid.uuid4()) != channel + + +class _Connection: + """A connection that reports whether it is closed, like asyncpg's.""" + + def __init__(self, *, closed: bool) -> None: + self._closed = closed + + def is_closed(self) -> bool: + return self._closed + + +@pytest.mark.asyncio +async def test_a_backend_that_dies_mid_statement_is_transient() -> None: + # What asyncpg raises when the backend goes away with a query in flight. + with pytest.raises(TransientInfraError, match="database unavailable"): + async with db.transient_connection_errors(_Connection(closed=False)): + raise asyncpg.ConnectionDoesNotExistError("connection was closed") + + +@pytest.mark.asyncio +async def test_a_statement_issued_on_a_dead_connection_is_transient() -> None: + # What asyncpg raises when the backend was already gone: a bare + # InterfaceError, which is in no connection-error class at all. + with pytest.raises(TransientInfraError, match="database unavailable"): + async with db.transient_connection_errors(_Connection(closed=True)): + raise asyncpg.InterfaceError("connection is closed") + + +@pytest.mark.asyncio +async def test_a_programming_error_is_not_disguised_as_an_outage() -> None: + # asyncpg raises InterfaceError for a bad call too ("the server expects 2 + # arguments for this query, 1 was passed"), on a connection that is still + # perfectly healthy. Retrying that spends the job's budget on a bug. + with pytest.raises(asyncpg.InterfaceError, match="expects 2 arguments"): + async with db.transient_connection_errors(_Connection(closed=False)): + raise asyncpg.InterfaceError("the server expects 2 arguments") + + +@pytest.mark.asyncio +async def test_an_unrelated_error_passes_through_untouched() -> None: + with pytest.raises(ValueError, match="not a connection problem"): + async with db.transient_connection_errors(_Connection(closed=True)): + raise ValueError("not a connection problem") + + +@pytest.mark.asyncio +async def test_a_pool_proxy_that_cannot_answer_is_treated_as_gone() -> None: + # A pooled connection whose backend died is terminated and taken back by + # the pool, and every method on the proxy then raises -- is_closed() too. + class _ReleasedProxy: + def is_closed(self) -> bool: + raise asyncpg.InterfaceError( + "cannot call Connection.is_closed(): " + "connection has been released back to the pool" + ) + + with pytest.raises(TransientInfraError, match="database unavailable"): + async with db.transient_connection_errors(_ReleasedProxy()): + raise asyncpg.InterfaceError("cannot call Connection.execute()") + + +@pytest.mark.parametrize( + ("exception", "transient"), + [ + # 53300: the pool or the server ran out of connections. Not a + # PostgresConnectionError -- the reason an exception-type allowlist + # classified this as permanent and threw away retryable jobs. + (asyncpg.TooManyConnectionsError("sorry, too many clients already"), True), + (asyncpg.OutOfMemoryError("out of memory"), True), # 53200 + (asyncpg.ConnectionDoesNotExistError("connection was closed"), True), # 08003 + (asyncpg.CannotConnectNowError("the database is starting up"), True), # 57P03 + (asyncpg.AdminShutdownError("terminating connection"), True), # 57P01 + (asyncpg.DeadlockDetectedError("deadlock detected"), True), # 40P01 + (asyncpg.SerializationError("could not serialize"), True), # 40001 + (ConnectionResetError("peer reset"), True), + (TimeoutError("timed out"), True), + # Application bugs: a retry re-runs a whole simulation to fail the same + # way, so these have to stay permanent. + (asyncpg.UndefinedColumnError("column does not exist"), False), # 42703 + (asyncpg.UniqueViolationError("duplicate key"), False), # 23505 + (asyncpg.InvalidTextRepresentationError("invalid uuid"), False), # 22P02 + (ValueError("not a database problem"), False), + ], +) +def test_transient_classification_follows_the_sqlstate_class( + exception: BaseException, transient: bool +) -> None: + assert db.is_transient(exception) is transient + + +@pytest.mark.asyncio +async def test_pool_exhaustion_is_retryable(monkeypatch: pytest.MonkeyPatch) -> None: + class _Pool: + async def acquire(self) -> object: + raise asyncpg.TooManyConnectionsError("sorry, too many clients already") - monkeypatch.setattr(db_module.psycopg, "connect", fail) + monkeypatch.setattr(db_module, "get_pool", lambda: _Pool()) - with pytest.raises(TransientInfraError, match="database connection failed"): - db.connect() + with pytest.raises(TransientInfraError, match="database unavailable"): + async with db.acquire(): + pass diff --git a/packages/api/tests/test_jobs.py b/packages/api/tests/test_jobs.py index 9a2dab1..17ded1f 100644 --- a/packages/api/tests/test_jobs.py +++ b/packages/api/tests/test_jobs.py @@ -1,63 +1,64 @@ -"""Database-free tests for retry, failure, and notification behavior.""" +"""Database-free tests for retry policy and payload decoding.""" import uuid -from procrastinate.jobs import Job as ProcrastinateJob +import pytest +from rqueue import PermanentFailure -from api.core.db import notify_channel from api.core.errors import TransientInfraError -from api.core.tasks import MAX_ATTEMPTS, TRANSIENT_RETRY, reap_action - - -def _job(attempts: int, job_id: int | None = None) -> ProcrastinateJob: - return ProcrastinateJob( - id=job_id, - queue="simulations", - lock=None, - queueing_lock=None, - task_name="api.run_simulation", - task_kwargs={"compute_job_id": "11111111-1111-4111-8111-111111111111"}, - attempts=attempts, - ) +from api.core.tasks import MAX_ATTEMPTS, RUN_SIMULATION, TRANSIENT_RETRY, decode_payload def test_transient_retry_retries_transient_infra_error_within_budget() -> None: - decision = TRANSIENT_RETRY.get_retry_decision( - exception=TransientInfraError("db down"), job=_job(attempts=0) + assert TRANSIENT_RETRY.should_retry(TransientInfraError("db down"), attempt=1) + assert TRANSIENT_RETRY.should_retry( + TransientInfraError("db down"), attempt=MAX_ATTEMPTS - 1 ) - assert decision is not None def test_transient_retry_gives_up_once_attempts_exhausted() -> None: - decision = TRANSIENT_RETRY.get_retry_decision( - exception=TransientInfraError("db down"), - job=_job(attempts=MAX_ATTEMPTS), + assert not TRANSIENT_RETRY.should_retry( + TransientInfraError("db down"), attempt=MAX_ATTEMPTS ) - assert decision is None def test_transient_retry_does_not_retry_other_exceptions() -> None: - decision = TRANSIENT_RETRY.get_retry_decision( - exception=RuntimeError("bad epicenter"), job=_job(attempts=0) - ) - assert decision is None + # retry_on is an allowlist: a pipeline error is terminal on attempt one. + assert not TRANSIENT_RETRY.should_retry(RuntimeError("bad epicenter"), attempt=1) + + +def test_transient_retry_never_retries_a_permanent_failure() -> None: + assert not TRANSIENT_RETRY.should_retry(PermanentFailure("unknown job"), attempt=1) + + +def test_transient_retry_backs_off_exponentially_from_fifteen_seconds() -> None: + first = TRANSIENT_RETRY.backoff_seconds(1) + second = TRANSIENT_RETRY.backoff_seconds(2) + + # Jitter is +/-10%, so these are ranges rather than exact values. + assert 13.5 <= first <= 16.5 + assert 27.0 <= second <= 33.0 -def test_notify_channel_is_a_bare_identifier_per_job() -> None: - simulation_id = uuid.UUID("4cfe522f-7e7d-46e0-96ca-7b98743fb9f5") - channel = notify_channel(simulation_id) +def test_the_task_name_is_the_one_already_deployed() -> None: + # Renaming this strands every job a running deployment has already queued. + assert RUN_SIMULATION == "api.run_simulation" - assert channel == "tsdhn_job_4cfe522f7e7d46e096ca7b98743fb9f5" - # No hyphens or quoting needed: it goes straight into LISTEN/NOTIFY. - assert channel.replace("_", "").isalnum() - assert notify_channel(uuid.uuid4()) != channel +def test_decode_payload_reads_the_compute_job_id() -> None: + compute_job_id = uuid.uuid4() -def test_reap_action_retries_a_stalled_job_within_budget() -> None: - assert reap_action(_job(attempts=0)) == "retry" - assert reap_action(_job(attempts=MAX_ATTEMPTS - 1)) == "retry" + assert decode_payload({"compute_job_id": str(compute_job_id)}) == compute_job_id -def test_reap_action_gives_up_once_the_budget_is_spent() -> None: - assert reap_action(_job(attempts=MAX_ATTEMPTS)) == "exhausted" - assert reap_action(_job(attempts=MAX_ATTEMPTS + 1)) == "exhausted" +@pytest.mark.parametrize( + "payload", + [None, [], "compute_job_id", {}, {"compute_job_id": "not-a-uuid"}], +) +def test_decode_payload_rejects_anything_it_cannot_name_a_job_from( + payload: object, +) -> None: + # A decode failure is durable in rqueue and costs no retry budget, so it + # has to raise rather than return a placeholder. + with pytest.raises((ValueError, KeyError, TypeError)): + decode_payload(payload) diff --git a/packages/api/tests/test_migrations_integration.py b/packages/api/tests/test_migrations_integration.py index 8870bdd..9124afd 100644 --- a/packages/api/tests/test_migrations_integration.py +++ b/packages/api/tests/test_migrations_integration.py @@ -1,18 +1,35 @@ """Database-role behavior for the compute migration.""" +import asyncio import uuid from collections.abc import Iterator +import asyncpg import psycopg import pytest from psycopg import sql +from rqueue import migrations -from api import migrate, queue_migrate, web_grants +from api import migrate, web_grants from api.core.schema import COMPUTE_SCHEMA_SQL +from api.core.settings import COMPUTE_QUEUE_SCHEMA pytestmark = pytest.mark.integration +def _apply_queue_schema(database_url: str) -> None: + """Apply rqueue's own migrations the way a deployment's CLI does.""" + + async def apply() -> None: + connection = await asyncpg.connect(database_url) + try: + await migrations.migrate(connection, schema=COMPUTE_QUEUE_SCHEMA) + finally: + await connection.close() + + asyncio.run(apply()) + + @pytest.fixture def temporary_web_role( isolated_database: str, monkeypatch: pytest.MonkeyPatch @@ -179,21 +196,15 @@ def test_provision_web_role_transfers_legacy_table_ownership( assert owner == (migration_user,) -def test_procrastinate_schema_migration_is_repeatable(isolated_database: str) -> None: - queue_migrate.apply_schema(isolated_database) - queue_migrate.apply_schema(isolated_database) - - with psycopg.connect(isolated_database) as conn: - assert all(queue_migrate.queue_schema_state(conn)) - - -def test_web_role_cannot_read_compute_queue_tables( +def test_web_role_cannot_read_the_queue_tables( isolated_database: str, temporary_web_role: str ) -> None: + # The queue lives in its own schema, which the web grants never mention. + # Stage 2 owns rqueue's own least-privilege roles; this asserts the floor. with psycopg.connect(isolated_database) as conn: web_grants.grant_web_tables(conn) conn.commit() - queue_migrate.apply_schema(isolated_database) + _apply_queue_schema(isolated_database) with ( psycopg.connect( @@ -201,4 +212,4 @@ def test_web_role_cannot_read_compute_queue_tables( ) as conn, pytest.raises(psycopg.errors.InsufficientPrivilege), ): - conn.execute("SELECT 1 FROM compute.procrastinate_jobs") + conn.execute(f"SELECT 1 FROM {COMPUTE_QUEUE_SCHEMA}.jobs") # noqa: S608 diff --git a/packages/api/tests/test_repository_integration.py b/packages/api/tests/test_repository_integration.py index 6a8e863..6962bfd 100644 --- a/packages/api/tests/test_repository_integration.py +++ b/packages/api/tests/test_repository_integration.py @@ -1,19 +1,19 @@ """Compute repository behavior against PostgreSQL.""" +import asyncio import uuid -from collections.abc import Iterator -from contextlib import contextmanager from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any -import psycopg import pytest -from psycopg.rows import dict_row +from rqueue import Queue -from api.core import repository +from api.core import db, repository from api.core.errors import TransientInfraError +from api.core.settings import COMPUTE_QUEUE, COMPUTE_QUEUE_SCHEMA from api.core.storage import output_store +from api.core.tasks import RUN_SIMULATION, enqueue_simulation from tsdhn.domain import CalculationResponse, EarthquakeInput, TsunamiTravelResponse from tsdhn.engine import OutputFile, SimulationOutputs, SimulationResult from tsdhn.runtime import RuntimeContext @@ -23,160 +23,187 @@ INPUT = EarthquakeInput(Mw=8.0, h=10.0, lat0=-20.5, lon0=-70.5, hhmm="0000", dia="23") -@pytest.fixture -def database( - isolated_database: str, - monkeypatch: pytest.MonkeyPatch, -) -> str: - @contextmanager - def pooled() -> Iterator[psycopg.Connection[dict[str, Any]]]: - with psycopg.connect(isolated_database, row_factory=dict_row) as conn: - yield conn - - # Keep the production query path while pointing it at the isolated database. - monkeypatch.setattr(repository, "pooled", pooled) - return isolated_database - - -def _create_job( - database: str, - *, - data: EarthquakeInput = INPUT, -) -> tuple[str, uuid.UUID]: +async def _create_job(*, data: EarthquakeInput = INPUT) -> tuple[str, uuid.UUID]: simulation_id = uuid.uuid4() deferred: list[uuid.UUID] = [] - repository.create_or_get_job( - data=data, - simulation_id=str(simulation_id), - defer=lambda _conn, compute_job_id: deferred.append(compute_job_id), + async def defer(_conn: Any, compute_job_id: uuid.UUID) -> None: + deferred.append(compute_job_id) + + await repository.create_or_get_job( + data=data, simulation_id=str(simulation_id), defer=defer ) assert len(deferred) == 1 return str(simulation_id), deferred[0] -def test_create_or_get_job_is_idempotent_and_rejects_changed_input( - database: str, +async def _queue_rows(queue: Queue) -> list[Any]: + async with db.acquire() as conn: + return list( + await conn.fetch( + f"SELECT * FROM {COMPUTE_QUEUE_SCHEMA}.jobs ORDER BY seq" # noqa: S608 + ) + ) + + +@pytest.mark.asyncio +async def test_create_or_get_job_is_idempotent_and_rejects_changed_input( + queue: Queue, ) -> None: - simulation_id, compute_job_id = _create_job(database) + simulation_id, compute_job_id = await _create_job() second_defer: list[uuid.UUID] = [] - same = repository.create_or_get_job( - data=INPUT, - simulation_id=simulation_id, - defer=lambda _conn, job_id: second_defer.append(job_id), + + async def defer(_conn: Any, job_id: uuid.UUID) -> None: + second_defer.append(job_id) + + same = await repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=defer ) assert same["status"] == "queued" assert second_defer == [] changed = INPUT.model_copy(update={"Mw": 8.1}) + + async def unreached(*_args: Any) -> None: + pytest.fail("a conflicting submission must not reach the queue") + with pytest.raises(ValueError, match="different input"): - repository.create_or_get_job( - data=changed, - simulation_id=simulation_id, - defer=lambda *_args: None, + await repository.create_or_get_job( + data=changed, simulation_id=simulation_id, defer=unreached ) - persisted = repository.get_job_status(simulation_id) + persisted = await repository.get_job_status(simulation_id) assert persisted["status"] == "queued" - with psycopg.connect(database) as conn: - persisted_id = conn.execute( - "SELECT id FROM compute.jobs WHERE simulation_id = %s", - [simulation_id], - ).fetchone() - assert persisted_id == (compute_job_id,) + async with db.acquire() as conn: + persisted_id = await conn.fetchval( + "SELECT id FROM compute.jobs WHERE simulation_id = $1", + uuid.UUID(simulation_id), + ) + assert persisted_id == compute_job_id -def test_create_or_get_job_rolls_back_when_enqueue_fails( - database: str, -) -> None: +@pytest.mark.asyncio +async def test_the_job_row_and_its_queue_entry_commit_together(queue: Queue) -> None: + simulation_id = str(uuid.uuid4()) + + job_status = await repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=enqueue_simulation + ) + assert job_status["status"] == "queued" + + (row,) = await _queue_rows(queue) + async with db.acquire() as conn: + compute_job_id = await conn.fetchval( + "SELECT id FROM compute.jobs WHERE simulation_id = $1", + uuid.UUID(simulation_id), + ) + assert row["task"] == RUN_SIMULATION + assert row["queue"] == COMPUTE_QUEUE + assert row["state"] == "pending" + assert row["dedupe_key"] == f"simulation:{compute_job_id}" + + +@pytest.mark.asyncio +async def test_a_rolled_back_producer_leaves_neither_row(queue: Queue) -> None: simulation_id = str(uuid.uuid4()) - def fail_to_enqueue(_conn: Any, _compute_job_id: uuid.UUID) -> None: + async def enqueue_then_fail(conn: Any, compute_job_id: uuid.UUID) -> None: + await enqueue_simulation(conn, compute_job_id) raise RuntimeError("queue unavailable") with pytest.raises(RuntimeError, match="queue unavailable"): - repository.create_or_get_job( - data=INPUT, - simulation_id=simulation_id, - defer=fail_to_enqueue, + await repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=enqueue_then_fail ) with pytest.raises(ValueError, match="Invalid or unknown job ID"): - repository.get_job_status(simulation_id) + await repository.get_job_status(simulation_id) + assert await _queue_rows(queue) == [] -def test_concurrent_identical_submissions_create_one_job( - database: str, -) -> None: - from concurrent.futures import ThreadPoolExecutor - from threading import Barrier +@pytest.mark.asyncio +async def test_a_repeated_submission_enqueues_exactly_one_job(queue: Queue) -> None: + simulation_id = str(uuid.uuid4()) + + for _ in range(2): + await repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=enqueue_simulation + ) + + assert len(await _queue_rows(queue)) == 1 + +@pytest.mark.asyncio +async def test_concurrent_identical_submissions_create_one_job(queue: Queue) -> None: simulation_id = str(uuid.uuid4()) - start = Barrier(2) deferred: list[uuid.UUID] = [] - def submit() -> dict[str, Any]: - start.wait() - return repository.create_or_get_job( - data=INPUT, - simulation_id=simulation_id, - defer=lambda _conn, compute_job_id: deferred.append(compute_job_id), - ) + async def defer(_conn: Any, compute_job_id: uuid.UUID) -> None: + deferred.append(compute_job_id) - with ThreadPoolExecutor(max_workers=2) as executor: - results = list(executor.map(lambda _index: submit(), range(2))) + results = await asyncio.gather( + *( + repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=defer + ) + for _ in range(2) + ) + ) assert {result["status"] for result in results} == {"queued"} assert len(deferred) == 1 - with psycopg.connect(database) as conn: - count = conn.execute( - "SELECT count(*) FROM compute.jobs WHERE simulation_id = %s", - [simulation_id], - ).fetchone() - assert count == (1,) + async with db.acquire() as conn: + count = await conn.fetchval( + "SELECT count(*) FROM compute.jobs WHERE simulation_id = $1", + uuid.UUID(simulation_id), + ) + assert count == 1 -def test_job_lookup_rejects_invalid_and_unknown_simulation_ids( - database: str, +@pytest.mark.asyncio +async def test_job_lookup_rejects_invalid_and_unknown_simulation_ids( + queue: Queue, ) -> None: with pytest.raises(ValueError, match="Invalid or unknown job ID"): - repository.get_job_status("not-a-uuid") + await repository.get_job_status("not-a-uuid") with pytest.raises(ValueError, match="Invalid or unknown job ID"): - repository.get_job_status(str(uuid.uuid4())) + await repository.get_job_status(str(uuid.uuid4())) -def test_outputs_are_empty_before_completion_and_unknown_jobs_are_rejected( - database: str, +@pytest.mark.asyncio +async def test_outputs_are_empty_before_completion_and_unknown_jobs_are_rejected( + queue: Queue, ) -> None: - simulation_id, _compute_job_id = _create_job(database) + simulation_id, _compute_job_id = await _create_job() - assert repository.get_outputs(simulation_id) == [] + assert await repository.get_outputs(simulation_id) == [] with pytest.raises(ValueError, match="Invalid or unknown job ID"): - repository.get_outputs(str(uuid.uuid4())) + await repository.get_outputs(str(uuid.uuid4())) -def test_job_state_updates_are_persisted_as_client_visible_behavior( - database: str, +@pytest.mark.asyncio +async def test_job_state_updates_are_persisted_as_client_visible_behavior( + queue: Queue, ) -> None: - database_url = database - simulation_id_text, compute_job_id = _create_job(database) + simulation_id_text, compute_job_id = await _create_job() simulation_id = uuid.UUID(simulation_id_text) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - repository.mark_started(conn, compute_job_id, simulation_id) - repository.record_progress( + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await repository.record_progress( conn, compute_job_id, simulation_id, "Processing tsunami", {"step": "tsunami", "step_index": 3, "total_steps": 8}, + 1, ) - assert repository.get_current_step(conn, compute_job_id) == "tsunami" + assert await repository.get_current_step(conn, compute_job_id) == "tsunami" - running = repository.get_job_status(simulation_id_text) + running = await repository.get_job_status(simulation_id_text) assert running["status"] == "running" assert running["details"] == "Processing tsunami" assert (running["step"], running["step_index"], running["total_steps"]) == ( @@ -185,96 +212,164 @@ def test_job_state_updates_are_persisted_as_client_visible_behavior( 8, ) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - repository.fail_job(conn, compute_job_id, simulation_id, "worker vanished") - failed = repository.get_job_status(simulation_id_text) - assert failed["status"] == "failed" - assert failed["error"] == "worker vanished" - assert failed["finished_at"] is not None +@pytest.mark.asyncio +async def test_progress_json_survives_the_text_jsonb_round_trip(queue: Queue) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + calculation = {"length": 1.5, "corners": [[1, 2], [3, 4]], "warning": None} + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await repository.record_progress( + conn, + compute_job_id, + simulation_id, + "Processing tsunami", + {"step": "tsunami", "calculation": calculation}, + 1, + ) + # A second write that mentions neither column must leave both alone. + await repository.record_progress( + conn, + compute_job_id, + simulation_id, + "Processing maxola", + {"step": "maxola"}, + 1, + ) -def test_list_abandoned_work_dirs_returns_only_old_failed_jobs( - database: str, + status = await repository.get_job_status(simulation_id_text) + assert status["calculation"] == calculation + assert status["travel_times"] is None + assert status["step"] == "maxola" + + +@pytest.mark.asyncio +async def test_a_notification_is_only_delivered_once_its_update_commits( + queue: Queue, ) -> None: - database_url = database - simulation_id_text, compute_job_id = _create_job(database) + simulation_id_text, compute_job_id = await _create_job() simulation_id = uuid.UUID(simulation_id_text) - fresh_simulation_id_text, fresh_job_id = _create_job(database) + seen: list[str] = [] + + listener = await db.connect() + try: + await listener.add_listener( + db.notify_channel(simulation_id), + lambda *_args: seen.append("notified"), + ) + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + # asyncpg dispatches the notification on the connection's own reads. + for _ in range(50): + if seen: + break + await asyncio.sleep(0.02) + assert seen, "mark_started must wake a listening SSE stream" + # The row the client is about to re-read is already committed. + assert (await repository.get_job_status(simulation_id_text))[ + "status" + ] == "running" + finally: + await listener.close() + + +@pytest.mark.asyncio +async def test_list_abandoned_work_dirs_returns_only_old_failed_jobs( + queue: Queue, +) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + fresh_simulation_id_text, fresh_job_id = await _create_job() fresh_simulation_id = uuid.UUID(fresh_simulation_id_text) cutoff = datetime.now(UTC) - timedelta(hours=24) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - repository.fail_job(conn, compute_job_id, simulation_id, "worker vanished") - repository.fail_job(conn, fresh_job_id, fresh_simulation_id, "worker vanished") - conn.execute( - "UPDATE compute.jobs SET finished_at = %s WHERE id = %s", - [cutoff - timedelta(minutes=1), compute_job_id], + async with db.acquire() as conn: + for job_id, external_id in ( + (compute_job_id, simulation_id), + (fresh_job_id, fresh_simulation_id), + ): + await repository.mark_started(conn, job_id, external_id, 1) + await repository.record_failure( + conn, + job_id, + external_id, + RuntimeError("worker vanished"), + step="tsunami", + will_retry=False, + attempt=1, + ) + await conn.execute( + "UPDATE compute.jobs SET finished_at = $1 WHERE id = $2", + cutoff - timedelta(minutes=1), + compute_job_id, ) - abandoned = repository.list_abandoned_work_dirs(cutoff) + abandoned = await repository.list_abandoned_work_dirs(cutoff) assert simulation_id_text in abandoned assert fresh_simulation_id_text not in abandoned -def test_record_failure_keeps_running_for_a_transient_retry( - database: str, +@pytest.mark.asyncio +async def test_record_failure_keeps_running_for_a_transient_retry( + queue: Queue, ) -> None: - database_url = database - simulation_id_text, compute_job_id = _create_job(database) + simulation_id_text, compute_job_id = await _create_job() simulation_id = uuid.UUID(simulation_id_text) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - repository.mark_started(conn, compute_job_id, simulation_id) - repository.record_failure( + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await repository.record_failure( conn, compute_job_id, simulation_id, TransientInfraError("minio down"), step="maxola", will_retry=True, + attempt=1, ) - status = repository.get_job_status(simulation_id_text) + status = await repository.get_job_status(simulation_id_text) assert status["status"] == "running" assert status["details"] == "Retrying after transient error (TransientInfraError)" assert status["finished_at"] is None -def test_record_failure_persists_a_sanitized_terminal_error( - database: str, +@pytest.mark.asyncio +async def test_record_failure_persists_a_sanitized_terminal_error( + queue: Queue, ) -> None: - database_url = database - simulation_id_text, compute_job_id = _create_job(database) + simulation_id_text, compute_job_id = await _create_job() simulation_id = uuid.UUID(simulation_id_text) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - repository.record_failure( + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await repository.record_failure( conn, compute_job_id, simulation_id, FileNotFoundError("/private/jobs/secret/tsunami"), step="tsunami", will_retry=False, + attempt=1, ) - status = repository.get_job_status(simulation_id_text) + status = await repository.get_job_status(simulation_id_text) assert status["status"] == "failed" assert status["details"] == "Pipeline failed - check error logs" assert status["error"] == "Simulation failed at step 'tsunami' (FileNotFoundError)" assert "/private/jobs" not in status["error"] -def test_complete_job_persists_the_uploaded_manifest_and_result( - database: str, - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - database_url = database - simulation_id_text, compute_job_id = _create_job(database) - output_path = tmp_path / "calculation.json" - output_path.write_text("{}", encoding="utf-8") +@pytest.mark.asyncio +async def test_is_database_connected_reports_a_closed_pool(queue: Queue) -> None: + assert await repository.is_database_connected() is True + await db.close_pool() + assert await repository.is_database_connected() is False + +def _result(tmp_path: Path, output_path: Path) -> SimulationResult: calculation = CalculationResponse( length=1.0, width=2.0, @@ -293,18 +388,29 @@ def test_complete_job_persists_the_uploaded_manifest_and_result( distances={"PORT": 100.0}, epicenter_info={"lat": "0.0"}, ) - outputs = SimulationOutputs( - root=tmp_path, - files=(OutputFile("calculation", output_path, "application/json"),), - ) - result = SimulationResult( + return SimulationResult( calculation=calculation, travel_times=travel_times, runtime=RuntimeContext( model_dir=tmp_path, model_version="test", capabilities={} ), - outputs=outputs, + outputs=SimulationOutputs( + root=tmp_path, + files=(OutputFile("calculation", output_path, "application/json"),), + ), ) + + +@pytest.mark.asyncio +async def test_complete_job_persists_the_uploaded_manifest_and_result( + queue: Queue, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + simulation_id_text, compute_job_id = await _create_job() + output_path = tmp_path / "calculation.json" + output_path.write_text("{}", encoding="utf-8") + result = _result(tmp_path, output_path) uploaded: dict[str, Any] = {} def fake_upload(**kwargs: Any) -> tuple[str, str]: @@ -313,19 +419,21 @@ def fake_upload(**kwargs: Any) -> tuple[str, str]: monkeypatch.setattr(output_store, "upload_simulation_result", fake_upload) - with psycopg.connect(database_url, row_factory=dict_row) as conn: - row = repository.fetch_by_id(conn, compute_job_id) + async with db.acquire() as conn: + simulation_id = uuid.UUID(simulation_id_text) + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + row = await repository.fetch_by_id(conn, compute_job_id) assert row is not None - repository.complete_job(conn, row, result) + await repository.complete_job(conn, row, result, 1) - status = repository.get_job_status(simulation_id_text) + status = await repository.get_job_status(simulation_id_text) assert status["status"] == "completed" - assert status["calculation"] == calculation.model_dump(mode="json") - assert status["travel_times"] == travel_times.model_dump(mode="json") + assert status["calculation"] == result.calculation.model_dump(mode="json") + assert status["travel_times"] == result.travel_times.model_dump(mode="json") assert status["outputs"] == ["calculation"] assert uploaded["simulation_id"] == simulation_id_text - assert repository.get_outputs(simulation_id_text) == [ + assert await repository.get_outputs(simulation_id_text) == [ { "name": "calculation", "key": f"simulations/{simulation_id_text}/outputs/calculation.json", @@ -333,3 +441,154 @@ def fake_upload(**kwargs: Any) -> tuple[str, str]: "content_type": "application/json", } ] + + +@pytest.mark.asyncio +async def test_a_backend_killed_mid_statement_reports_as_transient( + queue: Queue, +) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + + async with db.acquire() as conn, db.acquire() as killer: + backend_pid = await conn.fetchval("SELECT pg_backend_pid()") + await killer.execute("SELECT pg_terminate_backend($1)", backend_pid) + await asyncio.sleep(0.3) + + # A database restart under a running simulation lands here, not on + # acquire(). TRANSIENT_RETRY only retries TransientInfraError, so an + # untranslated asyncpg error would turn an outage into a dead job. + with pytest.raises(TransientInfraError, match="database unavailable"): + await repository.record_progress( + conn, compute_job_id, simulation_id, "Processing tsunami", {}, 1 + ) + + +@pytest.mark.asyncio +async def test_record_failure_never_overwrites_a_finished_job(queue: Queue) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + + async with db.acquire() as conn: + await conn.execute( + "UPDATE compute.jobs SET status = 'completed', " + "details = 'Simulation completed successfully' WHERE id = $1", + compute_job_id, + ) + # complete_job's UPDATE can commit and the connection still drop before + # the caller learns it did, so the caller reports a failure for a job + # that actually finished. Neither branch may act on that. + await repository.record_failure( + conn, + compute_job_id, + simulation_id, + TransientInfraError("output upload failed"), + step="copy_ttt_pdf", + will_retry=True, + attempt=1, + ) + await repository.record_failure( + conn, + compute_job_id, + simulation_id, + RuntimeError("bad epicenter"), + step="copy_ttt_pdf", + will_retry=False, + attempt=1, + ) + + status = await repository.get_job_status(simulation_id_text) + assert status["status"] == "completed" + assert status["details"] == "Simulation completed successfully" + assert status["error"] is None + + +@pytest.mark.asyncio +async def test_complete_job_reports_a_result_it_could_not_record( + queue: Queue, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + simulation_id_text, compute_job_id = await _create_job() + output_path = tmp_path / "calculation.json" + output_path.write_text("{}", encoding="utf-8") + result = _result(tmp_path, output_path) + monkeypatch.setattr( + output_store, + "upload_simulation_result", + lambda **_kwargs: ("results", "simulations/orphan/metadata.json"), + ) + + async with db.acquire() as conn: + simulation_id = uuid.UUID(simulation_id_text) + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + row = await repository.fetch_by_id(conn, compute_job_id) + assert row is not None + # A newer attempt owns the row by the time this one finishes. + await repository.mark_started(conn, compute_job_id, simulation_id, 2) + + with caplog.at_level("ERROR"): + recorded = await repository.complete_job(conn, row, result, 1) + + # The upload already happened, so a silent no-op here would leave objects + # in MinIO with nothing in compute.jobs pointing at them and no trace why. + assert recorded is False + assert "simulations/orphan/metadata.json" in caplog.text + assert (await repository.get_job_status(simulation_id_text))["status"] == "running" + + +@pytest.mark.asyncio +async def test_a_claim_cannot_reopen_a_completed_job(queue: Queue) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await conn.execute( + "UPDATE compute.jobs SET status = 'completed', finished_at = now() " + "WHERE id = $1", + compute_job_id, + ) + # A redelivered attempt. The refusal has to be part of the UPDATE: a + # separate SELECT before it leaves room for a job to complete in between. + assert not await repository.mark_started(conn, compute_job_id, simulation_id, 2) + + status = await repository.get_job_status(simulation_id_text) + assert status["status"] == "completed" + assert status["finished_at"] is not None + + +@pytest.mark.asyncio +async def test_a_reclaimed_job_does_not_keep_the_previous_failure( + queue: Queue, +) -> None: + simulation_id_text, compute_job_id = await _create_job() + simulation_id = uuid.UUID(simulation_id_text) + + async with db.acquire() as conn: + await repository.mark_started(conn, compute_job_id, simulation_id, 1) + await repository.record_failure( + conn, + compute_job_id, + simulation_id, + RuntimeError("bad epicenter"), + step="tsunami", + will_retry=False, + attempt=1, + ) + failed = await repository.get_job_status(simulation_id_text) + assert failed["error"] is not None + assert failed["finished_at"] is not None + + # An operator retry, or a later attempt, reclaims it. + assert await repository.mark_started(conn, compute_job_id, simulation_id, 2) + + # `running` alongside a previous attempt's error and finish time is not a + # state any client can read correctly -- status_from_row hands all three to + # the status endpoint and to SSE together. + running = await repository.get_job_status(simulation_id_text) + assert running["status"] == "running" + assert running["error"] is None + assert running["finished_at"] is None + assert running["started_at"] is not None diff --git a/packages/api/tests/test_tasks.py b/packages/api/tests/test_tasks.py index 057dbf7..ea58edd 100644 --- a/packages/api/tests/test_tasks.py +++ b/packages/api/tests/test_tasks.py @@ -1,20 +1,21 @@ """Worker lifecycle behavior at the repository and engine boundaries.""" +import asyncio +import threading import uuid -from collections.abc import Iterator -from contextlib import contextmanager from pathlib import Path -from types import SimpleNamespace from typing import Any import pytest -from procrastinate import exceptions as procrastinate_exceptions -from procrastinate.jobs import Job as ProcrastinateJob +from rqueue import PermanentFailure, TaskContext +from rqueue.testing import RecordingQueue +from api.core import queue as queue_module from api.core import tasks from api.core.errors import TransientInfraError -from api.core.tasks import MAX_ATTEMPTS -from tsdhn.domain import EarthquakeInput +from api.core.settings import COMPUTE_QUEUE +from api.core.tasks import MAX_ATTEMPTS, RUN_SIMULATION +from tsdhn.domain import EarthquakeInput, JobStatus tasks_module: Any = tasks @@ -22,73 +23,137 @@ class _Connection: - pass - - -class _TaskContext: - def __init__(self, attempts: int) -> None: - self.job = SimpleNamespace(attempts=attempts) + """Stands in for the caller's asyncpg connection; never used as one.""" + + +def _context(attempt: int) -> TaskContext: + async def heartbeat() -> bool: + return False + + return TaskContext( + job_id=uuid.uuid4(), + queue=COMPUTE_QUEUE, + task=RUN_SIMULATION, + attempt=attempt, + max_attempts=MAX_ATTEMPTS, + metadata={}, + heartbeat=heartbeat, + cancel_event=asyncio.Event(), + ) -def _row(job_id: uuid.UUID, simulation_id: uuid.UUID) -> dict[str, Any]: +def _row( + job_id: uuid.UUID, simulation_id: uuid.UUID, status: JobStatus = JobStatus.QUEUED +) -> dict[str, Any]: return { "id": job_id, "simulation_id": simulation_id, + "status": status.value, "input_params": INPUT.model_dump(mode="json"), } +@pytest.fixture +def recording_queue(monkeypatch: pytest.MonkeyPatch) -> RecordingQueue: + """Install a queue that runs the real validation and records the result.""" + queue = RecordingQueue(name=COMPUTE_QUEUE) + tasks.register_tasks(queue) + monkeypatch.setattr(queue_module, "_queue", queue) + return queue + + @pytest.fixture def worker_job( monkeypatch: pytest.MonkeyPatch, tmp_path: Path -) -> tuple[uuid.UUID, uuid.UUID, Path, _Connection]: +) -> tuple[uuid.UUID, uuid.UUID, Path]: job_id = uuid.uuid4() simulation_id = uuid.uuid4() work_dir = tmp_path / "jobs" / str(simulation_id) work_dir.mkdir(parents=True) connection = _Connection() - @contextmanager - def connect() -> Iterator[_Connection]: - yield connection + class _Acquire: + async def __aenter__(self) -> _Connection: + return connection + + async def __aexit__(self, *_exc: object) -> None: + return None - monkeypatch.setattr(tasks_module, "connect", connect) + monkeypatch.setattr(tasks_module.db, "acquire", lambda: _Acquire()) monkeypatch.setattr(tasks_module, "JOBS_DIR", tmp_path / "jobs") - monkeypatch.setattr( - tasks_module.repository, - "fetch_by_id", - lambda _conn, _id: _row(job_id, simulation_id), - ) - return job_id, simulation_id, work_dir, connection + async def fetch_by_id(_conn: Any, _id: uuid.UUID) -> dict[str, Any]: + return _row(job_id, simulation_id) + + monkeypatch.setattr(tasks_module.repository, "fetch_by_id", fetch_by_id) + return job_id, simulation_id, work_dir + + +@pytest.mark.asyncio +async def test_enqueue_simulation_builds_the_job_the_worker_expects( + recording_queue: RecordingQueue, +) -> None: + connection = _Connection() + compute_job_id = uuid.uuid4() + + await tasks.enqueue_simulation(connection, compute_job_id) -def test_run_simulation_task_completes_and_removes_the_workspace( - worker_job: tuple[uuid.UUID, uuid.UUID, Path, _Connection], + (recorded,) = recording_queue.enqueued(RUN_SIMULATION) + assert recorded.connection is connection + assert recorded.queue == COMPUTE_QUEUE + assert recorded.payload == {"compute_job_id": str(compute_job_id)} + # queueing_lock and lock, ported one for one onto rqueue's two keys. + assert recorded.dedupe_key == f"simulation:{compute_job_id}" + assert recorded.concurrency_key == f"compute-job:{compute_job_id}" + # The registered retry policy, not the queue default, reaches the row. + assert recorded.spec.max_attempts == MAX_ATTEMPTS + + +@pytest.mark.asyncio +async def test_enqueue_simulation_fails_when_the_queue_was_never_built( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(queue_module, "_queue", None) + + with pytest.raises(RuntimeError, match="build_queue"): + await tasks.enqueue_simulation(_Connection(), uuid.uuid4()) + + +@pytest.mark.asyncio +async def test_run_simulation_task_completes_and_removes_the_workspace( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], monkeypatch: pytest.MonkeyPatch, ) -> None: - job_id, simulation_id, work_dir, connection = worker_job + job_id, simulation_id, work_dir = worker_job events: list[tuple[str, Any]] = [] result = object() - monkeypatch.setattr( - tasks_module.repository, - "mark_started", - lambda conn, current_id, current_external: events.append( - ("started", (conn, current_id, current_external)) - ), - ) - monkeypatch.setattr( - tasks_module.repository, - "record_progress", - lambda conn, current_id, current_external, message, details: events.append( - ("progress", (conn, current_id, current_external, message, details)) - ), - ) - monkeypatch.setattr( - tasks_module.repository, - "complete_job", - lambda conn, row, completed: events.append(("completed", (conn, completed))), - ) + async def mark_started( + _conn: Any, current_id: Any, current_external: Any, attempt: int + ) -> bool: + events.append(("started", (current_id, current_external, attempt))) + return True + + async def record_progress( + _conn: Any, + current_id: Any, + current_external: Any, + message: str, + details: dict[str, Any], + attempt: int, + ) -> bool: + events.append(("progress", (current_id, current_external, message, details))) + return True + + async def complete_job( + _conn: Any, _row: Any, completed: Any, _attempt: int + ) -> bool: + events.append(("completed", completed)) + return True + + monkeypatch.setattr(tasks_module.repository, "mark_started", mark_started) + monkeypatch.setattr(tasks_module.repository, "record_progress", record_progress) + monkeypatch.setattr(tasks_module.repository, "complete_job", complete_job) def run_simulation( data: EarthquakeInput, @@ -103,53 +168,85 @@ def run_simulation( monkeypatch.setattr(tasks_module, "run_simulation", run_simulation) - tasks_module.run_simulation_task.func(_TaskContext(attempts=0), str(job_id)) + await tasks.run_simulation_task(job_id, _context(attempt=1)) assert events == [ - ("started", (connection, job_id, simulation_id)), + ("started", (job_id, simulation_id, 1)), ("run", (INPUT, work_dir, True)), ( "progress", - ( - connection, - job_id, - simulation_id, - "Processing tsunami", - {"step": "tsunami"}, - ), + (job_id, simulation_id, "Processing tsunami", {"step": "tsunami"}), ), - ("completed", (connection, result)), + ("completed", result), ] assert not work_dir.exists() -@pytest.mark.parametrize(("attempts", "will_retry"), [(0, True), (MAX_ATTEMPTS, False)]) -def test_run_simulation_task_records_failure_and_preserves_workspace( - worker_job: tuple[uuid.UUID, uuid.UUID, Path, _Connection], +@pytest.mark.asyncio +async def test_the_simulation_runs_off_the_event_loop( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], + monkeypatch: pytest.MonkeyPatch, +) -> None: + job_id, _simulation_id, _work_dir = worker_job + threads: list[int] = [] + + async def claimed(*_args: Any, **_kwargs: Any) -> bool: + """mark_started and friends: the claim always succeeds in these tests.""" + return True + + monkeypatch.setattr(tasks_module.repository, "mark_started", claimed) + monkeypatch.setattr(tasks_module.repository, "record_progress", claimed) + monkeypatch.setattr(tasks_module.repository, "complete_job", claimed) + + def run_simulation(*_args: Any, on_progress: Any, **_kwargs: Any) -> object: + import threading + + threads.append(threading.get_ident()) + # The callback runs on this thread and has to reach the loop's pool. + on_progress("Processing tsunami", {"step": "tsunami"}) + return object() + + monkeypatch.setattr(tasks_module, "run_simulation", run_simulation) + + import threading + + await tasks.run_simulation_task(job_id, _context(attempt=1)) + + # A blocking kernel on the loop thread would stall rqueue's heartbeat and + # cost the worker its lease mid-run. + assert threads and threads[0] != threading.get_ident() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("attempt", "will_retry"), [(1, True), (MAX_ATTEMPTS, False)]) +async def test_run_simulation_task_records_failure_and_preserves_workspace( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], monkeypatch: pytest.MonkeyPatch, - attempts: int, + attempt: int, will_retry: bool, ) -> None: - job_id, simulation_id, work_dir, connection = worker_job + job_id, simulation_id, work_dir = worker_job failures: list[dict[str, Any]] = [] - monkeypatch.setattr(tasks_module.repository, "mark_started", lambda *_args: None) - monkeypatch.setattr( - tasks_module.repository, "get_current_step", lambda _conn, _job_id: "tsunami" - ) + async def claimed(*_args: Any, **_kwargs: Any) -> bool: + """mark_started and friends: the claim always succeeds in these tests.""" + return True - def record_failure( - conn: Any, + async def get_current_step(_conn: Any, _job_id: uuid.UUID) -> str: + return "tsunami" + + async def record_failure( + _conn: Any, current_id: Any, current_external: Any, exc: Exception, *, step: str | None, will_retry: bool, + attempt: int, ) -> None: failures.append( { - "conn": conn, "job_id": current_id, "simulation_id": current_external, "exception": exc, @@ -158,11 +255,9 @@ def record_failure( } ) - monkeypatch.setattr( - tasks_module.repository, - "record_failure", - record_failure, - ) + monkeypatch.setattr(tasks_module.repository, "mark_started", claimed) + monkeypatch.setattr(tasks_module.repository, "get_current_step", get_current_step) + monkeypatch.setattr(tasks_module.repository, "record_failure", record_failure) def fail(*_args: Any, **_kwargs: Any) -> object: raise TransientInfraError("storage unavailable") @@ -170,10 +265,9 @@ def fail(*_args: Any, **_kwargs: Any) -> object: monkeypatch.setattr(tasks_module, "run_simulation", fail) with pytest.raises(TransientInfraError, match="storage unavailable"): - tasks_module.run_simulation_task.func(_TaskContext(attempts), str(job_id)) + await tasks.run_simulation_task(job_id, _context(attempt)) assert len(failures) == 1 - assert failures[0]["conn"] is connection assert failures[0]["job_id"] == job_id assert failures[0]["simulation_id"] == simulation_id assert isinstance(failures[0]["exception"], TransientInfraError) @@ -182,258 +276,348 @@ def fail(*_args: Any, **_kwargs: Any) -> object: assert work_dir.exists() -def test_run_simulation_task_rejects_an_unknown_job( +@pytest.mark.asyncio +async def test_a_failed_failure_write_does_not_replace_the_original_error( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], monkeypatch: pytest.MonkeyPatch, ) -> None: - connection = _Connection() - - @contextmanager - def connect() -> Iterator[_Connection]: - yield connection - - monkeypatch.setattr(tasks_module, "connect", connect) - monkeypatch.setattr(tasks_module.repository, "fetch_by_id", lambda *_args: None) - job_id = uuid.uuid4() - - with pytest.raises(RuntimeError, match="Unknown compute job"): - tasks_module.run_simulation_task.func(_TaskContext(0), str(job_id)) + job_id, _simulation_id, _work_dir = worker_job + async def claimed(*_args: Any, **_kwargs: Any) -> bool: + """mark_started and friends: the claim always succeeds in these tests.""" + return True -def test_enqueue_simulation_configures_and_defers_the_worker_job( - monkeypatch: pytest.MonkeyPatch, -) -> None: - connection = _Connection() - compute_job_id = uuid.uuid4() - configured: list[dict[str, Any]] = [] - deferred: list[dict[str, Any]] = [] + async def unavailable(*_args: Any, **_kwargs: Any) -> None: + raise TransientInfraError("database connection failed") - class _ConfiguredTask: - def configure(self, **kwargs: Any) -> Any: - configured.append(kwargs) - return self + monkeypatch.setattr(tasks_module.repository, "mark_started", claimed) + monkeypatch.setattr(tasks_module.repository, "get_current_step", unavailable) - def defer(self, **kwargs: Any) -> None: - deferred.append(kwargs) + def fail(*_args: Any, **_kwargs: Any) -> object: + raise RuntimeError("bad epicenter") - monkeypatch.setattr(tasks_module, "run_simulation_task", _ConfiguredTask()) + monkeypatch.setattr(tasks_module, "run_simulation", fail) - tasks_module.enqueue_simulation(connection, compute_job_id) + # The pipeline error is what rqueue must see, so it applies retry_on to + # the real cause rather than to a second outage reporting it. + with pytest.raises(RuntimeError, match="bad epicenter"): + await tasks.run_simulation_task(job_id, _context(attempt=1)) - assert configured == [ - { - "connection": connection, - "queue": tasks_module.PROCRASTINATE_QUEUE, - "queueing_lock": f"simulation:{compute_job_id}", - "lock": f"compute-job:{compute_job_id}", - } - ] - assert deferred == [{"compute_job_id": str(compute_job_id)}] +@pytest.mark.asyncio +async def test_run_simulation_task_rejects_an_unknown_job( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Acquire: + async def __aenter__(self) -> _Connection: + return _Connection() -class _JobManager: - def __init__(self, jobs: list[ProcrastinateJob]) -> None: - self.jobs = jobs - self.retried: list[int] = [] - self.finished: list[int] = [] + async def __aexit__(self, *_exc: object) -> None: + return None - async def get_stalled_jobs( - self, *, seconds_since_heartbeat: int - ) -> list[ProcrastinateJob]: - assert seconds_since_heartbeat == tasks_module.STALLED_HEARTBEAT_SECONDS - return self.jobs + async def missing(*_args: Any) -> None: + return None - async def retry_job_by_id_async(self, *, job_id: int, retry_at: Any) -> None: - self.retried.append(job_id) + monkeypatch.setattr(tasks_module.db, "acquire", lambda: _Acquire()) + monkeypatch.setattr(tasks_module.repository, "fetch_by_id", missing) - async def finish_job_by_id_async( - self, *, job_id: int, status: Any, delete_job: bool - ) -> None: - assert status.value == "failed" - assert delete_job is False - self.finished.append(job_id) - - -def _queue_job(job_id: int | None, attempts: int) -> ProcrastinateJob: - return ProcrastinateJob( - id=job_id, - queue="simulations", - lock=None, - queueing_lock=None, - task_name="api.run_simulation", - task_kwargs={"compute_job_id": str(uuid.uuid4())}, - attempts=attempts, - ) + # No later attempt can find a row that was never written. + with pytest.raises(PermanentFailure, match="Unknown compute job"): + await tasks.run_simulation_task(uuid.uuid4(), _context(attempt=1)) @pytest.mark.asyncio -async def test_reap_stalled_jobs_retries_or_finishes_by_attempt_budget( - monkeypatch: pytest.MonkeyPatch, +async def test_sweep_abandoned_work_dirs_removes_only_selected_workspaces( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: - retry_job = _queue_job(1, 0) - exhausted_job = _queue_job(2, MAX_ATTEMPTS) - manager = _JobManager([retry_job, exhausted_job]) - exhausted_ids: list[str] = [] + old_id = str(uuid.uuid4()) + old_dir = tmp_path / old_id + old_dir.mkdir() + unrelated = tmp_path / "keep" + unrelated.mkdir() + + async def list_abandoned_work_dirs(_cutoff: Any) -> list[str]: + return [old_id] - monkeypatch.setattr(tasks_module.app, "job_manager", manager) monkeypatch.setattr( - tasks_module, - "_fail_exhausted", - lambda compute_job_id: exhausted_ids.append(compute_job_id), + tasks_module.repository, + "list_abandoned_work_dirs", + list_abandoned_work_dirs, ) + monkeypatch.setattr(tasks_module, "JOBS_DIR", tmp_path) - await tasks_module.reap_stalled_jobs_task.func(0) + await tasks.sweep_abandoned_work_dirs() - assert manager.retried == [1] - assert manager.finished == [2] - assert exhausted_ids == [exhausted_job.task_kwargs["compute_job_id"]] + assert not old_dir.exists() + assert unrelated.exists() @pytest.mark.asyncio -async def test_reap_stalled_jobs_returns_when_the_queue_is_empty( +async def test_the_periodic_sweep_keeps_running_after_a_failed_pass( monkeypatch: pytest.MonkeyPatch, ) -> None: - manager = _JobManager([]) - monkeypatch.setattr(tasks_module.app, "job_manager", manager) + passes: list[int] = [] + stop = asyncio.Event() + + async def sweep() -> None: + passes.append(len(passes)) + if len(passes) == 1: + raise RuntimeError("database unavailable") + stop.set() - await tasks_module.reap_stalled_jobs_task.func(0) + monkeypatch.setattr(tasks_module, "sweep_abandoned_work_dirs", sweep) - assert manager.retried == [] - assert manager.finished == [] + await tasks.run_periodic_sweep(stop, interval=0.01) + + assert len(passes) == 2 @pytest.mark.asyncio -async def test_reap_stalled_jobs_ignores_malformed_queue_entries( +async def test_the_periodic_reconcile_keeps_running_after_a_failed_pass( monkeypatch: pytest.MonkeyPatch, ) -> None: - missing_id = _queue_job(None, 0) - missing_compute_id = ProcrastinateJob( - id=3, - queue="simulations", - lock=None, - queueing_lock=None, - task_name="api.run_simulation", - task_kwargs={}, - attempts=0, - ) - manager = _JobManager([missing_id, missing_compute_id]) - monkeypatch.setattr(tasks_module.app, "job_manager", manager) + passes: list[int] = [] + stop = asyncio.Event() - await tasks_module.reap_stalled_jobs_task.func(0) + async def reconcile(**_kwargs: Any) -> None: + passes.append(len(passes)) + if len(passes) == 1: + raise TransientInfraError("database connection failed") + stop.set() - assert manager.retried == [] - assert manager.finished == [] + monkeypatch.setattr(tasks_module, "reconcile_terminal_jobs", reconcile) + + # A database outage is exactly when jobs get stranded, so it must not be + # the thing that stops the pass which unstrands them. + await tasks.run_periodic_reconcile(stop, interval=0.01) + + assert len(passes) == 2 @pytest.mark.asyncio -async def test_reap_stalled_jobs_tolerates_queue_connector_races( +async def test_an_abandoned_attempt_stops_writing_from_its_kernel_thread( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], monkeypatch: pytest.MonkeyPatch, ) -> None: - retry_job = _queue_job(1, 0) - exhausted_job = _queue_job(2, MAX_ATTEMPTS) - manager = _JobManager([retry_job, exhausted_job]) - retry_attempts: list[int] = [] - finish_attempts: list[int] = [] - exhausted_ids: list[str] = [] - - async def retry_failure(*, job_id: int, retry_at: Any) -> None: - retry_attempts.append(job_id) - raise procrastinate_exceptions.ConnectorException("job already resolved") - - async def finish_failure(*, job_id: int, status: Any, delete_job: bool) -> None: - finish_attempts.append(job_id) - raise procrastinate_exceptions.ConnectorException("job already resolved") + job_id, _simulation_id, _work_dir = worker_job + messages: list[str] = [] + entered = threading.Event() + release = threading.Event() + finished = threading.Event() + raised: list[BaseException] = [] + + async def claimed(*_args: Any, **_kwargs: Any) -> bool: + """mark_started and friends: the claim always succeeds in these tests.""" + return True + + async def record_progress( + _conn: Any, _id: Any, _external: Any, message: str, _details: Any, _attempt: int + ) -> bool: + messages.append(message) + return True + + monkeypatch.setattr(tasks_module.repository, "mark_started", claimed) + monkeypatch.setattr(tasks_module.repository, "record_progress", record_progress) + + def run_simulation(*_args: Any, on_progress: Any, **_kwargs: Any) -> object: + try: + on_progress("Processing tsunami", {}) + entered.set() + release.wait(5) + on_progress("Processing maxola", {}) + return object() + except BaseException as e: + raised.append(e) + raise + finally: + finished.set() - monkeypatch.setattr(tasks_module.app, "job_manager", manager) - monkeypatch.setattr(tasks_module, "_fail_exhausted", exhausted_ids.append) - monkeypatch.setattr(manager, "retry_job_by_id_async", retry_failure) - monkeypatch.setattr(manager, "finish_job_by_id_async", finish_failure) - - await tasks_module.reap_stalled_jobs_task.func(0) + monkeypatch.setattr(tasks_module, "run_simulation", run_simulation) - assert retry_attempts == [1] - assert exhausted_ids == [exhausted_job.task_kwargs["compute_job_id"]] - assert finish_attempts == [2] + task = asyncio.create_task(tasks.run_simulation_task(job_id, _context(attempt=1))) + await asyncio.to_thread(entered.wait, 5) + + # rqueue cancels the handler coroutine when the heartbeat finds the lease + # gone. Python cannot kill the thread underneath it, and that thread still + # holds a live reference to this loop. + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + release.set() + await asyncio.to_thread(finished.wait, 5) + + # The write from before the cancellation stands; the one after is refused, + # and the refusal unwinds the kernel instead of letting it run on. + assert messages == ["Processing tsunami"] + assert raised and isinstance(raised[0], tasks.AbandonedAttempt) + + +def test_a_workspace_held_by_a_live_thread_is_not_resumed(tmp_path: Path) -> None: + """The filesystem half of the fence, which the database fence cannot cover. + + A kernel thread abandoned mid-step keeps writing into a workspace keyed by + simulation id. A replacement attempt resuming from those checkpoints could + read a half-written one, which is worse than a confusing status: it yields a + wrong scientific result rather than an error. + """ + work_dir = tmp_path / "sim" + work_dir.mkdir() + holding = threading.Event() + release = threading.Event() + + def zombie() -> None: + claim, _resume = tasks.claim_workspace(work_dir, 1) + holding.set() + release.wait(5) + claim.release() + claim.release() + + thread = threading.Thread(target=zombie) + thread.start() + assert holding.wait(5) + + # flock is held against the open file description, so the replacement is + # refused even though it is another thread of this same process. + with pytest.raises(TransientInfraError, match="still held by an earlier attempt"): + tasks.claim_workspace(work_dir, 2) + + release.set() + thread.join(5) + + # Once the owner really is gone the workspace is resumable again, which is + # what an ordinary worker crash looks like: the kernel drops the lock with + # the process. + claim, resume = tasks.claim_workspace(work_dir, 2) + assert resume is True + claim.release() + claim.release() + + +def test_a_claim_survives_until_its_last_share_is_given_back(tmp_path: Path) -> None: + """The kernel thread and the upload hold the same claim, one share each. + + complete_job reads the result files back out of the workspace after the + kernel thread has returned. If the claim ended with the kernel, a lease lost + in between would let a redelivered attempt take the freed lock and write + into the directory being uploaded from -- the same corruption, a few lines + later. + """ + work_dir = tmp_path / "sim" + work_dir.mkdir() + claim, _resume = tasks.claim_workspace(work_dir, 1) + + # The kernel thread finishes; the upload has not started. + claim.release() + with pytest.raises(TransientInfraError, match="still held by an earlier attempt"): + tasks.claim_workspace(work_dir, 2) + + # The upload finishes too. + claim.release() + second, _ = tasks.claim_workspace(work_dir, 2) + second.release() + second.release() + + +def test_releasing_a_claim_more_than_twice_is_harmless(tmp_path: Path) -> None: + work_dir = tmp_path / "sim" + work_dir.mkdir() + claim, _resume = tasks.claim_workspace(work_dir, 1) + for _ in range(4): + claim.release() -def test_sweep_abandoned_work_dirs_removes_only_selected_workspaces( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path +@pytest.mark.asyncio +async def test_a_cancelled_attempt_keeps_holding_its_workspace( + worker_job: tuple[uuid.UUID, uuid.UUID, Path], + monkeypatch: pytest.MonkeyPatch, ) -> None: - old_id = str(uuid.uuid4()) - old_dir = tmp_path / old_id - old_dir.mkdir() - unrelated = tmp_path / "keep" - unrelated.mkdir() - monkeypatch.setattr( - tasks_module.repository, "list_abandoned_work_dirs", lambda _cutoff: [old_id] - ) - monkeypatch.setattr(tasks_module, "JOBS_DIR", tmp_path) + """The claim is tied to the thread, not to the coroutine. - tasks_module.sweep_abandoned_work_dirs_task.func(0) + Releasing it when the coroutine is cancelled would hand the workspace to a + replacement attempt while the thread it belongs to is still writing there. + """ + job_id, _simulation_id, work_dir = worker_job + entered = threading.Event() + release = threading.Event() - assert not old_dir.exists() - assert unrelated.exists() + async def claimed(*_args: Any, **_kwargs: Any) -> bool: + return True + monkeypatch.setattr(tasks_module.repository, "mark_started", claimed) -@pytest.mark.parametrize("status", ["completed", "failed"]) -def test_fail_exhausted_does_not_overwrite_terminal_jobs( - monkeypatch: pytest.MonkeyPatch, status: str -) -> None: - connection = _Connection() + def hang(*_args: Any, **_kwargs: Any) -> object: + entered.set() + release.wait(5) + return object() - @contextmanager - def pooled() -> Iterator[_Connection]: - yield connection + monkeypatch.setattr(tasks_module, "run_simulation", hang) - monkeypatch.setattr(tasks_module, "pooled", pooled) - fetched: list[tuple[Any, ...]] = [] + task = asyncio.create_task(tasks.run_simulation_task(job_id, _context(attempt=1))) + await asyncio.to_thread(entered.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task - def fetch(*args: Any) -> dict[str, Any]: - fetched.append(args) - return {"status": status, "simulation_id": uuid.uuid4()} + # The coroutine is gone, and gave back its own share; the thread is not, and + # still holds its own. + with pytest.raises(TransientInfraError, match="still held by an earlier attempt"): + tasks.claim_workspace(work_dir, 2) - monkeypatch.setattr(tasks_module.repository, "fetch_by_id", fetch) - monkeypatch.setattr( - tasks_module.repository, - "fail_job", - lambda *_args: pytest.fail("terminal job must not be overwritten"), - ) + release.set() - tasks_module._fail_exhausted(str(uuid.uuid4())) - assert len(fetched) == 1 +def test_a_fresh_workspace_reports_that_there_is_nothing_to_resume( + tmp_path: Path, +) -> None: + claim, resume = tasks.claim_workspace(tmp_path / "unstarted", 1) + assert resume is False + claim.release() + claim.release() + # The lock lives beside the workspace, not inside it: the engine removes the + # whole directory when a run starts without resuming. + assert (tmp_path / "unstarted.lock").exists() + assert not (tmp_path / "unstarted").exists() -def test_fail_exhausted_marks_an_active_job_failed( - monkeypatch: pytest.MonkeyPatch, -) -> None: - connection = _Connection() - simulation_id = uuid.uuid4() - failures: list[tuple[Any, ...]] = [] +def test_removing_a_workspace_takes_its_lock_file_with_it(tmp_path: Path) -> None: + work_dir = tmp_path / "sim" + work_dir.mkdir() + claim, _resume = tasks.claim_workspace(work_dir, 1) + claim.release() + claim.release() - @contextmanager - def pooled() -> Iterator[_Connection]: - yield connection + assert tasks.remove_workspace(work_dir) is True - monkeypatch.setattr(tasks_module, "pooled", pooled) - monkeypatch.setattr( - tasks_module.repository, - "fetch_by_id", - lambda *_args: {"status": "running", "simulation_id": simulation_id}, - ) - monkeypatch.setattr( - tasks_module.repository, - "fail_job", - lambda *args: failures.append(args), - ) + assert not work_dir.exists() + assert not (tmp_path / "sim.lock").exists() + # Idempotent: the sweep runs over jobs it may already have cleaned. + assert tasks.remove_workspace(work_dir) is True + + +def test_a_held_workspace_is_left_for_the_next_sweep(tmp_path: Path) -> None: + """The flock+unlink race, which is why removal takes the lock first. + + unlink() drops the directory entry while the holder keeps its lock on the + now-orphaned inode, so the next O_CREAT at that path makes a *new* inode + and locks it uncontended -- two threads each believing they own the + workspace. The sweep reaches jobs whose row is terminal while their kernel + thread is still alive, so this is reachable. + """ + work_dir = tmp_path / "sim" + work_dir.mkdir() + (work_dir / "checkpoint").write_text("half written", encoding="utf-8") + claim, _resume = tasks.claim_workspace(work_dir, 1) + + assert tasks.remove_workspace(work_dir) is False + assert work_dir.exists() + assert (tmp_path / "sim.lock").exists() - job_id = uuid.uuid4() - tasks_module._fail_exhausted(str(job_id)) + # And the holder's claim still means something afterwards. + with pytest.raises(TransientInfraError, match="still held by an earlier attempt"): + tasks.claim_workspace(work_dir, 2) - assert failures == [ - ( - connection, - job_id, - simulation_id, - tasks_module.CRASH_BUDGET_EXHAUSTED_ERROR, - ) - ] + claim.release() + claim.release() + assert tasks.remove_workspace(work_dir) is True + assert not work_dir.exists() diff --git a/packages/api/tests/test_worker_integration.py b/packages/api/tests/test_worker_integration.py new file mode 100644 index 0000000..ae41bd7 --- /dev/null +++ b/packages/api/tests/test_worker_integration.py @@ -0,0 +1,679 @@ +"""A real rqueue.Worker running the simulation task against PostgreSQL.""" + +import asyncio +import threading +import uuid +from datetime import datetime, timedelta +from pathlib import Path +from typing import Any + +import asyncpg +import pytest +from rqueue import Admin, Queue, Worker + +from api.core import db, repository, tasks +from api.core.errors import TransientInfraError +from api.core.settings import COMPUTE_QUEUE_SCHEMA +from api.core.storage import output_store +from api.core.tasks import RUN_SIMULATION, enqueue_simulation +from tsdhn.domain import CalculationResponse, EarthquakeInput, TsunamiTravelResponse +from tsdhn.engine import OutputFile, SimulationOutputs, SimulationResult +from tsdhn.runtime import RuntimeContext + +pytestmark = pytest.mark.integration + +db_module: Any = db + +INPUT = EarthquakeInput(Mw=8.0, h=10.0, lat0=-20.5, lon0=-70.5, hhmm="0000", dia="23") + + +def _result(root: Path) -> SimulationResult: + output_path = root / "calculation.json" + output_path.write_text("{}", encoding="utf-8") + return SimulationResult( + calculation=CalculationResponse( + length=1.0, + width=2.0, + dislocation=3.0, + seismic_moment=4.0, + tsunami_warning="none", + distance_to_coast=5.0, + azimuth=6.0, + dip=7.0, + epicenter_location="0.00/0.00", + rectangle_parameters={}, + rectangle_corners=[], + ), + travel_times=TsunamiTravelResponse( + arrival_times={"PORT": "01:00"}, + distances={"PORT": 100.0}, + epicenter_info={"lat": "0.0"}, + ), + runtime=RuntimeContext(model_dir=root, model_version="test", capabilities={}), + outputs=SimulationOutputs( + root=root, + files=(OutputFile("calculation", output_path, "application/json"),), + ), + ) + + +async def _submit(simulation_id: str) -> uuid.UUID: + await repository.create_or_get_job( + data=INPUT, simulation_id=simulation_id, defer=enqueue_simulation + ) + async with db.acquire() as conn: + compute_job_id: uuid.UUID = await conn.fetchval( + "SELECT id FROM compute.jobs WHERE simulation_id = $1", + uuid.UUID(simulation_id), + ) + return compute_job_id + + +async def _queue_row(compute_job_id: uuid.UUID) -> Any: + async with db.acquire() as conn: + return await conn.fetchrow( + f"SELECT * FROM {COMPUTE_QUEUE_SCHEMA}.jobs " # noqa: S608 + "WHERE payload->>'compute_job_id' = $1", + str(compute_job_id), + ) + + +def _worker(queue: Queue) -> Worker: + return Worker(queue, worker_id="test-worker", concurrency=2, lease_duration=30.0) + + +@pytest.mark.asyncio +async def test_a_queued_simulation_runs_to_completion_on_a_bounded_thread( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + (tmp_path / "jobs" / simulation_id).mkdir(parents=True) + kernel_threads: list[str] = [] + + def run_simulation( + _data: EarthquakeInput, + work_dir: Path, + *, + resume: bool, + on_progress: Any, + ) -> SimulationResult: + kernel_threads.append(threading.current_thread().name) + on_progress( + "Processing tsunami", + {"step": "tsunami", "step_index": 3, "total_steps": 8}, + ) + return _result(work_dir) + + monkeypatch.setattr(tasks, "run_simulation", run_simulation) + monkeypatch.setattr( + output_store, + "upload_simulation_result", + lambda **_kwargs: ("results", "simulations/x/metadata.json"), + ) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + + status = await repository.get_job_status(simulation_id) + assert status["status"] == "completed" + assert status["outputs"] == ["calculation"] + + row = await _queue_row(compute_job_id) + assert row["state"] == "succeeded" + assert row["attempt"] == 1 + assert row["error_type"] is None + + # rqueue installs its own bounded default executor, so a bare + # asyncio.to_thread lands there rather than on the interpreter's default. + assert kernel_threads + assert kernel_threads[0].startswith("rqueue-test-worker") + assert kernel_threads[0] != threading.current_thread().name + + # The workspace is removed once the result is durable. + assert not (tmp_path / "jobs" / simulation_id).exists() + + +@pytest.mark.asyncio +async def test_a_pipeline_error_fails_the_job_after_one_attempt( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + def fail(*_args: Any, **_kwargs: Any) -> SimulationResult: + raise RuntimeError("bad epicenter") + + monkeypatch.setattr(tasks, "run_simulation", fail) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + # retry_on is an allowlist of transient infrastructure errors only. + assert row["state"] == "failed" + assert row["attempt"] == 1 + assert row["error_type"] == "RuntimeError" + + status = await repository.get_job_status(simulation_id) + assert status["status"] == "failed" + assert status["error"] == "Simulation failed (RuntimeError)" + + +@pytest.mark.asyncio +async def test_a_transient_error_is_retried_within_its_budget( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + from api.core.errors import TransientInfraError + + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + def fail(*_args: Any, **_kwargs: Any) -> SimulationResult: + raise TransientInfraError("storage unavailable") + + monkeypatch.setattr(tasks, "run_simulation", fail) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + # Rescheduled, not terminal: the backoff puts the next attempt in the + # future, so drain() returns with the job pending rather than failed. + assert row["state"] == "pending" + assert row["attempt"] == 1 + assert row["error_type"] == "TransientInfraError" + + status = await repository.get_job_status(simulation_id) + assert status["status"] == "running" + assert status["details"] == "Retrying after transient error (TransientInfraError)" + + +@pytest.mark.asyncio +async def test_a_job_whose_compute_row_vanished_fails_permanently( + queue: Queue, tmp_path: Path +) -> None: + compute_job_id = uuid.uuid4() + async with db.acquire() as conn: + await tasks.enqueue_simulation(conn, compute_job_id) + + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + assert row["state"] == "failed" + assert row["attempt"] == 1 + assert row["error_type"] == "PermanentFailure" + + +@pytest.mark.asyncio +async def test_an_undecodable_payload_fails_before_the_handler_runs( + queue: Queue, monkeypatch: pytest.MonkeyPatch +) -> None: + def unreached(*_args: Any, **_kwargs: Any) -> SimulationResult: + pytest.fail("the handler must not run for a payload that cannot decode") + + monkeypatch.setattr(tasks, "run_simulation", unreached) + + async with db.acquire() as conn: + await queue.enqueue( + conn, task=RUN_SIMULATION, payload={"compute_job_id": "not-a-uuid"} + ) + + await _worker(queue).drain(timeout=60) + + async with db.acquire() as conn: + row = await conn.fetchrow( + f"SELECT state, attempt, error_type FROM {COMPUTE_QUEUE_SCHEMA}.jobs" # noqa: S608 + ) + assert row["state"] == "failed" + assert row["error_type"] == "PayloadDecodeError" + + +@pytest.mark.asyncio +async def test_the_queue_tables_stay_out_of_the_compute_schema(queue: Queue) -> None: + async with db.acquire() as conn: + compute_tables = { + record["tablename"] + for record in await conn.fetch( + "SELECT tablename FROM pg_tables WHERE schemaname = 'compute'" + ) + } + queue_tables = { + record["tablename"] + for record in await conn.fetch( + "SELECT tablename FROM pg_tables WHERE schemaname = $1", + COMPUTE_QUEUE_SCHEMA, + ) + } + + assert compute_tables == {"jobs"} + assert {"jobs", "job_attempts", "schema_migrations"} <= queue_tables + + +async def _put_on_its_last_attempt(compute_job_id: uuid.UUID) -> None: + """Shrink the job's budget so one lease expiry exhausts it. + + Three real expiries would cost three lease durations of wall clock and + prove nothing the first one does not: what is under test is what happens + when recovery has no attempt left to give, not how it counts to three. + """ + async with db.acquire() as conn: + await conn.execute( + f"UPDATE {COMPUTE_QUEUE_SCHEMA}.jobs SET max_attempts = 1 " # noqa: S608 + "WHERE payload->>'compute_job_id' = $1", + str(compute_job_id), + ) + + +@pytest.mark.asyncio +async def test_a_crash_that_exhausts_the_lease_budget_is_reconciled( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + def unreached(*_args: Any, **_kwargs: Any) -> SimulationResult: + # Lease recovery is a SQL update. It never calls the handler back, so + # nothing in run_simulation_task gets the chance to raise or to report. + pytest.fail("recovery must not re-invoke the simulation handler") + + monkeypatch.setattr(tasks, "run_simulation", unreached) + + compute_job_id = await _submit(simulation_id) + await _put_on_its_last_attempt(compute_job_id) + + # A worker claims the job and writes compute.jobs = running exactly as + # run_simulation_task's first statement does -- and then its process dies. + # A sub-second lease is how rqueue models a holder that never comes back: + # no heartbeat renews it and no finalizing write is ever made. + async with db.acquire() as conn: + claimed = await queue.storage.claim( + conn, + queue=queue.name, + worker_id="crashed-worker", + tasks=[RUN_SIMULATION], + limit=1, + lease_seconds=0.2, + ) + assert len(claimed) == 1 + await repository.mark_started(conn, compute_job_id, uuid.UUID(simulation_id), 1) + + await asyncio.sleep(0.4) + # Any live worker's poll recovers the expired lease; drain() is one poll. + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + assert row["state"] == "failed" + assert row["error_type"] == "LeaseExpired" + + # The gap this pass exists to close: the queue has given up, and nothing + # told compute.jobs, which would otherwise read "running" forever. + assert (await repository.get_job_status(simulation_id))["status"] == "running" + + await tasks.reconcile_terminal_jobs(grace=timedelta(0)) + + status = await repository.get_job_status(simulation_id) + assert status["status"] == "failed" + assert status["error"] == ( + "Simulation stopped without reporting a result; " + "status reconciled from the task queue (LeaseExpired)" + ) + assert status["finished_at"] is not None + + # Reconciled once, never again: the second pass finds nothing to do. + assert ( + await repository.reconcile_terminal_jobs( + queue_schema=queue.schema, + queue_name=queue.name, + task=RUN_SIMULATION, + cutoff=datetime.now().astimezone(), + ) + == [] + ) + + +@pytest.mark.asyncio +async def test_reconciliation_leaves_a_finished_job_alone( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + (tmp_path / "jobs" / simulation_id).mkdir(parents=True) + monkeypatch.setattr( + tasks, + "run_simulation", + lambda _data, work_dir, **_kwargs: _result(work_dir), + ) + monkeypatch.setattr( + output_store, + "upload_simulation_result", + lambda **_kwargs: ("results", "simulations/x/metadata.json"), + ) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + assert (await repository.get_job_status(simulation_id))["status"] == "completed" + + # A run that reported its own outcome is off limits, whatever the queue + # row next to it says. Nothing here is stuck, so nothing is reconciled. + async with db.acquire() as conn: + await conn.execute( + f"UPDATE {COMPUTE_QUEUE_SCHEMA}.jobs " # noqa: S608 + "SET state = 'failed', error_type = 'LeaseExpired' " + "WHERE payload->>'compute_job_id' = $1", + str(compute_job_id), + ) + await tasks.reconcile_terminal_jobs(grace=timedelta(0)) + + assert (await repository.get_job_status(simulation_id))["status"] == "completed" + + +@pytest.mark.asyncio +async def test_a_transient_finalization_failure_is_reported_as_retrying( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + (tmp_path / "jobs" / simulation_id).mkdir(parents=True) + monkeypatch.setattr( + tasks, "run_simulation", lambda _data, work_dir, **_kwargs: _result(work_dir) + ) + + def minio_is_down(**_kwargs: Any) -> tuple[str, str]: + raise TransientInfraError("output upload failed") + + monkeypatch.setattr(output_store, "upload_simulation_result", minio_is_down) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + assert row["state"] == "pending" + assert row["error_type"] == "TransientInfraError" + + # The kernel succeeded and the upload did not, so the failure lands in + # complete_job -- outside the block that used to be the only caller of + # _record_failure. crash_recovery_e2e.sh scenario 3 watches for exactly + # this text while MinIO is stopped. + status = await repository.get_job_status(simulation_id) + assert status["status"] == "running" + assert status["details"] == "Retrying after transient error (TransientInfraError)" + # The workspace survives, so the retry resumes rather than recomputing. + assert (tmp_path / "jobs" / simulation_id).exists() + + +@pytest.mark.asyncio +async def test_a_redelivered_completed_job_is_not_run_again( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + work_dir = tmp_path / "jobs" / simulation_id + work_dir.mkdir(parents=True) + monkeypatch.setattr( + tasks, "run_simulation", lambda _data, wd, **_kwargs: _result(wd) + ) + monkeypatch.setattr( + output_store, + "upload_simulation_result", + lambda **_kwargs: ("results", "simulations/x/metadata.json"), + ) + + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + finished = await repository.get_job_status(simulation_id) + assert finished["status"] == "completed" + + # rqueue is at-least-once: a worker that died after complete_job committed + # but before rqueue finalized gets the job again. Put the queue row back + # the way that leaves it, and leave a workspace behind too. + work_dir.mkdir(parents=True, exist_ok=True) + async with db.acquire() as conn, conn.transaction(): + queue_job_id = await conn.fetchval( + f"UPDATE {COMPUTE_QUEUE_SCHEMA}.jobs " # noqa: S608 + "SET state = 'pending', finished_at = NULL, attempt = 0 " + "WHERE payload->>'compute_job_id' = $1 RETURNING id", + str(compute_job_id), + ) + # The attempt record is closed only at finalization, which is exactly + # what the dead worker never reached. + await conn.execute( + f"DELETE FROM {COMPUTE_QUEUE_SCHEMA}.job_attempts " # noqa: S608 + "WHERE job_id = $1", + queue_job_id, + ) + + def unreached(*_args: Any, **_kwargs: Any) -> SimulationResult: + pytest.fail("a completed job must not be simulated again") + + monkeypatch.setattr(tasks, "run_simulation", unreached) + await _worker(queue).drain(timeout=60) + + # Not re-run, and -- the part that actually hurt -- mark_started did not + # flip a finished job back to `running` for everyone watching. + redelivered = await repository.get_job_status(simulation_id) + assert redelivered["status"] == "completed" + assert redelivered["finished_at"] == finished["finished_at"] + assert (await _queue_row(compute_job_id))["state"] == "succeeded" + # The step the dead attempt may never have reached. + assert not work_dir.exists() + + +@pytest.mark.asyncio +async def test_one_unreadable_payload_does_not_abort_the_reconciliation_pass( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + compute_job_id = await _submit(simulation_id) + await _put_on_its_last_attempt(compute_job_id) + async with db.acquire() as conn: + claimed = await queue.storage.claim( + conn, + queue=queue.name, + worker_id="crashed-worker", + tasks=[RUN_SIMULATION], + limit=1, + lease_seconds=0.2, + ) + assert len(claimed) == 1 + await repository.mark_started(conn, compute_job_id, uuid.UUID(simulation_id), 1) + + # A terminal queue row whose payload is not a uuid. Casting it aborts + # the whole statement, so without a guard this one row would silently + # un-reconcile every genuinely stuck job beside it. + await conn.execute( + f""" + INSERT INTO {COMPUTE_QUEUE_SCHEMA}.jobs + (queue, task, payload, state, attempt, max_attempts, finished_at) + VALUES ($1, $2, '{{"compute_job_id": "not-a-uuid"}}'::jsonb, + 'failed', 1, 1, now()) + """, # noqa: S608 + queue.name, + RUN_SIMULATION, + ) + + await asyncio.sleep(0.4) + await _worker(queue).drain(timeout=60) + + await tasks.reconcile_terminal_jobs(grace=timedelta(0)) + + assert (await repository.get_job_status(simulation_id))["status"] == "failed" + + +@pytest.mark.asyncio +async def test_a_stale_attempt_cannot_undo_a_reconciled_failure( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The exposure reconciliation itself had, closed in the same statement. + + A thread abandoned by attempt N keeps running after its coroutine is gone. + If reconciliation has already marked the job failed, a stale progress write + would set it back to `running` -- silently undoing the repair. The write is + still made by the row's *owner*, so nothing but the status predicate on the + UPDATE itself can refuse it. + """ + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + compute_job_id = await _submit(simulation_id) + await _put_on_its_last_attempt(compute_job_id) + + async with db.acquire() as conn: + claimed = await queue.storage.claim( + conn, + queue=queue.name, + worker_id="crashed-worker", + tasks=[RUN_SIMULATION], + limit=1, + lease_seconds=0.2, + ) + assert len(claimed) == 1 + attempt = claimed[0].job.attempt + assert await repository.mark_started( + conn, compute_job_id, uuid.UUID(simulation_id), attempt + ) + + await asyncio.sleep(0.4) + await _worker(queue).drain(timeout=60) + await tasks.reconcile_terminal_jobs(grace=timedelta(0)) + assert (await repository.get_job_status(simulation_id))["status"] == "failed" + + # The kernel thread from that attempt is still alive and still owns the row. + async with db.acquire() as conn: + written = await repository.record_progress( + conn, + compute_job_id, + uuid.UUID(simulation_id), + "Processing tsunami", + {"step": "tsunami", "step_index": 3, "total_steps": 8}, + attempt, + ) + + assert written is False + status = await repository.get_job_status(simulation_id) + assert status["status"] == "failed" + assert status["step"] != "tsunami" + assert status["error"] == ( + "Simulation stopped without reporting a result; " + "status reconciled from the task queue (LeaseExpired)" + ) + + +@pytest.mark.asyncio +async def test_a_superseded_attempt_cannot_write_over_the_newer_one( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + compute_job_id = await _submit(simulation_id) + job_uuid = uuid.UUID(simulation_id) + + async with db.acquire() as conn: + assert await repository.mark_started(conn, compute_job_id, job_uuid, 1) + # A newer attempt takes the row over, as it does after lease recovery. + assert await repository.mark_started(conn, compute_job_id, job_uuid, 2) + assert await repository.record_progress( + conn, compute_job_id, job_uuid, "Processing maxola", {"step": "maxola"}, 2 + ) + + # Attempt 1's thread is still running and still holds the loop. + assert not await repository.record_progress( + conn, compute_job_id, job_uuid, "Processing tsunami", {"step": "tsunami"}, 1 + ) + # And it cannot take the row back either. + assert not await repository.mark_started(conn, compute_job_id, job_uuid, 1) + + status = await repository.get_job_status(simulation_id) + assert status["step"] == "maxola" + assert status["details"] == "Processing maxola" + + +@pytest.mark.asyncio +async def test_an_operator_retry_of_a_failed_job_can_still_start( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Why `mark_started` refuses only COMPLETED, never FAILED. + + `rqueue.Admin.retry_job` is the documented way an operator restarts a job + rqueue considers terminally failed, and it works by putting the same row + back to `pending` with a raised budget -- so the retry arrives as an + ordinary claim of the same job, running the same handler. If `mark_started` + refused a `failed` compute row the way it refuses a `completed` one, that + retry would run a whole simulation whose every write was rejected. + """ + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + + def fail(*_args: Any, **_kwargs: Any) -> SimulationResult: + raise RuntimeError("bad epicenter") + + monkeypatch.setattr(tasks, "run_simulation", fail) + compute_job_id = await _submit(simulation_id) + await _worker(queue).drain(timeout=60) + + failed_row = await _queue_row(compute_job_id) + assert failed_row["state"] == "failed" + assert (await repository.get_job_status(simulation_id))["status"] == "failed" + + # The operator retry, through rqueue's own public API. + (tmp_path / "jobs" / simulation_id).mkdir(parents=True, exist_ok=True) + monkeypatch.setattr( + tasks, "run_simulation", lambda _data, wd, **_kwargs: _result(wd) + ) + monkeypatch.setattr( + output_store, + "upload_simulation_result", + lambda **_kwargs: ("results", "simulations/x/metadata.json"), + ) + retried = await Admin(queue.pool, schema=queue.schema).retry_job(failed_row["id"]) + assert retried.state == "pending" + + await _worker(queue).drain(timeout=60) + + assert (await _queue_row(compute_job_id))["state"] == "succeeded" + assert (await repository.get_job_status(simulation_id))["status"] == "completed" + + +@pytest.mark.asyncio +async def test_a_stand_down_is_not_counted_as_a_failed_job( + queue: Queue, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """A refused write must not look like a broken simulation to rqueue.""" + simulation_id = str(uuid.uuid4()) + monkeypatch.setattr(tasks, "JOBS_DIR", tmp_path / "jobs") + compute_job_id = await _submit(simulation_id) + + def superseded(*_args: Any, on_progress: Any, **_kwargs: Any) -> SimulationResult: + # A newer attempt takes the row over while this kernel is running. + asyncio.run(_take_over(compute_job_id, uuid.UUID(simulation_id))) + on_progress("Processing tsunami", {"step": "tsunami"}) + pytest.fail("the kernel must not continue past a refused write") + + monkeypatch.setattr(tasks, "run_simulation", superseded) + await _worker(queue).drain(timeout=60) + + row = await _queue_row(compute_job_id) + # Not `failed`: standing down is the fence working, and rqueue's failure + # accounting has to keep meaning genuine failures. + assert row["state"] == "cancelled" + assert row["error_type"] == "Cancelled" + assert "no longer owns this job" in row["error_message"] + + # And reconciliation treats a cancelled queue row as the queue giving up, + # so compute.jobs cannot be left reading `running` forever either. + await tasks.reconcile_terminal_jobs(grace=timedelta(0)) + assert (await repository.get_job_status(simulation_id))["status"] == "failed" + + +async def _take_over(compute_job_id: uuid.UUID, simulation_id: uuid.UUID) -> None: + """Claim the compute row for a later attempt, from off the loop.""" + # The queue fixture points this at the test's disposable database. + connection = await asyncpg.connect(db_module.COMPUTE_DATABASE_URL) + try: + await connection.execute( + "UPDATE compute.jobs SET owner_attempt = 99 WHERE id = $1", compute_job_id + ) + finally: + await connection.close() diff --git a/pyproject.toml b/pyproject.toml index 220694c..8435940 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,7 @@ exclude_lines = [ testpaths = ["packages/tsdhn/tests", "packages/api/tests", "packages/tsdhn-parity/tests"] pythonpath = ["."] addopts = "-ra" +asyncio_default_fixture_loop_scope = "function" markers = [ "golden: requires the real compiled toolchain; see mise run test-golden", "integration: uses disposable databases on the mise-managed PostgreSQL cluster; see mise run test-integration", diff --git a/scripts/e2e/compute_stack_smoke.sh b/scripts/e2e/compute_stack_smoke.sh index bf5d7ce..e318fc5 100755 --- a/scripts/e2e/compute_stack_smoke.sh +++ b/scripts/e2e/compute_stack_smoke.sh @@ -7,6 +7,8 @@ set -euo pipefail : "${MINIO_SECRET_KEY:=minioadmin}" : "${MINIO_BUCKET:=tsdhn-results}" : "${COMPUTE_API_URL:=http://localhost:8000}" +: "${COMPUTE_QUEUE_SCHEMA:=task_queue}" +: "${COMPUTE_QUEUE:=simulations}" wait_for_api() { for _ in {1..60}; do @@ -101,12 +103,12 @@ assert_queued_job() { test "$compute_count" = "1" docker compose exec -T postgres psql -U tsdhn -d tsdhn -c \ - "SELECT tablename FROM pg_tables WHERE schemaname = 'public' AND tablename LIKE 'procrastinate_%' ORDER BY tablename" + "SELECT tablename FROM pg_tables WHERE schemaname = '$COMPUTE_QUEUE_SCHEMA' ORDER BY tablename" # Idempotent submissions must not create a second queue row. queue_count="$( docker compose exec -T postgres psql -U tsdhn -d tsdhn -tAc \ - "SELECT count(*) FROM procrastinate_jobs WHERE task_name = 'api.run_simulation' AND queue_name = 'simulations'" + "SELECT count(*) FROM $COMPUTE_QUEUE_SCHEMA.jobs WHERE task = 'api.run_simulation' AND queue = '$COMPUTE_QUEUE'" )" test "$queue_count" = "1" } diff --git a/scripts/e2e/crash_recovery_e2e.sh b/scripts/e2e/crash_recovery_e2e.sh index d2f81d7..273b9ce 100755 --- a/scripts/e2e/crash_recovery_e2e.sh +++ b/scripts/e2e/crash_recovery_e2e.sh @@ -9,6 +9,7 @@ set -euo pipefail : "${MINIO_SECRET_KEY:=minioadmin}" : "${MINIO_BUCKET:=tsdhn-results}" : "${COMPUTE_API_URL:=http://localhost:8000}" +: "${COMPUTE_QUEUE_SCHEMA:=task_queue}" # The crash must happen after the resumable step writes a checkpoint. : "${TSUNAMI_CHECKPOINT_WAIT_SECONDS:=120}" @@ -110,27 +111,27 @@ scenario_crash_and_requeue() { echo "In tsunami step; waiting ${TSUNAMI_CHECKPOINT_WAIT_SECONDS}s for checkpoints to accumulate" sleep "$TSUNAMI_CHECKPOINT_WAIT_SECONDS" - # Requeueing increments the existing queue row's attempt counter. + # Recovering an expired lease increments the same queue row's attempt. local queue_job_id before_attempts queue_job_id="$( - psql_c "SELECT id FROM procrastinate_jobs - WHERE task_name = 'api.run_simulation' - AND task_kwargs->>'compute_job_id' = + psql_c "SELECT id FROM $COMPUTE_QUEUE_SCHEMA.jobs + WHERE task = 'api.run_simulation' + AND payload->>'compute_job_id' = (SELECT id::text FROM compute.jobs WHERE simulation_id = '$CRASH_SIMULATION_ID'::uuid)" )" test -n "$queue_job_id" - before_attempts="$(psql_c "SELECT attempts FROM procrastinate_jobs WHERE id = $queue_job_id")" + before_attempts="$(psql_c "SELECT attempt FROM $COMPUTE_QUEUE_SCHEMA.jobs WHERE id = '$queue_job_id'::uuid")" echo "Killing worker (SIGKILL) to simulate a crash" docker compose kill -s SIGKILL worker - echo "Waiting up to 240s for the reaper to detect the stale job and requeue it" + echo "Waiting up to 240s for rqueue to recover the expired lease and requeue it" local deadline requeued=0 deadline=$((SECONDS + 240)) while [ "$SECONDS" -lt "$deadline" ]; do local attempts status - attempts="$(psql_c "SELECT attempts FROM procrastinate_jobs WHERE id = $queue_job_id")" + attempts="$(psql_c "SELECT attempt FROM $COMPUTE_QUEUE_SCHEMA.jobs WHERE id = '$queue_job_id'::uuid")" status="$(job_status "$CRASH_SIMULATION_ID" | jq -r .status)" if [ "$status" = "failed" ]; then echo "::error::Job was marked FAILED instead of being requeued" @@ -146,7 +147,7 @@ scenario_crash_and_requeue() { echo "::error::Job was never requeued after the worker crash" return 1 } - echo "Confirmed: attempts incremented on the same queue job, not marked FAILED" + echo "Confirmed: attempt incremented on the same queue job, not marked FAILED" poll_job_to_completion "$CRASH_SIMULATION_ID" 1800 @@ -175,12 +176,27 @@ scenario_ttl_sweep() { " > /dev/null docker compose exec -T worker uv run --no-dev python -c " -from api.core.tasks import sweep_abandoned_work_dirs_task -sweep_abandoned_work_dirs_task(0) +import asyncio + +from api.core import db +from api.core.settings import worker_pool_size +from api.core.tasks import sweep_abandoned_work_dirs + + +async def main(): + minimum, maximum = worker_pool_size() + await db.open_pool(min_size=minimum, max_size=maximum) + try: + await sweep_abandoned_work_dirs() + finally: + await db.close_pool() + + +asyncio.run(main()) " if docker compose exec -T worker test -e "$work_dir"; then - echo "::error::sweep_abandoned_work_dirs_task did not remove $work_dir" + echo "::error::sweep_abandoned_work_dirs did not remove $work_dir" return 1 fi echo "Confirmed: expired FAILED job's work_dir was swept" diff --git a/scripts/integration.sh b/scripts/integration.sh index 7537ca3..621a930 100755 --- a/scripts/integration.sh +++ b/scripts/integration.sh @@ -6,6 +6,7 @@ base_url="postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn" database_name="tsdhn_integration_$(date +%s)_$$" app_role="${database_name}_role" app_password="tsdhn-web-test-password" +queue_schema="${COMPUTE_QUEUE_SCHEMA:-task_queue}" case "${1:-}" in "") coverage=0 ;; @@ -44,8 +45,7 @@ APP_DB_ROLE="$app_role" \ APP_DB_PASSWORD="$app_password" \ uv run tsdhn-compute-migrate -COMPUTE_DATABASE_URL="$admin_url" \ -uv run tsdhn-procrastinate-migrate +uv run rqueue --database-url "$admin_url" --schema "$queue_schema" migrate # Schema changes run with the database-owner connection. The app role below # is intentionally limited to runtime DML and compute-state reads. diff --git a/uv.lock b/uv.lock index 73ab0f7..31398dd 100644 --- a/uv.lock +++ b/uv.lock @@ -91,15 +91,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/42/b9/f8d6fa329ab25128b7e98fd83a3cb34d9db5b059a9847eddb840a0af45dd/argon2_cffi_bindings-25.1.0-cp39-abi3-win_arm64.whl", hash = "sha256:b0fdbcf513833809c882823f98dc2f931cf659d9a1429616ac3adebb49f5db94", size = 27149, upload-time = "2025-07-30T10:01:59.329Z" }, ] -[[package]] -name = "asgiref" -version = "3.11.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/63/40/f03da1264ae8f7cfdbf9146542e5e7e8100a4c66ab48e791df9a03d3f6c0/asgiref-3.11.1.tar.gz", hash = "sha256:5f184dc43b7e763efe848065441eac62229c9f7b0475f41f80e207a114eda4ce", size = 38550, upload-time = "2026-02-03T13:30:14.33Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" }, -] - [[package]] name = "ast-serialize" version = "0.6.0" @@ -142,12 +133,27 @@ wheels = [ ] [[package]] -name = "attrs" -version = "26.1.0" +name = "asyncpg" +version = "0.31.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fe/cc/d18065ce2380d80b1bcce927c24a2642efd38918e33fd724bc4bca904877/asyncpg-0.31.0.tar.gz", hash = "sha256:c989386c83940bfbd787180f2b1519415e2d3d6277a70d9d0f0145ac73500735", size = 993667, upload-time = "2025-11-24T23:27:00.812Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, + { url = "https://files.pythonhosted.org/packages/3c/36/e9450d62e84a13aea6580c83a47a437f26c7ca6fa0f0fd40b6670793ea30/asyncpg-0.31.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f6b56b91bb0ffc328c4e3ed113136cddd9deefdf5f79ab448598b9772831df44", size = 660867, upload-time = "2025-11-24T23:26:17.631Z" }, + { url = "https://files.pythonhosted.org/packages/82/4b/1d0a2b33b3102d210439338e1beea616a6122267c0df459ff0265cd5807a/asyncpg-0.31.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:334dec28cf20d7f5bb9e45b39546ddf247f8042a690bff9b9573d00086e69cb5", size = 638349, upload-time = "2025-11-24T23:26:19.689Z" }, + { url = "https://files.pythonhosted.org/packages/41/aa/e7f7ac9a7974f08eff9183e392b2d62516f90412686532d27e196c0f0eeb/asyncpg-0.31.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:98cc158c53f46de7bb677fd20c417e264fc02b36d901cc2a43bd6cb0dc6dbfd2", size = 3410428, upload-time = "2025-11-24T23:26:21.275Z" }, + { url = "https://files.pythonhosted.org/packages/6f/de/bf1b60de3dede5c2731e6788617a512bc0ebd9693eac297ee74086f101d7/asyncpg-0.31.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9322b563e2661a52e3cdbc93eed3be7748b289f792e0011cb2720d278b366ce2", size = 3471678, upload-time = "2025-11-24T23:26:23.627Z" }, + { url = "https://files.pythonhosted.org/packages/46/78/fc3ade003e22d8bd53aaf8f75f4be48f0b460fa73738f0391b9c856a9147/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19857a358fc811d82227449b7ca40afb46e75b33eb8897240c3839dd8b744218", size = 3313505, upload-time = "2025-11-24T23:26:25.235Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e9/73eb8a6789e927816f4705291be21f2225687bfa97321e40cd23055e903a/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ba5f8886e850882ff2c2ace5732300e99193823e8107e2c53ef01c1ebfa1e85d", size = 3434744, upload-time = "2025-11-24T23:26:26.944Z" }, + { url = "https://files.pythonhosted.org/packages/08/4b/f10b880534413c65c5b5862f79b8e81553a8f364e5238832ad4c0af71b7f/asyncpg-0.31.0-cp314-cp314-win32.whl", hash = "sha256:cea3a0b2a14f95834cee29432e4ddc399b95700eb1d51bbc5bfee8f31fa07b2b", size = 532251, upload-time = "2025-11-24T23:26:28.404Z" }, + { url = "https://files.pythonhosted.org/packages/d3/2d/7aa40750b7a19efa5d66e67fc06008ca0f27ba1bd082e457ad82f59aba49/asyncpg-0.31.0-cp314-cp314-win_amd64.whl", hash = "sha256:04d19392716af6b029411a0264d92093b6e5e8285ae97a39957b9a9c14ea72be", size = 604901, upload-time = "2025-11-24T23:26:30.34Z" }, + { url = "https://files.pythonhosted.org/packages/ce/fe/b9dfe349b83b9dee28cc42360d2c86b2cdce4cb551a2c2d27e156bcac84d/asyncpg-0.31.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bdb957706da132e982cc6856bb2f7b740603472b54c3ebc77fe60ea3e57e1bd2", size = 702280, upload-time = "2025-11-24T23:26:32Z" }, + { url = "https://files.pythonhosted.org/packages/6a/81/e6be6e37e560bd91e6c23ea8a6138a04fd057b08cf63d3c5055c98e81c1d/asyncpg-0.31.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6d11b198111a72f47154fa03b85799f9be63701e068b43f84ac25da0bda9cb31", size = 682931, upload-time = "2025-11-24T23:26:33.572Z" }, + { url = "https://files.pythonhosted.org/packages/a6/45/6009040da85a1648dd5bc75b3b0a062081c483e75a1a29041ae63a0bf0dc/asyncpg-0.31.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18c83b03bc0d1b23e6230f5bf8d4f217dc9bc08644ce0502a9d91dc9e634a9c7", size = 3581608, upload-time = "2025-11-24T23:26:35.638Z" }, + { url = "https://files.pythonhosted.org/packages/7e/06/2e3d4d7608b0b2b3adbee0d0bd6a2d29ca0fc4d8a78f8277df04e2d1fd7b/asyncpg-0.31.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e009abc333464ff18b8f6fd146addffd9aaf63e79aa3bb40ab7a4c332d0c5e9e", size = 3498738, upload-time = "2025-11-24T23:26:37.275Z" }, + { url = "https://files.pythonhosted.org/packages/7d/aa/7d75ede780033141c51d83577ea23236ba7d3a23593929b32b49db8ed36e/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3b1fbcb0e396a5ca435a8826a87e5c2c2cc0c8c68eb6fadf82168056b0e53a8c", size = 3401026, upload-time = "2025-11-24T23:26:39.423Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7a/15e37d45e7f7c94facc1e9148c0e455e8f33c08f0b8a0b1deb2c5171771b/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8df714dba348efcc162d2adf02d213e5fab1bd9f557e1305633e851a61814a7a", size = 3429426, upload-time = "2025-11-24T23:26:41.032Z" }, + { url = "https://files.pythonhosted.org/packages/13/d5/71437c5f6ae5f307828710efbe62163974e71237d5d46ebd2869ea052d10/asyncpg-0.31.0-cp314-cp314t-win32.whl", hash = "sha256:1b41f1afb1033f2b44f3234993b15096ddc9cd71b21a42dbd87fc6a57b43d65d", size = 614495, upload-time = "2025-11-24T23:26:42.659Z" }, + { url = "https://files.pythonhosted.org/packages/3c/d7/8fb3044eaef08a310acfe23dae9a8e2e07d305edc29a53497e52bc76eca7/asyncpg-0.31.0-cp314-cp314t-win_amd64.whl", hash = "sha256:bd4107bb7cdd0e9e65fae66a62afd3a249663b844fa34d479f6d5b3bef9c04c3", size = 706062, upload-time = "2025-11-24T23:26:44.086Z" }, ] [[package]] @@ -384,18 +390,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/52/30/21b2ad45959cd50e909e02ebac1e30b4ceb7162e91c11d4c570223a458b7/coverage-7.15.0-py3-none-any.whl", hash = "sha256:56da6a4cbe8f7e9e80bd072ca9cefe67d7106a440a7ec06519ec6507ac94ad19", size = 212632, upload-time = "2026-07-02T13:10:48.641Z" }, ] -[[package]] -name = "croniter" -version = "6.2.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "python-dateutil" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/03/35/96ad0a71eb0b27ab4476a7ed23facd0713d82da9c911edc8af7f34a62d6a/croniter-6.2.3.tar.gz", hash = "sha256:fb129986ef7e2c44e3f4c9f503da83ad914d2afa48f40a43ee3dca4b5c41d476", size = 166174, upload-time = "2026-07-02T14:34:22.166Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/dd/6466498a8b69754cffbd7237ce4c66446ca5ffcf53fb397d437666f056d2/croniter-6.2.3-py3-none-any.whl", hash = "sha256:137a97001b4d52fb71c10b750e303db79e6e42d40fff8ff77126102176c9f786", size = 46446, upload-time = "2026-07-02T14:34:20.889Z" }, -] - [[package]] name = "cyclonedx-python-lib" version = "11.11.0" @@ -957,24 +951,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] -[[package]] -name = "procrastinate" -version = "3.9.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "asgiref" }, - { name = "attrs" }, - { name = "croniter" }, - { name = "packaging" }, - { name = "psycopg", extra = ["pool"] }, - { name = "python-dateutil" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ab/bc/adfe64992725143791ea78c33e9a0778a97f67c0ca8761bb6fd269fbb0be/procrastinate-3.9.0.tar.gz", hash = "sha256:5805ab2af35eab12befa700ecd49e572c4f655d654151df9e5ce1ca07efb5e6e", size = 89448, upload-time = "2026-06-20T23:09:14.663Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ee/0c/17bfd406fe4bf85e8a0cfbb9744271df447004233da295fb6e2c10dc29e1/procrastinate-3.9.0-py3-none-any.whl", hash = "sha256:af4b9ccaeebbf2a1439e02ae1e8ce327f8436a81e9d68a0b1bd26d5e7ac697bd", size = 153597, upload-time = "2026-06-20T23:09:13.105Z" }, -] - [[package]] name = "psycopg" version = "3.3.4" @@ -991,9 +967,6 @@ wheels = [ binary = [ { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, ] -pool = [ - { name = "psycopg-pool" }, -] [[package]] name = "psycopg-binary" @@ -1013,18 +986,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/e6/5fff07a70d1f945ed90ae131c3bd76cab32beff7c58c6db15ad5820b6d1f/psycopg_binary-3.3.4-cp314-cp314-win_amd64.whl", hash = "sha256:c37e024c07308cd06cf3ec51bfd0e7f6157585a4d84d1bce4a7f5f7913719bf8", size = 3666849, upload-time = "2026-05-01T23:31:51.165Z" }, ] -[[package]] -name = "psycopg-pool" -version = "3.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/90/82/7a23d26039827ecd4ebe93905651029ddd307c5182ad59296dfb6f67b528/psycopg_pool-3.3.1.tar.gz", hash = "sha256:b10b10b7a175d5cc1592147dc5b7eec8a9e0834eb3ed2c4a92c858e2f51eb63c", size = 31661, upload-time = "2026-05-01T23:31:59.809Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/37/ed/89c2c620af0e1660354cd8aabf9f5b21f911597ce22acb37c805d6c86bc8/psycopg_pool-3.3.1-py3-none-any.whl", hash = "sha256:2af5b432941c4c9ad5c87b3fa410aec910ec8f7c122855897983a06c45f2e4b5", size = 40023, upload-time = "2026-05-01T23:31:53.136Z" }, -] - [[package]] name = "py-serializable" version = "2.1.0" @@ -1275,6 +1236,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, ] +[[package]] +name = "rqueue" +version = "0.2.0" +source = { git = "https://github.com/totallynotdavid/transactions.git?rev=545d67aa341372c663972427b2b520fcd07faf80#545d67aa341372c663972427b2b520fcd07faf80" } +dependencies = [ + { name = "asyncpg" }, +] + [[package]] name = "ruff" version = "0.15.20" @@ -1459,11 +1428,12 @@ version = "0.0.1" source = { editable = "packages/api" } dependencies = [ { name = "anyio" }, + { name = "asyncpg" }, { name = "fastapi" }, { name = "minio" }, - { name = "procrastinate" }, - { name = "psycopg", extra = ["binary", "pool"] }, + { name = "psycopg", extra = ["binary"] }, { name = "pydantic" }, + { name = "rqueue" }, { name = "tsdhn" }, { name = "uvicorn" }, ] @@ -1471,11 +1441,12 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "anyio", specifier = "==4.14.1" }, + { name = "asyncpg", specifier = ">=0.31.0" }, { name = "fastapi", specifier = "==0.139.0" }, { name = "minio", specifier = ">=7.2.20" }, - { name = "procrastinate", specifier = ">=3.9.0" }, - { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.3.4" }, + { name = "psycopg", extras = ["binary"], specifier = ">=3.3.4" }, { name = "pydantic", specifier = "==2.13.4" }, + { name = "rqueue", git = "https://github.com/totallynotdavid/transactions.git?rev=545d67aa341372c663972427b2b520fcd07faf80" }, { name = "tsdhn", editable = "packages/tsdhn" }, { name = "uvicorn", specifier = "==0.51.0" }, ]