Fix libtpu version for torch and do not pre-install tensorflow-tpu on TPU. (#1499)

We install a libtpu version compatible with both jax 0.7.2 and torch 2.8.0. Why? tunix latest -> flax 0.12 -> jax 0.7.2 -> libtpu 0.0.23, and that libtpu version causes pjrt api errors for torch 2.8.0: ``` pjrt_c_api_helpers.cc:258] Unexpected error status Unexpected PJRT_Plugin_Attributes_Args size: expe cted 32, got 24. The plugin is likely built with a later version than the framework. This plugin is built with PJRT API version 0.75. ``` * https://github.com/pytorch/xla/blob/d517649bdef6ab0519c30c704bde8779c8216502/setup.py#L111 * https://github.com/jax-ml/jax/blob/3489529b38d1f11d1e5caf4540775aadd5f2cdda/setup.py#L26 Of particular note, we no longer pre-install `tensorflow-tpu` as the newer libtpu causes issues finding the TPUs ``` external/local_xla/xla/stream_executor/tpu/tpu_platform_interface.cc:78] No TPU platform found. Platform manager status: OK ``` We also update how we install Python packages via `uv` for consistency and reproducibility. From a `requirements.in` file, we first generate a consistent dependency closure via `uv pip compile`, and then `uv pip install` the packages from the generated `requirements.txt`.

Author: metrizable

tpu/Dockerfile

@@ -34,20 +34,39 @@ RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
 # 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
+ADD tpu/requirements.in /kaggle_requirements.in
+RUN envsubst < /kaggle_requirements.in > /requirements.in
 
 # 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 --force-reinstall && \
+RUN export PATH="${HOME}/.local/bin:${PATH}" && \
+ uv pip compile --system --prerelease=allow \
+ --verbose \
+ --upgrade \
+ --find-links=https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+ --find-links=https://storage.googleapis.com/libtpu-releases/index.html \
+ --find-links=https://storage.googleapis.com/libtpu-wheels/index.html \
+ --find-links=https://download.pytorch.org/whl/torch_stable.html \
+ --emit-find-links \
+ --no-emit-package pip \
+ --no-emit-package setuptools \
+ --output-file /requirements.txt \
+ /requirements.in && \
+ uv pip install --system --prerelease=allow --force-reinstall \
+ -r /requirements.txt && \
+ uv cache clean && \
  /tmp/clean-layer.sh
 ENV PATH="~/.local/bin:${PATH}"
 
-# Try to force tensorflow to reliably install without breaking other installed deps
+# We install a libtpu version compatible with both jax 0.7.2 and torch 2.8.0.
+# Why? tunix latest -> flax 0.12 -> jax 0.7.2 -> libtpu 0.0.23. However, that
+# libtpu causes pjrt api errors for torch 2.8.0. screenshot/5heUtdyaJ4MmR3D
+# https://github.com/pytorch/xla/blob/d517649bdef6ab0519c30c704bde8779c8216502/setup.py#L111
+# https://github.com/jax-ml/jax/blob/3489529b38d1f11d1e5caf4540775aadd5f2cdda/setup.py#L26
 RUN export PATH="${HOME}/.local/bin:${PATH}" && \
- uv pip freeze --system > /tmp/constraints.txt && \
- uv pip install --system -c /tmp/constraints.txt tensorflow-tpu -f https://storage.googleapis.com/libtpu-tf-releases/index.html --force-reinstall && \
- rm /tmp/constraints.txt
+ uv pip install --system --force-reinstall libtpu==0.0.17 && \
+ uv cache clean && \
+ /tmp/clean-layer.sh
 
 # Kaggle Model Hub patches:
 ADD patches/kaggle_module_resolver.py /usr/local/lib/${PYTHON_VERSION_PATH}/site-packages/tensorflow_hub/kaggle_module_resolver.py

tpu/requirements.txt renamed to tpu/requirements.in

@@ -1,8 +1,8 @@
 # TPU Utils
 tpu-info
 # Tensorflow packages
-tensorflow-tpu==${TENSORFLOW_VERSION}
---find-links https://storage.googleapis.com/libtpu-tf-releases/index.html
+# TODO: b/447621961 - re-enable tensorflow-tpu when a compatible libtpu can be found.
+tensorflow-cpu==${TENSORFLOW_VERSION}
 tensorflow_hub
 tensorflow-io
 tensorflow-probability
@@ -13,8 +13,7 @@ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TOR
 torchaudio==${TORCHAUDIO_VERSION}
 torchvision==${TORCHVISION_VERSION}
 # Jax packages
-jax[tpu]>=0.5.2
---find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+jax[tpu]
 distrax
 flax
 git+https://github.com/deepmind/dm-haiku