fix: coordinate concurrent account capacity filling

modelhub_submmit_api/common.py

@@ -1,9 +1,15 @@
 from __future__ import annotations
 
+import fcntl
+import hashlib
 import json
+import os
+import socket
+import tempfile
+from contextlib import contextmanager
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable, Iterator
 
 
 def ensure_utc(value: datetime) -> datetime:
@@ -54,16 +60,57 @@ def json_dumps(value: Any) -> str:
     return json.dumps(value, ensure_ascii=False, sort_keys=True)
 
 
-def read_json(path: Path) -> Any:
-    return json.loads(path.read_text(encoding="utf-8"))
+def runtime_instance_id() -> str:
+    """Return a non-secret identity that differs between concurrent workers."""
+    base = (
+        os.getenv("MODELHUB_AGENT_INSTANCE_ID")
+        or os.getenv("STRATEGY_ID")
+        or os.getenv("HOSTNAME")
+        or socket.gethostname()
+    )
+    raw = f"{base}:{os.getpid()}"
+    return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:20]
 
 
-def write_json(path: Path, value: Any) -> None:
+def _lock_path(path: Path) -> Path:
+    return path.with_name(f".{path.name}.lock")
+
+
+@contextmanager
+def file_lock(path: Path, *, shared: bool = False) -> Iterator[None]:
+    """Coordinate access to a path between processes on the same filesystem."""
     path.parent.mkdir(parents=True, exist_ok=True)
-    path.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+    lock_path = _lock_path(path)
+    with lock_path.open("a+b") as handle:
+        fcntl.flock(handle.fileno(), fcntl.LOCK_SH if shared else fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
 
 
-def read_jsonl(path: Path) -> list[dict[str, Any]]:
+def _atomic_write_text(path: Path, content: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    descriptor, temporary_name = tempfile.mkstemp(
+        dir=str(path.parent),
+        prefix=f".{path.name}.",
+        suffix=".tmp",
+    )
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+            handle.write(content)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary_name, path)
+    except Exception:
+        try:
+            os.unlink(temporary_name)
+        except FileNotFoundError:
+            pass
+        raise
+
+
+def _read_jsonl_unlocked(path: Path) -> list[dict[str, Any]]:
     if not path.exists():
         return []
     rows: list[dict[str, Any]] = []
@@ -74,13 +121,45 @@ def read_jsonl(path: Path) -> list[dict[str, Any]]:
     return rows
 
 
+def _jsonl_content(rows: list[dict[str, Any]]) -> str:
+    return "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows)
+
+
+def read_json(path: Path) -> Any:
+    with file_lock(path, shared=True):
+        return json.loads(path.read_text(encoding="utf-8"))
+
+
+def write_json(path: Path, value: Any) -> None:
+    content = json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
+    with file_lock(path):
+        _atomic_write_text(path, content)
+
+
+def read_jsonl(path: Path) -> list[dict[str, Any]]:
+    with file_lock(path, shared=True):
+        return _read_jsonl_unlocked(path)
+
+
 def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    content = "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows)
-    path.write_text(content, encoding="utf-8")
+    with file_lock(path):
+        _atomic_write_text(path, _jsonl_content(rows))
 
 
 def append_jsonl(path: Path, row: dict[str, Any]) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("a", encoding="utf-8") as handle:
-        handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
+    line = json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n"
+    with file_lock(path):
+        with path.open("a", encoding="utf-8") as handle:
+            handle.write(line)
+            handle.flush()
+
+
+def update_jsonl(
+    path: Path,
+    updater: Callable[[list[dict[str, Any]]], list[dict[str, Any]]],
+) -> list[dict[str, Any]]:
+    """Atomically read, update, and replace a JSONL file under one lock."""
+    with file_lock(path):
+        updated = updater(_read_jsonl_unlocked(path))
+        _atomic_write_text(path, _jsonl_content(updated))
+        return updated