RohanExploit/Meta-hackathon
0
1"""FastAPI server for multi-channel retail environment."""2import asyncio3import json4import os5from pathlib import Path6import threading7import time8from typing import Any, Dict, List, Optional9 10import uvicorn11from fastapi import FastAPI, HTTPException, Request12from fastapi.responses import FileResponse, StreamingResponse13from pydantic import BaseModel, ValidationError14 15from environment.grader import evaluate_seeded_summaries, score_episode16from environment.models import (17 ActionType,18 AllocateAction,19 CompositeAction,20 NoOpAction,21 OrderAction,22 PromoteAction,23 RetailAction,24 SetPriceAction,25)26from environment.retail_env import MultiChannelRetailEnv27from environment.tasks import get_task_config, list_tasks28 29app = FastAPI(30 title="Multi-Channel Retail Environment",31 description="OpenEnv compliant multi-channel retail with disruption recovery",32)33 34BASE_DIR = Path(__file__).resolve().parent35UI_DIR = BASE_DIR / "ui"36 37# Global environment instance38env = MultiChannelRetailEnv()39 40 41class ResetRequest(BaseModel):42 task_config: Optional[Dict[str, Any]] = None43 task_name: Optional[str] = None44 seed: Optional[int] = None45 46 47class StepRequest(BaseModel):48 action: Dict[str, Any]49 50 51class EvaluateRequest(BaseModel):52 summary: Optional[Dict[str, Any]] = None53 summaries: Optional[List[Dict[str, Any]]] = None54 variance_penalty: float = 0.055 56 57class LiveStartRequest(BaseModel):58 task_name: Optional[str] = None59 seed: Optional[int] = None60 mode: str = "heuristic" # heuristic | noop61 interval_ms: int = 60062 63 64class LiveStopRequest(BaseModel):65 reason: Optional[str] = None66 67 68def _parse_action(payload: Dict[str, Any]) -> RetailAction:69 """Parse action from JSON payload.70 71 Supports both legacy single actions and the new CompositeAction72 format that allows multiple actions per timestep.73 """74 if not isinstance(payload, dict):75 raise ValueError("Action must be a JSON object")76 77 action_type = payload.get("action")78 if action_type is None:79 raise ValueError("Missing 'action' field")80 81 action_value = str(action_type).strip().lower()82 83 try:84 if action_value == ActionType.COMPOSITE.value:85 return CompositeAction(**payload)86 if action_value == ActionType.ALLOCATE.value:87 return AllocateAction(**payload)88 if action_value == ActionType.SET_PRICE.value:89 return SetPriceAction(**payload)90 if action_value == ActionType.ORDER.value:91 return OrderAction(**payload)92 if action_value == ActionType.PROMOTE.value:93 return PromoteAction(**payload)94 if action_value == ActionType.NOOP.value:95 return NoOpAction(**payload)96 except ValidationError as exc:97 raise ValueError(f"Invalid action: {exc}") from exc98 99 raise ValueError(f"Unknown action type: {action_type}")100 101 102def _resolve_task_config(request: ResetRequest) -> Dict[str, Any]:103 """Resolve task configuration.104 105 Falls back to the first available task ("easy") when the judge106 sends a POST /reset with no body — which is valid per OpenEnv spec.107 """108 if request.task_config is not None:109 return request.task_config110 if request.task_name is not None:111 return get_task_config(request.task_name)112 # No body provided — default to "easy" task so bare POST /reset works113 return get_task_config("easy")114 115 116env_lock = threading.Lock()117 118 119class LiveRunner:120 """Realtime step runner for continuous wall-clock execution."""121 122 def __init__(self, env_ref, lock_ref) -> None:123 self.env = env_ref124 self.env_lock = lock_ref125 self._thread: Optional[threading.Thread] = None126 self._running = False127 self._mode = "heuristic"128 self._interval_s = 0.6129 self._tick = 0130 self._latest: Dict[str, Any] = {131 "tick": 0,132 "running": False,133 "done": False,134 "reward": 0.0,135 "observation": None,136 "info": {},137 "timestamp": 0.0,138 "error": None,139 }140 self._status_lock = threading.Lock()141 142 def _heuristic_action(self, observation: Any) -> Dict[str, Any]:143 products = list(observation.inventory.keys())144 if not products:145 return {"action": "noop"}146 147 target = min(products, key=lambda p: int(observation.inventory.get(p, 0)))148 min_inv = int(observation.inventory.get(target, 0))149 if min_inv <= 2:150 return {"action": "order", "product": target, "quantity": 3}151 152 if getattr(observation, "disruption_active", False) and float(observation.cash) > 20.0:153 return {"action": "promote", "product": target, "budget_allocated": 8.0}154 155 return {"action": "noop"}156 157 def _run_loop(self) -> None:158 while self._running:159 try:160 with self.env_lock:161 if self.env.state is None:162 self._set_latest(error="Environment not initialized. Use /live/start with task_name.")163 time.sleep(self._interval_s)164 continue165 166 obs = self.env._get_observation()167 if self._mode == "noop":168 action_payload = {"action": "noop"}169 else:170 action_payload = self._heuristic_action(obs)171 172 action = _parse_action(action_payload)173 observation, reward, done, info = self.env.step(action)174 175 self._tick += 1176 self._set_latest(177 tick=self._tick,178 running=True,179 done=bool(done),180 reward=float(reward),181 observation=observation.model_dump(),182 info=info,183 timestamp=time.time(),184 error=None,185 )186 187 if done:188 self._running = False189 self._set_latest(running=False)190 break191 192 except Exception as exc:193 self._set_latest(error=str(exc), running=False)194 self._running = False195 break196 197 time.sleep(self._interval_s)198 199 def _set_latest(self, **kwargs: Any) -> None:200 with self._status_lock:201 self._latest.update(kwargs)202 203 def start(self, mode: str, interval_ms: int) -> None:204 if self._running:205 return206 self._mode = mode if mode in {"heuristic", "noop"} else "heuristic"207 self._interval_s = max(0.1, float(interval_ms) / 1000.0)208 self._running = True209 self._set_latest(running=True, error=None)210 self._thread = threading.Thread(target=self._run_loop, daemon=True)211 self._thread.start()212 213 def stop(self, reason: Optional[str] = None) -> None:214 self._running = False215 self._set_latest(running=False, info={"reason": reason or "stopped"})216 217 def status(self) -> Dict[str, Any]:218 with self._status_lock:219 return {220 "running": bool(self._running),221 "mode": self._mode,222 "interval_ms": int(self._interval_s * 1000),223 "tick": int(self._latest.get("tick", 0)),224 "done": bool(self._latest.get("done", False)),225 "timestamp": float(self._latest.get("timestamp", 0.0)),226 "error": self._latest.get("error"),227 }228 229 def latest(self) -> Dict[str, Any]:230 with self._status_lock:231 return dict(self._latest)232 233 234live_runner = LiveRunner(env, env_lock)235 236 237@app.post("/reset")238async def reset_environment(raw_request: Request):239 """Reset the environment.240 241 Accepts POST with any body (JSON, empty, or null) per OpenEnv spec.242 """243 try:244 body = await raw_request.body()245 if body and body.strip():246 try:247 data = json.loads(body)248 request = ResetRequest(**(data if isinstance(data, dict) else {}))249 except (json.JSONDecodeError, ValidationError):250 request = ResetRequest()251 else:252 request = ResetRequest()253 254 with env_lock:255 task_config = _resolve_task_config(request)256 if request.seed is not None:257 env.seed = int(request.seed)258 elif "seed" in task_config:259 env.seed = int(task_config["seed"])260 261 observation = env.reset(task_config)262 263 return {264 "observation": observation.model_dump(),265 "reward": 0.0,266 "done": False,267 "info": {268 "task_name": task_config.get("name", request.task_name),269 "horizon": int(task_config.get("horizon", 30)),270 },271 }272 except Exception as e:273 raise HTTPException(status_code=400, detail=str(e))274 275 276@app.post("/step")277async def step_environment(raw_request: Request):278 """Take a step in the environment.279 280 Accepts POST with any body (JSON, empty, or null) per OpenEnv spec.281 """282 try:283 body = await raw_request.body()284 if body and body.strip():285 try:286 data = json.loads(body)287 request = StepRequest(**(data if isinstance(data, dict) else {"action": {"action": "noop"}}))288 except (json.JSONDecodeError, ValidationError):289 request = StepRequest(action={"action": "noop"})290 else:291 request = StepRequest(action={"action": "noop"})292 293 with env_lock:294 action = _parse_action(request.action)295 observation, reward, done, info = env.step(action)296 297 return {298 "observation": observation.model_dump(),299 "reward": reward,300 "done": done,301 "info": info,302 }303 except Exception as e:304 raise HTTPException(status_code=400, detail=str(e))305 306 307@app.get("/state")308async def get_state():309 """Get current internal state and episode metrics."""310 try:311 with env_lock:312 state = env.get_state()313 metrics = dict(env.episode_metrics)314 return {"state": state, "episode_metrics": metrics}315 except Exception as e:316 raise HTTPException(status_code=400, detail=str(e))317 318 319@app.get("/tasks")320async def get_tasks_list():321 """List available tasks."""322 try:323 return {"tasks": list_tasks()}324 except Exception as e:325 raise HTTPException(status_code=400, detail=str(e))326 327 328@app.post("/evaluate")329async def evaluate_episode(request: EvaluateRequest):330 """Evaluate episode(s)."""331 try:332 if request.summaries is not None:333 result = evaluate_seeded_summaries(334 request.summaries,335 variance_penalty=request.variance_penalty,336 )337 return {338 "mode": "multi_seed",339 "result": result,340 }341 342 if request.summary is not None:343 result = score_episode(request.summary)344 return {345 "mode": "single_summary",346 "result": result,347 }348 349 raise ValueError("Either summary or summaries must be provided")350 except Exception as e:351 raise HTTPException(status_code=400, detail=str(e))352 353 354@app.post("/live/start")355async def live_start(request: LiveStartRequest):356 """Start realtime continuous stepping."""357 try:358 if request.task_name:359 with env_lock:360 task_cfg = get_task_config(request.task_name)361 if request.seed is not None:362 env.seed = int(request.seed)363 elif "seed" in task_cfg:364 env.seed = int(task_cfg["seed"])365 env.reset(task_cfg)366 367 live_runner.start(mode=request.mode, interval_ms=int(request.interval_ms))368 return {369 "ok": True,370 "status": live_runner.status(),371 }372 except Exception as e:373 raise HTTPException(status_code=400, detail=str(e))374 375 376@app.post("/live/stop")377async def live_stop(request: LiveStopRequest):378 """Stop realtime stepping."""379 live_runner.stop(request.reason)380 return {381 "ok": True,382 "status": live_runner.status(),383 }384 385 386@app.get("/live/status")387async def live_status():388 """Get realtime runner status."""389 return live_runner.status()390 391 392@app.get("/live/latest")393async def live_latest():394 """Get latest realtime tick payload."""395 return live_runner.latest()396 397 398@app.get("/live/stream")399async def live_stream():400 """Server-Sent Events stream for realtime live runner updates.401 402 Clients receive a ``data:`` event containing a JSON-encoded payload403 identical to ``/live/latest`` every 500 ms while the connection is open.404 This removes the need for client-side polling.405 """406 407 async def _event_generator():408 try:409 while True:410 data = live_runner.latest()411 yield f"data: {json.dumps(data)}\n\n"412 await asyncio.sleep(0.5)413 except GeneratorExit:414 return415 except Exception:416 # Yield a generic error event so connected clients can react, then close.417 yield f"data: {json.dumps({'error': 'stream error'})}\n\n"418 return419 420 return StreamingResponse(421 _event_generator(),422 media_type="text/event-stream",423 headers={424 "Cache-Control": "no-cache",425 "X-Accel-Buffering": "no",426 "Connection": "keep-alive",427 },428 )429 430 431@app.get("/health")432async def health_check():433 """Health check."""434 return {"status": "healthy"}435 436 437@app.get("/dashboard")438async def ui_dashboard():439 """Serve the modern visual dashboard UI."""440 ui_file = UI_DIR / "dashboard.html"441 if not ui_file.exists():442 raise HTTPException(status_code=404, detail="Dashboard UI not found")443 return FileResponse(ui_file, media_type="text/html")444 445 446@app.get("/static/chart.js")447async def serve_chartjs():448 """Serve bundled Chart.js for the dashboard."""449 js_file = UI_DIR / "chart.umd.min.js"450 if not js_file.exists():451 raise HTTPException(status_code=404, detail="Chart.js not found")452 return FileResponse(js_file, media_type="application/javascript")453 454 455@app.get("/")456async def ui_home():457 """Serve the terminal UI."""458 ui_file = UI_DIR / "index.html"459 if not ui_file.exists():460 raise HTTPException(status_code=404, detail="UI not found")461 return FileResponse(ui_file)462 463 464def main():465 port = int(os.environ.get("PORT", 8000))466 uvicorn.run(app, host="0.0.0.0", port=port)467 468if __name__ == "__main__":469 main()470