feat: add validation specification (#3075)

Co-authored-by: Basil Beirouti <[email protected]>

Author: BasilBeirouti

src/sagemaker/model.py

@@ -305,6 +305,7 @@ def register(
         description=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        validation_specification=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -360,6 +361,7 @@ def register(
             container_def_list=[container_def],
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            validation_specification=validation_specification,
         )
         model_package = self.sagemaker_session.create_model_package_from_containers(
             **model_pkg_args

src/sagemaker/session.py

@@ -2801,6 +2801,7 @@ def create_model_package_from_containers(
         description=None,
         drift_check_baselines=None,
         customer_metadata_properties=None,
+        validation_specification=None,
     ):
         """Get request dictionary for CreateModelPackage API.
 
@@ -2846,6 +2847,7 @@ def create_model_package_from_containers(
             description,
             drift_check_baselines=drift_check_baselines,
             customer_metadata_properties=customer_metadata_properties,
+            validation_specification=validation_specification,
         )
         if model_package_group_name is not None:
             try:
@@ -4206,6 +4208,7 @@ def get_model_package_args(
     container_def_list=None,
     drift_check_baselines=None,
     customer_metadata_properties=None,
+    validation_specification=None,
 ):
     """Get arguments for create_model_package method.
 
@@ -4275,6 +4278,8 @@ def get_model_package_args(
         model_package_args["tags"] = tags
     if customer_metadata_properties is not None:
         model_package_args["customer_metadata_properties"] = customer_metadata_properties
+    if validation_specification is not None:
+        model_package_args["validation_specification"] = validation_specification
     return model_package_args
 
 
@@ -4294,6 +4299,7 @@ def get_create_model_package_request(
     tags=None,
     drift_check_baselines=None,
     customer_metadata_properties=None,
+    validation_specification=None,
 ):
     """Get request dictionary for CreateModelPackage API.
 
@@ -4345,6 +4351,8 @@ def get_create_model_package_request(
         request_dict["MetadataProperties"] = metadata_properties
     if customer_metadata_properties is not None:
         request_dict["CustomerMetadataProperties"] = customer_metadata_properties
+    if validation_specification:
+        request_dict["ValidationSpecification"] = validation_specification
     if containers is not None:
         if not all([content_types, response_types, inference_instances, transform_instances]):
             raise ValueError(

tests/unit/sagemaker/model/test_model.py

@@ -48,6 +48,47 @@
 IMAGE_URI = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.9.0-gpu-py38"
 
 
+MODEL_DESCRIPTION = "a description"
+
+SUPPORTED_REALTIME_INFERENCE_INSTANCE_TYPES = ["ml.m4.xlarge"]
+SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES = ["ml.m4.xlarge"]
+
+SUPPORTED_CONTENT_TYPES = ["text/csv", "application/json", "application/jsonlines"]
+SUPPORTED_RESPONSE_MIME_TYPES = ["application/json", "text/csv", "application/jsonlines"]
+
+VALIDATION_FILE_NAME = "input.csv"
+VALIDATION_INPUT_PATH = "s3://" + BUCKET_NAME + "/validation-input-csv/"
+VALIDATION_OUTPUT_PATH = "s3://" + BUCKET_NAME + "/validation-output-csv/"
+
+VALIDATION_SPECIFICATION = {
+    "ValidationRole": "some_role",
+    "ValidationProfiles": [
+        {
+            "ProfileName": "Validation-test",
+            "TransformJobDefinition": {
+                "BatchStrategy": "SingleRecord",
+                "TransformInput": {
+                    "DataSource": {
+                        "S3DataSource": {
+                            "S3DataType": "S3Prefix",
+                            "S3Uri": VALIDATION_INPUT_PATH,
+                        }
+                    },
+                    "ContentType": SUPPORTED_CONTENT_TYPES[0],
+                },
+                "TransformOutput": {
+                    "S3OutputPath": VALIDATION_OUTPUT_PATH,
+                },
+                "TransformResources": {
+                    "InstanceType": SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES[0],
+                    "InstanceCount": 1,
+                },
+            },
+        },
+    ],
+}
+
+
 class DummyFrameworkModel(FrameworkModel):
     def __init__(self, **kwargs):
         super(DummyFrameworkModel, self).__init__(
@@ -687,3 +728,40 @@ def test_script_mode_model_uses_proper_sagemaker_submit_dir(repack_model, sagema
         ]
         == "/opt/ml/model/code"
     )
+
+
+@patch("sagemaker.get_model_package_args")
+def test_register_calls_model_package_args(get_model_package_args, sagemaker_session):
+
+    source_dir = "s3://blah/blah/blah"
+    t = Model(
+        entry_point=ENTRY_POINT_INFERENCE,
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        source_dir=source_dir,
+        image_uri=IMAGE_URI,
+        model_data=MODEL_DATA,
+    )
+
+    t.register(
+        SUPPORTED_CONTENT_TYPES,
+        SUPPORTED_RESPONSE_MIME_TYPES,
+        SUPPORTED_REALTIME_INFERENCE_INSTANCE_TYPES,
+        SUPPORTED_BATCH_TRANSFORM_INSTANCE_TYPES,
+        marketplace_cert=True,
+        description=MODEL_DESCRIPTION,
+        model_package_name=MODEL_NAME,
+        validation_specification=VALIDATION_SPECIFICATION,
+    )
+
+    # check that the kwarg validation_specification was passed to the internal method 'get_model_package_args'
+    assert (
+        "validation_specification" in get_model_package_args.call_args_list[0][1]
+    ), "validation_specification kwarg was not passed to get_model_package_args"
+
+    # check that the kwarg validation_specification is identical to the one passed into the method 'register'
+    assert (
+        VALIDATION_SPECIFICATION
+        == get_model_package_args.call_args_list[0][1]["validation_specification"]
+    ), """ValidationSpecification from model.register method is not identical to validation_spec from
+         get_model_package_args"""