from modal import App, Image, Volume
import modal
import json
model_name = "Learnable"
app = App(f"Generalization Model {model_name} with ImageNet")
# Build image with all local dependencies added directly
image = (
    Image.from_registry("nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04",add_python="3.10")  # <--- This is the required 'tag'
    .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", 
                  "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) 
    .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""")
    .add_local_file("pyproject.toml", "/root/pyproject.toml")
    .add_local_dir("src", "/root/src")
)
# Shared volume for saving outputs or checkpoints
volume1 = Volume.from_name("weights", create_if_missing=True)
volume2 = Volume.from_name("datasets", create_if_missing=True)
volume3 = Volume.from_name("plots", create_if_missing=True)
volume4 = Volume.from_name("results", create_if_missing=True)
volume5 = Volume.from_name("shot_noise", create_if_missing=True)
@app.function(
    image=image,
    gpu="H100",
    timeout=3600 * 24,
    volumes={
        "/root/weights":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4,
        "/root/shot_noise":volume5,
    },
)
def run_command():
    import os
    # os.system("pip install -e .")
    os.system("ls /root/shot_noise")
    # os.system("""
    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
    #     --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed0/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052\
    #     --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_5004\
    #     --data-original /root/datasets/imagenet\
    #     --data-generalization /root/shot_noise/shot_noise/1
    # """)
    os.system("""
    CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
        --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed0/finetune-learnable-indice0-heads12-shared1-routed0-topk0-seed0/best_65052\
        --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_5004\
        --data-original /root/datasets/imagenet\
        --data-generalization /root/shot_noise/shot_noise/4
    """)
    # os.system("""
    # CUDA_VISIBLE_DEVICES=0  python src/imagenet/generalization.py \
    #     --model-a /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed20/best_5004\
    #     --model-b /root/weights/lmc/imagenet/vit-finetune/finetune-learnable-indice0-heads12-shared1-routed0-topk0-mlpFalse-seed40/best_5004\
    #     --data-original /root/datasets/imagenet\
    #     --data-generalization /root/shot_noise/shot_noise/1
    # """)
if __name__ == "__main__":
    with app.run():
        run_command.remote()