Spaces:
Running on Zero
Running on Zero
AndrianBalanescu commited on
Commit ·
54cef29
1
Parent(s): 815f934
fix(telemetry): detect cold start inside ZeroGPU worker body
Browse filesZeroGPU 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.
- app.py +13 -12
- 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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):
|