File size: 5,115 Bytes
e84b953 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
#!/usr/bin/env python
# Sample usage:
# python env-setup.py --version 1.5 --apt-packages libomp5
import argparse
import collections
from datetime import datetime
import os
import platform
import re
import requests
import subprocess
import threading
import sys
VersionConfig = collections.namedtuple('VersionConfig',
['wheels', 'tpu', 'py_version', 'cuda_version'])
DEFAULT_CUDA_VERSION = '10.2'
OLDEST_VERSION = datetime.strptime('20200318', '%Y%m%d')
OLDEST_GPU_VERSION = datetime.strptime('20200707', '%Y%m%d')
DIST_BUCKET = 'gs://tpu-pytorch/wheels'
TORCH_WHEEL_TMPL = 'torch-{whl_version}-cp{py_version}-cp{py_version}m-linux_x86_64.whl'
TORCH_XLA_WHEEL_TMPL = 'torch_xla-{whl_version}-cp{py_version}-cp{py_version}m-linux_x86_64.whl'
TORCHVISION_WHEEL_TMPL = 'torchvision-{whl_version}-cp{py_version}-cp{py_version}m-linux_x86_64.whl'
def is_gpu_runtime():
return os.environ.get('COLAB_GPU', 0) == 1
def is_tpu_runtime():
return 'TPU_NAME' in os.environ
def update_tpu_runtime(tpu_name, version):
print(f'Updating TPU runtime to {version.tpu} ...')
try:
import cloud_tpu_client
except ImportError:
subprocess.call([sys.executable, '-m', 'pip', 'install', 'cloud-tpu-client'])
import cloud_tpu_client
client = cloud_tpu_client.Client(tpu_name)
client.configure_tpu_version(version.tpu)
print('Done updating TPU runtime')
def get_py_version():
version_tuple = platform.python_version_tuple()
return version_tuple[0] + version_tuple[1] # major_version + minor_version
def get_cuda_version():
if is_gpu_runtime():
# cuda available, install cuda wheels
return DEFAULT_CUDA_VERSION
def get_version(version):
cuda_version = get_cuda_version()
if version == 'nightly':
return VersionConfig(
'nightly', 'pytorch-nightly', get_py_version(), cuda_version)
version_date = None
try:
version_date = datetime.strptime(version, '%Y%m%d')
except ValueError:
pass # Not a dated nightly.
if version_date:
if cuda_version and version_date < OLDEST_GPU_VERSION:
raise ValueError(
f'Oldest nightly version build with CUDA available is {OLDEST_GPU_VERSION}')
elif not cuda_version and version_date < OLDEST_VERSION:
raise ValueError(f'Oldest nightly version available is {OLDEST_VERSION}')
return VersionConfig(f'nightly+{version}', f'pytorch-dev{version}',
get_py_version(), cuda_version)
version_regex = re.compile('^(\d+\.)+\d+$')
if not version_regex.match(version):
raise ValueError(f'{version} is an invalid torch_xla version pattern')
return VersionConfig(
version, f'pytorch-{version}', get_py_version(), cuda_version)
def install_vm(version, apt_packages, is_root=False):
dist_bucket = DIST_BUCKET
if version.cuda_version:
dist_bucket = os.path.join(
DIST_BUCKET, 'cuda/{}'.format(version.cuda_version.replace('.', '')))
torch_whl = TORCH_WHEEL_TMPL.format(
whl_version=version.wheels, py_version=version.py_version)
torch_whl_path = os.path.join(dist_bucket, torch_whl)
torch_xla_whl = TORCH_XLA_WHEEL_TMPL.format(
whl_version=version.wheels, py_version=version.py_version)
torch_xla_whl_path = os.path.join(dist_bucket, torch_xla_whl)
torchvision_whl = TORCHVISION_WHEEL_TMPL.format(
whl_version=version.wheels, py_version=version.py_version)
torchvision_whl_path = os.path.join(dist_bucket, torchvision_whl)
apt_cmd = ['apt-get', 'install', '-y']
apt_cmd.extend(apt_packages)
if not is_root:
# Colab/Kaggle run as root, but not GCE VMs so we need privilege
apt_cmd.insert(0, 'sudo')
installation_cmds = [
[sys.executable, '-m', 'pip', 'uninstall', '-y', 'torch', 'torchvision'],
['gsutil', 'cp', torch_whl_path, '.'],
['gsutil', 'cp', torch_xla_whl_path, '.'],
['gsutil', 'cp', torchvision_whl_path, '.'],
[sys.executable, '-m', 'pip', 'install', torch_whl],
[sys.executable, '-m', 'pip', 'install', torch_xla_whl],
[sys.executable, '-m', 'pip', 'install', torchvision_whl],
apt_cmd,
]
for cmd in installation_cmds:
subprocess.call(cmd)
def run_setup(args):
version = get_version(args.version)
# Update TPU
print('Updating... This may take around 2 minutes.')
if is_tpu_runtime():
update = threading.Thread(
target=update_tpu_runtime, args=(
args.tpu,
version,
))
update.start()
install_vm(version, args.apt_packages, is_root=not args.tpu)
if is_tpu_runtime():
update.join()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--version',
type=str,
default='20200515',
help='Versions to install (nightly, release version, or YYYYMMDD).',
)
parser.add_argument(
'--apt-packages',
nargs='+',
default=['libomp5'],
help='List of apt packages to install',
)
parser.add_argument(
'--tpu',
type=str,
help='[GCP] Name of the TPU (same zone, project as VM running script)',
)
args = parser.parse_args()
run_setup(args)
|