Merge pull request #12 from EPFL-AI-Team/marcus-server-implementation

Implement Server Websocket Endpoint (#1)

Author: AnnaelleMyriam

.gitignore

@@ -679,3 +679,4 @@ tags
 .ionide
 
 # End of https://www.toptal.com/developers/gitignore/api/vim,latex,linux,macos,synology,jetbrains+all,visualstudiocode,python,jupyternotebooks
+inference_results.txt

src/iris/cli/server.py

@@ -0,0 +1,4 @@
+from iris.server.app import main
+
+if __name__ == "__main__":
+    main()

src/iris/server/__init__.py

@@ -0,0 +1,109 @@
+"""IRIS Inference Server - receives frames, runs VLM inference."""
+
+import base64
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from io import BytesIO
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from PIL import Image
+
+from iris.server.dependencies import get_server_state
+from iris.vlm.inference.model_loader import load_model_and_processor
+from iris.vlm.inference.queue.jobs import SingleFrameJob
+from iris.vlm.inference.queue.queue import InferenceQueue
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+    """Manage startup and shutdown."""
+    # Startup
+    state = get_server_state()
+
+    logger.info("Loading model...")
+    state.model, state.processor = load_model_and_processor("smolvlm2")
+
+    logger.info("Starting inference queue...")
+    state.queue = InferenceQueue(max_queue_size=10, num_workers=1)
+    await state.queue.start()
+
+    state.model_loaded = True
+    logger.info("Server ready!")
+
+    yield
+
+    # Shutdown
+    if state.queue:
+        await state.queue.stop()
+    logger.info("Server stopped.")
+
+
+app = FastAPI(title="IRIS Inference Server", lifespan=lifespan)
+
+
+@app.get("/health")
+async def health() -> dict[str, str | bool]:
+    """Health check endpoint."""
+    state = get_server_state()
+    return {
+        "status": "healthy" if state.model_loaded else "loading",
+        "model_loaded": state.model_loaded,
+    }
+
+
+@app.websocket("/ws/stream")
+async def inference_endpoint(websocket: WebSocket) -> None:
+    """Receive frames and return inference results."""
+    await websocket.accept()
+    state = get_server_state()
+    logger.info("Client connected")
+
+    try:
+        while True:
+            data = await websocket.receive_json()
+
+            frame_b64 = data["frame"]
+            frame_id = data["frame_id"]
+
+            image_data = base64.b64decode(frame_b64)
+            image = Image.open(BytesIO(image_data))
+
+            job = SingleFrameJob(
+                job_id=f"frame-{frame_id}",
+                frame=image,
+                model=state.model,
+                processor=state.processor,
+                prompt="Describe what you see in one sentence.",
+                executor=state.queue.executor,
+            )
+
+            await state.queue.submit(job)
+            result_job = await state.queue.get_result(timeout=30.0)
+
+            if result_job:
+                await websocket.send_json({
+                    "job_id": result_job.job_id,
+                    "status": result_job.status.value,
+                    "result": result_job.result,
+                    "processing_time": result_job.processing_time,
+                })
+
+    except WebSocketDisconnect:
+        logger.info("Client disconnected")
+    except Exception as e:
+        logger.error(f"Error: {e}", exc_info=True)
+
+
+def main() -> None:
+    """Entry point for server."""
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8001)
+
+
+if __name__ == "__main__":
+    main()

src/iris/server/app.py

@@ -0,0 +1,22 @@
+"""Server-side state management with dependency injection."""
+
+from iris.vlm.inference.queue.queue import InferenceQueue
+
+
+class ServerState:
+    """Server application state."""
+
+    def __init__(self):
+        self.model = None
+        self.processor = None
+        self.queue: InferenceQueue | None = None
+        self.model_loaded = False
+
+
+# Singleton
+_server_state = ServerState()
+
+
+def get_server_state() -> ServerState:
+    """Get server state for dependency injection."""
+    return _server_state