feature: Add support for ModelMonitor/Clarify integration in model building pipelines

* feature: Add support for QualityCheckStep and ClarifyCheckStep in model building pipelines

* feature: Adding Baseline metrics for model packages

Co-authored-by: qidewenwhen <[email protected]>

Author: 2 people

doc/api/inference/model_monitor.rst

@@ -31,3 +31,13 @@ Model Monitor
     :members:
     :undoc-members:
     :show-inheritance:
+
+.. automodule:: sagemaker.model_metrics
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. automodule:: sagemaker.drift_check_baselines
+    :members:
+    :undoc-members:
+    :show-inheritance:

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -32,6 +32,11 @@ Conditions
 
 .. autoclass:: sagemaker.workflow.conditions.ConditionOr
 
+CheckJobConfig
+--------------
+
+.. autoclass:: sagemaker.workflow.check_job_config.CheckJobConfig
+
 Entities
 --------
 
@@ -128,3 +133,13 @@ Steps
 .. autoclass:: sagemaker.workflow.steps.CacheConfig
 
 .. autoclass:: sagemaker.workflow.lambda_step.LambdaStep
+
+.. autoclass:: sagemaker.workflow.steps.CompilationStep
+
+.. autoclass:: sagemaker.workflow.quality_check_step.QualityCheckConfig
+
+.. autoclass:: sagemaker.workflow.quality_check_step.QualityCheckStep
+
+.. autoclass:: sagemaker.workflow.clarify_check_step.ClarifyCheckConfig
+
+.. autoclass:: sagemaker.workflow.clarify_check_step.ClarifyCheckStep

src/sagemaker/__init__.py

@@ -49,7 +49,7 @@
 from sagemaker.local.local_session import LocalSession  # noqa: F401
 
 from sagemaker.model import Model, ModelPackage  # noqa: F401
-from sagemaker.model_metrics import ModelMetrics, MetricsSource  # noqa: F401
+from sagemaker.model_metrics import ModelMetrics, MetricsSource, FileSource  # noqa: F401
 from sagemaker.pipeline import PipelineModel  # noqa: F401
 from sagemaker.predictor import Predictor  # noqa: F401
 from sagemaker.processing import Processor, ScriptProcessor  # noqa: F401

src/sagemaker/clarify.py

@@ -58,7 +58,7 @@ def __init__(
             s3_data_distribution_type (str): Valid options are "FullyReplicated" or
                 "ShardedByS3Key".
             s3_compression_type (str): Valid options are "None" or "Gzip".
-            joinsource (str): The name or index of the column in the dataset that acts an
+            joinsource (str): The name or index of the column in the dataset that acts as an
                 identifier column (for instance, while performing a join). This column is only
                 used as an identifier, and not used for any other computations. This is an
                 optional field in all cases except when the dataset contains more than one file,

src/sagemaker/drift_check_baselines.py

@@ -0,0 +1,102 @@
+# Copyright 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 drift check baselines"""
+from __future__ import absolute_import
+
+
+class DriftCheckBaselines(object):
+    """Accepts drift check baselines parameters for conversion to request dict."""
+
+    def __init__(
+        self,
+        model_statistics=None,
+        model_constraints=None,
+        model_data_statistics=None,
+        model_data_constraints=None,
+        bias_config_file=None,
+        bias_pre_training_constraints=None,
+        bias_post_training_constraints=None,
+        explainability_constraints=None,
+        explainability_config_file=None,
+    ):
+        """Initialize a ``DriftCheckBaselines`` instance and turn parameters into dict.
+
+        Args:
+            model_statistics (MetricsSource): A metric source object that represents
+                model statistics (default: None).
+            model_constraints (MetricsSource): A metric source object that represents
+                model constraints (default: None).
+            model_data_statistics (MetricsSource): A metric source object that represents
+                model data statistics (default: None).
+            model_data_constraints (MetricsSource): A metric source object that represents
+                model data constraints (default: None).
+            bias_config_file (FileSource): A file source object that represents bias config
+                (default: None).
+            bias_pre_training_constraints (MetricsSource):
+                A metric source object that represents Pre-training constraints (default: None).
+            bias_post_training_constraints (MetricsSource):
+                A metric source object that represents Post-training constraits (default: None).
+            explainability_constraints (MetricsSource):
+                A metric source object that represents explainability constraints (default: None).
+            explainability_config_file (FileSource): A file source object that represents
+                explainability config (default: None).
+        """
+        self.model_statistics = model_statistics
+        self.model_constraints = model_constraints
+        self.model_data_statistics = model_data_statistics
+        self.model_data_constraints = model_data_constraints
+        self.bias_config_file = bias_config_file
+        self.bias_pre_training_constraints = bias_pre_training_constraints
+        self.bias_post_training_constraints = bias_post_training_constraints
+        self.explainability_constraints = explainability_constraints
+        self.explainability_config_file = explainability_config_file
+
+    def _to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        drift_check_baselines_request = {}
+
+        model_quality = {}
+        if self.model_statistics is not None:
+            model_quality["Statistics"] = self.model_statistics._to_request_dict()
+        if self.model_constraints is not None:
+            model_quality["Constraints"] = self.model_constraints._to_request_dict()
+        if model_quality:
+            drift_check_baselines_request["ModelQuality"] = model_quality
+
+        model_data_quality = {}
+        if self.model_data_statistics is not None:
+            model_data_quality["Statistics"] = self.model_data_statistics._to_request_dict()
+        if self.model_data_constraints is not None:
+            model_data_quality["Constraints"] = self.model_data_constraints._to_request_dict()
+        if model_data_quality:
+            drift_check_baselines_request["ModelDataQuality"] = model_data_quality
+
+        bias = {}
+        if self.bias_config_file is not None:
+            bias["ConfigFile"] = self.bias_config_file._to_request_dict()
+        if self.bias_pre_training_constraints is not None:
+            bias["PreTrainingConstraints"] = self.bias_pre_training_constraints._to_request_dict()
+        if self.bias_post_training_constraints is not None:
+            bias["PostTrainingConstraints"] = self.bias_post_training_constraints._to_request_dict()
+        if bias:
+            drift_check_baselines_request["Bias"] = bias
+
+        explainability = {}
+        if self.explainability_constraints is not None:
+            explainability["Constraints"] = self.explainability_constraints._to_request_dict()
+        if self.explainability_config_file is not None:
+            explainability["ConfigFile"] = self.explainability_config_file._to_request_dict()
+        if explainability:
+            drift_check_baselines_request["Explainability"] = explainability
+
+        return drift_check_baselines_request

src/sagemaker/estimator.py

@@ -977,6 +977,7 @@ def register(
         description=None,
         compile_model_family=None,
         model_name=None,
+        drift_check_baselines=None,
         **kwargs,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
@@ -1005,6 +1006,7 @@ def register(
             compile_model_family (str): Instance family for compiled model, if specified, a compiled
                 model will be used (default: None).
             model_name (str): User defined model name (default: None).
+            drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during
                 deploy. For more, see the implementation docs.
@@ -1034,6 +1036,7 @@ def register(
             marketplace_cert,
             approval_status,
             description,
+            drift_check_baselines=drift_check_baselines,
         )
 
     @property
@@ -1920,7 +1923,13 @@ def training_image_uri(self):
         return self.image_uri
 
     def set_hyperparameters(self, **kwargs):
-        """Placeholder docstring"""
+        """Sets the hyperparameter dictionary to use for training.
+
+        The hyperparameters are made accessible as a dict[str, str] to the
+        training code on SageMaker. For convenience, this accepts other types
+        for keys and values, but ``str()`` will be called to convert them before
+        training.
+        """
         for k, v in kwargs.items():
             self.hyperparam_dict[k] = v
 

src/sagemaker/huggingface/model.py

@@ -186,6 +186,7 @@ def register(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        drift_check_baselines=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -212,6 +213,7 @@ def register(
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval". Defaults to ``PendingManualApproval``.
             description (str): Model Package description. Defaults to ``None``.
+            drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -239,6 +241,7 @@ def register(
             marketplace_cert,
             approval_status,
             description,
+            drift_check_baselines=drift_check_baselines,
         )
 
     def prepare_container_def(self, instance_type=None, accelerator_type=None):

src/sagemaker/model.py

@@ -146,6 +146,7 @@ def register(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        drift_check_baselines=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -170,6 +171,7 @@ def register(
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval" (default: "PendingManualApproval").
             description (str): Model Package description (default: None).
+            drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -191,6 +193,7 @@ def register(
             marketplace_cert,
             approval_status,
             description,
+            drift_check_baselines=drift_check_baselines,
         )
         model_package = self.sagemaker_session.create_model_package_from_containers(
             **model_pkg_args

src/sagemaker/model_metrics.py

@@ -10,7 +10,7 @@
 # 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 model metrics, including metric source."""
+"""This file contains code related to model metrics, including metric source and file source."""
 from __future__ import absolute_import
 
 
@@ -25,22 +25,36 @@ def __init__(
         model_data_constraints=None,
         bias=None,
         explainability=None,
+        bias_pre_training=None,
+        bias_post_training=None,
     ):
         """Initialize a ``ModelMetrics`` instance and turn parameters into dict.
 
-        # TODO: flesh out docstrings
         Args:
-            model_constraints (MetricsSource):
-            model_data_constraints (MetricsSource):
-            model_data_statistics (MetricsSource):
-            bias (MetricsSource):
-            explainability (MetricsSource):
+            model_statistics (MetricsSource): A metric source object that represents
+                model statistics (default: None).
+            model_constraints (MetricsSource): A metric source object that represents
+                model constraints (default: None).
+            model_data_statistics (MetricsSource): A metric source object that represents
+                model data statistics (default: None).
+            model_data_constraints (MetricsSource): A metric source object that represents
+                model data constraints (default: None).
+            bias (MetricsSource): A metric source object that represents bias report
+                (default: None).
+            explainability (MetricsSource): A metric source object that represents
+                explainability report (default: None).
+            bias_pre_training (MetricsSource): A metric source object that represents
+                Pre-training report (default: None).
+            bias_post_training (MetricsSource): A metric source object that represents
+                Post-training report (default: None).
         """
         self.model_statistics = model_statistics
         self.model_constraints = model_constraints
         self.model_data_statistics = model_data_statistics
         self.model_data_constraints = model_data_constraints
         self.bias = bias
+        self.bias_pre_training = bias_pre_training
+        self.bias_post_training = bias_post_training
         self.explainability = explainability
 
     def _to_request_dict(self):
@@ -63,10 +77,20 @@ def _to_request_dict(self):
         if model_data_quality:
             model_metrics_request["ModelDataQuality"] = model_data_quality
 
+        bias = {}
         if self.bias is not None:
-            model_metrics_request["Bias"] = self.bias._to_request_dict()
+            bias["Report"] = self.bias._to_request_dict()
+        if self.bias_pre_training is not None:
+            bias["PreTrainingReport"] = self.bias_pre_training._to_request_dict()
+        if self.bias_post_training is not None:
+            bias["PostTrainingReport"] = self.bias_post_training._to_request_dict()
+        model_metrics_request["Bias"] = bias
+
+        explainability = {}
         if self.explainability is not None:
-            model_metrics_request["Explainability"] = self.explainability._to_request_dict()
+            explainability["Report"] = self.explainability._to_request_dict()
+        model_metrics_request["Explainability"] = explainability
+
         return model_metrics_request
 
 
@@ -81,11 +105,10 @@ def __init__(
     ):
         """Initialize a ``MetricsSource`` instance and turn parameters into dict.
 
-        # TODO: flesh out docstrings
         Args:
-            content_type (str):
-            s3_uri (str):
-            content_digest (str):
+            content_type (str): Specifies the type of content in S3 URI
+            s3_uri (str): The S3 URI of the metric
+            content_digest (str): The digest of the metric (default: None)
         """
         self.content_type = content_type
         self.s3_uri = s3_uri
@@ -97,3 +120,33 @@ def _to_request_dict(self):
         if self.content_digest is not None:
             metrics_source_request["ContentDigest"] = self.content_digest
         return metrics_source_request
+
+
+class FileSource(object):
+    """Accepts file source parameters for conversion to request dict."""
+
+    def __init__(
+        self,
+        s3_uri,
+        content_digest=None,
+        content_type=None,
+    ):
+        """Initialize a ``FileSource`` instance and turn parameters into dict.
+
+        Args:
+            s3_uri (str): The S3 URI of the metric
+            content_digest (str): The digest of the metric (default: None)
+            content_type (str): Specifies the type of content in S3 URI (default: None)
+        """
+        self.content_type = content_type
+        self.s3_uri = s3_uri
+        self.content_digest = content_digest
+
+    def _to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        file_source_request = {"S3Uri": self.s3_uri}
+        if self.content_digest is not None:
+            file_source_request["ContentDigest"] = self.content_digest
+        if self.content_type is not None:
+            file_source_request["ContentType"] = self.content_type
+        return file_source_request

src/sagemaker/mxnet/model.py

@@ -157,6 +157,7 @@ def register(
         marketplace_cert=False,
         approval_status=None,
         description=None,
+        drift_check_baselines=None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -181,6 +182,7 @@ def register(
             approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
                 or "PendingManualApproval" (default: "PendingManualApproval").
             description (str): Model Package description (default: None).
+            drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -208,6 +210,7 @@ def register(
             marketplace_cert,
             approval_status,
             description,
+            drift_check_baselines=drift_check_baselines,
         )
 
     def prepare_container_def(self, instance_type=None, accelerator_type=None):