Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| import os | |
| from typing import Any | |
| from fastapi import HTTPException | |
| from openenv.core.env_server.http_server import create_app | |
| from pydantic import BaseModel | |
| import uvicorn | |
| from graders import GRADERS, grade_easy | |
| from server.env import EcoLLMInferenceRoutingEnvironment | |
| from server.models import ModelChoice, RLAction, RLObservation, RLState, Strategy | |
| from tasks import TASKS as VALIDATOR_TASKS | |
| from tasks import TASK_REGISTRY | |
| app = create_app( | |
| EcoLLMInferenceRoutingEnvironment, | |
| RLAction, | |
| RLObservation, | |
| env_name="eco_llm_inference_routing", | |
| max_concurrent_envs=4, | |
| ) | |
| _TASK_META = [ | |
| { | |
| **task, | |
| "grader": task.get("grader_path", "unknown"), | |
| "graders": task.get("grader_paths", []), | |
| } | |
| for task in VALIDATOR_TASKS | |
| ] | |
| _ACTION_SPACE = { | |
| "strategy": [strategy.value for strategy in Strategy], | |
| "model_choice": [model.value for model in ModelChoice], | |
| "exit_flag": "boolean", | |
| } | |
| _ENDPOINTS = { | |
| "reset": {"method": "POST", "path": "/reset"}, | |
| "step": {"method": "POST", "path": "/step"}, | |
| "state": {"method": "GET", "path": "/state"}, | |
| "tasks": {"method": "GET", "path": "/tasks"}, | |
| "tasks_meta": {"method": "GET", "path": "/tasks/meta"}, | |
| "metadata": {"method": "GET", "path": "/metadata"}, | |
| "schema": {"method": "GET", "path": "/schema"}, | |
| "grader": {"method": "POST", "path": "/grader"}, | |
| "grade": {"method": "POST", "path": "/grade"}, | |
| "grader_replay": {"method": "POST", "path": "/grader/replay"}, | |
| "docs": {"method": "GET", "path": "/docs"}, | |
| } | |
| class ResetRequest(BaseModel): | |
| task_id: str | None = None | |
| seed: int | None = None | |
| episode_id: str | None = None | |
| class GradeRequest(BaseModel): | |
| task_id: str | |
| episode_id: str = "" | |
| rewards: list[float] = [] | |
| steps: int = 0 | |
| success: bool = False | |
| state: dict[str, Any] = {} | |
| class ReplayAction(BaseModel): | |
| strategy: str | |
| model_choice: str | |
| exit_flag: bool = False | |
| class ReplayRequest(BaseModel): | |
| task_id: str | |
| actions: list[ReplayAction] | |
| def root() -> dict[str, Any]: | |
| return { | |
| "status": "ok", | |
| "service": "eco_llm_inference_routing", | |
| "docs": "/docs", | |
| "metadata": "/metadata", | |
| "tasks": "/tasks", | |
| } | |
| def health() -> dict[str, str]: | |
| return {"status": "healthy"} | |
| def metadata() -> dict[str, Any]: | |
| return { | |
| "name": "eco_llm_inference_routing", | |
| "description": "Carbon-aware RL environment for multi-objective LLM query routing", | |
| "version": "1.0.0", | |
| "domain": "llm systems optimization", | |
| "deterministic": True, | |
| "observation_highlights": [ | |
| "query text", | |
| "carbon intensity", | |
| "cache contents", | |
| "knowledge-base availability", | |
| "structured reward breakdown", | |
| ], | |
| "action_space": _ACTION_SPACE, | |
| "endpoints": _ENDPOINTS, | |
| "tasks": _TASK_META, | |
| "reward_range": [0.0, 1.0], | |
| } | |
| def schema() -> dict[str, Any]: | |
| return { | |
| "action_schema": RLAction.model_json_schema(), | |
| "observation_schema": RLObservation.model_json_schema(), | |
| "state_schema": RLState.model_json_schema(), | |
| } | |
| def list_tasks() -> list[dict[str, Any]]: | |
| return _TASK_META | |
| def list_task_ids() -> list[str]: | |
| return list(TASK_REGISTRY.keys()) | |
| def list_tasks_meta() -> dict[str, Any]: | |
| return { | |
| "tasks": _TASK_META, | |
| "count": len(_TASK_META), | |
| "available_task_ids": [task["id"] for task in _TASK_META], | |
| } | |
| def get_task(task_id: str) -> dict[str, Any]: | |
| for task in _TASK_META: | |
| if task["id"] == task_id: | |
| return task | |
| raise HTTPException(status_code=404, detail=f"Unknown task '{task_id}'") | |
| def _grade_payload(request: GradeRequest) -> dict[str, Any]: | |
| fn = GRADERS.get(request.task_id) | |
| if fn is None: | |
| raise HTTPException( | |
| status_code=400, | |
| detail=f"Invalid task_id '{request.task_id}'. Valid: {sorted(GRADERS)}", | |
| ) | |
| total_reward = sum(request.rewards) | |
| score = float(fn(request.state, total_reward)) | |
| return { | |
| "score": round(score, 4), | |
| "task_id": request.task_id, | |
| "success": score >= 0.5, | |
| "total_reward": round(total_reward, 4), | |
| "steps": request.steps, | |
| "reward": round(score, 4), | |
| "normalized_reward": round(score, 4), | |
| } | |
| def grader_endpoint(request: GradeRequest) -> dict[str, Any]: | |
| return _grade_payload(request) | |
| def grade_endpoint(request: GradeRequest) -> dict[str, Any]: | |
| return _grade_payload(request) | |
| def grader_replay(request: ReplayRequest) -> dict[str, Any]: | |
| if request.task_id not in TASK_REGISTRY: | |
| raise HTTPException(status_code=400, detail=f"Invalid task_id: {request.task_id}") | |
| env = EcoLLMInferenceRoutingEnvironment() | |
| obs = env.reset(task_id=request.task_id) | |
| rewards: list[float] = [] | |
| for action in request.actions: | |
| if obs.done: | |
| break | |
| try: | |
| obs = env.step( | |
| RLAction( | |
| strategy=Strategy(action.strategy), | |
| model_choice=ModelChoice(action.model_choice), | |
| exit_flag=action.exit_flag, | |
| ) | |
| ) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail=f"Invalid replay action: {exc}") from exc | |
| rewards.append(float(obs.reward)) | |
| total_reward = sum(rewards) | |
| fn = GRADERS.get(request.task_id, grade_task_1) | |
| score = float(fn({}, total_reward)) | |
| return { | |
| "score": round(score, 4), | |
| "task_id": request.task_id, | |
| "success": score >= 0.5, | |
| "total_reward": round(total_reward, 4), | |
| "steps": len(rewards), | |
| } | |
| def main() -> None: | |
| port = int(os.getenv("PORT", "7860")) | |
| uvicorn.run("server.app:app", host="0.0.0.0", port=port) | |
| if __name__ == "__main__": | |
| main() | |
| # ── Gradio UI (mounted at /ui — does NOT override any API routes) ───────────── | |
| try: | |
| import gradio as gr | |
| from gradio.routes import mount_gradio_app | |
| from server.env import EcoLLMInferenceRoutingEnvironment | |
| from server.models import RLAction, Strategy, ModelChoice | |
| _env_store: dict = {} | |
| def ui_reset(task_id: str) -> str: | |
| env = EcoLLMInferenceRoutingEnvironment() | |
| obs = env.reset(task_id=task_id) | |
| _env_store["env"] = env | |
| _env_store["rewards"] = [] | |
| _env_store["task_id"] = task_id | |
| return ( | |
| f"✅ Reset — {task_id} ({env.current_task.difficulty})\n" | |
| f"Query : {obs.query}\n" | |
| f"Carbon : {obs.carbon_intensity:.2f}\n" | |
| f"Cache : {obs.cache_contents or 'empty'}\n" | |
| f"KB avail : {obs.kb_available}\n" | |
| f"Done : {obs.done}" | |
| ) | |
| def ui_step(strategy: str, model: str, exit_flag: bool) -> str: | |
| env = _env_store.get("env") | |
| if env is None: | |
| return "⚠️ Click Reset first." | |
| if env._episode_complete: | |
| rewards = _env_store.get("rewards", []) | |
| total = sum(rewards) | |
| return f"🏁 Episode complete.\nTotal reward: {total:.3f}\nRewards: {rewards}" | |
| try: | |
| obs = env.step(RLAction( | |
| strategy=Strategy(strategy), | |
| model_choice=ModelChoice(model), | |
| exit_flag=exit_flag, | |
| )) | |
| except Exception as exc: | |
| return f"❌ Error: {exc}" | |
| _env_store.setdefault("rewards", []).append(obs.reward_details.total_reward) | |
| rd = obs.reward_details | |
| out = ( | |
| f"step reward : {rd.total_reward:+.3f} ({'✓ correct' if rd.correct else '✗ wrong'})\n" | |
| f" score : {rd.score}\n" | |
| f" carbon pen: {rd.carbon_penalty:.3f}\n" | |
| f" lat pen : {rd.latency_penalty:.3f}\n" | |
| f" bonuses : {rd.bonuses:.3f}\n" | |
| f" trace : {rd.strategy_trace}\n\n" | |
| ) | |
| if obs.done: | |
| total = sum(_env_store["rewards"]) | |
| out += f"🏁 Done! Total reward: {total:.3f}" | |
| else: | |
| out += ( | |
| f"Next query : {obs.query}\n" | |
| f"Carbon : {obs.carbon_intensity:.2f}\n" | |
| f"Cache : {obs.cache_contents or 'empty'}\n" | |
| f"KB avail : {obs.kb_available}" | |
| ) | |
| return out | |
| with gr.Blocks(title="🌱 Eco-LLM Routing") as _demo: | |
| gr.Markdown( | |
| "# 🌱 Eco-LLM Inference Routing\n" | |
| "Carbon-aware LLM query routing — balance **accuracy**, **energy**, **latency** & **carbon**.\n\n" | |
| "**API endpoints:** `POST /reset` · `POST /step` · `GET /tasks` · `POST /grade`" | |
| ) | |
| with gr.Row(): | |
| task_dd = gr.Dropdown( | |
| choices=["task_1", "task_2", "task_3"], | |
| value="task_1", | |
| label="Task", | |
| ) | |
| reset_btn = gr.Button("🔄 Reset", variant="primary") | |
| state_box = gr.Textbox(label="Environment State", lines=8, interactive=False) | |
| gr.Markdown("### Take a Step") | |
| with gr.Row(): | |
| strat_dd = gr.Dropdown( | |
| choices=["NONE", "USE_CACHE", "DO_CASCADE", "EARLY_EXIT", "WAIT", "CALL_KB"], | |
| value="NONE", | |
| label="Strategy", | |
| ) | |
| model_dd = gr.Dropdown( | |
| choices=["SMALL", "MEDIUM", "LARGE"], | |
| value="SMALL", | |
| label="Model", | |
| ) | |
| exit_cb = gr.Checkbox(label="exit_flag", value=False) | |
| step_btn = gr.Button("▶️ Step", variant="secondary") | |
| result_box = gr.Textbox(label="Step Result", lines=10, interactive=False) | |
| gr.Markdown( | |
| "### Reward Formula\n" | |
| "`reward = score − (carbon × energy) − (0.1 × latency) + bonuses`\n\n" | |
| "| Model | Energy | Latency | | Strategy | Bonus |\n" | |
| "|-------|--------|---------|--|----------|-------|\n" | |
| "| SMALL | 0.1 | 1.0 | | USE_CACHE | +0.5 |\n" | |
| "| MEDIUM | 0.3 | 2.0 | | EARLY_EXIT | +0.1 |\n" | |
| "| LARGE | 0.6 | 5.0 | | CALL_KB | low cost |" | |
| ) | |
| reset_btn.click(fn=ui_reset, inputs=[task_dd], outputs=[state_box]) | |
| step_btn.click(fn=ui_step, inputs=[strat_dd, model_dd, exit_cb], outputs=[result_box]) | |
| app = mount_gradio_app(app, _demo, path="/ui") | |
| except ImportError: | |
| pass # Gradio optional — API works without it |