recovery/cloud/modal_embed.py

"""Run the Kd-only frozen-encoder pass on a temporary Modal L4.

This script is intentionally not run by tests. Build the sample locally first,
review the current Modal price, then invoke it with:

    modal run cloud/modal_embed.py

Retrieve the result with the `modal volume get` command printed on completion.
"""

from __future__ import annotations

import json
import os
import subprocess
from pathlib import Path

import modal

app = modal.App("mitointeract-v2-embeddings")
output_volume = modal.Volume.from_name(
    "mitointeract-v2-artifacts", create_if_missing=True
)
cache_volume = modal.Volume.from_name("mitointeract-hf-cache", create_if_missing=True)

recovery_dir = Path(__file__).parents[1]
image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "numpy>=2.0",
        "safetensors>=0.5",
        "torch>=2.4",
        "transformers>=4.45",
    )
    .add_local_dir(recovery_dir / "src", remote_path="/app/src")
    .add_local_file(
        recovery_dir / "scripts/embed_sample.py", remote_path="/app/embed_sample.py"
    )
    .add_local_file(
        recovery_dir / "artifacts/balm-kd/sample.jsonl",
        remote_path="/input/sample.jsonl",
    )
)


@app.function(
    image=image,
    gpu="L4",
    cpu=4,
    memory=16_384,
    timeout=3_600,
    volumes={"/output": output_volume, "/cache": cache_volume},
)
def embed() -> dict:
    env = os.environ.copy()
    env["PYTHONPATH"] = "/app/src"
    env["HF_HOME"] = "/cache/huggingface"
    subprocess.run(
        [
            "python",
            "/app/embed_sample.py",
            "--sample",
            "/input/sample.jsonl",
            "--target-key",
            "pkd",
            "--target-name",
            "pKd",
            "--batch-size",
            "32",
            "--device",
            "cuda",
            "--output",
            "/output/embeddings-balm-kd.npz",
        ],
        check=True,
        env=env,
    )
    output_volume.commit()
    cache_volume.commit()
    metadata = json.loads(Path("/output/embeddings-balm-kd.json").read_text())
    metadata["output_bytes"] = Path("/output/embeddings-balm-kd.npz").stat().st_size
    return metadata


@app.local_entrypoint()
def main() -> None:
    result = embed.remote()
    print(json.dumps(result, indent=2))
    print(
        "modal volume get mitointeract-v2-artifacts "
        "embeddings-balm-kd.npz artifacts/embeddings-balm-kd.npz"
    )
    print(
        "modal volume get mitointeract-v2-artifacts "
        "embeddings-balm-kd.json artifacts/embeddings-balm-kd.json"
    )