feat: Add optional monitoring_config_override parameter in suggest_baseline API (#3939)

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

Author: kusumbhatt

doc/api/inference/model_monitor.rst

@@ -41,3 +41,8 @@ Model Monitor
     :members:
     :undoc-members:
     :show-inheritance:
+
+.. automodule:: sagemaker.model_monitor.data_quality_monitoring_config
+    :members:
+    :undoc-members:
+    :show-inheritance:

src/sagemaker/model_monitor/__init__.py

@@ -46,3 +46,10 @@
 from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat  # noqa: F401
 
 from sagemaker.network import NetworkConfig  # noqa: F401
+
+from sagemaker.model_monitor.data_quality_monitoring_config import (  # noqa: F401
+    DataQualityMonitoringConfig,
+)
+from sagemaker.model_monitor.data_quality_monitoring_config import (  # noqa: F401
+    DataQualityDistributionConstraints,
+)

src/sagemaker/model_monitor/data_quality_monitoring_config.py

@@ -0,0 +1,66 @@
+# 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 module contains code related to the MonitoringConfig of constraints file.
+
+Code is used to represent the Monitoring Config object and its parameters suggested
+in constraints file by Model Monitor Container in data quality analysis.
+"""
+from __future__ import print_function, absolute_import
+
+CHI_SQUARED_METHOD = "ChiSquared"
+L_INFINITY_METHOD = "LInfinity"
+
+
+class DataQualityDistributionConstraints:
+    """Represents the distribution_constraints object of monitoring_config in constraints file."""
+
+    def __init__(self, categorical_drift_method: str = None):
+        self.categorical_drift_method = categorical_drift_method
+
+    @staticmethod
+    def valid_distribution_constraints(distribution_constraints):
+        """Checks whether distribution_constraints are valid or not."""
+
+        if not distribution_constraints:
+            return True
+
+        return DataQualityDistributionConstraints.valid_categorical_drift_method(
+            distribution_constraints.categorical_drift_method
+        )
+
+    @staticmethod
+    def valid_categorical_drift_method(categorical_drift_method):
+        """Checks whether categorical_drift_method is valid or not."""
+
+        if not categorical_drift_method:
+            return True
+
+        return categorical_drift_method in [CHI_SQUARED_METHOD, L_INFINITY_METHOD]
+
+
+class DataQualityMonitoringConfig:
+    """Represents monitoring_config object in constraints file."""
+
+    def __init__(self, distribution_constraints: DataQualityDistributionConstraints = None):
+        self.distribution_constraints = distribution_constraints
+
+    @staticmethod
+    def valid_monitoring_config(monitoring_config):
+        """Checks whether monitoring_config is valid or not."""
+
+        if not monitoring_config:
+            return True
+
+        return DataQualityDistributionConstraints.valid_distribution_constraints(
+            monitoring_config.distribution_constraints
+        )

src/sagemaker/model_monitor/model_monitoring.py

@@ -52,6 +52,7 @@
     MonitoringAlertActions,
     ModelDashboardIndicatorAction,
 )
+from sagemaker.model_monitor.data_quality_monitoring_config import DataQualityMonitoringConfig
 from sagemaker.model_monitor.dataset_format import MonitoringDatasetFormat
 from sagemaker.network import NetworkConfig
 from sagemaker.processing import Processor, ProcessingInput, ProcessingJob, ProcessingOutput
@@ -98,6 +99,7 @@
 _INFERENCE_ATTRIBUTE_ENV_NAME = "inference_attribute"
 _PROBABILITY_ATTRIBUTE_ENV_NAME = "probability_attribute"
 _PROBABILITY_THRESHOLD_ATTRIBUTE_ENV_NAME = "probability_threshold_attribute"
+_CATEGORICAL_DRIFT_METHOD_ENV_NAME = "categorical_drift_method"
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -1136,6 +1138,7 @@ def _generate_env_map(
         probability_attribute=None,
         ground_truth_attribute=None,
         probability_threshold_attribute=None,
+        categorical_drift_method=None,
     ):
         """Generate a list of environment variables from first-class parameters.
 
@@ -1157,6 +1160,9 @@ def _generate_env_map(
                 Only used for ModelQualityMonitor.
             probability_threshold_attribute (float): threshold to convert probabilities to binaries
                 Only used for ModelQualityMonitor.
+            categorical_drift_method (str): categorical_drift_method to override the
+                categorical_drift_method of global monitoring_config in constraints
+                suggested by Model Monitor container. Only used for DataQualityMonitor.
 
         Returns:
             dict: Dictionary of environment keys and values.
@@ -1206,6 +1212,9 @@ def _generate_env_map(
         if probability_threshold_attribute is not None:
             env[_PROBABILITY_THRESHOLD_ATTRIBUTE_ENV_NAME] = probability_threshold_attribute
 
+        if categorical_drift_method is not None:
+            env[_CATEGORICAL_DRIFT_METHOD_ENV_NAME] = categorical_drift_method
+
         return env
 
     @staticmethod
@@ -1647,6 +1656,7 @@ def suggest_baseline(
         wait=True,
         logs=True,
         job_name=None,
+        monitoring_config_override=None,
     ):
         """Suggest baselines for use with Amazon SageMaker Model Monitoring Schedules.
 
@@ -1666,12 +1676,18 @@ def suggest_baseline(
                 Only meaningful when wait is True (default: True).
             job_name (str): Processing job name. If not specified, the processor generates
                 a default job name, based on the image name and current timestamp.
-
+            monitoring_config_override (DataQualityMonitoringConfig): monitoring_config object to
+                override the global monitoring_config parameter of constraints suggested by
+                Model Monitor Container. If not specified, the values suggested by container is
+                set.
         Returns:
             sagemaker.processing.ProcessingJob: The ProcessingJob object representing the
                 baselining job.
 
         """
+        if not DataQualityMonitoringConfig.valid_monitoring_config(monitoring_config_override):
+            raise RuntimeError("Invalid value for monitoring_config_override.")
+
         self.latest_baselining_job_name = self._generate_baselining_job_name(job_name=job_name)
 
         normalized_baseline_dataset_input = self._upload_and_convert_to_processing_input(
@@ -1731,6 +1747,11 @@ def suggest_baseline(
 
         normalized_baseline_output = self._normalize_baseline_output(output_s3_uri=output_s3_uri)
 
+        categorical_drift_method = None
+        if monitoring_config_override and monitoring_config_override.distribution_constraints:
+            distribution_constraints = monitoring_config_override.distribution_constraints
+            categorical_drift_method = distribution_constraints.categorical_drift_method
+
         normalized_env = self._generate_env_map(
             env=self.env,
             dataset_format=dataset_format,
@@ -1739,6 +1760,7 @@ def suggest_baseline(
             dataset_source_container_path=baseline_dataset_container_path,
             record_preprocessor_script_container_path=record_preprocessor_script_container_path,
             post_processor_script_container_path=post_processor_script_container_path,
+            categorical_drift_method=categorical_drift_method,
         )
 
         baselining_processor = Processor(