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

Add haiku-ingester run-batch, remove run-once

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-29 17:00:59 +03:00
parent 62be23d130
commit 5ba7838b71
No known key found for this signature in database
8 changed files with 384 additions and 137 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

20
docs/ingester.md
View file

@ -307,19 +307,23 @@ 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.
### The queue

200
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,14 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
class BatchReport(BaseModel):
"""Outcome of a one-shot batch run: terminal job counts after the queue
drained."""
succeeded: int = 0
dead: int = 0
class IngesterApp:
"""Top-level lifecycle for the production ingester.
@ -32,11 +44,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 +72,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 +99,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 +109,109 @@ 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:
# 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:
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),
)
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."""

65
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,24 @@ 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."""
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:
typer.echo(
f"Ingested {uri}: document_id={result.document_id} "
f"revision={result.revision} md5={result.content_hash}"
)
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.dead:
raise typer.Exit(1)

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

@ -77,6 +77,13 @@ class PollerManager:
poller._stop.clear()
self._tasks.append(asyncio.create_task(poller.run()))
async def sweep_all(self) -> None:
"""Run one discover() sweep on every poller, sequentially. Used by
one-shot batch runs that drive discovery explicitly rather than
through the periodic loop."""
for poller in self._pollers:
await poller._sweep_once()
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)

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

@ -0,0 +1,205 @@
"""IngesterApp.run_batch: one discover sweep across configured sources,
drain the queue, then exit. The document store (HaikuRAG) is patched out —
the behavior under test is the sweep -> queue -> worker -> drain
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.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_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()