rosdiff / api /main.py
Chandrakiran-08
Deploy to HF Spaces
8501103
Raw History Blame Contribute Delete
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")
@dataclass
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
@app.on_event("startup")
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")
@app.get("/health")
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,
}
@app.post("/simulate/generate", response_model=SimulateResponse)
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),
)
@app.post("/code/generate", response_model=CodeResponse)
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
@app.post("/simulate/sweep", response_model=JobResponse)
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)
@app.post("/simulate/run/{run_id}/review", response_model=JobResponse)
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)
@app.post("/datasets/export", response_model=JobResponse)
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)
@app.post("/simulate/run/{run_id}/isaac", response_model=JobResponse)
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)
@app.post("/code/check", response_model=JobResponse)
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)
@app.post("/simulate/run", response_model=RunResponse)
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}",
)
@app.get("/simulate/run/{run_id}", response_model=RunStatus)
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),
)
@app.get("/report/{run_id}", response_class=HTMLResponse)
@app.get("/simulate/report/{run_id}", response_class=HTMLResponse)
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"},
)
@app.get("/jobs")
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
],
}
@app.get("/jobs/{job_id}")
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}
@app.get("/me/history")
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}
@app.get("/generations/{gen_id}")
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)),
}
@app.get("/code/{gen_id}/zip")
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"'},
)
@app.get("/", response_class=HTMLResponse, include_in_schema=False)
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"},
)
@app.post("/feedback")
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}
@app.get("/me/usage")
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
@app.post("/feedback/{gen_id}")
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()