Skip to content

fix: put the check steps' output KMS key on ProcessingOutputConfig and send CheckJobConfig tags #6497

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mohamedzeidan2021 merged 1 commit into from
Oct 19, 2026
Merged

fix: put the check steps' output KMS key on ProcessingOutputConfig and send CheckJobConfig tags #6497

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 132 additions & 0 deletions sagemaker-mlops/src/sagemaker/mlops/workflow/_check_job_request.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Shared CreateProcessingJob request builder for the check steps.

`QualityCheckStep` and `ClarifyCheckStep` both expose their baselining job as a
`CreateProcessingJob` request dict. The two requests are built identically apart
from the inputs, the output and the default job name, so the shaping lives here
to keep the field names in one place and in sync with the SageMaker API.
"""

from __future__ import absolute_import

from typing import List

from sagemaker.core.helper.pipeline_variable import RequestType
from sagemaker.core.shapes import ProcessingInput, ProcessingOutput

_DEFAULT_MAX_RUNTIME_IN_SECONDS = 86400
_DEFAULT_VOLUME_SIZE_IN_GB = 30


def _to_processing_input_dict(processing_input: ProcessingInput) -> dict:
"""Converts a ProcessingInput into its `ProcessingInputs` entry.

Args:
processing_input (ProcessingInput): The input to convert.

Returns:
dict: A `ProcessingInput` API entry.
"""
input_dict = {
"InputName": processing_input.input_name,
}
if processing_input.s3_input:
input_dict["S3Input"] = {
"S3Uri": processing_input.s3_input.s3_uri,
"LocalPath": processing_input.s3_input.local_path,
"S3DataType": getattr(processing_input.s3_input, "s3_data_type", "S3Prefix"),
"S3InputMode": getattr(processing_input.s3_input, "s3_input_mode", "File"),
}
return input_dict


def _build_check_job_request(
processing_inputs: List[ProcessingInput],
processing_output: ProcessingOutput,
processor,
check_job_config,
default_job_name: str,
) -> RequestType:
"""Builds the CreateProcessingJob request for a check step's baselining job.

Args:
processing_inputs (List[ProcessingInput]): The baselining job inputs.
processing_output (ProcessingOutput): The baselining job output.
processor (sagemaker.core.processing.Processor): The baselining processor.
check_job_config (CheckJobConfig): The check job config of the step.
default_job_name (str): The job name to fall back to when the processor
has not generated one yet.

Returns:
RequestType: A `CreateProcessingJob` request dict.
"""
# `KmsKeyId` belongs on ProcessingOutputConfig, not on ProcessingS3Output -
# the API applies it to every output. ProcessingS3Output accepts only S3Uri,
# LocalPath and S3UploadMode.
output_config = {
"Outputs": [
{
"OutputName": processing_output.output_name,
"S3Output": {
"S3Uri": processing_output.s3_output.s3_uri,
"LocalPath": processing_output.s3_output.local_path,
"S3UploadMode": processing_output.s3_output.s3_upload_mode,
},
}
],
}
if check_job_config.output_kms_key:
output_config["KmsKeyId"] = check_job_config.output_kms_key

cluster_config = {
"InstanceCount": processor.instance_count,
"InstanceType": processor.instance_type,
"VolumeSizeInGB": getattr(processor, "volume_size_in_gb", _DEFAULT_VOLUME_SIZE_IN_GB),
}
if check_job_config.volume_kms_key:
cluster_config["VolumeKmsKeyId"] = check_job_config.volume_kms_key

request_dict = {
"ProcessingInputs": [
_to_processing_input_dict(processing_input)
for processing_input in processing_inputs or []
],
"ProcessingOutputConfig": output_config,
"ProcessingJobName": processor._current_job_name or default_job_name,
"ProcessingResources": {
"ClusterConfig": cluster_config,
},
"AppSpecification": {
"ImageUri": processor.image_uri,
},
"RoleArn": processor.role,
"StoppingCondition": {
"MaxRuntimeInSeconds": getattr(processor, "max_runtime_in_seconds", None)
or _DEFAULT_MAX_RUNTIME_IN_SECONDS
},
}

# Add optional fields if they exist
if processor.env:
request_dict["Environment"] = processor.env
if processor.network_config:
request_dict["NetworkConfig"] = processor.network_config
if processor.entrypoint:
request_dict["AppSpecification"]["ContainerEntrypoint"] = processor.entrypoint
if processor.arguments:
request_dict["AppSpecification"]["ContainerArguments"] = processor.arguments
if check_job_config.tags:
request_dict["Tags"] = check_job_config.tags

return request_dict
81 changes: 10 additions & 71 deletions sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@

from sagemaker.mlops.workflow.steps import Step, StepTypeEnum, CacheConfig
from sagemaker.mlops.workflow.check_job_config import CheckJobConfig
from sagemaker.mlops.workflow._check_job_request import _build_check_job_request
from sagemaker.core.workflow.utilities import trim_request_dict

_DATA_BIAS_TYPE = "DATA_BIAS"
Expand Down Expand Up @@ -266,77 +267,15 @@ def arguments(self) -> RequestType:
"""The arguments dict that is used to define the ClarifyCheck step."""
from sagemaker.core.workflow.utilities import _pipeline_config

# Create request dictionary manually with correct AWS API field names
processing_inputs = []
for inp_key in ["config_input", "data_input"]:
inp = self._processing_params[inp_key]
input_dict = {
"InputName": inp.input_name,
}
if inp.s3_input:
input_dict["S3Input"] = {
"S3Uri": inp.s3_input.s3_uri,
"LocalPath": inp.s3_input.local_path,
"S3DataType": getattr(inp.s3_input, "s3_data_type", "S3Prefix"),
"S3InputMode": getattr(inp.s3_input, "s3_input_mode", "File"),
}
processing_inputs.append(input_dict)

s3_output_dict = {
"S3Uri": self._processing_params["result_output"].s3_output.s3_uri,
"LocalPath": self._processing_params["result_output"].s3_output.local_path,
"S3UploadMode": self._processing_params["result_output"].s3_output.s3_upload_mode,
}
if self.check_job_config.output_kms_key:
s3_output_dict["KmsKeyId"] = self.check_job_config.output_kms_key

processing_outputs = [
{
"OutputName": self._processing_params["result_output"].output_name,
"S3Output": s3_output_dict,
}
]

cluster_config = {
"InstanceCount": self._baselining_processor.instance_count,
"InstanceType": self._baselining_processor.instance_type,
"VolumeSizeInGB": getattr(self._baselining_processor, "volume_size_in_gb", 30),
}
if self.check_job_config.volume_kms_key:
cluster_config["VolumeKmsKeyId"] = self.check_job_config.volume_kms_key

request_dict = {
"ProcessingInputs": processing_inputs,
"ProcessingOutputConfig": {"Outputs": processing_outputs},
"ProcessingJobName": self._baselining_processor._current_job_name or "clarify-job",
"ProcessingResources": {
"ClusterConfig": cluster_config,
},
"AppSpecification": {
"ImageUri": self._baselining_processor.image_uri,
},
"RoleArn": self._baselining_processor.role,
"StoppingCondition": {
"MaxRuntimeInSeconds": getattr(
self._baselining_processor, "max_runtime_in_seconds", None
)
or 86400
},
}

# Add optional fields if they exist
if self._baselining_processor.env:
request_dict["Environment"] = self._baselining_processor.env
if self._baselining_processor.network_config:
request_dict["NetworkConfig"] = self._baselining_processor.network_config
if self._baselining_processor.entrypoint:
request_dict["AppSpecification"][
"ContainerEntrypoint"
] = self._baselining_processor.entrypoint
if self._baselining_processor.arguments:
request_dict["AppSpecification"][
"ContainerArguments"
] = self._baselining_processor.arguments
request_dict = _build_check_job_request(
processing_inputs=[
self._processing_params[inp_key] for inp_key in ["config_input", "data_input"]
],
processing_output=self._processing_params["result_output"],
processor=self._baselining_processor,
check_job_config=self.check_job_config,
default_job_name="clarify-job",
)
# Continue to pop job name if not explicitly opted-in via config
request_dict = trim_request_dict(request_dict, "ProcessingJobName", _pipeline_config)

Expand Down
79 changes: 8 additions & 71 deletions sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@

from sagemaker.mlops.workflow.steps import Step, StepTypeEnum, CacheConfig
from sagemaker.mlops.workflow.check_job_config import CheckJobConfig
from sagemaker.mlops.workflow._check_job_request import _build_check_job_request
from sagemaker.core.workflow.utilities import trim_request_dict

_CONTAINER_BASE_PATH = "/opt/ml/processing"
Expand Down Expand Up @@ -240,77 +241,13 @@ def arguments(self) -> RequestType:
"""The arguments dict that is used to define the QualityCheck step."""
from sagemaker.core.workflow.utilities import _pipeline_config

# Create request dictionary manually with correct AWS API field names
processing_inputs = []
if self._baseline_job_inputs:
for inp in self._baseline_job_inputs:
input_dict = {
"InputName": inp.input_name,
}
if inp.s3_input:
input_dict["S3Input"] = {
"S3Uri": inp.s3_input.s3_uri,
"LocalPath": inp.s3_input.local_path,
"S3DataType": getattr(inp.s3_input, "s3_data_type", "S3Prefix"),
"S3InputMode": getattr(inp.s3_input, "s3_input_mode", "File"),
}
processing_inputs.append(input_dict)

s3_output_dict = {
"S3Uri": self._baseline_output.s3_output.s3_uri,
"LocalPath": self._baseline_output.s3_output.local_path,
"S3UploadMode": self._baseline_output.s3_output.s3_upload_mode,
}
if self.check_job_config.output_kms_key:
s3_output_dict["KmsKeyId"] = self.check_job_config.output_kms_key

processing_outputs = [
{
"OutputName": self._baseline_output.output_name,
"S3Output": s3_output_dict,
}
]

cluster_config = {
"InstanceCount": self._baselining_processor.instance_count,
"InstanceType": self._baselining_processor.instance_type,
"VolumeSizeInGB": getattr(self._baselining_processor, "volume_size_in_gb", 30),
}
if self.check_job_config.volume_kms_key:
cluster_config["VolumeKmsKeyId"] = self.check_job_config.volume_kms_key

request_dict = {
"ProcessingInputs": processing_inputs,
"ProcessingOutputConfig": {"Outputs": processing_outputs},
"ProcessingJobName": self._baselining_processor._current_job_name or "baseline-job",
"ProcessingResources": {
"ClusterConfig": cluster_config,
},
"AppSpecification": {
"ImageUri": self._baselining_processor.image_uri,
},
"RoleArn": self._baselining_processor.role,
"StoppingCondition": {
"MaxRuntimeInSeconds": getattr(
self._baselining_processor, "max_runtime_in_seconds", None
)
or 86400
},
}

# Add optional fields if they exist
if self._baselining_processor.env:
request_dict["Environment"] = self._baselining_processor.env
if self._baselining_processor.network_config:
request_dict["NetworkConfig"] = self._baselining_processor.network_config
if self._baselining_processor.entrypoint:
request_dict["AppSpecification"][
"ContainerEntrypoint"
] = self._baselining_processor.entrypoint
if self._baselining_processor.arguments:
request_dict["AppSpecification"][
"ContainerArguments"
] = self._baselining_processor.arguments
request_dict = _build_check_job_request(
processing_inputs=self._baseline_job_inputs,
processing_output=self._baseline_output,
processor=self._baselining_processor,
check_job_config=self.check_job_config,
default_job_name="baseline-job",
)
# Continue to pop job name if not explicitly opted-in via config
request_dict = trim_request_dict(request_dict, "ProcessingJobName", _pipeline_config)

Expand Down
Loading
Loading