Coverage for src/lilbee/data/offload.py: 100%

1"""Dedicated thread pool for ingest-side blocking work. 

2 

3``asyncio.to_thread`` runs on the loop's default executor -- the same pool the 

4serving path (chat dispatch, handlers) offloads to. A fanned ingest can fill 

5that pool with extraction/OCR/embedding calls and starve request handling, so 

6ingest work runs on its own bounded executor instead. 

7""" 

8 

9from __future__ import annotations 

10 

11import asyncio 

12import contextvars 

13import functools 

14import logging 

15import os 

16from collections.abc import Callable, Coroutine 

17from concurrent.futures import Executor, ThreadPoolExecutor 

18from typing import Any, ParamSpec, TypeVar 

19 

20log = logging.getLogger(__name__) 

21 

22_P = ParamSpec("_P") 

23_R = TypeVar("_R") 

24 

25_MAX_WORKERS_ENV = "LILBEE_INGEST_MAX_WORKERS" 

26_POOL_THREAD_PREFIX = "lilbee-ingest" 

27 

28# Files per embed replica to keep in flight during ingest. A replica interleaves 

29# its file's extraction with embedding, so a handful of files per card must be 

30# admitted at once or the GPU sits idle between requests. Scaling admission by 

31# the detected replica count auto-sizes a multi-GPU fleet with no manual cap, 

32# instead of the CPU-bound default that pins an 8-GPU box at ~4 files/card. 

33# Tuned against the 8x A100 MS MARCO fleet; override per run with 

34# ``ingest_max_inflight`` / ``LILBEE_INGEST_MAX_INFLIGHT``. 

35_EMBED_INFLIGHT_PER_REPLICA = 8 

36 

37 

38def embed_inflight_target() -> int: 

39 """Admission that keeps every embed replica fed, from the detected fleet size. 

40 

41 ``embed replicas x _EMBED_INFLIGHT_PER_REPLICA`` when a multi-replica fleet is 

42 resolvable, else 0 (single card or no fleet: the CPU-bound sizing already 

43 fits). Never raises -- a fleet that cannot be probed yet returns 0. 

44 """ 

45 try: 

46 from lilbee.providers.fleet.replicas import gpu_device_count, resolve_replica_count 

47 from lilbee.providers.roles import WorkerRole 

48 

49 slots = resolve_replica_count(WorkerRole.EMBED, gpu_device_count()) 

50 except Exception: 

51 return 0 

52 return slots * _EMBED_INFLIGHT_PER_REPLICA if slots > 1 else 0 

53 

54 

55@functools.cache 

56def max_workers() -> int: 

57 """The ingest pool's worker count -- its hard concurrency ceiling. 

58 

59 The adaptive-concurrency controller uses this as the upper bound on in-flight 

60 documents, since each one needs a pool thread to run its blocking extraction. 

61 Resolved once and cached: the pool is built with this same value on first use, 

62 so re-reading ``LILBEE_INGEST_MAX_WORKERS`` per call would let the reported 

63 ceiling drift away from the pool that actually exists. 

64 

65 Extraction rasterizes PDFs and drives OCR on this pool. The default caps at 

66 ``min(32, cpu_count + 4)``: a worker-count sweep on forced-OCR multi-page PDFs 

67 (4x H100) held throughput flat from 16 to 32 workers and *declining* past it, 

68 with the GPUs already ~85-90% busy the whole time. OCR ingest is GPU-bound, not 

69 extraction-bound, so extra threads only rasterize ahead into a buffer the GPUs 

70 cannot drain any faster while oversubscribing the box. The ``+4`` keeps headroom 

71 for threads parked on OCR/embed I/O; small hosts scale below the cap. The 

72 override lifts the ceiling for genuinely CPU-bound work (e.g. bulk text 

73 extraction) that can use more threads; non-positive or unparseable values warn 

74 and fall back to the default. 

75 """ 

76 override = os.environ.get(_MAX_WORKERS_ENV) 

77 if override is not None: 

78 try: 

79 value = int(override) 

80 if value > 0: 

81 return value 

82 except ValueError: 

83 pass # bad override falls through to the warning + default below 

84 log.warning( 

85 "Ignoring %s=%r: must be a positive integer; using default.", 

86 _MAX_WORKERS_ENV, 

87 override, 

88 ) 

89 default = min(32, (os.cpu_count() or 4) + 4) 

90 # The pool (and thus the adaptive controller's permit_max, which is this 

91 # value) must be able to feed the admission ceiling, or in adaptive mode the 

92 # gate clamps back to 32 and a multi-GPU fleet stays starved. Size it to the 

93 # explicit ingest_max_inflight override, else auto-scale with the detected 

94 # embed fleet so no manual cap is needed on a multi-GPU box. 

95 from lilbee.core.config import active_config 

96 

97 inflight = active_config().ingest_max_inflight or embed_inflight_target() 

98 return max(default, inflight) 

99 

100 

101@functools.cache 

102def _ingest_executor() -> ThreadPoolExecutor: 

103 """The shared ingest pool for work outside an ingest run, created on first use.""" 

104 return ThreadPoolExecutor(max_workers=max_workers(), thread_name_prefix=_POOL_THREAD_PREFIX) 

105 

106 

107_run_pool: contextvars.ContextVar[ThreadPoolExecutor | None] = contextvars.ContextVar( 

108 "lilbee_ingest_run_pool", default=None 

109) 

110 

111 

112def owns_ingest_pool( 

113 run: Callable[_P, Coroutine[Any, Any, _R]], 

114) -> Callable[_P, Coroutine[Any, Any, _R]]: 

115 """Give each call of the async *run* its own ingest pool, shut down when the call ends.""" 

116 

117 @functools.wraps(run) 

118 async def _with_pool(*args: _P.args, **kwargs: _P.kwargs) -> _R: 

119 pool = ThreadPoolExecutor(max_workers=max_workers(), thread_name_prefix=_POOL_THREAD_PREFIX) 

120 token = _run_pool.set(pool) 

121 try: 

122 return await run(*args, **kwargs) 

123 finally: 

124 _run_pool.reset(token) 

125 pool.shutdown(wait=False, cancel_futures=True) 

126 

127 return _with_pool 

128 

129 

130async def to_executor( 

131 executor: Executor, fn: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs 

132) -> _R: 

133 """``asyncio.to_thread`` on *executor*, contextvars preserved. 

134 

135 Extraction relies on contextvar propagation into its workers (cancel, config 

136 and progress context), which ``run_in_executor`` alone would drop; copy the 

137 context exactly like ``asyncio.to_thread`` does. 

138 """ 

139 return await _submit(executor, fn, *args, **kwargs) 

140 

141 

142def _submit( 

143 executor: Executor, fn: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs 

144) -> asyncio.Future[_R]: 

145 loop = asyncio.get_running_loop() 

146 ctx = contextvars.copy_context() 

147 call = functools.partial(ctx.run, fn, *args, **kwargs) 

148 return loop.run_in_executor(executor, call) 

149 

150 

151async def to_ingest_thread(fn: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs) -> _R: 

152 """Run *fn* on the current ingest run's pool, or on the shared pool outside a run.""" 

153 return await ingest_thread_future(fn, *args, **kwargs) 

154 

155 

156def ingest_thread_future( 

157 fn: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs 

158) -> asyncio.Future[_R]: 

159 """Start *fn* on the ingest pool; the future is not a task, so no sweep of tasks cancels it.""" 

160 return _submit(_run_pool.get() or _ingest_executor(), fn, *args, **kwargs)