Build tpu image with uv (#1452)

This should significantly improve build times for this image.

http://b/350047073

Author: djherbis

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:

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