Add web topology viewer with hot path visualization

- FastAPI backend (server.py) with REST API + WebSocket for event streaming
- SVG-based topology viewer (index.html) with SIP/CUBE/PE drill-down views
- Event logging infrastructure (event_log.py) generating events from real
  probe cases (H2D/D2H/PE-DMA) and bench workloads (QKV GEMM single/multi-PE)
- Timeline replay engine with play/pause, speed control, and scrubbing
- Workload selector dropdown grouped by category (Probe/Bench)
- CLI entry points: kernbench web, ./kernbench wrapper scripts

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-19 03:19:19 -07:00
parent fc6abbc8ee
commit dcbc41571f
8 changed files with 2904 additions and 1 deletions
View File
+126
View File
@@ -0,0 +1,126 @@
"""Phase B: FastAPI backend with WebSocket streaming for KernBench viewer."""
import asyncio
import json
import pathlib
import webbrowser
import threading
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
import uvicorn
STATIC_DIR = pathlib.Path(__file__).parent / "static"
DEFAULT_PORT = 8765
app = FastAPI(title="KernBench Viewer")
# ── REST: workloads + events ───────────────────────────────────────
@app.get("/api/workloads")
async def list_workloads():
from kernbench.sim_engine.event_log import get_workload_list
return JSONResponse(content=get_workload_list())
@app.get("/api/events/demo")
async def get_demo_events():
from kernbench.sim_engine.event_log import generate_demo_events
return JSONResponse(content=generate_demo_events())
@app.get("/api/events/{workload_id}")
async def get_workload_events(workload_id: str):
from kernbench.sim_engine.event_log import generate_workload_events
try:
return JSONResponse(content=generate_workload_events(workload_id))
except ValueError as e:
return JSONResponse(content={"error": str(e)}, status_code=404)
# ── WebSocket: stream events with timing ───────────────────────────
@app.websocket("/ws")
async def ws_endpoint(ws: WebSocket):
await ws.accept()
try:
while True:
msg = await ws.receive_text()
data = json.loads(msg)
cmd = data.get("cmd")
if cmd == "load_demo" or cmd == "load_workload":
from kernbench.sim_engine.event_log import generate_workload_events
wid = data.get("workload_id", "probe-h2d")
events = generate_workload_events(wid)
await ws.send_json({"type": "events", "events": events})
elif cmd == "stream_demo" or cmd == "stream_workload":
from kernbench.sim_engine.event_log import generate_workload_events
wid = data.get("workload_id", "probe-h2d")
events = generate_workload_events(wid)
speed = data.get("speed", 1.0) # ns → real-time multiplier
# Send metadata first
duration = max(e["t_ns"] for e in events) if events else 0
await ws.send_json({
"type": "stream_start",
"total_events": len(events),
"duration_ns": duration,
})
prev_t = 0.0
for ev in events:
# Wait proportional to simulation time gap
dt = ev["t_ns"] - prev_t
if dt > 0 and speed > 0:
# Scale: 1 ns sim → (1/speed) ms real
await asyncio.sleep(dt / (speed * 1000.0))
await ws.send_json({"type": "event", **ev})
prev_t = ev["t_ns"]
await ws.send_json({"type": "stream_end"})
elif cmd == "ping":
await ws.send_json({"type": "pong"})
except WebSocketDisconnect:
pass
except Exception:
pass
# ── Static files (must be last) ────────────────────────────────────
@app.get("/")
async def index():
return FileResponse(STATIC_DIR / "index.html")
app.mount("/", StaticFiles(directory=str(STATIC_DIR)), name="static")
# ── Entry point (called from CLI) ──────────────────────────────────
def serve(port: int = DEFAULT_PORT, open_browser: bool = True) -> None:
import socket
# Check port availability before starting
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try:
s.bind(("127.0.0.1", port))
except OSError:
print(f"ERROR: Port {port} is already in use.")
print(f" Try: kernbench web --port {port + 1}")
print(f" Or kill the existing process using port {port}")
return
url = f"http://127.0.0.1:{port}"
print(f"KernBench Viewer: {url}")
if open_browser:
threading.Timer(0.5, webbrowser.open, args=(url,)).start()
uvicorn.run(app, host="127.0.0.1", port=port, log_level="warning")
if __name__ == "__main__":
serve()
File diff suppressed because it is too large Load Diff