chore: Add docstrings to functions and add check for the case where instead of endpoint name, the user provides an EndpointInput object

Author: Sushanth Sathish Kumar

src/sagemaker/dashboard/dashboard_variables.py

@@ -20,9 +20,33 @@
 import json
 
 class DashboardVariable:
+    """
+    Represents a dashboard variable used for dynamic configuration in CloudWatch Dashboards.
+
+    Attributes:
+        variable_type (str): Type of dashboard variable ('property' or 'pattern').
+        variable_property (str): Property affected by the variable, such as a JSON property or metric dimension.
+        inputType (str): Type of input field ('input', 'select', or 'radio') for user interaction.
+        id (str): Identifier for the variable, up to 32 characters.
+        label (str): Label displayed for the input field (optional, defaults based on context).
+        search (str): Metric search expression to populate input options (required for 'select' or 'radio').
+        populateFrom (str): Dimension name used to populate input options from search results.
+    """
     def __init__(
         self, variable_type, variable_property, inputType, variable_id, label, search, populateFrom
     ):
+        """
+        Initializes a DashboardVariable instance.
+
+        Args:
+            variable_type (str): Type of dashboard variable ('property' or 'pattern').
+            variable_property (str): Property affected by the variable, such as a JSON property or metric dimension.
+            inputType (str): Type of input field ('input', 'select', or 'radio') for user interaction.
+            variable_id (str): Identifier for the variable, up to 32 characters.
+            label (str, optional): Label displayed for the input field (default is None).
+            search (str, optional): Metric search expression to populate input options (required for 'select' or 'radio').
+            populateFrom (str, optional): Dimension name used to populate input options from search results.
+        """
         self.variable_type = variable_type
         self.variable_property = variable_property
         self.inputType = inputType
@@ -32,6 +56,12 @@ def __init__(
         self.populateFrom = populateFrom
 
     def to_dict(self):
+        """
+        Converts DashboardVariable instance to a dictionary representation.
+
+        Returns:
+            dict: Dictionary containing variable properties suitable for JSON serialization.
+        """
         variable_properties_dict = {}
         if self.variable_type is not None:
             variable_properties_dict["type"] = self.variable_type
@@ -50,5 +80,11 @@ def to_dict(self):
         return variable_properties_dict
 
     def to_json(self):
+        """
+        Converts DashboardVariable instance to a JSON string.
+
+        Returns:
+            str: JSON string representation of the variable properties.
+        """
         json.dumps(self.to_dict(), indent=4)
 

src/sagemaker/dashboard/dashboard_widgets.py

@@ -20,6 +20,18 @@
 import json 
 
 class DashboardWidgetProperties:
+    """
+    Represents properties of a dashboard widget used for metrics in CloudWatch.
+
+    Attributes:
+        view (str): Type of visualization ('timeSeries', 'singleValue', 'gauge', 'bar', 'pie', 'table').
+        stacked (bool): Whether to display the graph as stacked lines (applies to 'timeSeries' view).
+        metrics (list): Array of metrics configurations for the widget.
+        region (str): Region associated with the metrics.
+        period (int): Period in seconds for data points on the graph.
+        title (str): Title displayed for the graph or number (optional).
+        markdown (str): Markdown content to display within the widget (optional).
+    """
     def __init__(
         self,
         view=None,
@@ -30,6 +42,18 @@ def __init__(
         title=None,
         markdown=None,
     ):
+        """
+        Initializes DashboardWidgetProperties instance.
+
+        Args:
+            view (str, optional): Type of visualization ('timeSeries', 'singleValue', 'gauge', 'bar', 'pie', 'table').
+            stacked (bool, optional): Whether to display the graph as stacked lines (applies to 'timeSeries' view).
+            metrics (list, optional): Array of metrics configurations for the widget.
+            region (str, optional): Region associated with the metrics.
+            period (int, optional): Period in seconds for data points on the graph.
+            title (str, optional): Title displayed for the graph or number.
+            markdown (str, optional): Markdown content to display within the widget.
+        """
         self.view = view
         self.stacked = stacked
         self.metrics = metrics
@@ -39,6 +63,12 @@ def __init__(
         self.markdown = markdown
 
     def to_dict(self):
+        """
+        Converts DashboardWidgetProperties instance to a dictionary representation.
+
+        Returns:
+            dict: Dictionary containing widget properties suitable for JSON serialization.
+        """
         widget_properties_dict = {}
         if self.view is not None:
             widget_properties_dict["view"] = self.view
@@ -57,11 +87,35 @@ def to_dict(self):
         return widget_properties_dict
 
     def to_json(self):
+        """
+        Converts DashboardWidgetProperties instance to a JSON string.
+
+        Returns:
+            str: JSON string representation of the widget properties.
+        """
         json.dumps(self.to_dict(), indent=4)
 
 
 class DashboardWidget:
+    """
+    Represents a widget in a CloudWatch dashboard.
+
+    Attributes:
+        height (int): Height of the widget.
+        width (int): Width of the widget.
+        type (str): Type of the widget.
+        properties (DashboardWidgetProperties): Properties specific to the widget type.
+    """
     def __init__(self, height, width, widget_type, properties=None):
+        """
+        Initializes DashboardWidget instance.
+
+        Args:
+            height (int): Height of the widget.
+            width (int): Width of the widget.
+            widget_type (str): Type of the widget.
+            properties (DashboardWidgetProperties, optional): Properties specific to the widget type.
+        """
         self.height = height
         self.width = width
         self.type = widget_type
@@ -72,6 +126,12 @@ def __init__(self, height, width, widget_type, properties=None):
         )
 
     def to_dict(self):
+        """
+        Converts DashboardWidget instance to a dictionary representation.
+
+        Returns:
+            dict: Dictionary containing widget attributes suitable for JSON serialization.
+        """
         return {
             "height": self.height,
             "width": self.width,
@@ -80,4 +140,10 @@ def to_dict(self):
         }
 
     def to_json(self):
+        """
+        Converts DashboardWidget instance to a JSON string.
+
+        Returns:
+            str: JSON string representation of the widget attributes.
+        """
         return json.dumps(self.to_dict(), indent=4)

src/sagemaker/dashboard/data_quality_dashboard.py

@@ -17,6 +17,7 @@
 import json 
 from sagemaker.dashboard.dashboard_variables import DashboardVariable
 from sagemaker.dashboard.dashboard_widgets import DashboardWidget, DashboardWidgetProperties
+from sagemaker.model_monitor.model_monitoring import EndpointInput
 
 class AutomaticDataQualityDashboard:
     DATA_QUALITY_METRICS_ENDPOINT_NAMESPACE = (
@@ -27,7 +28,11 @@ class AutomaticDataQualityDashboard:
     )
 
     def __init__(self, endpoint_name, monitoring_schedule_name, batch_transform_input, region_name):
-        self.endpoint = endpoint_name
+        if type(endpoint_name) == EndpointInput: 
+            self.endpoint = endpoint_name.endpoint_name
+        else:
+            self.endpoint = endpoint_name
+            
         self.monitoring_schedule = monitoring_schedule_name
         self.batch_transform = batch_transform_input
         self.region = region_name

src/sagemaker/dashboard/model_quality_dashboard.py

@@ -1,7 +1,31 @@
 import json
 from sagemaker.dashboard.dashboard_widgets import DashboardWidget, DashboardWidgetProperties
-
+from sagemaker.model_monitor import EndpointInput
 class AutomaticModelQualityDashboard:   
+    """
+    Represents a dashboard for automatic model quality metrics in Amazon SageMaker.
+
+    Attributes:
+        MODEL_QUALITY_METRICS_ENDPOINT_NAMESPACE (str): Namespace for model metrics at endpoint level.
+        MODEL_QUALITY_METRICS_BATCH_NAMESPACE (str): Namespace for model metrics at batch transform level.
+        REGRESSION_MODEL_QUALITY_METRICS (list): List of regression model quality metrics and their graphs.
+        BINARY_CLASSIFICATION_MODEL_QUALITY_METRICS (list): List of binary classification model quality metrics and their graphs.
+        MULTICLASS_CLASSIFICATION_MODEL_QUALITY_METRICS (list): List of multiclass classification model quality metrics and their graphs.
+
+    Methods:
+        __init__(self, endpoint_name, monitoring_schedule_name, batch_transform_input, problem_type, region_name):
+            Initializes an AutomaticModelQualityDashboard instance.
+
+        _generate_widgets(self):
+            Generates widgets based on the specified problem type and metrics.
+
+        to_dict(self):
+            Converts the dashboard instance to a dictionary representation.
+
+        to_json(self):
+            Converts the dashboard instance to a JSON string.
+    """
+    
     MODEL_QUALITY_METRICS_ENDPOINT_NAMESPACE = (
         "{aws/sagemaker/Endpoints/model-metrics,Endpoint,MonitoringSchedule}"
     )
@@ -60,7 +84,21 @@ class AutomaticModelQualityDashboard:
     ]
 
     def __init__(self, endpoint_name, monitoring_schedule_name, batch_transform_input, problem_type, region_name):
-        self.endpoint = endpoint_name
+        """
+        Initializes an AutomaticModelQualityDashboard instance.
+
+        Args:
+            endpoint_name (str): Name of the SageMaker endpoint.
+            monitoring_schedule_name (str): Name of the monitoring schedule.
+            batch_transform_input (str): Batch transform input (can be None).
+            problem_type (str): Type of problem ('Regression', 'BinaryClassification', or 'MulticlassClassification').
+            region_name (str): AWS region name.
+        """
+        if type(endpoint_name) == EndpointInput: 
+            self.endpoint = endpoint_name.endpoint_name
+        else:
+            self.endpoint = endpoint_name
+        
         self.monitoring_schedule = monitoring_schedule_name
         self.batch_transform = batch_transform_input
         self.region = region_name
@@ -72,6 +110,12 @@ def __init__(self, endpoint_name, monitoring_schedule_name, batch_transform_inpu
 
 
     def _generate_widgets(self):
+        """
+        Generates widgets based on the specified problem type and metrics.
+
+        Returns:
+            list: List of DashboardWidget instances representing each metric graph.
+        """
         list_of_widgets = []
         metrics_to_graph = None
         if (self.problem_type == "Regression"):
@@ -82,11 +126,12 @@ def _generate_widgets(self):
             metrics_to_graph = AutomaticModelQualityDashboard.MULTICLASS_CLASSIFICATION_MODEL_QUALITY_METRICS
         else:
             raise ValueError("Parameter problem_type is invalid. Valid options are Regression, BinaryClassification, or MulticlassClassification.")
-            
+        
         for graphs_per_line in metrics_to_graph:
             for graph in graphs_per_line:
                 graph_title = graph[0]
                 graph_metrics = graph[1] 
+                metrics_string = " OR ".join(graph_metrics)
                 if self.batch_transform is not None:
                     graph_properties = DashboardWidgetProperties(
                         view="timeSeries",
@@ -96,7 +141,7 @@ def _generate_widgets(self):
                                 {
                                     "expression": (
                                         f"SEARCH( '{AutomaticModelQualityDashboard.MODEL_QUALITY_METRICS_BATCH_NAMESPACE} "
-                                        f"{" OR ".join(graph_metrics)}"
+                                        f"{metrics_string} "
                                         f"MonitoringSchedule=\"{self.monitoring_schedule}\" ', "
                                         f"'Average')"
                                     )
@@ -115,7 +160,7 @@ def _generate_widgets(self):
                                 {
                                     "expression": (
                                         f"SEARCH( '{AutomaticModelQualityDashboard.MODEL_QUALITY_METRICS_ENDPOINT_NAMESPACE} "
-                                        f"{" OR ".join(graph_metrics)}"
+                                        f"{metrics_string} "
                                         f"Endpoint=\"{self.endpoint}\" "
                                         f"MonitoringSchedule=\"{self.monitoring_schedule}\" ', "
                                         f"'Average')"
@@ -128,16 +173,28 @@ def _generate_widgets(self):
                     )
                 list_of_widgets.append(
                     DashboardWidget(
-                        height=8, width=24//len(graph_metrics), widget_type="metric", properties=graph_properties
+                        height=8, width=24//len(graphs_per_line), widget_type="metric", properties=graph_properties
                     )
                 )
 
         return list_of_widgets
 
     def to_dict(self):
+        """
+        Converts the AutomaticModelQualityDashboard instance to a dictionary representation.
+
+        Returns:
+            dict: Dictionary containing the dashboard widgets.
+        """
         return {
             "widgets": [widget.to_dict() for widget in self.dashboard["widgets"]],
         }
 
     def to_json(self):
+        """
+        Converts the AutomaticModelQualityDashboard instance to a JSON string.
+
+        Returns:
+            str: JSON string representation of the dashboard widgets.
+        """
         return json.dumps(self.to_dict(), indent=4)

src/sagemaker/model_monitor/model_monitoring.py

@@ -3290,10 +3290,11 @@ def create_monitoring_schedule(
             cw_client = self.sagemaker_session.boto_session.client("cloudwatch")
             cw_client.put_dashboard(
                 DashboardName=dashboard_name,
-                DashboardBody=AutomaticDataQualityDashboard(
+                DashboardBody=AutomaticModelQualityDashboard(
                     endpoint_name=endpoint_input,
                     monitoring_schedule_name=monitor_schedule_name,
                     batch_transform_input=batch_transform_input,
+                    problem_type=problem_type,
                     region_name=self.sagemaker_session.boto_region_name,
                 ).to_json(),
             )