feature: Adding serial inference pipeline support to RegisterModel Step (#2405)

* feature: Adding serial inference pipeline support to RegisterModel Step

* feature: Adding serial inference pipeline support to RegisterModel Step

* Using Model class to abstract the containers

* Using Model class to abstract the containers

* Formatting changes and lint fixes

* Repack steps for models

* Repack steps for models

* lint error corrections

* lint error corrections

* Unit test corrections

* Unit test corrections

* Integ test added

* Update test_workflow.py

* Update model.py

* Update model.py

* fix unit test failure

* fixing unit tests for repack

* Path fix for the integ tests

* Fix integ test paths

* Black tool fixes

* Test data files

* Updating the input to RegisterModelStep to Pipeline Model

* Fix for local tests

* Checking if the key is present

* black check fix

* Fix unit test failures

* Updating the input to model or pipeline model

* Review comment incorporation

* Fixing unit test failure

* Fixing unit test failure

Co-authored-by: Ahsan Khan <[email protected]>

Author: sreedes

src/sagemaker/__init__.py

@@ -55,6 +55,7 @@
 from sagemaker.processing import Processor, ScriptProcessor  # noqa: F401
 from sagemaker.session import Session  # noqa: F401
 from sagemaker.session import container_def, pipeline_container_def  # noqa: F401
+from sagemaker.session import get_model_package_args  # noqa: F401
 from sagemaker.session import production_variant  # noqa: F401
 from sagemaker.session import get_execution_role  # noqa: F401
 

src/sagemaker/model.py

@@ -158,14 +158,15 @@ def register(
         if self.model_data is None:
             raise ValueError("SageMaker Model Package cannot be created without model data.")
 
-        model_pkg_args = self._get_model_package_args(
+        model_pkg_args = sagemaker.get_model_package_args(
             content_types,
             response_types,
             inference_instances,
             transform_instances,
             model_package_name,
             model_package_group_name,
-            image_uri,
+            self.model_data,
+            image_uri or self.image_uri,
             model_metrics,
             metadata_properties,
             marketplace_cert,
@@ -181,80 +182,6 @@ def register(
             model_package_arn=model_package.get("ModelPackageArn"),
         )
 
-    def _get_model_package_args(
-        self,
-        content_types,
-        response_types,
-        inference_instances,
-        transform_instances,
-        model_package_name=None,
-        model_package_group_name=None,
-        image_uri=None,
-        model_metrics=None,
-        metadata_properties=None,
-        marketplace_cert=False,
-        approval_status=None,
-        description=None,
-        tags=None,
-    ):
-        """Get arguments for session.create_model_package method.
-
-        Args:
-            content_types (list): The supported MIME types for the input data.
-            response_types (list): The supported MIME types for the output data.
-            inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time.
-            transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed.
-            model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
-                using `model_package_name` makes the Model Package un-versioned (default: None).
-            model_package_group_name (str): Model Package Group name, exclusive to
-                `model_package_name`, using `model_package_group_name` makes the Model Package
-                versioned (default: None).
-            image_uri (str): Inference image uri for the container. Model class' self.image will
-                be used if it is None (default: None).
-            model_metrics (ModelMetrics): ModelMetrics object (default: None).
-            metadata_properties (MetadataProperties): MetadataProperties object (default: None).
-            marketplace_cert (bool): A boolean value indicating if the Model Package is certified
-                for AWS Marketplace (default: False).
-            approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
-                or "PendingManualApproval" (default: "PendingManualApproval").
-            description (str): Model Package description (default: None).
-        Returns:
-            dict: A dictionary of method argument names and values.
-        """
-        if image_uri:
-            self.image_uri = image_uri
-        container = {
-            "Image": self.image_uri,
-            "ModelDataUrl": self.model_data,
-        }
-
-        model_package_args = {
-            "containers": [container],
-            "content_types": content_types,
-            "response_types": response_types,
-            "inference_instances": inference_instances,
-            "transform_instances": transform_instances,
-            "marketplace_cert": marketplace_cert,
-        }
-
-        if model_package_name is not None:
-            model_package_args["model_package_name"] = model_package_name
-        if model_package_group_name is not None:
-            model_package_args["model_package_group_name"] = model_package_group_name
-        if model_metrics is not None:
-            model_package_args["model_metrics"] = model_metrics._to_request_dict()
-        if metadata_properties is not None:
-            model_package_args["metadata_properties"] = metadata_properties._to_request_dict()
-        if approval_status is not None:
-            model_package_args["approval_status"] = approval_status
-        if description is not None:
-            model_package_args["description"] = description
-        if tags is not None:
-            model_package_args["tags"] = tags
-        return model_package_args
-
     def _init_sagemaker_session_if_does_not_exist(self, instance_type):
         """Set ``self.sagemaker_session`` to ``LocalSession`` or ``Session`` if it's not already.
 
@@ -1128,6 +1055,10 @@ def _upload_code(self, key_prefix, repack=False):
             )
 
         if repack and self.model_data is not None and self.entry_point is not None:
+            if isinstance(self.model_data, sagemaker.workflow.properties.Properties):
+                # model is not yet there, defer repacking to later during pipeline execution
+                return
+
             bucket = self.bucket or self.sagemaker_session.default_bucket()
             repacked_model_data = "s3://" + "/".join([bucket, key_prefix, "model.tar.gz"])
 

src/sagemaker/session.py

@@ -2521,8 +2521,11 @@ def _create_model_request(
         request = {"ModelName": name, "ExecutionRoleArn": role}
         if isinstance(container_definition, list):
             request["Containers"] = container_definition
+        elif "ModelPackageName" in container_definition:
+            request["Containers"] = [container_definition]
         else:
             request["PrimaryContainer"] = container_definition
+
         if tags:
             request["Tags"] = tags
 
@@ -2731,7 +2734,7 @@ def create_model_package_from_containers(
             description (str): Model Package description (default: None).
         """
 
-        request = self._get_create_model_package_request(
+        request = get_create_model_package_request(
             model_package_name,
             model_package_group_name,
             containers,
@@ -2747,82 +2750,6 @@ def create_model_package_from_containers(
         )
         return self.sagemaker_client.create_model_package(**request)
 
-    def _get_create_model_package_request(
-        self,
-        model_package_name=None,
-        model_package_group_name=None,
-        containers=None,
-        content_types=None,
-        response_types=None,
-        inference_instances=None,
-        transform_instances=None,
-        model_metrics=None,
-        metadata_properties=None,
-        marketplace_cert=False,
-        approval_status="PendingManualApproval",
-        description=None,
-        tags=None,
-    ):
-        """Get request dictionary for CreateModelPackage API.
-
-        Args:
-            model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
-                using `model_package_name` makes the Model Package un-versioned (default: None).
-            model_package_group_name (str): Model Package Group name, exclusive to
-                `model_package_name`, using `model_package_group_name` makes the Model Package
-                versioned (default: None).
-            containers (list): A list of inference containers that can be used for inference
-                specifications of Model Package (default: None).
-            content_types (list): The supported MIME types for the input data (default: None).
-            response_types (list): The supported MIME types for the output data (default: None).
-            inference_instances (list): A list of the instance types that are used to
-                generate inferences in real-time (default: None).
-            transform_instances (list): A list of the instance types on which a transformation
-                job can be run or on which an endpoint can be deployed (default: None).
-            model_metrics (ModelMetrics): ModelMetrics object (default: None).
-            metadata_properties (MetadataProperties): MetadataProperties object (default: None).
-            marketplace_cert (bool): A boolean value indicating if the Model Package is certified
-                for AWS Marketplace (default: False).
-            approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
-                or "PendingManualApproval" (default: "PendingManualApproval").
-            description (str): Model Package description (default: None).
-        """
-        if all([model_package_name, model_package_group_name]):
-            raise ValueError(
-                "model_package_name and model_package_group_name cannot be present at the "
-                "same time."
-            )
-        request_dict = {}
-        if model_package_name is not None:
-            request_dict["ModelPackageName"] = model_package_name
-        if model_package_group_name is not None:
-            request_dict["ModelPackageGroupName"] = model_package_group_name
-        if description is not None:
-            request_dict["ModelPackageDescription"] = description
-        if tags is not None:
-            request_dict["Tags"] = tags
-        if model_metrics:
-            request_dict["ModelMetrics"] = model_metrics
-        if metadata_properties:
-            request_dict["MetadataProperties"] = metadata_properties
-        if containers is not None:
-            if not all([content_types, response_types, inference_instances, transform_instances]):
-                raise ValueError(
-                    "content_types, response_types, inference_inferences and transform_instances "
-                    "must be provided if containers is present."
-                )
-            inference_specification = {
-                "Containers": containers,
-                "SupportedContentTypes": content_types,
-                "SupportedResponseMIMETypes": response_types,
-                "SupportedRealtimeInferenceInstanceTypes": inference_instances,
-                "SupportedTransformInstanceTypes": transform_instances,
-            }
-            request_dict["InferenceSpecification"] = inference_specification
-        request_dict["CertifyForMarketplace"] = marketplace_cert
-        request_dict["ModelApprovalStatus"] = approval_status
-        return request_dict
-
     def wait_for_model_package(self, model_package_name, poll=5):
         """Wait for an Amazon SageMaker endpoint deployment to complete.
 
@@ -4097,6 +4024,160 @@ def account_id(self) -> str:
         return sts_client.get_caller_identity()["Account"]
 
 
+def get_model_package_args(
+    content_types,
+    response_types,
+    inference_instances,
+    transform_instances,
+    model_package_name=None,
+    model_package_group_name=None,
+    model_data=None,
+    image_uri=None,
+    model_metrics=None,
+    metadata_properties=None,
+    marketplace_cert=False,
+    approval_status=None,
+    description=None,
+    tags=None,
+    container_def_list=None,
+):
+    """Get arguments for create_model_package method.
+
+    Args:
+        content_types (list): The supported MIME types for the input data.
+        response_types (list): The supported MIME types for the output data.
+        inference_instances (list): A list of the instance types that are used to
+        generate inferences in real-time.
+        transform_instances (list): A list of the instance types on which a transformation
+        job can be run or on which an endpoint can be deployed.
+        model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
+        using `model_package_name` makes the Model Package un-versioned (default: None).
+        model_package_group_name (str): Model Package Group name, exclusive to
+        `model_package_name`, using `model_package_group_name` makes the Model Package
+        versioned (default: None).
+        image_uri (str): Inference image uri for the container. Model class' self.image will
+        be used if it is None (default: None).
+        model_metrics (ModelMetrics): ModelMetrics object (default: None).
+        metadata_properties (MetadataProperties): MetadataProperties object (default: None).
+        marketplace_cert (bool): A boolean value indicating if the Model Package is certified
+        for AWS Marketplace (default: False).
+        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.
+    """
+    if container_def_list is not None:
+        containers = container_def_list
+    else:
+        container = {
+            "Image": image_uri,
+            "ModelDataUrl": model_data,
+        }
+        containers = [container]
+
+    model_package_args = {
+        "containers": containers,
+        "content_types": content_types,
+        "response_types": response_types,
+        "inference_instances": inference_instances,
+        "transform_instances": transform_instances,
+        "marketplace_cert": marketplace_cert,
+    }
+
+    if model_package_name is not None:
+        model_package_args["model_package_name"] = model_package_name
+    if model_package_group_name is not None:
+        model_package_args["model_package_group_name"] = model_package_group_name
+    if model_metrics is not None:
+        model_package_args["model_metrics"] = model_metrics._to_request_dict()
+    if metadata_properties is not None:
+        model_package_args["metadata_properties"] = metadata_properties._to_request_dict()
+    if approval_status is not None:
+        model_package_args["approval_status"] = approval_status
+    if description is not None:
+        model_package_args["description"] = description
+    if tags is not None:
+        model_package_args["tags"] = tags
+    return model_package_args
+
+
+def get_create_model_package_request(
+    model_package_name=None,
+    model_package_group_name=None,
+    containers=None,
+    content_types=None,
+    response_types=None,
+    inference_instances=None,
+    transform_instances=None,
+    model_metrics=None,
+    metadata_properties=None,
+    marketplace_cert=False,
+    approval_status="PendingManualApproval",
+    description=None,
+    tags=None,
+):
+    """Get request dictionary for CreateModelPackage API.
+
+    Args:
+        model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
+        using `model_package_name` makes the Model Package un-versioned (default: None).
+        model_package_group_name (str): Model Package Group name, exclusive to
+        `model_package_name`, using `model_package_group_name` makes the Model Package
+        versioned (default: None).
+        containers (list): A list of inference containers that can be used for inference
+        specifications of Model Package (default: None).
+        content_types (list): The supported MIME types for the input data (default: None).
+        response_types (list): The supported MIME types for the output data (default: None).
+        inference_instances (list): A list of the instance types that are used to
+        generate inferences in real-time (default: None).
+        transform_instances (list): A list of the instance types on which a transformation
+        job can be run or on which an endpoint can be deployed (default: None).
+        model_metrics (ModelMetrics): ModelMetrics object (default: None).
+        metadata_properties (MetadataProperties): MetadataProperties object (default: None).
+        marketplace_cert (bool): A boolean value indicating if the Model Package is certified
+        for AWS Marketplace (default: False).
+        approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
+        or "PendingManualApproval" (default: "PendingManualApproval").
+        description (str): Model Package description (default: None).
+    """
+    if all([model_package_name, model_package_group_name]):
+        raise ValueError(
+            "model_package_name and model_package_group_name cannot be present at the " "same time."
+        )
+    request_dict = {}
+    if model_package_name is not None:
+        request_dict["ModelPackageName"] = model_package_name
+    if model_package_group_name is not None:
+        request_dict["ModelPackageGroupName"] = model_package_group_name
+    if description is not None:
+        request_dict["ModelPackageDescription"] = description
+    if tags is not None:
+        request_dict["Tags"] = tags
+    if model_metrics:
+        request_dict["ModelMetrics"] = model_metrics
+    if metadata_properties:
+        request_dict["MetadataProperties"] = metadata_properties
+    if containers is not None:
+        if not all([content_types, response_types, inference_instances, transform_instances]):
+            raise ValueError(
+                "content_types, response_types, inference_inferences and transform_instances "
+                "must be provided if containers is present."
+            )
+        inference_specification = {
+            "Containers": containers,
+            "SupportedContentTypes": content_types,
+            "SupportedResponseMIMETypes": response_types,
+            "SupportedRealtimeInferenceInstanceTypes": inference_instances,
+            "SupportedTransformInstanceTypes": transform_instances,
+        }
+        request_dict["InferenceSpecification"] = inference_specification
+    request_dict["CertifyForMarketplace"] = marketplace_cert
+    request_dict["ModelApprovalStatus"] = approval_status
+    return request_dict
+
+
 def update_args(args: Dict[str, Any], **kwargs):
     """Updates the request arguments dict with the value if populated.