From bbcdbbaffc2f0d5619f0e16c90a4c8b7e088c27f Mon Sep 17 00:00:00 2001 From: Valera Filev Date: Thu, 14 Sep 2023 16:31:55 +0400 Subject: [PATCH] Remove checkpoints only on main process (#1974) * Remove checkpoints only on main process shutil.rmtree might throw errors if called on multiply processes. Make a call only on main process * Apply style --- src/accelerate/accelerator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py index 881ada02e15..77e9c6ef962 100755 --- a/src/accelerate/accelerator.py +++ b/src/accelerate/accelerator.py @@ -2720,8 +2720,10 @@ def save_state(self, output_dir: str = None, **save_model_func_kwargs): os.makedirs(output_dir, exist_ok=True) if self.project_configuration.automatic_checkpoint_naming: folders = [os.path.join(output_dir, folder) for folder in os.listdir(output_dir)] - if self.project_configuration.total_limit is not None and ( - len(folders) + 1 > self.project_configuration.total_limit + if ( + self.project_configuration.total_limit is not None + and (len(folders) + 1 > self.project_configuration.total_limit) + and self.is_main_process ): def _inner(folder):