change: add method to Model class to check if repack is needed (#4213)

Author: trungleduc

src/sagemaker/model.py

@@ -623,14 +623,7 @@ def prepare_container_def(
         )
         deploy_env = copy.deepcopy(self.env)
         if self.source_dir or self.dependencies or self.entry_point or self.git_config:
-            is_repack = (
-                self.source_dir
-                and self.entry_point
-                and not (
-                    (self.key_prefix and issubclass(type(self), FrameworkModel)) or self.git_config
-                )
-            )
-            self._upload_code(deploy_key_prefix, repack=is_repack)
+            self._upload_code(deploy_key_prefix, repack=self.is_repack())
             deploy_env.update(self._script_mode_env_vars())
 
         return sagemaker.container_def(
@@ -640,6 +633,14 @@ def prepare_container_def(
             image_config=self.image_config,
         )
 
+    def is_repack(self) -> bool:
+        """Whether the source code needs to be repacked before uploading to S3.
+
+        Returns:
+            bool: if the source need to be repacked or not
+        """
+        return self.source_dir and self.entry_point and not self.git_config
+
     def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
         """Uploads code to S3 to be used with script mode with SageMaker inference.
 
@@ -1775,6 +1776,14 @@ def __init__(
             **kwargs,
         )
 
+    def is_repack(self) -> bool:
+        """Whether the source code needs to be repacked before uploading to S3.
+
+        Returns:
+            bool: if the source need to be repacked or not
+        """
+        return self.source_dir and self.entry_point and not (self.key_prefix or self.git_config)
+
 
 # works for MODEL_PACKAGE_ARN with or without version info.
 MODEL_PACKAGE_ARN_PATTERN = r"arn:aws:sagemaker:(.*?):(.*?):model-package/(.*?)(?:/(\d+))?$"

tests/unit/sagemaker/model/test_framework_model.py

@@ -48,6 +48,7 @@
 CODECOMMIT_REPO_SSH = "ssh://git-codecommit.us-west-2.amazonaws.com/v1/repos/test-repo/"
 CODECOMMIT_BRANCH = "master"
 REPO_DIR = "/tmp/repo_dir"
+IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
 
 
 class DummyFrameworkModel(FrameworkModel):
@@ -471,3 +472,66 @@ def test_git_support_codecommit_ssh_passphrase_required(
         )
         model.prepare_container_def(instance_type=INSTANCE_TYPE)
     assert "returned non-zero exit status" in str(error.value)
+
+
+@patch("sagemaker.utils.repack_model")
+def test_not_repack_code_location_with_key_prefix(repack_model, sagemaker_session):
+
+    code_location = "s3://my-bucket/code/location/"
+
+    t = FrameworkModel(
+        entry_point=ENTRY_POINT,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir="s3://codebucket/someprefix/sourcedir.tar.gz",
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+        code_location=code_location,
+    )
+    t.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=INSTANCE_COUNT)
+
+    repack_model.assert_not_called()
+
+
+@patch("sagemaker.utils.repack_model")
+def test_is_repack_with_code_location(repack_model, sagemaker_session):
+
+    code_location = "s3://my-bucket/code/location/"
+
+    model = FrameworkModel(
+        entry_point=ENTRY_POINT,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir="s3://codebucket/someprefix/sourcedir.tar.gz",
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+        code_location=code_location,
+    )
+
+    assert not model.is_repack()
+
+
+@patch("sagemaker.git_utils.git_clone_repo")
+@patch("sagemaker.model.fw_utils.tar_and_upload_dir")
+def test_is_repack_with_git_config(tar_and_upload_dir, git_clone_repo, sagemaker_session):
+    git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
+        "entry_point": "entry_point",
+        "source_dir": "/tmp/repo_dir/source_dir",
+        "dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
+    }
+
+    entry_point = "entry_point"
+    source_dir = "source_dir"
+    dependencies = ["foo", "bar"]
+    git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+    model = FrameworkModel(
+        sagemaker_session=sagemaker_session,
+        entry_point=entry_point,
+        source_dir=source_dir,
+        dependencies=dependencies,
+        git_config=git_config,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+    )
+
+    assert not model.is_repack()

tests/unit/sagemaker/model/test_model.py

@@ -704,6 +704,49 @@ def test_repack_code_location_with_key_prefix(repack_model, sagemaker_session):
     repack_model.assert_called_once()
 
 
+@patch("sagemaker.utils.repack_model")
+def test_is_repack_with_code_location(repack_model, sagemaker_session):
+
+    code_location = "s3://my-bucket/code/location/"
+
+    model = Model(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=SCRIPT_URI,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+        code_location=code_location,
+    )
+
+    assert model.is_repack()
+
+
+@patch("sagemaker.git_utils.git_clone_repo")
+@patch("sagemaker.model.fw_utils.tar_and_upload_dir")
+def test_is_repack_with_git_config(tar_and_upload_dir, git_clone_repo, sagemaker_session):
+    git_clone_repo.side_effect = lambda gitconfig, entrypoint, sourcedir, dependency: {
+        "entry_point": "entry_point",
+        "source_dir": "/tmp/repo_dir/source_dir",
+        "dependencies": ["/tmp/repo_dir/foo", "/tmp/repo_dir/bar"],
+    }
+
+    entry_point = "entry_point"
+    source_dir = "source_dir"
+    dependencies = ["foo", "bar"]
+    git_config = {"repo": GIT_REPO, "branch": BRANCH, "commit": COMMIT}
+    model = Model(
+        sagemaker_session=sagemaker_session,
+        entry_point=entry_point,
+        source_dir=source_dir,
+        dependencies=dependencies,
+        git_config=git_config,
+        image_uri=IMAGE_URI,
+    )
+
+    assert not model.is_repack()
+
+
 @patch("sagemaker.utils.repack_model")
 @patch("sagemaker.fw_utils.tar_and_upload_dir")
 def test_all_framework_models_add_jumpstart_base_name(