danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Renew job leases from the worker pool; lease-based reaping

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-06-25 13:20:28 +03:00
parent fd73f57649
commit 4707db780f
No known key found for this signature in database
6 changed files with 319 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
CHANGELOG.md
View file

@ -1,18 +1,16 @@
# Changelog
## [Unreleased]
### Added
### Changed
- Ingester queue schema v2 adds a `last_heartbeat_at` lease column on `jobs`; existing queue databases migrate in place on open.
- Ingester reaping is lease-based: a worker renews `last_heartbeat_at` on its in-flight jobs every `heartbeat_interval_s`, and the reaper reclaims a claim only once its lease is older than `lease_ttl_s`. A job slower than the timeout is no longer reaped and reprocessed while still running. Adds queue schema v2 (`last_heartbeat_at` on `jobs`); existing queue databases migrate in place on open.
- **Breaking:** `ingester.workers.claim_timeout_s` removed; set `lease_ttl_s` (default 120) and `heartbeat_interval_s` (default 30) instead. `WorkerConfig` rejects unknown keys.
- A bare `${VAR}` in YAML config now raises `MissingEnvVarError` when the variable is set but empty, matching the unset case. Use `${VAR:-default}` to allow an empty/absent value.
### Fixed
- A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible.
- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`) instead of `worker-{n}`, so the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue.
### Changed
- A bare `${VAR}` in YAML config now raises `MissingEnvVarError` when the variable is set but empty, matching the unset case. Use `${VAR:-default}` to allow an empty/absent value.
- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`); the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue.
## [0.61.2] - 2026-06-24

43
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -1,7 +1,7 @@
from pathlib import Path
from typing import Annotated, Any, Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from haiku.rag.utils import get_default_data_dir
@ -309,6 +309,10 @@ class CircuitBreakerConfig(BaseModel):
class WorkerConfig(BaseModel):
# Reject unknown keys so a renamed/removed setting (e.g. the former
# claim_timeout_s) fails loudly instead of being silently ignored.
model_config = ConfigDict(extra="forbid", validate_assignment=True)
worker_count: int = Field(
default=4,
description="Number of async worker tasks pulling from the queue. "
@ -323,14 +327,22 @@ class WorkerConfig(BaseModel):
"polls. Lower = lower latency picking up new jobs, higher = less "
"queue churn when the queue is usually empty.",
)
claim_timeout_s: int = Field(
default=1800,
description="A `claimed` job whose claimed_at is older than this is "
"presumed dead and reset to `queued` by the reaper. MUST exceed your "
"longest legitimate job duration — set it too short and the reaper "
"resurrects jobs still being processed, causing two workers to run "
"the same URI. Default (30min) covers typical docling conversions; "
"raise it if you ingest very large PDFs through docling-local.",
lease_ttl_s: int = Field(
default=120,
gt=0,
description="A `claimed` job whose lease has not been renewed within "
"this window is presumed dead and reset to `queued` by the reaper. A "
"live worker renews its lease every heartbeat_interval_s while "
"processing, so this need not exceed job duration — it only bounds how "
"long a crashed worker's job stays stuck before another worker takes "
"it over.",
)
heartbeat_interval_s: int = Field(
default=30,
gt=0,
description="How often a worker renews the lease on its in-flight "
"jobs. Must be comfortably shorter than lease_ttl_s so scheduler "
"jitter or a slow DB round-trip can't let a live job's lease lapse.",
)
reaper_interval_s: int = Field(
default=60,
@ -341,10 +353,19 @@ class WorkerConfig(BaseModel):
shutdown_grace_s: float = Field(
default=60.0,
description="On SIGINT/SIGTERM, how long to wait for in-flight jobs to "
"finish before forcing cancellation. Cancelled jobs stay 'claimed' in "
"the queue; the reaper resets them after claim_timeout_s.",
"finish before forcing cancellation. Cancelled jobs are released back "
"to `queued` for immediate re-claim.",
)
@model_validator(mode="after")
def _check_heartbeat_cadence(self) -> "WorkerConfig":
if self.heartbeat_interval_s > self.lease_ttl_s / 3:
raise ValueError(
"heartbeat_interval_s must be <= lease_ttl_s / 3 so a live "
"worker renews its lease several times before it could expire"
)
return self
class APIConfig(BaseModel):
"""HTTP control plane settings for the ingester."""

13
haiku_rag_slim/haiku/rag/ingester/app.py
View file

@ -133,7 +133,8 @@ class IngesterApp:
worker_count=ingester_cfg.workers.worker_count,
retry_policy=retry,
poll_idle_interval_s=ingester_cfg.workers.poll_idle_interval_s,
claim_timeout_s=ingester_cfg.workers.claim_timeout_s,
lease_ttl_s=ingester_cfg.workers.lease_ttl_s,
heartbeat_interval_s=ingester_cfg.workers.heartbeat_interval_s,
reaper_interval_s=ingester_cfg.workers.reaper_interval_s,
retention_s=(
ingester_cfg.queue.retention_days * 86400
@ -193,12 +194,14 @@ class IngesterApp:
try:
await asyncio.wait_for(self._pool.stop(), timeout=grace_s)
except TimeoutError:
# In-flight jobs stay 'claimed'; the reaper resets them after
# claim_timeout_s on the next start.
# Cancelling the workers triggers their cancel-cleanup, which
# releases each in-flight job back to `queued` (drained just below).
# Any release that doesn't land has its lease stop being renewed, so
# the reaper reclaims it after lease_ttl_s.
logger.warning(
"Shutdown grace of %.1fs elapsed with jobs still in flight; "
"cancelling — they'll be reclaimed after claim_timeout_s on "
"next start",
"cancelling — they'll be released back to the queue (or "
"reclaimed after lease_ttl_s if release doesn't land)",
grace_s,
)
landed = await self._pool.drain_pending_releases(timeout=2.0)

99
haiku_rag_slim/haiku/rag/ingester/workers/pool.py
View file

@ -28,9 +28,11 @@ _WORKER_BREAKER_COOLDOWN_S = 60.0
class WorkerPool:
"""`worker_count` async tasks each pull jobs from the queue and run them
through `run_job`. A reaper task resets claims older than
`claim_timeout_s` so a crashed worker doesn't strand its job. Lifecycle:
build it, await start(), let it run, await stop().
through `run_job`. While a worker processes a job it renews the job's lease
on a `heartbeat_interval_s` cadence; a reaper task resets claims whose lease
has gone stale (`lease_ttl_s`) so a crashed worker doesn't strand its job,
without reaping jobs that are merely slow. Lifecycle: build it, await
start(), let it run, await stop().
"""
def __init__(
@ -42,7 +44,8 @@ class WorkerPool:
worker_count: int = 4,
retry_policy: RetryPolicy | None = None,
poll_idle_interval_s: float = 1.0,
claim_timeout_s: int = 1800,
lease_ttl_s: int = 120,
heartbeat_interval_s: int = 30,
reaper_interval_s: int = 60,
retention_s: int | None = None,
sources: "list[Source] | None" = None,
@ -54,7 +57,8 @@ class WorkerPool:
self._worker_count = worker_count
self._retry = retry_policy or RetryPolicy()
self._poll_idle_s = poll_idle_interval_s
self._claim_timeout_s = claim_timeout_s
self._lease_ttl_s = lease_ttl_s
self._heartbeat_interval_s = heartbeat_interval_s
self._reaper_interval_s = reaper_interval_s
self._retention_s = retention_s
self._sources: list[Source] = list(sources) if sources else []
@ -69,18 +73,46 @@ class WorkerPool:
self._stop = asyncio.Event()
self._workers: list[asyncio.Task] = []
self._reaper: asyncio.Task | None = None
self._heartbeat: asyncio.Task | None = None
# job_id -> claimed_by for jobs currently being processed by this pool;
# the heartbeat renews exactly these leases.
self._inflight: dict[str, str] = {}
# Set whenever _inflight is empty so the heartbeat can leave promptly
# once stopping instead of sleeping out a full interval.
self._idle = asyncio.Event()
self._idle.set()
self._pending_releases: set[asyncio.Task] = set()
self._breakers: dict[str, CircuitBreaker] = {}
def _worker_id(self, i: int) -> str:
return f"{self._instance}-{i}"
def _track_inflight(self, job_id: str, worker_id: str) -> None:
self._inflight[job_id] = worker_id
self._idle.clear()
def _untrack_inflight(self, job_id: str, worker_id: str) -> None:
# Only drop our own entry: if the reaper reset this claim and a sibling
# worker re-claimed it, the entry now belongs to that worker and must
# keep being renewed — our exit must not evict it.
if self._inflight.get(job_id) == worker_id:
del self._inflight[job_id]
if not self._inflight:
self._idle.set()
@property
def live_workers(self) -> int:
"""Worker tasks that are still running. Equal to worker_count under
normal operation; less when a worker has crashed."""
return sum(1 for t in self._workers if not t.done())
@property
def heartbeat_alive(self) -> bool:
"""Whether the lease-renewal task is running. Once the pool is started
this stays True until stop(); a False here while workers are live means
in-flight leases are no longer being renewed and may be reaped."""
return self._heartbeat is not None and not self._heartbeat.done()
@property
def breaker_open(self) -> bool:
return any(b.is_open for b in self._breakers.values())
@ -112,7 +144,7 @@ class WorkerPool:
# reboot) so fresh workers can take them over. Scoped by the lease TTL,
# not 0: a peer process sharing the queue may hold live claims, and
# those must not be wiped out from under it on our startup.
reset = await self._jobs.reap_stale(lease_ttl_seconds=self._claim_timeout_s)
reset = await self._jobs.reap_stale(lease_ttl_seconds=self._lease_ttl_s)
if reset:
logger.info("Boot-reaped %d stale claim(s) from previous process", reset)
for i in range(self._worker_count):
@ -120,20 +152,35 @@ class WorkerPool:
asyncio.create_task(self._worker_loop(self._worker_id(i)))
)
self._reaper = asyncio.create_task(self._reaper_loop())
self._heartbeat = asyncio.create_task(self._heartbeat_loop())
self._heartbeat.add_done_callback(self._on_heartbeat_done)
def _on_heartbeat_done(self, task: asyncio.Task) -> None:
if task.cancelled():
return
exc = task.exception()
if exc is not None:
logger.error("Heartbeat task died: %r", exc)
elif not self._stop.is_set():
logger.error("Heartbeat task exited while the pool was still running")
async def stop(self) -> None:
self._stop.set()
# Wake workers parked on job_available.wait() so they notice _stop
# immediately instead of sleeping out the full poll_idle interval.
# _stop also wakes the heartbeat's interval sleep.
async with self._jobs.job_available:
self._jobs.job_available.notify_all()
tasks = list(self._workers)
if self._reaper is not None:
tasks.append(self._reaper)
if self._heartbeat is not None:
tasks.append(self._heartbeat)
if tasks:
await asyncio.gather(*tasks, return_exceptions=True)
self._workers.clear()
self._reaper = None
self._heartbeat = None
async def drain_pending_releases(self, timeout: float = 2.0) -> int:
"""Wait for any in-flight cancel-cleanup release Tasks to finish.
@ -180,7 +227,7 @@ class WorkerPool:
await self._sleep_or_stop(self._reaper_interval_s)
if self._stop.is_set():
return
reset = await self._jobs.reap_stale(self._claim_timeout_s)
reset = await self._jobs.reap_stale(self._lease_ttl_s)
if reset:
logger.info("Reaper reset %d stale claim(s)", reset)
if self._retention_s is not None:
@ -188,6 +235,35 @@ class WorkerPool:
if pruned:
logger.info("Reaper pruned %d terminal job(s)", pruned)
async def _heartbeat_loop(self) -> None:
"""Renew the lease on this pool's in-flight jobs so the reaper leaves
them alone while they are still being processed. Continues renewing
through graceful shutdown until the last job drains, so a peer process
doesn't reap a job we're still finishing. A forced shutdown cancels
this task along with the workers, which is correct — they are no longer
draining gracefully."""
while True:
if self._inflight:
try:
await self._jobs.renew_claims(dict(self._inflight))
except Exception:
logger.exception(
"Lease renewal failed; retrying on the next heartbeat"
)
if self._stop.is_set() and not self._inflight:
return
if self._stop.is_set():
# Draining: keep the renewal cadence, but leave as soon as the
# last in-flight job finishes instead of sleeping it out.
try:
await asyncio.wait_for(
self._idle.wait(), timeout=self._heartbeat_interval_s
)
except TimeoutError:
pass
else:
await self._sleep_or_stop(self._heartbeat_interval_s)
async def _sleep_or_stop(self, seconds: float) -> None:
try:
await asyncio.wait_for(self._stop.wait(), timeout=seconds)
@ -197,6 +273,15 @@ class WorkerPool:
async def _process(self, job: Job) -> None:
assert job.claimed_by is not None, "_process only runs on claimed jobs"
worker_id = job.claimed_by
# Track before any await so the heartbeat renews this job's lease for
# its whole lifetime; untrack on every exit (success, error, cancel).
self._track_inflight(job.id, worker_id)
try:
await self._run_job_lifecycle(job, worker_id)
finally:
self._untrack_inflight(job.id, worker_id)
async def _run_job_lifecycle(self, job: Job, worker_id: str) -> None:
started = time.monotonic()
logger.info("Processing %s %s (job %s)", job.op.value, job.uri, job.id)
try:

27
tests/ingester/test_config.py
View file

@ -12,6 +12,7 @@ from haiku.rag.config import (
IngesterConfig,
RetryPolicyConfig,
S3SourceConfig,
WorkerConfig,
)
@ -20,11 +21,37 @@ def test_default_ingester_config_has_sane_values():
assert cfg.sources == []
assert cfg.workers.worker_count == 4
assert cfg.workers.retry.max_attempts == 5
assert cfg.workers.lease_ttl_s == 120
assert cfg.workers.heartbeat_interval_s == 30
assert cfg.api.enabled is True
assert cfg.api.port == 8765
assert cfg.api.root_path == ""
def test_worker_config_rejects_heartbeat_too_close_to_lease():
with pytest.raises(ValidationError, match="heartbeat_interval_s"):
WorkerConfig(lease_ttl_s=60, heartbeat_interval_s=30)
def test_worker_config_accepts_heartbeat_within_a_third_of_lease():
cfg = WorkerConfig(lease_ttl_s=60, heartbeat_interval_s=20)
assert cfg.heartbeat_interval_s == 20
def test_worker_config_rejects_non_positive_timings():
with pytest.raises(ValidationError):
WorkerConfig(lease_ttl_s=0)
with pytest.raises(ValidationError):
WorkerConfig(heartbeat_interval_s=0)
def test_worker_config_rejects_unknown_field():
"""The former claim_timeout_s (and any typo) must fail loudly rather than
being silently ignored."""
with pytest.raises(ValidationError, match="claim_timeout_s"):
WorkerConfig(claim_timeout_s=1800) # ty: ignore[unknown-argument]
@pytest.mark.parametrize(
"raw, expected",
[

158
tests/ingester/test_workers.py
View file

@ -24,7 +24,8 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool:
worker_count=kwargs.pop("worker_count", 2),
poll_idle_interval_s=kwargs.pop("poll_idle_interval_s", 0.05),
reaper_interval_s=kwargs.pop("reaper_interval_s", 60),
claim_timeout_s=kwargs.pop("claim_timeout_s", 60),
lease_ttl_s=kwargs.pop("lease_ttl_s", 60),
heartbeat_interval_s=kwargs.pop("heartbeat_interval_s", 30),
retention_s=kwargs.pop("retention_s", None),
retry_policy=kwargs.pop("retry_policy", RetryPolicy()),
sources=kwargs.pop("sources", None),
@ -400,7 +401,7 @@ async def test_shutdown_grace_lets_inflight_job_complete(client, jobs, sync):
async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync):
"""When grace elapses and the worker is cancelled mid-job, the claim is
released back to 'queued' so the next process can pick it up immediately
— no waiting on the reaper's claim_timeout_s."""
— no waiting on the reaper's lease_ttl_s."""
cancelled = asyncio.Event()
async def _hangs_forever(*args, **kwargs):
@ -432,6 +433,157 @@ async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync):
assert refreshed.attempts == 0
# --- heartbeat / lease renewal ---
@pytest.mark.asyncio
async def test_heartbeat_keeps_long_job_from_being_reaped(client, jobs, sync):
"""A job that takes longer than lease_ttl_s is kept alive by the heartbeat,
so the reaper never resets it and it completes exactly once. Without
renewal the reaper would reset the claim mid-flight and the final
mark_succeeded would be a no-op, leaving the job back in 'queued'."""
calls = 0
async def _slow(*args, **kwargs):
nonlocal calls
calls += 1
await asyncio.sleep(0.8)
return Document(
id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"}
)
client.create_document_from_source.side_effect = _slow
job = await jobs.enqueue("src", "u", JobOp.UPSERT)
assert job is not None
pool = _pool(
client,
jobs,
sync,
worker_count=1,
lease_ttl_s=0.3,
heartbeat_interval_s=0.05,
reaper_interval_s=0.05,
)
await pool.start()
try:
for _ in range(100):
refreshed = await jobs.get_job(job.id)
if refreshed is not None and refreshed.status is JobStatus.SUCCEEDED:
break
await asyncio.sleep(0.05)
finally:
await pool.stop()
refreshed = await jobs.get_job(job.id)
assert refreshed is not None
assert refreshed.status is JobStatus.SUCCEEDED
assert calls == 1
@pytest.mark.asyncio
async def test_heartbeat_continues_during_graceful_stop(client, jobs, sync):
"""While stop() waits for an in-flight job to drain, the heartbeat keeps
renewing its lease so a peer process doesn't reap it mid-drain."""
release = asyncio.Event()
async def _block(*args, **kwargs):
await release.wait()
return Document(
id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"}
)
client.create_document_from_source.side_effect = _block
job = await jobs.enqueue("src", "u", JobOp.UPSERT)
assert job is not None
pool = _pool(
client,
jobs,
sync,
worker_count=1,
lease_ttl_s=5,
heartbeat_interval_s=0.05,
reaper_interval_s=60,
)
await pool.start()
stop_task: asyncio.Task | None = None
try:
for _ in range(100):
refreshed = await jobs.get_job(job.id)
if refreshed is not None and refreshed.status is JobStatus.CLAIMED:
break
await asyncio.sleep(0.02)
assert refreshed is not None and refreshed.status is JobStatus.CLAIMED
# Begin a graceful stop; it blocks until the job drains.
stop_task = asyncio.create_task(pool.stop())
await asyncio.sleep(0.05)
first = (await jobs.get_job(job.id)).last_heartbeat_at
await asyncio.sleep(0.2)
second = (await jobs.get_job(job.id)).last_heartbeat_at
assert first is not None and second is not None
assert second > first, "heartbeat did not renew during graceful drain"
assert not stop_task.done()
release.set()
await asyncio.wait_for(stop_task, timeout=2.0)
finally:
release.set()
if stop_task is not None and not stop_task.done():
await stop_task
final = await jobs.get_job(job.id)
assert final is not None and final.status is JobStatus.SUCCEEDED
@pytest.mark.asyncio
async def test_stale_worker_exit_keeps_a_siblings_renewal_entry(client, jobs, sync):
"""Same-pool reaper race: worker A is processing J, the reaper resets it,
sibling worker B re-claims J. A's eventual exit must not evict B's renewal
entry, or B's lease would stop being renewed and J could be reaped again."""
pool = _pool(client, jobs, sync, worker_count=2)
worker_a, worker_b = pool._worker_id(0), pool._worker_id(1)
pool._track_inflight("J", worker_a)
# Reaper reset + sibling re-claim: B now owns J's renewal entry.
pool._track_inflight("J", worker_b)
assert pool._inflight["J"] == worker_b
# A finally exits — must leave B's entry intact.
pool._untrack_inflight("J", worker_a)
assert pool._inflight.get("J") == worker_b
# B's own exit clears it.
pool._untrack_inflight("J", worker_b)
assert "J" not in pool._inflight
@pytest.mark.asyncio
async def test_forced_shutdown_cancels_heartbeat(client, jobs, sync):
"""When the shutdown grace elapses and stop() is cancelled, the heartbeat
task is cancelled along with the workers — they are no longer draining."""
async def _hangs_forever(*args, **kwargs):
await asyncio.sleep(60)
return Document(id="d", content="x", uri="u")
client.create_document_from_source.side_effect = _hangs_forever
job = await jobs.enqueue("src", "u", JobOp.UPSERT)
assert job is not None
pool = _pool(client, jobs, sync, worker_count=1)
await pool.start()
await asyncio.sleep(0.05)
assert pool.heartbeat_alive
with pytest.raises(TimeoutError):
await asyncio.wait_for(pool.stop(), timeout=0.2)
# Let the cancellation settle.
await asyncio.sleep(0.05)
assert not pool.heartbeat_alive
@pytest.mark.asyncio
async def test_worker_loses_claim_to_reaper_does_not_write_sync_state(
client, jobs, sync
@ -900,7 +1052,7 @@ async def test_reaper_resets_stale_claims(client, jobs, sync, conn):
sync,
worker_count=0,
reaper_interval_s=0.05,
claim_timeout_s=1,
lease_ttl_s=1,
)
await pool.start()
try: