Updated wheel packaging for building nightlies. (#1197)

Signed-off-by: Edward Z. Yang <[email protected]>

Author: ezyang

packaging/wheel/linux_manywheel.sh

@@ -1,14 +1,23 @@
+#!/bin/bash
+set -ex
+
 if [ "$#" -ne 1 ]; then
     echo "Illegal number of parameters. Pass cuda version"
-    echo "CUDA version should be cu90, cu100 or cpu"
+    echo "CUDA version should be cu92, cu100 or cpu"
     exit 1
 fi
-export CUVER="$1" # cu90 cu100 cpu
+export CUVER="$1" # cu92 cu100 cpu
+
+if [[ "$CUVER" == "cu100" ]]; then
+  cu_suffix=""
+else
+  cu_suffix="+$CUVER"
+fi
 
-export TORCHVISION_BUILD_VERSION="0.3.0"
+export TORCHVISION_BUILD_VERSION="0.4.0.dev$(date "+%Y%m%d")${cu_suffix}"
 export TORCHVISION_BUILD_NUMBER="1"
+export TORCHVISION_LOCAL_VERSION_LABEL="$CUVER"
 export OUT_DIR="/remote/$CUVER"
-export TORCH_WHEEL="torch -f https://download.pytorch.org/whl/$CUVER/stable.html --no-index"
 
 pushd /opt/python
 DESIRED_PYTHON=(*/)
@@ -18,12 +27,33 @@ for desired_py in "${DESIRED_PYTHON[@]}"; do
 done
 
 OLD_PATH=$PATH
-git clone https://github.com/pytorch/vision -b v${TORCHVISION_BUILD_VERSION}
-pushd vision
+cd /tmp
+rm -rf vision
+git clone https://github.com/pytorch/vision
+
+cd /tmp/vision
+
 for PYDIR in "${python_installations[@]}"; do
     export PATH=$PYDIR/bin:$OLD_PATH
+    pip install --upgrade pip
     pip install numpy pyyaml future
-    pip install $TORCH_WHEEL
+
+    pip uninstall -y torch || true
+    pip uninstall -y torch_nightly || true
+
+    export TORCHVISION_PYTORCH_DEPENDENCY_NAME=torch_nightly
+    pip install torch_nightly -f https://download.pytorch.org/whl/nightly/$CUVER/torch_nightly.html
+    # CPU/CUDA variants of PyTorch have ABI compatible PyTorch for
+    # the CPU only bits.  Therefore, we
+    # strip off the local package qualifier, but ONLY if we're
+    # doing a CPU build.
+    if [[ "$CUVER" == "cpu" ]]; then
+        export TORCHVISION_PYTORCH_DEPENDENCY_VERSION="$(pip show torch_nightly | grep ^Version: | sed 's/Version: \+//' | sed 's/+.\+//')"
+    else
+        export TORCHVISION_PYTORCH_DEPENDENCY_VERSION="$(pip show torch_nightly | grep ^Version: | sed 's/Version: \+//')"
+    fi
+    echo "Building against ${TORCHVISION_PYTORCH_DEPENDENCY_VERSION}"
+
     pip install ninja
     python setup.py clean
     python setup.py bdist_wheel

setup.py

@@ -29,7 +29,7 @@ def get_dist(pkgname):
         return None
 
 
-version = '0.3.0a0'
+version = '0.4.0a0'
 sha = 'Unknown'
 package_name = os.getenv('TORCHVISION_PACKAGE_NAME', 'torchvision')
 
@@ -43,9 +43,13 @@ def get_dist(pkgname):
 if os.getenv('TORCHVISION_BUILD_VERSION'):
     assert os.getenv('TORCHVISION_BUILD_NUMBER') is not None
     build_number = int(os.getenv('TORCHVISION_BUILD_NUMBER'))
-    version = os.getenv('TORCHVISION_BUILD_VERSION')
+    base_version = os.getenv('TORCHVISION_BUILD_VERSION')
+    version = base_version
     if build_number > 1:
         version += '.post' + str(build_number)
+    local_label = os.getenv('TORCHVISION_LOCAL_VERSION_LABEL')
+    if local_label is not None:
+        version += '+' + local_label
 elif sha != 'Unknown':
     version += '+' + sha[:7]
 print("Building wheel {}-{}".format(package_name, version))
@@ -65,12 +69,17 @@ def write_version_file():
 
 readme = open('README.rst').read()
 
-pytorch_package_name = os.getenv('TORCHVISION_PYTORCH_DEPENDENCY_NAME', 'torch')
+pytorch_dep = os.getenv('TORCHVISION_PYTORCH_DEPENDENCY_NAME', 'torch')
+if os.getenv('TORCHVISION_PYTORCH_DEPENDENCY_VERSION'):
+    pytorch_dep += "==" + os.getenv('TORCHVISION_PYTORCH_DEPENDENCY_VERSION')
+    # torchvision has CUDA bits, thus, we must specify a local dependency
+    if local_label is not None:
+        pytorch_dep += '+' + local_label
 
 requirements = [
     'numpy',
     'six',
-    pytorch_package_name,
+    pytorch_dep,
 ]
 
 pillow_ver = ' >= 4.1.1'