"""SSD read speed for expert streaming: random experts from experts.bin, direct (unbuffered) reads, 1 to 16 threads.
usage: python bench_ssd.py MODEL_DIR [SECONDS_PER_SETTING]"""
import json, os, random, sys, time
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from onw.xio import DirectReader


def main():
    d = sys.argv[1]
    secs = float(sys.argv[2]) if len(sys.argv) > 2 else 4
    idx = json.load(open(os.path.join(d, "experts.bin.json")))
    ents = [(w[0], w[1] + s[1]) for ex in idx["layers"].values() for (w, s) in ex.values()]    # w and s are adjacent
    size = max(n for _, n in ents)
    r = DirectReader(os.path.join(d, "experts.bin"), size)
    print(f"{len(ents)} experts, {size / 2**20:.2f} MB each, direct I/O: {r.direct}")
    bufs = {}

    def one(e):
        import threading
        b = bufs.setdefault(threading.get_ident(), np.empty(size, np.uint8))
        r.read_into(b[:e[1]], e[0], e[1])
        return e[1]

    for threads in (1, 2, 4, 8, 16):
        with ThreadPoolExecutor(threads) as pool:
            t0, total, lat = time.time(), 0, []
            while time.time() - t0 < secs:
                batch = random.sample(ents, threads * 4)
                t1 = time.perf_counter()
                total += sum(pool.map(one, batch))
                lat.append((time.perf_counter() - t1) / 4)
            dt = time.time() - t0
        print(f"threads {threads:2d}: {total / dt / 2**30:5.2f} GB/s, {total / dt / size:6.0f} experts/s, "
              f"{np.median(lat) * 1000:5.2f} ms per batch of {threads}", flush=True)


if __name__ == "__main__":
    main()