Handle PR comments

Author: mvsusp

src/sagemaker/utils.py

@@ -330,27 +330,24 @@ def repack_model(inference_script, source_directory, model_uri, sagemaker_sessio
         model_from_s3 = model_uri.startswith('s3://')
         if model_from_s3:
 
-            local_model_uri = os.path.join(tmp, 'tar_file')
-            download_file_from_url(model_uri, local_model_uri, sagemaker_session)
+            local_model_path = os.path.join(tmp, 'tar_file')
+            download_file_from_url(model_uri, local_model_path, sagemaker_session)
 
             new_model_path = os.path.join(tmp, new_model_name)
         else:
-            local_model_uri = model_uri.replace('file://', '')
-            new_model_path = os.path.join(os.path.dirname(local_model_uri), new_model_name)
+            local_model_path = model_uri.replace('file://', '')
+            new_model_path = os.path.join(os.path.dirname(local_model_path), new_model_name)
 
-        with tarfile.open(name=local_model_uri, mode='r:gz') as t:
+        with tarfile.open(name=local_model_path, mode='r:gz') as t:
             t.extractall(path=tmp_model_dir)
 
         code_dir = os.path.join(tmp_model_dir, 'code')
         if os.path.exists(code_dir):
             shutil.rmtree(code_dir, ignore_errors=True)
 
-        os.mkdir(code_dir)
+        dirname = source_directory if source_directory else os.path.dirname(inference_script)
 
-        source_files = _list_files(inference_script, source_directory)
-
-        for source_file in source_files:
-            shutil.copy(source_file, code_dir)
+        shutil.copytree(dirname, code_dir)
 
         files_to_compress = [os.path.join(tmp_model_dir, file)
                              for file in os.listdir(tmp_model_dir)]

tests/integ/test_tfs.py

@@ -50,21 +50,21 @@ def tfs_predictor(instance_type, sagemaker_session, tf_full_version):
         yield predictor
 
 
-def tar_dir(directory):
-
-    tmp = tempfile.mkdtemp()
+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(tmp, 'model.tar.gz'))
+    return sagemaker.utils.create_tar_file(source_files, os.path.join(str(tmpdir), 'model.tar.gz'))
 
 
 @pytest.fixture(scope='module')
 def tfs_predictor_with_model_and_entry_point_same_tar(instance_type,
                                                       sagemaker_session,
-                                                      tf_full_version):
+                                                      tf_full_version,
+                                                      tmpdir):
     endpoint_name = sagemaker.utils.unique_name_from_base('sagemaker-tensorflow-serving')
 
-    model_tar = tar_dir(os.path.join(tests.integ.DATA_DIR, 'tfs/tfs-test-model-with-inference'))
+    model_tar = tar_dir(os.path.join(tests.integ.DATA_DIR, 'tfs/tfs-test-model-with-inference'),
+                        tmpdir)
 
     model_data = sagemaker_session.upload_data(
         path=model_tar,