fix: breaking deviations in _create_sagemaker_model call (#3919)

Author: mufaddal-rohawala

src/sagemaker/model.py

@@ -17,6 +17,7 @@
 import json
 import logging
 import os
+import re
 import copy
 from typing import List, Dict, Optional, Union
 
@@ -1662,6 +1663,10 @@ def __init__(
         )
 
 
+# works for MODEL_PACKAGE_ARN with or without version info.
+MODEL_PACKAGE_ARN_PATTERN = r"arn:aws:sagemaker:(.*?):(.*?):model-package/(.*?)(?:/(\d+))?$"
+
+
 class ModelPackage(Model):
     """A SageMaker ``Model`` that can be deployed to an ``Endpoint``."""
 
@@ -1769,14 +1774,19 @@ def _create_sagemaker_model(self, *args, **kwargs):  # pylint: disable=unused-ar
             model_package_name = self._created_model_package_name
         else:
             # When a ModelPackageArn is provided we just create the Model
-            model_package_name = self.model_package_arn
+            match = re.match(MODEL_PACKAGE_ARN_PATTERN, self.model_package_arn)
+            if match:
+                model_package_name = match.group(3)
+            else:
+                # model_package_arn can be just the name if your account owns the Model Package
+                model_package_name = self.model_package_arn
 
-        container_def = {"ModelPackageName": model_package_name}
+        container_def = {"ModelPackageName": self.model_package_arn}
 
         if self.env != {}:
             container_def["Environment"] = self.env
 
-        self._ensure_base_name_if_needed(model_package_name.split("/")[-1])
+        self._ensure_base_name_if_needed(model_package_name)
         self._set_model_name_if_needed()
 
         self.sagemaker_session.create_model(
@@ -1785,6 +1795,7 @@ def _create_sagemaker_model(self, *args, **kwargs):  # pylint: disable=unused-ar
             container_def,
             vpc_config=self.vpc_config,
             enable_network_isolation=self.enable_network_isolation(),
+            tags=kwargs.get("tags"),
         )
 
     def _ensure_base_name_if_needed(self, base_name):

tests/unit/sagemaker/model/test_model_package.py

@@ -115,6 +115,42 @@ def test_create_sagemaker_model_uses_model_name(name_from_base, sagemaker_sessio
         {"ModelPackageName": model_package_name},
         vpc_config=None,
         enable_network_isolation=False,
+        tags=None,
+    )
+
+
+@pytest.mark.parametrize(
+    "model_package_arn",
+    [
+        "arn:aws:sagemaker:us-east-2:123:model-package/my-model-package-arn",
+        "arn:aws:sagemaker:us-east-2:123:model-package/my-model-package-arn/12",
+    ],
+)
+@patch("sagemaker.utils.name_from_base")
+def test_create_sagemaker_model_uses_model_package_arn(
+    name_from_base, sagemaker_session, model_package_arn
+):
+    model_name = "my-model"
+
+    model_package = ModelPackage(
+        role="role",
+        name=model_name,
+        model_package_arn=model_package_arn,
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_package._create_sagemaker_model()
+
+    assert model_name == model_package.name
+    name_from_base.assert_not_called()
+
+    sagemaker_session.create_model.assert_called_with(
+        model_name,
+        "role",
+        {"ModelPackageName": model_package_arn},
+        vpc_config=None,
+        enable_network_isolation=False,
+        tags=None,
     )
 
 
@@ -141,6 +177,35 @@ def test_create_sagemaker_model_include_environment_variable(sagemaker_session):
         {"ModelPackageName": model_package_name, "Environment": environment},
         vpc_config=None,
         enable_network_isolation=False,
+        tags=None,
+    )
+
+
+def test_create_sagemaker_model_include_tags(sagemaker_session):
+    model_name = "my-model"
+    model_package_name = "my-model-package"
+    env_key = "env_key"
+    env_value = "env_value"
+    environment = {env_key: env_value}
+    tags = {"Key": "foo", "Value": "bar"}
+
+    model_package = ModelPackage(
+        role="role",
+        name=model_name,
+        model_package_arn=model_package_name,
+        env=environment,
+        sagemaker_session=sagemaker_session,
+    )
+
+    model_package._create_sagemaker_model(tags=tags)
+
+    sagemaker_session.create_model.assert_called_with(
+        model_name,
+        "role",
+        {"ModelPackageName": model_package_name, "Environment": environment},
+        vpc_config=None,
+        enable_network_isolation=False,
+        tags=tags,
     )