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. 

5""" 

6 

7from __future__ import annotations 

8 

9import multiprocessing 

10import os 

11import sys 

12import time 

13from dataclasses import dataclass 

14from pathlib import Path 

15from typing import TYPE_CHECKING, Literal, Protocol 

16 

17from lilbee.catalog.models import CatalogModel 

18from lilbee.runtime.cancellation import CancelSignal, TaskCancelledError 

19 

20if TYPE_CHECKING: 

21 from multiprocessing.connection import Connection 

22 

23 from lilbee.catalog.download_progress import ProgressCallback 

24 

25_POLL_INTERVAL_S = 0.2 

26 

27_EXIT_GRACE_S = 10.0 

28 

29_PROGRESS_MIN_INTERVAL_S = 0.1 

30 

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

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

33 

34 

35@dataclass(frozen=True) 

36class _Progress: 

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

38 

39 kind: Literal["progress"] 

40 downloaded: int 

41 total: int 

42 

43 

44@dataclass(frozen=True) 

45class _Done: 

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

47 

48 kind: Literal["done"] 

49 path: str 

50 

51 

52@dataclass(frozen=True) 

53class _Failed: 

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

55 

56 kind: Literal["failed"] 

57 error_type: str 

58 message: str 

59 

60 

61_ChildMessage = _Progress | _Done | _Failed 

62 

63 

64class _Worker(Protocol): 

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

66 

67 @property 

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

69 

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

71 

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

73 

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

75 

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

77 

78 

79class _PipeProgress: 

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

81 

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

83 self._conn = conn 

84 self._last_sent: float | None = None 

85 

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

87 now = time.monotonic() 

88 final = total > 0 and downloaded >= total 

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

90 if not final and throttled: 

91 return 

92 self._last_sent = now 

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

94 

95 

96def download_in_subprocess( 

97 entry: CatalogModel, 

98 models_dir: Path, 

99 token: str | None, 

100 *, 

101 on_progress: ProgressCallback | None, 

102 cancel: CancelSignal, 

103) -> Path: 

104 """Run one download in its own process, relaying progress until it finishes. 

105 

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

107 the bandwidth of a running hf_xet transfer mid-flight. 

108 """ 

109 if cancel.is_set(): 

110 raise TaskCancelledError 

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

112 try: 

113 return _relay_until_done(entry, worker, receiver, on_progress, cancel) 

114 finally: 

115 _stop_worker(worker) 

116 receiver.close() 

117 

118 

119def _start_worker( 

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

121) -> tuple[_Worker, Connection]: 

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

123 

124 Daemonic on purpose: at interpreter exit multiprocessing terminates daemon 

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

126 must not wait for a multi-GB transfer. 

127 """ 

128 context = multiprocessing.get_context("spawn") 

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

130 worker = context.Process( 

131 target=_run_download_child, 

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

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

134 daemon=True, 

135 ) 

136 worker.start() 

137 sender.close() 

138 return worker, receiver 

139 

140 

141def _relay_until_done( 

142 entry: CatalogModel, 

143 worker: _Worker, 

144 receiver: Connection, 

145 on_progress: ProgressCallback | None, 

146 cancel: CancelSignal, 

147) -> Path: 

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

149 while True: 

150 if cancel.is_set(): 

151 raise TaskCancelledError 

152 if receiver.poll(_POLL_INTERVAL_S): 

153 try: 

154 message = receiver.recv() 

155 except (EOFError, OSError): 

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

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

158 # BrokenPipeError. 

159 raise _died_silently(entry, worker) from None 

160 verdict = _apply(message, on_progress) 

161 if verdict is not None: 

162 return verdict 

163 elif not worker.is_alive() and not receiver.poll(): 

164 raise _died_silently(entry, worker) 

165 

166 

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

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

169 return RuntimeError( 

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

171 ) 

172 

173 

174def _apply(message: _ChildMessage, on_progress: ProgressCallback | None) -> Path | None: 

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

176 if message.kind == "progress": 

177 if on_progress is not None: 

178 on_progress(message.downloaded, message.total) 

179 return None 

180 if message.kind == "done": 

181 return Path(message.path) 

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

183 

184 

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

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

187 if worker.is_alive(): 

188 worker.terminate() 

189 worker.join(_EXIT_GRACE_S) 

190 if worker.is_alive(): 

191 worker.kill() 

192 worker.join(_EXIT_GRACE_S) 

193 

194 

195def _run_download_child( 

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

197) -> None: 

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

199 _silence_output() 

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

201 from lilbee.catalog.download import fetch_model_files 

202 

203 try: 

204 path = fetch_model_files(entry, Path(models_dir), token, on_progress=_PipeProgress(conn)) 

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

206 except Exception as exc: 

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

208 

209 

210def _silence_output() -> None: 

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

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

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

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

215 os.close(devnull)