fix: Fix key prefix preventing jumpstart model repack (#3943)

Author: evakravi

src/sagemaker/model.py

@@ -592,7 +592,11 @@ 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 or self.git_config)
+                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)
             deploy_env.update(self._script_mode_env_vars())

tests/unit/sagemaker/model/test_model.py

@@ -686,6 +686,25 @@ def test_script_mode_model_uses_jumpstart_base_name(repack_model, sagemaker_sess
     ].startswith(JUMPSTART_RESOURCE_BASE_NAME)
 
 
+@patch("sagemaker.utils.repack_model")
+def test_repack_code_location_with_key_prefix(repack_model, sagemaker_session):
+
+    code_location = "s3://my-bucket/code/location/"
+
+    t = 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,
+    )
+    t.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=INSTANCE_COUNT)
+
+    repack_model.assert_called_once()
+
+
 @patch("sagemaker.utils.repack_model")
 @patch("sagemaker.fw_utils.tar_and_upload_dir")
 def test_all_framework_models_add_jumpstart_base_name(