Repository: getmoto/moto
Pull request #7165 (2023-12-28 21:00:46)

Sagemaker model: support versioned models
The current implementation for SageMaker Model Registry has some issues regarding the usage of both `ModelPackageGroup` and `ModelPackage` when using the API [CreateModelPackage](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModelPackage.html).

As the documentation states:

> [ModelPackageGroupName](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModelPackage.html#API_CreateModelPackage_RequestSyntax)
**This parameter is required for versioned models, and does not apply to unversioned models.**

> [ModelPackageName](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateModelPackage.html#API_CreateModelPackage_RequestSyntax)
This parameter is required for unversioned models. It is not applicable to versioned models.

The current behavior accounts that the `ModelPackageName` is **required** for this API, even though this is not the case and this raises some issues for my team when creating model packages for a model package group.

This MR addresses this and updates tests too.

ps: I was non-orthodox for the class ModelPackage and made some pretty unconventional changes on it. If this is too much I can refactor the class a bit better.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`e5f9621`)](https://app.codecov.io/gh/getmoto/moto/commit/e5f962193cb12e696bd6ec641844c4d09b5fb87c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) 95.90% compared to head [(`c98e348`)](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) 95.90%.

| [Files](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Patch % | Lines |
|---|---|---|
| [moto/sagemaker/models.py](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#diff-bW90by9zYWdlbWFrZXIvbW9kZWxzLnB5) | 83.33% | [1 Missing :warning: ](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #7165      +/-   ##
==========================================
- Coverage   95.90%   95.90%   -0.01%     
==========================================
  Files         840      840              
  Lines       82438    82444       +6     
==========================================
+ Hits        79060    79064       +4     
- Misses       3378     3380       +2     
```

| [Flag](https://app.codecov.io/gh/getmoto/moto/pull/7165/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Coverage Δ | |
|---|---|---|
| [servertests](https://app.codecov.io/gh/getmoto/moto/pull/7165/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `35.48% <0.00%> (-0.01%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/getmoto/moto/pull/7165/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `95.84% <83.33%> (-0.01%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/getmoto/moto/pull/7165?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).


Patch:
diff --git a/moto/sagemaker/models.py b/moto/sagemaker/models.py
index d4c932d3e..896dd3b66 100644
--- a/moto/sagemaker/models.py
+++ b/moto/sagemaker/models.py
@@ -3,7 +3,7 @@ import os
 import random
 import string
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Union, cast
 
 from dateutil.tz import tzutc
 
@@ -1004,6 +1004,7 @@ class ModelPackage(BaseObject):
         client_token: str,
         region_name: str,
         account_id: str,
+        model_package_type: str,
         tags: Optional[List[Dict[str, str]]] = None,
     ) -> None:
         fake_user_profile_name = "fake-user-profile-name"
@@ -1032,6 +1033,7 @@ class ModelPackage(BaseObject):
         self.inference_specification = inference_specification
         self.source_algorithm_specification = source_algorithm_specification
         self.validation_specification = validation_specification
+        self.model_package_type = model_package_type
         self.model_package_status_details = {
             "ValidationStatuses": [
                 {
@@ -1105,6 +1107,10 @@ class ModelPackage(BaseObject):
             "AdditionalInferenceSpecifications",
             "SkipModelValidation",
         ]
+        if self.model_package_type == "Versioned":
+            del response_object["ModelPackageName"]
+        elif self.model_package_type == "Unversioned":
+            del response_object["ModelPackageGroupName"]
         return {
             k: v
             for k, v in response_object.items()
@@ -3295,6 +3301,8 @@ class SageMakerModelBackend(BaseBackend):
             )
         if model_package_group_name is not None:
             model_package_type = "Versioned"
+            if model_package_group_name.startswith("arn:aws"):
+                model_package_group_name = model_package_group_name.split("/")[-1]
         model_package_summary_list = list(
             filter(
                 lambda x: (
@@ -3377,7 +3385,7 @@ class SageMakerModelBackend(BaseBackend):
 
     def create_model_package(
         self,
-        model_package_name: str,
+        model_package_name: Optional[str],
         model_package_group_name: Optional[str],
         model_package_description: Optional[str],
         inference_specification: Any,
@@ -3397,15 +3405,32 @@ class SageMakerModelBackend(BaseBackend):
         additional_inference_specifications: Any,
     ) -> str:
         model_package_version = None
-        if model_package_group_name is not None:
+        if model_package_group_name and model_package_name:
+            raise AWSValidationException(
+                "An error occurred (ValidationException) when calling the CreateModelPackage operation: Both ModelPackageName and ModelPackageGroupName are provided in the input. Cannot determine which one to use."
+            )
+        elif not model_package_group_name and not model_package_name:
+            raise AWSValidationException(
+                "An error ocurred (ValidationException) when calling the CreateModelPackag operation: Missing ARN."
+            )
+        elif model_package_group_name:
+            model_package_type = "Versioned"
+            model_package_name = model_package_group_name
             model_packages_for_group = [
                 x
                 for x in self.model_packages.values()
                 if x.model_package_group_name == model_package_group_name
             ]
+            if model_package_group_name not in self.model_package_groups:
+                raise AWSValidationException(
+                    "An error ocurred (ValidationException) when calling the CreateModelPackage operation: Model Package Group does not exist."
+                )
             model_package_version = len(model_packages_for_group) + 1
+        else:
+            model_package_type = "Unversioned"
+
         model_package = ModelPackage(
-            model_package_name=model_package_name,
+            model_package_name=cast(str, model_package_name),
             model_package_group_name=model_package_group_name,
             model_package_description=model_package_description,
             inference_specification=inference_specification,
@@ -3427,6 +3452,7 @@ class SageMakerModelBackend(BaseBackend):
             region_name=self.region_name,
             account_id=self.account_id,
             client_token=client_token,
+            model_package_type=model_package_type,
         )
         self.model_package_name_mapping[
             model_package.model_package_name


Test patch:
diff --git a/tests/test_sagemaker/test_sagemaker_model_packages.py b/tests/test_sagemaker/test_sagemaker_model_packages.py
index 1a8f303b6..5d6b0d240 100644
--- a/tests/test_sagemaker/test_sagemaker_model_packages.py
+++ b/tests/test_sagemaker/test_sagemaker_model_packages.py
@@ -4,6 +4,7 @@ from unittest import SkipTest, TestCase
 
 import boto3
 import pytest
+from botocore.exceptions import ClientError
 from dateutil.tz import tzutc  # type: ignore
 from freezegun import freeze_time
 
@@ -133,19 +134,17 @@ def test_list_model_packages_approval_status():
 def test_list_model_packages_model_package_group_name():
     client = boto3.client("sagemaker", region_name="eu-west-1")
     group1 = "test-model-package-group"
+    client.create_model_package_group(
+        ModelPackageGroupName=group1,
+        ModelPackageGroupDescription="test-model-package-description",
+    )
     client.create_model_package(
-        ModelPackageName="test-model-package",
-        ModelPackageDescription="test-model-package-description",
         ModelPackageGroupName=group1,
     )
     client.create_model_package(
-        ModelPackageName="test-model-package",
-        ModelPackageDescription="test-model-package-description-2",
         ModelPackageGroupName=group1,
     )
     client.create_model_package(
-        ModelPackageName="test-model-package-2",
-        ModelPackageDescription="test-model-package-description-3",
         ModelPackageGroupName=group1,
     )
     client.create_model_package(
@@ -170,8 +169,8 @@ def test_list_model_packages_model_package_group_name():
 @mock_sagemaker
 def test_list_model_packages_model_package_type():
     client = boto3.client("sagemaker", region_name="eu-west-1")
+    client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     client.create_model_package(
-        ModelPackageName="test-model-package",
         ModelPackageDescription="test-model-package-description",
         ModelPackageGroupName="test-model-package-group",
     )
@@ -233,18 +232,16 @@ def test_describe_model_package_default():
     client = boto3.client("sagemaker", region_name="eu-west-1")
     client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     with freeze_time("2015-01-01 00:00:00"):
-        client.create_model_package(
-            ModelPackageName="test-model-package",
+        model_package_arn = client.create_model_package(
             ModelPackageGroupName="test-model-package-group",
             ModelPackageDescription="test-model-package-description",
-        )
-    resp = client.describe_model_package(ModelPackageName="test-model-package")
-    assert resp["ModelPackageName"] == "test-model-package"
+        )["ModelPackageArn"]
+    resp = client.describe_model_package(ModelPackageName=model_package_arn)
     assert resp["ModelPackageGroupName"] == "test-model-package-group"
     assert resp["ModelPackageDescription"] == "test-model-package-description"
     assert (
         resp["ModelPackageArn"]
-        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package/1"
+        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package-group/1"
     )
     assert resp["CreationTime"] == datetime(2015, 1, 1, 0, 0, 0, tzinfo=tzutc())
     assert (
@@ -257,13 +254,13 @@ def test_describe_model_package_default():
     assert resp.get("ModelPackageStatusDetails") is not None
     assert resp["ModelPackageStatusDetails"]["ValidationStatuses"] == [
         {
-            "Name": "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package/1",
+            "Name": "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package-group/1",
             "Status": "Completed",
         }
     ]
     assert resp["ModelPackageStatusDetails"]["ImageScanStatuses"] == [
         {
-            "Name": "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package/1",
+            "Name": "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package-group/1",
             "Status": "Completed",
         }
     ]
@@ -273,10 +270,8 @@ def test_describe_model_package_default():
 @mock_sagemaker
 def test_describe_model_package_with_create_model_package_arguments():
     client = boto3.client("sagemaker", region_name="eu-west-1")
-    client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     client.create_model_package(
         ModelPackageName="test-model-package",
-        ModelPackageGroupName="test-model-package-group",
         ModelPackageDescription="test-model-package-description",
         ModelApprovalStatus="PendingManualApproval",
         MetadataProperties={
@@ -303,7 +298,6 @@ def test_update_model_package():
     client = boto3.client("sagemaker", region_name="eu-west-1")
     client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     model_package_arn = client.create_model_package(
-        ModelPackageName="test-model-package",
         ModelPackageGroupName="test-model-package-group",
         ModelPackageDescription="test-model-package-description",
         CustomerMetadataProperties={
@@ -317,7 +311,7 @@ def test_update_model_package():
         CustomerMetadataProperties={"test-key": "test-value"},
         CustomerMetadataPropertiesToRemove=["test-key-to-remove"],
     )
-    resp = client.describe_model_package(ModelPackageName="test-model-package")
+    resp = client.describe_model_package(ModelPackageName=model_package_arn)
     assert resp["ModelApprovalStatus"] == "Approved"
     assert resp["ApprovalDescription"] == "test-approval-description"
     assert resp["CustomerMetadataProperties"] is not None
@@ -330,7 +324,6 @@ def test_update_model_package_given_additional_inference_specifications_to_add()
     client = boto3.client("sagemaker", region_name="eu-west-1")
     client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     model_package_arn = client.create_model_package(
-        ModelPackageName="test-model-package",
         ModelPackageGroupName="test-model-package-group",
         ModelPackageDescription="test-model-package-description",
     )["ModelPackageArn"]
@@ -366,7 +359,7 @@ def test_update_model_package_given_additional_inference_specifications_to_add()
             additional_inference_specifications_to_add
         ],
     )
-    resp = client.describe_model_package(ModelPackageName="test-model-package")
+    resp = client.describe_model_package(ModelPackageName=model_package_arn)
     assert resp["AdditionalInferenceSpecifications"] is not None
     TestCase().assertDictEqual(
         resp["AdditionalInferenceSpecifications"][0],
@@ -375,13 +368,12 @@ def test_update_model_package_given_additional_inference_specifications_to_add()
 
 
 @mock_sagemaker
-def test_update_model_package_shoudl_update_last_modified_information():
+def test_update_model_package_should_update_last_modified_information():
     if settings.TEST_SERVER_MODE:
         raise SkipTest("Can't freeze time in ServerMode")
     client = boto3.client("sagemaker", region_name="eu-west-1")
     client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     model_package_arn = client.create_model_package(
-        ModelPackageName="test-model-package",
         ModelPackageGroupName="test-model-package-group",
     )["ModelPackageArn"]
     with freeze_time("2020-01-01 12:00:00"):
@@ -389,7 +381,7 @@ def test_update_model_package_shoudl_update_last_modified_information():
             ModelPackageArn=model_package_arn,
             ModelApprovalStatus="Approved",
         )
-    resp = client.describe_model_package(ModelPackageName="test-model-package")
+    resp = client.describe_model_package(ModelPackageName=model_package_arn)
     assert resp.get("LastModifiedTime") is not None
     assert resp["LastModifiedTime"] == datetime(2020, 1, 1, 12, 0, 0, tzinfo=tzutc())
     assert resp.get("LastModifiedBy") is not None
@@ -430,6 +422,42 @@ def test_create_model_package():
     )
 
 
+@mock_sagemaker
+def test_create_model_package_should_raise_error_when_model_package_group_name_and_model_package_group_name_are_not_provided():
+    client = boto3.client("sagemaker", region_name="eu-west-1")
+    with pytest.raises(ClientError) as exc:
+        _ = client.create_model_package(
+            ModelPackageDescription="test-model-package-description",
+        )
+    assert "Missing ARN." in str(exc.value)
+
+
+@mock_sagemaker
+def test_create_model_package_should_raise_error_when_package_name_and_group_are_provided():
+    client = boto3.client("sagemaker", region_name="eu-west-1")
+    with pytest.raises(ClientError) as exc:
+        _ = client.create_model_package(
+            ModelPackageGroupName="TestModelPackageGroup",
+            ModelPackageName="TestModelPackage",
+            ModelPackageDescription="test-model-package-description",
+        )
+    assert (
+        "Both ModelPackageName and ModelPackageGroupName are provided in the input."
+        in str(exc.value)
+    )
+
+
+@mock_sagemaker
+def test_create_model_package_should_raise_error_when_model_package_group_provided_not_exist():
+    client = boto3.client("sagemaker", region_name="eu-west-1")
+    with pytest.raises(ClientError) as exc:
+        _ = client.create_model_package(
+            ModelPackageGroupName="TestModelPackageGroupNotExist",
+            ModelPackageDescription="test-model-package-description",
+        )
+    assert "Model Package Group does not exist." in str(exc.value)
+
+
 @pytest.mark.parametrize(
     "model_approval_status",
     ["Approved", "Rejected", "PendingManualApproval"],
@@ -456,20 +484,18 @@ def test_create_model_package_in_model_package_group():
     client = boto3.client("sagemaker", region_name="eu-west-1")
     client.create_model_package_group(ModelPackageGroupName="test-model-package-group")
     resp_version_1 = client.create_model_package(
-        ModelPackageName="TestModelPackage",
         ModelPackageGroupName="test-model-package-group",
         ModelPackageDescription="test-model-package-description",
     )
     resp_version_2 = client.create_model_package(
-        ModelPackageName="TestModelPackage",
         ModelPackageGroupName="test-model-package-group",
         ModelPackageDescription="test-model-package-description",
     )
     assert (
         resp_version_1["ModelPackageArn"]
-        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/testmodelpackage/1"
+        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package-group/1"
     )
     assert (
         resp_version_2["ModelPackageArn"]
-        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/testmodelpackage/2"
+        == "arn:aws:sagemaker:eu-west-1:123456789012:model-package/test-model-package-group/2"
     )