From 691838fb6bef70a3af662821c6fc60a990582786 Mon Sep 17 00:00:00 2001 From: Will Cromar Date: Fri, 8 Sep 2023 13:00:06 -0700 Subject: [PATCH] Update r2.1 wheel to be compatible with PyPI (#5550) * Update project metadata and remove useless files * Update README * Add manylinux platform tag * formatting --- README.md | 6 +- .../ansible/roles/build_srcs/tasks/main.yaml | 2 +- scripts/fixup_binary.py | 65 ------------------- setup.py | 34 +++++++--- 4 files changed, 30 insertions(+), 77 deletions(-) delete mode 100755 scripts/fixup_binary.py diff --git a/README.md b/README.md index f2a46040908..3dc9da23cc0 100644 --- a/README.md +++ b/README.md @@ -24,12 +24,12 @@ started: To install PyTorch/XLA a new VM: ``` -pip install torch~=2.0.0 https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0-cp38-cp38-linux_x86_64.whl +pip install torch~=2.1.0 torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/libtpu_releases.html ``` To update your existing training loop, make the following changes: -``` +```diff -import torch.multiprocessing as mp +import torch_xla.core.xla_model as xm +import torch_xla.distributed.parallel_loader as pl @@ -64,7 +64,7 @@ To update your existing training loop, make the following changes: If you're using `DistributedDataParallel`, make the following changes: -``` +```diff import torch.distributed as dist -import torch.multiprocessing as mp +import torch_xla.core.xla_model as xm diff --git a/infra/ansible/roles/build_srcs/tasks/main.yaml b/infra/ansible/roles/build_srcs/tasks/main.yaml index 2cd7cff6f1c..a17366d59ce 100644 --- a/infra/ansible/roles/build_srcs/tasks/main.yaml +++ b/infra/ansible/roles/build_srcs/tasks/main.yaml @@ -31,7 +31,7 @@ - name: Build PyTorch/XLA ansible.builtin.command: - cmd: python setup.py bdist_wheel + cmd: python setup.py bdist_wheel -p manylinux_2_28_x86_64 chdir: "{{ (src_root, 'pytorch/xla') | path_join }}" environment: "{{ env_vars }}" diff --git a/scripts/fixup_binary.py b/scripts/fixup_binary.py deleted file mode 100755 index 8d12d1c78f2..00000000000 --- a/scripts/fixup_binary.py +++ /dev/null @@ -1,65 +0,0 @@ -#!/usr/bin/env python - -import argparse -import glob -import os -import site -import subprocess - - -def find_torch_xla_site(site_paths): - for site_path in site_paths: - # If there is one named 'torch_xla', this is what we pick. - path = os.path.join(site_path, 'torch_xla', 'lib') - if os.path.isdir(path): - return [site_path, path] - dirs = glob.glob(os.path.join(site_path, 'torch_xla*')) - # Get the most recent one. - for xpath in sorted(dirs, key=os.path.getmtime): - path = os.path.join(xpath, 'lib') - if os.path.isdir(path): - return [site_path, path] - if os.path.isfile(os.path.join(xpath, 'libptxla.so')): - return [site_path, xpath, os.path.join(xpath, 'torch_xla', 'lib')] - raise RuntimeError('Unable to find torch_xla package in {}'.format(site_path)) - - -def find_torch_site(site_paths): - for site_path in site_paths: - path = os.path.join(site_path, 'torch', 'lib') - if os.path.isdir(path): - return [path] - raise RuntimeError('Unable to find torch package in {}'.format(site_path)) - - -def list_rpaths(path): - if subprocess.call(['patchelf', '--shrink-rpath', path]) != 0: - raise RuntimeError('Failed to shrink RPATH folders: {}'.format(path)) - return subprocess.check_output(['patchelf', '--print-rpath', - path]).decode('utf-8').strip('\n').split(':') - - -def set_rpaths(path, rpaths): - if subprocess.call(['patchelf', '--set-rpath', ':'.join(rpaths), path]) != 0: - raise RuntimeError('Failed to set RPATH folders {}: {}'.format( - rpaths, path)) - - -def fixup_binary(args): - site_paths = site.getsitepackages() - xla_rpaths = find_torch_xla_site(site_paths) - torch_rpaths = find_torch_site(site_paths) - rpaths = list_rpaths(args.binary) - rpaths = xla_rpaths + torch_rpaths + rpaths - set_rpaths(args.binary, rpaths) - - -if __name__ == '__main__': - arg_parser = argparse.ArgumentParser() - arg_parser.add_argument( - 'binary', - type=str, - metavar='BINARY', - help='The path to the binary to be patched') - args, files = arg_parser.parse_known_args() - fixup_binary(args) diff --git a/setup.py b/setup.py index a01a3ac843c..2b99d02a401 100644 --- a/setup.py +++ b/setup.py @@ -73,7 +73,6 @@ base_dir = os.path.dirname(os.path.abspath(__file__)) _libtpu_version = '0.1.dev20230809' -_libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl' def _get_build_mode(): @@ -307,15 +306,37 @@ def run(self): super().run() +# Read in README.md for our long_description +cwd = os.path.dirname(os.path.abspath(__file__)) +with open(os.path.join(cwd, "README.md"), encoding="utf-8") as f: + long_description = f.read() + setup( name=os.environ.get('TORCH_XLA_PACKAGE_NAME', 'torch_xla'), version=version, description='XLA bridge for PyTorch', + long_description=long_description, + long_description_content_type="text/markdown", url='https://github.com/pytorch/xla', author='PyTorch/XLA Dev Team', author_email='pytorch-xla@googlegroups.com', - # Exclude the build files. - packages=find_packages(exclude=['build']), + classifiers=[ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: BSD License", + "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development", + "Topic :: Software Development :: Libraries", + "Topic :: Software Development :: Libraries :: Python Modules", + "Programming Language :: C++", + "Programming Language :: Python :: 3", + ], + python_requires=">=3.8.0", + packages=find_packages(include=['torch_xla*']), ext_modules=[ BazelExtension('//:_XLAC.so'), ], @@ -334,12 +355,9 @@ def run(self): }, extras_require={ # On Cloud TPU VM install with: - # $ sudo pip3 install torch_xla[tpuvm] -f https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.11-cp38-cp38-linux_x86_64.whl - 'tpuvm': [f'libtpu-nightly @ {_libtpu_storage_path}'], + # pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/libtpu_releases.html + 'tpu': [f'libtpu-nightly=={_libtpu_version}'], }, - data_files=[ - 'scripts/fixup_binary.py', - ], cmdclass={ 'build_ext': BuildBazelExtension, 'clean': Clean,