From 0f88219618b957bc14847770fec163da6203a350 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 15:32:57 +0000 Subject: [PATCH] docker: fix the unsloth CLI and the vLLM engine in the image Two defects found by running the built image rather than reading it. 1. Every unsloth_cli subcommand that touches the studio backend died on import. `unsloth list-checkpoints` on the published image: ModuleNotFoundError: No module named 'structlog' and the same for train / export / chat, since all four import studio.backend.core.*. structlog is a studio backend requirement, not an unsloth[huggingface] one, so nothing in the base install pulled it in. Added it to the base venv, and added a build-time `from studio.backend.core.export import ExportBackend` so a future missing dependency in that closure fails the build instead of the user's first CLI invocation. That guard has to live in the LAST builder verification block: the closure also needs starlette, which only arrives with vLLM two stages later. 2. flashinfer-jit-cache was pinned to a literal 0.6.6 while vLLM 0.26.0 resolves flashinfer-python 0.6.14. flashinfer raises at import when the two disagree, and that exception is thrown inside the vLLM EngineCore, so Unsloth's GRPO fast_inference path fails at engine start with no earlier warning. A literal pin drifts again on the next vLLM bump, so the version is now read back from the resolved flashinfer-python, and the build proves `import flashinfer` works. Verified on the rebuilt image: flashinfer-python 0.6.14 with flashinfer-jit-cache 0.6.14+cu128, structlog 26.1.0, the export backend importable, and `unsloth list-checkpoints` exiting 0. tests/python/test_docker_llama_cuda_backend.py gains two static cases pinning both: the jit-cache version must be derived rather than literal and the build must import flashinfer, and the base venv must ask for structlog with the CLI reachability guard present. --- docker/Dockerfile | 29 +++++++++++++++++-- .../python/test_docker_llama_cuda_backend.py | 28 ++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index dc29131350..9348c89d1a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -121,7 +121,11 @@ RUN set -eux \ "bitsandbytes>=0.49.2,!=0.46.0,!=0.48.0" \ "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo@${UNSLOTH_ZOO_REF}" \ "unsloth[${UNSLOTH_EXTRA}] @ git+https://github.com/unslothai/unsloth@${UNSLOTH_REF}" \ - "timm>=1.0.11" "addict" + `# structlog is a studio backend dep, not an unsloth[huggingface] dep,` \ + `# but unsloth_cli's train / export / chat / list-checkpoints all import` \ + `# studio.backend.core.*, so without it every one of them dies on` \ + `# ModuleNotFoundError. The last builder stage imports it as a guard.` \ + "timm>=1.0.11" "addict" "structlog" # vLLM: required by Unsloth's GRPO path (fast_inference=True). A SECOND uv pass so # torch 2.11.0 settles first; with torch held, uv picks the newest compatible vLLM @@ -165,11 +169,21 @@ RUN set -eux \ && ${VENV}/bin/python -c "import numba; print('numba', numba.__version__, 'imports ok')" \ # flashinfer-jit-cache: precompiled cubins so flashinfer ops skip the JIT # path (standalone `vllm serve` dies there for fmha_gen on sm_100a). ~1.5 GB. + # The version MUST equal the flashinfer-python vLLM resolved: flashinfer + # raises at import when the two disagree, which takes the vLLM EngineCore + # down with it and breaks Unsloth's GRPO fast_inference path. So read the + # resolved version instead of pinning a literal that drifts. + && FI_VER="$(${VENV}/bin/python -c 'from importlib.metadata import version; print(version("flashinfer-python"))')" \ + && echo ">> flashinfer-python ${FI_VER}, matching flashinfer-jit-cache" \ && { ${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"; } \ + "flashinfer-jit-cache==${FI_VER}" \ + || echo ">> flashinfer-jit-cache ${FI_VER} unavailable for ${TARGETARCH:-amd64}; vllm serve may require nvcc for uncached ops"; } \ + # Whatever happened above, flashinfer has to import: a version mismatch + # here is silent until the first vLLM engine start. + && ${VENV}/bin/python -c \ + "import flashinfer; print('OK: flashinfer', flashinfer.__version__, 'imports')" \ && echo ">> vLLM installed (numpy + numba re-upgraded post-vllm)"; \ } || { \ if [ "${TARGETARCH:-amd64}" != "amd64" ]; then \ @@ -342,6 +356,15 @@ if target == "amd64": for pkg in LIGHT_IMPORTS: importlib.import_module(pkg) print(f"OK: {' + '.join(LIGHT_IMPORTS)} import cleanly on no-GPU host") + +# `unsloth train` / `export` / `chat` / `list-checkpoints` all import +# studio.backend.core.*, whose dependency closure (structlog, and starlette by +# way of the logging handlers) is NOT part of unsloth[huggingface]. Missing any +# of it turns every one of those commands into a ModuleNotFoundError traceback, +# which no functional test here would otherwise catch. Runs last in the builder, +# after vLLM, because that is what pulls starlette in. +from studio.backend.core.export import ExportBackend # noqa: F401 +print("OK: the unsloth CLI can reach the studio export backend") PY # ============================================================================= diff --git a/tests/python/test_docker_llama_cuda_backend.py b/tests/python/test_docker_llama_cuda_backend.py index d043de7c45..a676ccd8ed 100644 --- a/tests/python/test_docker_llama_cuda_backend.py +++ b/tests/python/test_docker_llama_cuda_backend.py @@ -89,3 +89,31 @@ def test_guard_runs_after_the_prebuilt_is_fetched(dockerfile: str): fetch = dockerfile.index("fetch_llama_prebuilt.py") guard = dockerfile.index("CUDA_SO=") assert fetch < guard, "the guard can only inspect a bundle that already exists" + + +def test_flashinfer_jit_cache_tracks_flashinfer(dockerfile: str): + # flashinfer raises at import when flashinfer-jit-cache and flashinfer-python + # disagree, and that exception kills the vLLM EngineCore, which is what + # Unsloth's GRPO fast_inference path runs on. A literal pin drifts the moment + # vLLM bumps its flashinfer requirement, so the version has to be derived. + assert "flashinfer-jit-cache==${FI_VER}" in dockerfile, ( + "flashinfer-jit-cache must be pinned to the resolved flashinfer-python version" + ) + assert not re.search(r"flashinfer-jit-cache==[0-9]", dockerfile), ( + "a literal flashinfer-jit-cache version will drift away from flashinfer-python" + ) + assert "import flashinfer" in dockerfile, ( + "the build must prove flashinfer imports, or a mismatch stays silent " + "until the first vLLM engine start" + ) + + +def test_cli_can_reach_the_studio_backend(dockerfile: str): + # unsloth_cli's train / export / chat / list-checkpoints import + # studio.backend.core.*, which needs structlog. It is a studio backend + # requirement rather than an unsloth[huggingface] one, so the base venv has + # to ask for it explicitly or the whole CLI dies on ModuleNotFoundError. + assert '"structlog"' in dockerfile, "the base venv must install structlog for unsloth_cli" + assert "from studio.backend.core.export import ExportBackend" in dockerfile, ( + "a build-time import guard must prove the CLI can reach the studio backend" + )