From 02731cab7ed972871035456dd91d649ce0b74b27 Mon Sep 17 00:00:00 2001 From: Michael Shtelma Date: Wed, 19 Mar 2025 21:11:31 +0100 Subject: [PATCH 1/3] Using TMPDIR env var as a temporary directory. If it's not set, checking if /local_disk0 exists, and if yes, using it for temp dir. --- composer/callbacks/checkpoint_saver.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/composer/callbacks/checkpoint_saver.py b/composer/callbacks/checkpoint_saver.py index a4f94d541d..fc411f82f2 100644 --- a/composer/callbacks/checkpoint_saver.py +++ b/composer/callbacks/checkpoint_saver.py @@ -274,7 +274,13 @@ def __init__( is_remote_folder = backend != '' if is_remote_folder: # If uploading to a remote path, use a temporary directory to save local checkpoints. - local_folder = os.path.join(tempfile.mkdtemp(), local_folder) + root_temp_folder = None + if os.environ.get('TMPDIR') is not None: + root_temp_folder = os.environ.get('TMPDIR') + elif pathlib.Path('/local_disk0').exists(): # Probably we are on MLR, so we have to use /local_disk0 + root_temp_folder = '/local_disk0/temp' + pathlib.Path(root_temp_folder).mkdir(exist_ok=True) + local_folder = os.path.join(tempfile.mkdtemp(prefix=root_temp_folder), local_folder) filename = str(filename) remote_file_name = str(remote_file_name) if remote_file_name is not None else None From 83dcd9529654a07bd764ec3f51e9b9c815ace913 Mon Sep 17 00:00:00 2001 From: Michael Shtelma Date: Wed, 19 Mar 2025 23:17:54 +0100 Subject: [PATCH 2/3] Formatting --- composer/callbacks/checkpoint_saver.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/composer/callbacks/checkpoint_saver.py b/composer/callbacks/checkpoint_saver.py index fc411f82f2..79e1f3e382 100644 --- a/composer/callbacks/checkpoint_saver.py +++ b/composer/callbacks/checkpoint_saver.py @@ -277,9 +277,9 @@ def __init__( root_temp_folder = None if os.environ.get('TMPDIR') is not None: root_temp_folder = os.environ.get('TMPDIR') - elif pathlib.Path('/local_disk0').exists(): # Probably we are on MLR, so we have to use /local_disk0 + elif pathlib.Path('/local_disk0').exists(): # Probably we are on MLR, so we have to use /local_disk0 root_temp_folder = '/local_disk0/temp' - pathlib.Path(root_temp_folder).mkdir(exist_ok=True) + os.makedirs(root_temp_folder, exist_ok=True) local_folder = os.path.join(tempfile.mkdtemp(prefix=root_temp_folder), local_folder) filename = str(filename) From ae174eefadb31008af8efebf04c141570a0dcec9 Mon Sep 17 00:00:00 2001 From: Michael Shtelma Date: Wed, 19 Mar 2025 23:18:29 +0100 Subject: [PATCH 3/3] Formatting --- composer/callbacks/checkpoint_saver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/composer/callbacks/checkpoint_saver.py b/composer/callbacks/checkpoint_saver.py index 79e1f3e382..0fd97c4008 100644 --- a/composer/callbacks/checkpoint_saver.py +++ b/composer/callbacks/checkpoint_saver.py @@ -277,7 +277,7 @@ def __init__( root_temp_folder = None if os.environ.get('TMPDIR') is not None: root_temp_folder = os.environ.get('TMPDIR') - elif pathlib.Path('/local_disk0').exists(): # Probably we are on MLR, so we have to use /local_disk0 + elif os.path.exists('/local_disk0/'): # Probably we are on MLR, so we have to use /local_disk0 root_temp_folder = '/local_disk0/temp' os.makedirs(root_temp_folder, exist_ok=True) local_folder = os.path.join(tempfile.mkdtemp(prefix=root_temp_folder), local_folder)