launch_jupyter_llava_next.py

import modal
import subprocess
import time 
import os 

GPU_CONFIG = modal.gpu.A100(count=4, memory = 80)
MINUTES = 60  # seconds



MODEL_PATH = "lmms-lab/llava-next-72b"

volume = modal.Volume.from_name("hf-model-store", create_if_missing=True)

model_store_path = f"/vol/models/{MODEL_PATH}"



sgl_image_jupy = (
    modal.Image.from_registry("nvidia/cuda:12.2.0-devel-ubuntu22.04", add_python="3.11")
    .apt_install("git", "wget", "cmake")
    .pip_install(
        "wheel==0.43.0",
        "torch==2.3.0",
        "torchvision==0.18.0",
        "transformers==4.40.2",
        "vllm==0.4.2",
        "timm==0.9.12",
        "Pillow==10.3.0",
        "huggingface_hub==0.22.2",
        "requests==2.31.0",
        "einops",
        "accelerate"
    )
    .run_commands("pip install git+https://github.com/agyaatcoder/LLaVA-NeXT.git", gpu = GPU_CONFIG)
    .pip_install("jupyterlab")

)

stub = modal.Stub("app_new", image = sgl_image_jupy)
JUPYTER_TOKEN = "1234"  # Change me to something non-guessable!


#secrets=[Secret.from_name("mistral-secret")
@stub.function(concurrency_limit=1, timeout= 30000, volumes={model_store_path: volume}, gpu = GPU_CONFIG)
def run_jupyter(timeout: int):
    jupyter_port = 8888
    with modal.forward(jupyter_port) as tunnel:
        jupyter_process = subprocess.Popen(
            [
                "jupyter",
                "lab",
                "--no-browser",
                "--allow-root",
                "--ip=0.0.0.0",
                f"--port={jupyter_port}",
                "--NotebookApp.allow_origin='*'",
                "--NotebookApp.allow_remote_access=1",
            ],
            env={**os.environ, "JUPYTER_TOKEN": JUPYTER_TOKEN},
        )

        print(f"Jupyter available at => {tunnel.url}")

        try:
            end_time = time.time() + timeout
            while time.time() < end_time:
                time.sleep(5)
            print(f"Reached end of {timeout} second timeout period. Exiting...")
        except KeyboardInterrupt:
            print("Exiting...")
        finally:
            jupyter_process.kill()


@stub.local_entrypoint()
def main(timeout: int = 10_000):
    # Write some images to a volume, for demonstration purposes.
    # Run the Jupyter Notebook server
    run_jupyter.remote(timeout=timeout)