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:
@@ -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
Reference in New Issue
Block a user