# ssd_bench.py - CPU + SSD-only inference: codon table on disk, mmap, read 2 rows/token.
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, export_codons_to_disk, make_disk_reader, load_tok

DISK = '/root/dna/codons.u8'
tok = load_tok()
m, cfg = load_model('cpu')

# 1) export the codon table to SSD, then drop the in-RAM copy -> force disk reads
path, shape = export_codons_to_disk(m, DISK)
bytes_total = os.path.getsize(DISK)
read_rows, ncod = make_disk_reader(m, DISK)
m.codons = None  # free the ~512MB RAM buffer; table now lives only on SSD
import gc; gc.collect()

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_bytes_per_token': 2 * ncod, 'table_bytes_on_disk': bytes_total,
        'io': f'2 codon rows/token from mmap uint8 SSD table ({2*ncod} B/token)', 'gpu': False}
json.dump(meta, open('/root/dna/ssd_bench.json', 'w'), indent=2)
print('SSD_BENCH_DONE', json.dumps(meta), flush=True)