File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change 22import gc
33import logging
44import math
5+ import shutil
56import subprocess
67import tempfile
8+ import time
79from pathlib import Path
810from typing import List
911
2527)
2628
2729
30+ EXPORT_RETRIES = 3
31+
32+
33+ def _clear_export_dir (model_dir ):
34+ for path in Path (model_dir ).iterdir ():
35+ if path .is_dir () and not path .is_symlink ():
36+ shutil .rmtree (path )
37+ else :
38+ path .unlink ()
39+
40+
2841def cli_export (command , model_dir ):
2942 p = Path (model_dir )
3043 if p .exists ():
@@ -34,11 +47,19 @@ def cli_export(command, model_dir):
3447 raise Exception (
3548 f"Existing directory { model_dir } is non-empty. Please remove it first."
3649 )
37- try :
38- subprocess .run (command , check = True )
39- print ("Export completed successfully." )
40- except subprocess .CalledProcessError as e :
41- print (f"Export failed with error: { e } " )
50+
51+ for attempt in range (1 , EXPORT_RETRIES + 1 ):
52+ try :
53+ subprocess .run (command , check = True )
54+ print ("Export completed successfully." )
55+ return
56+ except subprocess .CalledProcessError as e :
57+ print (f"Export attempt { attempt } /{ EXPORT_RETRIES } failed with error: { e } " )
58+ if attempt == EXPORT_RETRIES :
59+ raise
60+ if p .exists ():
61+ _clear_export_dir (model_dir )
62+ time .sleep (attempt * 10 )
4263
4364
4465def check_causal_lm_output_quality (
You can’t perform that action at this time.
0 commit comments