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 from xberg.progress import ProgressCallback 

25 

26 from lilbee.data.types import ExtractMode 

27 

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

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

30# worth of files are ever admitted at once. 

31_FLUSH_WINDOW_S = 0.05 

32 

33 

34@dataclass 

35class _Pending: 

36 data: bytes 

37 filename: str 

38 ocr_token: str 

39 on_progress: ProgressCallback | None 

40 future: asyncio.Future[ExtractedDocument] 

41 

42 

43class ExtractBatcher: 

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

45 

46 def __init__( 

47 self, 

48 *, 

49 size: int, 

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

51 ocr_fn: Callable[[str], OcrConfig | None], 

52 batch_fn: Callable[ 

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

54 ], 

55 window: float = _FLUSH_WINDOW_S, 

56 ) -> None: 

57 self._size = size 

58 self._window = window 

59 self._config_fn = config_fn 

60 self._ocr_fn = ocr_fn 

61 self._batch_fn = batch_fn 

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

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

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

65 

66 async def submit( 

67 self, 

68 mode: ExtractMode, 

69 data: bytes, 

70 filename: str, 

71 ocr_token: str, 

72 on_progress: ProgressCallback | None = None, 

73 ) -> ExtractedDocument: 

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

75 

76 *on_progress* receives this file's OCR page events only. 

77 """ 

78 loop = asyncio.get_running_loop() 

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

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

81 group.append(_Pending(data, filename, ocr_token, on_progress, future)) 

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

83 self._flush(mode) 

84 elif mode not in self._timers: 

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

86 return await future 

87 

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

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

90 if timer is not None: 

91 timer.cancel() 

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

93 if not pending: 

94 return 

95 config = self._config_fn(mode) 

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

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

98 items = [ 

99 BatchItem(p.data, None, p.filename, self._ocr_fn(p.ocr_token), p.on_progress) 

100 for p in pending 

101 ] 

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

103 self._running.add(task) 

104 task.add_done_callback(self._running.discard) 

105 

106 async def _run( 

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

108 ) -> None: 

109 try: 

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

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

112 for p in pending: 

113 if not p.future.done(): 

114 p.future.set_exception(exc) 

115 return 

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

117 if p.future.done(): 

118 continue 

119 if isinstance(doc, BaseException): 

120 p.future.set_exception(doc) 

121 else: 

122 p.future.set_result(doc) 

123 

124 async def close(self) -> None: 

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

126 for mode in list(self._groups): 

127 self._flush(mode) 

128 if self._running: 

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

130 

131 

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

133 "lilbee_extract_batcher", default=None 

134) 

135 

136 

137def active_extract_batcher() -> ExtractBatcher | None: 

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

139 return _active.get() 

140 

141 

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

143 return _active.set(batcher) 

144 

145 

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

147 _active.reset(token)