Coverage for src/lilbee/data/ingest/fanout.py: 100%
198 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""One ingest worker process per GPU, each over its own slice of the corpus."""
3from __future__ import annotations
5import asyncio
6import contextlib
7import logging
8import multiprocessing
9import os
10import queue
11import sys
12import time
13from dataclasses import dataclass
14from pathlib import Path
15from typing import TYPE_CHECKING, Literal
17from rich.progress import (
18 BarColumn,
19 MofNCompleteColumn,
20 Progress,
21 SpinnerColumn,
22 TextColumn,
23 TimeElapsedColumn,
24)
26from lilbee.core.config import active_config
27from lilbee.data.ingest.errors import error_reason
28from lilbee.data.types import ShardId, SyncResult
29from lilbee.runtime.cpu import available_cpu_count, cpu_quota
30from lilbee.runtime.engine_lock import ENGINE_DIR_ENV
31from lilbee.runtime.progress import (
32 BatchProgressEvent,
33 BatchStatus,
34 DetailedProgressCallback,
35 EventType,
36 ProgressEvent,
37)
39if TYPE_CHECKING:
40 from collections.abc import Sequence
41 from multiprocessing.process import BaseProcess
42 from multiprocessing.queues import Queue
43 from multiprocessing.synchronize import Event
45 from lilbee.core.config.model import Config
46 from lilbee.runtime.cancellation import CancelSignal
48log = logging.getLogger(__name__)
50# Per-worker state (store, skip markers, engine slots) under the parent data root.
51SHARDS_DIRNAME = "shards"
52_DATA_ROOT_ENV = "LILBEE_DATA"
53_CPU_QUOTA_ENV = "LILBEE_CPU_QUOTA"
55# Below this many files on disk a fan-out costs more than it saves: every worker
56# pays a fresh interpreter, its own engine and a store of its own.
57_MIN_FILES_FOR_FANOUT = 2000
59# Under two workers there is nothing to fan out to.
60_MIN_FANOUT_WORKERS = 2
62# How often a worker reports its counters to the parent.
63_REPORT_INTERVAL_S = 0.25
65# How long the parent sleeps between drains of the worker message queue.
66_DRAIN_INTERVAL_S = 0.1
68# Grace for the queue's feeder thread to flush a dead worker's last messages.
69_FINAL_DRAIN_S = 1.0
71# How long a worker gets to exit on its own before it is killed.
72_WORKER_EXIT_GRACE_S = 30.0
74# Where a worker's console output lands, under its own data root.
75WORKER_LOG_NAME = "sync.log"
78@dataclass(frozen=True)
79class ShardSpec:
80 """One worker's slice, its card, and the private state it owns."""
82 shard: ShardId
83 device: int
84 config: Config
85 engine_dir: Path
86 cpu_share: int
87 visible_devices: dict[str, str]
90@dataclass(frozen=True)
91class ShardOptions:
92 """What every worker of one fan-out is told about the run it belongs to."""
94 parent_pid: int
95 force_rebuild: bool = False
96 retry_skipped: bool = False
99@dataclass(frozen=True)
100class ShardProgress:
101 """A worker's counters as it works."""
103 kind: Literal["progress"]
104 index: int
105 done: int
106 planned: int
107 file: str
108 status: BatchStatus
111@dataclass(frozen=True)
112class ShardDone:
113 """A worker's verdict; *error* set means it produced no usable shard."""
115 kind: Literal["done"]
116 index: int
117 result: SyncResult | None
118 error: str | None
121ShardMessage = ShardProgress | ShardDone
124def resolve_process_count(devices: int) -> int:
125 """Ingest worker processes for this run; 1 keeps ingest in this process.
127 Auto (``ingest_processes = 0``) is one worker per visible card. An explicit
128 count is honored past the card count, since two workers on one card is a
129 legitimate configuration; they share that card's engine slot rather than
130 putting a second fleet on it.
131 """
132 configured = active_config().ingest_processes
133 if configured:
134 return max(1, configured)
135 return devices
138def plan_fanout() -> list[ShardSpec]:
139 """The workers for this sync, empty when it runs in this process."""
140 from lilbee.data.ingest.discovery import corpus_has_at_least
141 from lilbee.providers.fleet.gpu_env import apply_fleet_gpu_env
142 from lilbee.providers.fleet.replicas import gpu_device_count
144 # Applied before the cards are counted, so a gpu_devices pin is the space the
145 # workers are dealt in: without it they would be dealt cards the pin excludes.
146 apply_fleet_gpu_env()
147 devices = gpu_device_count()
148 processes = resolve_process_count(devices)
149 if processes < _MIN_FANOUT_WORKERS or not corpus_has_at_least(_MIN_FILES_FOR_FANOUT):
150 return []
151 return shard_specs(active_config(), processes, devices)
154def shard_specs(config: Config, processes: int, devices: int) -> list[ShardSpec]:
155 """One spec per worker, dividing the corpus, the cards and the CPU pools."""
156 from lilbee.providers.fleet.gpu_env import shard_visible_devices
158 cpu_share = max(1, cpu_quota() // processes)
159 plan_share = max(1, available_cpu_count() // processes)
160 root = config.data_root / SHARDS_DIRNAME
161 return [
162 ShardSpec(
163 shard=ShardId(index=index, count=processes),
164 device=index % devices,
165 config=_shard_config(config, root / f"w{index}", plan_share),
166 # Keyed by card, not by worker: workers sharing a card share one
167 # fleet, workers on different cards never see each other's.
168 engine_dir=root / f"gpu{index % devices}" / "engine",
169 cpu_share=cpu_share,
170 visible_devices=shard_visible_devices(index % devices),
171 )
172 for index in range(processes)
173 ]
176def _shard_config(config: Config, root: Path, plan_share: int) -> Config:
177 """*config* with a private data root and this worker's share of the CPU pools.
179 ``documents_dir`` and ``linked_roots`` are inherited: every worker reads the
180 one shared corpus and only its own state is private.
181 """
182 return config.model_copy(
183 update={
184 "data_root": root,
185 "lancedb_dir": root / "data" / "lancedb",
186 "ingest_workers": plan_share,
187 }
188 )
191def _apply_shard_env(spec: ShardSpec) -> None:
192 """Pin this process to the worker's card, engine slot, CPU share and log."""
193 os.environ.update(spec.visible_devices)
194 os.environ[ENGINE_DIR_ENV] = str(spec.engine_dir)
195 os.environ[_DATA_ROOT_ENV] = str(spec.config.data_root)
196 os.environ[_CPU_QUOTA_ENV] = str(spec.cpu_share)
197 _redirect_output(spec.config.data_root / WORKER_LOG_NAME)
200def _redirect_output(path: Path) -> None:
201 """Send this process's console output to *path*.
203 At the file descriptor, so the engine this worker spawns follows it: N
204 workers logging onto the parent's terminal is the pile of log files the one
205 aggregated bar exists to replace.
206 """
207 path.parent.mkdir(parents=True, exist_ok=True)
208 with path.open("ab", buffering=0) as handle:
209 os.dup2(handle.fileno(), sys.stdout.fileno())
210 os.dup2(handle.fileno(), sys.stderr.fileno())
213class _ShardReporter:
214 """Throttled relay of a worker's counters onto the parent's queue.
216 The counters are the pipeline's own: how much of this worker's slice is done
217 and how big that slice is. Counting files here instead would only re-derive
218 the first, and the per-file events carry no slice size -- FILE_START's total
219 is the plan so far, which grows all run.
220 """
222 def __init__(self, index: int, messages: Queue[ShardMessage]) -> None:
223 self._index = index
224 self._messages = messages
225 self._done = 0
226 self._planned = 0
227 self._last_sent = 0.0
229 def __call__(self, event_type: EventType, data: ProgressEvent) -> None:
230 if event_type is not EventType.BATCH_PROGRESS or not isinstance(data, BatchProgressEvent):
231 return
232 self._done = data.current
233 self._planned = data.total
234 now = time.monotonic()
235 if now - self._last_sent < _REPORT_INTERVAL_S:
236 return
237 self._last_sent = now
238 self._send(data.file, data.status)
240 def flush(self) -> None:
241 """Send the final counters past the throttle, so the bar lands on its total."""
242 self._send("", BatchStatus.INGESTED)
244 def _send(self, file: str, status: BatchStatus) -> None:
245 self._messages.put(
246 ShardProgress(
247 kind="progress",
248 index=self._index,
249 done=self._done,
250 planned=self._planned,
251 file=file,
252 status=status,
253 )
254 )
257class _Aggregate:
258 """Every worker's latest counters, as one set of totals."""
260 def __init__(self, on_progress: DetailedProgressCallback) -> None:
261 self._latest: dict[int, ShardProgress] = {}
262 self._on_progress = on_progress
264 def update(self, message: ShardProgress) -> tuple[int, int]:
265 """Record *message* and return the corpus-wide (done, planned)."""
266 self._latest[message.index] = message
267 done = sum(p.done for p in self._latest.values())
268 planned = sum(p.planned for p in self._latest.values())
269 self._on_progress(
270 EventType.BATCH_PROGRESS,
271 BatchProgressEvent(
272 file=message.file, status=message.status, current=done, total=planned
273 ),
274 )
275 return done, planned
278def _drain(messages: Queue[ShardMessage]) -> list[ShardMessage]:
279 """Every message queued right now, without blocking."""
280 drained: list[ShardMessage] = []
281 with contextlib.suppress(queue.Empty):
282 while True:
283 drained.append(messages.get_nowait())
284 return drained
287def _shard_progress_bar(quiet: bool) -> Progress:
288 """The one bar a fan-out reports on, disabled when the caller wants no output."""
289 return Progress(
290 SpinnerColumn(),
291 TextColumn("[progress.description]{task.description}"),
292 BarColumn(),
293 MofNCompleteColumn(),
294 TimeElapsedColumn(),
295 disable=quiet,
296 )
299async def _supervise(
300 workers: Sequence[BaseProcess],
301 messages: Queue[ShardMessage],
302 stop: Event,
303 *,
304 quiet: bool,
305 on_progress: DetailedProgressCallback,
306 cancel: CancelSignal | None,
307) -> dict[int, ShardDone]:
308 """Drain worker messages until every worker has reported, keeping one bar current."""
309 verdicts: dict[int, ShardDone] = {}
310 aggregate = _Aggregate(on_progress)
311 with _shard_progress_bar(quiet) as progress:
312 task = progress.add_task(f"Ingesting on {len(workers)} workers", total=None)
313 while len(verdicts) < len(workers):
314 for message in _drain(messages):
315 if message.kind == "done":
316 verdicts[message.index] = message
317 else:
318 done, planned = aggregate.update(message)
319 progress.update(task, completed=done, total=planned or None)
320 if cancel is not None and cancel.is_set():
321 stop.set()
322 if not any(worker.is_alive() for worker in workers):
323 verdicts.update(_final_verdicts(workers, messages, verdicts))
324 break
325 await asyncio.sleep(_DRAIN_INTERVAL_S)
326 return verdicts
329def _final_verdicts(
330 workers: Sequence[BaseProcess],
331 messages: Queue[ShardMessage],
332 verdicts: dict[int, ShardDone],
333) -> dict[int, ShardDone]:
334 """Verdicts still in flight once every worker has exited, plus one per silent death.
336 A worker the kernel killed (out of memory is the usual reason) reports
337 nothing, so its shard is recorded as failed rather than silently missing from
338 the merge.
339 """
340 time.sleep(_FINAL_DRAIN_S)
341 late = {m.index: m for m in _drain(messages) if m.kind == "done"}
342 for index, worker in enumerate(workers):
343 if index in verdicts or index in late:
344 continue
345 late[index] = ShardDone(
346 kind="done",
347 index=index,
348 result=None,
349 error=f"worker exited with code {worker.exitcode} before reporting",
350 )
351 return late
354def _stop_workers(workers: Sequence[BaseProcess], stop: Event) -> None:
355 """Ask every live worker to stop, then wait for it, then insist.
357 A worker owns a GPU fleet, and its teardown can outlast a TERM; a plain join
358 would hang the sync behind it instead of returning a result it already has.
359 """
360 stop.set()
361 for worker in workers:
362 if worker.is_alive():
363 worker.terminate()
364 worker.join(_WORKER_EXIT_GRACE_S)
365 if worker.is_alive():
366 log.warning("Ingest worker %s did not exit; killing it", worker.name)
367 worker.kill()
368 worker.join()
371async def run_workers(
372 specs: list[ShardSpec],
373 *,
374 options: ShardOptions,
375 quiet: bool,
376 on_progress: DetailedProgressCallback,
377 cancel: CancelSignal | None,
378) -> list[ShardDone]:
379 """Run every worker to completion and return their verdicts, in shard order."""
380 context = multiprocessing.get_context("spawn")
381 messages: Queue[ShardMessage] = context.Queue()
382 stop = context.Event()
383 workers = [
384 context.Process(
385 target=run_shard,
386 args=(spec, options, messages, stop),
387 name=f"lilbee-shard-{spec.shard.index}",
388 )
389 for spec in specs
390 ]
391 log.warning("Ingesting across %d worker processes, one per GPU", len(workers))
392 for worker in workers:
393 worker.start()
394 try:
395 verdicts = await _supervise(
396 workers, messages, stop, quiet=quiet, on_progress=on_progress, cancel=cancel
397 )
398 finally:
399 _stop_workers(workers, stop)
400 return [verdicts[index] for index in sorted(verdicts)]
403def aggregate_results(verdicts: list[ShardDone]) -> SyncResult:
404 """The one result a fan-out reports, unioned from every worker's."""
405 results = [verdict.result for verdict in verdicts if verdict.result is not None]
406 return SyncResult(
407 added=[name for r in results for name in r.added],
408 updated=[name for r in results for name in r.updated],
409 relocated=[name for r in results for name in r.relocated],
410 failed=[name for r in results for name in r.failed],
411 skipped=[name for r in results for name in r.skipped],
412 unchanged=sum(r.unchanged for r in results),
413 truncated=sum(r.truncated for r in results),
414 )
417def run_shard(
418 spec: ShardSpec, options: ShardOptions, messages: Queue[ShardMessage], stop: Event
419) -> None:
420 """Ingest this worker's slice in a fresh process, reporting onto *messages*."""
421 from lilbee.app.services import build_services, services_scope
422 from lilbee.core.config.context import config_scope
423 from lilbee.data.ingest.pipeline import sync
424 from lilbee.providers.fleet.child_guard import bind_lifetime_to_parent
426 bind_lifetime_to_parent(options.parent_pid)
427 _apply_shard_env(spec)
428 index = spec.shard.index
429 reporter = _ShardReporter(index, messages)
430 try:
431 with config_scope(spec.config), services_scope(build_services(spec.config)):
432 result = asyncio.run(
433 sync(
434 force_rebuild=options.force_rebuild,
435 quiet=True,
436 on_progress=reporter,
437 cancel=stop,
438 retry_skipped=options.retry_skipped,
439 shard=spec.shard,
440 )
441 )
442 reporter.flush()
443 messages.put(ShardDone(kind="done", index=index, result=result, error=None))
444 except (Exception, asyncio.CancelledError) as exc:
445 messages.put(ShardDone(kind="done", index=index, result=None, error=error_reason(exc)))