agharsallah commited on
Commit
7d636f8
·
1 Parent(s): c6cdf25

feat: Replace llama.cpp backend with in-process transformers backend for local GPU inference

Browse files
docs/adr/0032-llama-cpp-local-backend.md CHANGED
@@ -2,8 +2,11 @@
2
 
3
  ## Status
4
 
5
- Accepted (extends ADR-0024 *second inference backend / unified registry*, amends
6
- ADR-0015 *LiteLLM gateway*, ADR-0022 *per-agent explicit model binding*)
 
 
 
7
 
8
  ## Context
9
 
 
2
 
3
  ## Status
4
 
5
+ Superseded by ADR-0033. llama.cpp's persistent `llama-server` cannot hold a GPU under
6
+ ZeroGPU's per-call grant model; replaced by an in-process `transformers` backend that
7
+ works on any HF Space hardware (ADR-0024 *second inference backend / unified registry*,
8
+ ADR-0015 *LiteLLM gateway*, ADR-0022 *per-agent explicit model binding* remain in
9
+ force; this ADR is retained as a historical record only).
10
 
11
  ## Context
12
 
docs/adr/0033-local-inproc-transformers-backend.md ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # ADR-0033: Local In-Process `transformers` Backend (Supersedes ADR-0032)
2
+
3
+ **Status:** Accepted
4
+ **Date:** 2026-06-13
5
+ **Deciders:** project maintainers
6
+
7
+ ## Context
8
+
9
+ ADR-0032 added a llama.cpp backend: a persistent `llama-server` process that exposes
10
+ an OpenAI-compatible HTTP endpoint and uses the operator's GPU for the lifetime of that
11
+ process. This model is structurally incompatible with Hugging Face **ZeroGPU**.
12
+
13
+ ZeroGPU grants a GPU *only for the duration of a `@spaces.GPU`-decorated function call*,
14
+ then reclaims it. A long-lived HTTP server needs to hold the GPU between requests — there
15
+ is no GPU to hold on ZeroGPU. The same mismatch rules out vLLM-as-a-server for the same
16
+ reason. Per the ZeroGPU documentation: the runtime is Gradio-SDK only; the GPU is an
17
+ NVIDIA RTX Pro 6000 Blackwell (48 GB `large` / 96 GB `xlarge`); anonymous users get ~2
18
+ minutes of free GPU time per day, authenticated users ~5 minutes; and the canonical usage
19
+ pattern is `transformers`/`diffusers` with model weights placed on `cuda` at module load
20
+ and the forward pass wrapped in `@spaces.GPU`. CUDA is *emulated* (no-op) outside the
21
+ decorated function and real inside it.
22
+
23
+ A second goal existed alongside ZeroGPU compatibility: the backend must be
24
+ **hardware-agnostic**. An HF Space can be assigned a **dedicated GPU** (T4, L4, L40S,
25
+ A100, …) or run on a local CUDA box. The solution should serve equally well in all three
26
+ environments — ZeroGPU, dedicated GPU, and local CUDA — without a code branch per
27
+ environment.
28
+
29
+ Replacing the llama.cpp GGUF path also means explicitly dropping the **Llama Champion**
30
+ bonus-quest badge (which required a real llama.cpp runtime in the cast). This is a
31
+ deliberate tradeoff: ZeroGPU compatibility and hardware-agnosticism are higher-value than
32
+ one badge.
33
+
34
+ No new Python dependencies are needed: `torch` and `transformers` already ship
35
+ transitively via `sentence-transformers`.
36
+
37
+ ## Decision
38
+
39
+ We will replace the llama.cpp backend with an in-process `transformers` backend that runs
40
+ model inference inside a `@spaces.GPU`-decorated call, caches loaded weights at module
41
+ level in the parent process so each ZeroGPU call inherits them without re-loading, and
42
+ gates availability on operator capability rather than a URL environment variable.
43
+
44
+ Four concrete sub-decisions:
45
+
46
+ **1. Non-HTTP router seam (`ProfileSpec.kind`).** `ProfileSpec` gains a `kind` field
47
+ (`"litellm"` | `"local"`). When `kind == "local"` the router dispatches to
48
+ `LocalTransformersProvider` directly, bypassing LiteLLM and HTTP entirely. This is the
49
+ first backend that does not go through the HTTP gateway — it is a clean extension of the
50
+ router seam left open by ADR-0024, not a hack around it.
51
+
52
+ **2. In-process forward pass with effect-free decorator.** `LocalTransformersProvider`
53
+ wraps the `transformers` forward pass in a module-level `@spaces.GPU(duration=<dynamic>)`
54
+ function. On ZeroGPU this decorator acquires and releases the GPU for that call's
55
+ duration. On a dedicated GPU or local CUDA box the decorator is a no-op (effect-free
56
+ passthrough), so the same code path is a persistent in-process provider on those
57
+ environments. No environment-specific branching in the provider.
58
+
59
+ **3. Parent-process model cache, lazy on first use.** The official ZeroGPU guidance is to
60
+ load model weights at module level (import time) so forked `@spaces.GPU` calls inherit
61
+ them via copy-on-write. We deviate deliberately: weights are loaded lazily on the first
62
+ `complete()` call and cached in a module-level `_LOADED` dict. This avoids loading unused
63
+ models at app boot and keeps the no-API-key stub fast. `torch` and `transformers` imports
64
+ are also lazy (never at module import) to prevent CUDA initialisation before the fork and
65
+ to avoid tripping the PyTorch multiprocessing fork guard. The tradeoff is that the *first*
66
+ call to a model incurs load time; all subsequent calls within the process inherit the
67
+ cache for free, matching the per-call efficiency of the module-level-load pattern.
68
+
69
+ **4. Capability gate and per-run opt-in.** `local_catalogue.has_credentials()` gates on
70
+ one of three signals: `SPACES_ZERO_GPU` is set in env (HF ZeroGPU Space), or
71
+ `LOCAL_INFERENCE=1` is set (explicit operator opt-in on a dedicated GPU or local CUDA
72
+ box), or a cached `torch.cuda.is_available()` probe is true — but the auto-probe runs
73
+ only when the env argument is `None` or `os.environ` itself, so tests passing an explicit
74
+ env dict never import torch and stay deterministic without a GPU. Picking "Local GPU" in
75
+ the Lab backend radio is the per-run opt-in; when none of the three signals is present the
76
+ backend stays inactive and the deterministic stub owns the no-config demo path.
77
+
78
+ ## Alternatives Considered
79
+
80
+ | Option | Pros | Cons |
81
+ |--------|------|------|
82
+ | Keep llama.cpp + server-per-request workaround | Preserves Llama Champion badge; GGUF models need no full-precision weights | Structurally incompatible with ZeroGPU; server startup latency per request; high complexity |
83
+ | vLLM-as-a-server on ZeroGPU | High throughput batching | Same persistent-process / per-call-GPU mismatch as llama.cpp |
84
+ | `llama-cpp-python` (in-process, no server) | Llama Champion badge preserved; GGUF quantisation | Additional heavy binary dep; GGUF format separate from HF model hub; weaker `transformers` ecosystem integration |
85
+ | In-process `transformers` (chosen) | Works on ZeroGPU, dedicated GPU, and local CUDA without branching; no new deps; full HF Hub model catalogue; keeps OpenBMB/MiniCPM and Tiny-Titan lanes | Drops Llama Champion badge; full-precision weights (larger VRAM footprint than GGUF quants) |
86
+
87
+ ## Consequences
88
+
89
+ **Positive:**
90
+ - The Space runs on ZeroGPU free tier with no code changes — every forward pass
91
+ naturally fits the per-call GPU grant model.
92
+ - The same binary runs on a dedicated GPU (T4/L4/L40S/A100) or local CUDA box with no
93
+ env-branch; the decorator is transparent.
94
+ - No new Python dependencies — `torch` and `transformers` are already transitive deps.
95
+ - The parent-process cache means each ZeroGPU call after the first is weight-load-free
96
+ within a session.
97
+ - **Prize-lane impact:** keeps the **OpenBMB / MiniCPM** track (MiniCPM4.1-8B is in the
98
+ local catalogue), the **Tiny Titan** lane (Qwen2.5-3B-Instruct at the `tiny` tier is
99
+ ≤4B), and the **Community Choice** on-device-inference story for the HF Space demo.
100
+
101
+ **Negative / Risks:**
102
+ - **Llama Champion badge is explicitly dropped.** No llama.cpp runtime in the cast means
103
+ this submission does not qualify for that bonus quest. This is a deliberate, accepted
104
+ tradeoff.
105
+ - First call to a model in a fresh process incurs full weight-load latency. On ZeroGPU
106
+ with the daily quota (~2–5 min GPU/day) this is a real cost on cold sessions.
107
+ - Full-precision (BF16/FP16) weights require more VRAM than GGUF quantised equivalents.
108
+ The ZeroGPU `large` tier (48 GB) is the practical ceiling; models above ~28B BF16 will
109
+ not fit without external quantisation.
110
+ - The ZeroGPU free-tier quota (≈5 min GPU/day authenticated) limits live demo length.
111
+ A dedicated GPU Space eliminates this limit but costs credits.
112
+ - Lazy torch import means the first call also pays Python import overhead for
113
+ `torch` + `transformers`. Acceptable for interactive demo pacing; unacceptable for
114
+ low-latency production workloads.
115
+
116
+ **Neutral / Notes:**
117
+ - `llamacpp_catalogue.py` and `llamacpp_server.py` are deleted; `app.py`'s
118
+ `gpu_selftest` `@spaces.GPU` guard is retained — it detects ZeroGPU availability at
119
+ startup and is unrelated to the inference path.
120
+ - The single catalogue entry tagged as a tier default is `Qwen/Qwen2.5-3B-Instruct`
121
+ (`tiny`). Additional models (MiniCPM4.1-8B, Qwen2.5-7B-Instruct) are in the catalogue
122
+ but untagged; ZeroGPU quota pressure justifies keeping the default footprint minimal.
123
+ - Tests live in `tests/test_local_backend.py`. All 676 tests pass; the capability-gate
124
+ logic is fully covered without a GPU or torch import in test processes.
125
+
126
+ ## Related ADRs
127
+
128
+ - ADR-0032: llama.cpp local backend — superseded by this ADR; retained as historical record.
129
+ - ADR-0024: Hugging Face inference backend / unified registry — the `ProfileSpec.kind` seam this ADR extends.
130
+ - ADR-0015: LiteLLM gateway — bypassed by `kind="local"`; all other backends still go through it.
131
+ - ADR-0022: Per-agent explicit model binding — unchanged; `local:<repo_id>` qualified keys follow the same binding contract.
132
+ - ADR-0019: Single model catalogue — local catalogue follows the same `binding_for` / `has_credentials` interface.
docs/architecture/model-routing.md CHANGED
@@ -77,7 +77,7 @@ Profiles map to the OpenAI-compatible vLLM endpoints served on Modal
77
  The LiteLLM model string for an OpenAI-compatible custom endpoint is
78
  `openai/<served_model_id>` with `api_base` pointing at the endpoint's `/v1` URL.
79
 
80
- ### Backends: Modal · Hugging Face · llama.cpp
81
 
82
  A *backend* is just a catalogue + a binding rule, unified behind one registry
83
  (`src/models/inference.py`, ADR-0024). Models are named by a **backend-qualified
@@ -88,25 +88,37 @@ working. Three backends ship today:
88
  |---|---|---|---|
89
  | Modal | *(bare)* | vLLM endpoints you deploy on Modal GPUs | `MODAL_WORKSPACE` / `MODAL_LLM_BASE_URL` |
90
  | Hugging Face | `hf:` | serverless Inference Providers router | `HF_TOKEN` |
91
- | llama.cpp | `llamacpp:` | **local** GGUF via `llama-server`, GPU when present | `LLAMACPP_BASE_URL` |
92
 
93
- The llama.cpp backend (ADR-0032) runs a cast fully on your own machine. Launch a
94
- model and export the URL it prints:
 
 
95
 
96
- ```bash
97
- uv run python -m src.models.llamacpp_server nemotron-3-nano-4b # GPU auto-detected
98
- export LLAMACPP_BASE_URL=http://127.0.0.1:8080/v1
99
- ```
 
 
 
 
 
 
 
 
 
 
 
 
 
 
100
 
101
- The launcher detects an accelerator Apple Metal on macOS, NVIDIA CUDA via
102
- `nvidia-smi`, else CPU — and offloads every layer to the GPU (`-ngl 999`) when one
103
- is present. `llama-server` downloads the GGUF on first run (`-hf`) and serves it
104
- under `--alias <key>`, so the engine binds to a stable id (`openai/<key>`) through
105
- the same LiteLLM transport. Bind a tier to a local model with a qualified key:
106
 
107
  ```yaml
108
  profiles:
109
- tiny: { endpoint: "llamacpp:nemotron-3-nano-4b", temperature: 0.7, max_tokens: 192 }
110
  ```
111
 
112
  ### Real cost → Governor
@@ -182,9 +194,9 @@ everything on the big model.
182
  - `src/core/registry.py` — `Registry.from_world()` (a UI/LLM-composed run on the same path)
183
  - `src/models/litellm_provider.py` — `LiteLLMProvider` (live transport, real cost)
184
  - `src/models/modal_catalogue.py` — engine view of the catalogue (key → binding)
185
- - `src/models/inference.py` — unified backend registry (Modal · HF · llama.cpp); qualified keys
186
- - `src/models/llamacpp_catalogue.py` — local GGUF catalogue (key binding)
187
- - `src/models/llamacpp_server.py` — `llama-server` launcher: GPU detection + command building
188
  - `src/core/manifest.py` — `resolve_model()` (env → catalogue default)
189
  - `src/core/registry.py` — `build_router()`, `_resolve_model_endpoints()`, `_expand_env()`
190
  - `src/models/provider.py` — `ModelProvider.last_usage`, `estimate_tokens()`
 
77
  The LiteLLM model string for an OpenAI-compatible custom endpoint is
78
  `openai/<served_model_id>` with `api_base` pointing at the endpoint's `/v1` URL.
79
 
80
+ ### Backends: Modal · Hugging Face · Local GPU
81
 
82
  A *backend* is just a catalogue + a binding rule, unified behind one registry
83
  (`src/models/inference.py`, ADR-0024). Models are named by a **backend-qualified
 
88
  |---|---|---|---|
89
  | Modal | *(bare)* | vLLM endpoints you deploy on Modal GPUs | `MODAL_WORKSPACE` / `MODAL_LLM_BASE_URL` |
90
  | Hugging Face | `hf:` | serverless Inference Providers router | `HF_TOKEN` |
91
+ | Local GPU | `local:` | in-process `transformers` model on the host's own GPU | `SPACES_ZERO_GPU` / `LOCAL_INFERENCE=1` / CUDA auto-detected |
92
 
93
+ The local backend (ADR-0033, supersedes ADR-0032) runs a cast fully in-process
94
+ no HTTP server, no extra process, no token. It loads a small instruct model via
95
+ `transformers` inside a `@spaces.GPU` function (`LocalTransformersProvider`,
96
+ `src/models/local_provider.py`). The hardware path is transparent:
97
 
98
+ - **ZeroGPU Space** — `@spaces.GPU` allocates a GPU per call from the shared pool
99
+ (~5 min/day free quota). Enabled when `SPACES_ZERO_GPU` is set.
100
+ - **Dedicated-GPU Space** (T4 / L4 / L40S / A100) — persistent GPU, no per-call
101
+ quota. Enabled with `LOCAL_INFERENCE=1`.
102
+ - **Local CUDA box** — the same `LOCAL_INFERENCE=1` flag, or CUDA auto-detected at
103
+ startup.
104
+
105
+ Off ZeroGPU and without `LOCAL_INFERENCE=1`, `@spaces.GPU` is a no-op and the
106
+ engine falls back to the deterministic stub — so the demo is always reproducible on
107
+ CPU-only hosts. Pick "Local GPU" in the Lab's backend radio to opt in per run.
108
+
109
+ Available models (all ≤32B; select via `local:<repo_id>`):
110
+
111
+ | Key | Model | Notes |
112
+ |---|---|---|
113
+ | `local:Qwen/Qwen2.5-3B-Instruct` | Qwen 2.5 3B | **tiny default** — latency + quota guardrail |
114
+ | `local:openbmb/MiniCPM4.1-8B` | MiniCPM 4.1 8B | alternate; OpenBMB prize lane |
115
+ | `local:Qwen/Qwen2.5-7B-Instruct` | Qwen 2.5 7B | alternate |
116
 
117
+ Bind a tier to a local model with a qualified key:
 
 
 
 
118
 
119
  ```yaml
120
  profiles:
121
+ tiny: { endpoint: "local:Qwen/Qwen2.5-3B-Instruct", temperature: 0.7, max_tokens: 192 }
122
  ```
123
 
124
  ### Real cost → Governor
 
194
  - `src/core/registry.py` — `Registry.from_world()` (a UI/LLM-composed run on the same path)
195
  - `src/models/litellm_provider.py` — `LiteLLMProvider` (live transport, real cost)
196
  - `src/models/modal_catalogue.py` — engine view of the catalogue (key → binding)
197
+ - `src/models/inference.py` — unified backend registry (Modal · HF · Local GPU); qualified keys
198
+ - `src/models/local_catalogue.py` — local model catalogue + capability gate (`has_credentials`)
199
+ - `src/models/local_provider.py` — `LocalTransformersProvider`: in-process `@spaces.GPU` inference
200
  - `src/core/manifest.py` — `resolve_model()` (env → catalogue default)
201
  - `src/core/registry.py` — `build_router()`, `_resolve_model_endpoints()`, `_expand_env()`
202
  - `src/models/provider.py` — `ModelProvider.last_usage`, `estimate_tokens()`
docs/architecture/next-steps/arena-roadmap.md CHANGED
@@ -24,7 +24,7 @@ drift from the log.
24
  | Competitive scenarios | ◐ only The Steeped is a real game; Open Table and Oracle Grove have **no judge at all** | `config/scenarios/*.yaml` |
25
  | Leaderboard / history UI | ❌ none | `src/ui/fishbowl/` |
26
  | Commentator | ❌ none (narrator feed is a transcript, not commentary) | `src/ui/fishbowl/show.py` |
27
- | Models | ◐ Modal vLLM (Nemotron 4B/30B/14B, MiniCPM 8B + o-4.5, Gemma 4 12B/26B), HF serverless (Arch-Router-1.5B), stub. No llama.cpp, no fine-tune, no gpt-oss | `modal/catalogue.py:144-296`, `src/models/hf_catalogue.py:71-85` |
28
 
29
  ---
30
 
@@ -177,10 +177,13 @@ engine change:
177
  OpenAI open model is in the catalogue. Add `gpt-oss-20b` (vLLM serves it) and
178
  make it the `strong`-tier default for the live demo path. Highest strategic
179
  priority in this workstream.
180
- 2. **llama.cpp backend** (🦙 *Llama Champion*) — third backend in
181
- `src/models/inference.py`'s registry: `LLAMA_CPP_BASE_URL` env var, a tiny
182
- static catalogue (`llamacpp:` key prefix), OpenAI-compatible transport
183
- already handled by LiteLLM. ~1 day, unlocks a badge.
 
 
 
184
  3. **Broaden the Modal catalogue** with strong ≤32B chat models so arena seats
185
  differ meaningfully: Qwen3 (4B/8B/14B/32B), Llama-3.1-8B-Instruct,
186
  Phi-4-14B, Mistral-Small-24B. Prefer models vLLM serves without
@@ -305,7 +308,8 @@ Phase D (reach) W4 models (gpt-oss → llama.cpp → catalogue → fine-tu
305
  no-API-key stub path kept fully working.
306
  - Prize coverage per phase: A unlocks 📡 trace export; B unlocks the Best
307
  Agent/Best Demo arena story; C is the delight criterion + 🐜 Tiny Titan; D is
308
- OpenAI track + 🦙 + 🎯.
 
309
 
310
  ## Beyond (post-arena ideas, in rough value order)
311
 
 
24
  | Competitive scenarios | ◐ only The Steeped is a real game; Open Table and Oracle Grove have **no judge at all** | `config/scenarios/*.yaml` |
25
  | Leaderboard / history UI | ❌ none | `src/ui/fishbowl/` |
26
  | Commentator | ❌ none (narrator feed is a transcript, not commentary) | `src/ui/fishbowl/show.py` |
27
+ | Models | ◐ Modal vLLM (Nemotron 4B/30B/14B, MiniCPM 8B + o-4.5, Gemma 4 12B/26B), HF serverless (Arch-Router-1.5B), in-process Local GPU (Qwen2.5-3B/7B, MiniCPM4.1-8B — ADR-0033), stub. No fine-tune, no gpt-oss | `modal/catalogue.py:144-296`, `src/models/hf_catalogue.py:71-85`, `src/models/local_catalogue.py` |
28
 
29
  ---
30
 
 
177
  OpenAI open model is in the catalogue. Add `gpt-oss-20b` (vLLM serves it) and
178
  make it the `strong`-tier default for the live demo path. Highest strategic
179
  priority in this workstream.
180
+ 2. **Local GPU backend** ( *shipped* — ADR-0033) — third backend in
181
+ `src/models/inference.py`'s registry: in-process `transformers` via
182
+ `@spaces.GPU` (`local:` key prefix, `LOCAL_INFERENCE=1` opt-in). Works on
183
+ ZeroGPU, dedicated-GPU Spaces, and local CUDA boxes. Supports the Community-
184
+ Choice / Tiny-Titan / OpenBMB lanes. **Note:** this replaces the earlier
185
+ llama.cpp design (ADR-0032) — the 🦙 Llama Champion runtime badge is not
186
+ pursued; on-device inference ships instead as the in-process GPU path.
187
  3. **Broaden the Modal catalogue** with strong ≤32B chat models so arena seats
188
  differ meaningfully: Qwen3 (4B/8B/14B/32B), Llama-3.1-8B-Instruct,
189
  Phi-4-14B, Mistral-Small-24B. Prefer models vLLM serves without
 
308
  no-API-key stub path kept fully working.
309
  - Prize coverage per phase: A unlocks 📡 trace export; B unlocks the Best
310
  Agent/Best Demo arena story; C is the delight criterion + 🐜 Tiny Titan; D is
311
+ OpenAI track + 🎯 (Local GPU / on-device inference already ships for Community-
312
+ Choice, OpenBMB, and Tiny-Titan lanes — 🦙 Llama Champion is not pursued).
313
 
314
  ## Beyond (post-arena ideas, in rough value order)
315
 
docs/runbook-live-mode.md CHANGED
@@ -74,13 +74,33 @@ catalogue key in `config/models.yaml` (source of truth: `modal/catalogue.py`).
74
  ### Option B — one explicit endpoint
75
 
76
  Point every profile at a single OpenAI-compatible base URL (one Modal-served
77
- model, or a local llama.cpp / vLLM box):
78
 
79
  ```ini
80
  # .env
81
  MODAL_LLM_BASE_URL=https://your-workspace--google-llms-gemma-4-12b.modal.run/v1
82
  ```
83
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  ### Per-profile overrides
85
 
86
  Highest priority. Override the model string bound to any profile — the cheapest
 
74
  ### Option B — one explicit endpoint
75
 
76
  Point every profile at a single OpenAI-compatible base URL (one Modal-served
77
+ model, or any other OpenAI-compatible endpoint):
78
 
79
  ```ini
80
  # .env
81
  MODAL_LLM_BASE_URL=https://your-workspace--google-llms-gemma-4-12b.modal.run/v1
82
  ```
83
 
84
+ ### Option C — Local GPU (in-process transformers)
85
+
86
+ Run inference in-process on the host's own GPU — no server to launch, no token.
87
+ The engine uses `LocalTransformersProvider` behind a `@spaces.GPU` function,
88
+ which works on ZeroGPU Spaces, dedicated-GPU Spaces (T4/L4/L40S/A100), and local
89
+ CUDA boxes. On a CPU-only host the call is a no-op and the stub remains active.
90
+
91
+ **On a CUDA box or dedicated-GPU Space:**
92
+
93
+ ```ini
94
+ # .env
95
+ LOCAL_INFERENCE=1
96
+ ```
97
+
98
+ Then pick **"Local GPU"** in the Lab's backend radio. On a ZeroGPU Space,
99
+ `SPACES_ZERO_GPU` is set automatically — no `.env` change needed, just select the
100
+ backend in the UI. See
101
+ [`docs/architecture/model-routing.md`](architecture/model-routing.md) for the
102
+ full model list and per-tier config syntax (`local:<repo_id>`).
103
+
104
  ### Per-profile overrides
105
 
106
  Highest priority. Override the model string bound to any profile — the cheapest
src/models/local_provider.py CHANGED
@@ -14,14 +14,16 @@ one decorated ``_generate`` covers every flavour:
14
  * **Dedicated GPU / local CUDA** — the decorator is a passthrough; the model runs on
15
  the persistent GPU.
16
 
17
- **Why the model loads in the parent, not inside ``@spaces.GPU``.** On ZeroGPU each call
18
- forks a GPU worker that inherits the parent's already-loaded model (CUDA is *emulated* in
19
- the parent, materialised on the real GPU inside the call). Loading inside the decorated
20
- function would reload the weights on every call the HF docs call this out as
21
- "significantly less efficient". So :meth:`complete` warms a module-level cache in the
22
- parent first (lazily, on first use — never at app boot), and the decorated ``_generate``
23
- only runs the forward pass. On a dedicated GPU the same cache simply keeps the model
24
- resident across calls.
 
 
25
 
26
  Heavy imports (``torch`` / ``transformers``) are lazy — confined to the functions that
27
  need them — so importing this module never initialises CUDA (which would trip ZeroGPU's
@@ -49,18 +51,20 @@ _LOADED: dict[str, tuple] = {}
49
 
50
 
51
  def _ensure_loaded(repo_id: str, trust_remote_code: bool) -> tuple:
52
- """Load (once, cached) the tokenizer + model for *repo_id*, placed on CUDA.
53
-
54
- Called from :meth:`LocalTransformersProvider.complete` in the parent process so the
55
- placement happens under ZeroGPU's CUDA emulation (or directly on a dedicated GPU) and
56
- every later ``@spaces.GPU`` call inherits it. ``dtype="auto"`` lets transformers pick
57
- the weights' native precision; we fall back to the legacy ``torch_dtype`` kwarg name
58
- for older transformers, and to CPU when no CUDA is present (the gate normally prevents
59
- that, but it keeps the path honest).
 
 
 
60
  """
61
  if repo_id in _LOADED:
62
  return _LOADED[repo_id]
63
- import torch
64
  from transformers import AutoModelForCausalLM, AutoTokenizer
65
 
66
  tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=trust_remote_code)
@@ -68,8 +72,6 @@ def _ensure_loaded(repo_id: str, trust_remote_code: bool) -> tuple:
68
  model = AutoModelForCausalLM.from_pretrained(repo_id, dtype="auto", trust_remote_code=trust_remote_code)
69
  except TypeError: # pragma: no cover - older transformers use the torch_dtype kwarg name
70
  model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype="auto", trust_remote_code=trust_remote_code)
71
- if torch.cuda.is_available():
72
- model = model.to("cuda")
73
  model.eval()
74
  _LOADED[repo_id] = (tokenizer, model)
75
  return _LOADED[repo_id]
@@ -90,13 +92,19 @@ def _generate(repo_id, trust_remote_code, system, prompt, max_new_tokens, temper
90
  """Run one chat completion on the GPU; return ``(text, prompt_tokens, completion_tokens)``.
91
 
92
  Module-level and decorated so ZeroGPU registers it and grants a GPU for the call. The
93
- model is fetched from the parent-warmed cache (a hit — never a reload here); only the
94
- input tensors are built and moved to the device inside the GPU window.
 
 
 
 
95
  """
96
  import torch
97
 
98
  tokenizer, model = _ensure_loaded(repo_id, trust_remote_code)
99
- device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
100
  messages = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
101
  inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(device)
102
  do_sample = temperature and float(temperature) > 0
 
14
  * **Dedicated GPU / local CUDA** — the decorator is a passthrough; the model runs on
15
  the persistent GPU.
16
 
17
+ **Where the weights load vs. where CUDA is touched.** ZeroGPU grants a real GPU *only*
18
+ for the duration of a ``@spaces.GPU`` call (each call forks a worker); the parent process
19
+ never gets one, and any low-level CUDA init outside such a call including a lazy
20
+ ``.to("cuda")`` at request time, which ZeroGPU's startup hook does not capture trips the
21
+ fork guard and kills the process. So the split is: :meth:`complete` warms a module-level
22
+ **CPU** cache in the parent first (lazily, on first use — never at app boot, never on
23
+ CUDA), and the decorated ``_generate`` moves those weights onto the GPU **inside the
24
+ granted window** and runs the forward pass. The forked worker inherits the parent's CPU
25
+ weights, so it pays a host→device copy, not a disk reload; on a dedicated GPU the cached
26
+ module simply stays resident across calls (the move is a no-op after the first).
27
 
28
  Heavy imports (``torch`` / ``transformers``) are lazy — confined to the functions that
29
  need them — so importing this module never initialises CUDA (which would trip ZeroGPU's
 
51
 
52
 
53
  def _ensure_loaded(repo_id: str, trust_remote_code: bool) -> tuple:
54
+ """Load (once, cached) the tokenizer + model for *repo_id* **on CPU**.
55
+
56
+ Called from :meth:`LocalTransformersProvider.complete` in the parent process to warm
57
+ the weights in host RAM. It deliberately **never touches CUDA**: under ZeroGPU the
58
+ parent process gets no GPU, and any low-level CUDA init outside a ``@spaces.GPU`` call
59
+ (a lazy ``.to("cuda")`` at request time is not captured by ZeroGPU's startup hook)
60
+ trips the fork guard and kills the process. The CPU→GPU transfer happens inside the
61
+ decorated :func:`_generate`, where a real GPU is granted; the forked worker inherits
62
+ these CPU weights, so it pays a device copy, not a disk reload. ``dtype="auto"`` lets
63
+ transformers pick the weights' native precision (fallback to the legacy ``torch_dtype``
64
+ kwarg name on older transformers).
65
  """
66
  if repo_id in _LOADED:
67
  return _LOADED[repo_id]
 
68
  from transformers import AutoModelForCausalLM, AutoTokenizer
69
 
70
  tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=trust_remote_code)
 
72
  model = AutoModelForCausalLM.from_pretrained(repo_id, dtype="auto", trust_remote_code=trust_remote_code)
73
  except TypeError: # pragma: no cover - older transformers use the torch_dtype kwarg name
74
  model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype="auto", trust_remote_code=trust_remote_code)
 
 
75
  model.eval()
76
  _LOADED[repo_id] = (tokenizer, model)
77
  return _LOADED[repo_id]
 
92
  """Run one chat completion on the GPU; return ``(text, prompt_tokens, completion_tokens)``.
93
 
94
  Module-level and decorated so ZeroGPU registers it and grants a GPU for the call. The
95
+ model weights are fetched from the parent-warmed CPU cache (a hit — never a disk reload
96
+ here) and moved onto the device **inside this GPU window** — the only place ZeroGPU
97
+ permits CUDA init. On ZeroGPU the forked worker inherits CPU weights and pays one
98
+ device copy per call; on a dedicated GPU the cached module is already resident after
99
+ the first call, so the move is a no-op. Input tensors are built and placed on the same
100
+ device.
101
  """
102
  import torch
103
 
104
  tokenizer, model = _ensure_loaded(repo_id, trust_remote_code)
105
+ if torch.cuda.is_available():
106
+ model = model.to("cuda")
107
+ device = next(model.parameters()).device
108
  messages = [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
109
  inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(device)
110
  do_sample = temperature and float(temperature) > 0
src/ui/fishbowl/assets/styles.css CHANGED
@@ -939,7 +939,7 @@ footer { display: none !important; }
939
  .fishbowl.fb-split .splitview { padding: 6px 0 0; max-width: 100%; }
940
 
941
  /* ---- FEED pane: read like a live transcript, comfortable scroll ---- */
942
- .fishbowl.fb-feed .feed { padding: 4px 4px 8px; max-height: 70vh; }
943
  .fishbowl .fe.narr p { font-size: 13px; }
944
 
945
  /* ---- NARRATOR feed entrance is already animated (feIn); add a hover tint ---- */
 
939
  .fishbowl.fb-split .splitview { padding: 6px 0 0; max-width: 100%; }
940
 
941
  /* ---- FEED pane: read like a live transcript, comfortable scroll ---- */
942
+ .fishbowl.fb-feed .feed { padding: 4px 4px 8px; max-height: 50vh; }
943
  .fishbowl .fe.narr p { font-size: 13px; }
944
 
945
  /* ---- NARRATOR feed entrance is already animated (feIn); add a hover tint ---- */
tests/test_local_backend.py CHANGED
@@ -155,3 +155,45 @@ def test_provider_resolves_trust_remote_code_from_catalogue():
155
  assert LocalTransformersProvider(model="Qwen/Qwen2.5-3B-Instruct")._trust_remote_code() is False
156
  # An off-catalogue repo defaults to the safe choice.
157
  assert LocalTransformersProvider(model="some/random-repo")._trust_remote_code() is False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
  assert LocalTransformersProvider(model="Qwen/Qwen2.5-3B-Instruct")._trust_remote_code() is False
156
  # An off-catalogue repo defaults to the safe choice.
157
  assert LocalTransformersProvider(model="some/random-repo")._trust_remote_code() is False
158
+
159
+
160
+ # ── ZeroGPU contract: CUDA only inside @spaces.GPU, never in the parent ───────────────
161
+ # Regression guard for the production crash "Low-level CUDA init (torch._C._cuda_init)
162
+ # reached … ZeroGPU's emulation did not intercept": the parent process gets no GPU, so any
163
+ # CUDA placement outside the @spaces.GPU window (a lazy .to("cuda") at request time) kills
164
+ # the worker. The forward pass can only be exercised with a GPU + weights (integration),
165
+ # so we pin the *structural* invariant — where CUDA may be touched — by source contract.
166
+
167
+
168
+ def test_parent_loader_never_initialises_cuda():
169
+ import ast
170
+ import inspect
171
+
172
+ from src.models import local_provider
173
+
174
+ # _ensure_loaded runs in the parent (warm CPU cache, inherited by the fork). It must
175
+ # not perform any CUDA operation — placement happens later, inside the decorated
176
+ # function. Check the executable body with the docstring stripped (the docstring
177
+ # explains the invariant in prose, so it legitimately mentions CUDA); the dangerous
178
+ # ops are the device move and any torch.cuda.* call.
179
+ fn = ast.parse(inspect.getsource(local_provider._ensure_loaded)).body[0]
180
+ if ast.get_docstring(fn):
181
+ fn.body = fn.body[1:]
182
+ code = ast.unparse(fn)
183
+ assert 'to("cuda")' not in code
184
+ assert "torch.cuda" not in code
185
+ assert ".cuda(" not in code
186
+
187
+
188
+ def test_gpu_transfer_lives_inside_the_spaces_gpu_function():
189
+ from pathlib import Path
190
+
191
+ from src.models import local_provider
192
+
193
+ # _generate is wrapped by @spaces.GPU, so read the module source and isolate its block.
194
+ module_src = Path(local_provider.__file__).read_text()
195
+ gen_block = module_src.split("def _generate(", 1)[1].split("\ndef ", 1)[0]
196
+ # The CPU→GPU move is here (the one place ZeroGPU grants a device)…
197
+ assert '.to("cuda")' in gen_block
198
+ # …and the function carries the decorator the platform registers.
199
+ assert "@spaces.GPU" in module_src.split("def _generate(", 1)[0].rsplit("\n\n", 1)[-1]