From 4707db780f1aad8aecc78fbe62f8e6d338f73f4d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 13:20:28 +0300 Subject: [PATCH] Renew job leases from the worker pool; lease-based reaping --- CHANGELOG.md | 12 +- haiku_rag_slim/haiku/rag/config/models.py | 43 +++-- haiku_rag_slim/haiku/rag/ingester/app.py | 13 +- .../haiku/rag/ingester/workers/pool.py | 99 ++++++++++- tests/ingester/test_config.py | 27 +++ tests/ingester/test_workers.py | 158 +++++++++++++++++- 6 files changed, 319 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 44fb73fb..7928b419 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 08622db7..7a10e6a4 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index dd08372c..321845c0 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 85876c9f..f270bcf5 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -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: diff --git a/tests/ingester/test_config.py b/tests/ingester/test_config.py index 36bcbf5b..066e68f6 100644 --- a/tests/ingester/test_config.py +++ b/tests/ingester/test_config.py @@ -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", [ diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 99837052..31a3c9ed 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -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: