add gpu test

pytorch · Apr 27, 2020 · 4b38cb3 · 4b38cb3
1 parent 41b53be
commit 4b38cb3
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 15 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -233,7 +233,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/manylinux-cuda100"
@@ -268,6 +268,20 @@ jobs:
       - store_test_results:
           path: test-results
 
+  unittest_linux_gpu:
+    <<: *binary_common
+    machine:
+      image: ubuntu-1604-cuda-10.1:201909-23
+    resource_class: gpu.large
+    steps:
+      - run: nvidia-smi
+      - run: docker run --gpus all nvidia/cuda:9.0-base nvidia-smi
+    # filters:
+    #   branches:
+    #     only:
+    #       - master
+    #       - nightly
+
 workflows:
   build:
     jobs:
@@ -313,14 +327,23 @@ workflows:
           python_version: "3.6"
   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:
+          name: unittest_linux_gpu_py3.6
           python_version: '3.6'
-      - unittest_linux:
-          name: unittest_linux_py3.7
+      - unittest_linux_gpu:
+          name: unittest_linux_gpu_py3.7
           python_version: '3.7'
-      - unittest_linux:
-          name: unittest_linux_py3.8
+      - unittest_linux_gpu:
+          name: unittest_linux_gpu_py3.8
           python_version: '3.8'
   nightly:
     jobs:

diff --git a/.circleci/config.yml.in b/.circleci/config.yml.in
@@ -233,7 +233,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/manylinux-cuda100"
@@ -268,6 +268,20 @@ jobs:
       - store_test_results:
           path: test-results
 
+  unittest_linux_gpu:
+    <<: *binary_common
+    machine:
+      image: ubuntu-1604-cuda-10.1:201909-23
+    resource_class: gpu.large
+    steps:
+      - run: nvidia-smi
+      - run: docker run --gpus all nvidia/cuda:9.0-base nvidia-smi
+    # filters:
+    #   branches:
+    #     only:
+    #       - master
+    #       - nightly
+
 workflows:
   build:
     jobs:

diff --git a/.circleci/regenerate.py b/.circleci/regenerate.py
@@ -113,13 +113,14 @@ def indent(indentation, data_list):
 def unittest_workflows(indentation=6):
     w = []
     for os_type in ["linux"]:
-        for python_version in PYTHON_VERSIONS:
-            w.append({
-                f"unittest_{os_type}": {
-                    "name": f"unittest_{os_type}_py{python_version}",
-                    "python_version": python_version,
-                }
-            })
+        for device_type in ["cpu", "gpu"]:
+            for python_version in PYTHON_VERSIONS:
+                w.append({
+                    f"unittest_{os_type}_{device_type}": {
+                        "name": f"unittest_{os_type}_{device_type}_py{python_version}",
+                        "python_version": python_version,
+                    }
+                })
     return indent(indentation, w)