Run tests in USE_PYTEST_LIST through run_tests (pytorch#95659)

Catherine Lee · pytorchmergebot · commit e3c5c369baff · 2023-02-28T22:09:01.000Z
Part of my effort to move everything to pytest and decrease the number of testrunner frameworks in ci Gives xmls but they might look a weird b/c module level tests vs tests in classes. Doesn't give skip/disable test infra because those are tied to classes. (for future ref, could either put tests in classes or move the check_if_enable stuff into a pytest hook) Tested in CI and checked that the same number of tests are run Pull Request resolved: pytorch#95659 Approved by: https://github.com/huydhn
diff --git a/pytest.ini b/pytest.ini
@@ -11,3 +11,5 @@ addopts =
 testpaths =
     test
 junit_logging_reruns = all
+filterwarnings =
+    ignore:Module already imported so cannot be rewritten.*hypothesis:pytest.PytestAssertRewriteWarning
diff --git a/test/distributed/elastic/events/lib_test.py b/test/distributed/elastic/events/lib_test.py
@@ -9,7 +9,6 @@
 
 import json
 import logging
-import unittest
 from dataclasses import asdict
 from unittest.mock import patch
 
@@ -21,10 +20,10 @@
     _get_or_create_logger,
     construct_and_record_rdzv_event,
 )
-from torch.testing._internal.common_utils import run_tests
+from torch.testing._internal.common_utils import run_tests, TestCase
 
 
-class EventLibTest(unittest.TestCase):
+class EventLibTest(TestCase):
     def assert_event(self, actual_event, expected_event):
         self.assertEqual(actual_event.name, expected_event.name)
         self.assertEqual(actual_event.source, expected_event.source)
@@ -59,7 +58,7 @@ def test_event_deser(self):
         deser_event = Event.deserialize(json_event)
         self.assert_event(event, deser_event)
 
-class RdzvEventLibTest(unittest.TestCase):
+class RdzvEventLibTest(TestCase):
     @patch("torch.distributed.elastic.events.record_rdzv_event")
     @patch("torch.distributed.elastic.events.get_logging_handler")
     def test_construct_and_record_rdzv_event(self, get_mock, record_mock):
diff --git a/test/distributed/pipeline/sync/skip/test_api.py b/test/distributed/pipeline/sync/skip/test_api.py
@@ -11,6 +11,7 @@
 from torch import nn
 
 from torch.distributed.pipeline.sync.skip import Namespace, skippable, stash
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_namespace_difference():
@@ -45,3 +46,7 @@ def forward(self, x):
 ))
 """.strip()
     )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_gpipe.py b/test/distributed/pipeline/sync/skip/test_gpipe.py
@@ -14,6 +14,7 @@
 from torch.distributed.pipeline.sync.skip import pop, skippable, stash
 from torch.distributed.pipeline.sync.skip.portal import PortalBlue, PortalCopy, PortalOrange
 from torch.distributed.pipeline.sync.utils import partition_model
+from torch.testing._internal.common_utils import run_tests
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@@ -108,3 +109,7 @@ def assert_grad_fn_is_not_portal(grad_fn, visited=None):
 
     output.local_value().sum().backward()
     assert input.grad.mean().item() == 1
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_inspect_skip_layout.py b/test/distributed/pipeline/sync/skip/test_inspect_skip_layout.py
@@ -10,6 +10,7 @@
 
 from torch.distributed.pipeline.sync.skip import Namespace, pop, skippable, stash
 from torch.distributed.pipeline.sync.skip.layout import inspect_skip_layout
+from torch.testing._internal.common_utils import run_tests
 
 
 class Pass(nn.Module):
@@ -111,3 +112,7 @@ def test_namespace():
 
     # p3 pops 'bar' before 'foo', but the plan is sorted by source partition index.
     assert policy == [[], [], [(0, ns1, "foo"), (1, ns2, "foo")]]
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_leak.py b/test/distributed/pipeline/sync/skip/test_leak.py
@@ -13,6 +13,7 @@
 from torch.distributed.pipeline.sync import Pipe, is_checkpointing, is_recomputing
 from torch.distributed.pipeline.sync.skip import pop, skippable, stash
 from torch.distributed.pipeline.sync.skip.tracker import current_skip_tracker
+from torch.testing._internal.common_utils import run_tests
 
 
 @skippable(stash=["skip"])
@@ -126,3 +127,7 @@ def deny(*args, **kwargs):
         model.eval()
         with torch.no_grad():
             model(input)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_portal.py b/test/distributed/pipeline/sync/skip/test_portal.py
@@ -12,6 +12,7 @@
 from torch.distributed.pipeline.sync.dependency import fork, join
 from torch.distributed.pipeline.sync.skip.portal import Portal
 from torch.distributed.pipeline.sync.stream import default_stream
+from torch.testing._internal.common_utils import run_tests
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@@ -155,3 +156,7 @@ def test_tensor_life_3_plus_1(self, new_portal):
         another_tensor = torch.rand(1, requires_grad=True)
         portal.put_tensor(another_tensor, tensor_life=1)
         portal.blue()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_stash_pop.py b/test/distributed/pipeline/sync/skip/test_stash_pop.py
@@ -12,6 +12,7 @@
 
 from torch.distributed.pipeline.sync.skip import pop, skippable, stash
 from torch.distributed.pipeline.sync.skip.tracker import SkipTracker, use_skip_tracker
+from torch.testing._internal.common_utils import run_tests
 
 
 @pytest.fixture(autouse=True)
@@ -136,3 +137,7 @@ def forward(self, input):
 
     l1 = Stash()
     l1(torch.tensor(42))
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_tracker.py b/test/distributed/pipeline/sync/skip/test_tracker.py
@@ -18,6 +18,7 @@
 from torch.distributed.pipeline.sync.skip import pop, skippable, stash
 from torch.distributed.pipeline.sync.skip.layout import SkipLayout
 from torch.distributed.pipeline.sync.skip.tracker import SkipTracker, SkipTrackerThroughPotals, current_skip_tracker
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_default_skip_tracker():
@@ -127,3 +128,7 @@ def test_tensor_life_with_checkpointing():
     with enable_recomputing():
         skip_tracker.save(batch, None, "test", tensor)
     assert skip_tracker.portals[(None, "test")].tensor_life == 0
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/skip/test_verify_skippables.py b/test/distributed/pipeline/sync/skip/test_verify_skippables.py
@@ -10,6 +10,7 @@
 from torch import nn
 
 from torch.distributed.pipeline.sync.skip import Namespace, skippable, verify_skippables
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_matching():
@@ -152,3 +153,7 @@ class Layer4(nn.Module):
     verify_skippables(
         nn.Sequential(Layer1().isolate(ns1), Layer2().isolate(ns1), Layer3().isolate(ns2), Layer4().isolate(ns2),)
     )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_balance.py b/test/distributed/pipeline/sync/test_balance.py
@@ -14,6 +14,7 @@
 
 from torch.distributed.pipeline.sync._balance import balance_by_size, balance_by_time, blockpartition
 from torch.distributed.pipeline.sync._balance.profile import layerwise_sandbox
+from torch.testing._internal.common_utils import run_tests
 
 skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
 
@@ -223,3 +224,7 @@ def test_already_has_grad():
 
     with pytest.raises(ValueError, match="some parameter already has gradient"):
         balance_by_time(1, model, sample, device="cpu")
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_bugs.py b/test/distributed/pipeline/sync/test_bugs.py
@@ -12,6 +12,7 @@
 import torch.nn.functional as F
 
 from torch.distributed.pipeline.sync import Pipe
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_python_autograd_function(setup_rpc):
@@ -137,3 +138,7 @@ def forward(self, x):
     y.norm().backward()
 
     assert y.to(torch.bool).tolist() == x.grad.to(torch.bool).tolist()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_checkpoint.py b/test/distributed/pipeline/sync/test_checkpoint.py
@@ -16,6 +16,7 @@
 from torch.distributed.pipeline.sync.checkpoint import Checkpointing, checkpoint, is_checkpointing, is_recomputing
 from torch.distributed.pipeline.sync.dependency import fork, join
 from torch.distributed.pipeline.sync.microbatch import Batch
+from torch.testing._internal.common_utils import run_tests
 
 devices = ["cpu"]
 if torch.cuda.is_available():
@@ -158,3 +159,7 @@ def forward(self, input):
 
     output = checkpoint(model, input)
     output[0].backward()
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_copy.py b/test/distributed/pipeline/sync/test_copy.py
@@ -11,6 +11,7 @@
 
 from torch.distributed.pipeline.sync.copy import Copy, Wait
 from torch.distributed.pipeline.sync.stream import CPUStream, current_stream, get_device, is_cuda, new_stream, use_stream
+from torch.testing._internal.common_utils import run_tests
 
 skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
 
@@ -68,3 +69,7 @@ def test_wait_multiple_tensors():
 
     assert a.grad_fn is b.grad_fn
     assert a.grad_fn.__class__ is Wait._backward_cls
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_deferred_batch_norm.py b/test/distributed/pipeline/sync/test_deferred_batch_norm.py
@@ -14,6 +14,7 @@
 from torch import nn, optim
 
 from torch.distributed.pipeline.sync.batchnorm import DeferredBatchNorm
+from torch.testing._internal.common_utils import run_tests
 
 CHUNKS = 4
 
@@ -192,3 +193,7 @@ def test_input_requiring_grad():
 
     assert not dbn.sum.requires_grad
     assert dbn.sum.grad_fn is None
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_dependency.py b/test/distributed/pipeline/sync/test_dependency.py
@@ -12,6 +12,7 @@
 import torch
 
 from torch.distributed.pipeline.sync.dependency import Fork, Join, fork, join
+from torch.testing._internal.common_utils import run_tests
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
@@ -144,3 +145,7 @@ def test_join_when_fork_requires_grad():
     assert not b.requires_grad
     b = join(b, p)
     assert b.requires_grad
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_inplace.py b/test/distributed/pipeline/sync/test_inplace.py
@@ -11,6 +11,7 @@
 from torch import nn
 
 from torch.distributed.pipeline.sync import Pipe
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_inplace_on_requires_grad(setup_rpc):
@@ -71,3 +72,7 @@ def forward(self, foo_bar):
     # The gradient of 'foo' should be 2, but it is 3 actually because
     # bar.add_(1) was executed twice due to checkpointing.
     assert foo.grad.item() == 2.0
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_microbatch.py b/test/distributed/pipeline/sync/test_microbatch.py
@@ -11,6 +11,7 @@
 import torch.cuda
 
 from torch.distributed.pipeline.sync.microbatch import Batch, check, gather, scatter
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_batch_atomic():
@@ -140,3 +141,7 @@ def test_scatter_multiple_tensors():
     assert list(b)[0].size() == (1, 1)
     assert list(a)[1].size() == (2, 2)
     assert list(b)[1].size() == (2, 2)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_phony.py b/test/distributed/pipeline/sync/test_phony.py
@@ -9,6 +9,7 @@
 import torch
 
 from torch.distributed.pipeline.sync.phony import get_phony
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_phony_size():
@@ -50,3 +51,7 @@ def forward(ctx, input):
     assert p1 is not p2
     assert p1.grad_fn is not None
     assert p2.grad_fn is None
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_pipe.py b/test/distributed/pipeline/sync/test_pipe.py
@@ -18,6 +18,7 @@
 
 from torch.distributed.pipeline.sync import Pipe, NoChunk, WithDevice
 from torch.distributed.pipeline.sync.pipe import PipeSequential
+from torch.testing._internal.common_utils import run_tests
 
 skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
 
@@ -819,3 +820,7 @@ def test_with_device_wrapper(setup_rpc):
     assert torch.device('cuda:0') == model(torch.rand(16, 16).cuda(0)).local_value().device
     assert [torch.device('cuda:0')] == model.devices
     assert torch.device('cuda:0') == fc2.weight.device
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_pipeline.py b/test/distributed/pipeline/sync/test_pipeline.py
@@ -7,6 +7,7 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 from torch.distributed.pipeline.sync.pipeline import _clock_cycles
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_clock_cycles():
@@ -29,3 +30,7 @@ def test_clock_cycles():
         [(3, 0), (2, 1)],
         [(3, 1)],
     ]
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_stream.py b/test/distributed/pipeline/sync/test_stream.py
@@ -21,6 +21,7 @@
     use_stream,
     wait_stream,
 )
+from torch.testing._internal.common_utils import run_tests
 
 skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
 
@@ -188,3 +189,7 @@ def test_record_stream_shifted_view(self, cuda_sleep):
         with torch.cuda.stream(stream_alloc):
             z = torch.rand(2, device=torch.device("cuda"))
         assert z.data_ptr() != data_ptr
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_transparency.py b/test/distributed/pipeline/sync/test_transparency.py
@@ -10,6 +10,7 @@
 from torch import nn
 
 from torch.distributed.pipeline.sync import Pipe
+from torch.testing._internal.common_utils import run_tests
 
 
 def test_simple_linears(setup_rpc):
@@ -43,3 +44,7 @@ def zero_grad(parameters):
 
     # Both grads should be identical.
     assert torch.allclose(grad_with_pipe, grad_without_pipe)
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributed/pipeline/sync/test_worker.py b/test/distributed/pipeline/sync/test_worker.py
@@ -14,6 +14,7 @@
 from torch.distributed.pipeline.sync.microbatch import Batch
 from torch.distributed.pipeline.sync.stream import CPUStream
 from torch.distributed.pipeline.sync.worker import Task, spawn_workers
+from torch.testing._internal.common_utils import run_tests
 
 
 class fake_device:
@@ -109,3 +110,7 @@ def test_worker_per_device():
         # 3: fake1, 4: fake2
         assert in_queues[3] is not in_queues[4]
         assert out_queues[3] is not out_queues[4]
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/test/distributions/test_constraints.py b/test/distributions/test_constraints.py
diff --git a/test/distributions/test_transforms.py b/test/distributions/test_transforms.py
diff --git a/test/distributions/test_utils.py b/test/distributions/test_utils.py
diff --git a/test/run_test.py b/test/run_test.py
diff --git a/test/test_typing.py b/test/test_typing.py
