fix: keep job name in pipeline request for ModelTrainer and HyperparameterTuner (#6323)

Both unconditionally popped training_job_name / hyper_parameter_tuning_job_name from
the request for a PipelineSession. The TrainingStep/TuningStep already strip it via
trim_request_dict (dropped by default, prefix preserved when
PipelineDefinitionConfig(use_custom_job_prefix=True)), so popping upstream left
use_custom_job_prefix nothing to preserve and made base_job_name have no effect in
pipelines. Stop popping; let trim_request_dict handle it.

Fixes #5776, #6299

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

Author: mohamedzeidan2021

sagemaker-train/src/sagemaker/train/model_trainer.py

@@ -33,6 +33,7 @@
 from sagemaker.core.shapes import AlgorithmSpecification, ModelPackageConfig
 from sagemaker.core.utils.utils import serialize
 from sagemaker.core.apiutils._boto_functions import to_pascal_case
+from sagemaker.core.common_utils import name_from_base
 
 from pydantic import BaseModel, ConfigDict, PrivateAttr, validate_call
 from sagemaker.core.config.config_schema import (
@@ -76,7 +77,6 @@
 from sagemaker.train.distributed import Torchrun, DistributedConfig
 from sagemaker.train.utils import (
     _default_s3_uri,
-    _get_unique_name,
     _is_valid_path,
     _is_valid_s3_uri,
     safe_serialize,
@@ -679,7 +679,13 @@ def _create_training_job_args(
             Dict[str, Any]: The training job arguments.
         """
         self._populate_intelligent_defaults()
-        current_training_job_name = _get_unique_name(self.base_job_name)
+        # Use the shared name_from_base timestamp format so that
+        # sagemaker.core.common_utils.base_from_name can strip it back to base_job_name.
+        # trim_request_dict relies on base_from_name to preserve the prefix when a pipeline
+        # sets PipelineDefinitionConfig(use_custom_job_prefix=True) (issues #5776, #6299).
+        # Underscores are replaced because SageMaker job names disallow them; _get_unique_name
+        # used to do this and name_from_base does not.
+        current_training_job_name = name_from_base(self.base_job_name.replace("_", "-"))
         # 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
@@ -880,8 +886,10 @@ def _create_training_job_args(
             training_request["model_package_config"] = self.model_package_config
 
         if boto3 or isinstance(self.sagemaker_session, PipelineSession):
-            if isinstance(self.sagemaker_session, PipelineSession):
-                training_request.pop("training_job_name", None)
+            # Keep training_job_name in the request. The TrainingStep strips it via
+            # trim_request_dict, which drops it by default but preserves the base_job_name
+            # prefix when PipelineDefinitionConfig(use_custom_job_prefix=True). Popping it here
+            # unconditionally left use_custom_job_prefix nothing to preserve (issue #5776, #6299).
             # ``serialize`` renames the fields of Tag models to the boto-style
             # {"Key": ..., "Value": ...} form the service requires, but passes plain dicts
             # through untouched, so hand it the validated Tag objects. ``TrainingJob.create``

sagemaker-train/src/sagemaker/train/tuner.py

@@ -1321,8 +1321,10 @@ def _start_tuning_job(self, inputs):
             from sagemaker.core.utils.utils import serialize
             from sagemaker.core.apiutils._boto_functions import to_pascal_case
 
-            # Remove job name for pipeline as it's auto-generated at execution time
-            tuning_request.pop("hyper_parameter_tuning_job_name", None)
+            # Keep hyper_parameter_tuning_job_name in the request. The TuningStep strips it via
+            # trim_request_dict, which drops it by default but preserves the base_tuning_job_name
+            # prefix when PipelineDefinitionConfig(use_custom_job_prefix=True). Popping it here
+            # unconditionally left use_custom_job_prefix nothing to preserve (issue #6299, #5776).
             # Convert snake_case to PascalCase for AWS API
             pipeline_request = {to_pascal_case(k): v for k, v in tuning_request.items()}
             serialized_request = serialize(pipeline_request)

sagemaker-train/tests/unit/train/test_model_trainer.py

@@ -257,6 +257,60 @@ def test_model_trainer_param_validation(test_case, modules_session):
         assert trainer.base_job_name == DEFAULT_BASE_NAME
 
 
+def test_pipeline_session_request_keeps_training_job_name(modules_session):
+    """Regression for #5776 / #6299.
+
+    Under a PipelineSession the request must KEEP TrainingJobName. The TrainingStep strips it via
+    trim_request_dict (dropped by default, prefix preserved when use_custom_job_prefix=True);
+    popping it here left use_custom_job_prefix nothing to preserve.
+    """
+    from unittest.mock import Mock
+    from sagemaker.core.workflow.pipeline_context import PipelineSession
+
+    session = Mock(spec=PipelineSession)
+    session.default_bucket.return_value = DEFAULT_BUCKET
+    session.default_bucket_prefix = DEFAULT_BUCKET_PREFIX
+    session.boto_region_name = DEFAULT_REGION
+    session.sagemaker_config = {}
+
+    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,
+        base_job_name="my-prefix",
+        sagemaker_session=session,
+    )
+
+    args = trainer._create_training_job_args(input_data_config=[])
+
+    assert "TrainingJobName" in args, "TrainingJobName must survive into the pipeline request"
+    assert args["TrainingJobName"].startswith("my-prefix")
+
+    # The generated name must use a timestamp format base_from_name can strip, otherwise
+    # trim_request_dict bakes a definition-time timestamp into the pipeline definition
+    # instead of preserving the base prefix.
+    from sagemaker.core.common_utils import base_from_name
+
+    assert base_from_name(args["TrainingJobName"]) == "my-prefix"
+
+    # End-to-end: with use_custom_job_prefix=True, trim_request_dict must leave exactly the base.
+    from types import SimpleNamespace
+    from sagemaker.core.workflow.utilities import trim_request_dict
+    from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig
+
+    custom_prefix_config = SimpleNamespace(
+        pipeline_definition_config=PipelineDefinitionConfig(use_custom_job_prefix=True)
+    )
+    trimmed = trim_request_dict(dict(args), "TrainingJobName", custom_prefix_config)
+    assert trimmed["TrainingJobName"] == "my-prefix"
+
+    # And by default (no config) the key is dropped entirely.
+    dropped = trim_request_dict(dict(args), "TrainingJobName", None)
+    assert "TrainingJobName" not in dropped
+
+
 @patch("sagemaker.train.model_trainer.TrainingJob")
 @patch("sagemaker.train.model_trainer.ModelTrainer.create_input_data_channel")
 def test_sm_drivers_channel_ignores_user_ignore_patterns(
@@ -1105,7 +1159,7 @@ def test_remote_debug_config(mock_training_job, modules_session):
         )
 
 
-@patch("sagemaker.train.model_trainer._get_unique_name")
+@patch("sagemaker.train.model_trainer.name_from_base")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_model_trainer_full_init(mock_training_job, mock_unique_name, modules_session):
     def mock_upload_data(path, bucket, key_prefix, extra_args=None):
@@ -1381,7 +1435,7 @@ def mock_upload_data(path, bucket, key_prefix, extra_args=None):
 
 # TODO: Re-Enable after local mode fully migrated to v3
 # @patch("sagemaker.train.model_trainer._LocalContainer")
-# @patch("sagemaker.train.model_trainer._get_unique_name")
+# @patch("sagemaker.train.model_trainer.name_from_base")
 # @patch("sagemaker.train.local.local_container.download_folder")
 # def test_model_trainer_local_full_init(
 #     mock_download_folder, mock_unique_name, mock_local_container, modules_session
@@ -1637,7 +1691,7 @@ def test_hyperparameters_invalid(mock_exists, modules_session):
             )
 
 
-@patch("sagemaker.train.model_trainer._get_unique_name")
+@patch("sagemaker.train.model_trainer.name_from_base")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_model_trainer_default_paths(mock_training_job, mock_unique_name, modules_session):
     def mock_upload_data(path, bucket, key_prefix, extra_args=None):
@@ -1828,7 +1882,7 @@ def test_metric_definitions(mock_training_job, modules_session):
         )
 
 
-@patch("sagemaker.train.model_trainer._get_unique_name")
+@patch("sagemaker.train.model_trainer.name_from_base")
 @patch("sagemaker.core.resources.TrainingJob")
 def test_nova_recipe(mock_training_job, mock_unique_name, modules_session):
     def mock_upload_data(path, bucket, key_prefix, extra_args=None):
@@ -2203,7 +2257,7 @@ def test_nova_recipe_model_package_config_only_mpg_from_recipe(modules_session):
         os.unlink(recipe.name)
 
 
-@patch("sagemaker.train.model_trainer._get_unique_name")
+@patch("sagemaker.train.model_trainer.name_from_base")
 @patch("sagemaker.train.model_trainer.TrainingJob")
 def test_llmft_recipe(mock_training_job, mock_unique_name, modules_session):
     def mock_upload_data(path, bucket, key_prefix, extra_args=None):

sagemaker-train/tests/unit/train/test_tuner.py

@@ -704,6 +704,40 @@ def test_build_training_job_definition_includes_spot_params(self):
             definition.stopping_condition.max_wait_time_in_seconds, int
         ), "Max wait time should be set"
 
+    def test_pipeline_session_request_keeps_tuning_job_name(self):
+        """Regression for #6299 / #5776.
+
+        Under a PipelineSession the tuning request must KEEP HyperParameterTuningJobName. The
+        TuningStep strips it via trim_request_dict (dropped by default, prefix preserved when
+        use_custom_job_prefix=True); popping it here left use_custom_job_prefix nothing to preserve.
+        """
+        from unittest.mock import Mock
+        from sagemaker.core.workflow.pipeline_context import PipelineSession
+
+        pipeline_session = Mock(spec=PipelineSession)
+        mock_trainer = _create_mock_model_trainer()
+        mock_trainer.sagemaker_session = pipeline_session
+
+        tuner = HyperparameterTuner(
+            model_trainer=mock_trainer,
+            objective_metric_name="accuracy",
+            hyperparameter_ranges=_create_single_hp_range(),
+        )
+        tuner._current_job_name = "my-prefix-2026-01-01-00-00-00-000"
+
+        tuner._start_tuning_job(inputs=None)
+
+        pipeline_session._intercept_create_request.assert_called_once()
+        serialized_request = pipeline_session._intercept_create_request.call_args.args[0]
+        assert "HyperParameterTuningJobName" in serialized_request
+        assert serialized_request["HyperParameterTuningJobName"].startswith("my-prefix")
+
+        # The name must be strippable back to the base, or trim_request_dict would preserve a
+        # definition-time timestamp instead of the prefix when use_custom_job_prefix=True.
+        from sagemaker.core.common_utils import base_from_name
+
+        assert base_from_name(serialized_request["HyperParameterTuningJobName"]) == "my-prefix"
+
     def test_build_training_job_definition_includes_environment_variables(self):
         """Test that _build_training_job_definition includes environment variables.