Use distutils.dir_util.copy_tree instead and add unit tests

Author: derekhh

.pylintrc

@@ -84,3 +84,7 @@ dummy-variables-rgx=_|unused_
 # Apply logging string format checks to calls on these modules.
 logging-modules=
     logging
+
+[TYPECHECK]
+ignored-modules=
+    distutils

src/sagemaker/local/image.py

@@ -27,6 +27,7 @@
 import tarfile
 import tempfile
 
+from distutils.dir_util import copy_tree
 from six.moves.urllib.parse import urlparse
 from threading import Thread
 
@@ -217,9 +218,9 @@ def retrieve_artifacts(self, compose_data, output_data_config, job_name):
             for volume in volumes:
                 host_dir, container_dir = volume.split(':')
                 if container_dir == '/opt/ml/model':
-                    sagemaker.local.utils.recursive_copy(host_dir, model_artifacts)
+                    copy_tree(host_dir, model_artifacts)
                 elif container_dir == '/opt/ml/output':
-                    sagemaker.local.utils.recursive_copy(host_dir, output_artifacts)
+                    copy_tree(host_dir, output_artifacts)
 
         # Tar Artifacts -> model.tar.gz and output.tar.gz
         model_files = [os.path.join(model_artifacts, name) for name in os.listdir(model_artifacts)]

src/sagemaker/local/utils.py

@@ -15,6 +15,7 @@
 import os
 import shutil
 
+from distutils.dir_util import copy_tree
 from six.moves.urllib.parse import urlparse
 
 
@@ -45,14 +46,15 @@ def move_to_destination(source, destination, job_name, sagemaker_session):
     Args:
         source (str): root directory to move
         destination (str): file:// or s3:// URI that source will be moved to.
+        job_name (str): SageMaker job name.
         sagemaker_session (sagemaker.Session): a sagemaker_session to interact with S3 if needed
 
     Returns:
         (str): destination URI
     """
     parsed_uri = urlparse(destination)
     if parsed_uri.scheme == 'file':
-        recursive_copy(source, parsed_uri.path)
+        copy_tree(source, parsed_uri.path)
         final_uri = destination
     elif parsed_uri.scheme == 's3':
         bucket = parsed_uri.netloc
@@ -64,25 +66,3 @@ def move_to_destination(source, destination, job_name, sagemaker_session):
 
     shutil.rmtree(source)
     return final_uri
-
-
-def recursive_copy(source, destination):
-    """Similar to shutil.copy but the destination directory can exist. Existing files will be overriden.
-    Args:
-        source (str): source path
-        destination (str): destination path
-    """
-    if not os.path.exists(destination):
-        os.makedirs(destination, exist_ok=True)
-
-    for root, dirs, files in os.walk(source):
-        root = os.path.relpath(root, source)
-        current_path = os.path.join(source, root)
-        target_path = os.path.join(destination, root)
-
-        for file in files:
-            shutil.copy(os.path.join(current_path, file), os.path.join(target_path, file))
-        for d in dirs:
-            new_dir = os.path.join(target_path, d)
-            if not os.path.exists(new_dir):
-                os.mkdir(os.path.join(target_path, d))

tests/unit/test_local_utils.py

@@ -17,15 +17,13 @@
 
 import sagemaker.local.utils
 
-BUCKET_NAME = 'some-nice-bucket'
-
 
 @patch('shutil.rmtree', Mock())
-@patch('sagemaker.local.utils.recursive_copy')
-def test_move_to_destination(recursive_copy):
-    # local files will just be recursive copied
-    sagemaker.local.utils.move_to_destination('/tmp/data', 'file:///target/dir/', 'job', None)
-    recursive_copy.assert_called()
+@patch('sagemaker.local.utils.copy_tree')
+def test_move_to_destination(copy_tree):
+    # local files will just be recursively copied
+    sagemaker.local.utils.move_to_destination('/tmp/data', 'file:///target/dir', 'job', None)
+    copy_tree.assert_called_with('/tmp/data', '/target/dir')
 
     # s3 destination will upload to S3
     sms = Mock()