Coverage for src/lilbee/data/ingest/fanout.py: 100%
200 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""One ingest worker process per GPU, each writing its slice of the corpus to the one index."""
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 TimeElapsedColumn,
23)
25from lilbee.core.config import active_config
26from lilbee.data.ingest.errors import error_reason
27from lilbee.data.types import ShardId, SyncResult
28from lilbee.runtime.console import PlainConsole
29from lilbee.runtime.cpu import available_cpu_count, cpu_quota
30from lilbee.runtime.progress import (
31 BatchProgressEvent,
32 BatchStatus,
33 DetailedProgressCallback,
34 EventType,
35 ProgressEvent,
36)
37from lilbee.runtime.progress.columns import literal_text_column
38from lilbee.runtime.shard_process import shard_process
40if TYPE_CHECKING:
41 from collections.abc import Sequence
42 from multiprocessing.process import BaseProcess
43 from multiprocessing.queues import Queue
44 from multiprocessing.synchronize import Event
46 from lilbee.core.config.model import Config
47 from lilbee.runtime.cancellation import CancelSignal
49log = logging.getLogger(__name__)
51# Per-worker state (engine slots, log) under the parent data root. The index
52# and the ``.lilbeeignore`` are the corpus's, at the parent data root itself.
53SHARDS_DIRNAME = "shards"
55# Below this many files on disk a fan-out costs more than it saves: every worker
56# pays a fresh interpreter and its own engine.
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 and reads of the cancel.
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 the workers get, together, to exit on a terminate before they are killed.
72_WORKER_EXIT_GRACE_S = 30.0
74# How often the parent looks for the workers' exit during that grace.
75_EXIT_POLL_S = 0.01
77# Where a worker's console output lands, under its own data root.
78WORKER_LOG_NAME = "sync.log"
81@dataclass(frozen=True)
82class ShardSpec:
83 """One worker's slice, its card, and the state it owns."""
85 shard: ShardId
86 device: int
87 config: Config
88 engine_dir: Path
89 cpu_share: int
90 visible_devices: dict[str, str]
93@dataclass(frozen=True)
94class ShardOptions:
95 """What every worker of one fan-out is told about the run it belongs to."""
97 parent_pid: int
98 rebuild: bool = False
101@dataclass(frozen=True)
102class ShardProgress:
103 """A worker's counters as it works."""
105 kind: Literal["progress"]
106 index: int
107 done: int
108 planned: int
109 file: str
110 status: BatchStatus
113@dataclass(frozen=True)
114class ShardDone:
115 """A worker's verdict; *error* set means it did not finish its slice."""
117 kind: Literal["done"]
118 index: int
119 result: SyncResult | None
120 error: str | None
123ShardMessage = ShardProgress | ShardDone
126def resolve_process_count(devices: int) -> int:
127 """Ingest worker processes for this run; 1 keeps ingest in this process.
129 Auto (``ingest_processes = 0``) is one worker per visible card. An explicit
130 count is honored past the card count, since two workers on one card is a
131 legitimate configuration; they share that card's engine slot rather than
132 putting a second fleet on it.
133 """
134 configured = active_config().ingest_processes
135 if configured:
136 return max(1, configured)
137 return devices
140def plan_fanout() -> list[ShardSpec]:
141 """The workers for this sync, empty when it runs in this process."""
142 from lilbee.data.ingest.discovery import corpus_has_at_least
143 from lilbee.providers.fleet.replicas import gpu_device_count
145 # The count runs the device probe under the pin, so a gpu_devices pin is the
146 # space the workers are dealt in: they are never dealt a card the pin excludes.
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, corpus_root=config.data_root),
164 device=index % devices,
165 config=_shard_config(config, root / f"w{index}", plan_share, processes),
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, processes: int) -> Config:
177 """*config* with a private data root and this worker's share of the CPU pools.
179 ``documents_dir``, ``linked_roots`` and ``lancedb_dir`` are inherited: every
180 worker reads the one corpus and writes the one index.
181 """
182 threads = config.extraction_threads
183 return config.model_copy(
184 update={
185 "data_root": root,
186 "ingest_workers": plan_share,
187 "extraction_threads": max(1, threads // processes) if threads else 0,
188 }
189 )
192def _pin_shard(spec: ShardSpec) -> None:
193 """Pin this process to the worker's card, engine slot, CPU share and log.
195 Nothing is written to ``os.environ``: the readers take the card, engine dir
196 and CPU share from the process's slice state.
197 """
198 shard_process.bind(spec.engine_dir, spec.cpu_share, spec.visible_devices)
199 _redirect_output(spec.config.data_root / WORKER_LOG_NAME)
202def _redirect_output(path: Path) -> None:
203 """Send this process's console output to *path*.
205 At the file descriptor, so the engine this worker spawns follows it: N
206 workers logging onto the parent's terminal is the pile of log files the one
207 aggregated bar exists to replace.
208 """
209 path.parent.mkdir(parents=True, exist_ok=True)
210 with path.open("ab", buffering=0) as handle:
211 os.dup2(handle.fileno(), sys.stdout.fileno())
212 os.dup2(handle.fileno(), sys.stderr.fileno())
215class _ShardReporter:
216 """Throttled relay of a worker's counters onto the parent's queue.
218 The counters are the pipeline's own: how much of this worker's slice is done
219 and how big that slice is. Counting files here instead would only re-derive
220 the first, and the per-file events carry no slice size -- FILE_START's total
221 is the plan so far, which grows all run.
222 """
224 def __init__(self, index: int, messages: Queue[ShardMessage]) -> None:
225 self._index = index
226 self._messages = messages
227 self._done = 0
228 self._planned = 0
229 self._last_sent = 0.0
231 def __call__(self, event_type: EventType, data: ProgressEvent) -> None:
232 if event_type is not EventType.BATCH_PROGRESS or not isinstance(data, BatchProgressEvent):
233 return
234 self._done = data.current
235 self._planned = data.total
236 now = time.monotonic()
237 if now - self._last_sent < _REPORT_INTERVAL_S:
238 return
239 self._last_sent = now
240 self._send(data.file, data.status)
242 def flush(self) -> None:
243 """Send the final counters past the throttle, so the bar lands on its total."""
244 self._send("", BatchStatus.INGESTED)
246 def _send(self, file: str, status: BatchStatus) -> None:
247 self._messages.put(
248 ShardProgress(
249 kind="progress",
250 index=self._index,
251 done=self._done,
252 planned=self._planned,
253 file=file,
254 status=status,
255 )
256 )
259class _Aggregate:
260 """Every worker's latest counters, as one set of totals."""
262 def __init__(self, on_progress: DetailedProgressCallback) -> None:
263 self._latest: dict[int, ShardProgress] = {}
264 self._on_progress = on_progress
266 def update(self, message: ShardProgress) -> tuple[int, int]:
267 """Record *message* and return the corpus-wide (done, planned)."""
268 self._latest[message.index] = message
269 done = sum(p.done for p in self._latest.values())
270 planned = sum(p.planned for p in self._latest.values())
271 self._on_progress(
272 EventType.BATCH_PROGRESS,
273 BatchProgressEvent(
274 file=message.file, status=message.status, current=done, total=planned
275 ),
276 )
277 return done, planned
280def _drain(messages: Queue[ShardMessage]) -> list[ShardMessage]:
281 """Every message queued right now, without blocking."""
282 drained: list[ShardMessage] = []
283 with contextlib.suppress(queue.Empty):
284 while True:
285 drained.append(messages.get_nowait())
286 return drained
289def _shard_progress_bar(quiet: bool) -> Progress:
290 """The one bar a fan-out reports on, disabled when the caller wants no output."""
291 return Progress(
292 SpinnerColumn(),
293 literal_text_column("{task.description}", style="progress.description"),
294 BarColumn(),
295 MofNCompleteColumn(),
296 TimeElapsedColumn(),
297 console=PlainConsole(),
298 disable=quiet,
299 )
302async def _supervise(
303 workers: Sequence[BaseProcess],
304 messages: Queue[ShardMessage],
305 *,
306 quiet: bool,
307 on_progress: DetailedProgressCallback,
308 cancel: CancelSignal | None,
309) -> dict[int, ShardDone]:
310 """Drain worker messages until every worker has reported, keeping one bar current.
312 Raises ``asyncio.CancelledError`` within one drain interval of a set *cancel*,
313 whether or not a worker reports.
314 """
315 verdicts: dict[int, ShardDone] = {}
316 aggregate = _Aggregate(on_progress)
317 with _shard_progress_bar(quiet) as progress:
318 task = progress.add_task(f"Ingesting on {len(workers)} workers", total=None)
319 while len(verdicts) < len(workers):
320 for message in _drain(messages):
321 if message.kind == "done":
322 verdicts[message.index] = message
323 else:
324 done, planned = aggregate.update(message)
325 progress.update(task, completed=done, total=planned or None)
326 if cancel is not None and cancel.is_set():
327 raise asyncio.CancelledError
328 if not any(worker.is_alive() for worker in workers):
329 verdicts.update(_final_verdicts(workers, messages, verdicts))
330 break
331 await asyncio.sleep(_DRAIN_INTERVAL_S)
332 return verdicts
335def _final_verdicts(
336 workers: Sequence[BaseProcess],
337 messages: Queue[ShardMessage],
338 verdicts: dict[int, ShardDone],
339) -> dict[int, ShardDone]:
340 """Verdicts still in flight once every worker has exited, plus one per silent death.
342 A worker the kernel killed (out of memory is the usual reason) reports
343 nothing, so it is recorded as failed and the sync does not end as a success.
344 """
345 time.sleep(_FINAL_DRAIN_S)
346 late = {m.index: m for m in _drain(messages) if m.kind == "done"}
347 for index, worker in enumerate(workers):
348 if index in verdicts or index in late:
349 continue
350 late[index] = ShardDone(
351 kind="done",
352 index=index,
353 result=None,
354 error=f"worker exited with code {worker.exitcode} before reporting",
355 )
356 return late
359async def _stop_workers(workers: Sequence[BaseProcess], stop: Event) -> None:
360 """Terminate every live worker, give them one grace period together, then kill the rest.
362 A worker owns a GPU fleet, and its teardown can outlast a TERM; a plain join
363 would hang the sync behind it instead of returning a result it already has.
364 The wait yields to the event loop, and a cancel during it kills at once.
365 """
366 stop.set()
367 for worker in workers:
368 if worker.is_alive():
369 worker.terminate()
370 try:
371 await _exited_or_grace_over(workers)
372 finally:
373 for worker in workers:
374 if worker.is_alive():
375 log.warning("Ingest worker %s did not exit; killing it", worker.name)
376 worker.kill()
377 worker.join()
380async def _exited_or_grace_over(workers: Sequence[BaseProcess]) -> None:
381 """Return once no worker is alive, or once the exit grace has passed."""
382 deadline = time.monotonic() + _WORKER_EXIT_GRACE_S
383 while time.monotonic() < deadline and any(worker.is_alive() for worker in workers):
384 await asyncio.sleep(_EXIT_POLL_S)
387async def run_workers(
388 specs: list[ShardSpec],
389 *,
390 options: ShardOptions,
391 quiet: bool,
392 on_progress: DetailedProgressCallback,
393 cancel: CancelSignal | None,
394) -> list[ShardDone]:
395 """Run every worker and return their verdicts, in shard order; a cancel terminates them."""
396 context = multiprocessing.get_context("spawn")
397 messages: Queue[ShardMessage] = context.Queue()
398 stop = context.Event()
399 workers = [
400 context.Process(
401 target=run_shard,
402 args=(spec, options, messages, stop),
403 name=f"lilbee-shard-{spec.shard.index}",
404 )
405 for spec in specs
406 ]
407 log.warning("Ingesting across %d worker processes, one per GPU", len(workers))
408 for worker in workers:
409 worker.start()
410 try:
411 verdicts = await _supervise(
412 workers, messages, quiet=quiet, on_progress=on_progress, cancel=cancel
413 )
414 finally:
415 await _stop_workers(workers, stop)
416 return [verdicts[index] for index in sorted(verdicts)]
419def aggregate_results(verdicts: list[ShardDone]) -> SyncResult:
420 """The one result a fan-out reports, unioned from every worker's."""
421 results = [verdict.result for verdict in verdicts if verdict.result is not None]
422 return SyncResult(
423 added=[name for r in results for name in r.added],
424 updated=[name for r in results for name in r.updated],
425 relocated=[name for r in results for name in r.relocated],
426 failed=[name for r in results for name in r.failed],
427 skipped=[name for r in results for name in r.skipped],
428 skipped_ocr={name: ocr for r in results for name, ocr in r.skipped_ocr.items()},
429 removed=[name for r in results for name in r.removed],
430 held_out=[held for r in results for held in r.held_out],
431 unchanged=sum(r.unchanged for r in results),
432 truncated=sum(r.truncated for r in results),
433 # Each worker walks the same tree, so each one meets the same folders.
434 unreadable_folders=list(
435 dict.fromkeys(folder for r in results for folder in r.unreadable_folders)
436 ),
437 )
440def run_shard(
441 spec: ShardSpec, options: ShardOptions, messages: Queue[ShardMessage], stop: Event
442) -> None:
443 """Ingest this worker's slice in a fresh process, reporting onto *messages*."""
444 from lilbee.app.services import build_services, services_scope
445 from lilbee.core.config.context import config_scope
446 from lilbee.data.ingest.pipeline import sync
447 from lilbee.providers.fleet.child_guard import bind_lifetime_to_parent
449 bind_lifetime_to_parent(options.parent_pid)
450 _pin_shard(spec)
451 index = spec.shard.index
452 reporter = _ShardReporter(index, messages)
453 try:
454 with config_scope(spec.config), services_scope(build_services(spec.config)):
455 result = asyncio.run(
456 sync(
457 quiet=True,
458 on_progress=reporter,
459 cancel=stop,
460 shard=spec.shard,
461 force_rebuild=options.rebuild,
462 )
463 )
464 reporter.flush()
465 messages.put(ShardDone(kind="done", index=index, result=result, error=None))
466 except (Exception, asyncio.CancelledError) as exc:
467 messages.put(ShardDone(kind="done", index=index, result=None, error=error_reason(exc)))