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 

17from concurrent.futures import Executor, ThreadPoolExecutor 

18from typing import 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 

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

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

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

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

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

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

33# ``ingest_max_inflight`` / ``LILBEE_INGEST_MAX_INFLIGHT``. 

34_EMBED_INFLIGHT_PER_REPLICA = 8 

35 

36 

37def embed_inflight_target() -> int: 

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

39 

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

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

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

43 """ 

44 try: 

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

46 from lilbee.providers.roles import WorkerRole 

47 

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

49 except Exception: 

50 return 0 

51 return slots * _EMBED_INFLIGHT_PER_REPLICA if slots > 1 else 0 

52 

53 

54@functools.cache 

55def max_workers() -> int: 

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

57 

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

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

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

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

62 ceiling drift away from the pool that actually exists. 

63 

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

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

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

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

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

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

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

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

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

73 and fall back to the default. 

74 """ 

75 override = os.environ.get(_MAX_WORKERS_ENV) 

76 if override is not None: 

77 try: 

78 value = int(override) 

79 if value > 0: 

80 return value 

81 except ValueError: 

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

83 log.warning( 

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

85 _MAX_WORKERS_ENV, 

86 override, 

87 ) 

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

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

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

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

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

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

94 from lilbee.core.config import active_config 

95 

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

97 return max(default, inflight) 

98 

99 

100@functools.cache 

101def _ingest_executor() -> ThreadPoolExecutor: 

102 """The shared ingest pool, created on first use (cache makes it a singleton).""" 

103 return ThreadPoolExecutor(max_workers=max_workers(), thread_name_prefix="lilbee-ingest") 

104 

105 

106async def to_executor( 

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

108) -> _R: 

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

110 

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

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

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

114 """ 

115 loop = asyncio.get_running_loop() 

116 ctx = contextvars.copy_context() 

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

118 return await loop.run_in_executor(executor, call) 

119 

120 

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

122 """Run *fn* on the shared ingest executor.""" 

123 return await to_executor(_ingest_executor(), fn, *args, **kwargs)