Skip to content

Upgrade TPU VM to python 3.10 #1314

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 2, 2023
Merged

Upgrade TPU VM to python 3.10 #1314

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
730cab3
Upgrade TPU VM to python 3.10
djherbis Oct 23, 2023
f3b758e
Merge branch 'main' into tpuvm-py310
Philmod Oct 25, 2023
5709b53
Add +libtpu to TORCH_VERSION
djherbis Nov 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 9 additions & 7 deletions tpu/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}+libtpu-${PYTHON_WHEEL_VERSION}-${PYTHON_WHEEL_VERSION}-${TORCH_LINUX_WHEEL_VERSION}.whl torchvision==${TORCHVISION_VERSION} torchtext==${TORCHTEXT_VERSION} torchaudio==${TORCHAUDIO_VERSION} \
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how did you figure this out?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://b.corp.google.com/issues/295050179#comment20
I asked Will and he noticed my mistake lol

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
Expand Down
23 changes: 12 additions & 11 deletions tpu/config.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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