flozi00 commited on
Commit
ddc9f01
·
verified ·
1 Parent(s): 86d2239

zero-gpu: report tokens/s so the rewrite budget can be sized to the slice

Browse files
Files changed (2) hide show
  1. BUILD.json +1 -1
  2. traceforge/space/local_llm.py +14 -0
BUILD.json CHANGED
@@ -1,5 +1,5 @@
1
  {
2
- "built_at": "2026-10-29T14:00:39",
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(),