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

1"""Cross-process lifecycle primitives for the shared engine. 

2 

3The engine (llama-swap + llama-server processes) is machine-level infrastructure: 

4any lilbee process may build it, every compatible process binds to it, and the 

5kernel arbitrates liveness through file locks so no pid bookkeeping can go stale. 

6The mechanics are agnostic to what they front. 

7""" 

8 

9from __future__ import annotations 

10 

11import hashlib 

12import logging 

13import math 

14import os 

15from contextlib import contextmanager, suppress 

16from dataclasses import dataclass, field 

17from functools import lru_cache 

18from pathlib import Path 

19from typing import TYPE_CHECKING 

20 

21from filelock import FileLock, SoftFileLock 

22from filelock import Timeout as FileLockTimeout 

23 

24from lilbee.runtime.shard_process import shard_process 

25 

26if TYPE_CHECKING: 

27 from collections.abc import Iterator 

28 

29log = logging.getLogger(__name__) 

30 

31ENGINE_DIR_ENV = "LILBEE_ENGINE_DIR" 

32_BUILD_LOCK_NAME = "engine.lock" 

33# A legitimate build holds the lock for the llama-swap spawn (a 30 s boot budget) 

34# and, before it, the planning sweep, which shells out to gguf-parser once per 

35# split candidate and per context probe. The model itself loads lazily on the 

36# first request, outside the lock. So honest contention can run well past the 

37# spawn budget on a wide box, and the bound is here to catch a wedged holder 

38# rather than to block forever, which would deadlock every startup and exit. 

39_BUILD_LOCK_TIMEOUT_S = 90.0 

40# What a caller that will proceed without the lock waits for it. Teardown and 

41# config-change are best-effort by definition, so spending the build timeout 

42# before giving up buys nothing and costs everything: an init system's TERM 

43# window is typically 90 seconds, and a shutdown that waits that long to then 

44# proceed anyway gets SIGKILLed mid-teardown, leaving the engines it was about 

45# to stop alive and holding VRAM. 

46_BEST_EFFORT_LOCK_TIMEOUT_S = 5.0 

47_USERS_DIRNAME = "engine-users" 

48# Held for its whole life by the one reaper process watching a dir's engine. 

49_REAPER_LOCK_NAME = "engine-reaper.lock" 

50_USER_LOCK_SUFFIX = ".lock" 

51# One opt-in file per installation, keyed by config root (not pid, which changes 

52# every run). Plain files, not locks: an opt-in outlives its process. Per-install 

53# so a restart reclaims its own prior mark while a peer's stays distinct. 

54_KEEP_WARM_SUFFIX = ".keep-warm" 

55_SECONDS_PER_MINUTE = 60 

56# Throwaway per-process file used to ask whether flock really works here. 

57_FLOCK_PROBE_PREFIX = ".flock-probe." 

58# Non-blocking probe: a live peer refuses instantly. 

59_PROBE_TIMEOUT_S = 0.0 

60# Finite: infinite acquires trip filelock's thread-local deadlock detection 

61# after a cross-thread release. The pid-named file has no live contender, so 

62# this never waits in practice. 

63_HOLD_TIMEOUT_S = 10.0 

64 

65 

66def machine_engine_dir() -> Path: 

67 """The per-OS-user engine slot every lilbee process scans first.""" 

68 from lilbee.core.system import default_state_dir 

69 

70 if shard_process.engine_dir is not None: 

71 return shard_process.engine_dir 

72 override = os.environ.get(ENGINE_DIR_ENV, "").strip() 

73 if override: 

74 return Path(override) 

75 return default_state_dir() / "engine" 

76 

77 

78def private_engine_dir(config_root: Path) -> Path: 

79 """The overflow engine dir for one config root, used when the slot is incompatible.""" 

80 return config_root / "data" / "engine" 

81 

82 

83@contextmanager 

84def build_lock(engine_dir: Path, *, best_effort: bool = False) -> Iterator[None]: 

85 """Serialize scan-or-build (and config-change restarts) for *engine_dir*. 

86 

87 Acquired with a finite timeout so a wedged holder cannot deadlock the machine: 

88 a blocking acquire would hang every startup behind it and, on the shutdown and 

89 config-change paths, leave processes unable to exit. A wait is logged so a 

90 stall is visible. On timeout a build caller raises (it could not acquire the 

91 engine, better than an unbounded hang); a *best_effort* caller -- teardown and 

92 config-change, which must not wedge a dying or reconfiguring process -- logs 

93 and proceeds without the lock. 

94 """ 

95 engine_dir.mkdir(parents=True, exist_ok=True) 

96 lock = FileLock(engine_dir / _BUILD_LOCK_NAME) 

97 wait_s = _BEST_EFFORT_LOCK_TIMEOUT_S if best_effort else _BUILD_LOCK_TIMEOUT_S 

98 try: 

99 lock.acquire(timeout=_PROBE_TIMEOUT_S) 

100 except FileLockTimeout: 

101 log.info( 

102 "Waiting up to %.0fs for another process to build the engine at %s", 

103 wait_s, 

104 engine_dir, 

105 ) 

106 try: 

107 lock.acquire(timeout=wait_s) 

108 except FileLockTimeout: 

109 if not best_effort: 

110 raise 

111 log.warning( 

112 "Engine build lock at %s held past %.0fs; proceeding without it.", 

113 engine_dir, 

114 wait_s, 

115 ) 

116 yield 

117 return 

118 try: 

119 yield 

120 finally: 

121 lock.release() 

122 

123 

124@lru_cache(maxsize=8) 

125def kernel_arbitrates_locks(engine_dir: Path) -> bool: 

126 """Whether *engine_dir*'s filesystem gives real kernel-arbitrated file locks. 

127 

128 The whole membership scheme rests on the kernel releasing a lock on any 

129 death, so no pid bookkeeping can go stale. On a filesystem where flock 

130 returns ENOSYS (FUSE, some NFS mounts) filelock silently rewrites itself to 

131 SoftFileLock with only a Python warning, and the guarantee is gone: that 

132 fallback path opens the lock file with O_TRUNC and unlinks it before 

133 re-acquiring, so a process merely *probing* a live member's lock destroys 

134 it. live_users_exist would then report an empty slot while members are 

135 serving, and the last-out stop would kill an engine in use. 

136 

137 Probed by acquiring a throwaway lock in the dir and checking what filelock 

138 turned it into; cached, since the answer is a property of the mount. 

139 """ 

140 engine_dir.mkdir(parents=True, exist_ok=True) 

141 # Named per process: the probe asks what the filesystem supports, which needs 

142 # no mutual exclusion. A shared probe file would make every lilbee queue on 

143 # one lock, and this runs while the in-process build lock is held. 

144 probe_path = engine_dir / f"{_FLOCK_PROBE_PREFIX}{os.getpid()}" 

145 lock = FileLock(probe_path, thread_local=False) 

146 try: 

147 lock.acquire(timeout=_HOLD_TIMEOUT_S) 

148 except (FileLockTimeout, OSError): 

149 # An unusable probe file says nothing about flock support; assume the 

150 # filesystem is fine rather than refuse the shared slot over it. 

151 return True 

152 try: 

153 return not isinstance(lock, SoftFileLock) 

154 finally: 

155 lock.release() 

156 probe_path.unlink(missing_ok=True) 

157 

158 

159def _users_dir(engine_dir: Path) -> Path: 

160 return engine_dir / _USERS_DIRNAME 

161 

162 

163def _keep_warm_path(engine_dir: Path, config_root: Path) -> Path: 

164 token = hashlib.blake2b(str(config_root).encode(), digest_size=8).hexdigest() 

165 return _users_dir(engine_dir) / f"{token}{_KEEP_WARM_SUFFIX}" 

166 

167 

168def request_keep_warm(engine_dir: Path, config_root: Path, idle_minutes: int) -> None: 

169 """Record that the installation at *config_root* wants *engine_dir*'s engine warm. 

170 

171 The slot is shared across installations whose configs differ, so any one's 

172 opt-in keeps the engine warm even when a default-config sibling is last out. 

173 Keyed by config root, not pid: a restart of the same installation reclaims 

174 its own mark. The marker holds *idle_minutes*, how long the weights stay 

175 loaded after the last user leaves. ``stop_engine`` clears the set, so a 

176 rebuilt engine starts unmarked. 

177 """ 

178 marker = _keep_warm_path(engine_dir, config_root) 

179 marker.parent.mkdir(parents=True, exist_ok=True) 

180 marker.write_text(str(idle_minutes), encoding="utf-8") 

181 

182 

183def withdraw_keep_warm(engine_dir: Path, config_root: Path) -> None: 

184 """Drop this installation's opt-in for *engine_dir*, leaving every peer's intact. 

185 

186 Reclaims the mark this installation left, so flipping the setting off (even 

187 across a restart) lets the engine stop instead of staying warm forever. 

188 """ 

189 _keep_warm_path(engine_dir, config_root).unlink(missing_ok=True) 

190 

191 

192def keep_warm_requested(engine_dir: Path) -> bool: 

193 """Whether any user of *engine_dir* asked for the engine to stay resident. 

194 

195 Not gated on liveness: an opt-in means "outlive me", so an exited user's 

196 marker is the case it exists for. 

197 """ 

198 return bool(_keep_warm_markers(engine_dir)) 

199 

200 

201def keep_warm_idle_s(engine_dir: Path, *, unset_minutes: int) -> float: 

202 """Seconds a kept-warm engine's weights stay loaded with no user; the longest opt-in wins. 

203 

204 A marker of 0 never unloads (infinite). An empty or unreadable marker counts 

205 as *unset_minutes*. 

206 """ 

207 minutes = [_marker_minutes(path, unset_minutes) for path in _keep_warm_markers(engine_dir)] 

208 longest = max(minutes, default=unset_minutes) 

209 return math.inf if 0 in minutes else longest * _SECONDS_PER_MINUTE 

210 

211 

212def _keep_warm_markers(engine_dir: Path) -> list[Path]: 

213 return list(_users_dir(engine_dir).glob(f"*{_KEEP_WARM_SUFFIX}")) 

214 

215 

216def _marker_minutes(marker: Path, unset_minutes: int) -> int: 

217 try: 

218 return int(marker.read_text(encoding="utf-8")) 

219 except (OSError, ValueError): 

220 return unset_minutes 

221 

222 

223def clear_keep_warm(engine_dir: Path) -> None: 

224 """Forget every persistence opt-in for *engine_dir*; the engine they applied to is gone.""" 

225 for marker in _keep_warm_markers(engine_dir): 

226 marker.unlink(missing_ok=True) 

227 

228 

229def _unlink_best_effort(path: Path) -> None: 

230 """Delete a released user lock file; a leftover is harmless. 

231 

232 Windows refuses to delete a file another handle has open (WinError 32), 

233 for example a concurrent probe. The next probe deletes the leftover. 

234 """ 

235 with suppress(OSError): 

236 path.unlink(missing_ok=True) 

237 

238 

239def _user_lock_path(engine_dir: Path, pid: int) -> Path: 

240 return _users_dir(engine_dir) / f"{pid}{_USER_LOCK_SUFFIX}" 

241 

242 

243@dataclass 

244class UserLockHold: 

245 """One process's held membership in an engine's user set.""" 

246 

247 engine_dir: Path 

248 path: Path 

249 _lock: FileLock = field(repr=False) 

250 

251 def release_and_check_last(self) -> bool: 

252 """Release this hold and report whether no live peers remain. 

253 

254 Idempotent. The lock file is removed when the last in-process hold 

255 releases; acquirable peer files belong to dead processes and are 

256 deleted in passing. 

257 """ 

258 if self._lock.is_locked: 

259 self._lock.release() 

260 if not self._lock.is_locked: 

261 _unlink_best_effort(self.path) 

262 return not live_users_exist(self.engine_dir) 

263 

264 

265def hold_user_lock(engine_dir: Path, pid: int | None = None) -> UserLockHold: 

266 """Hold this process's user lock for *engine_dir* until released or death. 

267 

268 *pid* names the lock file (defaults to this process); tests pass explicit 

269 pids to simulate peers from one process. 

270 """ 

271 path = _user_lock_path(engine_dir, os.getpid() if pid is None else pid) 

272 path.parent.mkdir(parents=True, exist_ok=True) 

273 lock = _user_file_lock(path) 

274 lock.acquire(timeout=_HOLD_TIMEOUT_S) 

275 return UserLockHold(engine_dir=engine_dir, path=path, _lock=lock) 

276 

277 

278def _user_file_lock(path: Path) -> FileLock: 

279 """One process-wide reentrant lock instance per user-lock path. 

280 

281 Not thread-local: acquire and release run on different threads. 

282 Singleton: two providers in one process hold the same pid-named file, 

283 and separate instances over fcntl falsely succeed or trip filelock's 

284 deadlock detection. 

285 """ 

286 return FileLock(path, thread_local=False, is_singleton=True) 

287 

288 

289def live_users_exist(engine_dir: Path, *, exclude_self: bool = False) -> bool: 

290 """Probe every user lock file; clean the dead, report whether any refused. 

291 

292 *exclude_self* skips this process's own lock file, so the answer is whether 

293 another process uses the engine. 

294 """ 

295 users = _users_dir(engine_dir) 

296 own = _user_lock_path(engine_dir, os.getpid()) if exclude_self else None 

297 for path in sorted(users.glob(f"*{_USER_LOCK_SUFFIX}")): 

298 if path == own: 

299 continue 

300 probe = _user_file_lock(path) 

301 if probe.is_locked: 

302 # Held by this process; acquiring would reentrantly succeed. 

303 return True 

304 try: 

305 probe.acquire(timeout=_PROBE_TIMEOUT_S) 

306 except FileLockTimeout: 

307 return True 

308 probe.release() 

309 _unlink_best_effort(path) 

310 return False 

311 

312 

313def try_hold_reaper_lock(engine_dir: Path, *, timeout: float = _PROBE_TIMEOUT_S) -> FileLock | None: 

314 """Hold *engine_dir*'s reaper lock, or return None when it stays held for *timeout*.""" 

315 engine_dir.mkdir(parents=True, exist_ok=True) 

316 lock = FileLock(engine_dir / _REAPER_LOCK_NAME, thread_local=False) 

317 try: 

318 lock.acquire(timeout=timeout) 

319 except FileLockTimeout: 

320 return None 

321 return lock 

322 

323 

324def reaper_running(engine_dir: Path) -> bool: 

325 """Whether a live reaper process watches *engine_dir*'s engine.""" 

326 lock = try_hold_reaper_lock(engine_dir) 

327 if lock is None: 

328 return True 

329 lock.release() 

330 return False