Improve huggingface robustness

.ci/scripts/test_huggingface_optimum_model.py

@@ -2,8 +2,10 @@
 import gc
 import logging
 import math
+import shutil
 import subprocess
 import tempfile
+import time
 from pathlib import Path
 from typing import List
 
@@ -25,6 +27,17 @@
 )
 
 
+EXPORT_RETRIES = 3
+
+
+def _clear_export_dir(model_dir):
+    for path in Path(model_dir).iterdir():
+        if path.is_dir() and not path.is_symlink():
+            shutil.rmtree(path)
+        else:
+            path.unlink()
+
+
 def cli_export(command, model_dir):
     p = Path(model_dir)
     if p.exists():
@@ -34,11 +47,19 @@ def cli_export(command, model_dir):
             raise Exception(
                 f"Existing directory {model_dir} is non-empty. Please remove it first."
             )
-    try:
-        subprocess.run(command, check=True)
-        print("Export completed successfully.")
-    except subprocess.CalledProcessError as e:
-        print(f"Export failed with error: {e}")
+
+    for attempt in range(1, EXPORT_RETRIES + 1):
+        try:
+            subprocess.run(command, check=True)
+            print("Export completed successfully.")
+            return
+        except subprocess.CalledProcessError as e:
+            print(f"Export attempt {attempt}/{EXPORT_RETRIES} failed with error: {e}")
+            if attempt == EXPORT_RETRIES:
+                raise
+            if p.exists():
+                _clear_export_dir(model_dir)
+            time.sleep(attempt * 10)
 
 
 def check_causal_lm_output_quality(