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 

47class WarmProgressTracker: 

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

49 

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

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

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

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

54 """ 

55 

56 def __init__(self) -> None: 

57 self._lock = threading.Lock() 

58 self._snapshot: WarmProgress | None = None 

59 self._started_at: float | None = None 

60 

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

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

63 with self._lock: 

64 self._started_at = time.monotonic() 

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

66 

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

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

69 self._advance( 

70 WarmPhase.READING_WEIGHTS, 

71 bytes_done=bytes_done, 

72 bytes_total=bytes_total, 

73 detail=detail, 

74 ) 

75 

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

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

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

79 

80 def ready(self) -> None: 

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

82 self._advance(WarmPhase.READY) 

83 

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

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

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

87 

88 def clear(self) -> None: 

89 """Drop any recorded warm state.""" 

90 with self._lock: 

91 self._snapshot = None 

92 self._started_at = None 

93 

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

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

96 with self._lock: 

97 if self._snapshot is None: 

98 return None 

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

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

101 

102 def _advance( 

103 self, 

104 phase: WarmPhase, 

105 *, 

106 bytes_done: int = 0, 

107 bytes_total: int = 0, 

108 detail: str | None = None, 

109 error: str | None = None, 

110 ) -> None: 

111 with self._lock: 

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

113 self._snapshot = WarmProgress( 

114 phase=phase, 

115 model_ref=model_ref, 

116 bytes_done=bytes_done, 

117 bytes_total=bytes_total, 

118 detail=detail, 

119 error=error, 

120 )