fix: Throw ModelFetch errors on lock timeout

Signed-off-by: Marcel Klehr <mklehr@gmx.net>

Author: marcelklehr

nc_py_api/ex_app/integration_fastapi.py

@@ -20,7 +20,7 @@
     status,
 )
 from fastapi.responses import JSONResponse, PlainTextResponse
-from filelock import SoftFileLock
+from filelock import SoftFileLock, Timeout
 from starlette.requests import HTTPConnection, Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
@@ -210,44 +210,47 @@ def __fetch_model_as_file(
 ) -> str:
     result_path = download_options.pop("save_path", urlparse(model_path).path.split("/")[-1])
     temp_path = result_path + ".tmp"
-    with SoftFileLock(result_path + ".lock", timeout=LOCK_TIMEOUT):
-        with niquests.get(model_path, stream=True) as response:
-            if not response.ok:
-                raise ModelFetchError(
-                    f"Downloading of '{model_path}' failed, returned ({response.status_code}) {response.text}"
-                )
-            downloaded_size = 0
-            linked_etag = ""
-            for each_history in response.history:
-                linked_etag = each_history.headers.get("X-Linked-ETag", "")
-                if linked_etag:
-                    break
-            if not linked_etag:
-                linked_etag = response.headers.get("X-Linked-ETag", response.headers.get("ETag", ""))
-            total_size = int(response.headers.get("Content-Length"))
-            try:
-                existing_size = os.path.getsize(result_path)
-            except OSError:
-                existing_size = 0
-            if linked_etag and total_size == existing_size:
-                with builtins.open(result_path, "rb") as file:
-                    sha256_hash = hashlib.sha256()
-                    for byte_block in iter(lambda: file.read(4096), b""):
-                        sha256_hash.update(byte_block)
-                    if f'"{sha256_hash.hexdigest()}"' == linked_etag:
-                        nc.set_init_status(min(current_progress + progress_for_task, 99))
-                        return result_path
-
-            with builtins.open(temp_path, "wb") as file:
-                last_progress = current_progress
-                for chunk in response.iter_raw(-1):
-                    downloaded_size += file.write(chunk)
-                    if total_size:
-                        new_progress = min(current_progress + int(progress_for_task * downloaded_size / total_size), 99)
-                        if new_progress != last_progress:
-                            nc.set_init_status(new_progress)
-                            last_progress = new_progress
-            os.replace(temp_path, result_path)
+    try:
+        with SoftFileLock(result_path + ".lock", timeout=LOCK_TIMEOUT):
+            with niquests.get(model_path, stream=True) as response:
+                if not response.ok:
+                    raise ModelFetchError(
+                        f"Downloading of '{model_path}' failed, returned ({response.status_code}) {response.text}"
+                    )
+                downloaded_size = 0
+                linked_etag = ""
+                for each_history in response.history:
+                    linked_etag = each_history.headers.get("X-Linked-ETag", "")
+                    if linked_etag:
+                        break
+                if not linked_etag:
+                    linked_etag = response.headers.get("X-Linked-ETag", response.headers.get("ETag", ""))
+                total_size = int(response.headers.get("Content-Length"))
+                try:
+                    existing_size = os.path.getsize(result_path)
+                except OSError:
+                    existing_size = 0
+                if linked_etag and total_size == existing_size:
+                    with builtins.open(result_path, "rb") as file:
+                        sha256_hash = hashlib.sha256()
+                        for byte_block in iter(lambda: file.read(4096), b""):
+                            sha256_hash.update(byte_block)
+                        if f'"{sha256_hash.hexdigest()}"' == linked_etag:
+                            nc.set_init_status(min(current_progress + progress_for_task, 99))
+                            return result_path
+
+                with builtins.open(temp_path, "wb") as file:
+                    last_progress = current_progress
+                    for chunk in response.iter_raw(-1):
+                        downloaded_size += file.write(chunk)
+                        if total_size:
+                            new_progress = min(current_progress + int(progress_for_task * downloaded_size / total_size), 99)
+                            if new_progress != last_progress:
+                                nc.set_init_status(new_progress)
+                                last_progress = new_progress
+                os.replace(temp_path, result_path)
+    except Timeout as e:
+        raise ModelFetchError(f"Timeout acquiring lock for '{model_path}'") from e
 
     return result_path
 
@@ -277,10 +280,13 @@ def display(self, msg=None, pos=None):
         if sep:
             safe_model_name = safe_model_name.replace(sep, "_")
     lock_path = os.path.join(cache, f"{safe_model_name}.lock")
-    with SoftFileLock(lock_path, timeout=LOCK_TIMEOUT):
-        return snapshot_download(
-            model_name, tqdm_class=TqdmProgress, **download_options, max_workers=workers, cache_dir=cache
-        )
+    try:
+        with SoftFileLock(lock_path, timeout=LOCK_TIMEOUT):
+            return snapshot_download(
+                model_name, tqdm_class=TqdmProgress, **download_options, max_workers=workers, cache_dir=cache
+            )
+    except Timeout as e:
+        raise ModelFetchError(f"Timeout acquiring lock for model '{model_name}'") from e
 
 
 def __nc_app(request: HTTPConnection) -> dict: