feat(k8s)

nc_py_api/ex_app/integration_fastapi.py

@@ -20,6 +20,7 @@
     status,
 )
 from fastapi.responses import JSONResponse, PlainTextResponse
+from filelock import SoftFileLock, Timeout
 from starlette.requests import HTTPConnection, Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
@@ -29,6 +30,9 @@
 from ..talk_bot import TalkBotMessage
 from .misc import persistent_storage
 
+# Lock timeout in seconds for model downloads (can be overridden in tests)
+LOCK_TIMEOUT = 3600
+
 
 def nc_app(request: HTTPConnection) -> NextcloudApp:
     """Authentication handler for requests from Nextcloud to the application."""
@@ -205,42 +209,50 @@ def __fetch_model_as_file(
     current_progress: int, progress_for_task: int, nc: NextcloudApp, model_path: str, download_options: dict
 ) -> str:
     result_path = download_options.pop("save_path", urlparse(model_path).path.split("/")[-1])
-    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(result_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
+    temp_path = result_path + ".tmp"
+    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
 
@@ -265,9 +277,18 @@ def display(self, msg=None, pos=None):
 
     workers = download_options.pop("max_workers", 2)
     cache = download_options.pop("cache_dir", persistent_storage())
-    return snapshot_download(
-        model_name, tqdm_class=TqdmProgress, **download_options, max_workers=workers, cache_dir=cache
-    )
+    safe_model_name = model_name
+    for sep in (os.sep, os.altsep):
+        if sep:
+            safe_model_name = safe_model_name.replace(sep, "_")
+    lock_path = os.path.join(cache, f"{safe_model_name}.lock")
+    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:

pyproject.toml

@@ -46,6 +46,7 @@ dynamic = [
 ]
 dependencies = [
   "fastapi>=0.109.2",
+  "filelock>=3.24.2,<4",
   "niquests>=3,<4",
   "pydantic>=2.1.1",
   "python-dotenv>=1",