Refactor version suffix so conda packages don't get suffixes. (#1218)

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

packaging/build_conda.sh

@@ -4,6 +4,7 @@ set -ex
 script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
 
+export BUILD_TYPE=conda
 setup_env 0.5.0
 export SOURCE_ROOT_DIR="$PWD"
 setup_conda_pytorch_constraint

packaging/build_wheel.sh

@@ -4,6 +4,7 @@ set -ex
 script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
 
+export BUILD_TYPE=wheel
 setup_env 0.5.0
 setup_wheel_python
 pip_install numpy pyyaml future ninja

packaging/pkg_helpers.bash

@@ -7,6 +7,7 @@
 # Inputs:
 #   CU_VERSION (cpu, cu92, cu100)
 #   NO_CUDA_PACKAGE (bool)
+#   BUILD_TYPE (conda, wheel)
 #
 # Outputs:
 #   VERSION_SUFFIX (e.g., "")
@@ -27,28 +28,31 @@
 # version of a Python package.  But that doesn't apply if you're on OS X,
 # since the default CU_VERSION on OS X is cpu.
 setup_cuda() {
-  if [[ "$(uname)" == Darwin ]] || [[ -n "$NO_CUDA_PACKAGE" ]]; then
-    if [[ "$CU_VERSION" != "cpu" ]]; then
-      echo "CU_VERSION on OS X / package with no CUDA must be cpu"
-      exit 1
-    fi
-    if [[ "$(uname)" == Darwin ]]; then
-      export PYTORCH_VERSION_SUFFIX=""
-    else
-      export PYTORCH_VERSION_SUFFIX="+cpu"
-    fi
+
+  # First, compute version suffixes.  By default, assume no version suffixes
   export VERSION_SUFFIX=""
-    # NB: When there is no CUDA package available, we put these
-    # packages in the top-level directory, so they are eligible
-    # for selection even if you are otherwise trying to install
-    # a cu100 stack.  This differs from when there ARE CUDA packages
-    # available; then we don't want the cpu package; we want
-    # to give you as much goodies as possible.
+  export PYTORCH_VERSION_SUFFIX=""
   export WHEEL_DIR=""
-  else
+  # Wheel builds need suffixes (but not if they're on OS X, which never has suffix)
+  if [[ "$BUILD_TYPE" == "wheel" ]] && [[ "$(uname)" != Darwin ]]; then
+    # The default CUDA has no suffix
+    if [[ "$CU_VERSION" != "cu100" ]]; then
+      export PYTORCH_VERSION_SUFFIX="+$CU_VERSION"
+    fi
+    # Match the suffix scheme of pytorch, unless this package does not have
+    # CUDA builds (in which case, use default)
+    if [[ -z "$NO_CUDA_PACKAGE" ]]; then
+      export VERSION_SUFFIX="$PYTORCH_VERSION_SUFFIX"
+      # If the suffix is non-empty, we will use a wheel subdirectory
+      if [[ -n "$PYTORCH_VERSION_SUFFIX" ]]; then
+        export WHEEL_DIR="$PYTORCH_VERSION_SUFFIX/"
+      fi
+    fi
+  fi
+
+  # Now work out the CUDA settings
   case "$CU_VERSION" in
     cu100)
-        export PYTORCH_VERSION_SUFFIX=""
       export CUDA_HOME=/usr/local/cuda-10.0/
       export FORCE_CUDA=1
       # Hard-coding gencode flags is temporary situation until
@@ -57,25 +61,23 @@ setup_cuda() {
       ;;
     cu92)
       export CUDA_HOME=/usr/local/cuda-9.2/
-        export PYTORCH_VERSION_SUFFIX="+cu92"
       export FORCE_CUDA=1
       export NVCC_FLAGS="-gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_50,code=compute_50"
       ;;
     cpu)
-        export PYTORCH_VERSION_SUFFIX="+cpu"
       ;;
     *)
       echo "Unrecognized CU_VERSION=$CU_VERSION"
+      exit 1
+      ;;
   esac
-    export VERSION_SUFFIX="$PYTORCH_VERSION_SUFFIX"
-    export WHEEL_DIR="$CU_VERSION/"
-  fi
 }
 
 # Populate build version if necessary, and add version suffix
 #
 # Inputs:
 #   BUILD_VERSION (e.g., 0.2.0 or empty)
+#   VERSION_SUFFIX (e.g., +cpu)
 #
 # Outputs:
 #   BUILD_VERSION (e.g., 0.2.0.dev20190807+cpu)