Coverage for src/lilbee/providers/warm_progress.py: 100%

1"""Server-side view of the chat model's cold load, for granular launch feedback. 

2 

3A launcher streams this while a chat model loads so the user sees real progress 

4(reading weights with a true byte percentage, then the engine load) instead of a 

5frozen "Warming..." line. The fleet provider drives the tracker through its warm 

6path; providers without a managed load report nothing and a launcher falls back 

7to a plain spinner. 

8""" 

9 

10from __future__ import annotations 

11 

12import threading 

13import time 

14from enum import StrEnum 

15 

16from pydantic import BaseModel 

17 

18 

19class WarmPhase(StrEnum): 

20 """The stage a chat-model cold load is in.""" 

21 

22 STARTING = "starting" 

23 """Warm thread has begun; the engine has not been touched yet.""" 

24 READING_WEIGHTS = "reading_weights" 

25 """Paging the GGUF shards off disk into the page cache; reports a true byte %.""" 

26 LOADING_ENGINE = "loading_engine" 

27 """The engine is loading the cached weights into VRAM; no byte signal, so 

28 surfaces show an indeterminate spinner bounded by readiness.""" 

29 READY = "ready" 

30 """The chat engine is loaded and can serve a first token.""" 

31 ERROR = "error" 

32 """The load failed; ``error`` carries the user-facing reason.""" 

33 

34 

35class WarmProgress(BaseModel): 

36 """A snapshot of the chat role's warm state, streamed to a launcher.""" 

37 

38 phase: WarmPhase 

39 model_ref: str | None = None 

40 bytes_done: int = 0 

41 bytes_total: int = 0 

42 detail: str | None = None 

43 error: str | None = None 

44 elapsed_s: float = 0.0 

45 

46 

47ACTIVE_WARM_PHASES = frozenset( 

48 {WarmPhase.STARTING, WarmPhase.READING_WEIGHTS, WarmPhase.LOADING_ENGINE} 

49) 

50"""Phases in which a load is still in flight. READY and ERROR are terminal.""" 

51 

52 

53def is_active_warm(snapshot: WarmProgress | None) -> bool: 

54 """Whether *snapshot* is a load still in flight, rather than absent or finished. 

55 

56 A terminal snapshot outlives the engine it describes: the tracker is not 

57 cleared when a role is swapped out, so READY can survive an eviction. 

58 """ 

59 return snapshot is not None and snapshot.phase in ACTIVE_WARM_PHASES 

60 

61 

62class WarmProgressTracker: 

63 """Thread-safe warm-state holder: the warm thread writes, handlers read. 

64 

65 The fleet warm-up runs on a daemon thread while the health / SSE handlers 

66 read concurrently, so every mutation and the snapshot read take the lock. 

67 ``elapsed_s`` is stamped at read time from the ``begin`` monotonic mark so 

68 callers always see a live elapsed without the writer ticking a clock. 

69 """ 

70 

71 def __init__(self) -> None: 

72 self._lock = threading.Lock() 

73 self._snapshot: WarmProgress | None = None 

74 self._started_at: float | None = None 

75 

76 def begin(self, model_ref: str | None) -> None: 

77 """Mark the start of a cold load; resets elapsed and clears prior state.""" 

78 with self._lock: 

79 self._started_at = time.monotonic() 

80 self._snapshot = WarmProgress(phase=WarmPhase.STARTING, model_ref=model_ref) 

81 

82 def reading(self, bytes_done: int, bytes_total: int, detail: str | None = None) -> None: 

83 """Report read-phase progress in bytes.""" 

84 self._advance( 

85 WarmPhase.READING_WEIGHTS, 

86 bytes_done=bytes_done, 

87 bytes_total=bytes_total, 

88 detail=detail, 

89 ) 

90 

91 def loading_engine(self, detail: str | None = None) -> None: 

92 """Mark the transition into the indeterminate VRAM-load phase.""" 

93 self._advance(WarmPhase.LOADING_ENGINE, detail=detail) 

94 

95 def ready(self) -> None: 

96 """Mark the chat engine ready to serve.""" 

97 self._advance(WarmPhase.READY) 

98 

99 def fail(self, message: str) -> None: 

100 """Mark the load as failed with a user-facing reason.""" 

101 self._advance(WarmPhase.ERROR, error=message) 

102 

103 def clear(self) -> None: 

104 """Drop any recorded warm state.""" 

105 with self._lock: 

106 self._snapshot = None 

107 self._started_at = None 

108 

109 def snapshot(self) -> WarmProgress | None: 

110 """Return a copy of the current state with live ``elapsed_s``, or None.""" 

111 with self._lock: 

112 if self._snapshot is None: 

113 return None 

114 elapsed = time.monotonic() - self._started_at if self._started_at is not None else 0.0 

115 return self._snapshot.model_copy(update={"elapsed_s": elapsed}) 

116 

117 def _advance( 

118 self, 

119 phase: WarmPhase, 

120 *, 

121 bytes_done: int = 0, 

122 bytes_total: int = 0, 

123 detail: str | None = None, 

124 error: str | None = None, 

125 ) -> None: 

126 with self._lock: 

127 model_ref = self._snapshot.model_ref if self._snapshot is not None else None 

128 self._snapshot = WarmProgress( 

129 phase=phase, 

130 model_ref=model_ref, 

131 bytes_done=bytes_done, 

132 bytes_total=bytes_total, 

133 detail=detail, 

134 error=error, 

135 )