fix: set full s3_input fields on QualityCheckStep baseline_dataset pipeline-variable input (#6322)

When baseline_dataset is a pipeline variable, _generate_baseline_job_inputs built a
ProcessingInput whose s3_input dict omitted s3_data_type (required by ProcessingS3Input
-> ValidationError) and also left s3_input_mode / s3_data_distribution_type as the
Unassigned() sentinel, which breaks pipeline.definition() serialization with
'TypeError: Object of type Unassigned is not JSON serializable'. Set all three to match
what the non-pipeline-variable branch (_upload_and_convert_to_processing_input) produces.

Fixes #6206

Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>

Author: mohamedzeidan2021

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:

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