commit 2a40fe2dbf021f95f4b02645c29ff5c19fc4eb23
author Michael Gschwind <mikekg@fb.com> Wed Sep 06 23:33:15 2023 +0000
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Wed Sep 06 23:33:18 2023 +0000
tree 69a0c709636bceded6e6809423ecdc096bf4a5dd
parent 6a304ed1f271cd249b4277d1fa5b1d219d0a1609

[experimental] use EXCEPT_FOR env to suppress CPU tests from GPU RE (#108672)

Summary:
[experimental] use EXCEPT_FOR env to suppress CPU tests from GPU RE -- alternative implementation to D48997976 using preexisting PYTORCH_TESTING_DEVICE_EXCEPT_FOR facility and building remaining logic (for assert-positive listers like test_transformers)  on top of that.

Goal: save ~100 GPU (10% of capacity), enables us to fund more aggressive PyPer unit testing on GPU RE

Test Plan: sandcastle, github

Differential Revision: D48998582

Pull Request resolved: https://github.com/pytorch/pytorch/pull/108672
Approved by: https://github.com/bertmaher

test/test_transformers.py

diff --git a/test/test_transformers.py b/test/test_transformers.py
index bbed844..0f53148 100644
--- a/test/test_transformers.py
+++ b/test/test_transformers.py
@@ -25,6 +25,7 @@
     set_default_dtype,
     gradcheck,
     make_tensor,
+    NOTEST_CPU
 )
 
 
@@ -2836,7 +2837,11 @@
         self.assertEqual(actual.contiguous(), math_ref.contiguous(), atol=1e-3, rtol=1e-2)
 
 
-device_types = ("cpu", "cuda")
+if NOTEST_CPU:
+    device_types = ("cuda", )
+else:
+    device_types = ("cpu", "cuda")
+
 instantiate_device_type_tests(TestTransformers, globals(), only_for=device_types)
 instantiate_device_type_tests(TestSDPAFailureModes, globals(), only_for=device_types)
 instantiate_device_type_tests(TestSDPA, globals(), only_for=device_types)

torch/testing/_internal/common_utils.py

diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
index 1f33d93..447f252 100644
--- a/torch/testing/_internal/common_utils.py
+++ b/torch/testing/_internal/common_utils.py
@@ -1074,6 +1074,11 @@
 
 BUILD_WITH_CAFFE2 = torch.onnx._CAFFE2_ATEN_FALLBACK
 
+def split_if_not_empty(x: str):
+    return x.split(",") if len(x) != 0 else []
+
+NOTEST_CPU = "cpu" in split_if_not_empty(os.getenv('PYTORCH_TESTING_DEVICE_EXCEPT_FOR', ''))
+
 # Python 2.7 doesn't have spawn
 TestEnvironment.def_flag("NO_MULTIPROCESSING_SPAWN", env_var="NO_MULTIPROCESSING_SPAWN")
 TestEnvironment.def_flag("TEST_WITH_ASAN", env_var="PYTORCH_TEST_WITH_ASAN")