Wf refactor

[bd-mkt]

Workflow engine refactor: class-based Worker, atomic claim + lease tokens, ResourceLock

Summary

This branch is primarily a foundational refactor of the workflow execution engine: it replaces a module-global runner with a class-based Worker that owns typed consumer pools, lease-gated terminal writes, and a cross-subsystem ResourceLock table. The new persistence layer eliminates several long-standing concurrency bugs in the original runner (process-local locking, TOCTOU on group completion, no protection against stale workers double-finalizing steps). A handful of smaller fixes ride along: compression for filesystem artifacts, a free-disk-space pre-check, async haiku.rag wiring, docling refactor, and broad documentation updates.

Workflow engine refactor

Implemented across phases 1–8 (see commits f0728ab, 545a1a7, 44631cd) plus a final atomic-claim race fix.

Persistence layer (wf/operations.py)

• Atomic claim — claim_next_step(worker_id, lease_token, allowed_types=…) replaces get_runnable_steps + set_step_status(RUNNING) with a single UPDATE-FROM-SELECT-FOR-UPDATE-SKIP-LOCKED. Per-type concurrency is bounded at the SQL layer, eliminating in-process semaphores.
• Lease tokens — every claim mints a per-claim UUID stamped on runstep.lease_token. complete_step / error_step / release_step are gated on the lease, so a worker reaped between claim and write can never bounce a fresh claimant.
• ResourceLock table — cross-subsystem rendezvous keyed by resource_key (typically rag:<abs-db-path>). All RAG-DB writers (workflow save_to_rag steps, web vacuum endpoint, si-diag CLI, end_group lifecycle vacuums) coordinate here. TTL-based; expired rows are swept by a background loop.
• Atomic claim + lock — when a step has a resource_key, the matching ResourceLock row is INSERTed in the same transaction as the claim. The unique PK constraint serializes concurrent claims targeting the same key (e.g. multiple save_to_rag steps for the same LanceDB); the loser rolls back the whole claim and returns None. Closes the residual race window that READ COMMITTED isolation left open between two claim transactions.
• try_complete_run_group — atomic conditional UPDATE so GROUP_END fires exactly once across N workers. Replaces the TOCTOU window in the old get_run_group_stats snapshot read.
• Self-skipping reaper — reap_dead_workers(my_id, threshold) always excludes the caller, eliminating the self-reaping race where a stalled checkin loop would let a worker reset its own in-flight steps.

Orchestration layer (wf/runner.py)

• class Worker — instance-scoped state; multiple workers can coexist in one process. Configurable per-type consumer pools: Worker(WorkerConfig(consumers={"parse": 4, "store": 8, "*": 2})).
• Graceful shutdown — Worker.stop(timeout) waits for in-flight steps, then cancels remaining consumers. Cancelled consumers call release_step so rows return to PENDING immediately rather than waiting WORKER_CHECKIN_TIMEOUT. The worker also deletes its own WorkerCheckin row so peers see the departure right away.
• Event-driven LifecycleBus — fire-and-forget tasks so a slow STEP_START handler can never block step execution.
• Pluggable Metrics protocol — counters (claim_attempts, claim_success, claim_idle, claim_lost_race, step_completed, step_error, step_failed, step_released, worker_reaped, lease_lost, resource_lock_swept) and histograms (claim_duration, step_duration). Default LoggingMetrics is a DEBUG no-op; production wires Prometheus.
• Backoff on race-release — when _run_step race-releases on a resource lock (rare defense-in-depth path after atomic claim), the consumer applies idle-style exponential backoff so it doesn't spin-claim until the holder finishes. poll_backoff_max tuned from 5s → 3s to match typical LanceDB write durations.

Caller wiring

• lib/rag.py — dropped the per-DB lock_map. Added resource_key_for(db_path) and hold_rag_lock(...) context manager for direct-from-Python writers. vacuum_db(db_name, sign, holder_kind, max_wait) wraps the operation in a ResourceLock; max_wait=0 for fail-fast.
• lib/docling.py — module-level HTTP semaphore removed; parse concurrency is now expressed by the parse consumer pool size.
• server/routes/lancedb.py — vacuum endpoint returns 200 / 404 / 409 (lock held by another writer) / 500. Routes through vacuum_db(holder_kind=WEB, max_wait=0).
• si-diag lancedb vacuum — --force flag calls force_release_resource_lock(...) (audit-logged) to break a stuck lock.
• si-cli vacuum / vacuum-all — holder_kind=CLI; vacuum-all skips DBs whose lock is currently held.
• FastAPI lifespan + si-cli worker construct a Worker directly with {"parse": docling_concurrency, "store": worker_task_count, "*": worker_task_count} and call Worker.stop(30s) on shutdown.

Database migration

New migration a17c4d9b3e02_wf_refactor adds:

• runstep.lease_token (nullable VARCHAR)
• runstep.resource_key (nullable VARCHAR, ix_runstep_resource_key index)
• resourcelock table + ix_resourcelock_expires_at index

A drop-in SQL upgrade script for environments where Alembic can't run in-place is committed at scripts/wf_refactor_upgrade.sql; the matching afsoc-rag init upgrade script is at afsoc-rag/docker/pgsql-init/upgrade_scripts/afsoc_ingester_002.sql.

Other improvements

• Artifact compression (c439bd6) — new FILE_COMPRESSION_ARTIFACTS / FILE_COMPRESSION_LEVEL settings opt-in zstd compression for filesystem-backed artifact types.
• Free disk space pre-check + I/O cleanup (f79b3df) — workflows refuse to start steps that would partially fill the disk; cleanup actions run on I/O errors to avoid partial writes.
• Splitting throughput tweaks — perf-oriented adjustments to chunk/split settings.
• Async haiku.rag (9a09980) — locked to >0.44; check_exists factored out of the rag client.
• Param-set sort by id (c439bd6) — predictable ordering in CLI/diag outputs.
• Dynamic /run/secrets handling — Settings only opts into the secrets directory if it actually exists, suppressing spurious warnings in non-container environments.
• uv upgrade in the Dockerfile.

Documentation

• docs/ARCHITECTURE.md — rewrote §3 Worker System and the execution-flow steps to cover typed pools, lease tokens, ResourceLock, graceful shutdown, metrics, and exactly-once group completion.
• docs/WORKFLOWS.md — replaced "Worker Processing" with the atomic-claim flow; added Consumer Pools, Graceful Shutdown, Metrics, and resource-lock troubleshooting sections.
• docs/DATABASE.md — lease_token and resource_key columns on RunStep, new ResourceLock table, ResourceLockKind enum, CANCELLED in RunStatus, indexes from the migration.
• docs/API.md — /lancedb/vacuum 404 / 409 / 500 response shapes and ResourceLock behavior.
• docs/CLI.md — si-diag lancedb subcommand docs (vacuum --force, vacuum-all, verify); updated si-cli vacuum to mention the cross-subsystem lock.
• docs/CONFIGURATION.md — rewrote DOCLING_CONCURRENCY / WORKER_TASK_COUNT / INGEST_WORKER_CONCURRENCY for the new consumer-pool semantics; added FILE_COMPRESSION_* settings.

Tests

• tests/unit/test_runner.py — rewritten against the new pure helpers (transition_allowed, elevate_terminal).
• tests/unit/test_runner_internals.py (new, +1487 lines) — lifts runner.py coverage from 26% to 100%. Covers LifecycleBus, Worker construction/start/stop/cancellation, consumer-loop backoff (including race-release), heartbeat / reaper / lock-sweeper loops, default lifecycle hooks, module-level shim parity.
• tests/unit/test_wf_claim.py (new, +586 lines) — claim_next_step priority/skip-lock/resource-key behavior; atomic lock insert on claim; acquire_resource_lock idempotence for the same holder; lease-gated terminal writes; reaper self-skip and orphan reset; end-to-end stale-worker block; recompute_run_status and try_complete_run_group semantics; _rag_resource_key_for_param stamping at workflow-run creation.
• tests/unit/test_commit_rollback_paths.py (new, +1066 lines) — exhaustive transaction-handling coverage.
• tests/unit/test_lancedb_routes.py — vacuum endpoint mocks the new rag.vacuum_db entry point and asserts 404 / 409 paths.

741 unit tests pass; full branch coverage on runner.py; suite-wide coverage up from 50% min → 72%.

Operator notes

• Migration required: existing deployments must apply a17c4d9b3e02_wf_refactor (or the SQL drop-in) before starting the new worker. Running the new worker against the old schema will error on the missing lease_token / resource_key columns and resourcelock table.
• DOCLING_CONCURRENCY semantics changed: was an in-process HTTP semaphore; is now the size of the dedicated parse consumer pool. Behavior is equivalent for the default parse workload, but operators who relied on the global cap (INGEST_WORKER_CONCURRENCY) should review their pool sizing.
• Vacuum coordination: workflow save_to_rag steps, web vacuum, CLI vacuum, and lifecycle vacuum now coordinate via ResourceLock. The web endpoint is fail-fast (409); the CLI defaults to wait-forever; si-diag lancedb vacuum --force is the escape hatch for a genuinely stuck lock.

Test plan

• Run a multi-document batch ingestion targeting a single LanceDB; verify save_to_rag steps serialize (no concurrent writers, no race-release log spam)
• Run a batch targeting two distinct LanceDBs; verify they process in parallel
• Restart the worker mid-batch; verify in-flight steps return to PENDING within seconds (not after the checkin timeout)
• [ x Hit /api/v1/lancedb/vacuum while a save_to_rag step is in flight; expect 409
• Run si-diag lancedb vacuum <db> --force on a healthy DB; expect audit log + successful vacuum
• Confirm unit suite passes locally (uv run pytest tests/unit)