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

1"""Shared 0-as-auto replica-count resolution used by planning, provider, and pipeline. 

2 

3The ``embed_replicas`` / ``vision_replicas`` knobs default to 0, meaning "auto: 

4one replica per GPU". Resolving that consistently in one place keeps the planning 

5hot path, the vision OCR gate, and the ingest fan-out from disagreeing on the 

6effective replica count. 

7""" 

8 

9from __future__ import annotations 

10 

11import functools 

12import logging 

13 

14from lilbee.providers.roles import ROLE_REGISTRY, WorkerRole 

15 

16log = logging.getLogger(__name__) 

17 

18# Auto resolves to at least one replica even when no GPU is enumerated. 

19_MIN_REPLICAS = 1 

20 

21 

22def resolve_replica_count(role: WorkerRole, device_count: int) -> int: 

23 """Requested data-parallel instances for *role* (0 = auto = one per GPU). 

24 

25 Embed and vision honor their ``*_replicas`` knob; an explicit value wins, 

26 0 means one replica per GPU (falling to one when GPU-less). Other roles run 

27 one instance. Capping to residual VRAM happens in placement. 

28 """ 

29 from lilbee.core.config import cfg 

30 

31 knob = ROLE_REGISTRY[role].replica_knob 

32 if knob is None: 

33 return _MIN_REPLICAS 

34 return getattr(cfg, knob) or max(_MIN_REPLICAS, device_count) 

35 

36 

37@functools.cache 

38def gpu_device_count() -> int: 

39 """Effective GPU count lilbee will use; fixed for the process lifetime (cached). 

40 

41 Resolved the same way planning does (binary ``--list-devices`` view), and 

42 floored at one so auto means "one replica" on a GPU-less host. Returns one 

43 when the engine binary is absent so callers that size concurrency (e.g. 

44 ingest) degrade gracefully instead of raising. 

45 """ 

46 from lilbee.providers.base import ProviderError 

47 from lilbee.providers.fleet.binary import resolve_llama_server 

48 from lilbee.providers.fleet.planning import resolve_devices 

49 

50 try: 

51 binary = resolve_llama_server() 

52 except ProviderError: 

53 log.debug("llama-server not found; treating host as GPU-less for replica sizing") 

54 return _MIN_REPLICAS 

55 devices = resolve_devices(binary) 

56 return max(_MIN_REPLICAS, len(devices))