Add uploading support (#220)

Also bugfixes and updates to the packaging scripts.

Add uploading support (#220)
Also bugfixes and updates to the packaging scripts.
1e1a7a4a · Edward Z. Yang · GitHub · 636e1499 · 1e1a7a4a · 1e1a7a4a
Unverified Commit 1e1a7a4a authored Aug 08, 2019 by Edward Z. Yang Committed by GitHub Aug 08, 2019
5 changed files
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -2,12 +2,13 @@ version: 2.1
 # How to test the Linux jobs:
 #   - Install CircleCI local CLI: https://circleci.com/docs/2.0/local-cli/
-#   - circleci config process .circleci/config.yml > gen.yml && circleci local execute -c gen.yml --job binary_linux_wheel
+#   - circleci config process .circleci/config.yml > gen.yml && circleci local execute -c gen.yml --job binary_linux_wheel_py3.7
-#     - Replace binary_linux_wheel with the name of the job you want to test
+#     - Replace binary_linux_wheel_py3.7 with the name of the job you want to test.
+#       Job names are 'name:' key.
 binary_common: &binary_common
  parameters:
-    # Edit these defaults to do a release`
+    # Edit these defaults to do a release
    build_version:
      description: "version number of release binary; by default, build a nightly"
      type: string
@@ -29,7 +30,7 @@ binary_common: &binary_common
    BUILD_VERSION: << parameters.build_version >>
    PYTORCH_VERSION: << parameters.pytorch_version >>
    UNICODE_ABI: << parameters.unicode_abi >>
-    CUDA_VERSION: cpu
+    CU_VERSION: cpu
 jobs:
  circleci_consistency:
@@ -53,6 +54,10 @@ jobs:
      - run: packaging/build_wheel.sh
      - store_artifacts:
          path: dist
+      - persist_to_workspace:
+          root: dist
+          paths:
+            - "*"
  binary_linux_conda:
    <<: *binary_common
@@ -64,6 +69,10 @@ jobs:
      - run: packaging/build_conda.sh
      - store_artifacts:
          path: /opt/conda/conda-bld/linux-64
+      - persist_to_workspace:
+          root: /opt/conda/conda-bld/linux-64
+          paths:
+            - "*"
  binary_macos_wheel:
    <<: *binary_common
@@ -82,6 +91,10 @@ jobs:
            packaging/build_wheel.sh
      - store_artifacts:
          path: dist
+      - persist_to_workspace:
+          root: dist
+          paths:
+            - "*"
  binary_macos_conda:
    <<: *binary_common
@@ -98,11 +111,55 @@ jobs:
            packaging/build_conda.sh
      - store_artifacts:
          path: /Users/distiller/miniconda3/conda-bld/osx-64
+      - persist_to_workspace:
+          root: /Users/distiller/miniconda3/conda-bld/osx-64
+          paths:
+            - "*"
+  # Requires org-member context
+  binary_conda_upload:
+    docker:
+      - image: continuumio/miniconda
+    steps:
+      - attach_workspace:
+          at: ~/workspace
+      - run:
+          command: |
+            # Prevent credential from leaking
+            conda install -yq anaconda-client
+            set +x
+            anaconda login \
+                --username "$PYTORCH_BINARY_PJH5_CONDA_USERNAME" \
+                --password "$PYTORCH_BINARY_PJH5_CONDA_PASSWORD"
+            set -x
+            anaconda upload ~/workspace/*.tar.bz2 -u pytorch-nightly --label main --no-progress --force
+  # Requires org-member context
+  binary_wheel_upload:
+    docker:
+      - image: circleci/python:3.7
+    steps:
+      - attach_workspace:
+          at: ~/workspace
+      - checkout
+      - run:
+          command: |
+            pip install --user awscli
+            export PATH="$HOME/.local/bin:$PATH"
+            # Prevent credential from leaking
+            set +x
+            export AWS_ACCESS_KEY_ID="${PYTORCH_BINARY_AWS_ACCESS_KEY_ID}"
+            export AWS_SECRET_ACCESS_KEY="${PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY}"
+            set -x
+            for pkg in ~/workspace/*.whl; do
+              aws s3 cp "$pkg" "s3://pytorch/whl/nightly/" --acl public-read
+            done
 workflows:
  build:
    jobs:
      - circleci_consistency
      - binary_linux_wheel:
          name: binary_linux_wheel_py2.7
          python_version: "2.7"
@@ -158,4 +215,162 @@ workflows:
          python_version: "3.6"
      - binary_macos_conda:
          name: binary_macos_conda_py3.7
          python_version: "3.7"
\ No newline at end of file
+  nightly:
+    triggers:
+      - schedule:
+          cron: "0 9 * * *"
+          filters:
+            branches:
+              only:
+                - master
+    jobs:
+      - circleci_consistency
+      - binary_linux_wheel:
+          name: nightly_binary_linux_wheel_py2.7
+          python_version: "2.7"
+      - binary_wheel_upload:
+          name: nightly_binary_linux_wheel_py2.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_wheel_py2.7
+      - binary_linux_wheel:
+          name: nightly_binary_linux_wheel_py2.7_unicode
+          python_version: "2.7"
+          unicode_abi: "1"
+      - binary_wheel_upload:
+          name: nightly_binary_linux_wheel_py2.7_unicode_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_wheel_py2.7_unicode
+      - binary_linux_wheel:
+          name: nightly_binary_linux_wheel_py3.5
+          python_version: "3.5"
+      - binary_wheel_upload:
+          name: nightly_binary_linux_wheel_py3.5_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_wheel_py3.5
+      - binary_linux_wheel:
+          name: nightly_binary_linux_wheel_py3.6
+          python_version: "3.6"
+      - binary_wheel_upload:
+          name: nightly_binary_linux_wheel_py3.6_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_wheel_py3.6
+      - binary_linux_wheel:
+          name: nightly_binary_linux_wheel_py3.7
+          python_version: "3.7"
+      - binary_wheel_upload:
+          name: nightly_binary_linux_wheel_py3.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_wheel_py3.7
+      - binary_macos_wheel:
+          name: nightly_binary_macos_wheel_py2.7
+          python_version: "2.7"
+      - binary_wheel_upload:
+          name: nightly_binary_macos_wheel_py2.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_wheel_py2.7
+      - binary_macos_wheel:
+          name: nightly_binary_macos_wheel_py2.7_unicode
+          python_version: "2.7"
+          unicode_abi: "1"
+      - binary_wheel_upload:
+          name: nightly_binary_macos_wheel_py2.7_unicode_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_wheel_py2.7_unicode
+      - binary_macos_wheel:
+          name: nightly_binary_macos_wheel_py3.5
+          python_version: "3.5"
+      - binary_wheel_upload:
+          name: nightly_binary_macos_wheel_py3.5_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_wheel_py3.5
+      - binary_macos_wheel:
+          name: nightly_binary_macos_wheel_py3.6
+          python_version: "3.6"
+      - binary_wheel_upload:
+          name: nightly_binary_macos_wheel_py3.6_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_wheel_py3.6
+      - binary_macos_wheel:
+          name: nightly_binary_macos_wheel_py3.7
+          python_version: "3.7"
+      - binary_wheel_upload:
+          name: nightly_binary_macos_wheel_py3.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_wheel_py3.7
+      - binary_linux_conda:
+          name: nightly_binary_linux_conda_py2.7
+          python_version: "2.7"
+      - binary_conda_upload:
+          name: nightly_binary_linux_conda_py2.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_conda_py2.7
+      - binary_linux_conda:
+          name: nightly_binary_linux_conda_py3.5
+          python_version: "3.5"
+      - binary_conda_upload:
+          name: nightly_binary_linux_conda_py3.5_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_conda_py3.5
+      - binary_linux_conda:
+          name: nightly_binary_linux_conda_py3.6
+          python_version: "3.6"
+      - binary_conda_upload:
+          name: nightly_binary_linux_conda_py3.6_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_conda_py3.6
+      - binary_linux_conda:
+          name: nightly_binary_linux_conda_py3.7
+          python_version: "3.7"
+      - binary_conda_upload:
+          name: nightly_binary_linux_conda_py3.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_linux_conda_py3.7
+      - binary_macos_conda:
+          name: nightly_binary_macos_conda_py2.7
+          python_version: "2.7"
+      - binary_conda_upload:
+          name: nightly_binary_macos_conda_py2.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_conda_py2.7
+      - binary_macos_conda:
+          name: nightly_binary_macos_conda_py3.5
+          python_version: "3.5"
+      - binary_conda_upload:
+          name: nightly_binary_macos_conda_py3.5_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_conda_py3.5
+      - binary_macos_conda:
+          name: nightly_binary_macos_conda_py3.6
+          python_version: "3.6"
+      - binary_conda_upload:
+          name: nightly_binary_macos_conda_py3.6_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_conda_py3.6
+      - binary_macos_conda:
+          name: nightly_binary_macos_conda_py3.7
+          python_version: "3.7"
+      - binary_conda_upload:
+          name: nightly_binary_macos_conda_py3.7_upload
+          context: org-member
+          requires:
+            - nightly_binary_macos_conda_py3.7
--- a/.circleci/config.yml.in
+++ b/.circleci/config.yml.in
@@ -2,12 +2,13 @@ version: 2.1
 # How to test the Linux jobs:
 #   - Install CircleCI local CLI: https://circleci.com/docs/2.0/local-cli/
-#   - circleci config process .circleci/config.yml > gen.yml && circleci local execute -c gen.yml --job binary_linux_wheel
+#   - circleci config process .circleci/config.yml > gen.yml && circleci local execute -c gen.yml --job binary_linux_wheel_py3.7
-#     - Replace binary_linux_wheel with the name of the job you want to test
+#     - Replace binary_linux_wheel_py3.7 with the name of the job you want to test.
+#       Job names are 'name:' key.
 binary_common: &binary_common
  parameters:
-    # Edit these defaults to do a release`
+    # Edit these defaults to do a release
    build_version:
      description: "version number of release binary; by default, build a nightly"
      type: string
@@ -29,7 +30,7 @@ binary_common: &binary_common
    BUILD_VERSION: << parameters.build_version >>
    PYTORCH_VERSION: << parameters.pytorch_version >>
    UNICODE_ABI: << parameters.unicode_abi >>
-    CUDA_VERSION: cpu
+    CU_VERSION: cpu
 jobs:
  circleci_consistency:
@@ -53,6 +54,10 @@ jobs:
      - run: packaging/build_wheel.sh
      - store_artifacts:
          path: dist
+      - persist_to_workspace:
+          root: dist
+          paths:
+            - "*"
  binary_linux_conda:
    <<: *binary_common
@@ -64,6 +69,10 @@ jobs:
      - run: packaging/build_conda.sh
      - store_artifacts:
          path: /opt/conda/conda-bld/linux-64
+      - persist_to_workspace:
+          root: /opt/conda/conda-bld/linux-64
+          paths:
+            - "*"
  binary_macos_wheel:
    <<: *binary_common
@@ -82,6 +91,10 @@ jobs:
            packaging/build_wheel.sh
      - store_artifacts:
          path: dist
+      - persist_to_workspace:
+          root: dist
+          paths:
+            - "*"
  binary_macos_conda:
    <<: *binary_common
@@ -98,23 +111,97 @@ jobs:
            packaging/build_conda.sh
      - store_artifacts:
          path: /Users/distiller/miniconda3/conda-bld/osx-64
+      - persist_to_workspace:
+          root: /Users/distiller/miniconda3/conda-bld/osx-64
+          paths:
+            - "*"
-workflows:
+  # Requires org-member context
-  build:
+  binary_conda_upload:
-    jobs:
+    docker:
-      - circleci_consistency
+      - image: continuumio/miniconda
-{%- for btype in ["wheel", "conda"] -%}
+    steps:
-{%-   for os in ["linux", "macos"] -%}
+      - attach_workspace:
-{%-     for python_version in ["2.7", "3.5", "3.6", "3.7"] %}
+          at: ~/workspace
-      - binary_{{os}}_{{btype}}:
+      - run:
-          name: binary_{{os}}_{{btype}}_py{{python_version}}
+          command: |
-          python_version: "{{python_version}}"
+            # Prevent credential from leaking
-{%-       if btype == "wheel" and python_version == "2.7" %}
+            conda install -yq anaconda-client
+            set +x
+            anaconda login \
+                --username "$PYTORCH_BINARY_PJH5_CONDA_USERNAME" \
+                --password "$PYTORCH_BINARY_PJH5_CONDA_PASSWORD"
+            set -x
+            anaconda upload ~/workspace/*.tar.bz2 -u pytorch-nightly --label main --no-progress --force
+  # Requires org-member context
+  binary_wheel_upload:
+    docker:
+      - image: circleci/python:3.7
+    steps:
+      - attach_workspace:
+          at: ~/workspace
+      - checkout
+      - run:
+          command: |
+            pip install --user awscli
+            export PATH="$HOME/.local/bin:$PATH"
+            # Prevent credential from leaking
+            set +x
+            export AWS_ACCESS_KEY_ID="${PYTORCH_BINARY_AWS_ACCESS_KEY_ID}"
+            export AWS_SECRET_ACCESS_KEY="${PYTORCH_BINARY_AWS_SECRET_ACCESS_KEY}"
+            set -x
+            for pkg in ~/workspace/*.whl; do
+              aws s3 cp "$pkg" "s3://pytorch/whl/nightly/" --acl public-read
+            done
+{%- macro workflow(btype, os, python_version, unicode, prefix='', upload=False) %}
      - binary_{{os}}_{{btype}}:
-          name: binary_{{os}}_{{btype}}_py{{python_version}}_unicode
+          name: {{prefix}}binary_{{os}}_{{btype}}_py{{python_version}}{{ "_unicode" if unicode }}
          python_version: "{{python_version}}"
+{%-   if unicode %}
          unicode_abi: "1"
-{%-       endif -%}
+{%-   endif %}
+{%-   if upload %}
+      - binary_{{btype}}_upload:
+          name: {{prefix}}binary_{{os}}_{{btype}}_py{{python_version}}{{ "_unicode" if unicode }}_upload
+          context: org-member
+          requires:
+            - {{prefix}}binary_{{os}}_{{btype}}_py{{python_version}}{{ "_unicode" if unicode }}
+{%-   endif %}
+{%- endmacro %}
+{%- macro workflows(prefix='', upload=False) %}
+{%- for btype in ["wheel", "conda"] -%}
+{%-   for os in ["linux", "macos"] -%}
+{%-     for python_version in ["2.7", "3.5", "3.6", "3.7"] -%}
+{%-       for unicode in ([False, True] if btype == "wheel" and python_version == "2.7" else [False]) -%}
+          {{ workflow(btype, os, python_version, unicode, prefix=prefix, upload=upload) }}
+{%-       endfor -%}
 {%-     endfor -%}
 {%-   endfor -%}
-{%-   endfor -%}
+{%- endfor %}
+{%- endmacro %}
+workflows:
+  build:
+{%- if True %}
+    jobs:
+      - circleci_consistency
+      {{ workflows() }}
+  nightly:
+    triggers:
+      - schedule:
+          cron: "0 9 * * *"
+          filters:
+            branches:
+              only:
+                - master
+{%- endif %}
+    jobs:
+      - circleci_consistency
+      {{ workflows(prefix="nightly_", upload=True) }}
--- a/packaging/build_conda.sh
+++ b/packaging/build_conda.sh
@@ -4,10 +4,8 @@ set -ex
 script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
-setup_build_version 0.4.0
+export NO_CUDA_PACKAGE=1
-setup_cuda_suffix
+setup_env 0.4.0
-setup_macos
 export SOURCE_ROOT_DIR="$PWD"
 setup_conda_pytorch_constraint
 conda build $CONDA_CHANNEL_FLAGS --no-anaconda-upload --python "$PYTHON_VERSION" packaging/torchaudio
--- a/packaging/build_wheel.sh
+++ b/packaging/build_wheel.sh
@@ -4,10 +4,9 @@ set -ex
 script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 . "$script_dir/pkg_helpers.bash"
-setup_python
+export NO_CUDA_PACKAGE=1
-setup_cuda_suffix
+setup_env 0.4.0
-setup_build_version 0.4.0
+setup_wheel_python
-setup_macos
 "$script_dir/build_from_source.sh" "$(pwd)"  # Build static dependencies
 pip_install numpy future
 setup_pip_pytorch_version

--- a/packaging/pkg_helpers.bash
+++ b/packaging/pkg_helpers.bash
 # A set of useful bash functions for common functionality we need to do in
 # many build scripts
-# Respecting PYTHON_VERSION and UNICODE_ABI, add (or install) the correct
-# version of Python to perform a build.  Relevant to wheel builds.
-setup_python() {
-  if [[ "$(uname)" == Darwin ]]; then
-    eval "$(conda shell.bash hook)"
-    conda env remove -n "env$PYTHON_VERSION" || true
-    conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
-    conda activate "env$PYTHON_VERSION"
-  else
-    case "$PYTHON_VERSION" in
-      2.7)
-        if [[ -n "$UNICODE_ABI" ]]; then
-          python_abi=cp27-cp27mu
-        else
-          python_abi=cp27-cp27m
-        fi
-        ;;
-      3.5) python_abi=cp35-cp35m ;;
-      3.6) python_abi=cp36-cp36m ;;
-      3.7) python_abi=cp37-cp37m ;;
-      *)
-        echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION"
-        exit 1
-        ;;
-    esac
-    export PATH="/opt/python/$python_abi/bin:$PATH"
-  fi
-}
-# Fill CUDA_SUFFIX with CUDA_VERSION, but don't fill it for the default
+# Setup CUDA environment variables, based on CU_VERSION
-# CUDA version (that's a blank suffix)
+#
-setup_cuda_suffix() {
+# Inputs:
-  if [[ "$(uname)" == Darwin ]]; then
+#   CU_VERSION (cpu, cu92, cu100)
-    if [[ "$CUDA_VERSION" != "cpu" ]]; then
+#   NO_CUDA_PACKAGE (bool)
-      echo "CUDA_VERSION on OS X must be cpu"
+#
+# Outputs:
+#   VERSION_SUFFIX (e.g., "")
+#   PYTORCH_VERSION_SUFFIX (e.g., +cpu)
+#   WHEEL_DIR (e.g., cu100/)
+#   CUDA_HOME (e.g., /usr/local/cuda-9.2, respected by torch.utils.cpp_extension)
+#   FORCE_CUDA (respected by torchvision setup.py)
+#   NVCC_FLAGS (respected by torchvision setup.py)
+#
+# Precondition: CUDA versions are installed in their conventional locations in
+# /usr/local/cuda-*
+#
+# NOTE: Why VERSION_SUFFIX versus PYTORCH_VERSION_SUFFIX?  If you're building
+# a package with CUDA on a platform we support CUDA on, VERSION_SUFFIX ==
+# PYTORCH_VERSION_SUFFIX and everyone is happy.  However, if you are building a
+# package with only CPU bits (e.g., torchaudio), then VERSION_SUFFIX is always
+# empty, but PYTORCH_VERSION_SUFFIX is +cpu (because that's how you get a CPU
+# 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
-    export CPU_SUFFIX=""
+    if [[ "$(uname)" == Darwin ]]; then
+      export PYTORCH_VERSION_SUFFIX=""
+    else
+      export PYTORCH_VERSION_SUFFIX="+cpu"
+    fi
+    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 WHEEL_DIR=""
  else
-    case "$CUDA_VERSION" in
+    case "$CU_VERSION" in
-      10.0)
+      cu100)
-        export CUDA_SUFFIX="" ;; # default!
+        export PYTORCH_VERSION_SUFFIX=""
-      9.2)
+        export CUDA_HOME=/usr/local/cuda-10.0/
-        export CUDA_SUFFIX="+cu92" ;;
+        export FORCE_CUDA=1
+        # Hard-coding gencode flags is temporary situation until
+        # https://github.com/pytorch/pytorch/pull/23408 lands
+        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_75,code=sm_75 -gencode=arch=compute_50,code=compute_50"
+        ;;
+      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 CUDA_SUFFIX="+cpu" ;;
+        export PYTORCH_VERSION_SUFFIX="+cpu"
+        ;;
      *)
-        echo "Unrecognized CUDA_VERSION=$CUDA_VERSION"
+        echo "Unrecognized CU_VERSION=$CU_VERSION"
    esac
-    export CPU_SUFFIX="+cpu"
+    export VERSION_SUFFIX="$PYTORCH_VERSION_SUFFIX"
+    export WHEEL_DIR="$CU_VERSION/"
  fi
 }
-# If a package is cpu-only, we never provide a cuda suffix
+# Populate build version if necessary, and add version suffix
-setup_cpuonly_cuda_suffix() {
+#
-  export CUDA_SUFFIX=""
+# Inputs:
-  export CPU_SUFFIX=""
+#   BUILD_VERSION (e.g., 0.2.0 or empty)
-}
+#
+# Outputs:
-# Fill BUILD_VERSION and BUILD_NUMBER if it doesn't exist already with a nightly string
+#   BUILD_VERSION (e.g., 0.2.0.dev20190807+cpu)
-# Usage: setup_build_version 0.2
+#
+# Fill BUILD_VERSION if it doesn't exist already with a nightly string
+# Usage: setup_build_version 0.2.0
 setup_build_version() {
  if [[ -z "$BUILD_VERSION" ]]; then
-    export BUILD_VERSION="$1.dev$(date "+%Y%m%d")"
+    export BUILD_VERSION="$1.dev$(date "+%Y%m%d")$VERSION_SUFFIX"
+  else
+    export BUILD_VERSION="$BUILD_VERSION$VERSION_SUFFIX"
  fi
 }
@@ -75,11 +97,55 @@ setup_macos() {
  fi
 }
+# Top-level entry point for things every package will need to do
+#
+# Usage: setup_env 0.2.0
+setup_env() {
+  setup_cuda
+  setup_build_version "$1"
+  setup_macos
+}
 # Function to retry functions that sometimes timeout or have flaky failures
 retry () {
    $*  || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
 }
+# Inputs:
+#   PYTHON_VERSION (2.7, 3.5, 3.6, 3.7)
+#   UNICODE_ABI (bool)
+#
+# Outputs:
+#   PATH modified to put correct Python version in PATH
+#
+# Precondition: If Linux, you are in a soumith/manylinux-cuda* Docker image
+setup_wheel_python() {
+  if [[ "$(uname)" == Darwin ]]; then
+    eval "$(conda shell.bash hook)"
+    conda env remove -n "env$PYTHON_VERSION" || true
+    conda create -yn "env$PYTHON_VERSION" python="$PYTHON_VERSION"
+    conda activate "env$PYTHON_VERSION"
+  else
+    case "$PYTHON_VERSION" in
+      2.7)
+        if [[ -n "$UNICODE_ABI" ]]; then
+          python_abi=cp27-cp27mu
+        else
+          python_abi=cp27-cp27m
+        fi
+        ;;
+      3.5) python_abi=cp35-cp35m ;;
+      3.6) python_abi=cp36-cp36m ;;
+      3.7) python_abi=cp37-cp37m ;;
+      *)
+        echo "Unrecognized PYTHON_VERSION=$PYTHON_VERSION"
+        exit 1
+        ;;
+    esac
+    export PATH="/opt/python/$python_abi/bin:$PATH"
+  fi
+}
 # Install with pip a bit more robustly than the default
 pip_install() {
  retry pip install --progress-bar off "$@"
@@ -89,15 +155,20 @@ pip_install() {
 # version into PYTORCH_VERSION, if applicable
 setup_pip_pytorch_version() {
  if [[ -z "$PYTORCH_VERSION" ]]; then
-    # Install latest prerelease CPU version of torch, per our nightlies.
+    # Install latest prerelease version of torch, per our nightlies, consistent
-    pip_install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+    # with the requested cuda version
-    # CPU/CUDA variants of PyTorch have ABI compatible PyTorch.  Therefore, we
+    pip_install --pre torch -f "https://download.pytorch.org/whl/nightly/${WHEEL_DIR}torch_nightly.html"
-    # strip off the local package qualifier.
+    if [[ "$CUDA_VERSION" == "cpu" ]]; then
-    export PYTORCH_VERSION="$(pip show torch | grep ^Version: | sed 's/Version:  *//' | sed 's/+.\+//')"
+      # CUDA and CPU are ABI compatible on the CPU-only parts, so strip
+      # in this case
+      export PYTORCH_VERSION="$(pip show torch | grep ^Version: | sed 's/Version:  *//' | sed 's/+.\+//')"
+    else
+      export PYTORCH_VERSION="$(pip show torch | grep ^Version: | sed 's/Version:  *//')"
+    fi
  else
-    # TODO: Maybe add staging too
+    pip_install "torch==$PYTORCH_VERSION$CUDA_SUFFIX" \
-    pip_install "torch==$PYTORCH_VERSION" \
+      -f https://download.pytorch.org/whl/torch_stable.html \
-      -f https://download.pytorch.org/whl/torch_stable.html
+      -f https://download.pytorch.org/whl/nightly/torch_nightly.html
  fi
 }
@@ -105,30 +176,39 @@ setup_pip_pytorch_version() {
 # CONDA_CHANNEL_FLAGS with appropriate flags to retrieve these versions
 #
 # You MUST have populated CUDA_SUFFIX before hand.
-#
-# TODO: This is currently hard-coded for CPU-only case
 setup_conda_pytorch_constraint() {
  if [[ -z "$PYTORCH_VERSION" ]]; then
    export CONDA_CHANNEL_FLAGS="-c pytorch-nightly"
    export PYTORCH_VERSION="$(conda search --json 'pytorch[channel=pytorch-nightly]' | python -c "import sys, json, re; print(re.sub(r'\\+.*$', '', json.load(sys.stdin)['pytorch'][-1]['version']))")"
  else
-    export CONDA_CHANNEL_FLAGS="-c pytorch"
+    export CONDA_CHANNEL_FLAGS="-c pytorch -c pytorch-nightly"
  fi
  if [[ "$CUDA_VERSION" == cpu ]]; then
-    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==$PYTORCH_VERSION${CPU_SUFFIX}"
+    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==$PYTORCH_VERSION${PYTORCH_VERSION_SUFFIX}"
    export CONDA_PYTORCH_CONSTRAINT="- pytorch==$PYTORCH_VERSION"
  else
-    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${CUDA_SUFFIX}"
+    export CONDA_PYTORCH_BUILD_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${PYTORCH_VERSION_SUFFIX}"
-    export CONDA_PYTORCH_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${CUDA_SUFFIX}"
+    export CONDA_PYTORCH_CONSTRAINT="- pytorch==${PYTORCH_VERSION}${PYTORCH_VERSION_SUFFIX}"
  fi
 }
 # Translate CUDA_VERSION into CUDA_CUDATOOLKIT_CONSTRAINT
 setup_conda_cudatoolkit_constraint() {
-  if [[ "$CUDA_VERSION" == cpu ]]; then
+  export CONDA_CPUONLY_FEATURE=""
+  if [[ "$(uname)" == Darwin ]]; then
    export CONDA_CUDATOOLKIT_CONSTRAINT=""
  else
-    echo
+    case "$CUDA_VERSION" in
-    # TODO
+      10.0)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=10.0,<10.1 # [not osx]"
+        ;;
+      9.2)
+        export CONDA_CUDATOOLKIT_CONSTRAINT="- cudatoolkit >=9.2,<9.3 # [not osx]"
+        ;;
+      cpu)
+        export CONDA_CUDATOOLKIT_CONSTRAINT=""
+        export CONDA_CPUONLY_FEATURE="- cpuonly"
+        ;;
+    esac
  fi
 }