Handle PR comments

Author: mvsusp

src/sagemaker/fw_utils.py

@@ -182,9 +182,16 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script,
     tmp = tempfile.mkdtemp()
 
     try:
-        source_files = _list_files_to_compress(script, directory) + dependencies
-        tar_file = sagemaker.utils.create_tar_file(source_files,
-                                                   os.path.join(tmp, _TAR_SOURCE_FILENAME))
+        if directory:
+            source_files = dependencies
+            dir_files = [directory]
+        else:
+            source_files = [script] + dependencies
+            dir_files = []
+
+        tar_file = sagemaker.utils.create_tar_file(source_files=source_files,
+                                                   dir_files=dir_files,
+                                                   target=os.path.join(tmp, _TAR_SOURCE_FILENAME))
 
         if kms_key:
             extra_args = {'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': kms_key}
@@ -198,14 +205,6 @@ def tar_and_upload_dir(session, bucket, s3_key_prefix, script,
     return UploadedCode(s3_prefix='s3://%s/%s' % (bucket, key), script_name=script_name)
 
 
-def _list_files_to_compress(script, directory):
-    if directory is None:
-        return [script]
-
-    basedir = directory if directory else os.path.dirname(script)
-    return [os.path.join(basedir, name) for name in os.listdir(basedir)]
-
-
 def framework_name_from_image(image_name):
     """Extract the framework and Python version from the image name.
 

src/sagemaker/local/image.py

@@ -234,11 +234,10 @@ def retrieve_artifacts(self, compose_data, output_data_config, job_name):
                 elif container_dir == '/opt/ml/output':
                     sagemaker.local.utils.recursive_copy(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)]
-        output_files = [os.path.join(output_artifacts, name) for name in os.listdir(output_artifacts)]
-        sagemaker.utils.create_tar_file(model_files, os.path.join(compressed_artifacts, 'model.tar.gz'))
-        sagemaker.utils.create_tar_file(output_files, os.path.join(compressed_artifacts, 'output.tar.gz'))
+        sagemaker.utils.create_tar_file(dir_files=[model_artifacts],
+                                        target=os.path.join(compressed_artifacts, 'model.tar.gz'))
+        sagemaker.utils.create_tar_file(dir_files=[output_artifacts],
+                                        target=os.path.join(compressed_artifacts, 'output.tar.gz'))
 
         if output_data_config['S3OutputPath'] == '':
             output_data = 'file://%s' % compressed_artifacts

src/sagemaker/utils.py

@@ -260,11 +260,14 @@ def download_folder(bucket_name, prefix, target, sagemaker_session):
         obj.download_file(file_path)
 
 
-def create_tar_file(source_files, target=None):
-    """Create a tar file containing all the source_files
+def create_tar_file(source_files=None, target=None, dir_files=None):
+    """Create a tar file containing all the source_files and the content of all dir_files
 
     Args:
         source_files (List[str]): List of file paths that will be contained in the tar file
+        target (str): target path of the tar file
+        dir_files (List[str]): List of directories which will have their contents copy into
+         the tar file
 
     Returns:
          (str): path to created tar file
@@ -275,10 +278,17 @@ def create_tar_file(source_files, target=None):
     else:
         _, filename = tempfile.mkstemp()
 
+    dir_files = dir_files or []
+    source_files = source_files or []
+
     with tarfile.open(filename, mode='w:gz') as t:
         for sf in source_files:
             # Add all files from the directory into the root of the directory structure of the tar
             t.add(sf, arcname=os.path.basename(sf))
+
+        for dir_file in dir_files:
+            t.add(dir_file, arcname=os.path.sep)
+
     return filename
 
 
@@ -323,13 +333,11 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
     new_model_name = 'model-%s.tar.gz' % sagemaker.utils.sagemaker_short_timestamp()
 
     with _tmpdir() as tmp:
-
         tmp_model_dir = os.path.join(tmp, 'model')
         os.mkdir(tmp_model_dir)
 
         model_from_s3 = model_uri.startswith('s3://')
         if model_from_s3:
-
             local_model_path = os.path.join(tmp, 'tar_file')
             download_file_from_url(model_uri, local_model_path, sagemaker_session)
 
@@ -349,10 +357,7 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
 
         shutil.copytree(dirname, code_dir)
 
-        files_to_compress = [os.path.join(tmp_model_dir, file)
-                             for file in os.listdir(tmp_model_dir)]
-
-        tar_file = sagemaker.utils.create_tar_file(files_to_compress, new_model_path)
+        tar_file = sagemaker.utils.create_tar_file(dir_files=[tmp_model_dir], target=new_model_path)
 
         if model_from_s3:
             url = parse.urlparse(model_uri)

tests/integ/test_tfs.py

@@ -51,8 +51,8 @@ def tfs_predictor(instance_type, sagemaker_session, tf_full_version):
 
 def tar_dir(directory, tmpdir):
 
-    source_files = [os.path.join(directory, name) for name in os.listdir(directory)]
-    return sagemaker.utils.create_tar_file(source_files, os.path.join(str(tmpdir), 'model.tar.gz'))
+    return sagemaker.utils.create_tar_file(dir_files=[directory],
+                                           target=os.path.join(str(tmpdir), 'model.tar.gz'))
 
 
 @pytest.fixture(scope='module')

tests/unit/test_utils.py

@@ -268,23 +268,57 @@ def test_download_file():
 
 @patch('tarfile.open')
 def test_create_tar_file_with_provided_path(open):
-    open.return_value = open
-    open.__enter__ = Mock()
-    open.__exit__ = Mock(return_value=None)
+    files = mock_tarfile(open)
+
     file_list = ['/tmp/a', '/tmp/b']
+
     path = sagemaker.utils.create_tar_file(file_list, target='/my/custom/path.tar.gz')
     assert path == '/my/custom/path.tar.gz'
+    assert files == [['/tmp/a', 'a'], ['/tmp/b', 'b']]
 
 
 @patch('tarfile.open')
-@patch('tempfile.mkstemp', Mock(return_value=(None, '/auto/generated/path')))
-def test_create_tar_file_with_auto_generated_path(open):
+def test_create_tar_file_with_directories(open):
+    files = mock_tarfile(open)
+
+    path = sagemaker.utils.create_tar_file(dir_files=['/tmp/a', '/tmp/b'],
+                                           target='/my/custom/path.tar.gz')
+    assert path == '/my/custom/path.tar.gz'
+    assert files == [['/tmp/a', '/'], ['/tmp/b', '/']]
+
+
+@patch('tarfile.open')
+def test_create_tar_file_with_files_and_directories(open):
+    files = mock_tarfile(open)
+
+    path = sagemaker.utils.create_tar_file(dir_files=['/tmp/a', '/tmp/b'],
+                                           source_files=['/tmp/c', '/tmp/d'],
+                                           target='/my/custom/path.tar.gz')
+    assert path == '/my/custom/path.tar.gz'
+    assert files == [['/tmp/c', 'c'], ['/tmp/d', 'd'], ['/tmp/a', '/'], ['/tmp/b', '/']]
+
+
+def mock_tarfile(open):
     open.return_value = open
+    files = []
+
+    def add_files(filename, arcname):
+        files.append([filename, arcname])
+
     open.__enter__ = Mock()
+    open.__enter__().add = add_files
     open.__exit__ = Mock(return_value=None)
-    file_list = ['/tmp/a', '/tmp/b']
-    path = sagemaker.utils.create_tar_file(file_list)
+    return files
+
+
+@patch('tarfile.open')
+@patch('tempfile.mkstemp', Mock(return_value=(None, '/auto/generated/path')))
+def test_create_tar_file_with_auto_generated_path(open):
+    files = mock_tarfile(open)
+
+    path = sagemaker.utils.create_tar_file(['/tmp/a', '/tmp/b'])
     assert path == '/auto/generated/path'
+    assert files == [['/tmp/a', 'a'], ['/tmp/b', 'b']]
 
 
 def write_file(path, content):