maxtext/src/dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile at 6ad79a9bad9bf7112d2f7732cbafcdb55292ffbd · AI-Hypercomputer/maxtext · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
# Copyright 2023–2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

ARG BASEIMAGE=maxtext_base_image
FROM ${BASEIMAGE}

ARG MODE
ENV MODE=$MODE

RUN echo "Installing Post-Training dependencies (tunix, vLLM, tpu-inference) with MODE=${MODE}"
RUN pip uninstall -y jax jaxlib libtpu

RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install aiohttp==3.12.15

# Install Python packages that enable pip to authenticate with Google Artifact Registry automatically.
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install keyring keyrings.google-artifactregistry-auth

RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install numba==0.61.2

COPY tunix /tunix
RUN pip uninstall -y google-tunix
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install -e /tunix

COPY vllm /vllm
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv VLLM_TARGET_DEVICE="tpu" pip install -e /vllm

COPY tpu-inference /tpu-inference
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install -e /tpu-inference

RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install --no-deps qwix==0.1.4

RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install math-verify==0.9.0

RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv if [ "$MODE" = "post-training-experimental" ]; then \
    echo "MODE=post-training-experimental: Re-installing JAX/libtpu"; \
    pip uninstall -y jax jaxlib libtpu && \
    pip install --pre -U jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ && \
    pip install -U --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html; \
    fi