Coverage for src/lilbee/data/ingest/adaptive.py: 100%
178 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Adaptive ingest concurrency: hill-climb to the hardware's throughput knee.
3The extraction-admission limit (how many documents are in their compute phase at
4once) can run in one of three modes, chosen by ``LILBEE_INGEST_CONCURRENCY``:
6- ``static`` -- a fixed limit (the pipeline's ``_max_concurrent()``); the proven
7 default and the guaranteed-safe fallback.
8- ``adaptive-conservative`` / ``adaptive-aggressive`` -- a background controller
9 resizes the limit every few seconds, climbing while throughput still improves
10 and backing off the instant a safety signal (CPU, free RAM, GPU temperature)
11 says the box is under pressure.
13The control law is a safety-gated AIMD hill-climb on smoothed throughput. The
14only thing that pushes the limit up is throughput still improving, so the fixed
15point is *this box's* real operating knee rather than a hardcoded utilization
16target (BBR / Kleinrock). GPU utilization is used only as a saturation veto, never
17as a setpoint. Multiplicative decrease on any danger signal is AIMD's proven
18fast-safe retreat; EWMA smoothing plus an asymmetric dead band and a one-step slew
19limit keep it from oscillating (CLR ThreadPool hill-climbing, TCP Vegas).
21``decide`` is a pure function -- the whole control law with no clock, asyncio, or
22hardware -- so the policy is unit-tested with plain numbers. ``ResizableGate`` and
23``AdaptiveController`` are the async machinery around it.
24"""
26from __future__ import annotations
28import asyncio
29import logging
30import math
31import os
32from collections.abc import Awaitable, Callable, Sequence
33from dataclasses import dataclass
34from enum import StrEnum
35from typing import TYPE_CHECKING
37if TYPE_CHECKING:
38 from lilbee.providers.fleet.gpu_stats import DeviceLike
40log = logging.getLogger(__name__)
42_MODE_ENV = "LILBEE_INGEST_CONCURRENCY"
45class ConcurrencyMode(StrEnum):
46 """How the extraction-admission limit is chosen for a sync run."""
48 STATIC = "static"
49 ADAPTIVE_CONSERVATIVE = "adaptive-conservative"
50 ADAPTIVE_AGGRESSIVE = "adaptive-aggressive"
53def resolve_mode() -> ConcurrencyMode:
54 """The concurrency mode from ``LILBEE_INGEST_CONCURRENCY``; ``static`` by default.
56 An unset or unrecognized value falls back to ``static`` (a warning is logged for
57 an unrecognized one), so a typo can never silently enable adaptive control.
58 """
59 raw = os.environ.get(_MODE_ENV, "").strip().lower()
60 if not raw:
61 return ConcurrencyMode.STATIC
62 try:
63 return ConcurrencyMode(raw)
64 except ValueError:
65 log.warning(
66 "Ignoring %s=%r: expected one of %s; using %s.",
67 _MODE_ENV,
68 raw,
69 [m.value for m in ConcurrencyMode],
70 ConcurrencyMode.STATIC.value,
71 )
72 return ConcurrencyMode.STATIC
75@dataclass(frozen=True)
76class SafetyLimits:
77 """Hard guardrails, identical across adaptive profiles -- safety is never relaxed.
79 Crossing a ``_crit`` line forces an immediate multiplicative backoff; a ``_warn``
80 / ``_soft`` line only vetoes further increases. Temperature and RAM are read raw
81 (never smoothed -- a fire alarm should not be averaged away).
82 """
84 gpu_sat_pct: float = 97.0
85 cpu_soft_pct: float = 90.0
86 cpu_crit_pct: float = 97.0
87 ram_soft_free: float = 0.20
88 ram_min_free: float = 0.10
89 temp_warn_c: float = 80.0
90 temp_crit_c: float = 85.0
91 decrease_factor: float = 0.5
94@dataclass(frozen=True)
95class ConcurrencyProfile:
96 """Climb-speed parameters for one adaptive profile (safety limits are shared)."""
98 name: str
99 interval_s: float
100 ewma_gamma: float
101 deadband_frac: float
102 cool_down_intervals: int
103 sqrt_step: bool
104 latency_veto_ratio: float # veto increases once residence time inflates past baseline x this
105 safety: SafetyLimits = SafetyLimits()
107 def step(self, permits: int) -> int:
108 """Additive step size at the current limit; larger early when ``sqrt_step``."""
109 if self.sqrt_step:
110 return 1 + math.isqrt(max(0, permits)) // 4
111 return 1
114CONSERVATIVE = ConcurrencyProfile(
115 name="conservative",
116 interval_s=5.0,
117 ewma_gamma=0.3,
118 deadband_frac=0.05,
119 cool_down_intervals=3,
120 sqrt_step=False,
121 latency_veto_ratio=1.5,
122)
123AGGRESSIVE = ConcurrencyProfile(
124 name="aggressive",
125 interval_s=2.0,
126 ewma_gamma=0.5,
127 deadband_frac=0.03,
128 cool_down_intervals=2,
129 sqrt_step=True,
130 latency_veto_ratio=2.0,
131)
134def profile_for(mode: ConcurrencyMode) -> ConcurrencyProfile | None:
135 """The profile for an adaptive mode, or None for ``static``."""
136 return {
137 ConcurrencyMode.ADAPTIVE_CONSERVATIVE: CONSERVATIVE,
138 ConcurrencyMode.ADAPTIVE_AGGRESSIVE: AGGRESSIVE,
139 }.get(mode)
142@dataclass(frozen=True)
143class Signals:
144 """One tick's measured state. ``gpu_*`` are None when no GPU telemetry is available."""
146 throughput: float # OCR pages completed since the previous tick (work done, not doc count)
147 gpu_util_pct: float | None
148 gpu_temp_c: float | None
149 cpu_pct: float
150 ram_free_frac: float
153@dataclass(frozen=True)
154class ControllerState:
155 """The controller's carried state between ticks."""
157 permits: int
158 ewma_tput: float | None # smoothed throughput from the previous tick; None at start
159 direction: int # +1 climbing, -1 backing off -- the last hill-climb step's sign
160 cool_down: int # intervals remaining during which increases are suppressed
161 w_min: float | None = None # smallest observed residence-time estimate (Little's Law baseline)
164def _is_critical(signals: Signals, s: SafetyLimits) -> bool:
165 """A signal that demands an immediate hard backoff (thermal / OOM / CPU meltdown)."""
166 temp = signals.gpu_temp_c
167 return (
168 (temp is not None and temp >= s.temp_crit_c)
169 or signals.ram_free_frac <= s.ram_min_free
170 or signals.cpu_pct >= s.cpu_crit_pct
171 )
174def _increase_vetoed(
175 profile: ConcurrencyProfile,
176 signals: Signals,
177 *,
178 cool_down: int,
179 gpu_saturated: bool,
180 w_est: float | None,
181 w_min: float | None,
182) -> bool:
183 """Whether soft pressure, saturation, cool-down, or inflating latency blocks a climb."""
184 s = profile.safety
185 temp = signals.gpu_temp_c
186 latency_inflated = (
187 w_est is not None and w_min is not None and w_est > w_min * profile.latency_veto_ratio
188 )
189 return (
190 cool_down > 0
191 or signals.cpu_pct >= s.cpu_soft_pct
192 or signals.ram_free_frac < s.ram_soft_free
193 or (temp is not None and temp >= s.temp_warn_c)
194 or gpu_saturated
195 or latency_inflated
196 )
199def _hill_climb(
200 profile: ConcurrencyProfile,
201 permits: int,
202 direction: int,
203 delta: float | None,
204 new_ewma: float,
205 clamp: Callable[[int], int],
206) -> tuple[int, int]:
207 """One dead-banded hill-climb step; returns the next (permits, direction).
209 No baseline yet -> a gentle probe up; a clear gain -> step in the same direction;
210 a clear loss -> reverse; inside the dead band -> hold.
211 """
212 if delta is None:
213 return clamp(permits + profile.step(permits)), 1
214 band = profile.deadband_frac * (new_ewma if new_ewma > 0 else 1.0)
215 if delta > band:
216 return clamp(permits + direction * profile.step(permits)), direction
217 if delta < -band:
218 direction = -direction
219 return clamp(permits + direction * profile.step(permits)), direction
220 return permits, direction
223def decide(
224 profile: ConcurrencyProfile,
225 state: ControllerState,
226 signals: Signals,
227 permit_min: int,
228 permit_max: int,
229) -> ControllerState:
230 """Pure control law: fold one tick of signals into the next permit target.
232 Priority: (1) any critical safety signal forces a multiplicative decrease and a
233 cool-down; (2) a soft-pressure, GPU-saturation, or latency-gradient signal vetoes
234 increases (with a single additive decrease when GPUs are saturated *and* throughput
235 is falling -- the Universal Scalability Law's retrograde region); (3) otherwise
236 hill-climb toward the knee. Always clamped to ``[permit_min, permit_max]``.
238 The veto blocks climbing only. A hill-climb step that goes *down* -- the reversal
239 on clearly falling throughput -- still runs under soft pressure, since backing off
240 is exactly what soft pressure should permit.
242 The latency veto is the leading knee indicator (TCP Vegas / Netflix Gradient2):
243 residence time ``W`` is estimated by Little's Law as ``permits / throughput``; once
244 ``W`` inflates past its observed minimum (the unloaded baseline) by the profile's
245 ratio, work is queueing at the bottleneck and increases stop -- before throughput
246 visibly rolls over.
247 """
248 if state.ewma_tput is None:
249 new_ewma: float = signals.throughput
250 delta: float | None = None
251 else:
252 gamma = profile.ewma_gamma
253 new_ewma = gamma * signals.throughput + (1.0 - gamma) * state.ewma_tput
254 delta = new_ewma - state.ewma_tput
256 permits = state.permits
257 cool_down = max(0, state.cool_down - 1)
259 # Little's Law residence-time estimate and its running minimum (the latency baseline).
260 w_est = permits / new_ewma if new_ewma > 0 else None
261 w_min = state.w_min if w_est is None else min(state.w_min or w_est, w_est)
263 def clamp(p: int) -> int:
264 return max(permit_min, min(permit_max, p))
266 def out(new_permits: int, new_direction: int, new_cool_down: int) -> ControllerState:
267 return ControllerState(new_permits, new_ewma, new_direction, new_cool_down, w_min)
269 if _is_critical(signals, profile.safety):
270 backoff = clamp(int(permits * profile.safety.decrease_factor))
271 return out(backoff, -1, profile.cool_down_intervals)
273 util = signals.gpu_util_pct
274 gpu_saturated = util is not None and util >= profile.safety.gpu_sat_pct
275 if gpu_saturated and delta is not None and delta < 0:
276 return out(clamp(permits - profile.step(permits)), -1, cool_down) # USL retrograde
278 vetoed = _increase_vetoed(
279 profile, signals, cool_down=cool_down, gpu_saturated=gpu_saturated, w_est=w_est, w_min=w_min
280 )
281 climbed, direction = _hill_climb(profile, permits, state.direction, delta, new_ewma, clamp)
282 if vetoed and climbed > permits:
283 # The veto blocks climbing, not retreating: a step down on falling
284 # throughput is the only graceful decrease, so it must still pass.
285 return out(permits, state.direction, cool_down)
286 return out(climbed, direction, cool_down)
289class ResizableGate:
290 """An async admission gate whose permit ceiling can change while it is in use.
292 Same ``async with`` shape as ``asyncio.Semaphore``, plus ``set_limit``: growing
293 wakes blocked acquirers, shrinking lowers the ceiling and lets the surplus drain
294 as active holders release. The limit never drops below one, so a shrink can never
295 deadlock a run. A plain counting gate, unlike ``anyio.CapacityLimiter``'s
296 per-borrower token model.
297 """
299 def __init__(self, limit: int) -> None:
300 self._limit = max(1, limit)
301 self._active = 0
302 self._cond = asyncio.Condition()
304 @property
305 def limit(self) -> int:
306 return self._limit
308 async def acquire(self) -> None:
309 async with self._cond:
310 await self._cond.wait_for(lambda: self._active < self._limit)
311 self._active += 1
313 async def release(self) -> None:
314 async with self._cond:
315 self._active -= 1
316 self._cond.notify_all()
318 async def __aenter__(self) -> ResizableGate:
319 await self.acquire()
320 return self
322 async def __aexit__(self, *exc: object) -> None:
323 await self.release()
325 async def set_limit(self, new_limit: int) -> None:
326 async with self._cond:
327 self._limit = max(1, new_limit)
328 self._cond.notify_all()
331class AdaptiveController:
332 """Drives a :class:`ResizableGate`'s limit from live signals until cancelled.
334 ``sample(throughput)`` returns the current :class:`Signals`; ``completed()`` is a
335 monotonic count of finished work units, from which per-interval throughput is
336 derived. The production wiring counts OCR pages, not documents (a per-document
337 count would bias the controller toward files of a given size). Both are injected
338 so the controller runs in tests with no clock or GPU.
339 """
341 def __init__(
342 self,
343 gate: ResizableGate,
344 profile: ConcurrencyProfile,
345 sample: Callable[[float], Signals],
346 completed: Callable[[], int],
347 *,
348 permit_min: int,
349 permit_max: int,
350 sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
351 ) -> None:
352 self._gate = gate
353 self._profile = profile
354 self._sample = sample
355 self._completed = completed
356 self._permit_min = permit_min
357 self._permit_max = permit_max
358 self._sleep = sleep
360 async def run(self) -> None:
361 """Sample-decide-resize loop. Runs until the task is cancelled.
363 A tuning tick is best-effort: a transient sampling/probe failure is logged and
364 skipped, never propagated, so the controller can never crash the ingest it is
365 only advising. Cancellation (``CancelledError``) still ends the loop cleanly.
366 """
367 state = ControllerState(self._gate.limit, None, 1, 0)
368 last_completed = self._completed()
369 while True:
370 await self._sleep(self._profile.interval_s)
371 try:
372 now_completed = self._completed()
373 throughput = float(max(0, now_completed - last_completed))
374 last_completed = now_completed
375 state = decide(
376 self._profile,
377 state,
378 self._sample(throughput),
379 self._permit_min,
380 self._permit_max,
381 )
382 if state.permits != self._gate.limit:
383 await self._gate.set_limit(state.permits)
384 log.debug("adaptive ingest: limit -> %d", state.permits)
385 except Exception:
386 log.debug("adaptive ingest: tuning tick failed; skipping", exc_info=True)
389def enumerate_fleet_devices() -> Sequence[DeviceLike]:
390 """The GPU devices to read telemetry from, or empty when none can be enumerated.
392 Any failure (no engine binary, probe error) degrades to an empty list, which the
393 caller treats as "no fleet to feed" and falls back to the static limit.
394 """
395 try:
396 from lilbee.providers.fleet.binary import resolve_llama_server
397 from lilbee.providers.fleet.planning import resolve_devices
399 return resolve_devices(resolve_llama_server())
400 except Exception:
401 log.debug("adaptive ingest: device enumeration failed; using static limit", exc_info=True)
402 return []
405def make_signal_sampler(devices: Sequence[DeviceLike]) -> Callable[[float], Signals]:
406 """Build a sampler that reads mean GPU util, max GPU temp, CPU %, and free RAM.
408 GPU util/temp are None when no device reports them (the controller then relies on
409 the CPU and RAM guards alone); throughput is supplied by the controller.
410 """
411 import psutil
413 from lilbee.providers.fleet.gpu_stats import probe_gpu_stats
415 def sample(throughput: float) -> Signals:
416 stats = probe_gpu_stats(devices)
417 utils = [g.utilization_pct for g in stats.values() if g.utilization_pct is not None]
418 temps = [g.temperature_c for g in stats.values() if g.temperature_c is not None]
419 vm = psutil.virtual_memory()
420 return Signals(
421 throughput=throughput,
422 gpu_util_pct=(sum(utils) / len(utils)) if utils else None,
423 gpu_temp_c=float(max(temps)) if temps else None,
424 cpu_pct=psutil.cpu_percent(interval=None),
425 ram_free_frac=vm.available / vm.total, # psutil total is always positive
426 )
428 return sample