atakan Claude Opus 5 commited on
Commit
ac301c2
·
1 Parent(s): 3e5d998

fix: Warm the retrieval embedder inside ZeroGPU's import window too

Browse files

The 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>

Files changed (2) hide show
  1. CLAUDE.md +4 -1
  2. 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. `device_map` fails independently, by
 
 
 
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
- print("[space] agent ready")
 
 
 
 
 
 
 
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