rewrite gpu/multi_gpu decorator with unittest.skipIf
+import os
+import unittest
+
+
try:
import pytest
_error = None
@@ -22,7 +26,8 @@ def get_error():
if _error is None:
- gpu = pytest.mark.gpu
+ _gpu_limit = int(os.getenv('CUPY_TEST_GPU_LIMIT', '-1'))
+
cudnn = pytest.mark.cudnn
slow = pytest.mark.slow
@@ -31,11 +36,32 @@ else:
check_available()
assert False # Not reachable
- gpu = _dummy_callable
cudnn = _dummy_callable
slow = _dummy_callable
def multi_gpu(gpu_num):
+ """Decorator to indicate number of GPUs required to run the test.
+
+ Tests can be annotated with this decorator (e.g., ``@multi_gpu(2)``) to
+ declare number of GPUs required to run. When running tests, if
+ ``CUPY_TEST_GPU_LIMIT`` environment variable is set to value greater
+ than or equals to 0, test cases that require GPUs more than the limit will
+ be skipped.
+ """
+
+ check_available()
+ return unittest.skipIf(
+ 0 <= _gpu_limit and _gpu_limit < gpu_num,
+ reason='{} GPUs required'.format(gpu_num))
+
+
+def gpu(f):
+ """Decorator to indicate that GPU is required to run the test.
+
+ Tests can be annotated with this decorator (e.g., ``@gpu``) to
+ declare that one GPU is required to run.
+ """
+
check_available()
- return pytest.mark.multi_gpu(gpu=gpu_num)
+ return multi_gpu(1)(pytest.mark.gpu(f))
