Merge pull request #101 from coreweave/jp/testing/slim-vllm-image

feat(vllm-tensorizer): Optimize Multi-Stage Build for Slimmer Inference Image

Author: Eta0

.github/configurations/vllm-tensorizer.yml

@@ -2,5 +2,7 @@ vllm-commit:
   - 'b6553be1bc75f046b00046a4ad7576364d03c835'
 flashinfer-commit:
   - 'v0.2.6.post1'
-base-image:
-  - 'ghcr.io/coreweave/ml-containers/torch-extras:es-compute-12.0-67208ca-nccl-cuda12.9.0-ubuntu22.04-nccl2.27.3-1-torch2.7.1-vision0.22.1-audio2.7.1-abi1'
+builder-base-image:
+  - 'ghcr.io/coreweave/ml-containers/torch-extras:es-cuda-12.9.1-74755e9-nccl-cuda12.9.1-ubuntu22.04-nccl2.27.5-1-torch2.7.1-vision0.22.1-audio2.7.1-abi1'
+final-base-image:
+  - 'ghcr.io/coreweave/ml-containers/torch-extras:es-cuda-12.9.1-74755e9-base-cuda12.9.1-ubuntu22.04-torch2.7.1-vision0.22.1-audio2.7.1-abi1'

.github/workflows/vllm-tensorizer.yml

@@ -25,4 +25,5 @@ jobs:
       build-args: |
         VLLM_COMMIT=${{ matrix.vllm-commit }}
         FLASHINFER_COMMIT=${{ matrix.flashinfer-commit }}
-        BASE_IMAGE=${{ matrix.base-image }}
+        BUILDER_BASE_IMAGE=${{ matrix.builder-base-image }}
+        FINAL_BASE_IMAGE=${{ matrix.final-base-image }}

vllm-tensorizer/Dockerfile

@@ -1,11 +1,13 @@
-ARG BASE_IMAGE="ghcr.io/coreweave/ml-containers/torch-extras:es-compute-12.0-67208ca-nccl-cuda12.9.0-ubuntu22.04-nccl2.27.3-1-torch2.7.1-vision0.22.1-audio2.7.1-abi1"
+ARG BUILDER_BASE_IMAGE="ghcr.io/coreweave/ml-containers/torch-extras:es-cuda-12.9.1-74755e9-nccl-cuda12.9.1-ubuntu22.04-nccl2.27.5-1-torch2.7.1-vision0.22.1-audio2.7.1-abi1"
+ARG FINAL_BASE_IMAGE="ghcr.io/coreweave/ml-containers/torch-extras:es-cuda-12.9.1-74755e9-base-cuda12.9.1-ubuntu22.04-torch2.7.1-vision0.22.1-audio2.7.1-abi1"
+
 FROM scratch AS freezer
 WORKDIR /
 COPY --chmod=755 freeze.sh /
 
-FROM ${BASE_IMAGE} AS builder-base
+FROM ${BUILDER_BASE_IMAGE} AS builder-base
 
-ARG MAX_JOBS="16"
+ARG MAX_JOBS="32"
 
 RUN ldconfig
 
@@ -81,7 +83,7 @@ RUN --mount=type=bind,from=flashinfer-downloader,source=/git/flashinfer,target=/
 WORKDIR /wheels
 
 
-FROM ${BASE_IMAGE} AS base
+FROM ${FINAL_BASE_IMAGE} AS base
 
 WORKDIR /workspace