Attempt to fix flaky test (#688)

* Fix transformers cache race conditions
mosaicml · Oct 23, 2023 · 091ddca · 091ddca
1 parent f65b07e
commit 091ddca
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
@@ -174,6 +174,10 @@ def check_hf_model_equivalence(model1: PreTrainedModel,
 
 
 def delete_transformers_cache():
+    # Only delete the files on local rank 0, otherwise race conditions are created
+    if not dist.get_local_rank() == 0:
+        return
+
     hf_cache_home = os.path.expanduser(
         os.getenv(
             'HF_HOME',
@@ -434,6 +438,7 @@ def test_huggingface_conversion_callback(model: str, tmp_path: pathlib.Path,
                 loaded_model)
             check_hf_tokenizer_equivalence(tokenizer, loaded_tokenizer)
 
+    dist.barrier()
     delete_transformers_cache()