fix: make repack_model only removes py file when new entry_point provided (#1352)

* fix: make repack_model only removes py files when new entry_point present

Author: chuyang-deng

src/sagemaker/utils.py

@@ -530,7 +530,9 @@ def _create_or_update_code_dir(
     """
     code_dir = os.path.join(model_dir, "code")
     if os.path.exists(code_dir):
-        shutil.rmtree(code_dir, ignore_errors=True)
+        for filename in os.listdir(code_dir):
+            if filename.endswith(".py"):
+                os.remove(os.path.join(code_dir, filename))
     if source_directory and source_directory.lower().startswith("s3://"):
         local_code_path = os.path.join(tmp, "local_code.tar.gz")
         download_file_from_url(source_directory, local_code_path, sagemaker_session)
@@ -539,9 +541,12 @@ def _create_or_update_code_dir(
             t.extractall(path=code_dir)
 
     elif source_directory:
+        if os.path.exists(code_dir):
+            shutil.rmtree(code_dir)
         shutil.copytree(source_directory, code_dir)
     else:
-        os.mkdir(code_dir)
+        if not os.path.exists(code_dir):
+            os.mkdir(code_dir)
         shutil.copy2(inference_script, code_dir)
 
     for dependency in dependencies:

tests/unit/test_utils.py

@@ -614,6 +614,35 @@ def test_repack_model_from_file_to_folder(tmp):
     }
 
 
+def test_repack_model_with_inference_code_and_requirements(tmp, fake_s3):
+    create_file_tree(
+        tmp,
+        [
+            "new-inference.py",
+            "model-dir/model",
+            "model-dir/code/old-inference.py",
+            "model-dir/code/requirements.txt",
+        ],
+    )
+
+    fake_s3.tar_and_upload("model-dir", "s3://fake/location")
+
+    sagemaker.utils.repack_model(
+        os.path.join(tmp, "new-inference.py"),
+        None,
+        None,
+        "s3://fake/location",
+        "s3://destination-bucket/repacked-model",
+        fake_s3.sagemaker_session,
+    )
+
+    assert list_tar_files(fake_s3.fake_upload_path, tmp) == {
+        "/code/requirements.txt",
+        "/code/new-inference.py",
+        "/model",
+    }
+
+
 class FakeS3(object):
     def __init__(self, tmp):
         self.tmp = tmp