Skip to content

Commit a0fceb5

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #3344 from AI-Hypercomputer:igorts/docker_prune
PiperOrigin-RevId: 879844454
2 parents 781d30a + e56309a commit a0fceb5

3 files changed

Lines changed: 13 additions & 13 deletions

File tree

‎dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ COPY src/install_maxtext_extra_deps/extra_deps_from_github.txt src/install_maxte
5353

5454
# Install dependencies - these steps are cached unless the copied files change
5555
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}"
56-
RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
56+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE}
5757

5858
# Now copy the remaining code (source files that may change frequently)
5959
COPY . .
@@ -68,4 +68,4 @@ RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
6868
fi
6969

7070
# Install (editable) MaxText
71-
RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .
71+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .

‎dependencies/dockerfiles/maxtext_post_training_local_dependencies.Dockerfile‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,28 +21,28 @@ ENV MODE=$MODE
2121
RUN echo "Installing Post-Training dependencies (tunix, vLLM, tpu-inference) with MODE=${MODE}"
2222
RUN pip uninstall -y jax jaxlib libtpu
2323

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

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

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

3131
COPY tunix /tunix
3232
RUN pip uninstall -y google-tunix
33-
RUN pip install -e /tunix --no-cache-dir
33+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install -e /tunix
3434

3535
COPY vllm /vllm
36-
RUN VLLM_TARGET_DEVICE="tpu" pip install -e /vllm --no-cache-dir
36+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv VLLM_TARGET_DEVICE="tpu" pip install -e /vllm
3737

3838
COPY tpu-inference /tpu-inference
39-
RUN pip install -e /tpu-inference --no-cache-dir
39+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv pip install -e /tpu-inference
4040

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

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

45-
RUN if [ "$MODE" = "post-training-experimental" ]; then \
45+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv if [ "$MODE" = "post-training-experimental" ]; then \
4646
echo "MODE=post-training-experimental: Re-installing JAX/libtpu"; \
4747
pip uninstall -y jax jaxlib libtpu && \
4848
pip install --pre -U jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ && \

‎dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ COPY libtpu.so* /root/custom_libtpu/
5454

5555
# Install dependencies - these steps are cached unless the copied files change
5656
RUN echo "Running command: bash setup.sh MODE=$ENV_MODE WORKFLOW=$ENV_WORKFLOW JAX_VERSION=$ENV_JAX_VERSION LIBTPU_VERSION=$ENV_LIBTPU_VERSION DEVICE=${ENV_DEVICE}"
57-
RUN --mount=type=cache,target=/root/.cache/pip bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
57+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv bash /deps/tools/setup/setup.sh MODE=${ENV_MODE} WORKFLOW=${ENV_WORKFLOW} JAX_VERSION=${ENV_JAX_VERSION} LIBTPU_VERSION=${ENV_LIBTPU_VERSION} DEVICE=${ENV_DEVICE}
5858

5959
# Now copy the remaining code (source files that may change frequently)
6060
COPY . .
@@ -69,4 +69,4 @@ RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
6969
fi
7070

7171
# Install (editable) MaxText
72-
RUN test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .
72+
RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=/root/.cache/uv test -f '/tmp/venv_created' && "$(tail -n1 /tmp/venv_created)"/bin/activate ; pip install --no-dependencies -e .

0 commit comments

Comments
 (0)