From 392457206af1b6a285c10664425b0d068cf5f9a5 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Sun, 18 Oct 2026 17:51:51 -0700 Subject: [PATCH] fix: put the check steps' output KMS key on ProcessingOutputConfig and send CheckJobConfig tags QualityCheckStep.arguments and ClarifyCheckStep.arguments build their CreateProcessingJob request by hand. Both put KmsKeyId inside ProcessingS3Output, which the API does not accept - KmsKeyId belongs on ProcessingOutputConfig, where it applies to every output. Both also never read check_job_config.tags, so configured tags were dropped. Validating the pre-fix request against the CreateProcessingJob service model reports: Unknown parameter in ProcessingOutputConfig.Outputs[0].S3Output: "KmsKeyId", must be one of: S3Uri, LocalPath, S3UploadMode Both are v2 -> v3 regressions. In v2 the request came from Session._get_process_request via Processor._get_process_args, which set output_config["KmsKeyId"] and request["Tags"]; the v3 hand-rolled builders lost both. The two builders were identical apart from the inputs, the output and the default job name, so the shaping moves into a shared _build_check_job_request helper. The fix then lands once and the two steps can no longer drift. Fixes #6495 --- .../mlops/workflow/_check_job_request.py | 132 ++++++++++ .../mlops/workflow/clarify_check_step.py | 81 +----- .../mlops/workflow/quality_check_step.py | 79 +----- .../integ/test_check_step_kms_propagation.py | 82 +++--- .../unit/workflow/test_check_job_request.py | 233 ++++++++++++++++++ .../workflow/test_clarify_check_step_kms.py | 89 +++++-- .../workflow/test_quality_check_step_kms.py | 90 +++++-- 7 files changed, 575 insertions(+), 211 deletions(-) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/_check_job_request.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_check_job_request.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_check_job_request.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_check_job_request.py new file mode 100644 index 0000000000..459d72db44 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_check_job_request.py @@ -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 diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py index e7f0077520..34a7cfbb06 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py @@ -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" @@ -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) 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 cda8bec8e0..b965a2aa1b 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/quality_check_step.py @@ -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" @@ -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) diff --git a/sagemaker-mlops/tests/integ/test_check_step_kms_propagation.py b/sagemaker-mlops/tests/integ/test_check_step_kms_propagation.py index 92dd80dc4f..d8f857b692 100644 --- a/sagemaker-mlops/tests/integ/test_check_step_kms_propagation.py +++ b/sagemaker-mlops/tests/integ/test_check_step_kms_propagation.py @@ -10,11 +10,11 @@ # 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. -"""Integration test for KMS key propagation in check steps. +"""Integration test for KMS key and tag propagation in check steps. This test constructs real QualityCheckStep and ClarifyCheckStep objects using the actual SDK classes with a real SageMaker Session, then inspects the compiled -step arguments to verify KmsKeyId and VolumeKmsKeyId are present. +step arguments to verify KmsKeyId, VolumeKmsKeyId and Tags are present. No SageMaker compute resources are launched. The only AWS interaction is a small S3 put_object for the Clarify analysis config (cleaned up in teardown). @@ -44,6 +44,7 @@ # the key appears in the compiled request dict. _TEST_OUTPUT_KMS_KEY = "arn:aws:kms:us-west-2:123456789012:key/test-output-key-id" _TEST_VOLUME_KMS_KEY = "arn:aws:kms:us-west-2:123456789012:key/test-volume-key-id" +_TEST_TAGS = [{"Key": "integ-test", "Value": "check-step"}] _S3_PREFIX = "integ-test-kms-check-step" @@ -71,7 +72,7 @@ def s3_client(sagemaker_session): @pytest.fixture def check_job_config_with_kms(role, sagemaker_session): - """CheckJobConfig with both output and volume KMS keys.""" + """CheckJobConfig with both output and volume KMS keys, and tags.""" return CheckJobConfig( role=role, instance_count=1, @@ -80,13 +81,14 @@ def check_job_config_with_kms(role, sagemaker_session): volume_kms_key=_TEST_VOLUME_KMS_KEY, output_kms_key=_TEST_OUTPUT_KMS_KEY, max_runtime_in_seconds=3600, + tags=_TEST_TAGS, sagemaker_session=sagemaker_session, ) @pytest.fixture def check_job_config_no_kms(role, sagemaker_session): - """CheckJobConfig without KMS keys.""" + """CheckJobConfig without KMS keys or tags.""" return CheckJobConfig( role=role, instance_count=1, @@ -130,13 +132,24 @@ def _build_step(self, check_job_config, bucket): ) def test_output_kms_key_in_arguments(self, check_job_config_with_kms, bucket): - """output_kms_key from CheckJobConfig appears as KmsKeyId in S3Output.""" + """output_kms_key from CheckJobConfig appears as KmsKeyId in ProcessingOutputConfig.""" step = self._build_step(check_job_config_with_kms, bucket) args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" in s3_output, f"Expected KmsKeyId in S3Output but got: {s3_output}" - assert s3_output["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY + output_config = args["ProcessingOutputConfig"] + assert ( + "KmsKeyId" in output_config + ), f"Expected KmsKeyId in ProcessingOutputConfig but got: {output_config}" + assert output_config["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY + # ProcessingS3Output does not accept KmsKeyId. + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] + + def test_tags_in_arguments(self, check_job_config_with_kms, bucket): + """tags from CheckJobConfig appear as top-level Tags.""" + step = self._build_step(check_job_config_with_kms, bucket) + args = step.arguments + + assert args["Tags"] == _TEST_TAGS def test_volume_kms_key_in_arguments(self, check_job_config_with_kms, bucket): """volume_kms_key from CheckJobConfig appears as VolumeKmsKeyId in ClusterConfig.""" @@ -149,17 +162,20 @@ def test_volume_kms_key_in_arguments(self, check_job_config_with_kms, bucket): ), f"Expected VolumeKmsKeyId in ClusterConfig but got: {cluster_config}" assert cluster_config["VolumeKmsKeyId"] == _TEST_VOLUME_KMS_KEY - def test_no_kms_keys_when_not_configured(self, check_job_config_no_kms, bucket): - """KMS keys are absent from arguments when not set in CheckJobConfig.""" + def test_no_kms_keys_or_tags_when_not_configured(self, check_job_config_no_kms, bucket): + """KMS keys and tags are absent from arguments when not set in CheckJobConfig.""" step = self._build_step(check_job_config_no_kms, bucket) args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + output_config = args["ProcessingOutputConfig"] + assert "KmsKeyId" not in output_config + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config + assert "Tags" not in args + def test_arguments_are_json_serializable(self, check_job_config_with_kms, bucket): """The compiled arguments dict is valid JSON (required for pipeline definitions).""" step = self._build_step(check_job_config_with_kms, bucket) @@ -167,13 +183,11 @@ def test_arguments_are_json_serializable(self, check_job_config_with_kms, bucket json_str = json.dumps(args, default=str) parsed = json.loads(json_str) - assert ( - parsed["ProcessingOutputConfig"]["Outputs"][0]["S3Output"]["KmsKeyId"] - == _TEST_OUTPUT_KMS_KEY - ) + assert parsed["ProcessingOutputConfig"]["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY assert ( parsed["ProcessingResources"]["ClusterConfig"]["VolumeKmsKeyId"] == _TEST_VOLUME_KMS_KEY ) + assert parsed["Tags"] == _TEST_TAGS class TestDataBiasCheckStepKms: @@ -208,13 +222,24 @@ def _build_step(self, check_job_config, bucket): ) def test_output_kms_key_in_arguments(self, check_job_config_with_kms, bucket): - """output_kms_key from CheckJobConfig appears as KmsKeyId in S3Output.""" + """output_kms_key from CheckJobConfig appears as KmsKeyId in ProcessingOutputConfig.""" step = self._build_step(check_job_config_with_kms, bucket) args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" in s3_output, f"Expected KmsKeyId in S3Output but got: {s3_output}" - assert s3_output["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY + output_config = args["ProcessingOutputConfig"] + assert ( + "KmsKeyId" in output_config + ), f"Expected KmsKeyId in ProcessingOutputConfig but got: {output_config}" + assert output_config["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY + # ProcessingS3Output does not accept KmsKeyId. + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] + + def test_tags_in_arguments(self, check_job_config_with_kms, bucket): + """tags from CheckJobConfig appear as top-level Tags.""" + step = self._build_step(check_job_config_with_kms, bucket) + args = step.arguments + + assert args["Tags"] == _TEST_TAGS def test_volume_kms_key_in_arguments(self, check_job_config_with_kms, bucket): """volume_kms_key from CheckJobConfig appears as VolumeKmsKeyId in ClusterConfig.""" @@ -227,17 +252,20 @@ def test_volume_kms_key_in_arguments(self, check_job_config_with_kms, bucket): ), f"Expected VolumeKmsKeyId in ClusterConfig but got: {cluster_config}" assert cluster_config["VolumeKmsKeyId"] == _TEST_VOLUME_KMS_KEY - def test_no_kms_keys_when_not_configured(self, check_job_config_no_kms, bucket): - """KMS keys are absent from arguments when not set in CheckJobConfig.""" + def test_no_kms_keys_or_tags_when_not_configured(self, check_job_config_no_kms, bucket): + """KMS keys and tags are absent from arguments when not set in CheckJobConfig.""" step = self._build_step(check_job_config_no_kms, bucket) args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + output_config = args["ProcessingOutputConfig"] + assert "KmsKeyId" not in output_config + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config + assert "Tags" not in args + def test_arguments_are_json_serializable(self, check_job_config_with_kms, bucket): """The compiled arguments dict is valid JSON (required for pipeline definitions).""" step = self._build_step(check_job_config_with_kms, bucket) @@ -245,10 +273,8 @@ def test_arguments_are_json_serializable(self, check_job_config_with_kms, bucket json_str = json.dumps(args, default=str) parsed = json.loads(json_str) - assert ( - parsed["ProcessingOutputConfig"]["Outputs"][0]["S3Output"]["KmsKeyId"] - == _TEST_OUTPUT_KMS_KEY - ) + assert parsed["ProcessingOutputConfig"]["KmsKeyId"] == _TEST_OUTPUT_KMS_KEY assert ( parsed["ProcessingResources"]["ClusterConfig"]["VolumeKmsKeyId"] == _TEST_VOLUME_KMS_KEY ) + assert parsed["Tags"] == _TEST_TAGS diff --git a/sagemaker-mlops/tests/unit/workflow/test_check_job_request.py b/sagemaker-mlops/tests/unit/workflow/test_check_job_request.py new file mode 100644 index 0000000000..e72d9823cb --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_check_job_request.py @@ -0,0 +1,233 @@ +# 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. +"""Unit tests for the shared check-step CreateProcessingJob request builder.""" + +from __future__ import absolute_import + +from unittest.mock import Mock + +import botocore.session +import pytest +from botocore.validate import ParamValidator + +from sagemaker.mlops.workflow._check_job_request import _build_check_job_request + +_OUTPUT_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/output-key-id" +_VOLUME_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/volume-key-id" +_TAGS = [{"Key": "test", "Value": "check-step"}] + + +@pytest.fixture +def processing_input(): + processing_input = Mock() + processing_input.input_name = "baseline_dataset_input" + processing_input.s3_input = Mock() + processing_input.s3_input.s3_uri = "s3://bucket/input/data.csv" + processing_input.s3_input.local_path = "/opt/ml/processing/input" + processing_input.s3_input.s3_data_type = "S3Prefix" + processing_input.s3_input.s3_input_mode = "File" + return processing_input + + +@pytest.fixture +def processing_output(): + processing_output = Mock() + processing_output.output_name = "quality_check_output" + processing_output.s3_output = Mock() + processing_output.s3_output.s3_uri = "s3://bucket/output" + processing_output.s3_output.local_path = "/opt/ml/processing/output" + processing_output.s3_output.s3_upload_mode = "EndOfJob" + return processing_output + + +@pytest.fixture +def processor(): + processor = Mock() + processor._current_job_name = "baseline-job-2026-01-01" + processor.instance_count = 1 + processor.instance_type = "ml.m5.xlarge" + processor.volume_size_in_gb = 30 + processor.image_uri = "123456789012.dkr.ecr.us-east-1.amazonaws.com/monitor:latest" + processor.role = "arn:aws:iam::123456789012:role/SageMakerRole" + processor.max_runtime_in_seconds = 3600 + processor.env = None + processor.network_config = None + processor.entrypoint = None + processor.arguments = None + return processor + + +def _check_job_config(output_kms_key=None, volume_kms_key=None, tags=None): + check_job_config = Mock() + check_job_config.output_kms_key = output_kms_key + check_job_config.volume_kms_key = volume_kms_key + check_job_config.tags = tags + return check_job_config + + +def _build(processing_input, processing_output, processor, check_job_config): + return _build_check_job_request( + processing_inputs=[processing_input], + processing_output=processing_output, + processor=processor, + check_job_config=check_job_config, + default_job_name="baseline-job", + ) + + +def test_request_matches_create_processing_job_shape( + processing_input, processing_output, processor +): + """The request validates against the CreateProcessingJob service model. + + This is the guard for issue #6495: `KmsKeyId` on ProcessingS3Output and a + `Tags` list in the wrong place are both rejected here. + """ + request = _build( + processing_input, + processing_output, + processor, + _check_job_config( + output_kms_key=_OUTPUT_KMS_KEY, + volume_kms_key=_VOLUME_KMS_KEY, + tags=_TAGS, + ), + ) + + input_shape = ( + botocore.session.get_session() + .get_service_model("sagemaker") + .operation_model("CreateProcessingJob") + .input_shape + ) + errors = ParamValidator().validate(request, input_shape) + assert not errors.has_errors(), errors.generate_report() + + +def test_output_kms_key_on_output_config(processing_input, processing_output, processor): + """output_kms_key lands on ProcessingOutputConfig, not on ProcessingS3Output.""" + request = _build( + processing_input, + processing_output, + processor, + _check_job_config(output_kms_key=_OUTPUT_KMS_KEY), + ) + + output_config = request["ProcessingOutputConfig"] + assert output_config["KmsKeyId"] == _OUTPUT_KMS_KEY + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] + + +def test_tags_on_request_root(processing_input, processing_output, processor): + """Tags land at the root of the request.""" + request = _build(processing_input, processing_output, processor, _check_job_config(tags=_TAGS)) + + assert request["Tags"] == _TAGS + + +@pytest.mark.parametrize("tags", [None, []]) +def test_tags_omitted_when_unset(processing_input, processing_output, processor, tags): + """Tags is omitted rather than sent empty.""" + request = _build(processing_input, processing_output, processor, _check_job_config(tags=tags)) + + assert "Tags" not in request + + +def test_inputs_and_output_are_shaped(processing_input, processing_output, processor): + """Inputs and the output keep their API field names.""" + request = _build(processing_input, processing_output, processor, _check_job_config()) + + assert request["ProcessingInputs"] == [ + { + "InputName": "baseline_dataset_input", + "S3Input": { + "S3Uri": "s3://bucket/input/data.csv", + "LocalPath": "/opt/ml/processing/input", + "S3DataType": "S3Prefix", + "S3InputMode": "File", + }, + } + ] + assert request["ProcessingOutputConfig"] == { + "Outputs": [ + { + "OutputName": "quality_check_output", + "S3Output": { + "S3Uri": "s3://bucket/output", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + }, + } + ], + } + + +def test_input_without_s3_input_is_name_only(processing_output, processor): + """An input with no s3_input contributes only its name.""" + processing_input = Mock() + processing_input.input_name = "no_s3_input" + processing_input.s3_input = None + + request = _build(processing_input, processing_output, processor, _check_job_config()) + + assert request["ProcessingInputs"] == [{"InputName": "no_s3_input"}] + + +def test_empty_inputs_are_allowed(processing_output, processor): + """A step with no inputs produces an empty ProcessingInputs list.""" + request = _build_check_job_request( + processing_inputs=[], + processing_output=processing_output, + processor=processor, + check_job_config=_check_job_config(), + default_job_name="baseline-job", + ) + + assert request["ProcessingInputs"] == [] + + +def test_default_job_name_used_when_processor_has_none( + processing_input, processing_output, processor +): + """The default job name is used when the processor has not generated one.""" + processor._current_job_name = None + + request = _build(processing_input, processing_output, processor, _check_job_config()) + + assert request["ProcessingJobName"] == "baseline-job" + + +def test_optional_processor_fields_are_added(processing_input, processing_output, processor): + """Environment, NetworkConfig and container overrides are added when set.""" + processor.env = {"KEY": "value"} + processor.network_config = {"EnableNetworkIsolation": True} + processor.entrypoint = ["python3"] + processor.arguments = ["--flag"] + + request = _build(processing_input, processing_output, processor, _check_job_config()) + + assert request["Environment"] == {"KEY": "value"} + assert request["NetworkConfig"] == {"EnableNetworkIsolation": True} + assert request["AppSpecification"]["ContainerEntrypoint"] == ["python3"] + assert request["AppSpecification"]["ContainerArguments"] == ["--flag"] + + +def test_defaults_when_processor_values_missing(processing_input, processing_output, processor): + """Volume size and max runtime fall back to their defaults.""" + processor.max_runtime_in_seconds = None + del processor.volume_size_in_gb + + request = _build(processing_input, processing_output, processor, _check_job_config()) + + assert request["ProcessingResources"]["ClusterConfig"]["VolumeSizeInGB"] == 30 + assert request["StoppingCondition"]["MaxRuntimeInSeconds"] == 86400 diff --git a/sagemaker-mlops/tests/unit/workflow/test_clarify_check_step_kms.py b/sagemaker-mlops/tests/unit/workflow/test_clarify_check_step_kms.py index 00ec687162..37d6649333 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_clarify_check_step_kms.py +++ b/sagemaker-mlops/tests/unit/workflow/test_clarify_check_step_kms.py @@ -10,7 +10,7 @@ # 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. -"""Unit tests for KMS key propagation in ClarifyCheckStep.""" +"""Unit tests for KMS key and tag propagation in ClarifyCheckStep.""" from __future__ import absolute_import @@ -22,6 +22,7 @@ _OUTPUT_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/output-key-id" _VOLUME_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/volume-key-id" +_TAGS = [{"Key": "test", "Value": "clarify-check"}] # Patch trim_request_dict to be a no-op since it only trims job names # and is not relevant to KMS key propagation testing. @@ -33,7 +34,7 @@ def _noop_trim(request_dict, *args, **kwargs): return request_dict -def _create_mock_clarify_check_step(output_kms_key=None, volume_kms_key=None): +def _create_mock_clarify_check_step(output_kms_key=None, volume_kms_key=None, tags=None): """Create a ClarifyCheckStep with mocked internals for testing arguments.""" step = object.__new__(ClarifyCheckStep) @@ -41,6 +42,7 @@ def _create_mock_clarify_check_step(output_kms_key=None, volume_kms_key=None): step.check_job_config = Mock() step.check_job_config.output_kms_key = output_kms_key step.check_job_config.volume_kms_key = volume_kms_key + step.check_job_config.tags = tags # Mock processing params (config_input, data_input, result_output) config_input = Mock() @@ -95,15 +97,17 @@ class TestClarifyCheckStepKmsKeyPropagation: """Tests for KMS key propagation in ClarifyCheckStep.arguments.""" @patch(_TRIM_PATCH, side_effect=_noop_trim) - def test_output_kms_key_propagated_in_s3_output(self, mock_trim): - """Test that output_kms_key from CheckJobConfig is included in S3Output.""" + def test_output_kms_key_propagated_in_output_config(self, mock_trim): + """Test that output_kms_key from CheckJobConfig is included in ProcessingOutputConfig.""" step = _create_mock_clarify_check_step(output_kms_key=_OUTPUT_KMS_KEY) args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" in s3_output - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + # KmsKeyId belongs on ProcessingOutputConfig per the CreateProcessingJob + # API - ProcessingS3Output does not accept it. + output_config = args["ProcessingOutputConfig"] + assert output_config["KmsKeyId"] == _OUTPUT_KMS_KEY + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_volume_kms_key_propagated_in_cluster_config(self, mock_trim): @@ -126,8 +130,7 @@ def test_both_kms_keys_propagated(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY cluster_config = args["ProcessingResources"]["ClusterConfig"] assert cluster_config["VolumeKmsKeyId"] == _VOLUME_KMS_KEY @@ -139,24 +142,26 @@ def test_no_kms_keys_when_not_set(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + output_config = args["ProcessingOutputConfig"] + assert "KmsKeyId" not in output_config + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_s3_output_retains_other_fields_with_kms(self, mock_trim): - """Test that S3Output still contains S3Uri, LocalPath, S3UploadMode alongside KmsKeyId.""" + """Test that S3Output still contains S3Uri, LocalPath and S3UploadMode.""" step = _create_mock_clarify_check_step(output_kms_key=_OUTPUT_KMS_KEY) args = step.arguments s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["S3Uri"] == "s3://bucket/output/results" - assert s3_output["LocalPath"] == "/opt/ml/processing/output" - assert s3_output["S3UploadMode"] == "EndOfJob" - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert s3_output == { + "S3Uri": "s3://bucket/output/results", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + } @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_cluster_config_retains_other_fields_with_kms(self, mock_trim): @@ -178,8 +183,7 @@ def test_output_kms_key_only_without_volume_kms(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config @@ -191,8 +195,53 @@ def test_volume_kms_key_only_without_output_kms(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + assert "KmsKeyId" not in args["ProcessingOutputConfig"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert cluster_config["VolumeKmsKeyId"] == _VOLUME_KMS_KEY + + +class TestClarifyCheckStepTagPropagation: + """Tests for tag propagation in ClarifyCheckStep.arguments.""" + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_propagated(self, mock_trim): + """Test that tags from CheckJobConfig are included in the request.""" + step = _create_mock_clarify_check_step(tags=_TAGS) + + args = step.arguments + + assert args["Tags"] == _TAGS + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_absent_when_not_set(self, mock_trim): + """Test that Tags is omitted when CheckJobConfig has no tags.""" + step = _create_mock_clarify_check_step(tags=None) + + args = step.arguments + + assert "Tags" not in args + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_absent_when_empty(self, mock_trim): + """Test that Tags is omitted rather than sent empty.""" + step = _create_mock_clarify_check_step(tags=[]) + + args = step.arguments + + assert "Tags" not in args + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_and_kms_keys_propagated_together(self, mock_trim): + """Test that tags and both KMS keys are propagated in the same request.""" + step = _create_mock_clarify_check_step( + output_kms_key=_OUTPUT_KMS_KEY, + volume_kms_key=_VOLUME_KMS_KEY, + tags=_TAGS, + ) + + args = step.arguments + + assert args["Tags"] == _TAGS + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingResources"]["ClusterConfig"]["VolumeKmsKeyId"] == _VOLUME_KMS_KEY diff --git a/sagemaker-mlops/tests/unit/workflow/test_quality_check_step_kms.py b/sagemaker-mlops/tests/unit/workflow/test_quality_check_step_kms.py index c421e56cfe..313480d7f3 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_quality_check_step_kms.py +++ b/sagemaker-mlops/tests/unit/workflow/test_quality_check_step_kms.py @@ -10,7 +10,7 @@ # 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. -"""Unit tests for KMS key propagation in QualityCheckStep.""" +"""Unit tests for KMS key and tag propagation in QualityCheckStep.""" from __future__ import absolute_import @@ -22,6 +22,7 @@ _OUTPUT_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/output-key-id" _VOLUME_KMS_KEY = "arn:aws:kms:us-east-1:123456789012:key/volume-key-id" +_TAGS = [{"Key": "test", "Value": "quality-check"}] # Patch trim_request_dict to be a no-op since it only trims job names # and is not relevant to KMS key propagation testing. @@ -33,7 +34,7 @@ def _noop_trim(request_dict, *args, **kwargs): return request_dict -def _create_mock_quality_check_step(output_kms_key=None, volume_kms_key=None): +def _create_mock_quality_check_step(output_kms_key=None, volume_kms_key=None, tags=None): """Create a QualityCheckStep with mocked internals for testing arguments.""" step = object.__new__(QualityCheckStep) @@ -41,6 +42,7 @@ def _create_mock_quality_check_step(output_kms_key=None, volume_kms_key=None): step.check_job_config = Mock() step.check_job_config.output_kms_key = output_kms_key step.check_job_config.volume_kms_key = volume_kms_key + step.check_job_config.tags = tags # Mock baseline output step._baseline_output = Mock() @@ -90,16 +92,17 @@ class TestQualityCheckStepKmsKeyPropagation: """Tests for KMS key propagation in QualityCheckStep.arguments.""" @patch(_TRIM_PATCH, side_effect=_noop_trim) - def test_output_kms_key_propagated_in_s3_output(self, mock_trim): - """Test that output_kms_key from CheckJobConfig is included in S3Output.""" + def test_output_kms_key_propagated_in_output_config(self, mock_trim): + """Test that output_kms_key from CheckJobConfig is included in ProcessingOutputConfig.""" step = _create_mock_quality_check_step(output_kms_key=_OUTPUT_KMS_KEY) args = step.arguments - # Verify KmsKeyId is present in S3Output - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" in s3_output - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + # KmsKeyId belongs on ProcessingOutputConfig per the CreateProcessingJob + # API - ProcessingS3Output does not accept it. + output_config = args["ProcessingOutputConfig"] + assert output_config["KmsKeyId"] == _OUTPUT_KMS_KEY + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_volume_kms_key_propagated_in_cluster_config(self, mock_trim): @@ -123,8 +126,7 @@ def test_both_kms_keys_propagated(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY cluster_config = args["ProcessingResources"]["ClusterConfig"] assert cluster_config["VolumeKmsKeyId"] == _VOLUME_KMS_KEY @@ -136,24 +138,26 @@ def test_no_kms_keys_when_not_set(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + output_config = args["ProcessingOutputConfig"] + assert "KmsKeyId" not in output_config + assert "KmsKeyId" not in output_config["Outputs"][0]["S3Output"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_s3_output_retains_other_fields_with_kms(self, mock_trim): - """Test that S3Output still contains S3Uri, LocalPath, S3UploadMode alongside KmsKeyId.""" + """Test that S3Output still contains S3Uri, LocalPath and S3UploadMode.""" step = _create_mock_quality_check_step(output_kms_key=_OUTPUT_KMS_KEY) args = step.arguments s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["S3Uri"] == "s3://bucket/output" - assert s3_output["LocalPath"] == "/opt/ml/processing/output" - assert s3_output["S3UploadMode"] == "EndOfJob" - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert s3_output == { + "S3Uri": "s3://bucket/output", + "LocalPath": "/opt/ml/processing/output", + "S3UploadMode": "EndOfJob", + } @patch(_TRIM_PATCH, side_effect=_noop_trim) def test_cluster_config_retains_other_fields_with_kms(self, mock_trim): @@ -175,8 +179,7 @@ def test_output_kms_key_only_without_volume_kms(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert s3_output["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY cluster_config = args["ProcessingResources"]["ClusterConfig"] assert "VolumeKmsKeyId" not in cluster_config @@ -188,8 +191,53 @@ def test_volume_kms_key_only_without_output_kms(self, mock_trim): args = step.arguments - s3_output = args["ProcessingOutputConfig"]["Outputs"][0]["S3Output"] - assert "KmsKeyId" not in s3_output + assert "KmsKeyId" not in args["ProcessingOutputConfig"] cluster_config = args["ProcessingResources"]["ClusterConfig"] assert cluster_config["VolumeKmsKeyId"] == _VOLUME_KMS_KEY + + +class TestQualityCheckStepTagPropagation: + """Tests for tag propagation in QualityCheckStep.arguments.""" + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_propagated(self, mock_trim): + """Test that tags from CheckJobConfig are included in the request.""" + step = _create_mock_quality_check_step(tags=_TAGS) + + args = step.arguments + + assert args["Tags"] == _TAGS + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_absent_when_not_set(self, mock_trim): + """Test that Tags is omitted when CheckJobConfig has no tags.""" + step = _create_mock_quality_check_step(tags=None) + + args = step.arguments + + assert "Tags" not in args + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_absent_when_empty(self, mock_trim): + """Test that Tags is omitted rather than sent empty.""" + step = _create_mock_quality_check_step(tags=[]) + + args = step.arguments + + assert "Tags" not in args + + @patch(_TRIM_PATCH, side_effect=_noop_trim) + def test_tags_and_kms_keys_propagated_together(self, mock_trim): + """Test that tags and both KMS keys are propagated in the same request.""" + step = _create_mock_quality_check_step( + output_kms_key=_OUTPUT_KMS_KEY, + volume_kms_key=_VOLUME_KMS_KEY, + tags=_TAGS, + ) + + args = step.arguments + + assert args["Tags"] == _TAGS + assert args["ProcessingOutputConfig"]["KmsKeyId"] == _OUTPUT_KMS_KEY + assert args["ProcessingResources"]["ClusterConfig"]["VolumeKmsKeyId"] == _VOLUME_KMS_KEY