Spaces:

chandrakiran06
/

rosdiff

Paused

Chandra Kiran

Keep the default Modal deploy free of stored data

916d469 unverified 27 days ago

8.47 kB

"""Modal deployment of the simulation worker (the default SIM_BACKEND).
Deploy from the repository root (the image reuses api/ and worker/)::
    modal secret create rosdiff-r2 R2_ACCOUNT_ID=... R2_ACCESS_KEY_ID=... R2_SECRET_ACCESS_KEY=... R2_BUCKET=...
    modal deploy -m worker.modal_app
These environment variables are read at deploy time:
    MODAL_APP              app name (default rosdiff-sim; the API's MODAL_APP must match)
    MODAL_GPU              GPU type (default T4, the cheapest); "none" renders on CPU instead
    MODAL_MAX_CONTAINERS   runs executing at once (default 2); more runs wait in Modal's queue
    RUN_EXECUTION_TIMEOUT_S  Modal stops a run after this long (default 600)
    MODAL_R2_SECRET        name of the Modal secret holding the R2 credentials (default rosdiff-r2)
"""

from __future__ import annotations

import os
from pathlib import Path

import modal

APP_NAME = os.environ.get("MODAL_APP", "rosdiff-sim")
_gpu = os.environ.get("MODAL_GPU", "T4").strip()
GPU = None if _gpu.lower() in ("", "none", "cpu") else _gpu
GL = "egl" if GPU else "osmesa"
MENAGERIE_DIR = "/opt/mujoco_menagerie"
ROBOT_DIRS = ("unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda")  # api.robots.ROBOTS

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("git", "libglvnd0", "libgl1", "libegl1", "libgles2", "libosmesa6")
    .run_commands(
        # Tell libglvnd where the NVIDIA EGL implementation is (the driver is provided by Modal's GPU hosts).
        "mkdir -p /usr/share/glvnd/egl_vendor.d && printf "
        '\'{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}\\n\' '
        "> /usr/share/glvnd/egl_vendor.d/10_nvidia.json",
        # Robot models (only the ones generated scenes may include).
        "git clone --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git "
        f"{MENAGERIE_DIR} && cd {MENAGERIE_DIR} && git sparse-checkout set {' '.join(ROBOT_DIRS)} && rm -rf .git",
    )
    .pip_install_from_requirements(str(Path(__file__).with_name("requirements.txt")))
    .env({"MENAGERIE_DIR": MENAGERIE_DIR, "MUJOCO_GL": GL, "PYOPENGL_PLATFORM": GL, "PYTHONUNBUFFERED": "1"})
    .add_local_python_source("api", "worker")
)

app = modal.App(APP_NAME)


@app.function(
    image=image,
    gpu=GPU,
    timeout=int(os.environ.get("RUN_EXECUTION_TIMEOUT_S", "600")),
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,  # stop paying for an idle container quickly (runs are sporadic)
    secrets=[modal.Secret.from_name(os.environ.get("MODAL_R2_SECRET", "rosdiff-r2"))],
)
def run_simulation(job_input: dict) -> dict:
    from worker.handler import handler

    return handler({"input": {**job_input, "function": "run_simulation"}})


@app.function(
    image=image,
    gpu=GPU,
    timeout=900,
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,
)
def sweep(job_input: dict) -> dict:
    """Robustness sweep with MuJoCo Warp; results return inline (no storage secret needed)."""
    from worker.handler import handler

    return handler({"input": {**job_input, "function": "sweep"}})


# --------------------------------------------------------------------------- ROS 2 sandbox (/code/check)


def _ros_image(distro: str) -> modal.Image:
    # Same Dockerfile CI builds and tests; Modal's own Python is added alongside the image's (the checks call
    # ROS through bash with a clean environment, so the two never mix).
    return modal.Image.from_dockerfile(
        str(Path(__file__).parent / "ros" / "Dockerfile"),
        context_dir=str(Path(__file__).parent.parent),
        build_args={"ROS_DISTRO": distro},
        add_python="3.11",
    ).add_local_python_source("worker")


# Generated code is untrusted: no network, no secrets. Results come back inline and the API stores them.
_ROS_SANDBOX = dict(
    cpu=2.0,
    memory=4096,
    timeout=900,
    max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
    scaledown_window=20,
    block_network=True,
)


@app.function(image=_ros_image("humble"), **_ROS_SANDBOX)
def ros_check_humble(job_input: dict) -> dict:
    return _ros_check(job_input, "humble")


@app.function(image=_ros_image("jazzy"), **_ROS_SANDBOX)
def ros_check_jazzy(job_input: dict) -> dict:
    return _ros_check(job_input, "jazzy")


def _ros_check(job_input: dict, distro: str) -> dict:
    import traceback

    from worker.ros_runner import run_check_job

    try:
        return run_check_job(job_input, distro)
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}


# --------------------------------------------------------------------------- Cosmos Reason review (L4)
# Opt-in (ENABLE_COSMOS=1 at deploy): it needs the rosdiff-hf secret and keeps the ~16 GB model in a Modal
# Volume, and volume storage is billed even when the spend limit is $0. Without it nothing is stored on Modal.
ENABLE_COSMOS = os.environ.get("ENABLE_COSMOS", "0") == "1"

if ENABLE_COSMOS:
    COSMOS_MODEL = os.environ.get("COSMOS_MODEL", "nvidia/Cosmos-Reason1-7B")
    models = modal.Volume.from_name("rosdiff-models", create_if_missing=True)
    cosmos_image = (
        modal.Image.debian_slim(python_version="3.11")
        .pip_install(
            "torch==2.7.1",
            "torchvision==0.22.1",
            "transformers>=4.51,<4.60",
            "accelerate>=1.6",
            "qwen-vl-utils>=0.0.8",
            "pillow",
            "imageio-ffmpeg==0.6.0",
            "huggingface_hub>=0.30",
        )
        .env({"HF_HOME": "/models/hf", "COSMOS_MODEL": COSMOS_MODEL})
        .add_local_python_source("worker")
    )

    @app.function(
        image=cosmos_image,
        volumes={"/models": models},
        timeout=3600,
        secrets=[modal.Secret.from_name(os.environ.get("MODAL_HF_SECRET", "rosdiff-hf"))],
    )
    def download_cosmos() -> str:
        """One-off, CPU only: `modal run -m worker.modal_app::download_cosmos` fills the models volume."""
        from huggingface_hub import snapshot_download

        path = snapshot_download(COSMOS_MODEL, local_dir=f"/models/{COSMOS_MODEL}")
        models.commit()
        return path

    @app.cls(
        image=cosmos_image,
        gpu="L4",
        volumes={"/models": models},
        timeout=600,
        max_containers=1,
        scaledown_window=20,
    )
    class Reviewer:
        @modal.enter()
        def load(self):
            from worker.cosmos_review import CosmosModel

            self.model = CosmosModel(f"/models/{COSMOS_MODEL}")

        @modal.method()
        def review(self, job_input: dict) -> dict:
            import traceback

            from worker.cosmos_review import run_review_job

            try:
                return run_review_job({**job_input, "model": COSMOS_MODEL}, self.model)
            except Exception as e:
                return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}

    @app.function(image=cosmos_image, timeout=600)
    def review_video(job_input: dict) -> dict:
        """Entry point the API spawns; hands the job to the (GPU, model-loaded) Reviewer."""
        return Reviewer().review.remote(job_input)


# --------------------------------------------------------------------------- LeRobot dataset export (CPU)

lerobot_image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("ffmpeg")
    .pip_install("lerobot==0.4.4", "imageio-ffmpeg==0.6.0")
    .add_local_python_source("worker")
)


@app.function(
    image=lerobot_image,
    cpu=2.0,
    memory=4096,
    timeout=900,
    max_containers=1,
    scaledown_window=20,
    # Only with MODAL_HF_SECRET=rosdiff-hf at deploy (a secret holding HF_TOKEN): datasets go to the Hub.
    # Without it, each dataset comes back as a zip stored in R2 next to the job.
    secrets=[modal.Secret.from_name(os.environ["MODAL_HF_SECRET"])] if os.environ.get("MODAL_HF_SECRET") else [],
)
def export_lerobot(job_input: dict) -> dict:
    """Runs -> LeRobot dataset, pushed to the Hub with the secret's HF_TOKEN (else returned as a zip)."""
    import traceback

    from worker.lerobot_export import run_export_job

    try:
        return run_export_job(job_input)
    except Exception as e:
        return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}

	"""Modal deployment of the simulation worker (the default SIM_BACKEND).

	Deploy from the repository root (the image reuses api/ and worker/)::

	modal secret create rosdiff-r2 R2_ACCOUNT_ID=... R2_ACCESS_KEY_ID=... R2_SECRET_ACCESS_KEY=... R2_BUCKET=...
	modal deploy -m worker.modal_app

	These environment variables are read at deploy time:

	MODAL_APP app name (default rosdiff-sim; the API's MODAL_APP must match)
	MODAL_GPU GPU type (default T4, the cheapest); "none" renders on CPU instead
	MODAL_MAX_CONTAINERS runs executing at once (default 2); more runs wait in Modal's queue
	RUN_EXECUTION_TIMEOUT_S Modal stops a run after this long (default 600)
	MODAL_R2_SECRET name of the Modal secret holding the R2 credentials (default rosdiff-r2)
	"""

	from __future__ import annotations

	import os
	from pathlib import Path

	import modal

	APP_NAME = os.environ.get("MODAL_APP", "rosdiff-sim")
	_gpu = os.environ.get("MODAL_GPU", "T4").strip()
	GPU = None if _gpu.lower() in ("", "none", "cpu") else _gpu
	GL = "egl" if GPU else "osmesa"
	MENAGERIE_DIR = "/opt/mujoco_menagerie"
	ROBOT_DIRS = ("unitree_g1", "unitree_h1", "unitree_go2", "franka_emika_panda") # api.robots.ROBOTS

	image = (
	modal.Image.debian_slim(python_version="3.11")
	.apt_install("git", "libglvnd0", "libgl1", "libegl1", "libgles2", "libosmesa6")
	.run_commands(
	# Tell libglvnd where the NVIDIA EGL implementation is (the driver is provided by Modal's GPU hosts).
	"mkdir -p /usr/share/glvnd/egl_vendor.d && printf "
	'\'{"file_format_version":"1.0.0","ICD":{"library_path":"libEGL_nvidia.so.0"}}\\n\' '
	"> /usr/share/glvnd/egl_vendor.d/10_nvidia.json",
	# Robot models (only the ones generated scenes may include).
	"git clone --depth 1 --filter=blob:none --sparse https://github.com/google-deepmind/mujoco_menagerie.git "
	f"{MENAGERIE_DIR} && cd {MENAGERIE_DIR} && git sparse-checkout set {' '.join(ROBOT_DIRS)} && rm -rf .git",
	)
	.pip_install_from_requirements(str(Path(__file__).with_name("requirements.txt")))
	.env({"MENAGERIE_DIR": MENAGERIE_DIR, "MUJOCO_GL": GL, "PYOPENGL_PLATFORM": GL, "PYTHONUNBUFFERED": "1"})
	.add_local_python_source("api", "worker")
	)

	app = modal.App(APP_NAME)


	@app.function(
	image=image,
	gpu=GPU,
	timeout=int(os.environ.get("RUN_EXECUTION_TIMEOUT_S", "600")),
	max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
	scaledown_window=20, # stop paying for an idle container quickly (runs are sporadic)
	secrets=[modal.Secret.from_name(os.environ.get("MODAL_R2_SECRET", "rosdiff-r2"))],
	)
	def run_simulation(job_input: dict) -> dict:
	from worker.handler import handler

	return handler({"input": {**job_input, "function": "run_simulation"}})


	@app.function(
	image=image,
	gpu=GPU,
	timeout=900,
	max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
	scaledown_window=20,
	)
	def sweep(job_input: dict) -> dict:
	"""Robustness sweep with MuJoCo Warp; results return inline (no storage secret needed)."""
	from worker.handler import handler

	return handler({"input": {**job_input, "function": "sweep"}})


	# --------------------------------------------------------------------------- ROS 2 sandbox (/code/check)


	def _ros_image(distro: str) -> modal.Image:
	# Same Dockerfile CI builds and tests; Modal's own Python is added alongside the image's (the checks call
	# ROS through bash with a clean environment, so the two never mix).
	return modal.Image.from_dockerfile(
	str(Path(__file__).parent / "ros" / "Dockerfile"),
	context_dir=str(Path(__file__).parent.parent),
	build_args={"ROS_DISTRO": distro},
	add_python="3.11",
	).add_local_python_source("worker")


	# Generated code is untrusted: no network, no secrets. Results come back inline and the API stores them.
	_ROS_SANDBOX = dict(
	cpu=2.0,
	memory=4096,
	timeout=900,
	max_containers=int(os.environ.get("MODAL_MAX_CONTAINERS", "2")),
	scaledown_window=20,
	block_network=True,
	)


	@app.function(image=_ros_image("humble"), **_ROS_SANDBOX)
	def ros_check_humble(job_input: dict) -> dict:
	return _ros_check(job_input, "humble")


	@app.function(image=_ros_image("jazzy"), **_ROS_SANDBOX)
	def ros_check_jazzy(job_input: dict) -> dict:
	return _ros_check(job_input, "jazzy")


	def _ros_check(job_input: dict, distro: str) -> dict:
	import traceback

	from worker.ros_runner import run_check_job

	try:
	return run_check_job(job_input, distro)
	except Exception as e:
	return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}


	# --------------------------------------------------------------------------- Cosmos Reason review (L4)
	# Opt-in (ENABLE_COSMOS=1 at deploy): it needs the rosdiff-hf secret and keeps the ~16 GB model in a Modal
	# Volume, and volume storage is billed even when the spend limit is $0. Without it nothing is stored on Modal.
	ENABLE_COSMOS = os.environ.get("ENABLE_COSMOS", "0") == "1"

	if ENABLE_COSMOS:
	COSMOS_MODEL = os.environ.get("COSMOS_MODEL", "nvidia/Cosmos-Reason1-7B")
	models = modal.Volume.from_name("rosdiff-models", create_if_missing=True)
	cosmos_image = (
	modal.Image.debian_slim(python_version="3.11")
	.pip_install(
	"torch==2.7.1",
	"torchvision==0.22.1",
	"transformers>=4.51,<4.60",
	"accelerate>=1.6",
	"qwen-vl-utils>=0.0.8",
	"pillow",
	"imageio-ffmpeg==0.6.0",
	"huggingface_hub>=0.30",
	)
	.env({"HF_HOME": "/models/hf", "COSMOS_MODEL": COSMOS_MODEL})
	.add_local_python_source("worker")
	)

	@app.function(
	image=cosmos_image,
	volumes={"/models": models},
	timeout=3600,
	secrets=[modal.Secret.from_name(os.environ.get("MODAL_HF_SECRET", "rosdiff-hf"))],
	)
	def download_cosmos() -> str:
	"""One-off, CPU only: `modal run -m worker.modal_app::download_cosmos` fills the models volume."""
	from huggingface_hub import snapshot_download

	path = snapshot_download(COSMOS_MODEL, local_dir=f"/models/{COSMOS_MODEL}")
	models.commit()
	return path

	@app.cls(
	image=cosmos_image,
	gpu="L4",
	volumes={"/models": models},
	timeout=600,
	max_containers=1,
	scaledown_window=20,
	)
	class Reviewer:
	@modal.enter()
	def load(self):
	from worker.cosmos_review import CosmosModel

	self.model = CosmosModel(f"/models/{COSMOS_MODEL}")

	@modal.method()
	def review(self, job_input: dict) -> dict:
	import traceback

	from worker.cosmos_review import run_review_job

	try:
	return run_review_job({**job_input, "model": COSMOS_MODEL}, self.model)
	except Exception as e:
	return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}

	@app.function(image=cosmos_image, timeout=600)
	def review_video(job_input: dict) -> dict:
	"""Entry point the API spawns; hands the job to the (GPU, model-loaded) Reviewer."""
	return Reviewer().review.remote(job_input)


	# --------------------------------------------------------------------------- LeRobot dataset export (CPU)

	lerobot_image = (
	modal.Image.debian_slim(python_version="3.11")
	.apt_install("ffmpeg")
	.pip_install("lerobot==0.4.4", "imageio-ffmpeg==0.6.0")
	.add_local_python_source("worker")
	)


	@app.function(
	image=lerobot_image,
	cpu=2.0,
	memory=4096,
	timeout=900,
	max_containers=1,
	scaledown_window=20,
	# Only with MODAL_HF_SECRET=rosdiff-hf at deploy (a secret holding HF_TOKEN): datasets go to the Hub.
	# Without it, each dataset comes back as a zip stored in R2 next to the job.
	secrets=[modal.Secret.from_name(os.environ["MODAL_HF_SECRET"])] if os.environ.get("MODAL_HF_SECRET") else [],
	)
	def export_lerobot(job_input: dict) -> dict:
	"""Runs -> LeRobot dataset, pushed to the Hub with the secret's HF_TOKEN (else returned as a zip)."""
	import traceback

	from worker.lerobot_export import run_export_job

	try:
	return run_export_job(job_input)
	except Exception as e:
	return {"error": f"{type(e).__name__}: {e}", "traceback": traceback.format_exc()[-2000:]}