run.py

import sys
from pathlib import Path

REMOTE_PROJECT_DIR = "/root/project"

sys.path.append(str(Path(__file__).resolve().parent / "src"))
sys.path.append(REMOTE_PROJECT_DIR + "/src")

from dotenv import load_dotenv

load_dotenv(".env")

import math
import os
import pathlib

import hydra
import modal
from omegaconf import OmegaConf

from trainer import train
from utils.sandbox import extract_dependencies

vol = modal.Volume.from_name(os.environ["MODAL_VOLUME_NAME"], create_if_missing=True)
chkpt_dir = pathlib.Path("/vol")
OmegaConf.register_new_resolver("eval", lambda expr: eval(expr, {"math": math}))

app = modal.App()
dependencies = extract_dependencies("pyproject.toml")
env_vars = dict(os.environ)
image = (
    modal.Image.debian_slim(python_version="3.13")
    .add_local_dir(
        ".",
        REMOTE_PROJECT_DIR,
        copy=True,
        ignore=["*.venv", "*__pycache__", "*.idea", "*.ruff_cache"],
    )
    .apt_install("git")
    .pip_install("uv")
    .run_commands(f"uv pip install --system --compile-bytecode {dependencies}")
    .env(env_vars)
)


@app.function(
    gpu="A100",
    image=image,
    timeout=10800,
    volumes={str(chkpt_dir): vol},
)
def run():
    """Entry point for launching llm fine-tuning job in a Modal cloud environment.

    Environment Variables:
        TRAIN_CONFIG: The name of the Hydra config to use (e.g., "lora", "fourier").
    """
    config_name = os.environ["TRAIN_CONFIG"]

    hydra.initialize(config_path="./project/configs", version_base=None)
    config = hydra.compose(config_name=config_name)

    train(config)
    vol.commit()