Coverage for src/lilbee/data/extract/batch.py: 100%

1"""Coalesce concurrent single-document extractions into one xberg batch call. 

2 

3Active only when ``cfg.batch_extraction`` is on. Intercepts the pipeline's per-file 

4extraction coroutines and groups in-flight ones into a single ``extract_batch``, 

5leaving the per-file contract (progress, admission, results) untouched. No 

6off-the-shelf async coalescer speaks xberg's per-input config, so the buffer-and- 

7flush below is hand-rolled. Inputs sharing a mode share a batch config; each keeps 

8its own OCR token as a per-file override. 

9""" 

10 

11from __future__ import annotations 

12 

13import asyncio 

14import contextvars 

15from dataclasses import dataclass 

16from typing import TYPE_CHECKING 

17 

18from .xberg import BatchItem 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Awaitable, Callable 

22 

23 from xberg import ExtractedDocument, ExtractionConfig, OcrConfig 

24 

25 from lilbee.data.types import ExtractMode 

26 

27# Delay before firing a partial batch, so concurrent extractions arriving in the 

28# same burst join it. Essential for the tail and when fewer than a full batch's 

29# worth of files are ever admitted at once. 

30_FLUSH_WINDOW_S = 0.05 

31 

32 

33@dataclass 

34class _Pending: 

35 data: bytes 

36 filename: str 

37 ocr_token: str 

38 future: asyncio.Future[ExtractedDocument] 

39 

40 

41class ExtractBatcher: 

42 """Buffers extraction requests per mode and flushes them as one batch call.""" 

43 

44 def __init__( 

45 self, 

46 *, 

47 size: int, 

48 config_fn: Callable[[ExtractMode], ExtractionConfig], 

49 ocr_fn: Callable[[str], OcrConfig], 

50 batch_fn: Callable[ 

51 [list[BatchItem], ExtractionConfig], Awaitable[list[ExtractedDocument | Exception]] 

52 ], 

53 window: float = _FLUSH_WINDOW_S, 

54 ) -> None: 

55 self._size = size 

56 self._window = window 

57 self._config_fn = config_fn 

58 self._ocr_fn = ocr_fn 

59 self._batch_fn = batch_fn 

60 self._groups: dict[ExtractMode, list[_Pending]] = {} 

61 self._timers: dict[ExtractMode, asyncio.TimerHandle] = {} 

62 self._running: set[asyncio.Task[None]] = set() 

63 

64 async def submit( 

65 self, mode: ExtractMode, data: bytes, filename: str, ocr_token: str 

66 ) -> ExtractedDocument: 

67 """Enqueue one extraction; resolves when its batch completes.""" 

68 loop = asyncio.get_running_loop() 

69 future: asyncio.Future[ExtractedDocument] = loop.create_future() 

70 group = self._groups.setdefault(mode, []) 

71 group.append(_Pending(data, filename, ocr_token, future)) 

72 if len(group) >= self._size: 

73 self._flush(mode) 

74 elif mode not in self._timers: 

75 self._timers[mode] = loop.call_later(self._window, self._flush, mode) 

76 return await future 

77 

78 def _flush(self, mode: ExtractMode) -> None: 

79 timer = self._timers.pop(mode, None) 

80 if timer is not None: 

81 timer.cancel() 

82 pending = self._groups.pop(mode, []) 

83 if not pending: 

84 return 

85 config = self._config_fn(mode) 

86 # mime=None: xberg detects the format from the filename, matching the 

87 # single-file path. Passing lilbee's bare content_type here is rejected. 

88 items = [BatchItem(p.data, None, p.filename, self._ocr_fn(p.ocr_token)) for p in pending] 

89 task = asyncio.ensure_future(self._run(items, config, pending)) 

90 self._running.add(task) 

91 task.add_done_callback(self._running.discard) 

92 

93 async def _run( 

94 self, items: list[BatchItem], config: ExtractionConfig, pending: list[_Pending] 

95 ) -> None: 

96 try: 

97 docs = await self._batch_fn(items, config) 

98 except Exception as exc: # whole-batch failure fails every awaiter 

99 for p in pending: 

100 if not p.future.done(): 

101 p.future.set_exception(exc) 

102 return 

103 for p, doc in zip(pending, docs, strict=True): 

104 if p.future.done(): 

105 continue 

106 if isinstance(doc, BaseException): 

107 p.future.set_exception(doc) 

108 else: 

109 p.future.set_result(doc) 

110 

111 async def close(self) -> None: 

112 """Flush every buffered group and await the batches in flight.""" 

113 for mode in list(self._groups): 

114 self._flush(mode) 

115 if self._running: 

116 await asyncio.gather(*self._running, return_exceptions=True) 

117 

118 

119_active: contextvars.ContextVar[ExtractBatcher | None] = contextvars.ContextVar( 

120 "lilbee_extract_batcher", default=None 

121) 

122 

123 

124def active_extract_batcher() -> ExtractBatcher | None: 

125 """The batcher for the current ingest run, or None when batching is off.""" 

126 return _active.get() 

127 

128 

129def set_active_batcher(batcher: ExtractBatcher) -> contextvars.Token[ExtractBatcher | None]: 

130 return _active.set(batcher) 

131 

132 

133def reset_active_batcher(token: contextvars.Token[ExtractBatcher | None]) -> None: 

134 _active.reset(token)