feature: Adding serial inference pipeline support to RegisterModel Step

Author: sreedes

src/sagemaker/model.py

@@ -195,6 +195,7 @@ def _get_model_package_args(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        container_def_list=None,
     ):
         """Get arguments for session.create_model_package method.
 
@@ -219,6 +220,7 @@ def _get_model_package_args(
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval" (default: "PendingManualApproval").
             description (str): Model Package description (default: None).
+            container_def_list (list): A list of container defintiions.
         Returns:
             dict: A dictionary of method argument names and values.
         """
@@ -229,8 +231,13 @@ def _get_model_package_args(
             "ModelDataUrl": self.model_data,
         }
 
+        if container_def_list is not None:
+            containers = container_def_list
+        else:
+            containers = [container]
+
         model_package_args = {
-            "containers": [container],
+            "containers": containers,
             "content_types": content_types,
             "response_types": response_types,
             "inference_instances": inference_instances,

src/sagemaker/workflow/_utils.py

@@ -205,18 +205,19 @@ class _RegisterModelStep(Step):
             Estimator's training container image will be used (default: None).
         compile_model_family (str): Instance family for compiled model, if specified, a compiled
             model will be used (default: None).
+        container_def_list (list): A list of container defintiions.
         **kwargs: additional arguments to `create_model`.
     """
 
     def __init__(
         self,
         name: str,
         estimator: EstimatorBase,
-        model_data,
         content_types,
         response_types,
         inference_instances,
         transform_instances,
+        model_data=None,
         model_package_group_name=None,
         model_metrics=None,
         metadata_properties=None,
@@ -225,6 +226,7 @@ def __init__(
         compile_model_family=None,
         description=None,
         depends_on: List[str] = None,
+        container_def_list=None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -254,6 +256,7 @@ def __init__(
             description (str): Model Package description (default: None).
             depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TrainingStep`
                 depends on
+            container_def_list (list): A list of container defintiions.
             **kwargs: additional arguments to `create_model`.
         """
         super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL, depends_on)
@@ -270,6 +273,7 @@ def __init__(
         self.image_uri = image_uri
         self.compile_model_family = compile_model_family
         self.description = description
+        self.container_def_list = container_def_list
         self.kwargs = kwargs
 
         self._properties = Properties(
@@ -301,7 +305,7 @@ def arguments(self) -> RequestType:
             self.estimator.output_path = output_path
 
             # yeah, there is some framework stuff going on that we need to pull in here
-            if model.image_uri is None:
+            if model.image_uri is None and self.container_def_list is None:
                 region_name = self.estimator.sagemaker_session.boto_session.region_name
                 model.image_uri = image_uris.retrieve(
                     model._framework_name,
@@ -324,6 +328,7 @@ def arguments(self) -> RequestType:
             metadata_properties=self.metadata_properties,
             approval_status=self.approval_status,
             description=self.description,
+            container_def_list=self.container_def_list,
         )
         request_dict = model.sagemaker_session._get_create_model_package_request(
             **model_package_args

src/sagemaker/workflow/step_collections.py

@@ -55,18 +55,19 @@ def __init__(
         self,
         name: str,
         estimator: EstimatorBase,
-        model_data,
         content_types,
         response_types,
         inference_instances,
         transform_instances,
+        model_data=None,
         depends_on: List[str] = None,
         model_package_group_name=None,
         model_metrics=None,
         approval_status=None,
         image_uri=None,
         compile_model_family=None,
         description=None,
+        container_def_list=None,
         **kwargs,
     ):
         """Construct steps `_RepackModelStep` and `_RegisterModelStep` based on the estimator.
@@ -94,6 +95,7 @@ def __init__(
             compile_model_family (str): The instance family for the compiled model. If
                 specified, a compiled model is used (default: None).
             description (str): Model Package description (default: None).
+            container_def_list (list): A list of container defintiions.
             **kwargs: additional arguments to `create_model`.
         """
         steps: List[Step] = []
@@ -134,6 +136,7 @@ def __init__(
             image_uri=image_uri,
             compile_model_family=compile_model_family,
             description=description,
+            container_def_list=container_def_list,
             **kwargs,
         )
         if not repack_model:

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -266,6 +266,76 @@ def test_register_model_tf(estimator_tf, model_metrics):
     )
 
 
+def test_register_model_sip(estimator, model_metrics):
+    container_def_list = [
+        {
+            "Image": "fakeimage1",
+            "ModelDataUrl": "Url1",
+            "Environment": [{"k1": "v1"}, {"k2": "v2"}],
+        },
+        {
+            "Image": "fakeimage2",
+            "ModelDataUrl": "Url2",
+            "Environment": [{"k3": "v3"}, {"k4": "v4"}],
+        },
+    ]
+
+    register_model = RegisterModel(
+        name="RegisterModelStep",
+        estimator=estimator,
+        content_types=["content_type"],
+        response_types=["response_type"],
+        inference_instances=["inference_instance"],
+        transform_instances=["transform_instance"],
+        model_package_group_name="mpg",
+        model_metrics=model_metrics,
+        approval_status="Approved",
+        description="description",
+        container_def_list=container_def_list,
+        depends_on=["TestStep"],
+    )
+    assert ordered(register_model.request_dicts()) == ordered(
+        [
+            {
+                "Name": "RegisterModelStep",
+                "Type": "RegisterModel",
+                "DependsOn": ["TestStep"],
+                "Arguments": {
+                    "InferenceSpecification": {
+                        "Containers": [
+                            {
+                                "Image": "fakeimage1",
+                                "ModelDataUrl": "Url1",
+                                "Environment": [{"k1": "v1"}, {"k2": "v2"}],
+                            },
+                            {
+                                "Image": "fakeimage2",
+                                "ModelDataUrl": "Url2",
+                                "Environment": [{"k3": "v3"}, {"k4": "v4"}],
+                            },
+                        ],
+                        "SupportedContentTypes": ["content_type"],
+                        "SupportedRealtimeInferenceInstanceTypes": ["inference_instance"],
+                        "SupportedResponseMIMETypes": ["response_type"],
+                        "SupportedTransformInstanceTypes": ["transform_instance"],
+                    },
+                    "ModelApprovalStatus": "Approved",
+                    "ModelMetrics": {
+                        "ModelQuality": {
+                            "Statistics": {
+                                "ContentType": "text/csv",
+                                "S3Uri": f"s3://{BUCKET}/metrics.csv",
+                            },
+                        },
+                    },
+                    "ModelPackageDescription": "description",
+                    "ModelPackageGroupName": "mpg",
+                },
+            },
+        ]
+    )
+
+
 def test_register_model_with_model_repack(estimator, model_metrics):
     model_data = f"s3://{BUCKET}/model.tar.gz"
     register_model = RegisterModel(