import os

import modal
import modal.experimental

cuda_version = "12.6.0"  # should be no greater than host CUDA version
flavor = "devel"  #  includes full CUDA toolkit
operating_sys = "ubuntu22.04"
tag = f"{cuda_version}-{flavor}-{operating_sys}"

LOCAL_CODE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
REMOTE_CODE_DIR = "/root/"
REMOTE_TRAIN_SCRIPT_PATH = "train_audios.py"
GPU_TYPE = "H100"

image = (
    modal.Image.from_registry(f"nvidia/cuda:{tag}", add_python="3.12")
    # Note that libibverbs-dev and libibverbs1 are required for RDMA.
    .apt_install("git", "libibverbs-dev", "libibverbs1")
    .pip_install(
        "torch==2.6.0",
        "wandb==0.19.11",
        "tqdm==4.67.1",
        "numpy==1.26.4",
        "torcheval==0.0.7",
        "funcy==2.0",
        "joblib==1.5.1",
        "ffmpeg-python>=0.2.0",
        "sox>=1.4.1",
        "tinytag>=1.8.1",
        "sentencepiece>=0.1.97",
        "transformers",
        "torchaudio",
        "tokenizers",
    )
    .apt_install("sox", "ffmpeg", "libsox-fmt-mp3")
    .workdir("/root/scripts")
    .add_local_dir("../../suno_utils", remote_path="/root/suno_utils")
    .add_local_dir(LOCAL_CODE_DIR, remote_path=REMOTE_CODE_DIR)
)

app = modal.App("suno-hoot-training", image=image)

model_volume = modal.Volume.from_name("suno-hoot-training")
data_volume = modal.Volume.from_name("suno-hoot-data")

# The number of containers (i.e. nodes) in the cluster. This can be between 1 and 8.
n_nodes = 2
# Typically this matches the number of GPUs per container.
n_proc_per_node = 8


def _train_multi_node(profile: bool = False) -> None:
    from torch.distributed.run import parse_args, run

    cluster_info = modal.experimental.get_cluster_info()
    # which container am I?
    container_rank: int = cluster_info.rank
    # what's the leader/master/main container's address?
    main_ip_addr: str = cluster_info.container_ips[0]
    container_id = os.environ["MODAL_TASK_ID"]

    print(f"hello from {container_id}, rank {container_rank} of {n_nodes}")
    if container_rank == 0:
        print(f"main container's address: {main_ip_addr}")

    args = [
        f"--nnodes={n_nodes}",
        f"--nproc-per-node={n_proc_per_node}",
        f"--node-rank={cluster_info.rank}",
        f"--master-addr={main_ip_addr}",
        REMOTE_TRAIN_SCRIPT_PATH,
        "--out_dir=/models/checkpoints",
        "--max_duration_s=280",
        "--max_iters=10000",
        "--warmup_iters=10",
        "--log_interval=1",
        "--eval_interval=5_000",
        "--eval_iters=50",
        "--learning_rate=1e-6",
        "--min_lr=1e-6",
        "--batch_size=2",
        "--decoder_type=v2",
        "--gradient_accumulation_steps=1",
        "--n_augment_freq_masks=4",
        "--preload_checkpoint=/models/hoot_v6_t2_300k.pt",
        "--preload_decoder=True",
        "--train_suno_input_path=/models/suno_hoot_cleaned.json",
        "--custom_seed_offset=8889",
        "--lyrics_key=text",
        "--wandb_run_name=hoot_v6_v2_sft_test",
    ]
    if profile:
        args.append("--profile=True")
    print(f"Running torchrun with args: {' '.join(args)}")
    run(parse_args(args))


@app.function(
    gpu=f"{GPU_TYPE}:{n_proc_per_node}",
    secrets=[
        # Required for connecting to Weights & Biases from within the Modal container.
        modal.Secret.from_name("wandb-secret"),
    ],
    volumes={
        "/models": model_volume,
        # This volume path matches the path in `suno_hoot_20250617_subset100k.json`. This should
        # be something like /data in a prod environment.
        "/app2/suno/data/dpo/audios/auk_t1": data_volume,
    },
    timeout=60 * 60 * 24,
)
@modal.experimental.clustered(n_nodes, rdma=True)
def train_multi_node(profile: bool = False):
    """
    Train the model on a multi-node cluster with N GPUs per node (typically 4).
    Good cluster scale performance should result in a ~linear speedup as the number of nodes
    is increased.
    """
    _train_multi_node(profile)