Pack whole current project folder, send to Modal, and execute a single compute-heavy CLI command

modal_runner.py

"""
Run repo commands in Modal with streamed logs and optional synthetic timeout.

Usage:
    uv run modal run runner.py --cmd '<bash command>'

Examples:
    # Run the main export command remotely
    uv run modal run runner.py --cmd \
      'uv run -m src.compute_heavy_entrypoint'

Notes:
    - The Modal function timeout is configured via FUNCTION_TIMEOUT_S.
    - cmd_timeout_s must be 0 (disabled) or less than FUNCTION_TIMEOUT_S.
    - Repository sync excludes .git, .venv, and common cache directories.
"""

from pathlib import Path
import subprocess
import time
import modal

app = modal.App("modal-runner")
FUNCTION_TIMEOUT_S = 12 * 60 * 60

SKIP_PATH_SEGMENTS = {
    ".git",
    ".venv",
    "__pycache__",
    ".mypy_cache",
    ".pytest_cache",
}


def should_ignore(path: Path) -> bool:
    return any(segment in SKIP_PATH_SEGMENTS for segment in path.parts)


def validate_cmd_timeout(cmd_timeout_s: int) -> None:
    if cmd_timeout_s < 0:
        raise ValueError("cmd_timeout_s must be >= 0")
    if cmd_timeout_s >= FUNCTION_TIMEOUT_S:
        raise ValueError(
            f"cmd_timeout_s ({cmd_timeout_s}) must be less than function timeout ({FUNCTION_TIMEOUT_S})"
        )

image = (
    modal.Image.debian_slim()
    .apt_install("git")
    .pip_install("uv")
    .uv_sync()
    .add_local_dir(".", remote_path="/root/repo", ignore=should_ignore)
)


@app.function(image=image, timeout=FUNCTION_TIMEOUT_S, memory=32768)
def run(cmd: str, cmd_timeout_s: int = 0):
    validate_cmd_timeout(cmd_timeout_s)
    print(f"[runner] starting command: {cmd}", flush=True)
    if cmd_timeout_s > 0:
        print(f"[runner] subprocess timeout: {cmd_timeout_s}s", flush=True)
    else:
        print("[runner] subprocess timeout: disabled", flush=True)

    start = time.perf_counter()
    process = subprocess.Popen(
        ["bash", "-lc", cmd],
        cwd="/root/repo",
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
    )
    assert process.stdout is not None

    try:
        for line in process.stdout:
            print(line, end="", flush=True)
            if cmd_timeout_s > 0 and (time.perf_counter() - start) > cmd_timeout_s:
                process.kill()
                raise TimeoutError(f"Command exceeded {cmd_timeout_s}s")
    finally:
        process.stdout.close()

    return_code = process.wait()
    elapsed_s = time.perf_counter() - start
    print(f"[runner] command finished in {elapsed_s:.2f}s with exit code {return_code}", flush=True)
    if return_code != 0:
        raise RuntimeError(f"Command failed: {return_code}")

@app.local_entrypoint()
def main(cmd: str, cmd_timeout_s: int = 0):
    validate_cmd_timeout(cmd_timeout_s)
    try:
        run.remote(cmd, cmd_timeout_s)
    except TimeoutError:
        print("[runner] timeout reached; see streamed logs above for latest output.", flush=True)
        raise