# cpu_bench.py - CPU-only inference throughput, codon table warm in RAM (compute-bound).
import os
os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ.setdefault('DNA_CK', '/root/dna/ckpt/base.pt')
import time, json, resource, numpy as np, torch
torch.set_num_threads(int(os.environ.get('THREADS', '12')))
from infer_dna import load_model, generate, make_ram_reader, load_tok

tok = load_tok()
m, cfg = load_model('cpu')
read_rows = make_ram_reader(m)

def run_once(prompt, n=64):
    t0 = time.time()
    _ = generate(m, tok, prompt, n=n, temp=0.0, device='cpu', read_rows=read_rows)
    return n / (time.time() - t0)

prompts = ["Hello there", "The weather today", "In the beginning", "Science is",
           "My favorite food", "The ocean is", "A long time ago", "Computers can",
           "The best way to", "People often say"]
run_once("warmup", 8)
res = []
for i, p in enumerate(prompts):
    ts = run_once(p, 64); res.append(ts)
    print(f'test {i+1}/10 prompt={p!r} tok_s={ts:.2f}', flush=True)
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
meta = {'tok_s': res, 'mean': float(np.mean(res)), 'std': float(np.std(res)),
        'rss_mb': rss, 'threads': torch.get_num_threads(),
        'io': 'codon table in RAM (warm, compute-bound)', 'gpu': False}
json.dump(meta, open('/root/dna/cpu_bench.json', 'w'), indent=2)
print('CPU_BENCH_DONE', json.dumps(meta), flush=True)