From 95318553122b9011a0514de48e1917ceb5cd2ba2 Mon Sep 17 00:00:00 2001
From: Akshay Shah <akshayshah0803@gmail.com>
Date: Fri, 9 Oct 2026 13:11:13 -0500
Subject: [PATCH] fix(mlops): keep S3CompressionType/S3DataDistributionType in
 ClarifyCheckStep inputs

ClarifyCheckStep builds the dataset ProcessingInput with
s3_compression_type and s3_data_distribution_type from the DataConfig, but
ClarifyCheckStep.arguments only serialized S3Uri, LocalPath, S3DataType and
S3InputMode. A DataConfig(s3_compression_type='Gzip') therefore produced a
processing job without S3CompressionType, so the gzipped dataset reached the
Clarify container still compressed.

Serialize both fields when they are set, as core's SageMakerClarifyProcessor
does.
---
 .../mlops/workflow/clarify_check_step.py      |  9 +++++++++
 .../workflow/test_clarify_check_step_kms.py   | 20 +++++++++++++++++++
 2 files changed, 29 insertions(+)

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..cbe6a6cb81 100644
--- a/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py
+++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/clarify_check_step.py
@@ -280,6 +280,15 @@ def arguments(self) -> RequestType:
                     "S3DataType": getattr(inp.s3_input, "s3_data_type", "S3Prefix"),
                     "S3InputMode": getattr(inp.s3_input, "s3_input_mode", "File"),
                 }
+                # Optional fields, e.g. DataConfig(s3_compression_type="Gzip"); skip unset
+                # (None/Unassigned) values.
+                for api_key, attr in (
+                    ("S3DataDistributionType", "s3_data_distribution_type"),
+                    ("S3CompressionType", "s3_compression_type"),
+                ):
+                    value = getattr(inp.s3_input, attr, None)
+                    if isinstance(value, str):
+                        input_dict["S3Input"][api_key] = value
             processing_inputs.append(input_dict)
 
         s3_output_dict = {
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..2fa163d4d4 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
@@ -91,6 +91,26 @@ def _create_mock_clarify_check_step(output_kms_key=None, volume_kms_key=None):
     return step
 
 
+class TestClarifyCheckStepDatasetInput:
+    """Tests for the dataset ProcessingInput in ClarifyCheckStep.arguments."""
+
+    @patch(_TRIM_PATCH, side_effect=_noop_trim)
+    def test_dataset_compression_and_distribution_type_propagated(self, mock_trim):
+        """S3CompressionType/S3DataDistributionType from DataConfig must reach the request."""
+        step = _create_mock_clarify_check_step()
+        data_s3_input = step._processing_params["data_input"].s3_input
+        data_s3_input.s3_compression_type = "Gzip"
+        data_s3_input.s3_data_distribution_type = "ShardedByS3Key"
+
+        args = step.arguments
+
+        dataset_input = next(
+            inp for inp in args["ProcessingInputs"] if inp["InputName"] == "dataset"
+        )
+        assert dataset_input["S3Input"]["S3CompressionType"] == "Gzip"
+        assert dataset_input["S3Input"]["S3DataDistributionType"] == "ShardedByS3Key"
+
+
 class TestClarifyCheckStepKmsKeyPropagation:
     """Tests for KMS key propagation in ClarifyCheckStep.arguments."""