Coverage for src/lilbee/providers/fleet/ingest_warmth.py: 100%

1"""Keep the fleet resident for the duration of a bulk ingest. 

2 

3A bulk ingest fans embed requests unevenly across replicas: a replica that goes 

4briefly idle hits ``engine_idle_ttl_minutes``, unloads its weights, and reloads 

5cold on the next request. That cold reload is a connection-refused to the 

6dispatcher, which retries onto the surviving replicas, piling more load on them 

7so they, too, fall behind and unload -- a positive-feedback collapse (8 GPUs 

8drop to 3). Holding the whole fleet resident (llama-swap ttl 0) for the sync 

9avoids it. The signal is a ContextVar so it scopes to exactly one ingest and 

10propagates into the ingest thread pool (``to_ingest_thread`` copies the context). 

11""" 

12 

13from __future__ import annotations 

14 

15import contextvars 

16from collections.abc import Iterator 

17from contextlib import contextmanager 

18 

19_INGEST_KEEP_WARM: contextvars.ContextVar[bool] = contextvars.ContextVar( 

20 "ingest_keep_warm", default=False 

21) 

22 

23 

24def ingest_keep_warm() -> bool: 

25 """Whether an ingest is currently holding the fleet resident.""" 

26 return _INGEST_KEEP_WARM.get() 

27 

28 

29@contextmanager 

30def keep_fleet_warm() -> Iterator[None]: 

31 """Hold the fleet resident (ttl 0) for the duration of the block.""" 

32 token = _INGEST_KEEP_WARM.set(True) 

33 try: 

34 yield 

35 finally: 

36 _INGEST_KEEP_WARM.reset(token)