Spaces:
Paused
Paused
Download api/main.py from chandrakiran06/rosdiff: direct link, hf CLI and curl.
- Browser
- Download file 42.7 kB
-
https://huggingface.co/spaces/chandrakiran06/rosdiff/resolve/main/api/main.py
- Command line
-
hf download hf://spaces/chandrakiran06/rosdiff/api/main.py
-
curl -L -o main.py https://huggingface.co/spaces/chandrakiran06/rosdiff/resolve/main/api/main.py
42.7 kB
| """RosDiff Stage 1 API (Orcanion internal, testers only). | |
| uvicorn api.main:app --host 0.0.0.0 --port 8000 | |
| Every endpoint except /health needs a tester key in the ``X-API-Key`` header | |
| (or ``Authorization: Bearer <key>``). | |
| """ | |
| from __future__ import annotations | |
| import datetime as dt | |
| import json | |
| import logging | |
| import shutil | |
| import time | |
| from collections.abc import Callable | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| from typing import Literal | |
| from fastapi import Depends, FastAPI, Header, HTTPException, Request | |
| from fastapi.responses import HTMLResponse, Response | |
| from fastapi.staticfiles import StaticFiles | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel, Field | |
| from rag.store import RagStore | |
| from . import budget | |
| from .auth import Allowlist, Tester | |
| from .codegen import DirectBackend, OpenCodeBackend, generate_code, mentioned_interfaces | |
| from .config import Settings, get_settings | |
| from .enhance import Enhanced | |
| from .enhance import enhance as enhance_prompt | |
| from .enhance import record as record_enhancement | |
| from .jobs import JobError, compute_usage, month_spend, refresh, start_isaac, start_review | |
| from .jobs import submit as submit_job | |
| from .llm import LLM, ChatClient, LLMError | |
| from .logstore import LogStore | |
| from .mjcf_check import validate_mjcf | |
| from .opencode_runner import OpenCodeWorkspace | |
| from .robots import ROBOT_NAMES, ROBOTS | |
| from .ros_check import load_interfaces | |
| from .routing import ModelSpec, all_models, is_free, plan_code, plan_scene | |
| from .runs import check_token, new_run_id, render_report, report_token | |
| from .sim_backend import SimBackend, make_backend | |
| from .simulate import generate_scene | |
| from .storage import R2Config, R2Storage, Storage | |
| from .zen import WithFreeModels, ZenLLM | |
| log = logging.getLogger("rosdiff") | |
| class Services: | |
| settings: Settings | |
| store: RagStore | |
| logs: LogStore | |
| allowlist: Allowlist | |
| llm_factory: Callable[[], LLM] # scene generation | |
| backend: SimBackend | None = None # None: /simulate/run is disabled | |
| storage: Storage | None = None | |
| code_llm_factory: Callable[[], LLM] | None = None # direct ROS 2 code backend; None: same as llm_factory | |
| enhancer: Callable[..., Enhanced] = enhance_prompt # free prompt enhancement through OpenCode Zen | |
| free_llm_factory: Callable[[], LLM] | None = None # free debugging model (OpenCode Zen); None: paid only | |
| def default_services() -> Services: | |
| settings = get_settings() | |
| r2 = R2Config.from_env() | |
| if not settings.report_secret: | |
| log.warning("REPORT_SECRET is not set: report links will stop working when the server restarts") | |
| return Services( | |
| settings=settings, | |
| store=RagStore(settings.data_dir / "chroma"), | |
| logs=LogStore(settings.log_db), | |
| allowlist=Allowlist(settings.testers_file), | |
| llm_factory=lambda: ChatClient(settings), | |
| code_llm_factory=lambda: ChatClient(settings, model=settings.code_llm_model), | |
| free_llm_factory=lambda: ZenLLM(settings), | |
| backend=make_backend(settings), | |
| storage=R2Storage(r2) if r2 else None, | |
| ) | |
| # --------------------------------------------------------------------------- request/response models | |
| class SimulateRequest(BaseModel): | |
| description: str = Field(..., min_length=10, max_length=4000, description="The scenario, in plain language") | |
| robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Robot to include, or null for none") | |
| confirm_cost: bool = Field(False, description="Run even if the estimate is above CONFIRM_ABOVE_INR") | |
| enhance: bool = Field(True, description="Sharpen the request first with a free OpenCode Zen model") | |
| class SimulateResponse(BaseModel): | |
| id: str | |
| valid: bool | |
| mjcf: str | None | |
| robot: str | None | |
| validation: dict | |
| attempts: int | |
| sources: list[str] | |
| model: str | |
| how_to_run: str | |
| routing: dict | None = None # what RosDiff chose and why (tier, reason, simulator) | |
| enhanced_prompt: dict | None = None # {used, text, model, note}: the brief the generator actually got | |
| cost: dict | None = None # rupees this task used, the estimate, and the remaining budget | |
| class CodeRequest(BaseModel): | |
| request: str = Field(..., min_length=10, max_length=4000, description="What the ROS 2 package should do") | |
| distro: Literal["humble", "jazzy"] = "humble" | |
| language: Literal["python", "cpp"] = "python" | |
| package_name: str | None = Field(None, pattern=r"^[a-z][a-z0-9_]{1,62}$") | |
| check: bool = Field( | |
| False, description="Also build, lint, test and launch it in the ROS 2 sandbox (see /code/check)" | |
| ) | |
| confirm_cost: bool = Field(False, description="Run even if the estimate is above CONFIRM_ABOVE_INR") | |
| enhance: bool = Field(True, description="Sharpen the request first with a free OpenCode Zen model") | |
| class CodeResponse(BaseModel): | |
| id: str | |
| valid: bool | |
| files: dict[str, str] | |
| validation: dict | |
| attempts: int | |
| sources: list[str] | |
| backend: str | |
| model: str | |
| how_to_build: str | |
| check: dict | None = None # the sandbox job, when requested: {id, status, report_url} or {error} | |
| routing: dict | None = None | |
| cost: dict | None = None | |
| enhanced_prompt: dict | None = None | |
| class RunRequest(BaseModel): | |
| generation_id: str | None = Field(None, description="A valid scene from /simulate/generate (yours)") | |
| mjcf: str | None = Field(None, max_length=300_000, description="Or: an MJCF scene to validate and run") | |
| robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Only used with mjcf") | |
| duration_s: float = Field(10.0, gt=0, le=60) | |
| controller: Literal["hold", "passive"] = Field("hold", description="hold: keep the starting pose; passive: limp") | |
| class RunResponse(BaseModel): | |
| run_id: str | |
| status: str | |
| report_url: str | |
| status_url: str | |
| class RunStatus(BaseModel): | |
| run_id: str | |
| status: str | |
| error: str | None | |
| summary: dict | None | |
| report_url: str | |
| class SweepRequest(BaseModel): | |
| run_id: str | None = Field(None, description="Sweep the scene of one of your runs (linked on its report)") | |
| generation_id: str | None = Field(None, description="Or: a valid scene from /simulate/generate") | |
| mjcf: str | None = Field(None, max_length=300_000, description="Or: an MJCF scene to validate and sweep") | |
| robot: Literal[ROBOT_NAMES] | None = Field("unitree_g1", description="Only used with mjcf") | |
| n_worlds: int = Field(64, ge=1, le=256, description="Randomized copies, simulated in parallel") | |
| duration_s: float = Field(5.0, gt=0, le=10) | |
| controller: Literal["hold", "passive"] = "hold" | |
| push_max_n: float = Field(80.0, ge=0, le=1000, description="Each world is pushed with 0..this many newtons") | |
| friction: tuple[float, float] = Field((0.5, 1.5), description="Friction scale range") | |
| object_mass: tuple[float, float] = Field((0.5, 2.0), description="Mass scale range for free objects") | |
| seed: int = 0 | |
| class ExportRequest(BaseModel): | |
| run_ids: list[str] = Field(..., min_length=1, max_length=100, description="Your finished runs, one robot") | |
| name: str = Field(..., pattern=r"^[a-z0-9][a-z0-9-]{1,60}$", description="Dataset name, e.g. g1-table-pushes") | |
| private: bool = True | |
| class FeedbackItem(BaseModel): | |
| target_type: Literal["generation", "job"] | |
| target_id: str | |
| rating: Literal["helpful", "not_helpful"] | |
| note: str = Field("", max_length=4000) | |
| class CheckRequest(BaseModel): | |
| generation_id: str = Field(..., description="A package from /code/generate (yours)") | |
| class JobResponse(BaseModel): | |
| id: str | |
| kind: str | |
| status: str | |
| report_url: str | |
| status_url: str | |
| class FeedbackRequest(BaseModel): | |
| rating: Literal["works", "partly", "broken"] | |
| note: str = Field("", max_length=4000) | |
| def _how_to_run(robot: str | None) -> str: | |
| if not robot: | |
| return "Save as scene.xml and run: python -m mujoco.viewer --mjcf scene.xml" | |
| spec = ROBOTS[robot] | |
| return ( | |
| f"Save as scene_rosdiff.xml inside mujoco_menagerie/{spec.directory}/ (next to {spec.model_file}) and run: " | |
| f"python -m mujoco.viewer --mjcf mujoco_menagerie/{spec.directory}/scene_rosdiff.xml" | |
| ) | |
| def _check_budget(s: Services, tester: str, plan, confirmed: bool) -> None: | |
| try: | |
| budget.check(s.logs, s.settings, tester, plan.estimate_inr, confirmed) | |
| except budget.BudgetError as e: | |
| raise HTTPException(status_code=e.status, detail={"message": str(e), **e.detail}) from e | |
| def _enhance(s: Services, wanted: bool, kind: str, text: str, **ctx) -> Enhanced: | |
| if not wanted: | |
| return Enhanced(text, text, False, note="not requested") | |
| try: | |
| return s.enhancer(s.settings, kind, text, **ctx) | |
| except Exception as e: # enhancement is a nicety: never fail a generation because of it | |
| log.warning("prompt enhancement crashed: %s", e) | |
| return Enhanced(text, text, False, note="enhancer error") | |
| def _enhanced_view(e: Enhanced) -> dict: | |
| return {"used": e.used, "text": e.text if e.used else None, "model": e.model, "note": e.note, "ms": e.ms} | |
| def _spec(model_id: str | None) -> ModelSpec: | |
| if is_free(model_id): | |
| return ModelSpec(model_id, "free", "OpenCode Zen", 0.0, 0.0, free=True) | |
| return all_models().get(model_id) or ModelSpec(model_id or "unknown", "?", "unknown", 0.0, 0.0) | |
| def _llm(s: Services, factory: Callable[[], LLM]) -> LLM: | |
| """The paid client, with `opencode/...` (free debugging) models routed to OpenCode Zen.""" | |
| paid = factory() | |
| return WithFreeModels(paid, s.free_llm_factory()) if s.free_llm_factory else paid | |
| def _book(s: Services, tester: str, gen_id: str, plan, attempts, brief: Enhanced | None = None) -> dict: | |
| """Record every model call of a generation in the rupee ledger; return what this task cost.""" | |
| before = budget.summary(s.logs, s.settings, tester)["used"] | |
| if brief is not None: | |
| record_enhancement(s.logs, tester, gen_id, plan.category, brief) | |
| for i, a in enumerate(attempts): | |
| spec = _spec(a.model) | |
| budget.record_call( | |
| s.logs, | |
| s.settings, | |
| tester=tester, | |
| generation_id=gen_id, | |
| category=plan.category if i == 0 else "debugging", | |
| model=spec, | |
| usage=a.usage, | |
| estimate_inr=plan.estimate_inr[0] if i == 0 else None, | |
| ) | |
| after = budget.summary(s.logs, s.settings, tester) | |
| return { | |
| "inr": round(after["used"] - before, 2), | |
| "estimate_inr": list(plan.estimate_inr), | |
| "budget": {"used": after["used"], "remaining": after["remaining"], "allocation": after["allocation"]}, | |
| } | |
| def _routing_view(plan, attempts) -> dict: | |
| tiers = [_spec(a.model).tier for a in attempts] | |
| view = { | |
| "complexity": plan.complexity, | |
| "tier": plan.tier, | |
| "reason": plan.reason, | |
| "tiers_used": tiers, | |
| "escalated": len({t for t in tiers if t != "free"}) > 1, | |
| "free_debug": "free" in tiers, | |
| "models": [a.model for a in attempts], | |
| } | |
| if plan.simulator: | |
| view.update(simulator=plan.simulator, simulator_reason=plan.simulator_reason) | |
| return view | |
| def _submit(s: Services, tester: str, kind: str, job_input: dict, **kw) -> dict: | |
| try: | |
| return submit_job(s, tester, kind, job_input, **kw) | |
| except JobError as e: | |
| raise HTTPException(status_code=e.status, detail=str(e)) from e | |
| # --------------------------------------------------------------------------- app | |
| def create_app(services: Services | None = None) -> FastAPI: | |
| app = FastAPI( | |
| title="RosDiff Stage 1", | |
| version="0.1.0", | |
| description="Internal tester build: RAG-grounded MJCF and ROS 2 code generation.", | |
| ) | |
| app.state.services = services | |
| app.mount("/static", StaticFiles(directory=Path(__file__).with_name("static")), name="static") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| def svc() -> Services: | |
| if app.state.services is None: | |
| app.state.services = default_services() | |
| return app.state.services | |
| def tester( | |
| x_api_key: str | None = Header(None), | |
| authorization: str | None = Header(None), | |
| s: Services = Depends(svc), | |
| ) -> Tester: | |
| key = x_api_key | |
| if not key and authorization and authorization.lower().startswith("bearer "): | |
| key = authorization[7:] | |
| if s.settings.owner_api_key and key == s.settings.owner_api_key: | |
| return Tester("owner", "", active=True) | |
| t = s.allowlist.identify(key) | |
| if t is None: | |
| raise HTTPException(status_code=401, detail="missing or unknown API key") | |
| return t | |
| def _startup_checks(): | |
| s = svc() | |
| if not s.settings.llm_api_key: | |
| log.warning("LLM_API_KEY is not set: generation endpoints will fail") | |
| if not s.settings.report_secret: | |
| log.warning("REPORT_SECRET is not set: report links will break on restart") | |
| def health(s: Services = Depends(svc)) -> dict: | |
| return { | |
| "status": "ok", | |
| "model": s.settings.llm_model, | |
| "code_model": s.settings.code_llm_model, | |
| "llm_configured": bool(s.settings.llm_api_key), | |
| "codegen_backend": s.settings.codegen_backend, | |
| "opencode_found": shutil.which(s.settings.opencode_bin) is not None, | |
| "prompt_enhance": s.settings.prompt_enhance and shutil.which(s.settings.opencode_bin) is not None, | |
| "enhance_model": s.settings.enhance_model, | |
| "gpu_share_usd": s.settings.compute_share_usd or None, | |
| "interface_catalogue": s.settings.interfaces_file.is_file(), | |
| "index": {name: s.store.count(name) for name in ("mujoco", "ros2")}, | |
| "robots": sorted(ROBOTS), | |
| "runs_enabled": s.backend is not None and s.storage is not None, | |
| "sim_backend": s.backend.name if s.backend else None, | |
| } | |
| def simulate(req: SimulateRequest, t: Tester = Depends(tester), s: Services = Depends(svc)) -> SimulateResponse: | |
| start = time.monotonic() | |
| plan = plan_scene(req.description, req.robot, s.settings) # routed on the tester's own words | |
| _check_budget(s, t.name, plan, req.confirm_cost) | |
| brief = _enhance(s, req.enhance, "scene", req.description, robot=req.robot) | |
| logged_request = { | |
| **req.model_dump(exclude={"confirm_cost"}), | |
| "routing": plan.to_dict(), | |
| "enhanced": _enhanced_view(brief), | |
| } | |
| try: | |
| result = generate_scene(brief.text, req.robot, _llm(s, s.llm_factory), s.store, s.settings, plan=plan) | |
| except LLMError as e: | |
| gen_id = s.logs.record( | |
| tester=t.name, | |
| kind="simulate", | |
| request=logged_request, | |
| error=str(e), | |
| duration_ms=int((time.monotonic() - start) * 1000), | |
| ) | |
| record_enhancement(s.logs, t.name, gen_id, "simulation_ai", brief) | |
| raise HTTPException(status_code=502, detail=f"model call failed (log id {gen_id}): {e}") from e | |
| gen_id = s.logs.record( | |
| tester=t.name, | |
| kind="simulate", | |
| request=logged_request, | |
| model=result.model, | |
| sources=result.sources, | |
| attempts=[asdict(a) for a in result.attempts], | |
| output=result.mjcf, | |
| valid=result.valid, | |
| validation=result.validation.to_dict(), | |
| duration_ms=int((time.monotonic() - start) * 1000), | |
| ) | |
| return SimulateResponse( | |
| id=gen_id, | |
| valid=result.valid, | |
| mjcf=result.mjcf, | |
| robot=req.robot, | |
| validation=result.validation.to_dict(), | |
| attempts=len(result.attempts), | |
| sources=result.sources, | |
| model=result.model, | |
| how_to_run=_how_to_run(req.robot), | |
| routing=_routing_view(plan, result.attempts), | |
| cost=_book(s, t.name, gen_id, plan, result.attempts, brief), | |
| enhanced_prompt=_enhanced_view(brief), | |
| ) | |
| def code( | |
| req: CodeRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc) | |
| ) -> CodeResponse: | |
| start = time.monotonic() | |
| catalogue = load_interfaces(str(s.settings.interfaces_file)).get(req.distro, {}) | |
| plan = plan_code( | |
| req.request, req.distro, req.language, s.settings, len(mentioned_interfaces(req.request, catalogue)) | |
| ) | |
| _check_budget(s, t.name, plan, req.confirm_cost) | |
| brief = _enhance(s, req.enhance, "code", req.request, distro=req.distro, language=req.language) | |
| logged_request = { | |
| **req.model_dump(exclude={"confirm_cost"}), | |
| "routing": plan.to_dict(), | |
| "enhanced": _enhanced_view(brief), | |
| } | |
| try: | |
| if s.settings.codegen_backend == "opencode": | |
| with OpenCodeWorkspace(s.settings) as ws: | |
| result = generate_code( | |
| brief.text, | |
| req.distro, | |
| req.language, | |
| OpenCodeBackend(ws), | |
| s.store, | |
| s.settings, | |
| req.package_name, | |
| plan=plan, | |
| ) | |
| else: | |
| result = generate_code( | |
| brief.text, | |
| req.distro, | |
| req.language, | |
| DirectBackend(_llm(s, s.code_llm_factory or s.llm_factory)), | |
| s.store, | |
| s.settings, | |
| req.package_name, | |
| plan=plan, | |
| ) | |
| except (LLMError, Exception) as e: | |
| if not isinstance(e, LLMError): | |
| log.exception("code generation crashed") | |
| gen_id = s.logs.record( | |
| tester=t.name, | |
| kind="code", | |
| request=logged_request, | |
| error=str(e), | |
| duration_ms=int((time.monotonic() - start) * 1000), | |
| ) | |
| record_enhancement(s.logs, t.name, gen_id, "coding", brief) | |
| raise HTTPException(status_code=502, detail=f"generation failed (log id {gen_id}): {e}") from e | |
| gen_id = s.logs.record( | |
| tester=t.name, | |
| kind="code", | |
| request=logged_request, | |
| model=f"{result.model} via {result.backend}", | |
| sources=result.sources, | |
| attempts=[asdict(a) for a in result.attempts], | |
| output=json.dumps(result.files), | |
| valid=result.valid, | |
| validation=result.validation.to_dict(), | |
| duration_ms=int((time.monotonic() - start) * 1000), | |
| ) | |
| check = None | |
| if req.check and result.files: | |
| try: | |
| job = _start_check(s, t.name, gen_id, result.files, req.distro, req.language) | |
| check = {"id": job["id"], "status": job["status"], "report_url": _report_url(request, s, job["id"])} | |
| except HTTPException as e: | |
| check = {"error": e.detail} | |
| return CodeResponse( | |
| id=gen_id, | |
| check=check, | |
| routing=_routing_view(plan, result.attempts), | |
| cost=_book(s, t.name, gen_id, plan, result.attempts, brief), | |
| enhanced_prompt=_enhanced_view(brief), | |
| valid=result.valid, | |
| files=result.files, | |
| validation=result.validation.to_dict(), | |
| attempts=len(result.attempts), | |
| sources=result.sources, | |
| backend=result.backend, | |
| model=result.model, | |
| how_to_build=( | |
| f"Write the files into <ws>/src/, then: source /opt/ros/{req.distro}/setup.bash && cd <ws> && " | |
| "colcon build && source install/setup.bash && ros2 launch <package> <launch file>" | |
| ), | |
| ) | |
| def _report_url(request: Request, s: Services, run_id: str) -> str: | |
| base = s.settings.public_base_url or str(request.base_url).rstrip("/") | |
| return f"{base}/report/{run_id}?token={report_token(s.settings.report_secret, run_id)}" | |
| def _job_response(request: Request, s: Services, job: dict) -> JobResponse: | |
| base = s.settings.public_base_url or str(request.base_url).rstrip("/") | |
| return JobResponse( | |
| id=job["id"], | |
| kind=job["kind"], | |
| status=job["status"], | |
| report_url=_report_url(request, s, job["id"]), | |
| status_url=f"{base}/jobs/{job['id']}", | |
| ) | |
| def _start_check(s: Services, tester: str, gen_id: str, files: dict, distro: str, language: str) -> dict: | |
| return _submit( | |
| s, | |
| tester, | |
| "ros_check", | |
| {"files": files, "distro": distro}, | |
| run_id=new_run_id(), | |
| generation_id=gen_id, | |
| params={"distro": distro, "language": language}, | |
| ) | |
| def _scene_for(s: Services, t: Tester, run_id, generation_id, mjcf, robot) -> tuple[str, str | None, str | None]: | |
| """(mjcf, robot, generation_id) from exactly one of a run, a generation, or raw MJCF (validated here).""" | |
| if sum(x is not None for x in (run_id, generation_id, mjcf)) != 1: | |
| raise HTTPException(status_code=422, detail="give exactly one of run_id, generation_id or mjcf") | |
| if run_id: | |
| run = s.logs.get_run(run_id) | |
| if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": | |
| raise HTTPException(status_code=404, detail="no simulation run with that id for this key") | |
| if run.get("generation_id"): | |
| generation_id = run["generation_id"] | |
| else: | |
| key = ((run.get("output") or {}).get("files") or {}).get("scene") | |
| text = s.storage.read_text(key) if key and s.storage else None | |
| if not text: | |
| raise HTTPException(status_code=422, detail="that run's scene is not available (not finished?)") | |
| return text, run.get("robot"), None | |
| if generation_id: | |
| gen = s.logs.get(generation_id) | |
| if not gen or gen["tester"] != t.name or gen["kind"] != "simulate": | |
| raise HTTPException(status_code=404, detail="no scene generation with that id for this key") | |
| if not gen["valid"] or not gen["output"]: | |
| raise HTTPException(status_code=422, detail="that generation did not produce a valid scene") | |
| return gen["output"], gen["request"].get("robot"), generation_id | |
| check = validate_mjcf( | |
| mjcf, robot, s.settings.menagerie_dir, s.settings.mjcf_sim_seconds, s.settings.validate_timeout | |
| ) | |
| if not check.valid: | |
| raise HTTPException( | |
| status_code=422, detail={"message": "scene failed validation", "validation": check.to_dict()} | |
| ) | |
| return mjcf, robot, None | |
| def simulate_sweep(req: SweepRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| """Run a scene many times with randomized friction, masses and pushes (MuJoCo Warp on the GPU).""" | |
| mjcf, robot, gen_id = _scene_for(s, t, req.run_id, req.generation_id, req.mjcf, req.robot) | |
| params = { | |
| "n_worlds": req.n_worlds, | |
| "duration_s": req.duration_s, | |
| "controller": req.controller, | |
| "push_n": [0.0, req.push_max_n], | |
| "friction": list(req.friction), | |
| "object_mass": list(req.object_mass), | |
| "seed": req.seed, | |
| } | |
| job = _submit( | |
| s, | |
| t.name, | |
| "sweep", | |
| {"mjcf": mjcf, "robot": robot, **params}, | |
| run_id=new_run_id(), | |
| parent_id=req.run_id, | |
| generation_id=gen_id, | |
| robot=robot, | |
| duration_s=req.duration_s, | |
| controller=req.controller, | |
| params=params, | |
| ) | |
| return _job_response(request, s, job) | |
| def review_run(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| """Ask Cosmos Reason to watch the run's video (done automatically when AUTO_REVIEW is on).""" | |
| run = s.logs.get_run(run_id) | |
| if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": | |
| raise HTTPException(status_code=404, detail="no simulation run with that id for this key") | |
| run = refresh(s, run) | |
| try: | |
| job = start_review(s, run) | |
| except JobError as e: | |
| raise HTTPException(status_code=e.status, detail=str(e)) from e | |
| return _job_response(request, s, job) | |
| def export_dataset(req: ExportRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| """Turn finished runs into a LeRobot dataset (state, velocity, action, video, task) on the Hugging Face Hub.""" | |
| if s.storage is None: | |
| raise HTTPException(status_code=503, detail="storage is not configured (R2_* variables)") | |
| episodes, robots = [], set() | |
| for run_id in dict.fromkeys(req.run_ids): | |
| run = s.logs.get_run(run_id) | |
| if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": | |
| raise HTTPException(status_code=404, detail=f"no simulation run {run_id} for this key") | |
| run = refresh(s, run) | |
| files = (run.get("output") or {}).get("files") or {} | |
| if run["status"] != "COMPLETED" or "trajectory" not in files: | |
| raise HTTPException(status_code=422, detail=f"run {run_id} is not finished or has no trajectory") | |
| robots.add(run.get("robot")) | |
| task = "" | |
| if run.get("generation_id"): | |
| task = ((s.logs.get(run["generation_id"]) or {}).get("request") or {}).get("description", "") | |
| episodes.append( | |
| { | |
| "run_id": run_id, | |
| "task": task, | |
| "trajectory_url": s.storage.url(files["trajectory"]), | |
| "video_url": s.storage.url(files["video"]) if "video" in files else None, | |
| } | |
| ) | |
| if len(robots) > 1: | |
| raise HTTPException(status_code=422, detail="all runs in one dataset must use the same robot") | |
| robot = robots.pop() | |
| repo_id = f"{s.settings.hf_namespace}/rosdiff-{req.name}" | |
| job = _submit( | |
| s, | |
| t.name, | |
| "export", | |
| {"repo_id": repo_id, "robot": robot, "episodes": episodes, "private": req.private}, | |
| run_id=new_run_id(), | |
| robot=robot, | |
| params={"repo_id": repo_id, "runs": [e["run_id"] for e in episodes]}, | |
| ) | |
| return _job_response(request, s, job) | |
| def isaac_render(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| """Re-render a finished run photorealistically with Isaac Sim (limited to ISAAC_DAILY_LIMIT per day).""" | |
| run = s.logs.get_run(run_id) | |
| if not run or run["tester"] != t.name or (run.get("kind") or "sim") != "sim": | |
| raise HTTPException(status_code=404, detail="no simulation run with that id for this key") | |
| run = refresh(s, run) | |
| try: | |
| job = start_isaac(s, run) | |
| except JobError as e: | |
| raise HTTPException(status_code=e.status, detail=str(e)) from e | |
| return _job_response(request, s, job) | |
| def code_check(req: CheckRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| """Build, lint, test and launch a generated package in the ROS 2 sandbox. Returns at once; see report_url.""" | |
| gen = s.logs.get(req.generation_id) | |
| if not gen or gen["tester"] != t.name or gen["kind"] != "code": | |
| raise HTTPException(status_code=404, detail="no code generation with that id for this key") | |
| files = json.loads(gen["output"]) if gen.get("output") else {} | |
| if not files: | |
| raise HTTPException(status_code=422, detail="that generation produced no files") | |
| job = _start_check( | |
| s, t.name, req.generation_id, files, gen["request"]["distro"], gen["request"].get("language", "python") | |
| ) | |
| return _job_response(request, s, job) | |
| def simulate_run( | |
| req: RunRequest, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc) | |
| ) -> RunResponse: | |
| """Run a validated scene on a cloud GPU. Returns immediately; follow report_url.""" | |
| if (req.generation_id is None) == (req.mjcf is None): | |
| raise HTTPException(status_code=422, detail="give exactly one of generation_id or mjcf") | |
| if req.duration_s > s.settings.run_max_duration_s: | |
| raise HTTPException(status_code=422, detail=f"duration_s must be <= {s.settings.run_max_duration_s}") | |
| if req.generation_id: | |
| gen = s.logs.get(req.generation_id) | |
| if not gen or gen["tester"] != t.name or gen["kind"] != "simulate": | |
| raise HTTPException(status_code=404, detail="no scene generation with that id for this key") | |
| if not gen["valid"] or not gen["output"]: | |
| raise HTTPException(status_code=422, detail="that generation did not produce a valid scene") | |
| mjcf, robot = gen["output"], gen["request"].get("robot") | |
| else: | |
| mjcf, robot = req.mjcf, req.robot | |
| check = validate_mjcf( | |
| mjcf, robot, s.settings.menagerie_dir, s.settings.mjcf_sim_seconds, s.settings.validate_timeout | |
| ) | |
| if not check.valid: | |
| raise HTTPException( | |
| status_code=422, detail={"message": "scene failed validation", "validation": check.to_dict()} | |
| ) | |
| # Usage visibility, not a cap: warn loudly when a key starts runs in a burst. | |
| since = (dt.datetime.now(dt.timezone.utc) - dt.timedelta(minutes=10)).isoformat(timespec="seconds") | |
| recent = s.logs.runs_since(t.name, since) | |
| if recent >= s.settings.run_burst_warn: | |
| log.warning("tester %s has started %d runs in the last 10 minutes", t.name, recent + 1) | |
| run_id = new_run_id() | |
| job_input = {"mjcf": mjcf, "robot": robot, "duration_s": req.duration_s, "controller": req.controller} | |
| job = _submit( | |
| s, | |
| t.name, | |
| "sim", | |
| job_input, | |
| run_id=run_id, | |
| generation_id=req.generation_id, | |
| robot=robot, | |
| duration_s=req.duration_s, | |
| controller=req.controller, | |
| ) | |
| base = s.settings.public_base_url or str(request.base_url).rstrip("/") | |
| return RunResponse( | |
| run_id=run_id, | |
| status=job["status"], | |
| report_url=_report_url(request, s, run_id), | |
| status_url=f"{base}/simulate/run/{run_id}", | |
| ) | |
| def run_status(run_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)): | |
| run = s.logs.get_run(run_id) | |
| if not run or run["tester"] != t.name: | |
| raise HTTPException(status_code=404, detail="no run with that id for this key") | |
| run = refresh(s, run) | |
| return RunStatus( | |
| run_id=run_id, | |
| status=run["status"], | |
| error=run.get("error"), | |
| summary=(run.get("output") or {}).get("summary"), | |
| report_url=_report_url(request, s, run_id), | |
| ) | |
| def run_report(run_id: str, request: Request, token: str | None = None, s: Services = Depends(svc)): | |
| """Browser page for any job; its link carries its own token, so no API key is needed.""" | |
| run = s.logs.get_run(run_id) | |
| if not run or not check_token(s.settings.report_secret, run_id, token): | |
| raise HTTPException(status_code=404, detail="report not found") | |
| run = refresh(s, run) | |
| children = [refresh(s, c) for c in s.logs.children(run_id)] | |
| return HTMLResponse( | |
| render_report(run, s.storage, s.settings.foxglove_org, children, link=lambda i: _report_url(request, s, i)), | |
| headers={"Cache-Control": "no-store", "Referrer-Policy": "no-referrer"}, | |
| ) | |
| def my_jobs(request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| """Your recent jobs, and how much of this month's compute budget is used (all testers together).""" | |
| jobs = s.logs.list_runs(t.name, 30) | |
| return { | |
| "budget": { | |
| "month_spend_usd_estimate": month_spend(s.settings, s.logs), | |
| "monthly_limit_usd": s.settings.monthly_credit_usd or None, | |
| "mine": compute_usage(s.settings, s.logs, t.name), | |
| }, | |
| "jobs": [ | |
| { | |
| "id": j["id"], | |
| "kind": j.get("kind") or "sim", | |
| "status": j["status"], | |
| "created_at": j["created_at"], | |
| "parent_id": j.get("parent_id"), | |
| "report_url": _report_url(request, s, j["id"]), | |
| } | |
| for j in jobs | |
| ], | |
| } | |
| def job_status(job_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| job = s.logs.get_run(job_id) | |
| if not job or job["tester"] != t.name: | |
| raise HTTPException(status_code=404, detail="no job with that id for this key") | |
| job = refresh(s, job) | |
| output = job.get("output") or {} | |
| return { | |
| "id": job_id, | |
| "kind": job.get("kind") or "sim", | |
| "status": job["status"], | |
| "error": job.get("error"), | |
| "parent_id": job.get("parent_id"), | |
| "summary": output.get("summary"), | |
| "files": sorted((output.get("files") or {}).keys()), | |
| "report_url": _report_url(request, s, job_id), | |
| "media": _media(s, job), | |
| } | |
| def _media(s: Services, job: dict) -> dict: | |
| """Presigned links for what the web app shows inline (videos), once a job has finished.""" | |
| files = (job.get("output") or {}).get("files") or {} | |
| if job["status"] != "COMPLETED" or s.storage is None: | |
| return {} | |
| names = {"video": "video", "video_webm": "video_webm", "isaac.mp4": "isaac_video"} | |
| return {label: s.storage.url(files[key]) for key, label in names.items() if key in files} | |
| def history(t: Tester = Depends(tester), s: Services = Depends(svc), limit: int = 60) -> dict: | |
| """Your recent scenes and packages, newest first (for the web app's sidebar).""" | |
| items = [] | |
| for g in s.logs.list(tester=t.name, limit=min(max(limit, 1), 200)): | |
| req = g.get("request") or {} | |
| text = req.get("description") or req.get("request") or "" | |
| items.append( | |
| { | |
| "id": g["id"], | |
| "kind": g["kind"], | |
| "created_at": g["created_at"], | |
| "valid": bool(g["valid"]) if g.get("valid") is not None else None, | |
| "title": text[:90], | |
| "robot": req.get("robot"), | |
| "distro": req.get("distro"), | |
| "language": req.get("language"), | |
| "rating": g.get("tester_rating"), | |
| } | |
| ) | |
| return {"tester": t.name, "items": items} | |
| def generation(gen_id: str, request: Request, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| """One of your scenes or packages, with its runs and checks (to reopen it in the web app).""" | |
| g = s.logs.get(gen_id) | |
| if not g or g["tester"] != t.name: | |
| raise HTTPException(status_code=404, detail="no generation with that id for this key") | |
| req = g.get("request") or {} | |
| output = g.get("output") | |
| if g["kind"] == "code": | |
| output = json.loads(output) if output else {} | |
| jobs = [ | |
| { | |
| "id": j["id"], | |
| "kind": j.get("kind") or "sim", | |
| "status": j["status"], | |
| "created_at": j["created_at"], | |
| "report_url": _report_url(request, s, j["id"]), | |
| } | |
| for j in s.logs.list_runs(t.name, 200) | |
| if j.get("generation_id") == gen_id | |
| ] | |
| return { | |
| "id": gen_id, | |
| "kind": g["kind"], | |
| "created_at": g["created_at"], | |
| "request": req, | |
| "valid": bool(g["valid"]) if g.get("valid") is not None else None, | |
| "validation": g.get("validation") or {}, | |
| "output": output, | |
| "error": g.get("error"), | |
| "model": g.get("model"), | |
| "rating": g.get("tester_rating"), | |
| "jobs": list(reversed(jobs)), | |
| } | |
| def code_zip(gen_id: str, t: Tester = Depends(tester), s: Services = Depends(svc)) -> Response: | |
| """The generated package as a .zip, ready to unpack into a colcon workspace's src/.""" | |
| import io | |
| import zipfile | |
| g = s.logs.get(gen_id) | |
| if not g or g["tester"] != t.name or g["kind"] != "code" or not g.get("output"): | |
| raise HTTPException(status_code=404, detail="no code generation with that id for this key") | |
| buf = io.BytesIO() | |
| with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z: | |
| for path, content in json.loads(g["output"]).items(): | |
| z.writestr(path, content) | |
| return Response( | |
| buf.getvalue(), | |
| media_type="application/zip", | |
| headers={"Content-Disposition": f'attachment; filename="rosdiff-{gen_id}.zip"'}, | |
| ) | |
| def web_app() -> HTMLResponse: | |
| """The web app for testers (api/static/app).""" | |
| return HTMLResponse( | |
| (Path(__file__).with_name("static") / "app" / "index.html").read_text(), | |
| headers={"Cache-Control": "no-cache"}, | |
| ) | |
| def feedback_item(req: FeedbackItem, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| """๐/๐ (plus an optional note) on a generation or a job, stored with what RosDiff did for it.""" | |
| if req.target_type == "generation": | |
| g = s.logs.get(req.target_id) | |
| if not g or g["tester"] != t.name: | |
| raise HTTPException(status_code=404, detail="nothing with that id for this key") | |
| validation = g.get("validation") or {} | |
| context = { | |
| "kind": g["kind"], | |
| "valid": bool(g["valid"]) if g.get("valid") is not None else None, | |
| "model": g.get("model"), | |
| "routing": (g.get("request") or {}).get("routing"), | |
| "error": g.get("error"), | |
| "validation_error": validation.get("error") or validation.get("errors"), | |
| } | |
| s.logs.set_feedback(req.target_id, t.name, "works" if req.rating == "helpful" else "broken", req.note) | |
| else: | |
| j = s.logs.get_run(req.target_id) | |
| if not j or j["tester"] != t.name: | |
| raise HTTPException(status_code=404, detail="nothing with that id for this key") | |
| kind = j.get("kind") or "sim" | |
| context = { | |
| "kind": kind, | |
| "status": j["status"], | |
| "error": j.get("error"), | |
| "generation_id": j.get("generation_id"), | |
| "parent_id": j.get("parent_id"), | |
| "simulator": {"sim": "mujoco", "sweep": "mujoco_warp", "isaac": "isaac_sim"}.get(kind), | |
| "hardware": j.get("hardware"), | |
| "params": j.get("params"), | |
| "headline": ((j.get("output") or {}).get("summary") or {}).get("headline"), | |
| } | |
| s.logs.record_feedback(t.name, req.target_type, req.target_id, req.rating, req.note, context) | |
| return {"ok": True} | |
| def my_usage(t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| """Your AI budget in rupees (and the team's, for admins). GPU jobs run on the separate compute credit.""" | |
| is_owner = t.name == "owner" | |
| out = { | |
| "tester": t.name, | |
| "is_owner": is_owner, | |
| "ai": budget.summary(s.logs, s.settings, t.name), | |
| "compute": { | |
| **compute_usage(s.settings, s.logs, t.name), | |
| "month_spend_usd_estimate": month_spend(s.settings, s.logs), | |
| "monthly_limit_usd": s.settings.monthly_credit_usd or None, | |
| "note": "Simulations and checks run on the cloud GPU credit, split equally between testers; " | |
| "it is separate from the rupee AI budget.", | |
| }, | |
| "admin": is_owner or t.name in s.settings.admin_testers, | |
| } | |
| if out["admin"]: | |
| names = [x.name for x in s.allowlist.testers()] | |
| out["team"] = budget.team(s.logs, s.settings, names) | |
| for row in out["team"]: | |
| row["gpu"] = compute_usage(s.settings, s.logs, row["tester"])["used"] | |
| return out | |
| def feedback(gen_id: str, req: FeedbackRequest, t: Tester = Depends(tester), s: Services = Depends(svc)) -> dict: | |
| """Tell us whether a generation actually worked when you ran it.""" | |
| if not s.logs.set_feedback(gen_id, t.name, req.rating, req.note): | |
| raise HTTPException(status_code=404, detail="no generation with that id for this key") | |
| return {"id": gen_id, "rating": req.rating} | |
| return app | |
| app = create_app() | |