diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py index 0686bc8f33..cda8bec8e0 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py @@ -362,6 +362,15 @@ def _generate_baseline_job_inputs(self): s3_input={ "s3_uri": self.quality_check_config.baseline_dataset, "local_path": baseline_dataset_des, + # Mirror the fields _upload_and_convert_to_processing_input sets on the + # non-pipeline-variable branch (issue #6206). s3_data_type is required by + # ProcessingS3Input (omitting it raised a ValidationError); s3_input_mode + # and s3_data_distribution_type must be set too, otherwise they stay the + # Unassigned() sentinel and break pipeline.definition() serialization + # (TypeError: Object of type Unassigned is not JSON serializable). + "s3_data_type": "S3Prefix", + "s3_input_mode": "File", + "s3_data_distribution_type": "FullyReplicated", }, ) else: diff --git a/sagemaker-mlops/tests/unit/workflow/test_quality_check_step.py b/sagemaker-mlops/tests/unit/workflow/test_quality_check_step.py index 2166340bb7..16dffce705 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_quality_check_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_quality_check_step.py @@ -14,10 +14,14 @@ from __future__ import absolute_import +from unittest.mock import Mock +from sagemaker.core.workflow.functions import Join +from sagemaker.core.workflow.parameters import ParameterString from sagemaker.mlops.workflow.quality_check_step import ( DataQualityCheckConfig, ModelQualityCheckConfig, + QualityCheckStep, ) @@ -36,3 +40,48 @@ def test_model_quality_check_config_init(): problem_type="BinaryClassification", ) assert config.problem_type == "BinaryClassification" + + +def test_generate_baseline_job_inputs_with_pipeline_variable_baseline_dataset(): + """Regression for #6206. + + When baseline_dataset is a pipeline variable, the ProcessingInput's s3_input must still + include the required s3_data_type; omitting it raised a pydantic ValidationError. + """ + baseline_dataset = Join( + on="/", + values=["s3:/", "my-bucket", ParameterString(name="EndpointName"), "baseline/data.parquet"], + ) + config = DataQualityCheckConfig( + baseline_dataset=baseline_dataset, + dataset_format={"parquet": {}}, + output_s3_uri="s3://my-bucket/output/", + ) + + step = object.__new__(QualityCheckStep) + step.quality_check_config = config + # Isolate the pipeline-variable baseline branch; the script inputs go through the monitor. + step._model_monitor = Mock() + step._model_monitor._upload_and_convert_to_processing_input.return_value = Mock() + + inputs = step._generate_baseline_job_inputs() + + baseline_input = inputs["baseline_dataset_input"] + s3_input = baseline_input.s3_input + assert s3_input.s3_data_type == "S3Prefix" + assert s3_input.s3_uri == baseline_dataset + # These must be concrete values, not the Unassigned() sentinel: the arguments serializer + # emits them into the pipeline definition, and json.dumps chokes on Unassigned (#6206). + import json + + from sagemaker.core.utils.utils import Unassigned + + assert not isinstance(s3_input.s3_input_mode, Unassigned) + assert not isinstance(s3_input.s3_data_distribution_type, Unassigned) + serialized = { + "S3Uri": "s3://resolved/at/runtime", # a pipeline var resolves to an expr; stub for json + "LocalPath": s3_input.local_path, + "S3DataType": getattr(s3_input, "s3_data_type", "S3Prefix"), + "S3InputMode": getattr(s3_input, "s3_input_mode", "File"), + } + json.dumps(serialized) # must not raise TypeError on an Unassigned sentinel