fix: override register method in framework model class (#457)

Author: chuyang-deng

src/sagemaker/estimator.py

@@ -866,6 +866,7 @@ def register(
             transform_instances,
             model_package_name,
             model_package_group_name,
+            image_uri,
             model_metrics,
             marketplace_cert,
             approval_status,

src/sagemaker/model.py

@@ -112,6 +112,7 @@ def register(
         transform_instances,
         model_package_name=None,
         model_package_group_name=None,
+        image_uri=None,
         model_metrics=None,
         marketplace_cert=False,
         approval_status=None,
@@ -131,6 +132,8 @@ def register(
             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).
             marketplace_cert (bool): A boolean value indicating if the Model Package is certified
                 for AWS Marketplace (default: False).
@@ -151,6 +154,7 @@ def register(
             transform_instances,
             model_package_name,
             model_package_group_name,
+            image_uri,
             model_metrics,
             marketplace_cert,
             approval_status,
@@ -159,7 +163,11 @@ def register(
         model_package = self.sagemaker_session.create_model_package_from_containers(
             **model_pkg_args
         )
-        return model_package.get("ModelPackageArn")
+        return ModelPackage(
+            role=self.role,
+            model_data=self.model_data,
+            model_package_arn=model_package.get("ModelPackageArn"),
+        )
 
     def _get_model_package_args(
         self,
@@ -169,6 +177,7 @@ def _get_model_package_args(
         transform_instances,
         model_package_name=None,
         model_package_group_name=None,
+        image_uri=None,
         model_metrics=None,
         marketplace_cert=False,
         approval_status=None,
@@ -187,6 +196,8 @@ def _get_model_package_args(
             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).
             marketplace_cert (bool): A boolean value indicating if the Model Package is certified
                 for AWS Marketplace (default: False).
@@ -196,6 +207,8 @@ def _get_model_package_args(
         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,

src/sagemaker/mxnet/model.py

@@ -128,9 +128,72 @@ def __init__(
         super(MXNetModel, self).__init__(
             model_data, image_uri, role, entry_point, predictor_cls=predictor_cls, **kwargs
         )
-
         self.model_server_workers = model_server_workers
 
+    def register(
+        self,
+        content_types,
+        response_types,
+        inference_instances,
+        transform_instances,
+        model_package_name=None,
+        model_package_group_name=None,
+        image_uri=None,
+        model_metrics=None,
+        marketplace_cert=False,
+        approval_status=None,
+        description=None,
+    ):
+        """Creates a model package for creating SageMaker models or listing on Marketplace.
+
+        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).
+            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:
+            str: A string of SageMaker Model Package ARN.
+        """
+        instance_type = inference_instances[0]
+        self._init_sagemaker_session_if_does_not_exist(instance_type)
+
+        if image_uri:
+            self.image_uri = image_uri
+        if not self.image_uri:
+            self.image_uri = self.serving_image_uri(
+                region_name=self.sagemaker_session.boto_session.region_name,
+                instance_type=instance_type,
+            )
+        return super(MXNetModel, self).register(
+            content_types,
+            response_types,
+            inference_instances,
+            transform_instances,
+            model_package_name,
+            model_package_group_name,
+            image_uri,
+            model_metrics,
+            marketplace_cert,
+            approval_status,
+            description,
+        )
+
     def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.

src/sagemaker/pytorch/model.py

@@ -130,6 +130,70 @@ def __init__(
 
         self.model_server_workers = model_server_workers
 
+    def register(
+        self,
+        content_types,
+        response_types,
+        inference_instances,
+        transform_instances,
+        model_package_name=None,
+        model_package_group_name=None,
+        image_uri=None,
+        model_metrics=None,
+        marketplace_cert=False,
+        approval_status=None,
+        description=None,
+    ):
+        """Creates a model package for creating SageMaker models or listing on Marketplace.
+
+        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).
+            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:
+            str: A string of SageMaker Model Package ARN.
+        """
+        instance_type = inference_instances[0]
+        self._init_sagemaker_session_if_does_not_exist(instance_type)
+
+        if image_uri:
+            self.image_uri = image_uri
+        if not self.image_uri:
+            self.image_uri = self.serving_image_uri(
+                region_name=self.sagemaker_session.boto_session.region_name,
+                instance_type=instance_type,
+            )
+        return super(PyTorchModel, self).register(
+            content_types,
+            response_types,
+            inference_instances,
+            transform_instances,
+            model_package_name,
+            model_package_group_name,
+            image_uri,
+            model_metrics,
+            marketplace_cert,
+            approval_status,
+            description,
+        )
+
     def prepare_container_def(self, instance_type=None, accelerator_type=None):
         """Return a container definition with framework configuration set in
         model environment variables.

src/sagemaker/session.py

@@ -2634,7 +2634,7 @@ def _get_create_model_package_request(
                 "SupportedTransformInstanceTypes": transform_instances,
             }
             request_dict["InferenceSpecification"] = inference_specification
-        request_dict["CertifyForMarketPlace"] = marketplace_cert
+        request_dict["CertifyForMarketplace"] = marketplace_cert
         request_dict["ModelApprovalStatus"] = approval_status
         return request_dict
 

src/sagemaker/tensorflow/model.py

@@ -202,6 +202,70 @@ def __init__(
         )
         self._container_log_level = container_log_level
 
+    def register(
+        self,
+        content_types,
+        response_types,
+        inference_instances,
+        transform_instances,
+        model_package_name=None,
+        model_package_group_name=None,
+        image_uri=None,
+        model_metrics=None,
+        marketplace_cert=False,
+        approval_status=None,
+        description=None,
+    ):
+        """Creates a model package for creating SageMaker models or listing on Marketplace.
+
+        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).
+            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:
+            str: A string of SageMaker Model Package ARN.
+        """
+        instance_type = inference_instances[0]
+        self._init_sagemaker_session_if_does_not_exist(instance_type)
+
+        if image_uri:
+            self.image_uri = image_uri
+        if not self.image_uri:
+            self.image_uri = self.serving_image_uri(
+                region_name=self.sagemaker_session.boto_session.region_name,
+                instance_type=instance_type,
+            )
+        return super(TensorFlowModel, self).register(
+            content_types,
+            response_types,
+            inference_instances,
+            transform_instances,
+            model_package_name,
+            model_package_group_name,
+            image_uri,
+            model_metrics,
+            marketplace_cert,
+            approval_status,
+            description,
+        )
+
     def deploy(
         self,
         initial_instance_count,

src/sagemaker/workflow/_utils.py

@@ -305,8 +305,8 @@ def arguments(self) -> RequestType:
             **model_package_args
         )
         # these are not available in the workflow service
-        if "CertifyForMarketPlace" in request_dict:
-            request_dict.pop("CertifyForMarketPlace")
+        if "CertifyForMarketplace" in request_dict:
+            request_dict.pop("CertifyForMarketplace")
         if "Description" in request_dict:
             request_dict.pop("Description")
         if "ModelApprovalStatus" in request_dict:

tests/integ/test_mxnet.py

@@ -19,6 +19,7 @@
 import pytest
 
 import tests.integ
+from sagemaker import ModelPackage
 from sagemaker.mxnet.estimator import MXNet
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.utils import sagemaker_timestamp
@@ -160,6 +161,45 @@ def test_deploy_model(
         assert "Could not find model" in str(exception.value)
 
 
+def test_register_model_package(
+    mxnet_training_job,
+    sagemaker_session,
+    mxnet_inference_latest_version,
+    mxnet_inference_latest_py_version,
+    cpu_instance_type,
+):
+    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        desc = sagemaker_session.sagemaker_client.describe_training_job(
+            TrainingJobName=mxnet_training_job
+        )
+        model_data = desc["ModelArtifacts"]["S3ModelArtifacts"]
+        script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist.py")
+        model = MXNetModel(
+            model_data,
+            "SageMakerRole",
+            entry_point=script_path,
+            py_version=mxnet_inference_latest_py_version,
+            sagemaker_session=sagemaker_session,
+            framework_version=mxnet_inference_latest_version,
+        )
+        model_package_name = "register-model-package-{}".format(sagemaker_timestamp())
+        model_pkg = model.register(
+            content_types=["application/json"],
+            response_types=["application/json"],
+            inference_instances=["ml.m5.large"],
+            transform_instances=["ml.m5.large"],
+            model_package_name=model_package_name,
+        )
+        assert isinstance(model_pkg, ModelPackage)
+        predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)
+        data = numpy.zeros(shape=(1, 1, 28, 28))
+        result = predictor.predict(data)
+        assert result is not None
+        sagemaker_session.sagemaker_client.delete_model_package(ModelPackageName=model_package_name)
+
+
 def test_deploy_model_with_tags_and_kms(
     mxnet_training_job,
     sagemaker_session,