|
20 | 20 | status, |
21 | 21 | ) |
22 | 22 | from fastapi.responses import JSONResponse, PlainTextResponse |
| 23 | +from filelock import FileLock |
| 24 | +from filelock import Timeout as FileLockTimeout |
23 | 25 | from starlette.requests import HTTPConnection, Request |
24 | 26 | from starlette.types import ASGIApp, Receive, Scope, Send |
25 | 27 |
|
@@ -158,7 +160,7 @@ def fetch_models_task(nc: NextcloudApp, models: dict[str, dict], progress_init_s |
158 | 160 | """Use for cases when you want to define custom `/init` but still need to easy download models. |
159 | 161 |
|
160 | 162 | :param nc: NextcloudApp instance. |
161 | | - :param models_to_fetch: Dictionary describing which models should be downloaded of the form: |
| 163 | + :param models: Dictionary describing which models should be downloaded of the form: |
162 | 164 | .. code-block:: python |
163 | 165 | { |
164 | 166 | "model_url_1": { |
@@ -205,42 +207,56 @@ def __fetch_model_as_file( |
205 | 207 | current_progress: int, progress_for_task: int, nc: NextcloudApp, model_path: str, download_options: dict |
206 | 208 | ) -> str: |
207 | 209 | result_path = download_options.pop("save_path", urlparse(model_path).path.split("/")[-1]) |
208 | | - with niquests.get(model_path, stream=True) as response: |
209 | | - if not response.ok: |
210 | | - raise ModelFetchError( |
211 | | - f"Downloading of '{model_path}' failed, returned ({response.status_code}) {response.text}" |
212 | | - ) |
213 | | - downloaded_size = 0 |
214 | | - linked_etag = "" |
215 | | - for each_history in response.history: |
216 | | - linked_etag = each_history.headers.get("X-Linked-ETag", "") |
217 | | - if linked_etag: |
218 | | - break |
219 | | - if not linked_etag: |
220 | | - linked_etag = response.headers.get("X-Linked-ETag", response.headers.get("ETag", "")) |
221 | | - total_size = int(response.headers.get("Content-Length")) |
222 | | - try: |
223 | | - existing_size = os.path.getsize(result_path) |
224 | | - except OSError: |
225 | | - existing_size = 0 |
226 | | - if linked_etag and total_size == existing_size: |
227 | | - with builtins.open(result_path, "rb") as file: |
228 | | - sha256_hash = hashlib.sha256() |
229 | | - for byte_block in iter(lambda: file.read(4096), b""): |
230 | | - sha256_hash.update(byte_block) |
231 | | - if f'"{sha256_hash.hexdigest()}"' == linked_etag: |
232 | | - nc.set_init_status(min(current_progress + progress_for_task, 99)) |
233 | | - return result_path |
234 | | - |
235 | | - with builtins.open(result_path, "wb") as file: |
236 | | - last_progress = current_progress |
237 | | - for chunk in response.iter_raw(-1): |
238 | | - downloaded_size += file.write(chunk) |
239 | | - if total_size: |
240 | | - new_progress = min(current_progress + int(progress_for_task * downloaded_size / total_size), 99) |
241 | | - if new_progress != last_progress: |
242 | | - nc.set_init_status(new_progress) |
243 | | - last_progress = new_progress |
| 210 | + tmp_path = result_path + ".tmp" |
| 211 | + try: |
| 212 | + with FileLock(result_path + ".lock", timeout=7200), niquests.get(model_path, stream=True) as response: |
| 213 | + if not response.ok: |
| 214 | + raise ModelFetchError( |
| 215 | + f"Downloading of '{model_path}' failed, returned ({response.status_code}) {response.text}" |
| 216 | + ) |
| 217 | + downloaded_size = 0 |
| 218 | + linked_etag = "" |
| 219 | + for redirect_resp in response.history: |
| 220 | + linked_etag = redirect_resp.headers.get("X-Linked-ETag", "") |
| 221 | + if linked_etag: |
| 222 | + break |
| 223 | + if not linked_etag: |
| 224 | + linked_etag = response.headers.get("X-Linked-ETag", response.headers.get("ETag", "")) |
| 225 | + total_size = int(response.headers.get("Content-Length", 0)) |
| 226 | + try: |
| 227 | + existing_size = os.path.getsize(result_path) |
| 228 | + except OSError: |
| 229 | + existing_size = 0 |
| 230 | + if linked_etag and total_size and total_size == existing_size: |
| 231 | + with builtins.open(result_path, "rb") as file: |
| 232 | + sha256_hash = hashlib.sha256() |
| 233 | + for byte_block in iter(lambda: file.read(4096), b""): |
| 234 | + sha256_hash.update(byte_block) |
| 235 | + if f'"{sha256_hash.hexdigest()}"' == linked_etag: |
| 236 | + nc.set_init_status(min(current_progress + progress_for_task, 99)) |
| 237 | + return result_path |
| 238 | + |
| 239 | + try: |
| 240 | + with builtins.open(tmp_path, "wb") as file: |
| 241 | + last_progress = current_progress |
| 242 | + for chunk in response.iter_raw(-1): |
| 243 | + downloaded_size += file.write(chunk) |
| 244 | + if total_size: |
| 245 | + new_progress = min( |
| 246 | + current_progress + int(progress_for_task * downloaded_size / total_size), 99 |
| 247 | + ) |
| 248 | + if new_progress != last_progress: |
| 249 | + nc.set_init_status(new_progress) |
| 250 | + last_progress = new_progress |
| 251 | + os.replace(tmp_path, result_path) |
| 252 | + except BaseException: |
| 253 | + if os.path.exists(tmp_path): |
| 254 | + os.remove(tmp_path) |
| 255 | + raise |
| 256 | + except FileLockTimeout as exc: |
| 257 | + raise ModelFetchError( |
| 258 | + f"Timed out waiting for lock on '{result_path}' after 7200s — another process may be stuck downloading" |
| 259 | + ) from exc |
244 | 260 |
|
245 | 261 | return result_path |
246 | 262 |
|
|
0 commit comments