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

1"""The engine reaper: a process that ends a shared engine's use once no user holds it. 

2 

3One runs per engine dir. Users can die in any order and by any means, so the 

4reaper applies the last-user-out rule after every death, crash included: it 

5stops the engine, or unloads a kept-warm engine's weights after the idle time. 

6""" 

7 

8from __future__ import annotations 

9 

10import logging 

11import sys 

12import time 

13from collections.abc import Callable 

14from dataclasses import dataclass 

15from pathlib import Path 

16 

17from filelock import FileLock 

18from filelock import Timeout as FileLockTimeout 

19 

20from lilbee.runtime.engine_lock import ( 

21 build_lock, 

22 keep_warm_idle_s, 

23 keep_warm_requested, 

24 live_users_exist, 

25 try_hold_reaper_lock, 

26) 

27 

28log = logging.getLogger(__name__) 

29 

30_POLL_S = 1.0 

31# A starting user probes the reaper lock for an instant; a claim that waits this 

32# long outlasts any probe, while a live reaper holds the lock for its whole life. 

33_CLAIM_TIMEOUT_S = 2.0 

34 

35 

36@dataclass 

37class _IdleCount: 

38 """When the last user of a kept-warm engine left, and whether its weights were unloaded.""" 

39 

40 since: float | None = None 

41 unloaded: bool = False 

42 

43 def reset(self) -> None: 

44 self.since = None 

45 self.unloaded = False 

46 

47 

48def _watch_is_over( 

49 engine_dir: Path, reaper_lock: FileLock, idle: _IdleCount, clock: Callable[[], float] 

50) -> bool: 

51 """Under the build lock, end the use of an engine nobody holds; True when nothing is left. 

52 

53 The reaper lock is released before the build lock, so a build that follows 

54 always finds the dir unwatched and starts a new reaper. 

55 """ 

56 # Heavy (about 290 ms and 40 MB): loaded only once no user is left, not for the 

57 # reaper's whole life beside a running engine. 

58 from lilbee.providers.fleet.swap_manager import engine_record_exists, stop_engine 

59 

60 with build_lock(engine_dir): 

61 if live_users_exist(engine_dir): 

62 idle.reset() 

63 return False 

64 recorded = engine_record_exists(engine_dir) 

65 kept_warm = recorded and keep_warm_requested(engine_dir) 

66 if kept_warm and not _unload_when_idle(engine_dir, idle, clock()): 

67 return False 

68 if recorded: 

69 stop_engine(engine_dir) 

70 log.info( 

71 "Engine stopped at %s (%s)", 

72 engine_dir, 

73 "it stopped answering" if kept_warm else "every user is gone", 

74 ) 

75 reaper_lock.release() 

76 return True 

77 

78 

79def _unload_when_idle(engine_dir: Path, idle: _IdleCount, now: float) -> bool: 

80 """Unload a kept-warm engine's weights after its idle time with no user. 

81 

82 A refused unload starts a new idle window, so the next try waits one window. 

83 Returns True when no recorded engine answers: it is gone and nothing is left to watch. 

84 """ 

85 # Loaded with the heavy swap_manager import above, once no user is left. 

86 from lilbee.core.config.defaults import ENGINE_IDLE_TTL_MINUTES_DEFAULT 

87 from lilbee.providers.fleet.swap_manager import engine_answers, unload_engine_models 

88 

89 if idle.since is None: 

90 idle.since = now 

91 idle_s = keep_warm_idle_s(engine_dir, unset_minutes=ENGINE_IDLE_TTL_MINUTES_DEFAULT) 

92 if idle.unloaded or now - idle.since < idle_s: 

93 return False 

94 if unload_engine_models(engine_dir): 

95 idle.unloaded = True 

96 log.info("Engine weights unloaded at %s (idle with no user)", engine_dir) 

97 return False 

98 idle.since = now 

99 return not engine_answers(engine_dir) 

100 

101 

102def watch( 

103 engine_dir: Path, 

104 reaper_lock: FileLock, 

105 *, 

106 poll_s: float = _POLL_S, 

107 clock: Callable[[], float] = time.monotonic, 

108) -> None: 

109 """Poll *engine_dir* until its engine is stopped or gone.""" 

110 idle = _IdleCount() 

111 while True: 

112 if live_users_exist(engine_dir): 

113 idle.reset() 

114 else: 

115 try: 

116 if _watch_is_over(engine_dir, reaper_lock, idle, clock): 

117 return 

118 except FileLockTimeout: 

119 log.info("Engine build lock at %s still held; checking again.", engine_dir) 

120 time.sleep(poll_s) 

121 

122 

123def main() -> None: 

124 """Watch the engine dir named in argv, unless another reaper already does.""" 

125 engine_dir = Path(sys.argv[1]) 

126 reaper_lock = try_hold_reaper_lock(engine_dir, timeout=_CLAIM_TIMEOUT_S) 

127 if reaper_lock is not None: 

128 watch(engine_dir, reaper_lock) 

129 

130 

131if __name__ == "__main__": # pragma: no cover - process entry glue 

132 main()