video: fix VRAM teardown leak + track in-flight loads with reservations

Teardown leak: _free_pipeline_vram now records the pipeline's CUDA
storage pointers up front and, after the normal teardown+gc, walks
gc.get_objects() and nulls any surviving tensors out of their list/dict
referrers (the proven pass from the CUDA text backend's cleanup) — the
"~22 GB untracked (teardown leak; referenced elsewhere)" VRAM is now
reclaimed instead of poisoning every later load until a restart. Names
the holders when it fires so future leak sources identify themselves.

Tracking: a video pipeline only registers in manager.models AFTER its
multi-minute load, so concurrent loads raced the free-VRAM check into
mutual OOM. note_loading()/clear_loading() reserve the model's estimate
for the whole load window; _get_free_vram_gb subtracts reservations and
the orphan-VRAM check counts them (no more false "teardown leak" for an
in-progress load). The embeddings OOM-retry waits out active
reservations (bounded 300s) before evicting and retrying.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014S8VtAvG499SsCbeESRK7V

codai/__init__.py

@@ -16,7 +16,7 @@
 
 # Canonical product version for CoderAI — single source of truth. Both the API
 # metadata and the admin web UI read from here.
-__version__ = "0.1.35"
+__version__ = "0.1.36"
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
 # expandable_segments lets the allocator return freed pages to the driver even

codai/api/embeddings.py

@@ -549,6 +549,19 @@ async def _run_embeddings(request: EmbeddingsRequest, http_request: Request = No
                         torch.cuda.empty_cache()
                     except Exception:
                         pass
+                    # If another model (e.g. a video pipeline) is mid-load, its
+                    # VRAM reservation is active — wait for that load to finish
+                    # (bounded) before evicting/retrying, exactly like queued
+                    # requests for any other model wait out a load.
+                    try:
+                        import time as _time
+                        _deadline = _time.time() + 300
+                        while (getattr(multi_model_manager,
+                                       '_loading_reservations', None)
+                               and _time.time() < _deadline):
+                            await asyncio.sleep(2)
+                    except Exception:
+                        pass
                     # extra_vram_gb keeps the eviction meaningful even when the
                     # model has no VRAM estimate yet (first-ever load).
                     await asyncio.get_event_loop().run_in_executor(

codai/api/video.py

@@ -737,6 +737,40 @@ def _free_pipeline_vram(pipe) -> None:
         import torch as _t
     except Exception:
         _t = None
+    # Record the CUDA storage pointers of every tensor this pipeline owns
+    # BEFORE breaking it apart. After teardown+gc, any of these pointers still
+    # alive means an EXTERNAL structure (accelerate's tied_params_map, the
+    # attention dispatcher, a stray cache…) kept a strong reference — the
+    # "referenced elsewhere" teardown leak that empty_cache() can't reclaim.
+    # We then null exactly those tensors out of their list/dict referrers (the
+    # same targeted pass the CUDA text backend's cleanup uses), never touching
+    # tensors of other resident models.
+    _orig_cuda_ptrs = set()
+    _vram_used0 = -1.0
+    if _t is not None and pipe is not None:
+        try:
+            if _t.cuda.is_available():
+                _vram_used0 = sum(
+                    _t.cuda.memory_allocated(i)
+                    for i in range(_t.cuda.device_count())) / 1e9
+        except Exception:
+            pass
+        try:
+            for _cv in (getattr(pipe, 'components', {}) or {}).values():
+                if not hasattr(_cv, 'parameters'):
+                    continue
+                for _tens in list(_cv.parameters()) + list(getattr(_cv, 'buffers', lambda: [])()):
+                    try:
+                        if _tens.is_cuda:
+                            _orig_cuda_ptrs.add(_tens.untyped_storage().data_ptr())
+                        # bnb 4-bit params keep quantized data on .data
+                        _d = getattr(_tens, 'data', None)
+                        if _d is not None and getattr(_d, 'is_cuda', False):
+                            _orig_cuda_ptrs.add(_d.untyped_storage().data_ptr())
+                    except Exception:
+                        continue
+        except Exception:
+            pass
     try:
         if pipe is not None:
             _comps = getattr(pipe, 'components', {}) or {}
@@ -821,6 +855,65 @@ def _free_pipeline_vram(pipe) -> None:
                 _t.cuda.empty_cache()
         except Exception:
             pass
+    # Referenced-leak breaker: any of THIS pipeline's tensors still on the GPU
+    # after gc are pinned by an external list/dict (accelerate maps, dispatcher
+    # caches…). Null them out of their referrers so the next gc+empty_cache can
+    # actually return the VRAM — this is what turns the "~22 GB untracked
+    # (teardown leak)" card back into a usable one without a restart.
+    if _t is not None and _orig_cuda_ptrs:
+        try:
+            if _t.cuda.is_available():
+                _broken = 0
+                _still_gb = 0.0
+                _holders = []
+                for _obj in _gc.get_objects():
+                    if not (isinstance(_obj, _t.Tensor) and _obj.is_cuda):
+                        continue
+                    try:
+                        if _obj.untyped_storage().data_ptr() not in _orig_cuda_ptrs:
+                            continue
+                    except Exception:
+                        continue
+                    _gb = _obj.numel() * _obj.element_size() / 1e9
+                    _still_gb += _gb
+                    for _ref in _gc.get_referrers(_obj):
+                        try:
+                            if isinstance(_ref, list):
+                                for _i, _it in enumerate(_ref):
+                                    if _it is _obj:
+                                        _ref[_i] = None
+                                        _broken += 1
+                            elif isinstance(_ref, dict):
+                                for _k, _v in list(_ref.items()):
+                                    if _v is _obj:
+                                        _ref[_k] = None
+                                        _broken += 1
+                                if _gb > 0.05 and len(_holders) < 6:
+                                    try:
+                                        _holders.append(
+                                            f"{_gb:.2f} GB dict{list(_ref.keys())[:3]}")
+                                    except Exception:
+                                        pass
+                        except Exception:
+                            pass
+                if _broken or _still_gb > 0.5:
+                    print(f"  [video] teardown-leak breaker: {_still_gb:.1f} GB of this "
+                          f"pipeline's tensors survived teardown; nulled {_broken} "
+                          f"external reference(s)"
+                          + (f"; holders: {_holders}" if _holders else ""))
+                    for _ in range(2):
+                        _gc.collect()
+                    _t.cuda.synchronize()
+                    _t.cuda.empty_cache()
+                    _t.cuda.synchronize()
+                if _vram_used0 >= 0:
+                    _vram_used1 = sum(
+                        _t.cuda.memory_allocated(i)
+                        for i in range(_t.cuda.device_count())) / 1e9
+                    print(f"  [video] pipeline free: torch-allocated "
+                          f"{_vram_used0:.1f} → {_vram_used1:.1f} GB")
+        except Exception as _lbe:
+            print(f"  [video] teardown-leak breaker failed: {_lbe}")
     try:
         from codai.models.manager import _trim_cpu_ram
         _trim_cpu_ram()
@@ -3499,6 +3592,10 @@ async def video_generations(request: VideoGenerationRequest,
         except Exception:
             _ram_before = -1.0
         from codai.tasks import loading_task
+        # Reserve the model's VRAM estimate for the whole load window so
+        # concurrent loads (e.g. an embedding model) see the card as spoken-for
+        # instead of racing this multi-minute load into a mutual OOM.
+        multi_model_manager.note_loading(model_key)
         try:
             with loading_task(model_name, model_type="video"):
                 pipe = await asyncio.get_event_loop().run_in_executor(
@@ -3544,6 +3641,7 @@ async def video_generations(request: VideoGenerationRequest,
                       + _tb2.format_exc())
                 _retried_fresh = False
             if not _retried_fresh:
+                multi_model_manager.clear_loading(model_key)
                 raise HTTPException(status_code=500, detail=f"Failed to load video model: {e}")
         # Fuse any configured acceleration/distillation LoRA (Lightning / Lightx2v /
         # LCM) into the freshly loaded pipeline. Done once at load; cached pipes keep
@@ -3565,6 +3663,8 @@ async def video_generations(request: VideoGenerationRequest,
             print(f"  [video][accel] skipped: {_e}")
         multi_model_manager.models[model_key] = pipe
         multi_model_manager.current_model_key = model_key
+        # Registered — normal tracking takes over from the load reservation.
+        multi_model_manager.clear_loading(model_key)
         # Record the model's FULL footprint (GPU-resident + offloaded-to-RAM). Pass
         # ram_before so an offloaded load's host-RAM weights are counted — otherwise
         # measured_vram_gb collapses to the tiny GPU slice (~0.3 GB) and eviction
@@ -3625,6 +3725,7 @@ async def video_generations(request: VideoGenerationRequest,
             print(f"  [video] fallback {_ri}/{len(_rungs)-1}: reloading "
                   f"'{model_name}' with strategy='{_strat}' (incremental={_incr}) — "
                   f"NOT returning to client…")
+            multi_model_manager.note_loading(model_key)
             try:
                 with _loading_task(model_name, model_type="video"):
                     pipe = await _loop.run_in_executor(
@@ -3632,12 +3733,14 @@ async def video_generations(request: VideoGenerationRequest,
                         _strat, _model_cfg, _incr)
                 multi_model_manager.models[model_key] = pipe
                 multi_model_manager.current_model_key = model_key
+                multi_model_manager.clear_loading(model_key)
             except Exception as _le:
                 print(f"  [video] fallback {_ri} load FAILED: "
                       f"{str(_le).splitlines()[0] if str(_le) else type(_le).__name__}\n"
                       + _tb.format_exc())
                 _last_err = _le
                 _le.__traceback__ = None
+                multi_model_manager.clear_loading(model_key)
                 multi_model_manager.models.pop(model_key, None)
                 multi_model_manager.model_pools.pop(model_key, None)
                 _free_pipeline_vram(pipe)

codai/models/manager.py

@@ -829,6 +829,15 @@ class MultiModelManager:
         self._model_ready_event = threading.Event()
         self._model_ready_event.set()   # initially ready (nothing loading)
         self.model_pools: Dict[str, ModelInstancePool] = {}  # per-key instance pools
+        # In-flight load reservations: model_key -> estimated GB. A big pipeline
+        # (video) takes minutes to load and is only registered in self.models
+        # AFTER the load completes — during that window the card is filling up
+        # but nothing is tracked, so a concurrent load (e.g. an embedding model)
+        # passes the free-VRAM check and both race into OOM. Reserving the
+        # estimate up front makes _get_free_vram_gb conservative for the whole
+        # load window and keeps sweep_orphan_vram from mislabeling an
+        # in-progress load as a teardown leak.
+        self._loading_reservations: Dict[str, float] = {}
         # Set once a CUDA device-side assert / unrecoverable CUDA error is seen.
         # The CUDA context is corrupted process-wide after such an error, so all
         # further GPU work is futile until the server is restarted.  We surface
@@ -2329,8 +2338,30 @@ class MultiModelManager:
                 total += self._get_model_used_vram_gb(key)
             except Exception:
                 pass
+        # A load in progress isn't in self.models yet but its weights are already
+        # streaming onto the card — count the reservation so the orphan check
+        # doesn't call an in-flight load a teardown leak.
+        for key, gb in list(self._loading_reservations.items()):
+            if key not in self.models:
+                total += gb
         return total
 
+    def note_loading(self, model_key: str, gb: float = 0.0) -> None:
+        """Reserve VRAM for a load that is about to start (see
+        _loading_reservations). Pass the model's estimate, or 0 to look it up."""
+        try:
+            if not gb or gb <= 0:
+                gb = self._get_model_used_vram_gb(model_key)
+            if gb and gb > 0:
+                self._loading_reservations[model_key] = float(gb)
+        except Exception:
+            pass
+
+    def clear_loading(self, model_key: str) -> None:
+        """Drop a load reservation (call from finally once the model is
+        registered in self.models — or the load failed)."""
+        self._loading_reservations.pop(model_key, None)
+
     def sweep_orphan_vram(self, context: str = "") -> None:
         """Reclaim any UNREFERENCED orphaned GPU memory and surface a diagnostic for
         the referenced kind. gc.collect() makes a just-dropped-but-not-collected
@@ -2404,7 +2435,12 @@ class MultiModelManager:
             except Exception:
                 pass
         if got:
-            return total_free
+            # Subtract in-flight load reservations: that VRAM is spoken for even
+            # if the loading model hasn't materialized all its weights yet.
+            # max(0, …) — late in a load the reservation overlaps memory the
+            # card already reports as used; being conservative here is the point.
+            _reserved = sum(self._loading_reservations.values())
+            return max(0.0, total_free - _reserved)
         return 999.0  # Unknown — assume enough
 
     @staticmethod