k1_dev/profile_host.py

"""Directly profile the HOST/CPU per-step overhead of the frontier stack on
A10G during steady single-stream spec-decode — the ~1.6ms gap nobody has
measured directly (only inferred from ablations). py-spy samples the whole
process tree (API server + EngineCore) while a load thread keeps decode busy.
Output: collapsed-stack self-time leaderboard of host functions -> where the
gap actually goes (scheduler loop / metadata build / detok / sampler / accept)."""
import json, os, subprocess, sys, threading, time, urllib.request, collections

SUB="/tmp/sub"; base="http://127.0.0.1:8000"
RUN=[True]

def boot():
    m=json.load(open(f"{SUB}/manifest.json")); env=dict(os.environ)
    for k,v in m.get("env",{}).items(): env[k]=v
    env["PORT"]="8000"
    log=open("/tmp/srv.log","w")
    p=subprocess.Popen([sys.executable,f"{SUB}/serve.py"],cwd=SUB,env=env,stdout=log,stderr=subprocess.STDOUT)
    dl=time.time()+1000
    while time.time()<dl:
        if p.poll() is not None:
            print(f"server EXITED {p.returncode}")
            for ln in open('/tmp/srv.log',errors='replace').readlines()[-20:]: print("[srv]",ln.rstrip()[:200])
            return None
        try: urllib.request.urlopen(base+"/v1/models",timeout=2); return p
        except Exception: time.sleep(5)
    return None

def load_thread():
    while RUN[0]:
        body=json.dumps({"model":"gemma-4-e4b-it","messages":[{"role":"user","content":"Write a long detailed technical essay on memory hierarchies and explain each level."}],"max_tokens":512,"temperature":0.0,"stream":False}).encode()
        try:
            urllib.request.urlopen(urllib.request.Request(base+"/v1/chat/completions",data=body,headers={"content-type":"application/json"}),timeout=180).read()
        except Exception: time.sleep(1)

def main():
    p=boot()
    if p is None: return
    # find EngineCore + API server PIDs (children of serve.py launcher)
    print("[prof] server ready; starting load + warmup")
    for _ in range(2): load_thread.__wrapped__ if False else None
    t=threading.Thread(target=load_thread,daemon=True); t.start()
    time.sleep(40)  # warmup: let onegraph capture + steady state settle
    # py-spy record the whole tree (raw/folded stacks) for 120s
    print("[prof] attaching py-spy --subprocesses for 120s")
    rec=subprocess.run(["py-spy","record","--subprocesses","--pid",str(p.pid),
                        "--duration","120","--rate","250","--format","raw","--output","/tmp/prof.folded","--nonblocking"],
                       capture_output=True,text=True)
    print("[prof] py-spy rc",rec.returncode, (rec.stderr or "")[-400:])
    RUN[0]=False
    # aggregate self-time (leaf frame) from folded stacks
    self_t=collections.Counter(); total=0
    try:
        for line in open("/tmp/prof.folded",errors="replace"):
            line=line.strip()
            if not line: continue
            stack,_,cnt=line.rpartition(" ")
            try: c=int(cnt)
            except: continue
            leaf=stack.split(";")[-1] if stack else "?"
            self_t[leaf]+=c; total+=c
        print(f"[prof] total samples={total}")
        print("[prof] TOP host self-time frames (leaf):")
        for fn,c in self_t.most_common(40):
            print(f"[prof] {100*c/max(1,total):5.1f}%  {fn[:110]}")
    except Exception as e:
        print("[prof] aggregate err",e)
        # fallback: raw head
        try:
            for ln in open("/tmp/prof.folded").readlines()[:30]: print("[raw]",ln.rstrip()[:160])
        except Exception: pass
    p.terminate()

if __name__=="__main__": main()