From a106b901cc698d4d427af3f858d266a85b13679f Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:33:08 -0700 Subject: [PATCH] feat: allow configurable S3 key prefix for ModelTrainer input uploads (#5638) Add an optional ModelTrainer.input_s3_key_prefix field. When set, it is used as the leading segment of the input-data S3 key prefix (e.g. a pipeline name) instead of base_job_name; when None, the upload layout is unchanged. Covers the SDK-managed code/sm_drivers/recipe channels and local input channels via _create_training_job_args, and the standalone create_input_data_channel default branch. --- .../src/sagemaker/train/model_trainer.py | 16 +++- .../tests/unit/train/test_model_trainer.py | 81 +++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 35b1395cce..7d3b40faae 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -206,6 +206,12 @@ class ModelTrainer(BaseModel): The base name for the training job. If not specified, a default name will be generated using the algorithm name or training image name. + input_s3_key_prefix (Optional[str]): + The leading S3 key prefix under which locally-staged input data (source code, + SDK drivers, and local input channels) is uploaded. This is useful for grouping + a training job's artifacts under a custom prefix such as a pipeline name, e.g. + ``s3://///input/``. If not + specified, ``base_job_name`` is used as the leading prefix (unchanged behavior). source_code (Optional[SourceCode]): The source code configuration. This is used to configure the source code for running the training job. @@ -274,6 +280,7 @@ class ModelTrainer(BaseModel): sagemaker_session: Optional[Session] = None role: Optional[str] = None base_job_name: Optional[str] = None + input_s3_key_prefix: Optional[str] = None source_code: Optional[SourceCode] = None distributed: Optional[DistributedConfig] = None compute: Optional[Compute] = None @@ -317,6 +324,7 @@ class ModelTrainer(BaseModel): CONFIGURABLE_ATTRIBUTES: ClassVar[List[str]] = [ "role", "base_job_name", + "input_s3_key_prefix", "source_code", "compute", "networking", @@ -620,7 +628,10 @@ def _create_training_job_args( """ self._populate_intelligent_defaults() current_training_job_name = _get_unique_name(self.base_job_name) - input_data_key_prefix = f"{self.base_job_name}/{current_training_job_name}/input" + # Use input_s3_key_prefix (e.g. a pipeline name) as the leading key prefix when set; + # otherwise fall back to base_job_name to preserve the default upload layout. + input_data_key_prefix_root = self.input_s3_key_prefix or self.base_job_name + input_data_key_prefix = f"{input_data_key_prefix_root}/{current_training_job_name}/input" final_input_data_config = self.input_data_config.copy() if self.input_data_config else [] @@ -1122,7 +1133,8 @@ def create_input_data_channel( key_prefix = ( f"{key_prefix}/{channel_name}" if key_prefix - else f"{self.base_job_name}/input/{channel_name}" + else f"{self.input_s3_key_prefix or self.base_job_name}" + f"/input/{channel_name}" ) if self.sagemaker_session.default_bucket_prefix: key_prefix = f"{self.sagemaker_session.default_bucket_prefix}/{key_prefix}" diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index ceb371dfa6..7ce3640cd1 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -582,6 +582,87 @@ def test_create_input_data_channel_with_instance_group_names(model_trainer): ] +@patch("sagemaker.train.model_trainer.ModelTrainer._resolve_staging_bucket") +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_create_input_data_channel_default_prefix_uses_base_job_name( + mock_default_bucket, mock_upload_data, mock_staging, model_trainer +): + """Without input_s3_key_prefix, the upload key prefix is derived from base_job_name (unchanged).""" + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_staging.return_value = (DEFAULT_BUCKET, None) + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/code" + + assert model_trainer.input_s3_key_prefix is None + model_trainer.create_input_data_channel("code", DEFAULT_SOURCE_DIR) + + # Leading segment is derived from base_job_name when input_s3_key_prefix is unset. + assert mock_upload_data.call_args.kwargs["key_prefix"] == f"{DEFAULT_BASE_NAME}/input/code" + + +@patch("sagemaker.train.model_trainer.ModelTrainer._resolve_staging_bucket") +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_create_input_data_channel_custom_input_s3_key_prefix( + mock_default_bucket, mock_upload_data, mock_staging +): + """input_s3_key_prefix replaces base_job_name as the leading key prefix (issue #5638).""" + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_staging.return_value = (DEFAULT_BUCKET, None) + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/code" + + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + input_s3_key_prefix="my-pipeline/my-step", + ) + trainer.create_input_data_channel("code", DEFAULT_SOURCE_DIR) + + # input_s3_key_prefix replaces base_job_name as the leading segment. + assert mock_upload_data.call_args.kwargs["key_prefix"] == "my-pipeline/my-step/input/code" + assert f"{DEFAULT_BASE_NAME}/input/code" not in mock_upload_data.call_args.kwargs["key_prefix"] + + +@patch("sagemaker.train.model_trainer.ModelTrainer._resolve_staging_bucket") +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_create_training_job_args_input_s3_key_prefix( + mock_default_bucket, mock_upload_data, mock_staging +): + """Managed sm_drivers/code channels use input_s3_key_prefix in their S3 URIs (issue #5638).""" + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_staging.return_value = (DEFAULT_BUCKET, None) + + def _echo(path, bucket, key_prefix, extra_args=None): + return f"s3://{bucket}/{key_prefix}" + + mock_upload_data.side_effect = _echo + + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + source_code=DEFAULT_SOURCE_CODE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + input_s3_key_prefix="my-pipeline", + ) + args = trainer._create_training_job_args() + uris = { + channel.channel_name: channel.data_source.s3_data_source.s3_uri + for channel in args["input_data_config"] + if channel.data_source and channel.data_source.s3_data_source + } + + # The SDK-managed sm_drivers channel is staged from a local temp dir, so its S3 URI + # is built from the input key prefix -> it must lead with the custom pipeline prefix. + assert "/my-pipeline/" in uris["sm_drivers"] + assert f"/{DEFAULT_BASE_NAME}/input" not in uris["sm_drivers"] + + HETEROGENEOUS_INSTANCE_GROUPS = [ InstanceGroup( instance_type="ml.t3.large", instance_count=1, instance_group_name="head-instance-group"