| # remdm-planner-craftax: vast.ai image. | |
| # | |
| # Build and push (amd64 is mandatory: vast GPU hosts are x86, an arm64 image | |
| # from an Apple Silicon Mac will not run): | |
| # | |
| # docker buildx build --platform linux/amd64 \ | |
| # -t <ANON_DOCKER_NAMESPACE>/remdm-planner-craftax:cuda13 \ | |
| # -t <ANON_DOCKER_NAMESPACE>/remdm-planner-craftax:latest --push . | |
| # | |
| # Bakes only the slow part: the jax[cuda13] wheel set (~3 GB of downloads). The | |
| # repo is cloned at boot by scripts/vast_onstart.sh, so code changes never need | |
| # a rebuild. Rebuild only when uv.lock changes. | |
| # | |
| # Nothing here starts training. You SSH in and run commands yourself. | |
| # Extending a vastai/* image keeps SSH, Jupyter and the instance portal working, | |
| # and vast hosts already cache these layers, so only your own layers get pulled. | |
| # Current tags: https://hub.docker.com/r/vastai/base-image/tags | |
| ARG BASE_IMAGE=vastai/base-image:cuda-13.3.1-auto | |
| FROM ${BASE_IMAGE} | |
| # NOTE: PATH is deliberately NOT modified. The vast base image runs its | |
| # supervisor, Jupyter and portal scripts against /venv/main, and shadowing bare | |
| # `python` / `pip` globally can break them. Everything below addresses | |
| # /opt/venv explicitly; interactive shells get it via the on-start script. | |
| ENV DEBIAN_FRONTEND=noninteractive \ | |
| UV_PROJECT_ENVIRONMENT=/opt/venv \ | |
| UV_LINK_MODE=copy \ | |
| UV_COMPILE_BYTECODE=1 \ | |
| UV_NO_CACHE=1 | |
| # uv. Pinned so a rebuild months from now is reproducible; bump deliberately. | |
| COPY --from=ghcr.io/astral-sh/uv:0.9.2 /uv /uvx /usr/local/bin/ | |
| # Tools the on-start script needs. The base image normally has these, but an | |
| # explicit install costs ~nothing and turns a confusing boot failure into a | |
| # build failure. | |
| RUN apt-get update \ | |
| && apt-get install -y --no-install-recommends git git-lfs tmux rsync \ | |
| && rm -rf /var/lib/apt/lists/* | |
| # --- the expensive layer: dependencies only ------------------------------- | |
| # README.md is copied because pyproject.toml references it. | |
| # --no-install-project keeps this layer keyed on the lockfile alone, so editing | |
| # source code never invalidates it. | |
| WORKDIR /opt/build | |
| COPY pyproject.toml uv.lock README.md ./ | |
| RUN uv sync --frozen --extra cuda13 --no-install-project | |
| # Fail the build now rather than after a 10 GB pull. Reports CPU-only devices | |
| # at build time, which is expected: there is no GPU in the builder. | |
| RUN /opt/venv/bin/python -c "import jax; print('jax', jax.__version__)" | |
| # Record the lockfile the venv was built from. The on-start script compares it | |
| # against the cloned repo's lockfile and warns loudly on drift, which otherwise | |
| # shows up as a silent 3 GB download at boot. | |
| RUN cp uv.lock /opt/venv/baked-uv.lock | |
| WORKDIR /workspace | |
| # --- optional: bake the released checkpoints ------------------------------ | |
| # Left out on purpose. 470 MB on the Hub unpacks into a large small-file layer, | |
| # which is the slowest kind to pull and a likely cause of stalled image pulls. | |
| # Fetching them at boot takes under a minute. If you do want them baked, use an | |
| # isolated env so the pinned huggingface_hub in /opt/venv is not disturbed: | |
| # | |
| # RUN HF_HOME=/tmp/hf uvx --from "huggingface_hub[cli]" \ | |
| # hf download AnonMLuser/remdm-planner-craftax \ | |
| # --include "checkpoints/**" --local-dir /opt/checkpoints \ | |
| # && rm -rf /tmp/hf |