Spaces:
Running
Running
zero-gpu: report tokens/s so the rewrite budget can be sized to the slice
Browse files- BUILD.json +1 -1
- traceforge/space/local_llm.py +14 -0
BUILD.json
CHANGED
|
@@ -1,5 +1,5 @@
|
|
| 1 |
{
|
| 2 |
-
"built_at": "2026-10-29T14:
|
| 3 |
"commit": "c778551",
|
| 4 |
"files": 53,
|
| 5 |
"with_store": false,
|
|
|
|
| 1 |
{
|
| 2 |
+
"built_at": "2026-10-29T14:15:09",
|
| 3 |
"commit": "c778551",
|
| 4 |
"files": 53,
|
| 5 |
"with_store": false,
|
traceforge/space/local_llm.py
CHANGED
|
@@ -66,6 +66,9 @@ LAST_RESPONSE: Optional[str] = None
|
|
| 66 |
# needed. A request that arrives before the checkpoint is on disk must fail
|
| 67 |
# fast: downloading 48 GB inside a 120 s GPU slice would burn the whole slice.
|
| 68 |
PREFETCH: dict = {"state": "idle", "error": "", "source": ""}
|
|
|
|
|
|
|
|
|
|
| 69 |
|
| 70 |
# A Space can carry HTTP(S)_PROXY for outbound traffic; urllib would then
|
| 71 |
# route the loopback call through it and the local server would never see
|
|
@@ -303,6 +306,7 @@ def generate(messages: list[dict], *, max_tokens: int = 512,
|
|
| 303 |
enc = tok(prompt, return_tensors="pt", truncation=True,
|
| 304 |
max_length=LOCAL_MAX_INPUT_TOKENS).to(LOCAL_DEVICE)
|
| 305 |
n = max(1, min(int(max_tokens or 512), LOCAL_MAX_NEW_TOKENS))
|
|
|
|
| 306 |
with torch.inference_mode():
|
| 307 |
out = model.generate(
|
| 308 |
**enc,
|
|
@@ -315,6 +319,15 @@ def generate(messages: list[dict], *, max_tokens: int = 512,
|
|
| 315 |
repetition_penalty=1.05,
|
| 316 |
pad_token_id=tok.eos_token_id,
|
| 317 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 318 |
text = tok.decode(out[0][enc["input_ids"].shape[1]:],
|
| 319 |
skip_special_tokens=True)
|
| 320 |
return _strip_reasoning(text)
|
|
@@ -583,6 +596,7 @@ def health() -> dict:
|
|
| 583 |
"gpu_duration_s": GPU_DURATION_S,
|
| 584 |
"gpu_size": GPU_SIZE,
|
| 585 |
"prefetch": dict(PREFETCH),
|
|
|
|
| 586 |
"volume_read": _volume_read_mbps(),
|
| 587 |
"server": _server_url or "",
|
| 588 |
"ping": _ping(),
|
|
|
|
| 66 |
# needed. A request that arrives before the checkpoint is on disk must fail
|
| 67 |
# fast: downloading 48 GB inside a 120 s GPU slice would burn the whole slice.
|
| 68 |
PREFETCH: dict = {"state": "idle", "error": "", "source": ""}
|
| 69 |
+
# Throughput of the last generation: the GPU slice is 120 s, so tokens/s
|
| 70 |
+
# decides whether a rewrite can finish inside one call.
|
| 71 |
+
LAST_GEN: dict = {}
|
| 72 |
|
| 73 |
# A Space can carry HTTP(S)_PROXY for outbound traffic; urllib would then
|
| 74 |
# route the loopback call through it and the local server would never see
|
|
|
|
| 306 |
enc = tok(prompt, return_tensors="pt", truncation=True,
|
| 307 |
max_length=LOCAL_MAX_INPUT_TOKENS).to(LOCAL_DEVICE)
|
| 308 |
n = max(1, min(int(max_tokens or 512), LOCAL_MAX_NEW_TOKENS))
|
| 309 |
+
t0 = time.time()
|
| 310 |
with torch.inference_mode():
|
| 311 |
out = model.generate(
|
| 312 |
**enc,
|
|
|
|
| 319 |
repetition_penalty=1.05,
|
| 320 |
pad_token_id=tok.eos_token_id,
|
| 321 |
)
|
| 322 |
+
elapsed = time.time() - t0
|
| 323 |
+
new_tokens = int(out.shape[1] - enc["input_ids"].shape[1])
|
| 324 |
+
LAST_GEN.update(
|
| 325 |
+
prompt_tokens=int(enc["input_ids"].shape[1]),
|
| 326 |
+
new_tokens=new_tokens,
|
| 327 |
+
seconds=round(elapsed, 1),
|
| 328 |
+
tok_s=round(new_tokens / elapsed, 1) if elapsed > 0 else 0.0,
|
| 329 |
+
capped=new_tokens >= n,
|
| 330 |
+
)
|
| 331 |
text = tok.decode(out[0][enc["input_ids"].shape[1]:],
|
| 332 |
skip_special_tokens=True)
|
| 333 |
return _strip_reasoning(text)
|
|
|
|
| 596 |
"gpu_duration_s": GPU_DURATION_S,
|
| 597 |
"gpu_size": GPU_SIZE,
|
| 598 |
"prefetch": dict(PREFETCH),
|
| 599 |
+
"last_gen": dict(LAST_GEN),
|
| 600 |
"volume_read": _volume_read_mbps(),
|
| 601 |
"server": _server_url or "",
|
| 602 |
"ping": _ping(),
|