fix: make instance_type optional for model registry model package and mandatory for marketplace model package

Author: rahven14

src/sagemaker/session.py

@@ -4500,18 +4500,28 @@ def get_create_model_package_request(
             "SupportedContentTypes": content_types,
             "SupportedResponseMIMETypes": response_types,
         }
-        if inference_instances is not None:
-            inference_specification.update(
-                {
-                    "SupportedRealtimeInferenceInstanceTypes": inference_instances,
-                }
-            )
-        if transform_instances is not None:
-            inference_specification.update(
-                {
-                    "SupportedTransformInstanceTypes": transform_instances,
-                }
-            )
+        if model_package_group_name is not None:
+            if inference_instances is not None:
+                inference_specification.update(
+                    {
+                        "SupportedRealtimeInferenceInstanceTypes": inference_instances,
+                    }
+                )
+            if transform_instances is not None:
+                inference_specification.update(
+                    {
+                        "SupportedTransformInstanceTypes": transform_instances,
+                    }
+                )
+        else:
+            if not all([inference_instances, transform_instances]):
+                raise ValueError(
+                    "inference_instances and transform_instances " "must be provided if model_package_group_name is not present."
+                )
+            inference_specification.update({
+                "SupportedRealtimeInferenceInstanceTypes": inference_instances,
+                "SupportedTransformInstanceTypes": transform_instances,
+            })
         request_dict["InferenceSpecification"] = inference_specification
     request_dict["CertifyForMarketplace"] = marketplace_cert
     request_dict["ModelApprovalStatus"] = approval_status

tests/unit/sagemaker/workflow/test_pipeline_session.py

@@ -298,3 +298,33 @@ def test_pipeline_session_context_for_model_step_with_one_instance_types(
     }
 
     assert register_step_args.create_model_package_request == expected_output
+
+def test_pipeline_session_context_for_model_step_without_model_package_group_name(
+    pipeline_session_mock,
+):
+    model = Model(
+        name="MyModel",
+        image_uri="fakeimage",
+        model_data=ParameterString(name="ModelData", default_value="s3://my-bucket/file"),
+        sagemaker_session=pipeline_session_mock,
+        entry_point=f"{DATA_DIR}/dummy_script.py",
+        source_dir=f"{DATA_DIR}",
+        role=_ROLE,
+    )
+    with pytest.raises(ValueError) as error:
+        model.register(
+            content_types=["text/csv"],
+            response_types=["text/csv"],
+            inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
+            model_package_name="MyModelPackageName",
+            task="IMAGE_CLASSIFICATION",
+            sample_payload_url="s3://test-bucket/model",
+            framework="TENSORFLOW",
+            framework_version="2.9",
+            nearest_model_name="resnet50",
+            data_input_configuration='{"input_1":[1,224,224,3]}',
+        )
+        assert (
+            "inference_inferences and transform_instances "
+            "must be provided if model_package_group_name is not present." == str(error)
+        )

tests/unit/test_session.py

@@ -2355,10 +2355,27 @@ def test_create_model_package_from_containers_incomplete_args(sagemaker_session)
             containers=containers,
         )
         assert (
-            "content_types, response_types, inference_inferences and transform_instances "
+            "content_types and response_types "
             "must be provided if containers is present." == str(error)
         )
 
+def test_create_model_package_from_containers_without_model_package_group_name(sagemaker_session):
+    model_package_name = "sagemaker-model-package"
+    containers = ["dummy-container"]
+    content_types = ["application/json"]
+    response_types = ["application/json"]
+    with pytest.raises(ValueError) as error:
+        sagemaker_session.create_model_package_from_containers(
+            model_package_name=model_package_name,
+            containers=containers,
+            content_types=content_types,
+            response_types=response_types,
+        )
+        assert (
+            "inference_inferences and transform_instances "
+            "must be provided if model_package_group_name is not present." == str(error)
+        )
+
 
 def test_create_model_package_from_containers_all_args(sagemaker_session):
     model_package_name = "sagemaker-model-package"
@@ -2437,7 +2454,7 @@ def test_create_model_package_from_containers_all_args(sagemaker_session):
 
 
 def test_create_model_package_from_containers_without_instance_types(sagemaker_session):
-    model_package_name = "sagemaker-model-package"
+    model_package_group_name = "sagemaker-model-package-group-name-1.0"
     containers = ["dummy-container"]
     content_types = ["application/json"]
     response_types = ["application/json"]
@@ -2470,7 +2487,7 @@ def test_create_model_package_from_containers_without_instance_types(sagemaker_s
         containers=containers,
         content_types=content_types,
         response_types=response_types,
-        model_package_name=model_package_name,
+        model_package_group_name=model_package_group_name,
         model_metrics=model_metrics,
         metadata_properties=metadata_properties,
         marketplace_cert=marketplace_cert,
@@ -2480,7 +2497,7 @@ def test_create_model_package_from_containers_without_instance_types(sagemaker_s
         customer_metadata_properties=customer_metadata_properties,
     )
     expected_args = {
-        "ModelPackageName": model_package_name,
+        "ModelPackageGroupName": model_package_group_name,
         "InferenceSpecification": {
             "Containers": containers,
             "SupportedContentTypes": content_types,
@@ -2498,7 +2515,7 @@ def test_create_model_package_from_containers_without_instance_types(sagemaker_s
 
 
 def test_create_model_package_from_containers_with_one_instance_types(sagemaker_session):
-    model_package_name = "sagemaker-model-package"
+    model_package_group_name = "sagemaker-model-package-group-name-1.0"
     containers = ["dummy-container"]
     content_types = ["application/json"]
     response_types = ["application/json"]
@@ -2533,7 +2550,7 @@ def test_create_model_package_from_containers_with_one_instance_types(sagemaker_
         content_types=content_types,
         response_types=response_types,
         transform_instances=transform_instances,
-        model_package_name=model_package_name,
+        model_package_group_name=model_package_group_name,
         model_metrics=model_metrics,
         metadata_properties=metadata_properties,
         marketplace_cert=marketplace_cert,
@@ -2543,7 +2560,7 @@ def test_create_model_package_from_containers_with_one_instance_types(sagemaker_
         customer_metadata_properties=customer_metadata_properties,
     )
     expected_args = {
-        "ModelPackageName": model_package_name,
+        "ModelPackageGroupName": model_package_group_name,
         "InferenceSpecification": {
             "Containers": containers,
             "SupportedContentTypes": content_types,