[test] skip tests when not enough GPUs are detected (#1090)

* [test] skip tests when not enough GPUs are detected

* polish code

* polish code

tests/test_data_pipeline_tensor_parallel/test_cifar_with_data_pipeline_tensor.py

@@ -10,7 +10,7 @@ import torch.multiprocessing as mp
 from colossalai.amp import AMP_TYPE
 from colossalai.trainer import Trainer, hooks
 from colossalai.context import ParallelMode
-from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus
 from colossalai.utils import free_port
 from colossalai.core import global_context as gpc
 from colossalai.logging import get_dist_logger
@@ -83,7 +83,7 @@ def run_trainer(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test requires 8 GPUs to execute")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_if_address_is_in_use()
 def test_hybrid_parallel():
     world_size = 8

tests/test_layers/test_3d/test_3d.py

@@ -9,7 +9,7 @@ from colossalai.core import global_context as gpc
 from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.utils import free_port
-from colossalai.testing import rerun_if_address_is_in_use
+from colossalai.testing import rerun_if_address_is_in_use, skip_if_not_enough_gpus
 from checks_3d.check_layer_3d import (check_classifier_given_embed_weight, check_classifier_no_given_weight,
                                       check_embed, check_layernorm, check_linear, check_loss, check_patch_embed,
                                       check_vocab_parallel_classifier_given_embed_weight,
@@ -38,7 +38,6 @@ def check_layer():
     check_loss()
     check_vocab_parallel_loss()
 
-
 def check_layer_and_operation(rank, world_size, port):
     disable_existing_loggers()
     launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
@@ -51,7 +50,7 @@ def check_layer_and_operation(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test requires 8 GPUs to execute")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_if_address_is_in_use()
 def test_3d():
     world_size = 8

tests/test_utils/test_checkpoint/test_checkpoint_1d.py

@@ -15,7 +15,7 @@ from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.utils import free_port, is_using_pp
 from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint
-from colossalai.testing import rerun_on_exception
+from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus
 
 
 def build_pipeline(model):
@@ -67,7 +67,7 @@ def check_checkpoint_1d(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test should be invoked with 8 GPUs")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
 def test_checkpoint_1d():
     world_size = 8

tests/test_utils/test_checkpoint/test_checkpoint_2d.py

@@ -15,7 +15,7 @@ from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.utils import free_port, get_current_device, is_using_pp
 from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint
-from colossalai.testing import rerun_on_exception
+from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus
 
 
 def build_pipeline(model):
@@ -67,7 +67,7 @@ def check_checkpoint_2d(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test should be invoked with 8 GPUs")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
 def test_checkpoint_2d():
     world_size = 8

tests/test_utils/test_checkpoint/test_checkpoint_2p5d.py

@@ -15,7 +15,7 @@ from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.utils import free_port, get_current_device, is_using_pp
 from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint
-from colossalai.testing import rerun_on_exception
+from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus
 
 
 def build_pipeline(model):
@@ -37,7 +37,6 @@ def build_pipeline(model):
 def check_equal(A, B):
     assert torch.allclose(A, B, rtol=1e-3, atol=1e-2)
 
-
 def check_checkpoint_2p5d(rank, world_size, port):
     config = dict(parallel=dict(pipeline=dict(size=2), tensor=dict(size=4, depth=1, mode="2.5d")),)
 
@@ -67,7 +66,7 @@ def check_checkpoint_2p5d(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test should be invoked with 8 GPUs")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
 def test_checkpoint_2p5d():
     world_size = 8

tests/test_utils/test_checkpoint/test_checkpoint_3d.py

@@ -15,7 +15,7 @@ from colossalai.initialize import launch
 from colossalai.logging import disable_existing_loggers
 from colossalai.utils import free_port, get_current_device, is_using_pp
 from colossalai.utils.checkpointing import gather_pipeline_parallel_state_dict, load_checkpoint, save_checkpoint
-from colossalai.testing import rerun_on_exception
+from colossalai.testing import rerun_on_exception, skip_if_not_enough_gpus
 
 
 def build_pipeline(model):
@@ -37,7 +37,6 @@ def build_pipeline(model):
 def check_equal(A, B):
     assert torch.allclose(A, B, rtol=1e-3, atol=1e-2)
 
-
 def check_checkpoint_3d(rank, world_size, port):
     config = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=8, mode="3d")),)
 
@@ -67,7 +66,7 @@ def check_checkpoint_3d(rank, world_size, port):
 
 
 @pytest.mark.dist
-@pytest.mark.skip("This test requires 8 GPUs to execute")
+@skip_if_not_enough_gpus(min_gpus=8)
 @rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
 def test_checkpoint_3d():
     world_size = 8