From 091ddca0d456125ce9f4243680766259c0a7a7ca Mon Sep 17 00:00:00 2001
From: Daniel King <43149077+dakinggg@users.noreply.github.com>
Date: Mon, 23 Oct 2023 14:45:12 -0700
Subject: [PATCH] Attempt to fix flaky test (#688)

* Fix transformers cache race conditions
---
 tests/test_hf_conversion_script.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/test_hf_conversion_script.py b/tests/test_hf_conversion_script.py
index 5bc3ed6d5d..e7787754de 100644
--- 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()