"""Follow one run of the toy through the lesson's six steps. Run it from this folder: python3 trace.py

It uses pd_sim.py's own functions and adds no rules of its own. It prints one line for each step, then one line for the limit: how the split does at each link speed on offer.
"""
from pd_sim import CACHE_BYTES_PER_TOKEN, LINKS_GBPS, READ_S_PER_TOKEN, STEP_S, compare, handover_seconds, make_requests, measure, shared, split

RPS, PROMPT, REPLY = 4, 2000, 200   # requests per second, tokens in a typical prompt, words in a typical reply
LINK_GBPS = 25      # the link between the two GPUs, in 10**9 bytes per second

ms = lambda seconds: f"{seconds * 1000:.1f}"

requests = make_requests(RPS, PROMPT, REPLY)
print(f"step 1  {len(requests)} requests in {requests[-1].arrival:.1f} s, {RPS} a second; prompts of {', '.join(str(p) for p in sorted({r.prompt for r in requests}))} tokens, "
      f"replies of {', '.join(str(w) for w in sorted({r.reply for r in requests}))} words")

read = PROMPT * READ_S_PER_TOKEN
print(f"step 2  shared: a {PROMPT}-token prompt takes {ms(read)} ms to read and no reply on that GPU moves meanwhile; a word due {ms(STEP_S)} ms after the last comes {ms(read + STEP_S)} ms after it")

both = measure(requests, shared(requests))
print(f"step 3  shared: first word p50 {ms(both['first_p50'])} ms, p99 {ms(both['first_p99'])} ms; gap between words p50 {ms(both['gap_p50'])} ms, p99 {ms(both['gap_p99'])} ms")

load = sum(r.prompt for r in requests) * READ_S_PER_TOKEN / requests[-1].arrival
print(f"step 4  split: the reader GPU has {load:.2f} s of prompts to read each second; the writer GPU only writes, {ms(STEP_S)} ms a step")

cache = PROMPT * CACHE_BYTES_PER_TOKEN
print(f"step 5  hand-over: {PROMPT} tokens x {CACHE_BYTES_PER_TOKEN} bytes = {cache / 1e6:.1f} MB; at {LINK_GBPS} GB/s that takes {ms(handover_seconds(PROMPT, LINK_GBPS))} ms, "
      f"against {ms(read)} ms to read (equal at {CACHE_BYTES_PER_TOKEN / READ_S_PER_TOKEN / 1e9:.2f} GB/s)")

apart = measure(requests, split(requests, LINK_GBPS))
print(f"step 6  split: first word p50 {ms(apart['first_p50'])} ms, p99 {ms(apart['first_p99'])} ms; gap p50 {ms(apart['gap_p50'])} ms, p99 {ms(apart['gap_p99'])} ms; "
      f"goodput {apart['goodput']:.1%} against {both['goodput']:.1%} shared")

by_link = ", ".join(f"{link} {compare(RPS, PROMPT, REPLY, link)['split']['goodput']:.1%}" for link in LINKS_GBPS)
print(f"limit   goodput of the split by link speed, GB/s: {by_link}")