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

Merge pull request #390 from ggozad/feat/ingester-run-batch

Browse source 
Add haiku-ingester run-batch, remove run-once
This commit is contained in:
Yiorgis Gozadinos 2026-06-01 10:34:57 +03:00 committed by GitHub
commit a0029e9ab9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 544 additions and 135 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -1,6 +1,14 @@
# Changelog
## [Unreleased]
### Added
- `haiku-ingester run-batch`: one discover sweep across every configured source, drains the queue, then exits. Orphan deletion requires a persisted `ingester.db`.
### Removed
- `haiku-ingester run-once`. Use `run-batch` to ingest configured sources, or `serve` for continuous operation.
## [0.51.0] - 2026-05-29
### Added

21
docs/ingester.md
View file

@ -307,19 +307,24 @@ ingester:
## Operating
### Smoke-test a single URI
### One-shot batch build
`run-once` bypasses the queue and runs a single Job through the
pipeline. Useful for sanity-checking a source before starting the
service.
`run-batch` runs a single discover sweep across every configured source,
drains the queue, then exits. New and changed resources are ingested,
resources that vanished from a source are deleted. The periodic poller
loops never start, so the run is deterministic and finishes as soon as the
queue is empty. This is the mode for building a database in CI or on a
schedule rather than running the service continuously.
```bash
haiku-ingester run-once /path/to/test.pdf
haiku-ingester run-once https://example.com/spec.pdf
haiku-ingester run-once s3://my-bucket/key.pdf
haiku-ingester run-batch
haiku-ingester run-batch --db rag.lancedb
```
Exit codes: `0` success, `1` transient error, `2` permanent error.
Orphan deletion compares each source against `sync_state` in the queue DB,
so persist `ingester.db` between runs for deletions to be detected. It exits
non-zero if any job dead-letters or a source's discovery sweep does not
complete.
### The queue

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

@ -1,9 +1,13 @@
import asyncio
import logging
import signal
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING
from pydantic import BaseModel
from haiku.rag.config import AppConfig
from haiku.rag.ingester.pollers.manager import PollerManager
from haiku.rag.ingester.queue.migrations import open_queue
@ -17,6 +21,15 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
class BatchReport(BaseModel):
"""Outcome of a one-shot batch run: terminal job counts after the queue
drained, plus any sources whose discovery sweep did not complete."""
succeeded: int = 0
dead: int = 0
failed_sweeps: list[str] = []
class IngesterApp:
"""Top-level lifecycle for the production ingester.
@ -32,11 +45,13 @@ class IngesterApp:
self._sync: SyncStateRepo | None = None
self._pool: WorkerPool | None = None
self._pollers: PollerManager | None = None
self._client = None
async def serve(self, *, api: bool = True) -> None:
"""Run pollers + workers (and the HTTP API when enabled) until a
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
@asynccontextmanager
async def _resources(self):
"""Open the queue connection and construct the repos, client, pollers
and worker pool. Yields with everything built but nothing started —
callers own the start/stop lifecycle. Closes the client and queue
connection on exit."""
from haiku.rag.client import HaikuRAG
from haiku.rag.converters import get_converter
@ -58,13 +73,12 @@ class IngesterApp:
jitter=ingester_cfg.workers.retry.jitter,
)
# The ingester is the sole writer for its LanceDB target;
# create on first start so docker-compose / fresh deployments
# don't require a manual `haiku-rag init`.
# The ingester is the sole writer for its LanceDB target; create on
# first start so docker-compose / fresh deployments don't require a
# manual `haiku-rag init`.
async with HaikuRAG(
self._db_path, config=self._config, create=True
) as client:
self._client = client
self._pollers = PollerManager(
configs=ingester_cfg.sources,
job_repo=self._jobs,
@ -86,73 +100,7 @@ class IngesterApp:
# HTTP / WebDAV / S3 fetches reuse credentials.
sources=self._pollers.sources,
)
stop_event = asyncio.Event()
loop = asyncio.get_running_loop()
for sig in (signal.SIGINT, signal.SIGTERM):
try:
loop.add_signal_handler(sig, stop_event.set)
except NotImplementedError:
# Windows; signal handlers unavailable in asyncio.
pass
await self._pollers.start()
await self._pool.start()
# Log the docling-serve fleet size when relevant so the
# operator can eyeball the worker/instance ratio. The convert
# phase is usually the throughput ceiling.
proc = self._config.processing
uses_docling_serve = (
proc.converter == "docling-serve" or proc.chunker == "docling-serve"
)
if uses_docling_serve:
logger.info(
"Ingester running: %d worker(s), %d source(s), "
"%d docling-serve instance(s)",
ingester_cfg.workers.worker_count,
len(ingester_cfg.sources),
len(self._config.providers.docling_serve.base_urls),
)
else:
logger.info(
"Ingester running: %d worker(s), %d source(s)",
ingester_cfg.workers.worker_count,
len(ingester_cfg.sources),
)
api_task, api_server = await self._maybe_start_api(api)
try:
await stop_event.wait()
finally:
logger.info("Shutting down ingester")
if api_server is not None:
api_server.should_exit = True
if api_task is not None:
await asyncio.gather(api_task, return_exceptions=True)
await self._pollers.stop()
grace_s = ingester_cfg.workers.shutdown_grace_s
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.
logger.warning(
"Shutdown grace of %.1fs elapsed with jobs still "
"in flight; cancelling — they'll be reclaimed after "
"claim_timeout_s on next start",
grace_s,
)
# Drain any cancel-cleanup release Tasks the worker pool
# spawned but didn't get to await (timeout path). They
# need the queue connection that the outer finally is
# about to close.
landed = await self._pool.drain_pending_releases(timeout=2.0)
if landed:
logger.info(
"Drained %d cancel-cleanup release(s) before close",
landed,
)
yield
finally:
# Close the queue connection unconditionally. aiosqlite runs the
# underlying sqlite3 in a background thread; leaving it open holds
@ -162,6 +110,110 @@ class IngesterApp:
await self._queue_conn.close()
self._queue_conn = None
async def _stop_pool(self) -> None:
"""Stop the worker pool, honouring the shutdown grace, then drain any
cancel-cleanup release tasks before the queue connection closes."""
assert self._pool is not None
grace_s = self._config.ingester.workers.shutdown_grace_s
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.
logger.warning(
"Shutdown grace of %.1fs elapsed with jobs still in flight; "
"cancelling — they'll be reclaimed after claim_timeout_s on "
"next start",
grace_s,
)
landed = await self._pool.drain_pending_releases(timeout=2.0)
if landed:
logger.info("Drained %d cancel-cleanup release(s) before close", landed)
async def serve(self, *, api: bool = True) -> None:
"""Run pollers + workers (and the HTTP API when enabled) until a
SIGINT/SIGTERM is received. Drains in-flight work on shutdown."""
ingester_cfg = self._config.ingester
async with self._resources():
assert self._pollers is not None and self._pool is not None
await self._pollers.start()
await self._pool.start()
# Log the docling-serve fleet size when relevant so the operator
# can eyeball the worker/instance ratio. The convert phase is
# usually the throughput ceiling.
proc = self._config.processing
uses_docling_serve = (
proc.converter == "docling-serve" or proc.chunker == "docling-serve"
)
if uses_docling_serve:
logger.info(
"Ingester running: %d worker(s), %d source(s), "
"%d docling-serve instance(s)",
ingester_cfg.workers.worker_count,
len(ingester_cfg.sources),
len(self._config.providers.docling_serve.base_urls),
)
else:
logger.info(
"Ingester running: %d worker(s), %d source(s)",
ingester_cfg.workers.worker_count,
len(ingester_cfg.sources),
)
api_task, api_server = await self._maybe_start_api(api)
stop_event = asyncio.Event()
loop = asyncio.get_running_loop()
for sig in (signal.SIGINT, signal.SIGTERM):
try:
loop.add_signal_handler(sig, stop_event.set)
except NotImplementedError: # pragma: no cover - Windows only
# Windows; signal handlers unavailable in asyncio.
pass
try:
await stop_event.wait()
finally:
logger.info("Shutting down ingester")
if api_server is not None:
api_server.should_exit = True
if api_task is not None:
await asyncio.gather(api_task, return_exceptions=True)
await self._pollers.stop()
await self._stop_pool()
async def run_batch(self) -> BatchReport:
"""Run one discover() sweep across every configured source, drain the
queue to completion, then stop. Unlike `serve`, the periodic poller
loops never start — discovery is driven explicitly, so the run is
deterministic and exits as soon as the queue is empty."""
async with self._resources():
assert (
self._pollers is not None
and self._pool is not None
and self._jobs is not None
)
# A persisted queue carries terminal rows from previous runs, and
# a recovered URI's dead row gets pruned mid-run, so the report
# counts only jobs that completed at or after this start instant.
started_at = datetime.now(UTC)
await self._pool.start()
try:
failed_sweeps = await self._pollers.sweep_all()
while True:
counts = await self._jobs.counts_by_status()
if not counts.get("queued") and not counts.get("claimed"):
break
await asyncio.sleep(0.1)
completed = await self._jobs.counts_by_status_since(started_at)
return BatchReport(
succeeded=completed.get("succeeded", 0),
dead=completed.get("dead", 0),
failed_sweeps=failed_sweeps,
)
finally:
await self._stop_pool()
async def _maybe_start_api(self, api: bool):
"""Spin up the FastAPI control plane on an asyncio task. Returns
(task, server) or (None, None) when the API is disabled."""

67
haiku_rag_slim/haiku/rag/ingester/cli.py
View file

@ -1,7 +1,5 @@
import asyncio
import sys
import uuid
from datetime import UTC, datetime
from pathlib import Path
import typer
@ -9,7 +7,6 @@ from dotenv import find_dotenv, load_dotenv
load_dotenv(find_dotenv(usecwd=True))
from haiku.rag.client import HaikuRAG # noqa: E402
from haiku.rag.config import ( # noqa: E402
AppConfig,
find_config_file,
@ -18,10 +15,7 @@ from haiku.rag.config import ( # noqa: E402
set_config,
)
from haiku.rag.ingester.app import IngesterApp # noqa: E402
from haiku.rag.ingester.exceptions import PermanentError, TransientError # noqa: E402
from haiku.rag.ingester.queue.migrations import open_queue # noqa: E402
from haiku.rag.ingester.queue.models import Job, JobOp, JobStatus # noqa: E402
from haiku.rag.ingester.workers.pipeline import run_job # noqa: E402
from haiku.rag.logging import configure_cli_logging # noqa: E402
from haiku.rag.store.exceptions import ( # noqa: E402
MigrationRequiredError,
@ -163,59 +157,30 @@ def serve(
asyncio.run(app.serve(api=not no_api))
@_cli.command("run-once")
def run_once(
uri: str = typer.Argument(..., help="URI to ingest (file://, http(s)://, s3://)."),
@_cli.command("run-batch")
def run_batch(
db: Path | None = typer.Option(
None,
"--db",
help="LanceDB path (overrides config.storage.data_dir).",
),
delete: bool = typer.Option(
False, "--delete", help="Run a delete op instead of upsert."
),
) -> None:
"""Build an ad-hoc Job and run it through the worker pipeline once.
Bypasses the queue — does NOT enqueue. Useful for smoke-testing the
Source adapter + pipeline path without spinning up the full pool.
"""
asyncio.run(_run_once(get_config(), uri, db, delete))
"""Run one discover sweep across every configured source, drain the queue,
then exit. New and changed resources are ingested, resources that vanished
from a source are deleted. Exits non-zero if any job dead-letters or a
source's sweep does not complete."""
asyncio.run(_run_batch(get_config(), db))
async def _run_once(
app_config: AppConfig, uri: str, db_path: Path | None, delete: bool
) -> None:
async def _run_batch(app_config: AppConfig, db_path: Path | None) -> None:
db = _resolve_db_path(app_config, db_path)
now = datetime.now(UTC)
job = Job(
id=f"adhoc-{uuid.uuid4()}",
source_id="adhoc",
uri=uri,
op=JobOp.DELETE if delete else JobOp.UPSERT,
status=JobStatus.CLAIMED,
attempts=1,
max_attempts=1,
enqueued_at=now,
scheduled_at=now,
claimed_at=now,
claimed_by="run-once",
)
async with HaikuRAG(db, config=app_config) as client:
try:
result = await run_job(client, job)
except PermanentError as e:
typer.echo(f"PERMANENT: {e}", err=True)
raise typer.Exit(2) from e
except TransientError as e:
typer.echo(f"TRANSIENT: {e}", err=True)
raise typer.Exit(1) from e
if result.deleted:
typer.echo(f"Deleted document at {uri}")
else:
app = IngesterApp(config=app_config, db_path=db)
report = await app.run_batch()
typer.echo(f"Batch complete: {report.succeeded} succeeded, {report.dead} dead")
if report.failed_sweeps:
typer.echo(
f"Ingested {uri}: document_id={result.document_id} "
f"revision={result.revision} md5={result.content_hash}"
f"Sources that failed to sweep: {', '.join(report.failed_sweeps)}",
err=True,
)
if report.dead or report.failed_sweeps:
raise typer.Exit(1)

12
haiku_rag_slim/haiku/rag/ingester/pollers/manager.py
View file

@ -77,6 +77,18 @@ class PollerManager:
poller._stop.clear()
self._tasks.append(asyncio.create_task(poller.run()))
async def sweep_all(self) -> list[str]:
"""Run one discover() sweep on every poller, sequentially. Used by
one-shot batch runs that drive discovery explicitly rather than
through the periodic loop. Returns the source ids whose sweep did not
complete (discovery failed, circuit open, or pending work already
queued) so callers can treat a one-shot run as failed."""
failed: list[str] = []
for poller in self._pollers:
if not await poller._sweep_once():
failed.append(poller.source_id)
return failed
async def stop(self) -> None:
for poller in self._pollers:
await poller.stop()

14
haiku_rag_slim/haiku/rag/ingester/queue/repository.py
View file

@ -294,6 +294,20 @@ class JobRepo:
rows = await cursor.fetchall()
return {row["status"]: row["n"] for row in rows}
async def counts_by_status_since(self, since: datetime) -> dict[str, int]:
"""status -> count of jobs that reached a terminal state at or after
`since` (by completed_at). Only succeeded/dead set completed_at, so
those are the only keys returned. Lets a one-shot batch report the
work it finished, independent of terminal rows from earlier runs."""
async with self._lock:
async with self._conn.execute(
"SELECT status, COUNT(*) AS n FROM jobs "
"WHERE completed_at >= ? GROUP BY status",
(since.isoformat(),),
) as cursor:
rows = await cursor.fetchall()
return {row["status"]: row["n"] for row in rows}
async def count_succeeded_since(self, seconds: int) -> int:
"""How many jobs reached `succeeded` in the last `seconds` seconds.
Drives the dashboard's rolling-throughput chips."""

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

@ -114,7 +114,7 @@ class WorkerPool:
async def drain_once(self, worker_id: str = "drain") -> int:
"""Drain every currently-claimable job to completion on the calling
coroutine. Used by run-once and by tests; not by `start()`."""
coroutine. Used by tests; not by `start()`."""
processed = 0
while True:
job = await self._jobs.claim_next(worker_id)

46
tests/ingester/test_cli.py Normal file
View file

@ -0,0 +1,46 @@
"""haiku-ingester run-batch CLI: echoes the batch report and exits non-zero
when any job dead-letters. IngesterApp is patched so no real ingestion runs."""
from unittest.mock import AsyncMock
from typer.testing import CliRunner
from haiku.rag.ingester.app import BatchReport
from haiku.rag.ingester.cli import _cli as cli
runner = CliRunner()
def _fake_app(report: BatchReport, monkeypatch) -> AsyncMock:
fake = AsyncMock()
fake.run_batch.return_value = report
monkeypatch.setattr("haiku.rag.ingester.cli.IngesterApp", lambda **_: fake)
return fake
def test_run_batch_reports_and_exits_zero(monkeypatch):
fake = _fake_app(BatchReport(succeeded=3, dead=0), monkeypatch)
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
assert result.exit_code == 0
assert "3 succeeded, 0 dead" in result.output
fake.run_batch.assert_awaited_once()
def test_run_batch_exits_nonzero_when_dead(monkeypatch):
_fake_app(BatchReport(succeeded=1, dead=2), monkeypatch)
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
assert result.exit_code == 1
assert "2 dead" in result.output
def test_run_batch_exits_nonzero_when_sweep_fails(monkeypatch):
_fake_app(BatchReport(succeeded=2, dead=0, failed_sweeps=["docs"]), monkeypatch)
result = runner.invoke(cli, ["run-batch", "--db", "x.lancedb"])
assert result.exit_code == 1
assert "failed to sweep: docs" in result.output

307
tests/ingester/test_run_batch.py Normal file
View file

@ -0,0 +1,307 @@
"""IngesterApp lifecycle: run_batch (one sweep -> drain -> exit) and serve
(pollers + workers + API until shutdown). The document store (HaikuRAG) is
patched out — the behavior under test is the orchestration and orphan
pruning, not embedding."""
import asyncio
from contextlib import asynccontextmanager
from pathlib import Path
from unittest.mock import AsyncMock
from urllib.parse import unquote, urlparse
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.client.exceptions import UnsupportedSourceError
from haiku.rag.config import (
APIConfig,
AppConfig,
FSSourceConfig,
IngesterConfig,
QueueConfig,
RetryPolicyConfig,
WorkerConfig,
)
from haiku.rag.ingester.app import IngesterApp
from haiku.rag.ingester.workers.pool import WorkerPool
from haiku.rag.store.models.document import Document
def _config(tmp_path, **worker_kwargs) -> AppConfig:
workers = WorkerConfig(
worker_count=2,
poll_idle_interval_s=0.05,
retry=RetryPolicyConfig(max_attempts=1),
**worker_kwargs,
)
return AppConfig(
ingester=IngesterConfig(
queue=QueueConfig(path=tmp_path / "queue.db"),
sources=[
FSSourceConfig(
type="fs",
id="local",
root=tmp_path,
poll_interval_s=3600.0,
)
],
workers=workers,
api=APIConfig(enabled=False),
),
)
def _mock_client() -> AsyncMock:
"""A HaikuRAG mock whose upserts return a fresh Document (with an md5 so
the worker records sync_state) and whose lookups return a Document with a
deterministic id so the DELETE path resolves."""
client = AsyncMock(spec=HaikuRAG)
counter = {"n": 0}
async def _create(uri, *_, **__):
counter["n"] += 1
# Mirror the real client: persist the FS revision (mtime_ns) so the
# poller's change-detection skips unchanged files on the next sweep.
path = Path(unquote(urlparse(uri).path))
return Document(
id=f"doc-{counter['n']}",
content="x",
uri=uri,
metadata={
"content_type": "text/markdown",
"md5": f"md5-{counter['n']}",
"source_revision": str(path.stat().st_mtime_ns),
},
)
async def _get_by_uri(uri):
return Document(id=f"doc-for-{uri}", content="x", uri=uri, metadata={})
client.create_document_from_source.side_effect = _create
client.get_document_by_uri.side_effect = _get_by_uri
return client
@pytest.fixture
def use_client(monkeypatch):
"""Make IngesterApp's internally-created HaikuRAG resolve to the given
mock client. Call again to swap the client for a later run in the same
test."""
def _install(client):
@asynccontextmanager
async def _cm(*_, **__):
yield client
monkeypatch.setattr("haiku.rag.client.HaikuRAG", lambda *a, **k: _cm())
return _install
@pytest.mark.asyncio
async def test_run_batch_drains_upserts(tmp_path, use_client):
(tmp_path / "a.md").write_text("hello")
(tmp_path / "b.md").write_text("world")
client = _mock_client()
use_client(client)
report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
).run_batch()
assert report.succeeded == 2
assert report.dead == 0
assert client.create_document_from_source.await_count == 2
ingested = {
call.args[0] for call in client.create_document_from_source.await_args_list
}
assert ingested == {
(tmp_path / "a.md").as_uri(),
(tmp_path / "b.md").as_uri(),
}
@pytest.mark.asyncio
async def test_run_batch_prunes_orphans(tmp_path, use_client):
(tmp_path / "a.md").write_text("hello")
(tmp_path / "b.md").write_text("world")
client = _mock_client()
use_client(client)
config = _config(tmp_path)
db_path = tmp_path / "db.lancedb"
# First batch ingests both files and records sync_state for each.
first = await IngesterApp(config=config, db_path=db_path).run_batch()
assert first.succeeded == 2
client.delete_document.assert_not_awaited()
# b.md disappears from the source tree. The next sweep sees it in
# sync_state but not on disk -> enqueues a DELETE for it.
(tmp_path / "b.md").unlink()
second = await IngesterApp(config=config, db_path=db_path).run_batch()
# a.md is unchanged (same mtime) so it's not re-ingested; only the orphan
# delete runs.
assert client.create_document_from_source.await_count == 2
client.delete_document.assert_awaited_once()
assert second.succeeded == 1
assert second.dead == 0
@pytest.mark.asyncio
async def test_run_batch_reports_dead_on_permanent_failure(tmp_path, use_client):
(tmp_path / "a.md").write_text("hello")
client = _mock_client()
client.create_document_from_source.side_effect = UnsupportedSourceError("nope")
use_client(client)
report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
).run_batch()
assert report.succeeded == 0
assert report.dead == 1
@pytest.mark.asyncio
async def test_run_batch_recovered_doc_is_not_counted_as_dead(tmp_path, use_client):
"""A doc that dead-lettered in an earlier run and succeeds in this one
prunes its dead row. The report counts only this run's terminal jobs, so
the recovery shows as a success, never a negative or stale dead count."""
(tmp_path / "a.md").write_text("hello")
config = _config(tmp_path)
db_path = tmp_path / "db.lancedb"
failing = _mock_client()
failing.create_document_from_source.side_effect = UnsupportedSourceError("nope")
use_client(failing)
first = await IngesterApp(config=config, db_path=db_path).run_batch()
assert first.dead == 1
healthy = _mock_client()
use_client(healthy)
second = await IngesterApp(config=config, db_path=db_path).run_batch()
assert second.dead == 0
assert second.succeeded == 1
@pytest.mark.asyncio
async def test_run_batch_reports_failed_sweep(
tmp_path, use_client, monkeypatch, caplog
):
"""A source whose discover() raises is reported in failed_sweeps so the
run can be treated as failed rather than a silent empty success."""
(tmp_path / "a.md").write_text("hello")
use_client(_mock_client())
async def _failing_discover(self, **kwargs):
raise RuntimeError("discover blew up")
yield # unreachable; makes this an async generator
monkeypatch.setattr(
"haiku.rag.ingester.sources.fs.FSSource.discover", _failing_discover
)
with caplog.at_level("ERROR", logger="haiku.rag.ingester.pollers.base"):
report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
).run_batch()
assert report.failed_sweeps == ["local"]
assert report.succeeded == 0
assert report.dead == 0
assert "discover() failed" in caplog.text
@pytest.mark.asyncio
async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
client = _mock_client()
use_client(client)
report = await asyncio.wait_for(
IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb"
).run_batch(),
timeout=5.0,
)
assert report.succeeded == 0
assert report.dead == 0
client.create_document_from_source.assert_not_awaited()
async def _wait_until(predicate, *, timeout: float = 5.0):
deadline = asyncio.get_running_loop().time() + timeout
while not predicate():
if asyncio.get_running_loop().time() > deadline:
raise AssertionError("condition not reached within timeout")
await asyncio.sleep(0.02)
@pytest.mark.asyncio
@pytest.mark.parametrize("api", [True, False])
async def test_serve_starts_workers_pollers_and_shuts_down(tmp_path, use_client, api):
"""serve brings up pollers, workers and (when enabled) the HTTP API, then
tears them all down. Shutdown is driven here by cancelling the serve task,
which runs the same drain-and-close path as a SIGINT/SIGTERM."""
use_client(_mock_client())
config = _config(tmp_path)
config.ingester.api = APIConfig(enabled=api, host="127.0.0.1", port=0)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
task = asyncio.create_task(app.serve(api=api))
try:
await _wait_until(
lambda: (
app._pool is not None
and app._pool.live_workers > 0
and app._pollers is not None
and app._pollers.live_pollers > 0
)
)
finally:
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=5.0)
# Workers and pollers are stopped after the shutdown path runs.
assert app._pool is not None and app._pollers is not None
assert app._pool.live_workers == 0
assert app._pollers.live_pollers == 0
class _SlowPool(WorkerPool):
"""A WorkerPool whose stop never finishes within the grace, to exercise
_stop_pool's timeout path. The real wiring is bypassed since _stop_pool
only calls stop() and drain_pending_releases()."""
def __init__(self) -> None:
self.released = 0
async def stop(self) -> None:
await asyncio.sleep(1.0)
async def drain_pending_releases(self, timeout: float = 2.0) -> int:
self.released += 1
return 2
@pytest.mark.asyncio
async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog):
"""When a worker doesn't stop within the shutdown grace, _stop_pool logs a
warning and still drains any pending cancel-cleanup releases."""
config = _config(tmp_path, shutdown_grace_s=0.01)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb")
pool = _SlowPool()
app._pool = pool
with caplog.at_level("WARNING", logger="haiku.rag.ingester.app"):
await app._stop_pool()
assert pool.released == 1
assert "Shutdown grace" in caplog.text