From c02edc917b8aec48c549f1afef2563c3bbde20f4 Mon Sep 17 00:00:00 2001 From: pintaoz Date: Tue, 25 Feb 2025 14:27:21 -0800 Subject: [PATCH 1/2] Ensure Model.is_repack() returns a boolean --- src/sagemaker/model.py | 4 ++++ .../unit/sagemaker/model/test_framework_model.py | 15 +++++++++++++++ tests/unit/sagemaker/model/test_model.py | 15 +++++++++++++++ 3 files changed, 34 insertions(+) diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 5cc260f3ef..e5ea1ea314 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -745,6 +745,8 @@ def is_repack(self) -> bool: Returns: bool: if the source need to be repacked or not """ + if self.source_dir is None or self.entry_point is None: + return False return self.source_dir and self.entry_point and not self.git_config def _upload_code(self, key_prefix: str, repack: bool = False) -> None: @@ -2143,6 +2145,8 @@ def is_repack(self) -> bool: Returns: bool: if the source need to be repacked or not """ + if self.source_dir is None or self.entry_point is None: + return False return self.source_dir and self.entry_point and not (self.key_prefix or self.git_config) diff --git a/tests/unit/sagemaker/model/test_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py index d41dd6f821..d52da4c8ad 100644 --- a/tests/unit/sagemaker/model/test_framework_model.py +++ b/tests/unit/sagemaker/model/test_framework_model.py @@ -511,6 +511,21 @@ def test_is_repack_with_code_location(repack_model, sagemaker_session): assert not model.is_repack() +@patch("sagemaker.utils.repack_model") +def test_is_repack_with_none_type(repack_model, sagemaker_session): + """Test is_repack() returns a boolean value when source_dir and entry_point are None""" + + model = FrameworkModel( + role=ROLE, + sagemaker_session=sagemaker_session, + source_dir="s3://codebucket/someprefix/sourcedir.tar.gz", + image_uri=IMAGE_URI, + model_data=MODEL_DATA, + ) + + assert model.is_repack() is False + + @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): diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py index 9175613662..cecf08e71b 100644 --- a/tests/unit/sagemaker/model/test_model.py +++ b/tests/unit/sagemaker/model/test_model.py @@ -1046,6 +1046,21 @@ def test_is_repack_with_code_location(repack_model, sagemaker_session): assert model.is_repack() +@patch("sagemaker.utils.repack_model") +def test_is_repack_with_none_type(repack_model, sagemaker_session): + """Test is_repack() returns a boolean value when source_dir and entry_point are None""" + + model = Model( + role=ROLE, + sagemaker_session=sagemaker_session, + source_dir=SCRIPT_URI, + image_uri=IMAGE_URI, + model_data=MODEL_DATA, + ) + + assert model.is_repack() is False + + @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): From 377cf1557c5a554d0307533cce09eaf41e03781d Mon Sep 17 00:00:00 2001 From: pintaoz Date: Wed, 26 Feb 2025 16:10:48 -0800 Subject: [PATCH 2/2] update test --- tests/unit/sagemaker/model/test_framework_model.py | 1 - tests/unit/sagemaker/model/test_model.py | 1 - 2 files changed, 2 deletions(-) diff --git a/tests/unit/sagemaker/model/test_framework_model.py b/tests/unit/sagemaker/model/test_framework_model.py index d52da4c8ad..432d90bd37 100644 --- a/tests/unit/sagemaker/model/test_framework_model.py +++ b/tests/unit/sagemaker/model/test_framework_model.py @@ -518,7 +518,6 @@ def test_is_repack_with_none_type(repack_model, sagemaker_session): model = FrameworkModel( role=ROLE, sagemaker_session=sagemaker_session, - source_dir="s3://codebucket/someprefix/sourcedir.tar.gz", image_uri=IMAGE_URI, model_data=MODEL_DATA, ) diff --git a/tests/unit/sagemaker/model/test_model.py b/tests/unit/sagemaker/model/test_model.py index cecf08e71b..3d498dfc59 100644 --- a/tests/unit/sagemaker/model/test_model.py +++ b/tests/unit/sagemaker/model/test_model.py @@ -1053,7 +1053,6 @@ def test_is_repack_with_none_type(repack_model, sagemaker_session): model = Model( role=ROLE, sagemaker_session=sagemaker_session, - source_dir=SCRIPT_URI, image_uri=IMAGE_URI, model_data=MODEL_DATA, )