fixes and docs

Author: kukushking

modules/sagemaker/sagemaker-model-monitoring/README.md

@@ -13,20 +13,7 @@ Available monitoring types:
 * Model Bias
 * Model Explainability
 
-### Baseline Generation
-
-The module includes an optional automated baseline generation feature that creates baseline statistics and constraints for your monitoring jobs. When you provide training data, the module will:
-
-1. Deploy a Step Functions state machine that orchestrates baseline generation
-2. Run SageMaker Processing jobs to analyze your training data
-3. Generate baseline statistics and constraints files
-4. Store the baseline artifacts in your specified S3 location
-5. Schedule automatic baseline regeneration (default: daily at 2 AM UTC)
-
-The baseline generation uses a Lambda function deployed as a Docker container image to handle the SageMaker SDK dependencies efficiently.
-
-Note that updating parameters will require replacing resources. Deployments may be delayed until any
-running monitoring jobs complete (and the resources can be destroyed).
+The module includes an optional automated baseline generation feature that creates baseline statistics and constraints for your monitoring jobs.
 
 ### Architecture
 
@@ -58,6 +45,22 @@ running monitoring jobs complete (and the resources can be destroyed).
     - **CloudWatch Metrics**: Some monitoring types emit metrics (e.g., data quality drift)
     - **CloudWatch Alarms**: Can be configured based on emitted metrics for automated alerting
 
+### Baseline Generation
+
+The module includes an optional automated baseline generation feature that creates baseline statistics and constraints for your monitoring jobs. 
+
+![SageMaker Model Monitoring Baseline Generation](docs/_static/sagemaker-model-monitoring-baseline.png "SageMaker Model Monitoring Baseline Generation")
+
+When you provide training data, the module will:
+
+1. Deploy a Step Functions state machine that orchestrates baseline generation
+2. Run SageMaker Processing jobs to analyze your training data
+3. Generate baseline statistics and constraints files
+4. Store the baseline artifacts in your specified S3 location
+5. Schedule automatic baseline regeneration (default: daily at 2 AM UTC)
+
+## Inputs/Outputs
+
 ### Input Parameters
 
 #### Required

modules/sagemaker/sagemaker-model-monitoring/docs/_static/sagemaker-model-monitoring-baseline.png

@@ -108,6 +108,7 @@ def __init__(
                     "baseline_output_uri": baseline_output_data_s3_uri,
                 }
             ),
+            output_path="$.Payload",
         )
 
         wait_task = sfn.Wait(self, "WaitForBaselining", time=sfn.WaitTime.duration(Duration.minutes(2)))
@@ -116,7 +117,8 @@ def __init__(
             self,
             "CheckBaselining",
             lambda_function=baselining_lambda,
-            payload=sfn.TaskInput.from_object({"action": "check", "job_name.$": "$.Payload.job_name"}),
+            payload=sfn.TaskInput.from_object({"action": "check", "job_name.$": "$.job_name"}),
+            output_path="$.Payload",
         )
 
         # State machine
@@ -126,7 +128,7 @@ def __init__(
             .next(
                 sfn.Choice(self, "IsJobComplete")
                 .when(
-                    sfn.Condition.string_equals("$.Payload.status", "COMPLETED"),
+                    sfn.Condition.string_equals("$.status", "COMPLETED"),
                     sfn.Succeed(self, "JobCompleted"),
                 )
                 .otherwise(wait_task)

modules/sagemaker/sagemaker-model-monitoring/sagemaker_model_monitoring/baselining_construct.py

@@ -1,6 +1,5 @@
 # mypy: disable-error-code="attr-defined,no-untyped-call,assignment,call-arg"
 import os
-from datetime import datetime
 from typing import Any, Dict
 
 import boto3
@@ -36,9 +35,7 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
         return {"statusCode": 500, "error": f"Internal error: {str(e)}"}
 
 
-def start_data_quality_baseline(
-    job_name: str, training_data_uri: str, baseline_output_uri: str, event: Dict[str, Any]
-) -> None:
+def start_data_quality_baseline(training_data_uri: str, baseline_output_uri: str, event: Dict[str, Any]) -> str:
     """Start data quality baseline job."""
     params = DataQualityParams(**event.get("data_quality_params", {}))
     monitor = DefaultModelMonitor(
@@ -55,11 +52,10 @@ def start_data_quality_baseline(
         wait=False,
         logs=False,
     )
+    return monitor.latest_baselining_job.job_name
 
 
-def start_model_quality_baseline(
-    job_name: str, training_data_uri: str, baseline_output_uri: str, event: Dict[str, Any]
-) -> None:
+def start_model_quality_baseline(training_data_uri: str, baseline_output_uri: str, event: Dict[str, Any]) -> str:
     """Start model quality baseline job."""
     params = ModelQualityParams(**event.get("model_quality_params", {}))
     monitor = ModelQualityMonitor(
@@ -70,7 +66,6 @@ def start_model_quality_baseline(
         max_runtime_in_seconds=params.max_runtime_seconds,
     )
     monitor.suggest_baseline(
-        job_name=job_name,
         baseline_dataset=training_data_uri,
         dataset_format=DatasetFormat.csv(header=True),
         output_s3_uri=baseline_output_uri,
@@ -81,11 +76,12 @@ def start_model_quality_baseline(
         wait=False,
         logs=False,
     )
+    return monitor.latest_baselining_job.job_name
 
 
 def start_model_bias_baseline(
-    job_name: str, training_data_uri: str, baseline_output_uri: str, endpoint_name: str, event: Dict[str, Any]
-) -> None:
+    training_data_uri: str, baseline_output_uri: str, endpoint_name: str, event: Dict[str, Any]
+) -> str:
     """Start model bias baseline job."""
     params = ModelBiasParams(**event.get("model_bias_params", {}))
     monitor = ModelBiasMonitor(role=SAGEMAKER_ROLE_ARN, max_runtime_in_seconds=params.max_runtime_seconds)
@@ -118,11 +114,12 @@ def start_model_bias_baseline(
         wait=False,
         logs=False,
     )
+    return monitor.latest_baselining_job.job_name
 
 
 def start_model_explainability_baseline(
-    job_name: str, training_data_uri: str, baseline_output_uri: str, endpoint_name: str, event: Dict[str, Any]
-) -> None:
+    training_data_uri: str, baseline_output_uri: str, endpoint_name: str, event: Dict[str, Any]
+) -> str:
     """Start model explainability baseline job."""
     params = ModelExplainabilityParams(**event.get("model_explainability_params", {}))
     monitor = ModelExplainabilityMonitor(role=SAGEMAKER_ROLE_ARN, max_runtime_in_seconds=params.max_runtime_seconds)
@@ -152,6 +149,7 @@ def start_model_explainability_baseline(
         model_config=model_config,
         explainability_config=shap_config,
     )
+    return monitor.latest_baselining_job.job_name
 
 
 def start_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
@@ -164,18 +162,16 @@ def start_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
     except KeyError as e:
         return {"statusCode": 400, "error": f"Missing required parameter: {e}"}
 
-    timestamp = datetime.utcnow().strftime("%Y%m%d-%H%M%S")
-    job_name = f"baseline-{monitor_type}-{timestamp}"
-
     try:
+        job_name = None
         if monitor_type == "data_quality":
-            start_data_quality_baseline(job_name, training_data_uri, baseline_output_uri, event)
+            job_name = start_data_quality_baseline(training_data_uri, baseline_output_uri, event)
         elif monitor_type == "model_quality":
-            start_model_quality_baseline(job_name, training_data_uri, baseline_output_uri, event)
+            job_name = start_model_quality_baseline(training_data_uri, baseline_output_uri, event)
         elif monitor_type == "model_bias":
-            start_model_bias_baseline(job_name, training_data_uri, baseline_output_uri, endpoint_name, event)
+            job_name = start_model_bias_baseline(training_data_uri, baseline_output_uri, endpoint_name, event)
         elif monitor_type == "model_explainability":
-            start_model_explainability_baseline(job_name, training_data_uri, baseline_output_uri, endpoint_name, event)
+            job_name = start_model_explainability_baseline(training_data_uri, baseline_output_uri, endpoint_name, event)
         else:
             return {"statusCode": 400, "error": f"Unsupported monitor_type: {monitor_type}"}