Fix race condition in Fabric test (#17002)

Lightning-AI · Mar 8, 2023 · aa7f252 · aa7f252
1 parent d7e2ab9
commit aa7f252
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tests/tests_fabric/parity/test_parity_ddp.py b/tests/tests_fabric/parity/test_parity_ddp.py
@@ -125,7 +125,6 @@ def train_fabric_ddp(fabric):
     return model.state_dict(), torch.tensor(iteration_timings), memory_stats
 
 
-@pytest.mark.flaky(reruns=3)
 @RunIf(standalone=True)
 @pytest.mark.usefixtures("reset_deterministic_algorithm", "reset_cudnn_benchmark")
 @pytest.mark.parametrize(
@@ -148,6 +147,9 @@ def test_parity_ddp(accelerator, devices, tolerance):
     fabric.barrier()
     cuda_reset()
     torch.distributed.destroy_process_group()
+    # sleep for a bit to avoid race conditions, since the very first call in `train_torch_ddp`
+    # is initializing a new process group
+    time.sleep(3)
 
     # Train with raw PyTorch
     state_dict_torch, timings_torch, memory_torch = train_torch_ddp(