Coverage for src/lilbee/catalog/download_process.py: 100%

1"""One model download per child process, so cancelling terminates the child. 

2 

3hf_xet cancels only at session granularity within a process (one session per 

4PID), so a terminatable child is what makes per-download cancel real, and it 

5is also the only stop a wedged transfer cannot refuse. 

6""" 

7 

8from __future__ import annotations 

9 

10import logging 

11import multiprocessing 

12import os 

13import sys 

14import time 

15from dataclasses import dataclass 

16from pathlib import Path 

17from typing import TYPE_CHECKING, Literal, Protocol 

18 

19from lilbee.catalog.download_progress import ProgressCallback 

20from lilbee.catalog.models import CatalogModel 

21from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError 

22 

23if TYPE_CHECKING: 

24 from multiprocessing.connection import Connection 

25 

26log = logging.getLogger(__name__) 

27 

28_POLL_INTERVAL_S = 0.2 

29 

30_EXIT_GRACE_S = 10.0 

31 

32_PROGRESS_MIN_INTERVAL_S = 0.1 

33 

34STALL_FLOOR_BYTES = 256 * 1024 

35"""Minimum new bytes a child must report for its transfer to count as alive. 

36 

37A wedged connection can trickle a few bytes a minute, which an any-activity 

38check reads as progress; ~2 KB/s is far below any usable model download.""" 

39 

40STALL_ATTEMPTS = 3 

41"""Children the parent spawns for one download before it reports a stall.""" 

42 

43_STALL_DEADLINE_S = 90.0 

44"""Seconds of no reported progress after which the child is terminated. 

45 

46Well past the hub's own 10 second read timeout and the retries it makes under 

47it, so a child only expires on a transfer those mechanisms cannot wake.""" 

48 

49_STARTUP_DEADLINE_S = 120.0 

50"""Seconds a child may stay silent before its first bytes. 

51 

52It spawns, imports huggingface_hub, lists the repo, and resolves the 

53filename, the dominant term, bounded by lilbee's own 30 second listing 

54timeout and 10 second header-probe timeout, not the hub's. Every file the 

55child then resolves restarts this clock, so the budget does not grow with 

56the shard count.""" 

57 

58# The child's translated errors, rebuilt in the parent by type name. 

59_ERRORS_BY_NAME: dict[str, type[Exception]] = {PermissionError.__name__: PermissionError} 

60 

61 

62@dataclass(frozen=True) 

63class _Progress: 

64 """A child's byte counters as it transfers.""" 

65 

66 kind: Literal["progress"] 

67 downloaded: int 

68 total: int 

69 

70 

71@dataclass(frozen=True) 

72class _Probed: 

73 """A child's note that it resolved one file with the Hub before transferring. 

74 

75 *blob* names the cache file that transfer writes, so the parent knows which 

76 temporary files the child owns and can clear them after terminating it. 

77 """ 

78 

79 kind: Literal["probed"] 

80 blob: str | None 

81 

82 

83@dataclass(frozen=True) 

84class _Done: 

85 """A child's success verdict carrying the downloaded model's path.""" 

86 

87 kind: Literal["done"] 

88 path: str 

89 

90 

91@dataclass(frozen=True) 

92class _Failed: 

93 """A child's failure, serialized because exception objects may not unpickle.""" 

94 

95 kind: Literal["failed"] 

96 error_type: str 

97 message: str 

98 

99 

100_ChildMessage = _Progress | _Probed | _Done | _Failed 

101 

102 

103class _ChildStalledError(Exception): 

104 """A child stopped reporting bytes, so its transfer is retried in a new child.""" 

105 

106 

107class _StallDeadline: 

108 """Times a child against the bytes it reports, so a wedged transfer expires.""" 

109 

110 def __init__( 

111 self, 

112 deadline_s: float = _STALL_DEADLINE_S, 

113 floor_bytes: int = STALL_FLOOR_BYTES, 

114 startup_s: float = _STARTUP_DEADLINE_S, 

115 ) -> None: 

116 self._deadline_s = deadline_s 

117 self._floor_bytes = floor_bytes 

118 self._startup_s = startup_s 

119 self._since = time.monotonic() 

120 self._base = 0 

121 self._moving = False 

122 

123 def saw(self, downloaded: int) -> None: 

124 """Restart the clock once the byte floor of new data has arrived.""" 

125 if downloaded - self._base < self._floor_bytes: 

126 return 

127 self._base = downloaded 

128 self._since = time.monotonic() 

129 self._moving = True 

130 

131 def probed(self) -> None: 

132 """Restart the clock for a child that resolved a file before its first bytes.""" 

133 self._since = time.monotonic() 

134 

135 def expired(self) -> bool: 

136 """Whether the child has reported nothing for the budget of its phase.""" 

137 budget = self._deadline_s if self._moving else self._startup_s 

138 return time.monotonic() - self._since >= budget 

139 

140 

141class _Worker(Protocol): 

142 """The slice of ``multiprocessing.Process`` the parent relay drives.""" 

143 

144 @property 

145 def exitcode(self) -> int | None: ... 

146 

147 def is_alive(self) -> bool: ... 

148 

149 def terminate(self) -> None: ... 

150 

151 def kill(self) -> None: ... 

152 

153 def join(self, timeout: float | None = None) -> None: ... 

154 

155 

156class _PipeProgress: 

157 """Byte-progress callback that relays over the pipe at ~10 Hz plus the final event.""" 

158 

159 def __init__(self, conn: Connection) -> None: 

160 self._conn = conn 

161 self._last_sent: float | None = None 

162 

163 def __call__(self, downloaded: int, total: int) -> None: 

164 now = time.monotonic() 

165 final = total > 0 and downloaded >= total 

166 throttled = self._last_sent is not None and now - self._last_sent < _PROGRESS_MIN_INTERVAL_S 

167 if not final and throttled: 

168 return 

169 self._last_sent = now 

170 self._conn.send(_Progress(kind="progress", downloaded=downloaded, total=total)) 

171 

172 

173class _PipeProbe: 

174 """Callback that tells the parent one file was resolved with the Hub.""" 

175 

176 def __init__(self, conn: Connection) -> None: 

177 self._conn = conn 

178 

179 def __call__(self, blob: str | None) -> None: 

180 self._conn.send(_Probed(kind="probed", blob=blob)) 

181 

182 

183def download_in_subprocess( 

184 entry: CatalogModel, 

185 models_dir: Path, 

186 token: str | None, 

187 *, 

188 on_progress: ProgressCallback | None, 

189 cancel: CancelSignal, 

190) -> Path: 

191 """Run one download in its own process, retrying in a fresh child after a stall. 

192 

193 A set *cancel* signal terminates the child, which is the only way to free 

194 the bandwidth of a running hf_xet transfer mid-flight. The same terminate 

195 ends a child that stops reporting bytes, so stopping a wedged transfer does 

196 not depend on it answering an abort. 

197 """ 

198 if cancel.is_set(): 

199 raise TaskCancelledError 

200 for attempt in range(STALL_ATTEMPTS): 

201 try: 

202 return _run_one_attempt(entry, models_dir, token, on_progress, cancel) 

203 except _ChildStalledError: 

204 log.warning( 

205 "Transfer of %s stalled (attempt %d/%d); starting again in a new process.", 

206 entry.hf_repo, 

207 attempt + 1, 

208 STALL_ATTEMPTS, 

209 ) 

210 raise RuntimeError(_stalled_download_message(entry.hf_repo)) 

211 

212 

213def _stalled_download_message(hf_repo: str) -> str: 

214 """The user-facing error for a transfer that never started moving again.""" 

215 return ( 

216 f"Download of {hf_repo} kept stalling with almost no data arriving. " 

217 "Check the network connection and retry. Files that finished are kept; " 

218 "the file in flight starts again." 

219 ) 

220 

221 

222def _run_one_attempt( 

223 entry: CatalogModel, 

224 models_dir: Path, 

225 token: str | None, 

226 on_progress: ProgressCallback | None, 

227 cancel: CancelSignal, 

228) -> Path: 

229 """Spawn one child, relay it, and clear the partials it owned on any failure. 

230 

231 ``KeyboardInterrupt`` is caught because Ctrl-C during a download leaves the 

232 same partial file any other interrupted attempt leaves. ``SystemExit`` is 

233 not, so interpreter shutdown skips the filesystem work. 

234 """ 

235 worker, receiver = _start_worker(entry, models_dir, token) 

236 owned: list[str] = [] 

237 try: 

238 try: 

239 return _relay_until_done(entry, worker, receiver, on_progress, cancel, owned) 

240 finally: 

241 _stop_worker(worker) 

242 receiver.close() 

243 # The inner finally must terminate the child before this delete runs, or the 

244 # delete races the child's own writer. 

245 except (Exception, KeyboardInterrupt): 

246 _discard_partial_blobs(models_dir, entry.hf_repo, owned) 

247 raise 

248 

249 

250def _discard_partial_blobs(models_dir: Path, hf_repo: str, owned: list[str]) -> None: 

251 """Delete the partial blobs the child reported owning, which nothing reads.""" 

252 # heavy: lilbee.catalog.download (>50ms; huggingface_hub fanout) 

253 from lilbee.catalog.download import discard_partial_blobs 

254 

255 discard_partial_blobs(models_dir, hf_repo, owned) 

256 

257 

258def _start_worker( 

259 entry: CatalogModel, models_dir: Path, token: str | None 

260) -> tuple[_Worker, Connection]: 

261 """Spawn the download child; fork is unsafe under the parent's threads. 

262 

263 Daemonic on purpose: at interpreter exit multiprocessing terminates daemon 

264 children but joins live non-daemon ones, and quitting the app mid-download 

265 must not wait for a multi-GB transfer. 

266 """ 

267 context = multiprocessing.get_context("spawn") 

268 receiver, sender = context.Pipe(duplex=False) 

269 worker = context.Process( 

270 target=_run_download_child, 

271 args=(sender, entry, str(models_dir), token), 

272 name=f"lilbee-download-{entry.hf_repo}", 

273 daemon=True, 

274 ) 

275 worker.start() 

276 sender.close() 

277 return worker, receiver 

278 

279 

280def _relay_until_done( 

281 entry: CatalogModel, 

282 worker: _Worker, 

283 receiver: Connection, 

284 on_progress: ProgressCallback | None, 

285 cancel: CancelSignal, 

286 owned: list[str], 

287) -> Path: 

288 """Forward child messages until its verdict, polling *cancel* between them. 

289 

290 The budget is tested only when the pipe holds nothing, so a child whose 

291 verdict is already queued is never terminated with its answer unread. 

292 """ 

293 deadline = _StallDeadline() 

294 while True: 

295 if cancel.is_set(): 

296 raise TaskCancelledError 

297 if receiver.poll(_POLL_INTERVAL_S): 

298 try: 

299 message = receiver.recv() 

300 except (EOFError, OSError): 

301 # No verdict is reachable once the read fails, however it 

302 # fails: POSIX raises EOFError at the closed pipe, Windows a 

303 # BrokenPipeError. 

304 raise _died_silently(entry, worker) from None 

305 verdict = _apply(message, on_progress, deadline, owned) 

306 if verdict is not None: 

307 return verdict 

308 elif receiver.poll(): 

309 continue # a message landed while the poll was expiring; read it next pass 

310 elif not worker.is_alive(): 

311 raise _died_silently(entry, worker) 

312 elif deadline.expired(): 

313 raise _ChildStalledError 

314 

315 

316def _died_silently(entry: CatalogModel, worker: _Worker) -> RuntimeError: 

317 """The error for a child that exited without reporting a verdict.""" 

318 return RuntimeError( 

319 f"Download of {entry.hf_repo} stopped: its process exited with code {worker.exitcode}." 

320 ) 

321 

322 

323def _apply( 

324 message: _ChildMessage, 

325 on_progress: ProgressCallback | None, 

326 deadline: _StallDeadline, 

327 owned: list[str], 

328) -> Path | None: 

329 """Act on one child message, returning the path once the child reports done.""" 

330 if message.kind == "progress": 

331 deadline.saw(message.downloaded) 

332 if on_progress is not None: 

333 on_progress(message.downloaded, message.total) 

334 return None 

335 if message.kind == "probed": 

336 deadline.probed() 

337 if message.blob is not None: 

338 owned.append(message.blob) 

339 return None 

340 if message.kind == "done": 

341 return Path(message.path) 

342 raise _ERRORS_BY_NAME.get(message.error_type, RuntimeError)(message.message) 

343 

344 

345def _stop_worker(worker: _Worker) -> None: 

346 """Terminate a live child and reap it, escalating to kill if TERM is ignored.""" 

347 if worker.is_alive(): 

348 worker.terminate() 

349 worker.join(_EXIT_GRACE_S) 

350 if worker.is_alive(): 

351 worker.kill() 

352 worker.join(_EXIT_GRACE_S) 

353 

354 

355def _run_download_child( 

356 conn: Connection, entry: CatalogModel, models_dir: str, token: str | None 

357) -> None: 

358 """Child-process entry: fetch the files and report the verdict over *conn*.""" 

359 _silence_output() 

360 # heavy: lilbee.catalog.download (>50ms; huggingface_hub fanout) 

361 from lilbee.catalog.download import fetch_model_files 

362 

363 try: 

364 path = fetch_model_files( 

365 entry, 

366 Path(models_dir), 

367 token, 

368 on_progress=_PipeProgress(conn), 

369 on_probe=_PipeProbe(conn), 

370 ) 

371 conn.send(_Done(kind="done", path=str(path))) 

372 except Exception as exc: 

373 conn.send(_Failed(kind="failed", error_type=type(exc).__name__, message=str(exc))) 

374 

375 

376def _silence_output() -> None: 

377 """Point stdout/stderr at devnull; the parent may own a Textual screen.""" 

378 devnull = os.open(os.devnull, os.O_WRONLY) 

379 os.dup2(devnull, sys.stdout.fileno()) 

380 os.dup2(devnull, sys.stderr.fileno()) 

381 os.close(devnull)