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