worker/handler.py

"""Simulation job: run one validated scene and upload the results to R2 (Modal or RunPod).

Job input (sent by POST /simulate/run)::

    {"run_id": "...", "mjcf": "<mujoco>...</mujoco>", "robot": "unitree_g1" | null,
     "duration_s": 10, "controller": "hold" | "passive"}

Output::

    {"run_id": ..., "summary": {...}, "files": {"mcap": "runs/<id>/run.mcap", "video": ..., ...},
     "worker_seconds": 41.2}

Entry points: ``handler`` for RunPod (``python -m worker.handler`` in worker/Dockerfile) and
``run_simulation`` in worker/modal_app.py for Modal. Environment on the worker: R2_ACCOUNT_ID,
R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET, MENAGERIE_DIR (baked into the image), MUJOCO_GL=egl.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import time
import traceback
from pathlib import Path

from api.mjcf_check import SCENE_FILE, prepare_scene_dir, static_check
from api.storage import LocalStorage, R2Config, R2Storage, Storage, run_prefix
from worker.sim_runner import RunConfig, config_for, simulate

MAX_DURATION_S = 60.0


def _storage() -> Storage:
    config = R2Config.from_env()
    if config:
        return R2Storage(config)
    local = os.environ.get("ROSDIFF_LOCAL_STORAGE")
    if local:
        return LocalStorage(Path(local))
    raise RuntimeError("R2 is not configured (R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET)")


def run_job(
    job_input: dict, storage: Storage | None = None, menagerie_dir: Path | None = None, render: bool = True
) -> dict:
    started = time.monotonic()
    run_id = str(job_input["run_id"])
    if not run_id.isalnum():
        raise ValueError("run_id must be alphanumeric")
    mjcf = job_input["mjcf"]
    robot = job_input.get("robot")
    duration = min(float(job_input.get("duration_s", 10.0)), MAX_DURATION_S)
    controller = job_input.get("controller", "hold")

    # Defence in depth: the API validated this scene already.
    failed = static_check(mjcf, robot)
    if failed:
        raise ValueError(f"scene rejected: {failed.error}")

    storage = storage or _storage()
    menagerie = Path(menagerie_dir or os.environ.get("MENAGERIE_DIR", "/opt/mujoco_menagerie"))
    scene_dir = prepare_scene_dir(mjcf, robot, menagerie)
    out = Path(tempfile.mkdtemp(prefix=f"run_{run_id}_"))
    try:
        config = config_for(robot, duration, controller, render=render)
        summary = simulate(scene_dir / SCENE_FILE, out, config)
        if render and summary.get("render_error") and os.environ.get("MUJOCO_GL") == "egl":
            summary = _rerun_with_osmesa(scene_dir / SCENE_FILE, out, config, robot, summary["render_error"])
        prefix = run_prefix(run_id)
        (out / "scene.xml").write_text(mjcf)
        files = {}
        for kind, name in (
            ("mcap", "run.mcap"),
            ("video", "video.mp4"),
            ("video_webm", "video.webm"),
            ("summary", "summary.json"),
            ("trajectory", "trajectory.npz"),
            ("scene", "scene.xml"),
        ):
            if (out / name).is_file():
                storage.upload(out / name, f"{prefix}/{name}")
                files[kind] = f"{prefix}/{name}"
        return {
            "run_id": run_id,
            "summary": summary,
            "files": files,
            "worker_seconds": round(time.monotonic() - started, 1),
        }
    finally:
        shutil.rmtree(scene_dir, ignore_errors=True)
        shutil.rmtree(out, ignore_errors=True)


def _rerun_with_osmesa(scene: Path, out: Path, config: RunConfig, robot: str | None, egl_error: str) -> dict:
    """GPU (EGL) rendering failed on this machine: redo the run with CPU rendering in a fresh process."""
    import json
    import subprocess
    import sys

    env = {**os.environ, "MUJOCO_GL": "osmesa", "PYOPENGL_PLATFORM": "osmesa"}
    cmd = [
        sys.executable,
        "-m",
        "worker.sim_runner",
        str(scene),
        str(out),
        "--duration",
        str(config.duration_s),
        "--controller",
        config.controller,
        "--robot",
        robot or "",
    ]
    subprocess.run(cmd, env=env, check=False, capture_output=True, timeout=3600)
    summary = json.loads((out / "summary.json").read_text())
    summary["egl_error"] = egl_error
    return summary


def _functions() -> dict:
    from worker.sweep import run_sweep_job

    return {"run_simulation": run_job, "sweep": run_sweep_job}


def handler(job: dict) -> dict:
    """RunPod entry point. Errors are returned (not raised) so they are recorded as the job's output."""
    try:
        job_input = dict(job["input"])
        fn = _functions()[job_input.pop("function", "run_simulation")]
        return fn(job_input)
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}


if __name__ == "__main__":
    import runpod

    runpod.serverless.start({"handler": handler})