Replace shutil.copytree() to fix permission error (#251)

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* added unit tests, fixed other issues raised by review.

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes 1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

* Closes #1089

Signed-off-by: Joe Olson <joe.olson@ibm.com>

---------

Signed-off-by: Joe Olson <joe.olson@ibm.com>
Co-authored-by: Anh Uong <anh.uong@ibm.com>

Author: olson-ibm

build/accelerate_launch.py

@@ -35,6 +35,7 @@
     process_accelerate_launch_args,
     serialize_args,
     get_highest_checkpoint,
+    copy_checkpoint,
 )
 from tuning.utils.config_utils import get_json_config
 from tuning.config.tracker_configs import FileLoggingTrackerConfig
@@ -124,10 +125,8 @@ def main():
                 pt_checkpoint_dir,
                 original_output_dir,
             )
-            shutil.copytree(
-                os.path.join(tempdir, pt_checkpoint_dir),
-                original_output_dir,
-                dirs_exist_ok=True,
+            copy_checkpoint(
+                os.path.join(tempdir, pt_checkpoint_dir), original_output_dir
             )
         except Exception as e:  # pylint: disable=broad-except
             logging.error(traceback.format_exc())

build/utils.py

@@ -21,6 +21,23 @@
 # Third Party
 import torch
 from accelerate.commands.launch import launch_command_parser
+import shutil
+
+
+def copy_checkpoint(source, destination):
+    if not os.path.exists(destination):
+        os.makedirs(destination)
+        shutil.copystat(source, destination)
+    # Have a list of directory objects, now iterate over them.
+    for item in os.listdir(source):
+        source_file = os.path.join(source, item)
+        destination_file = os.path.join(destination, item)
+        if os.path.isdir(source_file):
+            # recursive call for subdirectories
+            copy_checkpoint(source_file, destination_file)
+        else:
+            # straight copy.
+            shutil.copy2(source_file, destination_file)
 
 
 def get_highest_checkpoint(dir_path):

tests/build/test_utils.py

@@ -17,12 +17,15 @@
 import json
 import os
 from unittest.mock import patch
+import tempfile
 
 # Third Party
 import pytest
+import filecmp
 
 # Local
 from build.utils import process_accelerate_launch_args
+from build.utils import copy_checkpoint
 
 HAPPY_PATH_DUMMY_CONFIG_PATH = os.path.join(
     os.path.dirname(__file__), "dummy_job_config.json"
@@ -108,3 +111,132 @@ def test_process_accelerate_launch_custom_config_file(patch_path_exists):
     temp_job_config = {"accelerate_launch_args": {"config_file": dummy_config_path}}
     args = process_accelerate_launch_args(temp_job_config)
     assert args.config_file == dummy_config_path
+
+
+class CopyCheckpointTestConfig:
+    def __init__(self, temp_root):
+
+        # Create the following file tree for testing:
+        # test_root
+        #   test_copytree_source
+        #      tf1.txt
+        #      tf2.txt
+        #      tf3.txt
+        #      subdir1
+        #         tf4.txt
+        #         tf5.txt
+        #         tf6.txt
+
+        self.test_root = temp_root
+        self.source_dir = os.path.join(self.test_root, "test_copytree_source")
+        self.source_sub_dir = os.path.join(self.source_dir, "subdir1")
+
+        os.mkdir(self.source_dir)
+        for file_number in range(2):
+            with open(
+                os.path.join(self.source_dir, f"tf{file_number+1}.txt"),
+                "a",
+                encoding="utf-8",
+            ) as f:
+                f.close()
+
+        os.mkdir(self.source_sub_dir)
+        for file_number in range(2):
+            with open(
+                os.path.join(self.source_sub_dir, f"tf{file_number+4}.txt"),
+                "a",
+                encoding="utf-8",
+            ) as f:
+                f.close()
+
+    def are_dir_trees_equal(self, dir1, dir2):
+
+        dirs_cmp = filecmp.dircmp(dir1, dir2)
+        if (
+            len(dirs_cmp.left_only) > 0
+            or len(dirs_cmp.right_only) > 0
+            or len(dirs_cmp.funny_files) > 0
+        ):
+            return False
+        (_, mismatch, errors) = filecmp.cmpfiles(
+            dir1, dir2, dirs_cmp.common_files, shallow=False
+        )
+        if len(mismatch) > 0 or len(errors) > 0:
+            return False
+        for common_dir in dirs_cmp.common_dirs:
+            new_dir1 = os.path.join(dir1, common_dir)
+            new_dir2 = os.path.join(dir2, common_dir)
+            if not self.are_dir_trees_equal(new_dir1, new_dir2):
+                return False
+        return True
+
+
+def test_copy_checkpoint_dest_dir_does_not_exist():
+
+    # Init source directory
+    with tempfile.TemporaryDirectory() as test_root:
+        config = CopyCheckpointTestConfig(test_root)
+
+        target_dir_does_not_exist = os.path.join(
+            config.test_root, "test_copytree_target"
+        )
+
+        # Execute the copy
+        copy_checkpoint(config.source_dir, target_dir_does_not_exist)
+        assert config.are_dir_trees_equal(config.source_dir, target_dir_does_not_exist)
+
+
+def test_copy_checkpoint_dest_dir_does_exist():
+
+    # Init source directory
+    with tempfile.TemporaryDirectory() as test_root:
+        config = CopyCheckpointTestConfig(test_root)
+
+        # Init target directory
+        target_dir_does_exist = os.path.join(config.test_root, "test_copytree_target2")
+        os.mkdir(target_dir_does_exist)
+        # Add a file to the target. This file will be overwritten during the copy.
+        with open(
+            os.path.join(target_dir_does_exist, "tf1.txt"),
+            "a",
+            encoding="utf-8",
+        ) as f:
+            f.close()
+        # Add a file to the target. This file does not exist in source.
+        with open(
+            os.path.join(target_dir_does_exist, "tf9.txt"),
+            "a",
+            encoding="utf-8",
+        ) as f:
+            f.close()
+        # Execute the copy
+        copy_checkpoint(config.source_dir, target_dir_does_exist)
+        assert os.path.exists(os.path.join(target_dir_does_exist, "tf9.txt"))
+        # Remove it so we can validate the dir trees are equal.
+        os.remove(os.path.join(target_dir_does_exist, "tf9.txt"))
+        assert config.are_dir_trees_equal(config.source_dir, target_dir_does_exist)
+
+
+def test_copy_checkpoint_dest_dir_not_writeable():
+
+    # Init source directory
+    with tempfile.TemporaryDirectory() as test_root:
+        config = CopyCheckpointTestConfig(test_root)
+
+        # Init target directory
+        target_dir_not_writeable = os.path.join(
+            config.test_root, "test_copytree_notwriteable"
+        )
+
+        os.makedirs(target_dir_not_writeable, mode=0o446)
+
+        # Execute the copy. Should FAIL
+        with pytest.raises(PermissionError) as e:
+            copy_checkpoint(config.source_dir, target_dir_not_writeable)
+        assert "Permission denied:" in str(e.value)
+
+
+def test_copy_checkpoint_source_dir_does_not_exist():
+    with pytest.raises(FileNotFoundError) as e:
+        copy_checkpoint("/doesnotexist", "/tmp")
+    assert "No such file or directory" in str(e.value)