Coverage for src/lilbee/runtime/hardware.py: 100%

1"""Hardware-fit signaling and per-row size-variant grouping for the catalog.""" 

2 

3from __future__ import annotations 

4 

5from dataclasses import dataclass 

6from enum import StrEnum 

7 

8from pydantic import BaseModel 

9 

10from lilbee.catalog.models import ModelFamily 

11from lilbee.core.config import cfg 

12 

13_BYTES_PER_GB = 1024**3 

14_FITS_HEADROOM_BYTES = 1 * _BYTES_PER_GB 

15 

16 

17class FitLevel(StrEnum): 

18 FITS = "fits" 

19 TIGHT = "tight" 

20 WONT_RUN = "wont_run" 

21 

22 

23@dataclass(frozen=True) 

24class FitChip: 

25 level: FitLevel 

26 headroom_gb: float 

27 

28 

29def compute_fit(model_size_bytes: int, available_bytes: int) -> FitChip: 

30 """Classify how a model footprint fits the available memory budget. 

31 

32 Headroom_gb is positive when the model fits and negative when it 

33 won't. The 1 GB band between FITS and TIGHT leaves room for the 

34 inference runtime, KV cache, and OS overhead beyond the raw weight 

35 file. 

36 """ 

37 headroom_bytes = available_bytes - model_size_bytes 

38 headroom_gb = headroom_bytes / _BYTES_PER_GB 

39 if headroom_bytes >= _FITS_HEADROOM_BYTES: 

40 level = FitLevel.FITS 

41 elif headroom_bytes >= 0: 

42 level = FitLevel.TIGHT 

43 else: 

44 level = FitLevel.WONT_RUN 

45 return FitChip(level=level, headroom_gb=headroom_gb) 

46 

47 

48def available_memory_for_fit() -> int | None: 

49 """Bytes available to a model after ``cfg.gpu_memory_fraction``, or None on probe failure. 

50 

51 Sums every GPU's memory (``total=True``) because lilbee tensor-splits a model 

52 too large for one card across the whole fleet; sizing the fit chip against a 

53 single card would wrongly mark a runnable split model "won't run". The actual 

54 per-card placement is decided precisely by the fleet planner at load time. 

55 

56 Single entry point so the TUI and the HTTP catalog handler classify fit 

57 against the same number; otherwise the same model would chip differently in 

58 each surface. 

59 """ 

60 try: 

61 from lilbee.providers.model_cache import get_available_memory 

62 

63 budget = get_available_memory(cfg.gpu_memory_fraction, total=True) 

64 except Exception: 

65 return None 

66 return budget + _expert_offload_headroom() 

67 

68 

69def _expert_offload_headroom() -> int: 

70 """System memory the fit budget may borrow when expert offload is configured. 

71 

72 A sparse model's experts live in system RAM under offload, so a host whose 

73 budget is discrete VRAM can run a model larger than that VRAM and must not 

74 be told otherwise. Zero unless the budget really is device memory: every 

75 other path (Apple unified memory, a non-NVIDIA or CPU-only host) already 

76 reports system RAM, and adding it twice would invent capacity. Zero too for a 

77 non-positive ``n_cpu_moe``, which offloads nothing. The chip is per-family and 

78 this budget is global, so it reads optimistically for a dense model pulled on 

79 an offload-enabled host (a sparse model gains the room, a dense one still 

80 fails to place); the planner sizes the real placement at load time. 

81 

82 Scaled from installed RAM, not from what is free this instant, to match the 

83 capacity basis of the VRAM budget it is added to. Mixing the two made a 

84 catalog entry fit or not fit depending on whatever else the machine happened 

85 to be doing when the page was drawn, and shrank the budget exactly when 

86 another model was already resident. 

87 """ 

88 from lilbee.providers.model_cache import has_nvidia_gpu, total_system_memory 

89 

90 if not (cfg.cpu_moe or (cfg.n_cpu_moe is not None and cfg.n_cpu_moe >= 1)): 

91 return 0 

92 try: 

93 if not has_nvidia_gpu(): 

94 return 0 

95 return int(total_system_memory() * cfg.gpu_memory_fraction) 

96 except Exception: 

97 return 0 

98 

99 

100class SizeVariantInfo(BaseModel): 

101 """One size/quant of a model family, serialised for HTTP responses.""" 

102 

103 size_label: str 

104 params: str 

105 size_gb: float 

106 ref: str 

107 

108 

109def family_size_variants(family: ModelFamily) -> list[SizeVariantInfo]: 

110 """Build the per-row size-variant strip for a featured ModelFamily, smallest first.""" 

111 variants = sorted(family.variants, key=lambda v: v.size_mb) 

112 return [ 

113 SizeVariantInfo( 

114 size_label=_size_variant_label(v.param_count, v.quant), 

115 params=v.param_count, 

116 size_gb=v.size_mb / 1024, 

117 ref=v.hf_repo, 

118 ) 

119 for v in variants 

120 ] 

121 

122 

123def _size_variant_label(param_count: str, quant: str) -> str: 

124 """Render the compact label for one size variant (``8B Q4_K_M``).""" 

125 pieces = [p for p in (param_count, quant) if p] 

126 return " ".join(pieces) if pieces else "--"