From 3b86333de953712ed5d837525f9fa596e43f564e Mon Sep 17 00:00:00 2001 From: Yifei Teng Date: Mon, 18 Nov 2024 19:40:49 +0000 Subject: [PATCH] Switch to libtpu package from the libtpu-wheels registry No new nightly libtpu builds will be uploaded to the libtpu-releases registry on Dec 31. The new registry hosts both nightly and stable builds of libtpu, but they are under a different name: `libtpu` instead of `libtpu-nightly`. This PR updates both the package name and the registry. In order to smoothly migrate to the new package, a cleanup package is also installed to remove the libtpu binary from the old package name. New registry: https://storage.googleapis.com/libtpu-wheels/index.html Deprecated registry: https://storage.googleapis.com/libtpu-releases/index.html --- .github/workflows/_tpu_ci.yml | 2 +- CONTRIBUTING.md | 4 +++- README.md | 8 +++++--- docs/source/contribute/configure-environment.md | 4 +++- infra/ansible/roles/build_plugin/tasks/main.yaml | 2 +- scripts/build_developer.sh | 4 +++- scripts/build_torch_wheels.sh | 4 +++- setup.py | 14 ++++++++++---- 8 files changed, 29 insertions(+), 13 deletions(-) diff --git a/.github/workflows/_tpu_ci.yml b/.github/workflows/_tpu_ci.yml index facb3aa50b2..f04c2c3b099 100644 --- a/.github/workflows/_tpu_ci.yml +++ b/.github/workflows/_tpu_ci.yml @@ -25,7 +25,7 @@ jobs: pip install rich # Jax nightly is needed for pallas tests. pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html - pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html + pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html pip install --upgrade protobuf - name: Run Tests env: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 33c59a57cc9..06aca135e37 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -63,7 +63,9 @@ We recommend you to use our prebuilt Docker image to start your development work cd pytorch/xla python setup.py develop # Optional: if you're using TPU, install libtpu - pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html + pip install torch_xla[tpu] \ + -f https://storage.googleapis.com/libtpu-wheels/index.html \ + -f https://storage.googleapis.com/libtpu-releases/index.html ``` * Test your build diff --git a/README.md b/README.md index 96ffe9e6deb..6c4994637c5 100644 --- a/README.md +++ b/README.md @@ -26,14 +26,14 @@ started: To install PyTorch/XLA stable build in a new TPU VM: ``` -pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html +pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html ``` To install PyTorch/XLA nightly build in a new TPU VM: ``` pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu -pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html +pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html ``` ### GPU Plugin @@ -147,7 +147,9 @@ can now install the main build with `pip install torch_xla`. To also install the Cloud TPU plugin corresponding to your installed `torch_xla`, install the optional `tpu` dependencies after installing the main build with ``` -pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html +pip install torch_xla[tpu] \ + -f https://storage.googleapis.com/libtpu-wheels/index.html \ + -f https://storage.googleapis.com/libtpu-releases/index.html ``` GPU and nightly builds are available in our public GCS bucket. diff --git a/docs/source/contribute/configure-environment.md b/docs/source/contribute/configure-environment.md index 972765108b0..4a5bf35d02d 100644 --- a/docs/source/contribute/configure-environment.md +++ b/docs/source/contribute/configure-environment.md @@ -87,7 +87,9 @@ via the Command Palette (`Python: Create Environment`). Install the latest PyTorch and PyTorch/XLA releases: ``` bash -pip install numpy torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html +pip install numpy torch torch_xla[tpu] \ + -f https://storage.googleapis.com/libtpu-wheels/index.html \ + -f https://storage.googleapis.com/libtpu-releases/index.html ``` Create a file `test.py`: diff --git a/infra/ansible/roles/build_plugin/tasks/main.yaml b/infra/ansible/roles/build_plugin/tasks/main.yaml index 2e2590b150a..142d29c3718 100644 --- a/infra/ansible/roles/build_plugin/tasks/main.yaml +++ b/infra/ansible/roles/build_plugin/tasks/main.yaml @@ -28,5 +28,5 @@ - name: Install libtpu ansible.builtin.pip: name: torch_xla[tpu] - extra_args: -f https://storage.googleapis.com/libtpu-releases/index.html + extra_args: -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html when: accelerator == "tpuvm" diff --git a/scripts/build_developer.sh b/scripts/build_developer.sh index 2590860424e..680c4a3e8f7 100755 --- a/scripts/build_developer.sh +++ b/scripts/build_developer.sh @@ -23,7 +23,9 @@ python3 setup.py develop # libtpu is needed to talk to the TPUs. If TPUs are not present, # installing this wouldn't hurt either. -pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html +pip install torch_xla[tpu] \ + -f https://storage.googleapis.com/libtpu-wheels/index.html \ + -f https://storage.googleapis.com/libtpu-releases/index.html # Test that the library is installed correctly. python3 -c 'import torch_xla as xla; print(xla.device())' diff --git a/scripts/build_torch_wheels.sh b/scripts/build_torch_wheels.sh index 6258f4b9124..5e3ada94cd2 100755 --- a/scripts/build_torch_wheels.sh +++ b/scripts/build_torch_wheels.sh @@ -280,7 +280,9 @@ function build_and_install_torch_xla() { python setup.py bdist_wheel pip install dist/*.whl if [ "$TPUVM_MODE" == "1" ]; then - pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html + pip install torch_xla[tpu] \ + -f https://storage.googleapis.com/libtpu-wheels/index.html \ + -f https://storage.googleapis.com/libtpu-releases/index.html sudo apt-get install -y google-perftools fi diff --git a/setup.py b/setup.py index efee197180a..cee23c8f44d 100644 --- a/setup.py +++ b/setup.py @@ -65,8 +65,8 @@ base_dir = os.path.dirname(os.path.abspath(__file__)) _date = '20241118' -_libtpu_version = f'0.1.dev{_date}' -_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}+nightly-py3-none-any.whl' +_libtpu_version = f'0.0.5.dev{_date}' +_libtpu_storage_path = f'https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-{_libtpu_version}+nightly-py3-none-linux_x86_64.whl' _jax_version = f'0.4.36.dev{_date}' @@ -312,8 +312,14 @@ def run(self): }, extras_require={ # On Cloud TPU VM install with: - # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html - 'tpu': [f'libtpu-nightly=={_libtpu_version}', 'tpu-info'], + # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-wheels/index.html -f https://storage.googleapis.com/libtpu-releases/index.html + 'tpu': [ + f'libtpu=={_libtpu_version}', + 'tpu-info', + # This special version removes `libtpu.so` from any `libtpu-nightly` installations, + # since we have migrated to using the `libtpu.so` from the `libtpu` package. + "libtpu-nightly==0.1.dev20241010+nightly.cleanup" + ], # pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html 'pallas': [f'jaxlib=={_jax_version}', f'jax=={_jax_version}'], },