From adffc5d565c0b122a7d5d60c6cf1e1acad360f95 Mon Sep 17 00:00:00 2001
From: Dustin H <herbison@google.com>
Date: Wed, 4 Dec 2024 16:46:16 -0500
Subject: [PATCH] Build tpu image with uv (#1452)

This should significantly improve build times for this image.

http://b/350047073
---
 tpu/Dockerfile       | 20 ++++++++++----------
 tpu/requirements.txt | 45 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 10 deletions(-)
 create mode 100644 tpu/requirements.txt

diff --git a/tpu/Dockerfile b/tpu/Dockerfile
index b94619da..6a45fa60 100644
--- a/tpu/Dockerfile
+++ b/tpu/Dockerfile
@@ -56,16 +56,16 @@ RUN apt-get update && apt-get install ffmpeg libsm6 libxext6  -y
 
 # Packages needed by the Notebook editor
 
-# Additional useful packages should be added here
-
-RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TF_LINUX_WHEEL_VERSION}.whl tensorflow-probability tensorflow-io \
-    torch~=${TORCH_VERSION} https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}+libtpu-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} \
-    jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax git+https://github.com/deepmind/dm-haiku jraph distrax \
-    papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
-    pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers \
-    seaborn timm albumentations einops pyarrow fastparquet opencv-python \
-    "keras>3" keras-cv keras-nlp \
-    kagglehub && \
+# Additional useful packages should be added in the requirements.txt
+
+# Bring in the requirements.txt and replace variables in it:
+RUN apt-get install -y gettext
+ADD tpu/requirements.txt /kaggle_requirements.txt
+RUN envsubst < /kaggle_requirements.txt > /requirements.txt
+
+# Install uv and then install the requirements:
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+RUN export PATH="${HOME}/.local/bin:${PATH}" && uv pip install --system -r /requirements.txt --prerelease=allow --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html && \
     /tmp/clean-layer.sh
 
 # Tensorflow libtpu:
diff --git a/tpu/requirements.txt b/tpu/requirements.txt
new file mode 100644
index 00000000..99db64a7
--- /dev/null
+++ b/tpu/requirements.txt
@@ -0,0 +1,45 @@
+# Tensorflow packages
+https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-${TENSORFLOW_VERSION}/tensorflow-${TENSORFLOW_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TF_LINUX_WHEEL_VERSION}.whl
+tensorflow_hub
+tensorflow-io
+tensorflow-probability
+# Torch packages
+torch~=${TORCH_VERSION}
+https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}+libtpu-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl
+torchaudio==${TORCHAUDIO_VERSION}
+torchvision==${TORCHVISION_VERSION}
+# Jax packages
+jax[tpu]==${JAX_VERSION}
+distrax
+flax
+git+https://github.com/deepmind/dm-haiku
+jraph
+optax
+trax
+# Jupyter packages
+jupyter-lsp==1.5.1
+jupyterlab
+papermill
+python-lsp-server[all]
+# Keras Packages
+keras>3
+keras-cv
+keras-nlp
+# Kaggle Packages
+kagglehub
+# Other useful packages, add more here
+accelerate
+albumentations
+diffusers
+einops
+fastparquet
+librosa
+matplotlib
+opencv-python
+opencv-python-headless
+pandas
+pyarrow
+scikit-learn
+seaborn
+timm
+transformers