Training script for lerobot using modal

modal_deploy.py

import modal

# how to run:
# modal run modal_deploy.py
# modal run --detach modal_deploy.py

image = (
    modal.Image
    .debian_slim()
    .pip_install_from_pyproject("pyproject.toml")
    .add_local_python_source("lerobot")

)

app = modal.App(name="lerobot-train", image=image)

@app.function(gpu="A10G", volumes={"/out": modal.Volume.from_name("lerobot")}, timeout=86400)
def train_lerobot():
    import subprocess
    command = [
        "python", "lerobot/scripts/train.py",
        f"--dataset.repo_id=shreyasgite/so100",
        "--policy.type=act",
        "--output_dir=/out/so100",
        "--job_name=act_so100_base_env",
        "--device=cuda",
        # "--wandb.enable=true"
    ]

    with subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1) as process:
        for line in process.stdout:
            print(line, end='')  # Print each line as it is received
        process.wait()  # Wait for the process to complete




@app.local_entrypoint()
def main():
    return train_lerobot.remote()