Enable CUDA tests for Windows (#637)

* Enable CUDA tests for Windows * Add back branch filters
pytorch · May 13, 2020 · 2dd0402 · 2dd0402
1 parent 36a1e90
commit 2dd0402
Show file tree

Hide file tree

Showing 4 changed files with 116 additions and 4 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -13,6 +13,12 @@ executors:
       image: windows-server-2019-vs2019:stable
       shell: bash.exe
 
+  windows-gpu:
+    machine:
+      resource_class: windows.gpu.nvidia.medium
+      image: windows-server-2019-nvidia:stable
+      shell: bash.exe
+
 binary_common: &binary_common
   parameters:
     # Edit these defaults to do a release
@@ -336,6 +342,45 @@ jobs:
       - store_test_results:
           path: test-results
 
+  unittest_windows_gpu:
+    <<: *binary_common
+    executor:
+      name: windows-gpu
+    environment:
+      CUDA_VERSION: "10.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-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+      - run:
+          name: Setup
+          command: .circleci/unittest/windows/scripts/setup_env.sh
+      - save_cache:
+
+          key: env-v1-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install torchaudio
+          command: .circleci/unittest/windows/scripts/install.sh
+      - run:
+          name: Run tests
+          command: .circleci/unittest/windows/scripts/run_test.sh
+      - run:
+          name: Post process
+          command: .circleci/unittest/windows/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
 workflows:
   build:
     jobs:
@@ -414,6 +459,24 @@ workflows:
       - unittest_windows_cpu:
           name: unittest_windows_cpu_py3.8
           python_version: '3.8'
+      - unittest_windows_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_windows_gpu_py3.6
+          python_version: '3.6'
+      - unittest_windows_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_windows_gpu_py3.7
+          python_version: '3.7'
+      - unittest_windows_gpu:
+          filters:
+            branches:
+              only: master
+          name: unittest_windows_gpu_py3.8
+          python_version: '3.8'
   nightly:
     jobs:
       - circleci_consistency:

diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
@@ -13,6 +13,12 @@ executors:
       image: windows-server-2019-vs2019:stable
       shell: bash.exe
 
+  windows-gpu:
+    machine:
+      resource_class: windows.gpu.nvidia.medium
+      image: windows-server-2019-nvidia:stable
+      shell: bash.exe
+
 binary_common: &binary_common
   parameters:
     # Edit these defaults to do a release
@@ -336,6 +342,45 @@ jobs:
       - store_test_results:
           path: test-results
 
+  unittest_windows_gpu:
+    <<: *binary_common
+    executor:
+      name: windows-gpu
+    environment:
+      CUDA_VERSION: "10.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-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+          {% endraw %}
+      - run:
+          name: Setup
+          command: .circleci/unittest/windows/scripts/setup_env.sh
+      - save_cache:
+          {% raw %}
+          key: env-v1-windows-{{ arch }}-py<< parameters.python_version >>-{{ checksum ".circleci/unittest/windows/scripts/environment.yml" }}-{{ checksum ".circleci-weekly" }}
+          {% endraw %}
+          paths:
+            - conda
+            - env
+      - run:
+          name: Install torchaudio
+          command: .circleci/unittest/windows/scripts/install.sh
+      - run:
+          name: Run tests
+          command: .circleci/unittest/windows/scripts/run_test.sh
+      - run:
+          name: Post process
+          command: .circleci/unittest/windows/scripts/post_process.sh
+      - store_test_results:
+          path: test-results
+
 workflows:
   build:
     jobs:

diff --git a/.circleci/regenerate.py b/.circleci/regenerate.py
@@ -114,8 +114,6 @@ def unittest_workflows(indentation=6):
     jobs = []
     for os_type in ["linux", "windows"]:
         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}",

diff --git a/.circleci/unittest/windows/scripts/install.sh b/.circleci/unittest/windows/scripts/install.sh
@@ -10,8 +10,14 @@ set -e
 eval "$(./conda/Scripts/conda.exe '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"
 IS_CONDA=true python setup.py develop