From 2f29d791fc31efc132ac70c4ab6562d7514708d5 Mon Sep 17 00:00:00 2001 From: gauravkriplani <127468958+gauravkriplani@users.noreply.github.com> Date: Sun, 11 Oct 2026 22:09:20 -0700 Subject: [PATCH] feature: support local sources for processing inputs --- .../src/sagemaker/core/shapes/__init__.py | 1 + .../src/sagemaker/core/shapes/processing.py | 95 +++++++++++++++++++ sagemaker-core/tests/unit/test_processing.py | 35 +++++++ 3 files changed, 131 insertions(+) create mode 100644 sagemaker-core/src/sagemaker/core/shapes/processing.py diff --git a/sagemaker-core/src/sagemaker/core/shapes/__init__.py b/sagemaker-core/src/sagemaker/core/shapes/__init__.py index c3c7df1b76..87cf78780d 100644 --- a/sagemaker-core/src/sagemaker/core/shapes/__init__.py +++ b/sagemaker-core/src/sagemaker/core/shapes/__init__.py @@ -1,3 +1,4 @@ from sagemaker.core.shapes.shapes import * # noqa: F401,F403 +from sagemaker.core.shapes.processing import ProcessingInput # noqa: F401 from sagemaker.core.shapes.model_card_shapes import * # noqa: F401,F403 diff --git a/sagemaker-core/src/sagemaker/core/shapes/processing.py b/sagemaker-core/src/sagemaker/core/shapes/processing.py new file mode 100644 index 0000000000..8e37dacb63 --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/shapes/processing.py @@ -0,0 +1,95 @@ +# 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. +"""Convenience shapes for SageMaker Processing jobs.""" + +from typing import Optional + +from sagemaker.core.helper.pipeline_variable import BoolPipeVar, StrPipeVar +from sagemaker.core.shapes.shapes import ( + DatasetDefinition, + ProcessingInput as _ProcessingInput, + ProcessingS3Input, +) +from sagemaker.core.utils.utils import Unassigned + + +class ProcessingInput(_ProcessingInput): + """Configure an input for a SageMaker Processing job. + + ``source`` and ``destination`` provide the convenient v2-style interface. A local + ``source`` is uploaded to S3 when the processor normalizes its inputs. Callers can + alternatively supply the service-level ``s3_input`` or ``dataset_definition`` shape. + + Args: + source: Local path or S3 URI containing the input data. + destination: Absolute path where the input is made available in the container. + input_name: Name of the processing input. Generated by the processor when omitted. + s3_data_type: How SageMaker interprets the S3 data, such as ``S3Prefix``. + s3_input_mode: How data is made available to the container, such as ``File``. + s3_data_distribution_type: How data is distributed across processing instances. + s3_compression_type: Compression format of the input data. + app_managed: Whether the processing application manages input operations. + s3_input: Service-level S3 input configuration. Mutually exclusive with ``source``. + dataset_definition: Athena or Redshift dataset configuration. + """ + + input_name: Optional[StrPipeVar] = None + + def __init__( + self, + *, + source: Optional[StrPipeVar] = None, + destination: Optional[StrPipeVar] = None, + input_name: Optional[StrPipeVar] = None, + s3_data_type: StrPipeVar = "S3Prefix", + s3_input_mode: StrPipeVar = "File", + s3_data_distribution_type: StrPipeVar = "FullyReplicated", + s3_compression_type: StrPipeVar = "None", + app_managed: Optional[BoolPipeVar] = Unassigned(), + s3_input: Optional[ProcessingS3Input] = Unassigned(), + dataset_definition: Optional[DatasetDefinition] = Unassigned(), + ): + if source is not None and s3_input: + raise ValueError("source and s3_input are mutually exclusive") + if source is not None and destination is None: + raise ValueError("destination is required when source is provided") + + if source is not None: + s3_input = ProcessingS3Input( + s3_uri=source, + local_path=destination, + s3_data_type=s3_data_type, + s3_input_mode=s3_input_mode, + s3_data_distribution_type=s3_data_distribution_type, + s3_compression_type=s3_compression_type, + ) + + input_config = {"input_name": input_name} + if not isinstance(app_managed, Unassigned): + input_config["app_managed"] = app_managed + if not isinstance(s3_input, Unassigned): + input_config["s3_input"] = s3_input + if not isinstance(dataset_definition, Unassigned): + input_config["dataset_definition"] = dataset_definition + + super().__init__(**input_config) + + @property + def source(self): + """Return the configured local path or S3 URI.""" + return self.s3_input.s3_uri if self.s3_input else None + + @property + def destination(self): + """Return the configured destination path in the processing container.""" + return self.s3_input.local_path if self.s3_input else None diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index f08450e5ee..54af571e0e 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -64,6 +64,41 @@ def test_normalize_args_with_pipeline_variable_code(self, mock_session): class TestProcessorNormalizeInputs: + def test_normalize_inputs_uploads_local_source(self, mock_session, tmp_path): + processor = Processor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + ) + processor._current_job_name = "test-job" + local_source = tmp_path / "input.csv" + local_source.write_text("feature,label\n1,0\n", encoding="utf-8") + processing_input = ProcessingInput( + source=str(local_source), + destination="/opt/ml/processing/input", + input_name="training-data", + ) + + with patch( + "sagemaker.core.s3.S3Uploader.upload", return_value="s3://bucket/uploaded/input.csv" + ) as mock_upload: + result = processor._normalize_inputs([processing_input]) + + mock_upload.assert_called_once_with( + local_path=str(local_source), + desired_s3_uri="s3://test-bucket/sagemaker/test-job/input/training-data", + sagemaker_session=mock_session, + kms_key=None, + ) + assert result[0].source == "s3://bucket/uploaded/input.csv" + assert result[0].destination == "/opt/ml/processing/input" + + def test_processing_input_source_rejects_missing_destination(self): + with pytest.raises(ValueError, match="destination is required"): + ProcessingInput(source="./data") + def test_normalize_inputs_with_dataset_definition(self, mock_session): from sagemaker.core.shapes import DatasetDefinition, AthenaDatasetDefinition