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

Author: Keshav Chandak

src/sagemaker/session.py

@@ -4347,7 +4347,29 @@ def submit(request):
             if model_package_group_name is not None and not model_package_group_name.startswith(
                 "arn:"
             ):
-                _create_resource(
+                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"]
                     )

tests/unit/test_chainer.py

@@ -50,7 +50,6 @@
 
 LIST_TAGS_RESULT = {"Tags": [{"Key": "TagtestKey", "Value": "TagtestValue"}]}
 
-
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)

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,
@@ -5093,6 +5096,10 @@ def test_create_model_package_from_containers_with_source_uri_and_inference_spec
     approval_status = ("Approved",)
     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,
@@ -5220,6 +5230,10 @@ def test_create_model_package_from_containers_with_source_uri_set_to_mp(sagemake
     sagemaker_session.sagemaker_client.create_model_package = Mock(
         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,
@@ -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,21 @@ 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")
+    
+