Spaces:
Running on Zero
Running on Zero
atakan Claude Opus 5 commited on
Commit ·
ac301c2
1
Parent(s): 3e5d998
fix: Warm the retrieval embedder inside ZeroGPU's import window too
Browse filesThe import-scope fix worked for the language model -- the Space starts, ZeroGPU
packs 17.6G of tensors, /api/status reports ready with all 80,370 chunks -- but
the embedder is a second model and was still lazy. It loads on the first query,
which is a request, outside the import window, so it hit the identical
_cuda_init wall and every answer came back with "[agent] retrieval failed":
the demo would have run with no retrieval and no citations, quietly.
_build_agent_at_import now embeds one throwaway string, which is what actually
triggers get_embedder()'s load, so both models land on the GPU during import.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- CLAUDE.md +4 -1
- app_space.py +9 -1
CLAUDE.md
CHANGED
|
@@ -201,7 +201,10 @@ and attaches real hardware only inside a `@spaces.GPU` call; only CUDA operation
|
|
| 201 |
window are intercepted. `app.py` normally builds the model in FastAPI's `lifespan`, on a
|
| 202 |
`ThreadPoolExecutor` worker — after import, on another thread — and `.to("cuda")` there reaches real
|
| 203 |
CUDA init and raises. `app_space.py::_build_agent_at_import` builds the agent during import and
|
| 204 |
-
assigns `app._agent`, so `lifespan`'s `get_agent()` is a no-op.
|
|
|
|
|
|
|
|
|
|
| 205 |
routing transformers through `caching_allocator_warmup` and its direct
|
| 206 |
`torch.empty(..., device="cuda")`; bitsandbytes requires `device_map`, so **4-bit quantisation is
|
| 207 |
unavailable**, which is what forces a model small enough to carry in bf16: `Qwen/Qwen3-8B`, ~16GB
|
|
|
|
| 201 |
window are intercepted. `app.py` normally builds the model in FastAPI's `lifespan`, on a
|
| 202 |
`ThreadPoolExecutor` worker — after import, on another thread — and `.to("cuda")` there reaches real
|
| 203 |
CUDA init and raises. `app_space.py::_build_agent_at_import` builds the agent during import and
|
| 204 |
+
assigns `app._agent`, so `lifespan`'s `get_agent()` is a no-op. **The retrieval embedder is a second
|
| 205 |
+
model with the same problem** — it loads lazily on the first query, which is a request, outside the
|
| 206 |
+
window — so the same function embeds one throwaway string to force it onto the GPU during import.
|
| 207 |
+
Without that, the Space starts fine and every answer carries `[agent] retrieval failed`. `device_map` fails independently, by
|
| 208 |
routing transformers through `caching_allocator_warmup` and its direct
|
| 209 |
`torch.empty(..., device="cuda")`; bitsandbytes requires `device_map`, so **4-bit quantisation is
|
| 210 |
unavailable**, which is what forces a model small enough to carry in bf16: `Qwen/Qwen3-8B`, ~16GB
|
app_space.py
CHANGED
|
@@ -76,10 +76,18 @@ def _build_agent_at_import() -> None:
|
|
| 76 |
"""
|
| 77 |
from controlai_agent.agent import ControlAgent
|
| 78 |
from controlai_agent.engine_torch import TorchEngine
|
|
|
|
| 79 |
|
| 80 |
print("[space] building agent at import scope (ZeroGPU CUDA window)")
|
| 81 |
app_module._agent = ControlAgent(engine=TorchEngine())
|
| 82 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 83 |
|
| 84 |
|
| 85 |
_fetch_index() # the agent prewarms retrieval, so the index must precede it
|
|
|
|
| 76 |
"""
|
| 77 |
from controlai_agent.agent import ControlAgent
|
| 78 |
from controlai_agent.engine_torch import TorchEngine
|
| 79 |
+
from controlai_rag.embeddings import get_embedder
|
| 80 |
|
| 81 |
print("[space] building agent at import scope (ZeroGPU CUDA window)")
|
| 82 |
app_module._agent = ControlAgent(engine=TorchEngine())
|
| 83 |
+
|
| 84 |
+
# The retrieval embedder is a second model, and it loads lazily on first
|
| 85 |
+
# query -- which is a request, outside the import window, so it hit exactly
|
| 86 |
+
# the same _cuda_init wall and every answer came back with
|
| 87 |
+
# "[agent] retrieval failed". Force it onto the GPU here too. Embedding one
|
| 88 |
+
# throwaway string is what actually triggers the load.
|
| 89 |
+
get_embedder().encode_query("warmup")
|
| 90 |
+
print("[space] agent and embedder ready")
|
| 91 |
|
| 92 |
|
| 93 |
_fetch_index() # the agent prewarms retrieval, so the index must precede it
|