Merge pull request #38 from flyingcircusio/PL-135530-save-state-race-condition

Pl 135530 save state race condition

Author: ctheune

src/skvaider/__init__.py

@@ -93,6 +93,8 @@ async def lifespan(
 
     loop.set_exception_handler(global_exception_handler)
 
+    skvaider.auth.start_verify_pool()
+
     backends: list[skvaider.proxy.backends.Backend] = []
     for backend_config in config.backend:
         if backend_config.type == "skvaider":
@@ -135,6 +137,7 @@ async def lifespan(
     if aramaki:
         aramaki.stop()
     pool.close()
+    skvaider.auth.stop_verify_pool()
 
 
 def app_factory(config: Config, lifespan: Any) -> FastAPI:

src/skvaider/auth.py

@@ -1,9 +1,12 @@
+import asyncio
 import base64
 import binascii
+import concurrent.futures
 import json
+import multiprocessing
 import time
 from json import JSONDecodeError
-from typing import Annotated, cast
+from typing import Annotated, Any, cast
 
 import svcs
 from argon2 import PasswordHasher
@@ -16,6 +19,41 @@
 
 hasher = PasswordHasher()
 
+VERIFY_POOL_WORKERS = 10
+
+# argon2 verification is CPU-bound and holds the GIL, so running it on the
+# event loop (or in a thread) stalls every other request on the worker. Offload
+# it to a process pool, which has independent GILs. "spawn" avoids forking a
+# process that already runs the asyncio loop and aiosqlite threads.
+verify_pool: concurrent.futures.ProcessPoolExecutor | None = None
+
+
+def start_verify_pool() -> None:
+    global verify_pool
+    verify_pool = concurrent.futures.ProcessPoolExecutor(
+        max_workers=VERIFY_POOL_WORKERS,
+        mp_context=multiprocessing.get_context("spawn"),
+    )
+
+
+def stop_verify_pool() -> None:
+    global verify_pool
+    if verify_pool is None:
+        return
+    verify_pool.shutdown(wait=False, cancel_futures=True)
+    verify_pool = None
+
+
+def verify_password(secret_hash: str, secret: str) -> bool:
+    # Runs in a pool worker process. Catch everything (not just mismatch) and
+    # report failure, so an unexpected library error can't leak into the caller
+    # as anything other than a rejected token.
+    try:
+        hasher.verify(secret_hash, secret)
+        return True
+    except Exception:
+        return False
+
 
 class AuthTokens(aramaki.Collection):
     collection = "fc.directory.ai.token"
@@ -36,19 +74,20 @@ class Cache:
     TTL = 300
 
     def __init__(self):
-        self.cache: dict[str, float] = {}
+        self.cache: dict[str, tuple[float, Any]] = {}
 
-    def __contains__(self, key: str):
+    def __getitem__(self, key: str):
         now = time.time()
         if key not in self.cache:
-            return False
-        if self.cache[key] < now:
+            raise KeyError(key)
+        expiry, payload = self.cache[key]
+        if expiry < now:
             del self.cache[key]
-            return False
-        return True
+            raise KeyError(key)
+        return payload
 
-    def add(self, key: str):
-        self.cache[key] = time.time() + self.TTL
+    def add(self, key: str, payload: Any):
+        self.cache[key] = (time.time() + self.TTL, payload)
 
 
 cache = Cache()  # XXX turn into service
@@ -59,16 +98,20 @@ async def verify_token(
     credentials: Annotated[HTTPAuthorizationCredentials, Depends(_bearer_auth)],
     services: svcs.fastapi.DepContainer,
 ) -> None:
-    if credentials.credentials in cache:
-        return
+    try:
+        token_id = cache[credentials.credentials]
+    except KeyError:
+        pass
+    else:
+        request.state.token_id = token_id
 
     try:
         admin_tokens = services.get(AdminTokens)
     except svcs.exceptions.ServiceNotFoundError:
         pass
     else:
         if credentials.credentials in admin_tokens.tokens:
-            request.state.token_id = "admin-token"
+            request.state.token_id = "<admin-token>"
             return
 
     # XXX There's a lot of type issues going on here, because the mechanics of passing through
@@ -92,14 +135,20 @@ async def verify_token(
         )
         if not db_token:
             raise HTTPException(401, detail="Bad authentication")
-        try:
-            assert isinstance(db_token["secret_hash"], str)
-            hasher.verify(db_token["secret_hash"], client_token["secret"])
-            request.state.token_id = client_token["id"]
-            cache.add(credentials.credentials)
-        # We could specify explicit exceptions here but go the safe route and just catch all in case the lib addes one
-        except Exception:
+        assert isinstance(db_token["secret_hash"], str)
+        loop = asyncio.get_running_loop()
+        # verify_pool is None outside a running app (e.g. tests); run_in_executor
+        # then falls back to the loop's default thread executor.
+        valid = await loop.run_in_executor(
+            verify_pool,
+            verify_password,
+            db_token["secret_hash"],
+            client_token["secret"],
+        )
+        if not valid:
             raise HTTPException(401, detail="Bad authentication")
+        request.state.token_id = client_token["id"]
+        cache.add(credentials.credentials, client_token["id"])
 
 
 async def verify_admin_token(
@@ -114,4 +163,4 @@ async def verify_admin_token(
         raise HTTPException(401, detail="No admin tokens configured")
     if credentials.credentials not in admin_tokens.tokens:
         raise HTTPException(401, detail="Bad authentication")
-    request.state.token_id = "admin-token"
+    request.state.token_id = "<admin-token>"

src/skvaider/debug.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Any
 
+import aiofiles
 from fastapi import Request
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
@@ -142,7 +143,6 @@ def data(self) -> bytes:
 
 
 class DebugRecorder:
-    # XXX async file io!
     temp_file: Path | None
     triggers: list[str]
     time_start: float
@@ -211,11 +211,13 @@ def trigger_flag(self):
 
     async def write_request(self, stem: str, data: dict[str, Any]) -> None:
         path = self.directory / f"{stem}.request"
-        path.write_text(_format_request(data))
+        async with aiofiles.open(path, mode="w") as f:
+            await f.write(_format_request(data))
 
     async def write_response(self, stem: str, data: dict[str, Any]) -> None:
         path = self.directory / f"{stem}.response"
-        path.write_text(_format_response(data))
+        async with aiofiles.open(path, mode="w") as f:
+            await f.write(_format_response(data))
 
     async def record(self) -> None:
         if 400 <= self.status_code < 500:

src/skvaider/inference/manager.py

@@ -133,13 +133,20 @@ def manifest(self, value: set[str]) -> None:
         self.manifest_changed.set()
 
     def update_manifest(self, model_ids: set[str], serial: Serial) -> None:
-        if serial <= self.manifest_serial:
+        if serial < self.manifest_serial:
             log.info(
                 "ignoring manifest with stale serial",
                 serial=serial,
                 current=self.manifest_serial,
             )
             return
+        if serial == self.manifest_serial:
+            log.info(
+                "ignoring manifest with current serial",
+                serial=serial,
+                current=self.manifest_serial,
+            )
+            return
         self.manifest_serial = serial
         self.manifest = model_ids
 

src/skvaider/proxy/pool.py

@@ -161,6 +161,7 @@ def __init__(
         self.model_configs = {m.id: m for m in model_configs}
 
         self.model_management_lock = asyncio.Lock()
+        self.save_state_lock = asyncio.Lock()
 
         for model_id in self.model_configs:
             self.semaphores[model_id] = ModelSemaphore(model_id, self)
@@ -206,33 +207,36 @@ async def save_state(self) -> None:
         if not self.state_file:
             return
 
-        records = {
-            backend.url: BackendStateRecord(
-                url=backend.url,
-                healthy=backend.healthy,
-                memory=backend.memory,
-                map_up=backend.map_up.state,
-                map_up_last_change=backend.map_up.last_change,
-                map_in=backend.map_in.state,
-                map_in_last_change=backend.map_in.last_change,
-                models={
-                    model_id: ModelStateRecord(
-                        id=model_id,
-                        memory_usage=model.memory_usage,
-                    )
-                    for model_id, model in backend.models.items()
-                },
-            )
-            for backend in self.backends
-        }
-        state = ClusterState(backends=records)
-        data = state.model_dump_json(indent=2)
-
-        tmp = self.state_file.with_suffix(".tmp")
-        async with aiofiles.open(tmp, mode="w") as f:
-            await f.write(data)
-            await f.flush()
-        tmp.rename(self.state_file)
+        # Serialize concurrent saves: multiple backend health monitors call
+        # this, and a shared temp file would race on the rename otherwise.
+        async with self.save_state_lock:
+            records = {
+                backend.url: BackendStateRecord(
+                    url=backend.url,
+                    healthy=backend.healthy,
+                    memory=backend.memory,
+                    map_up=backend.map_up.state,
+                    map_up_last_change=backend.map_up.last_change,
+                    map_in=backend.map_in.state,
+                    map_in_last_change=backend.map_in.last_change,
+                    models={
+                        model_id: ModelStateRecord(
+                            id=model_id,
+                            memory_usage=model.memory_usage,
+                        )
+                        for model_id, model in backend.models.items()
+                    },
+                )
+                for backend in self.backends
+            }
+            state = ClusterState(backends=records)
+            data = state.model_dump_json(indent=2)
+
+            tmp = self.state_file.with_suffix(".tmp")
+            async with aiofiles.open(tmp, mode="w") as f:
+                await f.write(data)
+                await f.flush()
+            tmp.rename(self.state_file)
 
     def placement_map(self) -> ModelMap:
         return placement.placement_map(

src/skvaider/proxy/tests/test_pool.py

@@ -1,17 +1,20 @@
 import asyncio
+import contextlib
 import datetime
 from pathlib import Path
-from typing import Callable
+from typing import Any, AsyncGenerator, Callable
 from unittest.mock import patch
 
+import aiofiles
+
 from skvaider.config import ModelInstanceConfig, parse_size
 from skvaider.conftest import registered_model_factory
 from skvaider.manifest import Serial
 from skvaider.proxy.backends import DummyBackend
 from skvaider.proxy.models import AIModel
 from skvaider.utils import TaskManager
 
-from ..pool import Pool
+from ..pool import ClusterState, Pool
 
 
 async def test_maps_only_includes_desired_models(
@@ -520,3 +523,57 @@ async def test_save_state_writes_state_file(
     assert pool.state_file.exists()
     content = pool.state_file.read_text()
     assert "m1" in content
+
+
+async def test_save_state_serializes_concurrent_calls(
+    dummy_backend: DummyBackend,
+    tmp_path: Path,
+):
+    """Concurrent save_state() calls must not race on the shared temp file.
+
+    Multiple backend health monitors call save_state() concurrently. Without
+    serialization they share one temp file and the second rename() fails with
+    FileNotFoundError once the first has moved the temp file into place.
+    """
+    dummy_backend.healthy = True
+    dummy_backend.memory = {"ram": {"free": 924, "total": 1024}}
+    registered_model_factory("m1", dummy_backend, ram=100)
+
+    pool = Pool(
+        [
+            ModelInstanceConfig(
+                id="m1", instances=1, memory={"ram": 100}, task="chat"
+            )
+        ],
+        [dummy_backend],
+        data_dir=tmp_path,
+    )
+
+    active = 0
+    max_active = 0
+    real_open: Callable[..., Any] = aiofiles.open
+
+    @contextlib.asynccontextmanager
+    async def tracking_open(
+        *args: Any, **kwargs: Any
+    ) -> AsyncGenerator[Any, None]:
+        nonlocal active, max_active
+        async with real_open(*args, **kwargs) as f:
+            active += 1
+            max_active = max(max_active, active)
+            # Force a scheduling point inside the critical section so any
+            # unserialized concurrency becomes observable.
+            await asyncio.sleep(0)
+            try:
+                yield f
+            finally:
+                active -= 1
+
+    with patch("skvaider.proxy.pool.aiofiles.open", tracking_open):
+        await asyncio.gather(*(pool.save_state() for _ in range(10)))
+
+    assert max_active == 1
+
+    assert pool.state_file is not None
+    # The final file is complete, valid JSON (no torn write left behind).
+    ClusterState.model_validate_json(pool.state_file.read_text())