Add distributed barrier test fixture to ensure pytest cleans up resou…

…rces properly (#2694)

tests/fixtures/autouse_fixtures.py

@@ -47,6 +47,15 @@ def disable_wandb(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureReques
             monkeypatch.setenv('WANDB_PROJECT', 'pytest')
 
 
+@pytest.fixture(scope='session')
+def cleanup_dist():
+    """Ensure all dist tests clean up resources properly."""
+    yield
+    # Avoid race condition where a test is still writing to a file on one rank
+    # while the file system is being torn down on another rank.
+    dist.barrier()
+
+
 @pytest.fixture(autouse=True, scope='session')
 def configure_dist(request: pytest.FixtureRequest):
     # Configure dist globally when the world size is greater than 1,