From 11430aaab55fa21b8f2089d793ae33554f6ace31 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 12 Jun 2026 11:04:13 +0000
Subject: [PATCH] docker: unbreak standalone vllm serve (ninja-build +
 flashinfer-jit-cache)

The notebook validation matrix caught the synthetic-data notebook dying
because the vllm server SyntheticDataKit launches never came up. Two
layers to the failure:

1. flashinfer's cpp_ext JIT shells out to ninja. The pip ninja lives in
   the venv bin, which subprocesses like vllm serve do not always
   inherit on PATH, so the JIT failed with exit 127. Install ninja-build
   so the binary is reachable from any PATH.
2. With ninja present the JIT still cannot succeed for device code: the
   runtime image deliberately ships no nvcc. Bake flashinfer-jit-cache
   (cu128) so ops missing from the cubin package (fmha_gen on sm_100a
   was the repro) come precompiled. In-process GRPO never hit this
   because unsloth-zoo blocks the FlashInfer JIT path; standalone
   vllm serve gets no zoo patches.

Fail-soft on the jit-cache for arches without a wheel; the vLLM chain
itself stays fail-loud on amd64.
---
 docker/Dockerfile | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index b0157e22e7..919f1cf21c 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -229,6 +229,17 @@ RUN set -eux \
        && ${VENV}/bin/python -c "import vllm; print('vllm', vllm.__version__)" \
        && ${VENV}/bin/python -c "import numpy.testing, numpy; print('numpy', numpy.__version__, 'testing ok')" \
        && ${VENV}/bin/python -c "import numba; print('numba', numba.__version__, 'imports ok')" \
+        # flashinfer-jit-cache: the runtime image ships no nvcc, so any
+        # flashinfer op missing from the cubin package would hit the JIT
+        # path and die (standalone `vllm serve` does exactly this for
+        # fmha_gen on sm_100a; in-process GRPO survives because zoo blocks
+        # the FlashInfer JIT). The precompiled cache removes the entire
+        # runtime-compile failure class for ~1.5 GB.
+       && { ${VENV}/bin/uv pip install \
+            --python ${VENV}/bin/python \
+            --index-url https://flashinfer.ai/whl/cu128 \
+            "flashinfer-jit-cache==0.6.6" \
+         || echo ">> flashinfer-jit-cache unavailable for ${TARGETARCH:-amd64}; vllm serve may require nvcc for uncached ops"; } \
        && echo ">> vLLM installed (numpy + numba re-upgraded post-vllm)"; \
         } || { \
             if [ "${TARGETARCH:-amd64}" != "amd64" ]; then \
@@ -411,9 +422,12 @@ ENV DEBIAN_FRONTEND=noninteractive \
 # wget: notebooks fetch sample assets with `!wget URL`; without it the cell
 # "succeeds" with sh's not-found on stderr and the next cell crashes on the
 # missing file (the Whisper notebook died exactly this way).
+# ninja-build: flashinfer's cpp_ext JIT shells out to ninja; subprocesses
+# (e.g. `vllm serve` launched by unsloth.dataprep) do not always inherit the
+# venv bin on PATH, so the pip ninja alone is not reachable there.
 RUN apt-get update && apt-get install -y --no-install-recommends \
         software-properties-common ca-certificates curl wget git libgomp1 \
-        gcc g++ zstd ffmpeg \
+        gcc g++ zstd ffmpeg ninja-build \
     && add-apt-repository -y ppa:deadsnakes/ppa \
     && apt-get update && apt-get install -y --no-install-recommends \
         python${PYTHON_VERSION} python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-dev \