Function make_vllm_image

src/fused_mm_sampling/modal_lib/utils.py:65–112

()

VLLM_FORK_BRANCH = "feature/fmms-sampler"
VLLM_FORK_SHA = "7a74973e4dc727df979f2a5ec9fff64ac5319467"
# Upstream-main parent of the latest main-to-feature merge in VLLM_FORK_SHA.
# Determine it with:
# merge=$(git rev-list --first-parent --merges -n 1 "$VLLM_FORK_SHA")
# git show -s --format='%P' "$merge"  # use the second parent
VLLM_PRECOMPILED_WHEEL_SHA = "1a2c17634eccc4e68d9e1ab654f702d55361c754"


def make_vllm_image() -> modal.Image:
    return (
        # Base image must match vLLM's pinned torch version (see requirements/cuda.txt
        # in the fork: torch==2.11.0). cuda13.0 is the newest CUDA tag pytorch/pytorch
        # publishes for 2.11.0. Because vLLM pins torch without a local version, uv
        # treats the base image's torch==2.11.0+cu130 as satisfying the pin and skips
        # reinstalling, so we keep cu13 throughout and the precompiled .so stays ABI-
        # compatible.
        modal.Image.from_registry("pytorch/pytorch:2.11.0-cuda13.0-cudnn9-devel")
        .apt_install("git", "curl")
        .run_commands("pip install --break-system-packages uv")
        .run_commands(
            f"git clone --depth 1 -b {VLLM_FORK_BRANCH}"
            " https://github.com/tomasruizt/vllm.git /opt/vllm"
            " && cd /opt/vllm"
            f" && git fetch --depth 1 origin {VLLM_FORK_SHA}"
            f" && git checkout {VLLM_FORK_SHA}",
            # Pre-install numpy at a vllm[bench]-compatible version before
            # installing vllm. The base image ships numpy 2.4.3, but
            # vllm[bench]'s deps (numba, mistral-common) need numpy<2.3.
            # If we let uv downgrade during vllm install, it leaves a
            # broken `numpy-2.4.3.dist-info` entry that makes
            # importlib.metadata.version("numpy") return None, which
            # aborts the vllm CLI on import.
            "pip install --break-system-packages 'numpy<2.3'",
            # Install build dependencies into the base environment because the
            # precompiled build below deliberately disables isolation.
            "uv pip install --system -r /opt/vllm/requirements/build.txt",
            # Let vLLM's resolver pick the torch version that matches the precompiled .so.
            # Earlier we pinned torch==2.10.0 to match a 2.10.0+cu130 .so, but the upstream
            # precompiled wheel is now built against torch 2.11.0, so any pin breaks the ABI.
            "cd /opt/vllm"
            f" && VLLM_PRECOMPILED_WHEEL_COMMIT={VLLM_PRECOMPILED_WHEEL_SHA}"
            " VLLM_USE_PRECOMPILED=1 uv pip install"
            " --system --no-build-isolation '.[bench]'",
            "test -f /usr/local/lib/python3.12/dist-packages/vllm/_C.abi3.so",
        )
        .add_local_dir(
            str(_repo_root / "src"),
            remote_path="/opt/fused-mm-sample/src",
            copy=True,
            ignore=["__pycache__", "*.pyc"],
        )
        .add_local_file(