aws / aws/sagemaker-python-sdk
QualityCheckStep fails with ValidationError when baseline_dataset is a pipeline variable
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 1.3k
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 35
Description
**PySDK Version**
- [ ] PySDK V2 (2.x)
- [x] PySDK V3 (3.x)
**Describe the bug**
`QualityCheckStep._generate_baseline_job_inputs()` creates a `ProcessingInput` with an incomplete `s3_input` dict when `baseline_dataset` is a pipeline variable (e.g. `Join`, `ParameterString`). The dict is missing the required `s3_data_type` field, causing a Pydantic `ValidationError`.
The bug is in [`sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py` line 354](https://github.com/aws/sagemaker-python-sdk/blob/a24812666d258d9d19c9a82d0c9385ca0bd3ba01/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py#L354):
```python
if is_pipeline_variable(baseline_dataset):
baseline_dataset_input = ProcessingInput(
input_name=_BASELINE_DATASET_INPUT_NAME,
s3_input={
"s3_uri": self.quality_check_config.baseline_dataset,
"local_path": baseline_dataset_des,
}
)
```
`ProcessingS3Input` (from `sagemaker.core.shapes`) requires `s3_data_type` as a mandatory field with no default. The `else` branch correctly provides it via `_upload_and_convert_to_processing_input()`, but the pipeline variable branch does not.
**To reproduce**
```python
import boto3
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.workflow.functions import Join
from sagemaker.core.workflow.parameters import ParameterString
from sagemaker.core.workflow.pipeline_context import PipelineSession
from sagemaker.mlops.workflow.quality_check_step import DataQualityCheckConfig, QualityCheckStep
from sagemaker.mlops.workflow.check_job_config import CheckJobConfig
pipeline_session = PipelineSession(boto_session=boto3.Session())
param_endpoint_name = ParameterString(name="EndpointName")
# baseline_dataset is a pipeline variable — resolved at execution time
baseline_dataset_uri = Join(
on="/",
values=["s3:/", "my-bucket", param_endpoint_name, "baseline/dataset.parquet"],
)
quality_check_config = DataQualityCheckConfig(
baseline_dataset=baseline_dataset_uri,
dataset_format={"parquet": {}},
output_s3_uri="s3://my-bucket/output/",
)
check_job_config = CheckJobConfig(
role="arn:aws:iam::123456789012:role/SageMakerRole",
instance_count=1,
instance_type="ml.m5.xlarge",
sagemaker_session=pipeline_session,
)
# This raises ValidationError
step = QualityCheckStep(
name="compute-baseline",
quality_check_config=quality_check_config,
check_job_config=check_job_config,
skip_check=True,
register_new_baseline=True,
)
```
**Expected behavior**
`QualityCheckStep` should instantiate successfully when `baseline_dataset` is a pipeline variable. The fix is to include `s3_data_type` in the dict:
```python
if is_pipeline_variable(baseline_dataset):
baseline_dataset_input = ProcessingInput(
input_name=_BASELINE_DATASET_INPUT_NAME,
s3_input={
"s3_uri": self.quality_check_config.baseline_dataset,
"local_path": baseline_dataset_des,
"s3_data_type": "S3Prefix", # <-- missing
}
)
```
**Screenshots or logs**
```
ValidationError: 1 validation error for ProcessingInput
s3_input.s3_data_type
Field required [type=missing, input_value={'s3_uri': Join(on='/', v...baseline_dataset_input'}, input_type=dict]
For further information visit https://errors.pydantic.dev/2.13/v/missing
```
**System information**
- **SageMaker Python SDK version**: sagemaker 3.20.0, sagemaker-mlops 1.20.0, sagemaker-core 2.3.0
- **Framework name**: N/A (SageMaker Model Monitor)
- **Framework version**: N/A
- **Python version**: 3.13.5
- **CPU or GPU**: CPU
- **Custom Docker image (Y/N)**: N
**Additional context**
The bug is present on the latest `main` branch as well as all released versions of `sagemaker-mlops` (1.0–1.20.0). It only manifests when `baseline_dataset` is a pipeline variable (`Join`, `JsonGet`, `ParameterString`, etc.) — static string paths work fine because they take the `else` branch which uses `_upload_and_convert_to_processing_input()`.
Workaround: patch `ProcessingS3Input` to make `s3_data_type` optional before constructing the step:
```python
from sagemaker.core.shapes import ProcessingInput, ProcessingS3Input
ProcessingS3Input.model_fields["s3_data_type"].default = "S3Prefix"
ProcessingS3Input.model_rebuild(force=True)
ProcessingInput.model_rebuild(force=True)
```
Contributor guide
Research direction
Start at sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py around line 354 and reproduce the QualityCheckStep construction with a Join or ParameterString baseline_dataset. Compare the pipeline-variable branch with _upload_and_convert_to_processing_input(); done means the step instantiates successfully and a regression check covers the required ProcessingInput fields.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 1/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 88/100