feat: allow configurable S3 key prefix for ModelTrainer input uploads (#5638) (#6354)

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.

Co-authored-by: Mohamed Zeidan <zeidmo@amazon.com>

Author: mohamedzeidan2021

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://<default_bucket_path>/<input_s3_key_prefix>/<job_name>/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}"

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"