Fix source_directory handling

Author: Payton Staub

src/sagemaker/workflow/_repack_model.py

@@ -68,15 +68,18 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
         # copy source_dir to code/
         if source_dir:
-            actual_source_dir_path = os.path.join("/opt/ml/code", source_dir)
-            if os.path.exists(actual_source_dir_path):
-                for item in os.listdir(actual_source_dir_path):
-                    s = os.path.join(actual_source_dir_path, item)
-                    d = os.path.join(os.path.join(code_dir, source_dir), item)
-                    if os.path.isdir(s):
-                        shutil.copytree(s, d)
-                    else:
-                        shutil.copy2(s, d)
+            if os.path.exists(code_dir):
+                shutil.rmtree(code_dir)
+                shutil.copytree(source_dir, code_dir)
+            # actual_source_dir_path = os.path.join("/opt/ml/code", source_dir)
+            # if os.path.exists(actual_source_dir_path):
+            #     for item in os.listdir(actual_source_dir_path):
+            #         s = os.path.join(actual_source_dir_path, item)
+            #         d = os.path.join(os.path.join(code_dir, source_dir), item)
+            #         if os.path.isdir(s):
+            #             shutil.copytree(s, d)
+            #         else:
+            #             shutil.copy2(s, d)
 
         # copy any dependencies to code/lib/
         if dependencies:

tests/unit/sagemaker/workflow/test_repack_model_script.py

@@ -155,8 +155,8 @@ def test_repack_with_source_dir_and_dependencies(tmp):
     assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "a"))
     assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "bb"))
     assert os.path.exists(os.path.join("/opt/ml/model/code/lib/dir", "b"))
-    assert os.path.exists(os.path.join("/opt/ml/model/code/sourcedir", "foo.py"))
-    assert os.path.exists(os.path.join("/opt/ml/model/code/sourcedir/some/dir", "a"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/", "foo.py"))
+    assert os.path.exists(os.path.join("/opt/ml/model/code/some/dir", "a"))
 
 
 def create_file_tree(root, tree):