Fix failing mps tests (#3145)

* Fix failing mps tests

* Fix failing tests

* tests

* remove unnecessary changes

* Fix flake8 errors

---------

Co-authored-by: vfdev <vfdev.5@gmail.com>

Author: pranavvp16

tests/ignite/distributed/comp_models/test_base.py

@@ -4,9 +4,6 @@
 from ignite.distributed.comp_models.base import _SerialModel, _torch_version_le_112, ComputationModel
 
 
-@pytest.mark.skipif(
-    _torch_version_le_112 and torch.backends.mps.is_available(), reason="Temporary skip if MPS is available"
-)
 def test_serial_model():
     _SerialModel.create_from_backend()
     model = _SerialModel.create_from_context()
@@ -19,6 +16,8 @@ def test_serial_model():
     assert model.get_node_rank() == 0
     if torch.cuda.is_available():
         assert model.device().type == "cuda"
+    elif _torch_version_le_112 and torch.backends.mps.is_available():
+        assert model.device().type == "mps"
     else:
         assert model.device().type == "cpu"
     assert model.backend() is None

tests/ignite/distributed/test_auto.py

@@ -12,7 +12,6 @@
 
 import ignite.distributed as idist
 from ignite.distributed.auto import auto_dataloader, auto_model, auto_optim, DistributedProxySampler
-from ignite.distributed.comp_models.base import _torch_version_le_112
 
 
 class DummyDS(Dataset):
@@ -180,16 +179,13 @@ def _test_auto_model_optimizer(ws, device):
         assert optimizer.backward_passes_per_step == backward_passes_per_step
 
 
-@pytest.mark.skipif(
-    _torch_version_le_112 and torch.backends.mps.is_available(), reason="Temporary skip if MPS is available"
-)
 def test_auto_methods_no_dist():
     _test_auto_dataloader(1, 1, batch_size=1)
     _test_auto_dataloader(1, 1, batch_size=10, num_workers=2)
     _test_auto_dataloader(1, 1, batch_size=10, sampler_name="WeightedRandomSampler")
     _test_auto_dataloader(1, 1, batch_size=10, sampler_name="DistributedSampler")
-
-    _test_auto_model_optimizer(1, "cuda" if torch.cuda.is_available() else "cpu")
+    device = idist.device()
+    _test_auto_model_optimizer(1, device)
 
 
 @pytest.mark.distributed

tests/ignite/distributed/test_launcher.py

@@ -8,7 +8,6 @@
 from packaging.version import Version
 
 import ignite.distributed as idist
-from ignite.distributed.comp_models.base import _torch_version_le_112
 from ignite.distributed.utils import has_hvd_support, has_native_dist_support, has_xla_support
 
 
@@ -258,11 +257,8 @@ def test_idist_parallel_n_procs_native(init_method, backend, get_fixed_dirname,
 
 
 @pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
-@pytest.mark.skipif(
-    _torch_version_le_112 and torch.backends.mps.is_available(), reason="Temporary skip if MPS is available"
-)
 def test_idist_parallel_no_dist():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = idist.device()
     with idist.Parallel(backend=None) as parallel:
         parallel.run(_test_func, ws=1, device=device, backend=None, true_init_method=None)
 

tests/ignite/distributed/utils/test_serial.py

@@ -1,4 +1,3 @@
-import pytest
 import torch
 
 import ignite.distributed as idist
@@ -15,13 +14,12 @@
 )
 
 
-@pytest.mark.skipif(
-    _torch_version_le_112 and torch.backends.mps.is_available(), reason="Temporary skip if MPS is available"
-)
 def test_no_distrib(capsys):
     assert idist.backend() is None
     if torch.cuda.is_available():
         assert idist.device().type == "cuda"
+    elif _torch_version_le_112 and torch.backends.mps.is_available():
+        assert idist.device().type == "mps"
     else:
         assert idist.device().type == "cpu"
     assert idist.get_rank() == 0
@@ -43,6 +41,8 @@ def test_no_distrib(capsys):
     assert "ignite.distributed.utils INFO: backend: None" in out[-1]
     if torch.cuda.is_available():
         assert "ignite.distributed.utils INFO: device: cuda" in out[-1]
+    elif _torch_version_le_112 and torch.backends.mps.is_available():
+        assert "ignite.distributed.utils INFO: device: mps" in out[-1]
     else:
         assert "ignite.distributed.utils INFO: device: cpu" in out[-1]
     assert "ignite.distributed.utils INFO: rank: 0" in out[-1]