bugix: Added check for the presence of model package group before creating one

Author: Keshav Chandak

src/sagemaker/session.py

@@ -4347,11 +4347,39 @@ def submit(request):
             if model_package_group_name is not None and not model_package_group_name.startswith(
                 "arn:"
             ):
-                _create_resource(
-                    lambda: self.sagemaker_client.create_model_package_group(
-                        ModelPackageGroupName=request["ModelPackageGroupName"]
+                model_package_groups = []
+                model_package_groups_response = self.sagemaker_client.list_model_package_groups(
+                    NameContains=request["ModelPackageGroupName"],
+                )
+                model_package_groups = (
+                    model_package_groups
+                    + model_package_groups_response["ModelPackageGroupSummaryList"]
+                )
+                next_token = model_package_groups_response.get("NextToken")
+
+                while next_token is not None and next_token != "":
+                    model_package_groups_response = self.sagemaker_client.list_model_package_groups(
+                        NameContains=request["ModelPackageGroupName"], NextToken=next_token
+                    )
+                    model_package_groups = (
+                        model_package_groups
+                        + model_package_groups_response["ModelPackageGroupSummaryList"]
+                    )
+                    next_token = model_package_groups_response.get("NextToken")
+
+                filtered_model_package_group = list(
+                    filter(
+                        lambda mpg: mpg.get("ModelPackageGroupName")
+                        == request["ModelPackageGroupName"],
+                        model_package_groups,
                     )
                 )
+                if not filtered_model_package_group:
+                    _create_resource(
+                        lambda: self.sagemaker_client.create_model_package_group(
+                            ModelPackageGroupName=request["ModelPackageGroupName"]
+                        )
+                    )
             if "SourceUri" in request and request["SourceUri"] is not None:
                 # Remove inference spec from request if the
                 # given source uri can lead to auto-population of it

tests/unit/test_session.py

@@ -5006,6 +5006,9 @@ def test_create_model_package_with_sagemaker_config_injection(sagemaker_session)
     domain = "COMPUTER_VISION"
     task = "IMAGE_CLASSIFICATION"
     sample_payload_url = "s3://test-bucket/model"
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
     sagemaker_session.create_model_package_from_containers(
         containers=containers,
         content_types=content_types,
@@ -5094,6 +5097,10 @@ def test_create_model_package_from_containers_with_source_uri_and_inference_spec
     skip_model_validation = "All"
     source_uri = "dummy-source-uri"
 
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
+
     created_versioned_mp_arn = (
         "arn:aws:sagemaker:us-west-2:123456789123:model-package/unit-test-package-version/1"
     )
@@ -5149,6 +5156,9 @@ def test_create_model_package_from_containers_with_source_uri_for_unversioned_mp
     approval_status = ("Approved",)
     skip_model_validation = "All"
     source_uri = "dummy-source-uri"
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
 
     with pytest.raises(
         ValueError,
@@ -5221,6 +5231,10 @@ def test_create_model_package_from_containers_with_source_uri_set_to_mp(sagemake
         return_value={"ModelPackageArn": created_versioned_mp_arn}
     )
 
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
+
     sagemaker_session.create_model_package_from_containers(
         model_package_group_name=model_package_group_name,
         containers=containers,
@@ -5443,6 +5457,9 @@ def test_create_model_package_from_containers_without_instance_types(sagemaker_s
     approval_status = ("Approved",)
     description = "description"
     customer_metadata_properties = {"key1": "value1"}
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
     sagemaker_session.create_model_package_from_containers(
         containers=containers,
         content_types=content_types,
@@ -5510,6 +5527,9 @@ def test_create_model_package_from_containers_with_one_instance_types(
     approval_status = ("Approved",)
     description = "description"
     customer_metadata_properties = {"key1": "value1"}
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
     sagemaker_session.create_model_package_from_containers(
         containers=containers,
         content_types=content_types,
@@ -7183,3 +7203,27 @@ def test_delete_hub_content_reference(sagemaker_session):
     }
 
     sagemaker_session.sagemaker_client.delete_hub_content_reference.assert_called_with(**request)
+
+
+def test_create_model_package_from_containers_to_create_mpg_if_not_present(sagemaker_session):
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": [{"ModelPackageGroupName": "mock-mpg"}]
+    }
+    sagemaker_session.create_model_package_from_containers(
+        source_uri="mock-source-uri", model_package_group_name="mock-mpg"
+    )
+    sagemaker_session.sagemaker_client.create_model_package_group.assert_not_called()
+    sagemaker_session.create_model_package_from_containers(
+        source_uri="mock-source-uri",
+        model_package_group_name="arn:aws:sagemaker:us-east-1:215995503607:model-package-group/mock-mpg",
+    )
+    sagemaker_session.sagemaker_client.create_model_package_group.assert_not_called()
+    sagemaker_session.sagemaker_client.list_model_package_groups.return_value = {
+        "ModelPackageGroupSummaryList": []
+    }
+    sagemaker_session.create_model_package_from_containers(
+        source_uri="mock-source-uri", model_package_group_name="mock-mpg"
+    )
+    sagemaker_session.sagemaker_client.create_model_package_group.assert_called_with(
+        ModelPackageGroupName="mock-mpg"
+    )