Restored checkpoint predictions' usage counted in totals (was 0 for fully-restored benches); token stats now correct on resume

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
sora 2026-09-11 08:27:55 +00:00
parent b8ed18771a
commit 3a4be917e6

View File

@ -443,6 +443,20 @@ async def generate_predictions(
preds_by_key[keys[i]] = pred preds_by_key[keys[i]] = pred
preds = [preds_by_key[k] for k in keys] preds = [preds_by_key[k] for k in keys]
usages = [p.get('usage', {}) for p in preds] usages = [p.get('usage', {}) for p in preds]
# include RESTORED predictions' usage (they carry it in the ckpt);
# previously only fresh generations counted -> restored benches showed 0
fresh_keys = {keys[i] for i, _ in pending}
for k, p in preds_by_key.items():
if k in fresh_keys:
continue # already counted via one()'s total_usage updates
u = p.get('usage') or {}
if not any(u.get(kk) for kk in ('input_tokens', 'output_tokens')):
continue
total_usage = total_usage + Usage(
input_tokens=int(u.get('input_tokens', 0) or 0),
output_tokens=int(u.get('output_tokens', 0) or 0),
total_tokens=int(u.get('total_tokens', 0) or 0),
latency_s=float(u.get('latency_s', 0) or 0))
if status_callback and pending: if status_callback and pending:
status_callback(f'Generation complete: {len(preds)} responses collected') status_callback(f'Generation complete: {len(preds)} responses collected')
return preds, usages, total_usage return preds, usages, total_usage