[hackathon] ci: Only generate cuda tests for cuda configurations (#55522)

Summary:
Signed-off-by: Eli Uriegas <eliuriegas@fb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/55522

Reviewed By: walterddr

Differential Revision: D27634951

Pulled By: seemethere

fbshipit-source-id: 1dccaeb4bc8d0d53d61e467ba676c5c538fd4cf2
This commit is contained in:
Eli Uriegas 2021-04-08 17:29:22 -07:00 committed by Facebook GitHub Bot
parent 3498fde20e
commit 2ca45cb9e8
2 changed files with 19 additions and 3 deletions

View File

@ -23,6 +23,13 @@ if [[ "$BUILD_ENVIRONMENT" == *coverage* ]]; then
export PYTORCH_COLLECT_COVERAGE=1
fi
if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
# Used so that only cuda specific versions of tests are generated
# mainly used so that we're not spending extra cycles testing cpu
# devices on expensive gpu machines
export PYTORCH_TESTING_ONLY_FOR="cuda"
fi
if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then
export BUILD_SPLIT_CUDA=ON
fi

View File

@ -476,15 +476,24 @@ def instantiate_device_type_tests(generic_test_class, scope, except_for=None, on
generic_members = set(generic_test_class.__dict__.keys()) - set(empty_class.__dict__.keys())
generic_tests = [x for x in generic_members if x.startswith('test')]
# Derive defaults from environment variables if available, default is still none
# Usage:
# export PYTORCH_TESTING_DEVICE_ONLY_FOR=cuda,cpu
# export PYTORCH_TESTING_DEVICE_EXCEPT_FOR=xla
if only_for is None:
only_for = os.getenv("PYTORCH_TESTING_DEVICE_ONLY_FOR", ""). split(",")
if except_for is None:
except_for = os.getenv("PYTORCH_TESTING_DEVICE_EXCEPT_FOR", ""). split(",")
# Creates device-specific test cases
for base in device_type_test_bases:
# Skips bases listed in except_for
if except_for is not None and only_for is not None:
if except_for and only_for:
assert base.device_type not in except_for or base.device_type not in only_for,\
"same device cannot appear in except_for and only_for"
if except_for is not None and base.device_type in except_for:
if except_for and base.device_type in except_for:
continue
if only_for is not None and base.device_type not in only_for:
if only_for and base.device_type not in only_for:
continue
# Special-case for ROCm testing -- only test for 'cuda' i.e. ROCm device by default
# The except_for and only_for cases were already checked above. At this point we only need to check 'cuda'.