diff --git a/sagemaker-core/src/sagemaker/core/model_monitor/clarify_model_monitoring.py b/sagemaker-core/src/sagemaker/core/model_monitor/clarify_model_monitoring.py index 07d646cbe8..e017c6a3b6 100644 --- a/sagemaker-core/src/sagemaker/core/model_monitor/clarify_model_monitoring.py +++ b/sagemaker-core/src/sagemaker/core/model_monitor/clarify_model_monitoring.py @@ -158,10 +158,10 @@ def list_executions(self): return [ ClarifyMonitoringExecution( sagemaker_session=execution.sagemaker_session, - job_name=execution.job_name, - inputs=execution.inputs, + job_name=execution.processing_job_name, + inputs=execution.processing_inputs, output=execution.output, - output_kms_key=execution.output_kms_key, + output_kms_key=execution.processing_output_config.kms_key_id or None, ) for execution in executions ] @@ -557,7 +557,8 @@ def suggest_baseline( self.latest_baselining_job_config = latest_baselining_job_config self.latest_baselining_job_name = baselining_job_name self.latest_baselining_job = ClarifyBaseliningJob( - processing_job=baselining_processor.latest_job + processing_job=baselining_processor.latest_job, + sagemaker_session=self.sagemaker_session, ) self.baselining_jobs.append(self.latest_baselining_job) @@ -1007,7 +1008,8 @@ def suggest_baseline( ) self.latest_baselining_job_name = baselining_job_name self.latest_baselining_job = ClarifyBaseliningJob( - processing_job=baselining_processor.latest_job + processing_job=baselining_processor.latest_job, + sagemaker_session=self.sagemaker_session, ) self.baselining_jobs.append(self.latest_baselining_job) @@ -1417,19 +1419,25 @@ class ClarifyBaseliningJob(mm.BaseliningJob): def __init__( self, processing_job, + sagemaker_session=None, ): """Initializes a ClarifyBaseliningJob that tracks a baselining job by suggest_baseline() Args: - processing_job (sagemaker.processing.ProcessingJob): The ProcessingJob used for + processing_job (sagemaker.core.resources.ProcessingJob): The ProcessingJob used for baselining instance. + sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object + which manages interactions with Amazon SageMaker APIs. The v3 ProcessingJob + resource does not carry a session, so callers should pass the monitor's. """ + output_config = processing_job.processing_output_config or None super(ClarifyBaseliningJob, self).__init__( - sagemaker_session=processing_job.sagemaker_session, - job_name=processing_job.job_name, - inputs=processing_job.inputs, - outputs=processing_job.outputs, - output_kms_key=processing_job.output_kms_key, + sagemaker_session=sagemaker_session + or getattr(processing_job, "sagemaker_session", None), + job_name=processing_job.processing_job_name, + inputs=processing_job.processing_inputs or [], + outputs=output_config.outputs if output_config else [], + output_kms_key=(output_config.kms_key_id or None) if output_config else None, ) def baseline_statistics(self, **_): diff --git a/sagemaker-core/tests/unit/model_monitor/test_clarify_model_monitoring.py b/sagemaker-core/tests/unit/model_monitor/test_clarify_model_monitoring.py index 607f76153e..ac20bd9859 100644 --- a/sagemaker-core/tests/unit/model_monitor/test_clarify_model_monitoring.py +++ b/sagemaker-core/tests/unit/model_monitor/test_clarify_model_monitoring.py @@ -20,7 +20,9 @@ ClarifyMonitoringExecution, ClarifyBaseliningConfig, BiasAnalysisConfig, + ClarifyBaseliningJob, ) +from sagemaker.core.model_monitor.model_monitoring import MonitoringExecution from sagemaker.core.clarify import BiasConfig, DataConfig, ModelConfig, ModelPredictedLabelConfig @@ -87,11 +89,7 @@ def test_list_executions(self, mock_from_arn, mock_list_executions, mock_session ] } - mock_execution = Mock() - mock_execution.sagemaker_session = mock_session - mock_execution.job_name = "test-job" - mock_execution.inputs = [] - mock_execution.output = ProcessingOutput( + output = ProcessingOutput( output_name="output", s3_output=ProcessingS3Output( s3_uri="s3://bucket/output", @@ -99,13 +97,22 @@ def test_list_executions(self, mock_from_arn, mock_list_executions, mock_session s3_upload_mode="EndOfJob", ), ) - mock_execution.output_kms_key = None - mock_from_arn.return_value = mock_execution + mock_from_arn.return_value = MonitoringExecution( + sagemaker_session=mock_session, + job_name="test-job", + inputs=[], + output=output, + output_kms_key="kms-key", + ) executions = monitor.list_executions() assert len(executions) == 1 assert isinstance(executions[0], ClarifyMonitoringExecution) + assert executions[0].processing_job_name == "test-job" + assert executions[0].sagemaker_session is mock_session + assert executions[0].output == output + assert executions[0].processing_output_config.kms_key_id == "kms-key" @patch("sagemaker.core.model_monitor.clarify_model_monitoring.boto_list_monitoring_executions") @patch("sagemaker.core.model_monitor.clarify_model_monitoring.logs_for_processing_job") @@ -285,6 +292,54 @@ def test_create_monitoring_schedule_success(self, mock_name, mock_upload, mock_s assert monitor.monitoring_schedule_name is not None +class TestClarifyBaseliningJob: + """Test cases for ClarifyBaseliningJob class""" + + def test_init_from_v3_processing_job(self, mock_session): + """Test ClarifyBaseliningJob reads the v3 ProcessingJob resource attributes""" + from sagemaker.core.resources import ProcessingJob + from sagemaker.core.shapes import ( + ProcessingInput, + ProcessingOutput, + ProcessingOutputConfig, + ProcessingS3Input, + ProcessingS3Output, + ) + + processing_input = ProcessingInput( + input_name="dataset", + s3_input=ProcessingS3Input( + s3_uri="s3://bucket/input", + local_path="/opt/ml/processing/input/data", + s3_data_type="S3Prefix", + s3_input_mode="File", + ), + ) + processing_output = ProcessingOutput( + output_name="analysis_result", + s3_output=ProcessingS3Output( + s3_uri="s3://bucket/output", + local_path="/opt/ml/processing/output", + s3_upload_mode="EndOfJob", + ), + ) + processing_job = ProcessingJob( + processing_job_name="baseline-job", + processing_inputs=[processing_input], + processing_output_config=ProcessingOutputConfig( + outputs=[processing_output], kms_key_id="kms-key" + ), + ) + + job = ClarifyBaseliningJob(processing_job=processing_job, sagemaker_session=mock_session) + + assert job.sagemaker_session is mock_session + assert job.job_name == "baseline-job" + assert job.inputs == [processing_input] + assert job.outputs == [processing_output] + assert job.output_kms_key == "kms-key" + + class TestClarifyBaseliningConfig: """Test cases for ClarifyBaseliningConfig"""