fix: add MetadataProperties to create_model_package request (aws#533)

Author: chuyang-deng

src/sagemaker/estimator.py

@@ -813,6 +813,7 @@ def register(
         model_package_name=None,
         model_package_group_name=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -837,6 +838,7 @@ def register(
                 `model_package_name`, using `model_package_group_name` makes the Model Package
                 versioned (default: None).
             model_metrics (ModelMetrics): ModelMetrics object (default: None).
+            metadata_properties (MetadataProperties): MetadataProperties (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",
@@ -868,6 +870,7 @@ def register(
             model_package_group_name,
             image_uri,
             model_metrics,
+            metadata_properties,
             marketplace_cert,
             approval_status,
             description,

src/sagemaker/metadata_properties.py

@@ -0,0 +1,52 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This file contains code related to metadata properties."""
+from __future__ import absolute_import
+
+
+class MetadataProperties(object):
+    """Accepts metadata properties parameters for conversion to request dict."""
+
+    def __init__(
+        self,
+        commit_id=None,
+        repository=None,
+        generated_by=None,
+        project_id=None,
+    ):
+        """Initialize a ``MetadataProperties`` instance and turn parameters into dict.
+
+        # TODO: flesh out docstrings
+        Args:
+            commit_id (str):
+            repository (str):
+            generated_by (str):
+            project_id (str):
+        """
+        self.commit_id = commit_id
+        self.repository = repository
+        self.generated_by = generated_by
+        self.project_id = project_id
+
+    def _to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        metadata_properties_request = dict()
+        if self.commit_id:
+            metadata_properties_request["CommitId"] = self.commit_id
+        if self.repository:
+            metadata_properties_request["Repository"] = self.repository
+        if self.generated_by:
+            metadata_properties_request["GeneratedBy"] = self.generated_by
+        if self.project_id:
+            metadata_properties_request["ProjectId"] = self.project_id
+        return metadata_properties_request

src/sagemaker/model.py

@@ -114,6 +114,7 @@ def register(
         model_package_group_name=None,
         image_uri=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -135,6 +136,7 @@ def register(
             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",
@@ -156,6 +158,7 @@ def register(
             model_package_group_name,
             image_uri,
             model_metrics,
+            metadata_properties,
             marketplace_cert,
             approval_status,
             description,
@@ -179,6 +182,7 @@ def _get_model_package_args(
         model_package_group_name=None,
         image_uri=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -199,6 +203,7 @@ def _get_model_package_args(
             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",
@@ -229,6 +234,8 @@ def _get_model_package_args(
             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:

src/sagemaker/model_metrics.py

@@ -1,4 +1,4 @@
-# Copyright 2019-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License"). You
 # may not use this file except in compliance with the License. A copy of

src/sagemaker/mxnet/model.py

@@ -140,6 +140,7 @@ def register(
         model_package_group_name=None,
         image_uri=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -161,6 +162,7 @@ def register(
             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 (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",
@@ -189,6 +191,7 @@ def register(
             model_package_group_name,
             image_uri,
             model_metrics,
+            metadata_properties,
             marketplace_cert,
             approval_status,
             description,

src/sagemaker/pytorch/model.py

@@ -140,6 +140,7 @@ def register(
         model_package_group_name=None,
         image_uri=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -161,6 +162,7 @@ def register(
             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",
@@ -189,6 +191,7 @@ def register(
             model_package_group_name,
             image_uri,
             model_metrics,
+            metadata_properties,
             marketplace_cert,
             approval_status,
             description,

src/sagemaker/tensorflow/model.py

@@ -212,6 +212,7 @@ def register(
         model_package_group_name=None,
         image_uri=None,
         model_metrics=None,
+        metadata_properties=None,
         marketplace_cert=False,
         approval_status=None,
         description=None,
@@ -233,6 +234,7 @@ def register(
             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",
@@ -261,6 +263,7 @@ def register(
             model_package_group_name,
             image_uri,
             model_metrics,
+            metadata_properties,
             marketplace_cert,
             approval_status,
             description,

src/sagemaker/workflow/_utils.py

@@ -216,6 +216,7 @@ def __init__(
         transform_instances,
         model_package_group_name=None,
         model_metrics=None,
+        metadata_properties=None,
         approval_status="PendingManualApproval",
         image_uri=None,
         compile_model_family=None,
@@ -238,6 +239,7 @@ def __init__(
                 `model_package_name`, using `model_package_group_name` makes the Model Package
                 versioned (default: None).
             model_metrics (ModelMetrics): ModelMetrics object (default: None).
+            metadata_properties (MetadataProperties): MetadataProperties object (default: None).
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval" (default: "PendingManualApproval").
             image_uri (str): The container image uri for Model Package, if not specified,
@@ -255,6 +257,7 @@ def __init__(
         self.transform_instances = transform_instances
         self.model_package_group_name = model_package_group_name
         self.model_metrics = model_metrics
+        self.metadata_properties = metadata_properties
         self.approval_status = approval_status
         self.image_uri = image_uri
         self.compile_model_family = compile_model_family
@@ -309,6 +312,7 @@ def arguments(self) -> RequestType:
             transform_instances=self.transform_instances,
             model_package_group_name=self.model_package_group_name,
             model_metrics=self.model_metrics,
+            metadata_properties=self.metadata_properties,
             approval_status=self.approval_status,
         )
         request_dict = model.sagemaker_session._get_create_model_package_request(

tests/unit/test_mxnet.py

@@ -22,6 +22,7 @@
 from pkg_resources import parse_version
 
 from sagemaker.fw_utils import UploadedCode
+from sagemaker.metadata_properties import MetadataProperties
 from sagemaker.model_metrics import MetricsSource, ModelMetrics
 from sagemaker.mxnet import defaults
 from sagemaker.mxnet import MXNet
@@ -544,13 +545,20 @@ def test_model_register_all_args(
         bias=dummy_metrics_source,
         explainability=dummy_metrics_source,
     )
+    metadata_properties = MetadataProperties(
+        commit_id="test-commit-id",
+        repository="test-repository",
+        generated_by="sagemaker-python-sdk-test",
+        project_id="test-project-id",
+    )
     model.register(
         content_types,
         response_types,
         inference_instances,
         transform_instances,
         model_package_name=model_package_name,
         model_metrics=model_metrics,
+        metadata_properties=metadata_properties,
         marketplace_cert=True,
         approval_status="Approved",
         description="description",
@@ -563,6 +571,7 @@ def test_model_register_all_args(
         "transform_instances": transform_instances,
         "model_package_name": model_package_name,
         "model_metrics": model_metrics._to_request_dict(),
+        "metadata_properties": metadata_properties._to_request_dict(),
         "marketplace_cert": True,
         "approval_status": "Approved",
         "description": "description",