lora-train: cross-engine GPU lock so nothing reloads mid-training (fix intermittent OOM)

evict_cosited_siblings() frees the co-located sibling's VRAM ONCE at training
start, but training runs for minutes as an in-engine background job that the front
swap-gate can't cover (its POST returns a job_id immediately). So a concurrent LLM
request reloaded the gguf text model mid-training and OOM'd the trainer (one
fighter LoRA failed while others succeeded).

Add codai/models/gpu_lock.py: a cross-engine GPU reservation. Training reserves the
card for its whole duration — locally AND on co-located siblings via new
/internal/gpu-reserve + /internal/gpu-release endpoints — and every ordinary
model-load path (manager.request_model, video _load_video_pipeline, image
_load_diffusers_pipeline) calls wait_until_free() first. A request that needs a
load during training now BLOCKS until training releases, then loads and serves —
the same queue-behind-the-owner behaviour the swap-gate gives request-level work,
now extended to cover training. The training thread is exempt from its own
reservation, so loading the base model never self-deadlocks; waits are bounded
(900s) so a stuck reservation can't hang forever.

Validated: sibling reservation blocks a loader thread until release; the reserving
thread never waits on itself.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce

codai/api/app.py

@@ -393,6 +393,41 @@ async def internal_evict_vram(request: Request):
         return {"ok": False, "error": str(e), "freed_gb": 0.0}
 
 
+@app.post("/internal/gpu-reserve", include_in_schema=False)
+async def internal_gpu_reserve(request: Request):
+    """A co-located sibling engine (same GPU) reserves the card for an exclusive op
+    (LoRA training). While reserved, this engine's model-load paths wait, so it
+    doesn't load a model and contend with the sibling's training VRAM."""
+    import json as _json
+    try:
+        data = _json.loads((await request.body()) or b"{}") or {}
+    except Exception:
+        data = {}
+    try:
+        from codai.models import gpu_lock
+        gpu_lock.set_sibling(str(data.get("from") or "sibling"),
+                             data.get("reason") or "sibling-reserved")
+        return {"ok": True}
+    except Exception as e:
+        return {"ok": False, "error": str(e)}
+
+
+@app.post("/internal/gpu-release", include_in_schema=False)
+async def internal_gpu_release(request: Request):
+    """Lift a sibling engine's GPU reservation (its exclusive op finished)."""
+    import json as _json
+    try:
+        data = _json.loads((await request.body()) or b"{}") or {}
+    except Exception:
+        data = {}
+    try:
+        from codai.models import gpu_lock
+        gpu_lock.set_sibling(str(data.get("from") or "sibling"), None)
+        return {"ok": True}
+    except Exception as e:
+        return {"ok": False, "error": str(e)}
+
+
 @app.post("/internal/thermal-pause", include_in_schema=False)
 async def internal_thermal_pause(request: Request):
     """Front-driven thermal pause: the supervisor (which monitors temps centrally)

codai/api/images.py

@@ -361,6 +361,14 @@ def _load_diffusers_pipeline(model_name: str, global_args, model_config: dict =
     from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DiffusionPipeline
     import torch
 
+    # Wait if the GPU is reserved for an exclusive op (LoRA training here or on a
+    # co-located sibling) so this load doesn't contend with it.
+    try:
+        from codai.models import gpu_lock
+        gpu_lock.wait_until_free(context=f"image load '{model_name}'")
+    except Exception:
+        pass
+
     _mc = model_config or {}
 
     def _cfg(key, default=None):

codai/api/loras.py

@@ -2101,6 +2101,14 @@ def _train_lora_blocking(req: LoraTrainRequest, job_id: Optional[str] = None) ->
             _active_job_id = None
             _train_lock.release()
             raise
+    # Reserve the GPU for the whole training: block ordinary model loads on this
+    # engine (image/video) AND on co-located siblings (gguf text), so nothing
+    # reloads mid-training and OOMs the trainer. Released in the finally below.
+    try:
+        from codai.models import gpu_lock
+        gpu_lock.reserve(f"lora-training:{getattr(req, 'name', '') or '?'}")
+    except Exception as _gle:
+        print(f"  [lora] could not reserve GPU lock: {_gle}")
     try:
         result = _train_lora_sync(req)
         if job_id:
@@ -2139,6 +2147,11 @@ def _train_lora_blocking(req: LoraTrainRequest, job_id: Optional[str] = None) ->
         _drop_wan_cache()
         raise
     finally:
+        try:
+            from codai.models import gpu_lock
+            gpu_lock.release()
+        except Exception:
+            pass
         _active_job_id = None
         _active_train_model = None
         if job_id:

codai/api/video.py

@@ -973,6 +973,13 @@ def _apply_attention_backend(pipe, model_cfg: dict) -> None:
 
 
 def _load_video_pipeline(model_name: str, device: str, mode: str, offload: str = None, model_cfg: dict = None, use_incremental_cache: bool = True):
+    # Wait if the GPU is reserved for an exclusive op (LoRA training here or on a
+    # co-located sibling) so this load doesn't contend with it.
+    try:
+        from codai.models import gpu_lock
+        gpu_lock.wait_until_free(context=f"video load '{model_name}'")
+    except Exception:
+        pass
     # GGUF models go through stable-diffusion.cpp, not diffusers
     from codai.api.images import _is_gguf_model
     if _is_gguf_model(model_name):

codai/models/gpu_lock.py

+"""Cross-engine GPU exclusivity for long in-engine background ops (LoRA training).
+
+The front swap-gate serializes REQUEST-level GPU use, but a LoRA training runs as
+an in-engine BACKGROUND job (its POST returns a job_id immediately), so no front
+request is in flight to hold the gate for the training's whole duration. This lock
+closes that gap: training reserves the GPU here AND on co-located sibling engines
+(POST /internal/gpu-reserve), and every ordinary model-load path calls
+``wait_until_free()`` first — so neither the same engine (image/video) nor a
+sibling (gguf text) loads a model and contends with training's VRAM (the OOM this
+prevents). A waiting request isn't dropped: it blocks until training releases, then
+loads and serves — the same queue-behind-the-owner behaviour the swap-gate gives
+request-level work.
+
+The training thread is exempt from its OWN reservation, so training loading its
+base model never deadlocks on the lock it holds.
+"""
+import os
+import threading
+import time
+
+_cv = threading.Condition()
+_local_reason = None            # this engine holds it (e.g. "lora-training:<name>")
+_local_thread = None            # tid that reserved locally — exempt from waiting
+_sibling_reasons = {}           # sibling key -> reason (set via internal endpoint)
+
+
+def reserved() -> bool:
+    return _local_reason is not None or bool(_sibling_reasons)
+
+
+def current_reason() -> str:
+    with _cv:
+        if _local_reason:
+            return _local_reason
+        for r in _sibling_reasons.values():
+            return r
+        return ""
+
+
+def _set_local(reason) -> None:
+    global _local_reason, _local_thread
+    with _cv:
+        _local_reason = reason
+        _local_thread = threading.get_ident() if reason else None
+        _cv.notify_all()
+
+
+def set_sibling(key: str, reason) -> None:
+    """Set (reason truthy) or clear (reason falsy) a co-located sibling's reservation."""
+    with _cv:
+        if reason:
+            _sibling_reasons[str(key)] = reason
+        else:
+            _sibling_reasons.pop(str(key), None)
+        _cv.notify_all()
+
+
+def wait_until_free(timeout: float = 900.0, context: str = "") -> bool:
+    """Block until the GPU isn't reserved (by this engine or a sibling). Returns
+    True if free, False on timeout (caller proceeds anyway — bounded so a stuck
+    reservation can't hang forever). The thread that holds the local reservation
+    (the training thread) never waits on itself."""
+    if _local_thread is not None and threading.get_ident() == _local_thread:
+        return True
+    if not reserved():
+        return True
+    _r = current_reason()
+    print(f"  [gpu-lock] {context or 'model load'} waiting — GPU reserved ({_r})",
+          flush=True)
+    deadline = time.time() + max(0.0, timeout)
+    with _cv:
+        while reserved():
+            if _local_thread is not None and threading.get_ident() == _local_thread:
+                return True
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                print(f"  [gpu-lock] wait timed out after {timeout:.0f}s "
+                      f"({_r}) — proceeding", flush=True)
+                return False
+            _cv.wait(timeout=min(remaining, 5.0))
+    print(f"  [gpu-lock] GPU free — {context or 'model load'} proceeding", flush=True)
+    return True
+
+
+def _cosited_urls():
+    return [u for u in (os.environ.get("CODERAI_COSITED_URLS") or "").split(",") if u]
+
+
+def _notify_siblings(reason) -> None:
+    urls = _cosited_urls()
+    if not urls:
+        return
+    try:
+        import httpx
+    except Exception:
+        return
+    tok = os.environ.get("CODERAI_INTERNAL_TOKEN") or ""
+    me = os.environ.get("CODERAI_ENGINE_NAME") or str(os.getpid())
+    path = "/internal/gpu-reserve" if reason else "/internal/gpu-release"
+    for u in urls:
+        try:
+            httpx.post(f"{u}{path}", json={"from": me, "reason": reason or ""},
+                       headers={"x-coderai-internal": tok}, timeout=15.0)
+        except Exception as e:
+            print(f"  [gpu-lock] sibling {u}{path} failed: {e}", flush=True)
+
+
+def reserve(reason: str) -> None:
+    """Reserve the GPU for an exclusive local op (training): block this engine's
+    ordinary loads AND tell co-located siblings to hold theirs."""
+    _set_local(reason or "reserved")
+    _notify_siblings(reason or "reserved")
+
+
+def release() -> None:
+    """Release the local reservation and clear it on siblings."""
+    _set_local(None)
+    _notify_siblings(None)

codai/models/manager.py

@@ -3899,6 +3899,13 @@ class MultiModelManager:
                 - 'config': The stored configuration for this model
                 - 'already_loaded': True if the model is already loaded and ready in VRAM
         """
+        # Wait if the GPU is reserved for an exclusive op (LoRA training, here or on
+        # a co-located sibling): don't load a model that would contend with it.
+        try:
+            from codai.models import gpu_lock
+            gpu_lock.wait_until_free(context=f"request '{requested_model}'")
+        except Exception:
+            pass
         # Fail fast if the CUDA context is already corrupted — loading anything
         # onto a poisoned context just re-triggers the same async assert.
         if self.cuda_context_poisoned: