Something went wrong. Try again.
OR-1 dataflow CPU sketch
Something went wrong. Try again.
15 kB · 381 lines
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382"""FastAPI server for the OR1 Monitor with bidirectional WebSocket protocol.
Exposes the SimulationBackend to a browser-based frontend via:- WebSocket /ws for bidirectional command/result communication- REST endpoints for non-WebSocket clients- Static file serving for the frontend"""
from __future__ import annotations
import asyncioimport jsonimport loggingfrom contextlib import asynccontextmanagerfrom pathlib import Path
from fastapi import FastAPI, WebSocket, WebSocketDisconnectfrom fastapi.staticfiles import StaticFiles
from cm_inst import MemOp, Portfrom asm.ir import IRGraphfrom monitor.backend import SimulationBackendfrom monitor.commands import ( ErrorResult, GraphLoaded, LoadCmd, ResetCmd, RunUntilCmd, SendCmd, InjectCmd, StepEventCmd, StepResult, StepTickCmd,)from monitor.graph_json import graph_loaded_json, graph_to_monitor_jsonfrom tokens import DyadToken, SMToken
logger = logging.getLogger(__name__)
class ConnectionManager: """Manages WebSocket connections and broadcasts to all connected clients."""
def __init__(self) -> None: self.active_connections: list[WebSocket] = []
async def connect(self, websocket: WebSocket) -> None: """Accept and register a new WebSocket connection.""" await websocket.accept() self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket) -> None: """Unregister a disconnected WebSocket.""" try: self.active_connections.remove(websocket) except ValueError: # Already removed pass
async def broadcast(self, message: dict) -> None: """Broadcast a message to all connected clients.""" disconnected: list[WebSocket] = [] for connection in self.active_connections: try: await connection.send_json(message) except Exception: disconnected.append(connection) for conn in disconnected: self.disconnect(conn)
def create_app(backend: SimulationBackend) -> FastAPI: """Create and configure the FastAPI application.
Args: backend: SimulationBackend instance (shared with REPL if both running)
Returns: Configured FastAPI application """ manager = ConnectionManager() current_graph_loaded: dict | None = None current_ir_graph = None
@asynccontextmanager async def lifespan(app: FastAPI): """Lifespan context manager for startup/shutdown.""" yield
app = FastAPI(lifespan=lifespan)
@app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket) -> None: """WebSocket endpoint for bidirectional command/result communication.
On connect, sends the current graph state if loaded. Receives JSON commands and dispatches to backend. Broadcasts results to all connected clients. """ nonlocal current_graph_loaded, current_ir_graph await manager.connect(websocket) try: # Send initial state if available if current_graph_loaded is not None: await websocket.send_json(current_graph_loaded)
# Receive loop: parse commands and dispatch to backend while True: data = await websocket.receive_text() cmd_json = json.loads(data) cmd_type = cmd_json.get("cmd")
# Dispatch command to backend in thread executor loop = asyncio.get_running_loop() result = None
try: if cmd_type == "load": source = cmd_json.get("source", "") result = await loop.run_in_executor( None, backend.send_command, LoadCmd(source=source), 10.0 )
elif cmd_type == "load_file": path = cmd_json.get("path", "") try: source = Path(path).read_text() result = await loop.run_in_executor( None, backend.send_command, LoadCmd(source=source), 10.0 ) except Exception as e: result = ErrorResult(message=f"Failed to read file: {e}")
elif cmd_type == "step_tick": result = await loop.run_in_executor( None, backend.send_command, StepTickCmd(), 10.0 )
elif cmd_type == "step_event": result = await loop.run_in_executor( None, backend.send_command, StepEventCmd(), 10.0 )
elif cmd_type == "run_until": until = cmd_json.get("until", 0.0) result = await loop.run_in_executor( None, backend.send_command, RunUntilCmd(until=until), 30.0 )
elif cmd_type == "inject": # Construct a CMToken from command data target = cmd_json.get("target", 0) offset = cmd_json.get("offset", 0) act_id = cmd_json.get("act_id", 0) data = cmd_json.get("data", 0) token = DyadToken( target=target, offset=offset, act_id=act_id, data=data, port=Port.L ) result = await loop.run_in_executor( None, backend.send_command, InjectCmd(token=token), 10.0 )
elif cmd_type == "send": # Same as inject but respects backpressure target = cmd_json.get("target", 0) offset = cmd_json.get("offset", 0) act_id = cmd_json.get("act_id", 0) data = cmd_json.get("data", 0) token = DyadToken( target=target, offset=offset, act_id=act_id, data=data, port=Port.L ) result = await loop.run_in_executor( None, backend.send_command, SendCmd(token=token), 10.0 )
elif cmd_type == "inject_sm": # Inject an SM token target = cmd_json.get("target", 0) addr = cmd_json.get("addr", 0) op_name = cmd_json.get("op", "READ") data = cmd_json.get("data", 0) op = MemOp[op_name] token = SMToken( target=target, addr=addr, op=op, flags=0, data=data, ret=None ) result = await loop.run_in_executor( None, backend.send_command, InjectCmd(token=token), 10.0 )
elif cmd_type == "send_sm": # Send an SM token with backpressure target = cmd_json.get("target", 0) addr = cmd_json.get("addr", 0) op_name = cmd_json.get("op", "READ") data = cmd_json.get("data", 0) op = MemOp[op_name] token = SMToken( target=target, addr=addr, op=op, flags=0, data=data, ret=None ) result = await loop.run_in_executor( None, backend.send_command, SendCmd(token=token), 10.0 )
elif cmd_type == "reset": reload = cmd_json.get("reload", False) result = await loop.run_in_executor( None, backend.send_command, ResetCmd(reload=reload), 10.0 )
else: result = ErrorResult(message=f"Unknown command: {cmd_type}")
except Exception as e: logger.exception("Error processing command %s", cmd_type) result = ErrorResult(message=str(e))
# Convert result to JSON and broadcast if result is not None: try: is_reset = cmd_type == "reset" response = _result_to_json(result, current_ir_graph, is_reset) except Exception as e: logger.exception("Error serializing result") response = { "type": "error", "message": f"Serialization error: {e}", } await manager.broadcast(response)
# Update current state if this was a GraphLoaded result if isinstance(result, GraphLoaded): current_ir_graph = result.ir_graph current_graph_loaded = graph_loaded_json( result.ir_graph, result.snapshot ) # If this was a reset without reload, clear cached state elif isinstance(result, StepResult) and cmd_type == "reset": reload = cmd_json.get("reload", False) if not reload: current_ir_graph = None current_graph_loaded = None
except WebSocketDisconnect: manager.disconnect(websocket) except Exception as e: logger.exception("WebSocket error") manager.disconnect(websocket)
@app.post("/load") async def load_endpoint(request_data: dict) -> dict: """REST endpoint for loading a program.
Body: {"source": "...dfasm source..."} """ nonlocal current_graph_loaded, current_ir_graph source = request_data.get("source", "") loop = asyncio.get_running_loop() result = await loop.run_in_executor( None, backend.send_command, LoadCmd(source=source), 10.0 ) try: response = _result_to_json(result, current_ir_graph, is_reset=False) except Exception as e: logger.exception("Error serializing result") response = { "type": "error", "message": f"Serialization error: {e}", } if isinstance(result, GraphLoaded): current_ir_graph = result.ir_graph current_graph_loaded = graph_loaded_json( result.ir_graph, result.snapshot ) await manager.broadcast(response) return response
@app.post("/reset") async def reset_endpoint(request_data: dict | None = None) -> dict: """REST endpoint for reset.
Body: {"reload": false} """ nonlocal current_graph_loaded, current_ir_graph reload = request_data.get("reload", False) if request_data else False loop = asyncio.get_running_loop() result = await loop.run_in_executor( None, backend.send_command, ResetCmd(reload=reload), 10.0 ) try: response = _result_to_json(result, current_ir_graph, is_reset=True) except Exception as e: logger.exception("Error serializing result") response = { "type": "error", "message": f"Serialization error: {e}", } if isinstance(result, GraphLoaded): current_ir_graph = result.ir_graph current_graph_loaded = graph_loaded_json( result.ir_graph, result.snapshot ) # If reset without reload, clear cached state elif isinstance(result, StepResult) and not reload: current_ir_graph = None current_graph_loaded = None await manager.broadcast(response) return response
@app.get("/state") async def state_endpoint() -> dict: """REST endpoint for current state snapshot.""" if current_graph_loaded is None: return {"error": "No program loaded"} return current_graph_loaded
# Mount static files for frontend frontend_dir = Path(__file__).parent / "frontend" if (frontend_dir / "dist").exists(): app.mount( "/dist", StaticFiles(directory=str(frontend_dir / "dist")), name="dist", ) if frontend_dir.exists(): app.mount( "/", StaticFiles(directory=str(frontend_dir), html=True), name="frontend", )
return app
def _result_to_json(result: GraphLoaded | StepResult | ErrorResult, ir_graph: IRGraph | None, is_reset: bool = False) -> dict: """Convert a command result to JSON for transmission.
Args: result: Result dataclass from backend ir_graph: Current IRGraph (for context in responses) is_reset: True if this result is from a reset command
Returns: JSON-serializable dict """ if isinstance(result, GraphLoaded): return graph_loaded_json(result.ir_graph, result.snapshot)
elif isinstance(result, StepResult): # Handle reset without reload: snapshot is None but this is not an error if result.snapshot is None: if is_reset and result.finished and result.sim_time == 0.0: return { "type": "reset", "sim_time": 0.0, "message": "Simulation reset", } else: return { "type": "error", "message": "No program loaded", } return graph_to_monitor_json(ir_graph, result.snapshot, list(result.events))
elif isinstance(result, ErrorResult): return { "type": "error", "message": result.message, "errors": result.errors, }
else: return { "type": "error", "message": f"Unknown result type: {type(result).__name__}", }