Upgrade TPU VM to python 3.10

Currently removing some packages which fail to install.

http://b/295050179

Author: djherbis

tpu/Dockerfile

@@ -6,7 +6,8 @@ FROM $BASE_IMAGE
 # See: https://docs.docker.com/engine/reference/builder/#understand-how-arg-and-from-interact
 ARG PYTHON_WHEEL_VERSION
 ARG PYTHON_VERSION_PATH
-ARG LINUX_WHEEL_VERSION
+ARG TF_LINUX_WHEEL_VERSION
+ARG TORCH_LINUX_WHEEL_VERSION
 ARG TORCH_VERSION
 ARG TENSORFLOW_VERSION
 ARG TF_LIBTPU_VERSION
@@ -52,12 +53,13 @@ ADD patches/kaggle_datasets.py /root/.local/lib/${PYTHON_VERSION_PATH}/site-pack
 
 # 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}-${LINUX_WHEEL_VERSION}.whl tensorflow-addons tensorflow-probability tensorflow-io \
-    torch==${TORCH_VERSION} https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-${TORCH_VERSION%.*}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${LINUX_WHEEL_VERSION}.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION} \
-    jax[tpu]==${JAX_VERSION} -f https://storage.googleapis.com/jax-releases/libtpu_releases.html trax flax optax elegy git+https://github.com/deepmind/dm-haiku jraph distrax \
-    numpy==1.23.5 \
-    papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1" \
-    pandas matplotlib opencv-python-headless librosa scikit-learn accelerate diffusers transformers
+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-addons tensorflow-probability tensorflow-io \
+    torch==${TORCH_VERSION} https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_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 \
+    # numpy==1.23.5 \
+    papermill jupyterlab python-lsp-server[all] "jupyter-lsp==1.5.1"
+    # elegy
+    # pandas matplotlib opencv-python-headless librosa accelerate diffusers scikit-learn transformers
 
 # Tensorflow libtpu:
 RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so

tpu/config.txt

@@ -1,16 +1,17 @@
-BASE_IMAGE=python:3.8
-PYTHON_WHEEL_VERSION=cp38
-PYTHON_VERSION_PATH=python3.8
+BASE_IMAGE=python:3.10
+PYTHON_WHEEL_VERSION=cp310
+PYTHON_VERSION_PATH=python3.10
 # https://cloud.google.com/tpu/docs/supported-tpu-configurations#tpu_software_versions:~:text=TensorFlow%20version-,libtpu.so%20version,-2.13.0
-TENSORFLOW_VERSION=2.12.0
-TF_LIBTPU_VERSION=1.6.0
-JAX_VERSION=0.4.13
+TENSORFLOW_VERSION=2.14.0
+TF_LIBTPU_VERSION=1.8.0
+TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64
+JAX_VERSION=0.4.17
 # Supports nightly
-TORCH_VERSION=2.0.0
+TORCH_VERSION=2.1.0
 # https://github.com/pytorch/audio supports nightly
-TORCHAUDIO_VERSION=2.0.0
+TORCHAUDIO_VERSION=2.1.0
 # https://github.com/pytorch/text supports main
-TORCHTEXT_VERSION=0.15.1
+TORCHTEXT_VERSION=0.16.0
 # https://github.com/pytorch/vision supports nightly
-TORCHVISION_VERSION=0.15.1
-LINUX_WHEEL_VERSION=linux_x86_64
+TORCHVISION_VERSION=0.16.0
+TORCH_LINUX_WHEEL_VERSION=manylinux_2_28_x86_64