AndrianBalanescu commited on
Commit
54cef29
·
1 Parent(s): 815f934

fix(telemetry): detect cold start inside ZeroGPU worker body

Browse files

ZeroGPU executes @spaces.GPU bodies in a worker process; get_model sets
_llm/_loaded_file there and the globals never reach the web process, so
the outer residency pre-check always saw no model and the -1.0 sentinel
leaked into records (summary filtered it to 0.0). Residency is now
evaluated inside the worker and returned with the result, which does
propagate back. Regression test pins the worker-side semantics.

Files changed (2) hide show
  1. app.py +13 -12
  2. tests/unit/test_gpu_usage_tracker.py +30 -0
app.py CHANGED
@@ -404,34 +404,35 @@ def gpu_tracked_call(kind: str, fn, *args, model: str = "", **kwargs):
404
  """
405
  t_submit = time.time()
406
  call_id = GPU_USAGE.lease_start(kind)
407
- # Cold-start detection: the model was not resident before this lease, so
408
- # this lease includes the full load cost (first lease after boot or after
409
- # a model switch freed it).
410
- was_resident = _llm is not None and _loaded_file == model
411
- cold_start_s = 0.0 if was_resident else -1.0 # -1 => resolved after body
412
 
413
  @spaces.GPU(duration=120)
414
  def _tracked_inner():
 
 
 
 
415
  t_body_start = time.time()
416
  try:
417
  result = fn(*args, **kwargs)
418
  error = ""
419
  except Exception as exc:
420
  result, error = None, str(exc)
421
- return result, error, t_body_start
 
 
 
 
 
 
 
422
 
423
  try:
424
- result, error, t_body_start = _tracked_inner()
425
  except Exception as exc:
426
  # Lease-level failure (e.g. 429 quota exhaustion): body never ran.
427
  GPU_USAGE.lease_end(call_id, model=model, kind=kind, error=str(exc))
428
  raise
429
  queue_wait = max(0.0, t_body_start - t_submit)
430
- if cold_start_s < 0 and _llm is not None and _loaded_file == model:
431
- # Body loaded the model inside this lease. The load share is not
432
- # directly observable, so record the whole post-body span as the
433
- # cold-start upper bound.
434
- cold_start_s = max(0.0, time.time() - t_body_start)
435
  GPU_USAGE.lease_end(
436
  call_id, model=model, kind=kind,
437
  queue_wait_s=queue_wait, cold_start_s=cold_start_s, error=error,
 
404
  """
405
  t_submit = time.time()
406
  call_id = GPU_USAGE.lease_start(kind)
 
 
 
 
 
407
 
408
  @spaces.GPU(duration=120)
409
  def _tracked_inner():
410
+ # Residency must be evaluated INSIDE the ZeroGPU worker: module
411
+ # globals set there (_llm/_loaded_file via get_model) never propagate
412
+ # back to the web process, so cold-start state is invisible outside.
413
+ resident_before = _llm is not None and _loaded_file == model
414
  t_body_start = time.time()
415
  try:
416
  result = fn(*args, **kwargs)
417
  error = ""
418
  except Exception as exc:
419
  result, error = None, str(exc)
420
+ if not resident_before and _llm is not None and _loaded_file == model:
421
+ # The body loaded the model inside this lease. The load share is
422
+ # not directly observable, so the whole post-body span is the
423
+ # cold-start upper bound.
424
+ cold_start_s = max(0.0, time.time() - t_body_start)
425
+ else:
426
+ cold_start_s = 0.0
427
+ return result, error, t_body_start, cold_start_s
428
 
429
  try:
430
+ result, error, t_body_start, cold_start_s = _tracked_inner()
431
  except Exception as exc:
432
  # Lease-level failure (e.g. 429 quota exhaustion): body never ran.
433
  GPU_USAGE.lease_end(call_id, model=model, kind=kind, error=str(exc))
434
  raise
435
  queue_wait = max(0.0, t_body_start - t_submit)
 
 
 
 
 
436
  GPU_USAGE.lease_end(
437
  call_id, model=model, kind=kind,
438
  queue_wait_s=queue_wait, cold_start_s=cold_start_s, error=error,
tests/unit/test_gpu_usage_tracker.py CHANGED
@@ -145,6 +145,36 @@ def test_gpu_tracked_call_records_usage(fresh_usage):
145
  assert snap["profile"]["total_calls"] == 1
146
 
147
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
148
  def test_gpu_tracked_call_records_error_path(fresh_usage):
149
  app = fresh_usage
150
  with pytest.raises(RuntimeError):
 
145
  assert snap["profile"]["total_calls"] == 1
146
 
147
 
148
+ def test_gpu_tracked_call_cold_start_worker_side(fresh_usage):
149
+ """ZeroGPU runs @spaces.GPU bodies in a worker process: get_model sets
150
+ _llm/_loaded_file there and those globals never reach the web process.
151
+ Cold-start detection must therefore run inside the body (returns DO
152
+ propagate back). Regression for the live finding where cold_start was
153
+ always 0.0 (the -1.0 sentinel was silently dropped by summary())."""
154
+ app = fresh_usage
155
+ app._llm = None
156
+ app._loaded_file = None
157
+
158
+ def loader(msgs, model_file, *a, **k):
159
+ # mimic get_model side effect inside the worker
160
+ app._llm = object()
161
+ app._loaded_file = model_file
162
+ return {"choices": [{"message": {"content": "ok"}}]}
163
+
164
+ model = "Qwen3.8-9B-Q8_0.gguf"
165
+ app.gpu_tracked_call("chat", loader, [], model, 0.7, 30, model=model)
166
+ cold_rec = app.GPU_USAGE._calls[-1]
167
+ assert cold_rec["cold_start_seconds"] >= 0.0
168
+ assert cold_rec["lease_seconds"] >= cold_rec["cold_start_seconds"]
169
+
170
+ def warm(msgs, model_file, *a, **k):
171
+ return {"choices": [{"message": {"content": "ok"}}]}
172
+
173
+ app.gpu_tracked_call("chat", warm, [], model, 0.7, 30, model=model)
174
+ warm_rec = app.GPU_USAGE._calls[-1]
175
+ assert warm_rec["cold_start_seconds"] == 0.0
176
+
177
+
178
  def test_gpu_tracked_call_records_error_path(fresh_usage):
179
  app = fresh_usage
180
  with pytest.raises(RuntimeError):