fix: baselining error handling (#370)

* fix: baselining error handling

* changelog

* test cases

Author: kukushking

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - updated `sagemaker-model-monitor` module and tested 4 types of monitoring end-to-end
 - added baseline generation step function to `sagemaker-model-monitor` module
 - updated lambda runtime and depdencies in `sagemaker-templates` module
+- fixed baseline error handling in `sagemaker-model-monitor` module
 
 ## v2.0.0
 

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

@@ -50,6 +50,7 @@ def __init__(
             memory_size=2048,
             environment={
                 "SAGEMAKER_ROLE_ARN": sagemaker_role_arn,
+                "LOG_LEVEL": "INFO",
             },
         )
 
@@ -131,6 +132,15 @@ def __init__(
                     sfn.Condition.string_equals("$.status", "COMPLETED"),
                     sfn.Succeed(self, "JobCompleted"),
                 )
+                .when(
+                    sfn.Condition.string_equals("$.status", "FAILED"),
+                    sfn.Fail(
+                        self,
+                        "JobFailed",
+                        cause="Baselining job failed",
+                        error="ProcessingJobFailed",
+                    ),
+                )
                 .otherwise(wait_task)
             )
         )

modules/sagemaker/sagemaker-model-monitoring/sagemaker_model_monitoring/lambda/baselining_handler.py

@@ -1,4 +1,6 @@
 # mypy: disable-error-code="attr-defined,no-untyped-call,assignment,call-arg,no-any-return,union-attr"
+import json
+import logging
 import os
 from typing import Any, Dict
 
@@ -13,15 +15,22 @@
 )
 from sagemaker.model_monitor.dataset_format import DatasetFormat
 
+logger = logging.getLogger()
+logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
+log = logging.getLogger(__name__)
+
 SAGEMAKER_ROLE_ARN = os.environ["SAGEMAKER_ROLE_ARN"]
 
 sagemaker = boto3.client("sagemaker")
 
 
 def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
     """Lambda handler for SageMaker Model Monitor baselining jobs."""
+    logger.info(f"Received event: {json.dumps(event)}")
+
     try:
         action = event.get("action", "start")
+        logger.info(f"Action: {action}")
 
         if action == "start":
             return start_baselining_job(event)
@@ -30,14 +39,18 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
 
         return {"statusCode": 200}
     except KeyError as e:
+        logger.error(f"Missing required parameter: {e}")
         return {"statusCode": 400, "error": f"Missing required parameter: {e}"}
     except Exception as e:
+        logger.exception(f"Internal error: {str(e)}")
         return {"statusCode": 500, "error": f"Internal error: {str(e)}"}
 
 
 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", {}))
+    logger.info(f"Data quality baseline params: {params.model_dump_json()}")
+
     monitor = DefaultModelMonitor(
         role=SAGEMAKER_ROLE_ARN,
         instance_count=params.instance_count,
@@ -52,12 +65,16 @@ def start_data_quality_baseline(training_data_uri: str, baseline_output_uri: str
         wait=False,
         logs=False,
     )
-    return monitor.latest_baselining_job.job_name
+    job_name = monitor.latest_baselining_job.job_name
+    logger.info(f"Started data quality baseline job: {job_name}")
+    return job_name
 
 
 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", {}))
+    logger.info(f"Model quality baseline params: {params.model_dump_json()}")
+
     monitor = ModelQualityMonitor(
         role=SAGEMAKER_ROLE_ARN,
         instance_count=params.instance_count,
@@ -76,14 +93,18 @@ def start_model_quality_baseline(training_data_uri: str, baseline_output_uri: st
         wait=False,
         logs=False,
     )
-    return monitor.latest_baselining_job.job_name
+    job_name = monitor.latest_baselining_job.job_name
+    logger.info(f"Started model quality baseline job: {job_name}")
+    return job_name
 
 
 def start_model_bias_baseline(
     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", {}))
+    logger.info(f"Model bias baseline params: {params.model_dump_json()}")
+
     monitor = ModelBiasMonitor(role=SAGEMAKER_ROLE_ARN, max_runtime_in_seconds=params.max_runtime_seconds)
 
     model_bias_data_config = DataConfig(
@@ -114,14 +135,18 @@ def start_model_bias_baseline(
         wait=False,
         logs=False,
     )
-    return monitor.latest_baselining_job.job_name
+    job_name = monitor.latest_baselining_job.job_name
+    logger.info(f"Started model bias baseline job: {job_name}")
+    return job_name
 
 
 def start_model_explainability_baseline(
     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", {}))
+    logger.info(f"Model explainability baseline params: {params.model_dump_json()}")
+
     monitor = ModelExplainabilityMonitor(role=SAGEMAKER_ROLE_ARN, max_runtime_in_seconds=params.max_runtime_seconds)
 
     model_explainability_data_config = DataConfig(
@@ -149,7 +174,9 @@ def start_model_explainability_baseline(
         model_config=model_config,
         explainability_config=shap_config,
     )
-    return monitor.latest_baselining_job.job_name
+    job_name = monitor.latest_baselining_job.job_name
+    logger.info(f"Started model explainability baseline job: {job_name}")
+    return job_name
 
 
 def start_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
@@ -159,7 +186,13 @@ def start_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
         endpoint_name = event["endpoint_name"]
         training_data_uri = event["training_data_uri"]
         baseline_output_uri = event["baseline_output_uri"]
+
+        logger.info(
+            f"Starting baselining job - Monitor type: {monitor_type}, Endpoint: {endpoint_name}, "
+            f"Training URI: {training_data_uri}, Output URI: {baseline_output_uri}"
+        )
     except KeyError as e:
+        logger.error(f"Missing required parameter: {e}")
         return {"statusCode": 400, "error": f"Missing required parameter: {e}"}
 
     try:
@@ -173,29 +206,48 @@ def start_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
         elif monitor_type == "model_explainability":
             job_name = start_model_explainability_baseline(training_data_uri, baseline_output_uri, endpoint_name, event)
         else:
+            logger.error(f"Unsupported monitor_type: {monitor_type}")
             return {"statusCode": 400, "error": f"Unsupported monitor_type: {monitor_type}"}
 
+        logger.info(f"Successfully started baselining job: {job_name}")
         return {"statusCode": 200, "job_name": job_name, "status": "STARTED"}
     except Exception as e:
+        logger.exception(f"Failed to start baseline job: {str(e)}")
         return {"statusCode": 500, "error": f"Failed to start baseline job: {str(e)}"}
 
 
 def check_baselining_job(event: Dict[str, Any]) -> Dict[str, Any]:
     """Check the status of a baselining processing job."""
     try:
         job_name = event["job_name"]
+        logger.info(f"Checking status for job: {job_name}")
 
         response = sagemaker.describe_processing_job(ProcessingJobName=job_name)
         status = response["ProcessingJobStatus"]
-
-        return {
-            "statusCode": 200,
-            "job_name": job_name,
-            "status": "COMPLETED" if status == "Completed" else "IN_PROGRESS",
-        }
+        logger.info(f"Job {job_name} status: {status}")
+
+        # Map SageMaker processing job statuses to simplified states
+        if status == "Completed":
+            logger.info(f"Job {job_name} completed successfully")
+            return {"statusCode": 200, "job_name": job_name, "status": "COMPLETED"}
+        elif status in ["Failed", "Stopped", "Stopping"]:
+            failure_reason = response.get("FailureReason", "Unknown failure reason")
+            logger.error(f"Job {job_name} failed: {failure_reason}")
+            return {
+                "statusCode": 200,
+                "job_name": job_name,
+                "status": "FAILED",
+                "failure_reason": failure_reason,
+            }
+        else:  # InProgress
+            logger.info(f"Job {job_name} still in progress")
+            return {"statusCode": 200, "job_name": job_name, "status": "IN_PROGRESS"}
     except KeyError as e:
+        logger.error(f"Missing required parameter: {e}")
         return {"statusCode": 400, "error": f"Missing required parameter: {e}"}
     except sagemaker.exceptions.ResourceNotFound:
+        logger.error(f"Processing job not found: {event.get('job_name', 'unknown')}")
         return {"statusCode": 404, "error": f"Processing job not found: {event.get('job_name', 'unknown')}"}
     except Exception as e:
+        logger.exception(f"Failed to check job status: {str(e)}")
         return {"statusCode": 500, "error": f"Failed to check job status: {str(e)}"}

modules/sagemaker/sagemaker-model-monitoring/tests/test_baselining_handler.py

@@ -84,31 +84,29 @@ def test_lambda_handler_check_action(mock_sagemaker_modules, mock_env, mock_sage
     assert result["job_name"] == "test-job"
 
 
-def test_check_baselining_job_completed(mock_sagemaker_modules, mock_env, mock_sagemaker_client):
-    """Test check_baselining_job with completed job."""
+@pytest.mark.parametrize(
+    "processing_job_status,expected_status",
+    [
+        ("Completed", "COMPLETED"),
+        ("Failed", "FAILED"),
+        ("InProgress", "IN_PROGRESS"),
+        ("Stopping", "FAILED"),
+        ("Stopped", "FAILED"),
+    ],
+)
+def test_check_baselining_job_status(
+    mock_sagemaker_modules, mock_env, mock_sagemaker_client, processing_job_status, expected_status
+):
+    """Test check_baselining_job with various job statuses."""
     check_baselining_job, lambda_handler = mock_sagemaker_modules
     event = {"job_name": "test-job"}
 
-    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": "Completed"}
-
-    result = check_baselining_job(event)
-
-    assert result["statusCode"] == 200
-    assert result["status"] == "COMPLETED"
-    assert result["job_name"] == "test-job"
-
-
-def test_check_baselining_job_in_progress(mock_sagemaker_modules, mock_env, mock_sagemaker_client):
-    """Test check_baselining_job with in-progress job."""
-    check_baselining_job, lambda_handler = mock_sagemaker_modules
-    event = {"job_name": "test-job"}
-
-    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": "InProgress"}
+    mock_sagemaker_client.describe_processing_job.return_value = {"ProcessingJobStatus": processing_job_status}
 
     result = check_baselining_job(event)
 
     assert result["statusCode"] == 200
-    assert result["status"] == "IN_PROGRESS"
+    assert result["status"] == expected_status
     assert result["job_name"] == "test-job"
 
 
@@ -122,3 +120,48 @@ def test_check_baselining_job_missing_job_name(mock_sagemaker_modules, mock_env)
     assert result["statusCode"] == 400
     assert "error" in result
     assert "Missing required parameter" in result["error"]
+
+
+@pytest.mark.parametrize(
+    "monitor_type,monitor_class_path,expected_job_name",
+    [
+        ("data_quality", "baselining_handler.DefaultModelMonitor", "data-quality-job-123"),
+        ("model_quality", "baselining_handler.ModelQualityMonitor", "model-quality-job-456"),
+        ("model_bias", "baselining_handler.ModelBiasMonitor", "model-bias-job-789"),
+        ("model_explainability", "baselining_handler.ModelExplainabilityMonitor", "model-explainability-job-012"),
+    ],
+)
+def test_start_baselining_job_types(
+    mock_sagemaker_modules, mock_env, monitor_type, monitor_class_path, expected_job_name
+):
+    """Test starting different types of baselining jobs."""
+    check_baselining_job, lambda_handler = mock_sagemaker_modules
+
+    event = {
+        "action": "start",
+        "monitor_type": monitor_type,
+        "endpoint_name": "test-endpoint",
+        "training_data_uri": "s3://bucket/training-data",
+        "baseline_output_uri": "s3://bucket/baseline-output",
+    }
+
+    with patch(monitor_class_path) as mock_monitor_class:
+        mock_monitor_instance = MagicMock()
+        mock_monitor_class.return_value = mock_monitor_instance
+
+        # Mock the baselining job
+        mock_baselining_job = MagicMock()
+        mock_baselining_job.job_name = expected_job_name
+        mock_monitor_instance.latest_baselining_job = mock_baselining_job
+
+        result = lambda_handler(event, None)
+
+        assert result["statusCode"] == 200
+        assert result["status"] == "STARTED"
+        assert result["job_name"] == expected_job_name
+
+        # Verify the monitor was instantiated
+        mock_monitor_class.assert_called_once()
+
+        # Verify suggest_baseline was called
+        mock_monitor_instance.suggest_baseline.assert_called_once()