atakan Claude Sonnet 5 commited on
Commit
8edeb7f
·
1 Parent(s): 01c9d8d

fix: Don't pass the model-holding agent as a @spaces.GPU argument

Browse files

Root cause of "_share_cuda_: only available on CUDA": ZeroGPU's docs
example always accesses the model via closure/module scope from
inside the decorated function -- never as a parameter. We were passing
the whole ControlAIAgent (holding the loaded CUDA model) as an
argument on every call, which appears to force ZeroGPU's argument
marshaling to try to share the model's CUDA tensors across its process
boundary and fail. Both GPU-wrapped functions now call get_agent()
internally instead, matching the documented pattern exactly.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Files changed (1) hide show
  1. app.py +6 -9
app.py CHANGED
@@ -83,13 +83,13 @@ def get_agent() -> ControlAIAgent:
83
  # releases it afterward). Outside of a ZeroGPU Space this decorator is a
84
  # harmless no-op, so it's safe to always wrap these.
85
  @spaces.GPU(duration=300)
86
- def _run_stream_on_gpu(agent: ControlAIAgent, message: str, history: list[dict[str, str]]):
87
- yield from agent.run_stream(message, history=history)
88
 
89
 
90
  @spaces.GPU(duration=300)
91
- def _run_on_gpu(agent: ControlAIAgent, message: str, history: list[dict[str, str]]):
92
- return agent.run(message, history=history, verbose=False)
93
 
94
 
95
  # ZeroGPU's startup check statically looks for a @spaces.GPU function wired
@@ -202,12 +202,10 @@ async def chat_stream_endpoint(req: ChatRequest):
202
  if not req.message.strip():
203
  raise HTTPException(status_code=400, detail="Message cannot be empty")
204
 
205
- agent = get_agent()
206
-
207
  def event_generator():
208
  try:
209
  with inference_lock:
210
- for event in _run_stream_on_gpu(agent, req.message.strip(), req.history):
211
  yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
212
  except Exception as exc:
213
  yield f"data: {json.dumps({'type': 'error', 'error': str(exc)}, ensure_ascii=False)}\n\n"
@@ -230,9 +228,8 @@ async def chat_endpoint(req: ChatRequest) -> ChatResponse:
230
 
231
  t0 = time.time()
232
  try:
233
- agent = get_agent()
234
  with inference_lock:
235
- result = _run_on_gpu(agent, req.message.strip(), req.history)
236
  elapsed = time.time() - t0
237
 
238
  # Collect tool traces
 
83
  # releases it afterward). Outside of a ZeroGPU Space this decorator is a
84
  # harmless no-op, so it's safe to always wrap these.
85
  @spaces.GPU(duration=300)
86
+ def _run_stream_on_gpu(message: str, history: list[dict[str, str]]):
87
+ yield from get_agent().run_stream(message, history=history)
88
 
89
 
90
  @spaces.GPU(duration=300)
91
+ def _run_on_gpu(message: str, history: list[dict[str, str]]):
92
+ return get_agent().run(message, history=history, verbose=False)
93
 
94
 
95
  # ZeroGPU's startup check statically looks for a @spaces.GPU function wired
 
202
  if not req.message.strip():
203
  raise HTTPException(status_code=400, detail="Message cannot be empty")
204
 
 
 
205
  def event_generator():
206
  try:
207
  with inference_lock:
208
+ for event in _run_stream_on_gpu(req.message.strip(), req.history):
209
  yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
210
  except Exception as exc:
211
  yield f"data: {json.dumps({'type': 'error', 'error': str(exc)}, ensure_ascii=False)}\n\n"
 
228
 
229
  t0 = time.time()
230
  try:
 
231
  with inference_lock:
232
+ result = _run_on_gpu(req.message.strip(), req.history)
233
  elapsed = time.time() - t0
234
 
235
  # Collect tool traces