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 os 

14from contextlib import contextmanager 

15from dataclasses import dataclass, field 

16from functools import lru_cache 

17from pathlib import Path 

18from typing import TYPE_CHECKING 

19 

20from filelock import FileLock, SoftFileLock 

21from filelock import Timeout as FileLockTimeout 

22 

23if TYPE_CHECKING: 

24 from collections.abc import Iterator 

25 

26log = logging.getLogger(__name__) 

27 

28ENGINE_DIR_ENV = "LILBEE_ENGINE_DIR" 

29_BUILD_LOCK_NAME = "engine.lock" 

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

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

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

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

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

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

36_BUILD_LOCK_TIMEOUT_S = 90.0 

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

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

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

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

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

42# to stop alive and holding VRAM. 

43_BEST_EFFORT_LOCK_TIMEOUT_S = 5.0 

44_USERS_DIRNAME = "engine-users" 

45_USER_LOCK_SUFFIX = ".lock" 

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

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

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

49_KEEP_WARM_SUFFIX = ".keep-warm" 

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

51_FLOCK_PROBE_PREFIX = ".flock-probe." 

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

53_PROBE_TIMEOUT_S = 0.0 

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

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

56# this never waits in practice. 

57_HOLD_TIMEOUT_S = 10.0 

58 

59 

60def machine_engine_dir() -> Path: 

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

62 from lilbee.core.system import default_state_dir 

63 

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

65 if override: 

66 return Path(override) 

67 return default_state_dir() / "engine" 

68 

69 

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

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

72 return config_root / "data" / "engine" 

73 

74 

75@contextmanager 

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

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

78 

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

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

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

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

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

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

85 and proceeds without the lock. 

86 """ 

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

88 lock = FileLock(engine_dir / _BUILD_LOCK_NAME) 

89 wait_s = _BEST_EFFORT_LOCK_TIMEOUT_S if best_effort else _BUILD_LOCK_TIMEOUT_S 

90 try: 

91 lock.acquire(timeout=_PROBE_TIMEOUT_S) 

92 except FileLockTimeout: 

93 log.info( 

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

95 wait_s, 

96 engine_dir, 

97 ) 

98 try: 

99 lock.acquire(timeout=wait_s) 

100 except FileLockTimeout: 

101 if not best_effort: 

102 raise 

103 log.warning( 

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

105 engine_dir, 

106 wait_s, 

107 ) 

108 yield 

109 return 

110 try: 

111 yield 

112 finally: 

113 lock.release() 

114 

115 

116@lru_cache(maxsize=8) 

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

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

119 

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

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

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

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

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

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

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

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

128 

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

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

131 """ 

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

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

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

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

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

137 lock = FileLock(probe_path, thread_local=False) 

138 try: 

139 lock.acquire(timeout=_HOLD_TIMEOUT_S) 

140 except (FileLockTimeout, OSError): 

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

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

143 return True 

144 try: 

145 return not isinstance(lock, SoftFileLock) 

146 finally: 

147 lock.release() 

148 probe_path.unlink(missing_ok=True) 

149 

150 

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

152 return engine_dir / _USERS_DIRNAME 

153 

154 

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

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

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

158 

159 

160def request_keep_warm(engine_dir: Path, config_root: Path) -> None: 

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

162 

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

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

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

166 its own mark. ``stop_engine`` clears the set, so a rebuilt engine starts 

167 unmarked. 

168 """ 

169 marker = _keep_warm_path(engine_dir, config_root) 

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

171 marker.touch() 

172 

173 

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

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

176 

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

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

179 """ 

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

181 

182 

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

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

185 

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

187 marker is the case it exists for. 

188 """ 

189 return any(_users_dir(engine_dir).glob(f"*{_KEEP_WARM_SUFFIX}")) 

190 

191 

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

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

194 for marker in _users_dir(engine_dir).glob(f"*{_KEEP_WARM_SUFFIX}"): 

195 marker.unlink(missing_ok=True) 

196 

197 

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

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

200 

201 

202@dataclass 

203class UserLockHold: 

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

205 

206 engine_dir: Path 

207 path: Path 

208 _lock: FileLock = field(repr=False) 

209 

210 def release_and_check_last(self) -> bool: 

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

212 

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

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

215 deleted in passing. 

216 """ 

217 if self._lock.is_locked: 

218 self._lock.release() 

219 if not self._lock.is_locked: 

220 self.path.unlink(missing_ok=True) 

221 return not live_users_exist(self.engine_dir) 

222 

223 

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

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

226 

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

228 pids to simulate peers from one process. 

229 """ 

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

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

232 lock = _user_file_lock(path) 

233 lock.acquire(timeout=_HOLD_TIMEOUT_S) 

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

235 

236 

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

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

239 

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

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

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

243 deadlock detection. 

244 """ 

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

246 

247 

248def live_users_exist(engine_dir: Path) -> bool: 

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

250 users = _users_dir(engine_dir) 

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

252 probe = _user_file_lock(path) 

253 if probe.is_locked: 

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

255 return True 

256 try: 

257 probe.acquire(timeout=_PROBE_TIMEOUT_S) 

258 except FileLockTimeout: 

259 return True 

260 probe.release() 

261 path.unlink(missing_ok=True) 

262 return False