fix: ensure model metrics is available to RegisterModel (aws#493)

Author: metrizable

src/sagemaker/workflow/_utils.py

@@ -215,9 +215,9 @@ def __init__(
         inference_instances,
         transform_instances,
         model_package_group_name=None,
-        image_uri=None,
         model_metrics=None,
         approval_status="PendingManualApproval",
+        image_uri=None,
         compile_model_family=None,
         **kwargs,
     ):
@@ -237,11 +237,11 @@ def __init__(
             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): The container image uri for Model Package, if not specified,
-                Estimator's training container image will be used (default: None).
             model_metrics (ModelMetrics): ModelMetrics 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,
+                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).
             **kwargs: additional arguments to `create_model`.
@@ -254,9 +254,9 @@ def __init__(
         self.inference_instances = inference_instances
         self.transform_instances = transform_instances
         self.model_package_group_name = model_package_group_name
-        self.image_uri = image_uri
         self.model_metrics = model_metrics
         self.approval_status = approval_status
+        self.image_uri = image_uri
         self.compile_model_family = compile_model_family
         self.kwargs = kwargs
 
@@ -314,7 +314,7 @@ def arguments(self) -> RequestType:
         request_dict = model.sagemaker_session._get_create_model_package_request(
             **model_package_args
         )
-        # these are not available in the workflow service
+        # these are not available in the workflow service and will cause rejection
         if "CertifyForMarketplace" in request_dict:
             request_dict.pop("CertifyForMarketplace")
         if "Description" in request_dict:

src/sagemaker/workflow/step_collections.py

@@ -65,6 +65,8 @@ def __init__(
         inference_instances,
         transform_instances,
         model_package_group_name=None,
+        model_metrics=None,
+        approval_status=None,
         image_uri=None,
         compile_model_family=None,
         **kwargs,
@@ -84,6 +86,9 @@ def __init__(
             model_package_group_name (str): The Model Package Group name, exclusive to
                 `model_package_name`, using `model_package_group_name` makes the Model Package
                 versioned (default: None).
+            model_metrics (ModelMetrics): ModelMetrics 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,
                 Estimator's training container image is used (default: None).
             compile_model_family (str): The instance family for the compiled model. If
@@ -115,6 +120,8 @@ def __init__(
             inference_instances=inference_instances,
             transform_instances=transform_instances,
             model_package_group_name=model_package_group_name,
+            model_metrics=model_metrics,
+            approval_status=approval_status,
             image_uri=image_uri,
             compile_model_family=compile_model_family,
             **kwargs,

tests/unit/sagemaker/workflow/test_step_collections.py

@@ -24,6 +24,10 @@
 
 from sagemaker.estimator import Estimator
 from sagemaker.inputs import CreateModelInput, TransformInput
+from sagemaker.model_metrics import (
+    MetricsSource,
+    ModelMetrics,
+)
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.steps import (
     Step,
@@ -108,6 +112,16 @@ def estimator(sagemaker_session):
     )
 
 
+@pytest.fixture
+def model_metrics():
+    return ModelMetrics(
+        model_statistics=MetricsSource(
+            s3_uri=f"s3://{BUCKET}/metrics.csv",
+            content_type="text/csv",
+        )
+    )
+
+
 def test_step_collection():
     step_collection = StepCollection(steps=[CustomStep("MyStep1"), CustomStep("MyStep2")])
     assert step_collection.request_dicts() == [
@@ -116,7 +130,7 @@ def test_step_collection():
     ]
 
 
-def test_register_model(estimator):
+def test_register_model(estimator, model_metrics):
     model_data = f"s3://{BUCKET}/model.tar.gz"
     register_model = RegisterModel(
         name="RegisterModelStep",
@@ -126,6 +140,9 @@ def test_register_model(estimator):
         response_types=["response_type"],
         inference_instances=["inference_instance"],
         transform_instances=["transform_instance"],
+        model_package_group_name="mpg",
+        model_metrics=model_metrics,
+        approval_status="Approved",
     )
     assert ordered(register_model.request_dicts()) == ordered(
         [
@@ -135,14 +152,23 @@ def test_register_model(estimator):
                 "Arguments": {
                     "InferenceSpecification": {
                         "Containers": [
-                            {"Image": "fakeimage", "ModelDataUrl": "s3://my-bucket/model.tar.gz"}
+                            {"Image": "fakeimage", "ModelDataUrl": f"s3://{BUCKET}/model.tar.gz"}
                         ],
                         "SupportedContentTypes": ["content_type"],
                         "SupportedRealtimeInferenceInstanceTypes": ["inference_instance"],
                         "SupportedResponseMIMETypes": ["response_type"],
                         "SupportedTransformInstanceTypes": ["transform_instance"],
                     },
-                    "ModelApprovalStatus": "PendingManualApproval",
+                    "ModelApprovalStatus": "Approved",
+                    "ModelMetrics": {
+                        "ModelQuality": {
+                            "Statistics": {
+                                "ContentType": "text/csv",
+                                "S3Uri": f"s3://{BUCKET}/metrics.csv",
+                            },
+                        },
+                    },
+                    "ModelPackageGroupName": "mpg",
                 },
             },
         ]