[test] fixed rerun_on_exception and adapted test cases (#487)

tests/test_layers/test_1d/test_1d.py

@@ -10,28 +10,15 @@ from colossalai.core import global_context as gpc
 from colossalai.logging import disable_existing_loggers
 from colossalai.initialize import launch
 from colossalai.utils import free_port
-
+from colossalai.testing import rerun_on_exception
 from checks_1d.check_layer_1d import *
 
-CONFIG = dict(
-    parallel=dict(
-        pipeline=dict(size=1),
-        tensor=dict(
-            size=4,
-            mode='1d'
-        )
-    ),
-)
+CONFIG = dict(parallel=dict(pipeline=dict(size=1), tensor=dict(size=4, mode='1d')),)
 
 
 def check_layer(rank, world_size, port):
     disable_existing_loggers()
-    launch(config=CONFIG,
-           rank=rank,
-           world_size=world_size,
-           host='localhost',
-           port=port,
-           backend='nccl')
+    launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
 
     check_linear_col()
     check_linear_row()
@@ -48,6 +35,7 @@ def check_layer(rank, world_size, port):
 
 
 @pytest.mark.dist
+@rerun_on_exception(exception_type=mp.ProcessRaisedException, pattern=".*Address already in use.*")
 def test_1d():
     world_size = 4
     run_func = partial(check_layer, world_size=world_size, port=free_port())