output_dir not in ckpt dir

tests/torchtune/training/checkpointing/test_checkpointer_utils.py

@@ -11,6 +11,7 @@
 import torch
 from torchtune.models.llama2 import llama2, llama2_classifier
 from torchtune.training.checkpointing._utils import (
+    check_outdir_not_in_ckptdir,
     FormattedCheckpointFiles,
     safe_torch_load,
     update_state_dict_for_classifier,
@@ -218,3 +219,47 @@ def test_build_checkpoint_filenames(self, expected_filenames):
         formatted_files = FormattedCheckpointFiles.from_dict(formatted_file_dict)
         actual_filenames = formatted_files.build_checkpoint_filenames()
         assert actual_filenames == expected_filenames
+
+
+class TestCheckOutdirNotInCkptdir:
+    def test_sibling_directories(self):
+        # Sibling directories should pass without raising an error
+        ckpt_dir = Path("/path/to/ckpt")
+        out_dir = Path("/path/to/output")
+        check_outdir_not_in_ckptdir(ckpt_dir, out_dir)
+
+    def test_ckpt_dir_in_output_dir(self):
+        # out_dir is a parent of ckpt_dir, should pass without raising an error
+        ckpt_dir = Path("/path/to/output/ckpt_dir")
+        out_dir = Path("/path/to/output")
+        check_outdir_not_in_ckptdir(ckpt_dir, out_dir)
+
+    def test_equal_directories(self):
+        # Equal directories should raise a ValueError
+        ckpt_dir = Path("/path/to/ckpt")
+        out_dir = Path("/path/to/ckpt")
+        with pytest.raises(
+            ValueError,
+            match="The output directory cannot be the same as or a subdirectory of the checkpoint directory.",
+        ):
+            check_outdir_not_in_ckptdir(ckpt_dir, out_dir)
+
+    def test_output_dir_in_ckpt_dir(self):
+        # out_dir is a subdirectory of ckpt_dir, should raise a ValueError
+        ckpt_dir = Path("/path/to/ckpt")
+        out_dir = Path("/path/to/ckpt/subdir")
+        with pytest.raises(
+            ValueError,
+            match="The output directory cannot be the same as or a subdirectory of the checkpoint directory.",
+        ):
+            check_outdir_not_in_ckptdir(ckpt_dir, out_dir)
+
+    def test_output_dir_ckpt_dir_few_levels_down(self):
+        # out_dir is a few levels down in ckpt_dir, should raise a ValueError
+        ckpt_dir = Path("/path/to/ckpt")
+        out_dir = Path("/path/to/ckpt/subdir/another_subdir")
+        with pytest.raises(
+            ValueError,
+            match="The output directory cannot be the same as or a subdirectory of the checkpoint directory.",
+        ):
+            check_outdir_not_in_ckptdir(ckpt_dir, out_dir)

torchtune/training/checkpointing/_checkpointer.py

@@ -20,6 +20,7 @@
     ADAPTER_CONFIG_FNAME,
     ADAPTER_MODEL_FNAME,
     BASE_MODEL_DIRNAME,
+    check_outdir_not_in_ckptdir,
     copy_files,
     get_adapter_checkpoint_path,
     get_model_checkpoint_path,
@@ -148,14 +149,17 @@ def __init__(
         # TODO: support loading more than one file
         if len(checkpoint_files) != 1:
             raise ValueError(
-                "Currently we only support reading from a single torchtune checkpoint file. "
+                "Currently we only support reading from a single checkpoint file. "
                 f"Got {len(checkpoint_files)} files instead."
             )
 
         self._checkpoint_dir = Path(checkpoint_dir)
         self._resume_from_checkpoint = resume_from_checkpoint
         self._model_type = ModelType[model_type]
         self._output_dir = Path(output_dir)
+        check_outdir_not_in_ckptdir(
+            ckpt_dir=self._checkpoint_dir, out_dir=self._output_dir
+        )
         self._output_dir.mkdir(parents=True, exist_ok=True)
 
         # save all files in input_dir, except model weights and mapping, to output_dir
@@ -391,6 +395,9 @@ def __init__(
         self._checkpoint_dir = Path(checkpoint_dir)
         self._model_type = ModelType[model_type]
         self._output_dir = Path(output_dir)
+        check_outdir_not_in_ckptdir(
+            ckpt_dir=self._checkpoint_dir, out_dir=self._output_dir
+        )
         self._output_dir.mkdir(parents=True, exist_ok=True)
 
         # weight_map contains the state_dict key -> checkpoint file mapping so we can correctly
@@ -914,14 +921,17 @@ def __init__(
         # TODO: support loading more than one file
         if len(checkpoint_files) != 1:
             raise ValueError(
-                "Currently we only support reading from a single torchtune checkpoint file. "
+                "Currently we only support reading from a single checkpoint file. "
                 f"Got {len(checkpoint_files)} files instead."
             )
 
         self._checkpoint_dir = Path(checkpoint_dir)
         self._resume_from_checkpoint = resume_from_checkpoint
         self._model_type = ModelType[model_type]
         self._output_dir = Path(output_dir)
+        check_outdir_not_in_ckptdir(
+            ckpt_dir=self._checkpoint_dir, out_dir=self._output_dir
+        )
         self._output_dir.mkdir(parents=True, exist_ok=True)
 
         # save all files in input_dir, except model weights and mapping, to output_dir

torchtune/training/checkpointing/_utils.py

@@ -559,3 +559,23 @@ def validate_checkpoint_files(
     )
 
     return checkpoint_paths
+
+
+def check_outdir_not_in_ckptdir(ckpt_dir: Path, out_dir: Path) -> bool:
+    """
+    Checks that the output directory is not a subdirectory of the checkpoint directory.
+    This is necessary to avoid making copies of copies when geting config files from ckpt_dir.
+    """
+
+    # Resolve the absolute paths to avoid issues with relative paths
+    _ckpt_dir = ckpt_dir.resolve()
+    _out_dir = out_dir.resolve()
+
+    # Check if out_dir is the same as ckpt_dir or a subdirectory of it
+    if _out_dir == _ckpt_dir or _ckpt_dir in _out_dir.parents:
+        raise ValueError(
+            "The output directory cannot be the same as or a subdirectory of the checkpoint directory. "
+            f"Found {ckpt_dir=} and {out_dir=}."
+        )
+
+    return True