vram: learn+persist stable offload split; front: graceful crash + task cancel routing

manager (record_vram_delta): record the FULL footprint after a load — including
offloaded loads (no longer skipped) — so the next load skips the OOM→move-to-RAM
retry dance:
  * measured_vram_gb      — GPU VRAM delta + reserve (GPU portion; eviction uses it)
  * measured_ram_gb       — host-RAM delta (the CPU-offloaded layers)   [new]
  * measured_n_gpu_layers — the layer split llama.cpp settled on         [new]
  total need = vram + ram. All gated by force_vram_update (else only when
  used_vram_gb unset). New _learned_n_gpu_layers(): when n_gpu_layers is auto/-1,
  reuse the learned split so a reload jumps straight to the config that fit. Both
  load sites snapshot host RAM and read the backend's settled n_gpu_layers.

frontproxy:
  * graceful streaming proxy: when an engine dies mid-stream (e.g. a CUDA
    ggml_abort SIGABRT on a VRAM-OOM decode), end the relayed stream cleanly
    instead of throwing an unhandled ASGI exception (noisy traceback).
  * task actions (cancel/interrupt/pause/resume/restart/DELETE) now route to the
    engine that OWNS the task — fan out to every live engine (each validates the
    same session cookie), first success wins — fixing 'Task not found' for tasks on
    a non-primary engine; front-only synthetic ids handled locally.

Bump version to 0.1.7.

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.6"
+__version__ = "0.1.7"
 
 # Configure the CUDA caching allocator BEFORE torch is imported anywhere.
 # expandable_segments lets the allocator return freed pages to the driver even

codai/frontproxy/app.py

@@ -1109,6 +1109,48 @@ class FrontProxy:
             target.loaded_models = set(target.loaded_models) | {path}
         return resp
 
+    async def task_action(self, request: Request) -> Response:
+        """Per-task actions (cancel / interrupt / pause / resume / restart, and
+        DELETE remove) must reach the engine that OWNS the task. A task can run on
+        ANY engine (or the system worker, for downloads), so the catch-all's
+        always-the-primary routing makes a task on another engine read 'Task not
+        found'. Front-only synthetic Tasks-page entries have no engine task at all.
+
+        Resolve it here: handle the front-synthetic ids, else fan the request out to
+        every engine — each validates the same session cookie (HMAC), so the engine
+        that owns the task acts on it and the rest 404; the first success wins."""
+        if not await self.is_admin(request):
+            return JSONResponse({"detail": "Unauthorized"}, status_code=401)
+        # The task id is the last or second-to-last path segment (…/{id} for DELETE,
+        # …/{id}/{action} for POST).
+        parts = [p for p in request.url.path.split("/") if p]
+        task_id = parts[-1] if request.method == "DELETE" else (
+            parts[-2] if len(parts) >= 2 else "")
+        # Front-only synthetic entries (see _merge_engine_tasks): no engine task.
+        if task_id.startswith("queued-"):
+            # A front-queued (not-yet-running) request is owned by the waiting client
+            # coroutine; it drops out when that client disconnects. There's nothing
+            # to cancel on an engine.
+            return JSONResponse(
+                {"detail": "Queued request — it cancels when the client disconnects."},
+                status_code=409)
+        if task_id.startswith(("inflight-", "loading-")):
+            return JSONResponse({"detail": "Task not found"}, status_code=404)
+        # Real task id: fan out to every live engine (incl. the system worker, which
+        # owns download/cache tasks). First 2xx wins; otherwise relay the last reply.
+        body = await request.body()
+        last = None
+        for e in self.registry.all():
+            if not e.is_alive():
+                continue
+            resp = await self._forward_to_engine(request, e, body)
+            code = getattr(resp, "status_code", 502)
+            if 200 <= code < 300:
+                return resp
+            last = resp
+        return last if last is not None else JSONResponse(
+            {"detail": "Task not found"}, status_code=404)
+
     def _cooling_engines(self) -> list:
         """Which engines are in thermal cooldown right now (for the Tasks banner).
 
@@ -1437,24 +1479,37 @@ class FrontProxy:
         _meas = (_rid is not None and rp_resp.status_code == 200
                  and "text/event-stream" in (rp_resp.headers.get("content-type") or ""))
 
-        async def _counting_iter():
+        async def _relay_iter():
             import time as _t
             t0 = _t.monotonic(); ntok = 0; last = 0.0
-            async for raw in rp_resp.aiter_raw():
-                ntok += raw.count(b"data:")
-                now = _t.monotonic()
-                if now - last >= 0.5:              # refresh ~2×/s, keep it cheap
-                    last = now
-                    dt = now - t0
-                    m = (engine.active or {}).get(_rid)
-                    if m is not None:
-                        m["step"] = ntok
-                        m["rate"] = round(ntok / dt, 1) if dt > 0 else 0.0
-                yield raw
+            try:
+                async for raw in rp_resp.aiter_raw():
+                    if _meas:
+                        ntok += raw.count(b"data:")
+                        now = _t.monotonic()
+                        if now - last >= 0.5:          # refresh ~2×/s, keep it cheap
+                            last = now
+                            dt = now - t0
+                            m = (engine.active or {}).get(_rid)
+                            if m is not None:
+                                m["step"] = ntok
+                                m["rate"] = round(ntok / dt, 1) if dt > 0 else 0.0
+                    yield raw
+            except httpx.HTTPError as exc:
+                # The upstream engine dropped the connection mid-stream — almost
+                # always because it CRASHED (e.g. a CUDA ggml_abort → SIGABRT on a
+                # VRAM-OOM during a large-context decode) or was restarted. The body
+                # is already partly sent, so we can't swap in a clean JSON error; just
+                # stop relaying. The supervisor respawns the engine. Without this the
+                # truncated read escapes as an unhandled ASGI exception (a noisy
+                # traceback) on every such crash.
+                print(f"[front] upstream engine#{engine.id} ({engine.name}) closed the "
+                      f"stream early on {path}: {exc!r}", flush=True)
+                return
 
         resp_headers = self._filter_headers(rp_resp.headers, _DROP_RESP)
         return StreamingResponse(
-            _counting_iter() if _meas else rp_resp.aiter_raw(),
+            _relay_iter(),
             status_code=rp_resp.status_code,
             headers=dict(resp_headers),
             media_type=rp_resp.headers.get("content-type"),
@@ -1664,6 +1719,18 @@ def build_app(config, config_dir=None) -> FastAPI:
     async def _tasks(request: Request):
         return await front.poll(request)
 
+    # Per-task actions must reach the engine that OWNS the task (any engine, not
+    # just the primary), else the UI reads "Task not found". Registered before the
+    # catch-all so they're routed here instead of proxied to the primary.
+    @app.api_route("/admin/api/tasks/{task_id}/{action}", methods=["POST"],
+                   include_in_schema=False)
+    async def _task_action(task_id: str, action: str, request: Request):
+        return await front.task_action(request)
+
+    @app.delete("/admin/api/tasks/{task_id}", include_in_schema=False)
+    async def _task_remove(task_id: str, request: Request):
+        return await front.task_action(request)
+
     @app.get("/admin/api/system-stats", include_in_schema=False)
     async def _system_stats(request: Request):
         return await front.system_stats(request)

codai/models/manager.py

@@ -1143,13 +1143,22 @@ class MultiModelManager:
 
                 print(f"Loading default model on demand: {self.default_model}")
                 _snap = self.vram_before_load()
+                _ram_snap = self._get_own_ram_gb()
                 kwargs['expected_vram_gb'] = self._get_model_used_vram_gb(self.default_model)
+                # Reuse a learned stable GPU-layer split so the load jumps straight
+                # to the config that fit (skips the OOM→move-to-RAM→retry dance).
+                _ngl = self._learned_n_gpu_layers(self.default_model, kwargs.get('n_gpu_layers'))
+                if _ngl is not None:
+                    kwargs['n_gpu_layers'] = _ngl
                 from codai.tasks import loading_task
                 with loading_task(self.default_model, model_type="text"):
                     model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs)
                 self._last_load_errors.pop(self.default_model, None)
                 self.add_model(self.default_model, model_manager)
-                self.record_vram_delta(self.default_model, _snap)
+                self.record_vram_delta(
+                    self.default_model, _snap, ram_before=_ram_snap,
+                    n_gpu_layers=getattr(getattr(model_manager, 'backend', None),
+                                         'n_gpu_layers', None))
                 self.current_model_key = self.default_model
                 print(f"Model loaded successfully: {self.default_model}")
                 self._model_ready_event.set()
@@ -1296,16 +1305,25 @@ class MultiModelManager:
                     except Exception as _ev_e:
                         print(f"  (ensure_vram_for warning: {_ev_e})")
                 _snap = self.vram_before_load()
+                _ram_snap = self._get_own_ram_gb()
                 # Tell the backend how much VRAM this model is expected to need so
                 # it can decide whether Flash-Attention-2 is safe (FA2 requires the
                 # whole model on GPU; it device-side-asserts when layers offload).
                 kwargs['expected_vram_gb'] = self._get_model_used_vram_gb(model_name)
+                # Reuse a learned stable GPU-layer split so the load jumps straight
+                # to the config that fit (skips the OOM→move-to-RAM→retry dance).
+                _ngl = self._learned_n_gpu_layers(model_name, kwargs.get('n_gpu_layers'))
+                if _ngl is not None:
+                    kwargs['n_gpu_layers'] = _ngl
                 from codai.tasks import loading_task
                 with loading_task(model_name, model_type="text"):
                     model_manager.load_model(model_name, backend_type=backend_type, **kwargs)
                 self._last_load_errors.pop(model_name, None)
                 self.add_model(model_name, model_manager)
-                self.record_vram_delta(model_name, _snap)
+                self.record_vram_delta(
+                    model_name, _snap, ram_before=_ram_snap,
+                    n_gpu_layers=getattr(getattr(model_manager, 'backend', None),
+                                         'n_gpu_layers', None))
                 self.current_model_key = model_name
                 print(f"Model loaded successfully: {model_name}"
                       + (f" (instance {inst_num})" if inst_num > 1 else ""))
@@ -2340,78 +2358,125 @@ class MultiModelManager:
         return self._free_vram_snapshot()
 
     def record_vram_delta(self, model_key: str, free_before: float,
-                          offloaded: bool = False) -> None:
-        """Call immediately after a model finishes loading to record actual VRAM consumed.
-
-        If the measured value exceeds the stored estimate by more than 10%, the real
-        value is written back into the model config and persisted to models.json so
-        future eviction decisions use the accurate figure.
-
-        ``offloaded`` MUST be True when the model was loaded with any CPU/disk
-        offload strategy (model/sequential/group/balanced/disk). In that case the
-        weights are not resident on the GPU, so the measured delta is a tiny,
-        meaningless lower bound — recording it would clobber a real full-GPU
-        measurement and make the next start under-estimate the footprint, pick a
-        full-GPU load, and OOM. So we skip recording entirely and keep the prior
-        estimate/measurement intact.
+                          offloaded: bool = False, ram_before: float = -1.0,
+                          n_gpu_layers=None) -> None:
+        """Call immediately after a model finishes loading to record its REAL memory
+        footprint on this hardware, so the next load reuses it instead of repeating
+        the OOM→move-layers-to-RAM retry dance.
+
+        Records up to three SYSTEM-OWNED fields (never the user's ``used_vram_gb``):
+          * ``measured_vram_gb``      — GPU VRAM delta + runtime reserve (the GPU
+                                        portion; eviction frees against this).
+          * ``measured_ram_gb``       — host-RAM delta (the CPU-offloaded layers).
+          * ``measured_n_gpu_layers`` — the layer split llama.cpp actually settled
+                                        on (after any progressive CPU-offload retry).
+        The TOTAL memory the model needs is ``measured_vram_gb + measured_ram_gb``.
+
+        An OFFLOADED load is recorded too (NOT skipped): paired with
+        ``measured_n_gpu_layers`` the figures are self-consistent — the next load
+        goes straight to the same split and uses the same VRAM — so there's no
+        under-estimate→full-GPU→OOM hazard the old skip was guarding against.
+
+        Per-field persist rule: with ``force_vram_update`` set, refresh every run;
+        otherwise only when ``used_vram_gb`` isn't configured (a configured value is
+        authoritative). ``offloaded`` is accepted for back-compat but no longer skips.
         """
         if free_before < 0:
             return
-        if offloaded:
-            return
         free_after = self._free_vram_snapshot()
         if free_after < 0:
             return
-        delta_gb = (free_before - free_after) / 1e9
-        if delta_gb <= 0:
-            return
-
-        # The delta is the resident WEIGHT footprint measured at load time. Add
-        # the runtime reserve (KV cache / activations / VAE-decode spike) so the
-        # value we cache and persist reflects the model's PEAK runtime need — not
-        # just its loaded weights — and future eviction frees enough headroom.
         cfg = self._config_for_model_key(model_key)
-        reserve_gb = self._runtime_reserve_gb(
-            cfg if isinstance(cfg, dict) else {}, model_key, delta_gb)
-        measured = round(delta_gb + reserve_gb, 3)
-        # In-memory truth for the rest of this run (used by eviction estimates).
-        self._measured_vram_gb[model_key] = measured
-        print(f"Measured VRAM for '{model_key}': {delta_gb:.2f} GB weights "
-              f"+ {reserve_gb:.2f} GB runtime reserve = {measured:.2f} GB")
-
-        # Persist the real value into the SEPARATE, system-owned
-        # `measured_vram_gb` field (never the user's `used_vram_gb`). Rules:
-        #   * force_vram_update set → refresh it EVERY run, even when
-        #     used_vram_gb is configured.
-        #   * otherwise → only when used_vram_gb is NOT configured (a configured
-        #     value is authoritative and must remain unchanged).
-        force_update = bool(cfg.get("force_vram_update")) if isinstance(cfg, dict) else False
-        if not force_update:
-            if isinstance(cfg, dict) and cfg.get("used_vram_gb") is not None:
-                return  # configured & not forced → leave it exactly as-is
-            # Avoid churning models.json when the measurement hasn't meaningfully
-            # changed from what's already persisted.
-            prev = cfg.get("measured_vram_gb") if isinstance(cfg, dict) else None
+        cfgd = cfg if isinstance(cfg, dict) else {}
+        force_update = bool(cfgd.get("force_vram_update"))
+        user_pinned = cfgd.get("used_vram_gb") is not None
+
+        # --- GPU VRAM footprint (weights resident on GPU + runtime reserve) ------
+        delta_gb = (free_before - free_after) / 1e9
+        measured_vram = None
+        if delta_gb > 0:
+            reserve_gb = self._runtime_reserve_gb(cfgd, model_key, delta_gb)
+            measured_vram = round(delta_gb + reserve_gb, 3)
+            self._measured_vram_gb[model_key] = measured_vram  # in-memory truth
+
+        # --- Host-RAM footprint (the CPU-offloaded layers) ----------------------
+        measured_ram = None
+        if ram_before is not None and ram_before >= 0:
             try:
-                if prev is not None and abs(measured - float(prev)) <= max(0.5, float(prev) * 0.10):
-                    return
-            except (TypeError, ValueError):
-                pass
+                rd = round(self._get_own_ram_gb() - ram_before, 3)
+                if rd > 0:
+                    measured_ram = rd
+            except Exception:
+                measured_ram = None
+
+        # --- Settled GPU-layer split (the value the retry loop landed on) -------
+        measured_layers = None
         try:
-            _why = "force_vram_update" if force_update else "no used_vram_gb configured"
-            print(f"  Saving measured_vram_gb for '{model_key}': {measured:.2f} GB ({_why})")
-            cfg = dict(cfg) if isinstance(cfg, dict) else {}
-            cfg["measured_vram_gb"] = measured
-            self.config[model_key] = cfg
-            # Persist to models.json via a SINGLE-field read-modify-write, so a
-            # secondary engine's stale in-memory models_data can't clobber a UI edit
-            # (e.g. n_ctx) made on the primary after this engine booted. Only
-            # measured_vram_gb on the matching entry is touched.
-            from codai.admin.routes import config_manager
-            if config_manager is not None:
-                config_manager.persist_model_field(model_key, "measured_vram_gb", measured)
-        except Exception as e:
-            print(f"  Warning: could not persist measured_vram_gb: {e}")
+            if n_gpu_layers is not None:
+                measured_layers = int(n_gpu_layers)
+        except (TypeError, ValueError):
+            measured_layers = None
+
+        _total = (measured_vram or 0.0) + (measured_ram or 0.0)
+        print(f"Measured footprint for '{model_key}': "
+              f"VRAM={'%.2f' % measured_vram if measured_vram is not None else 'n/a'} GB + "
+              f"RAM={'%.2f' % measured_ram if measured_ram is not None else 'n/a'} GB = "
+              f"{_total:.2f} GB total  "
+              f"(n_gpu_layers={measured_layers if measured_layers is not None else 'n/a'})")
+
+        # Persist each system-owned field, accumulating into the in-memory config so
+        # successive fields don't clobber each other. Each is written to models.json
+        # via a single-field read-modify-write (safe vs. a sibling engine's stale view).
+        working = dict(cfgd)
+
+        def _persist(field, value, *, vram_guard=False):
+            if value is None:
+                return
+            if not force_update and user_pinned:
+                return  # configured & not forced → authoritative, leave as-is
+            if vram_guard and not force_update:
+                # Avoid churning models.json on a <=10% no-op VRAM change.
+                prev = cfgd.get(field)
+                try:
+                    if prev is not None and abs(value - float(prev)) <= max(0.5, float(prev) * 0.10):
+                        return
+                except (TypeError, ValueError):
+                    pass
+            try:
+                working[field] = value
+                self.config[model_key] = dict(working)
+                from codai.admin.routes import config_manager
+                if config_manager is not None:
+                    config_manager.persist_model_field(model_key, field, value)
+                print(f"  Saved {field}={value} for '{model_key}' "
+                      f"({'force_vram_update' if force_update else 'no used_vram_gb'})")
+            except Exception as e:
+                print(f"  Warning: could not persist {field}: {e}")
+
+        _persist("measured_vram_gb", measured_vram, vram_guard=True)
+        _persist("measured_ram_gb", measured_ram)
+        _persist("measured_n_gpu_layers", measured_layers)
+
+    def _learned_n_gpu_layers(self, model_key: str, configured):
+        """Prefer a previously-learned stable GPU-layer split over an auto guess.
+
+        When ``n_gpu_layers`` is auto (-1 / None), reuse ``measured_n_gpu_layers``
+        (recorded by record_vram_delta after the offload-retry settled) so a reload
+        jumps straight to the split that fit — no OOM→move-to-RAM→retry on every
+        load. An explicit positive pin in the config always wins."""
+        if configured is not None and configured != -1:
+            return configured
+        cfg = self._config_for_model_key(model_key)
+        if isinstance(cfg, dict):
+            learned = cfg.get("measured_n_gpu_layers")
+            if learned is not None:
+                try:
+                    iv = int(learned)
+                    if iv >= 0:
+                        return iv
+                except (TypeError, ValueError):
+                    pass
+        return configured
 
     # In-process cache: model_id -> size_bytes (never stale within a run)
     _hf_size_cache: Dict[str, int] = {}