Add CUDA test on CCI (#586)

* Add GPU test * Run only on master
pytorch · May 6, 2020 · 867d669 · 867d669
1 parent b71b5ee
commit 867d669
Show file tree

Hide file tree

Showing 7 changed files with 156 additions and 29 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -221,7 +221,7 @@ jobs:
             docker tag ${image_name}:${CIRCLE_WORKFLOW_ID} 308535385114.dkr.ecr.us-east-1.amazonaws.com/${image_name}:${CIRCLE_WORKFLOW_ID}
             docker push 308535385114.dkr.ecr.us-east-1.amazonaws.com/${image_name}:${CIRCLE_WORKFLOW_ID}
 
-  unittest_linux:
+  unittest_linux_cpu:
     <<: *binary_common
     docker:
       - image: "pytorch/torchaudio_unittest_base:manylinux"
@@ -259,7 +259,47 @@ jobs:
       - store_test_results:
           path: test-results
 
-  unittest_windows:
+  unittest_linux_gpu:
+    <<: *binary_common
+    machine:
+      image: ubuntu-1604-cuda-10.1:201909-23
+    resource_class: gpu.small
+    environment:
+      image_name: "pytorch/torchaudio_unittest_base:manylinux-cuda10.1"
+    steps:
+      - checkout
+      - run:
+          name: Generate cache key
+          # This will refresh cache on Sundays, nightly build should generate new cache.
+          command: echo "$(date +"%Y-%U")" > .circleci-weekly
+      - restore_cache:
+
+          keys:
+            - env-v1-linux-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/linux/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+      - run:
+          name: Setup
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/setup_env.sh
+      - save_cache:
+
+          key: env-v1-linux-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/linux/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install torchaudio
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/install.sh
+      - run:
+          name: Run tests
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/run_test.sh
+      - run:
+          name: Post Process
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
+  unittest_windows_cpu:
     <<: *binary_common
     executor:
       name: windows-cpu
@@ -338,23 +378,41 @@ workflows:
           python_version: '3.8'
   unittest:
     jobs:
-      - unittest_linux:
-          name: unittest_linux_py3.6
+      - unittest_linux_cpu:
+          name: unittest_linux_cpu_py3.6
+          python_version: '3.6'
+      - unittest_linux_cpu:
+          name: unittest_linux_cpu_py3.7
+          python_version: '3.7'
+      - unittest_linux_cpu:
+          name: unittest_linux_cpu_py3.8
+          python_version: '3.8'
+      - unittest_linux_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_linux_gpu_py3.6
           python_version: '3.6'
-      - unittest_linux:
-          name: unittest_linux_py3.7
+      - unittest_linux_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_linux_gpu_py3.7
           python_version: '3.7'
-      - unittest_linux:
-          name: unittest_linux_py3.8
+      - unittest_linux_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_linux_gpu_py3.8
           python_version: '3.8'
-      - unittest_windows:
-          name: unittest_windows_py3.6
+      - unittest_windows_cpu:
+          name: unittest_windows_cpu_py3.6
           python_version: '3.6'
-      - unittest_windows:
-          name: unittest_windows_py3.7
+      - unittest_windows_cpu:
+          name: unittest_windows_cpu_py3.7
           python_version: '3.7'
-      - unittest_windows:
-          name: unittest_windows_py3.8
+      - unittest_windows_cpu:
+          name: unittest_windows_cpu_py3.8
           python_version: '3.8'
   nightly:
     jobs:

diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
@@ -221,7 +221,7 @@ jobs:
             docker tag ${image_name}:${CIRCLE_WORKFLOW_ID} 308535385114.dkr.ecr.us-east-1.amazonaws.com/${image_name}:${CIRCLE_WORKFLOW_ID}
             docker push 308535385114.dkr.ecr.us-east-1.amazonaws.com/${image_name}:${CIRCLE_WORKFLOW_ID}
 
-  unittest_linux:
+  unittest_linux_cpu:
     <<: *binary_common
     docker:
       - image: "pytorch/torchaudio_unittest_base:manylinux"
@@ -259,7 +259,47 @@ jobs:
       - store_test_results:
           path: test-results
 
-  unittest_windows:
+  unittest_linux_gpu:
+    <<: *binary_common
+    machine:
+      image: ubuntu-1604-cuda-10.1:201909-23
+    resource_class: gpu.small
+    environment:
+      image_name: "pytorch/torchaudio_unittest_base:manylinux-cuda10.1"
+    steps:
+      - checkout
+      - run:
+          name: Generate cache key
+          # This will refresh cache on Sundays, nightly build should generate new cache.
+          command: echo "$(date +"%Y-%U")" > .circleci-weekly
+      - restore_cache:
+          {% raw %}
+          keys:
+            - env-v1-linux-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/linux/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+          {% endraw %}
+      - run:
+          name: Setup
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/setup_env.sh
+      - save_cache:
+          {% raw %}
+          key: env-v1-linux-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/linux/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+          {% endraw %}
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install torchaudio
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/install.sh
+      - run:
+          name: Run tests
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/run_test.sh
+      - run:
+          name: Post Process
+          command: docker run -t --gpus all -v $PWD:$PWD -w $PWD "${image_name}" .circleci/unittest/linux/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
+  unittest_windows_cpu:
     <<: *binary_common
     executor:
       name: windows-cpu

diff --git a/.circleci/regenerate.py b/.circleci/regenerate.py
@@ -111,16 +111,21 @@ def indent(indentation, data_list):
 
 
 def unittest_workflows(indentation=6):
-    w = []
+    jobs = []
     for os_type in ["linux", "windows"]:
-        for python_version in PYTHON_VERSIONS:
-            w.append({
-                f"unittest_{os_type}": {
-                    "name": f"unittest_{os_type}_py{python_version}",
+        for device_type in ["cpu", "gpu"]:
+            if os_type == 'windows' and device_type == 'gpu':
+                continue
+            for python_version in PYTHON_VERSIONS:
+                job = {
+                    "name": f"unittest_{os_type}_{device_type}_py{python_version}",
                     "python_version": python_version,
                 }
-            })
-    return indent(indentation, w)
+
+                if device_type == 'gpu':
+                    job['filters'] = gen_filter_branch_tree('master')
+                jobs.append({f"unittest_{os_type}_{device_type}": job})
+    return indent(indentation, jobs)
 
 
 if __name__ == "__main__":

diff --git a/.circleci/unittest/linux/docker/.gitignore b/.circleci/unittest/linux/docker/.gitignore
@@ -1 +1,2 @@
 scripts/build_third_parties.sh
+Dockerfile.tmp
diff --git a/.circleci/unittest/linux/docker/Dockerfile b/.circleci/unittest/linux/docker/Dockerfile
@@ -45,7 +45,7 @@ RUN bash /scripts/build_third_parties.sh /
 ################################################################################
 # Build the final image
 ################################################################################
-FROM ubuntu:18.04
+FROM BASE_IMAGE
 RUN apt update && apt install -y \
         g++ \
         gfortran \

diff --git a/.circleci/unittest/linux/docker/build_and_push.sh b/.circleci/unittest/linux/docker/build_and_push.sh
@@ -2,12 +2,29 @@
 
 set -euo pipefail
 
+if [ $# -ne 1 ]; then
+    printf "Usage %s <CUDA_VERSION>\n\n" "$0"
+    exit 1
+fi
+
+if [ "$1" = "cpu" ]; then
+    base_image="ubuntu:18.04"
+    image="pytorch/torchaudio_unittest_base:manylinux"
+elif [[ "$1" =~ ^(9.2|10.1)$ ]]; then
+    base_image="nvidia/cuda:$1-runtime-ubuntu18.04"
+    image="pytorch/torchaudio_unittest_base:manylinux-cuda$1"
+else
+    printf "Unexpected <CUDA_VERSION> string: %s" "$1"
+    exit 1;
+fi
+
 cd "$( dirname "${BASH_SOURCE[0]}" )"
 
 root_dir="$(git rev-parse --show-toplevel)"
 cp "${root_dir}"/packaging/build_from_source.sh ./scripts/build_third_parties.sh
 
-tag="manylinux"
-image="pytorch/torchaudio_unittest_base:${tag}"
-docker build -t "${image}" .
+# docker build also accepts reading from STDIN
+# but in that case, no context (other files) can be passed, so we write out Dockerfile
+sed "s|BASE_IMAGE|${base_image}|g" Dockerfile > Dockerfile.tmp
+docker build -t "${image}" -f Dockerfile.tmp .
 docker push "${image}"
diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh
@@ -10,8 +10,14 @@ set -e
 eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
-printf "* Installing PyTorch nightly build"
-conda install -y -c pytorch-nightly pytorch cpuonly
+if [ -z "${CUDA_VERSION:-}" ] ; then
+    cudatoolkit="cpuonly"
+else
+    version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
+    cudatoolkit="cudatoolkit=${version}"
+fi
+printf "Installing PyTorch with %s\n" "${cudatoolkit}"
+conda install -y -c pytorch-nightly pytorch "${cudatoolkit}"
 
 printf "* Installing torchaudio\n"
 # Link codecs present at /third_party. See Dockerfile for how this is built