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.
This commit is contained in:
Daniel Han 2026-07-26 15:32:57 +00:00
commit 0f88219618
2 changed files with 54 additions and 3 deletions

View file

@ -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"
)