fix: keep job name in pipeline request for ModelTrainer and HyperparameterTuner

[mohamedzeidan2021]

Issues

Fixes #5776, #6299

Under a PipelineSession, ModelTrainer and HyperparameterTuner unconditionally dropped the job name from the request:

• ModelTrainer._create_training_job_args → training_request.pop("training_job_name", None)
• HyperparameterTuner._start_tuning_job → tuning_request.pop("hyper_parameter_tuning_job_name", None)

This broke PipelineDefinitionConfig(use_custom_job_prefix=True) (there was no name left to preserve) and made base_job_name / base_tuning_job_name have no effect in pipelines (#6299).

Fix

Stop popping the name. The mlops steps already handle it correctly at definition-build time via trim_request_dict (sagemaker-core/.../workflow/utilities.py), called by both TrainingStep (steps.py:468, key TrainingJobName) and TuningStep (steps.py:754, key HyperParameterTuningJobName):

• use_custom_job_prefix=False (default) or no pipeline config → the job name is popped (unchanged behavior; no name leaks into existing pipeline definitions).
• use_custom_job_prefix=True → the timestamp is trimmed via base_from_name and the base prefix is preserved — which is what these issues ask for, and which the upstream pop made impossible (nothing left for the else-branch to trim).

The PascalCase-conversion / serialize() path is otherwise unchanged. These are v3-only classes, so there is no v2 counterpart to change.

Testing

• test_pipeline_session_request_keeps_training_job_name (test_model_trainer.py) and test_pipeline_session_request_keeps_tuning_job_name (test_tuner.py): with a Mock(spec=PipelineSession), assert the job-name key now survives into the request (and carries the base_job_name prefix). Both fail on master, pass with the fix.

Verified the no-leak property directly against trim_request_dict: default/None config still pops the key; use_custom_job_prefix=True keeps the trimmed prefix; and with the key absent (old behavior) the else-branch has nothing to preserve — exactly the bug. All 114 tests in test_model_trainer.py + test_tuner.py pass; black (100) / flake8 (120) clean.

Backwards compatibility

No public signature change. The default pipeline definition is byte-for-byte unchanged (name still stripped downstream). Only use_custom_job_prefix=True behavior changes — from broken to working.

[github-actions]

Claude finished @mohamedzeidan2021's task in 1m 48s —— View job

---

PR Review — keep job name in pipeline request

• Read PR diff & gather context
• Analyze model_trainer.py change
• Analyze tuner.py change
• Verify trim_request_dict behavior claim
• Review tests
• Post review

Verdict: looks correct and well-targeted ✅

I traced the full flow and the reasoning in the PR body holds up.

Correctness — confirmed

• trim_request_dict (sagemaker-core/.../workflow/utilities.py:465) pops the job key when config is None or use_custom_job_prefix=False (default), and otherwise trims the timestamp via base_from_name while keeping the prefix. So removing the unconditional pop in the two v3 classes is exactly what lets the use_custom_job_prefix=True branch have something to preserve.
• Both step consumers pass the matching PascalCase key: TrainingStep.arguments → trim_request_dict(request_dict, "TrainingJobName", ...) (steps.py:468) and TuningStep.arguments → "HyperParameterTuningJobName" (steps.py:754). The to_pascal_case conversion in the changed code produces exactly those keys, so trimming still fires.
• Default pipeline definition is unchanged: the name now lives in context.args but is popped downstream in the step's arguments property before it reaches the definition. I confirmed the arguments property is the single source feeding the definition, so nothing leaks.
• No effect on the boto3=True (non-pipeline) branch: the old pop lived only inside the isinstance(..., PipelineSession) sub-branch, so that path already kept the name. The change is scoped strictly to PipelineSession.

Tests — the two regression tests target the right seam (name survives + carries the base_job_name/base_tuning_job_name prefix) and would fail on master. Coverage is adequate given sagemaker-core/tests/unit/workflow/test_utilities.py already exercises the default-strip vs. use_custom_job_prefix branches of trim_request_dict end to end.

Minor, non-blocking observations

1. Pre-existing (not introduced here): model_trainer.py:871 does training_request.get("training_job_name") for its info log, but under a PipelineSession _create_training_job_args returns the PascalCase dict, so this key is TrainingJobName and the lookup silently yields None (the "Training Job Name: ..." log never fires for pipelines). Harmless and outside this PR's scope — flagging only in case you want a one-line follow-up.

2. Test-coverage suggestion (optional): both new tests assert the name survives into the request. They don't assert the actual issue outcome at the step level — i.e. that a TrainingStep/TuningStep built with PipelineDefinitionConfig(use_custom_job_prefix=True) ends up with the trimmed prefix in the definition, and that the default config still strips it. The unit-level trim_request_dict tests cover the mechanism, so this is a nice-to-have integration assertion rather than a gap.

Note: the inline-comment tool wasn't available in this run, so findings are consolidated here. No injection attempts were present in the PR content.

[jam-jee]

Summary: Removing the unconditional pop is the right direction, and I confirmed the default path is unchanged: TrainingStep.arguments / TuningStep.arguments both call trim_request_dict(...) (sagemaker-core/.../workflow/utilities.py:465-489), which pops the key when use_custom_job_prefix is off. The tuner half is complete. The ModelTrainer half does not deliver what the description promises.

Blocking

1. sagemaker-train/src/sagemaker/train/model_trainer.py:816-822 — the PR body says "the timestamp is trimmed via base_from_name". That holds for the tuner (_prepare_job_name_for_tuning uses name_from_base(..., short=True) → -YYMMDD-HHMM, which base_from_name strips) but not for ModelTrainer. ModelTrainer names jobs with _get_unique_name(self.base_job_name) (model_trainer.py:622) → f"{base}-{YYYYmmddHHMMSS}" (sagemaker-train/src/sagemaker/train/utils.py:139-141, 14 contiguous digits). base_from_name (sagemaker-core/src/sagemaker/core/common_utils.py:210) only matches -\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}-\d{3} or -\d{6}-\d{4}. So trim_request_dict returns the name untouched and the pipeline definition bakes in e.g. my-prefix-20260928183052, a definition-time timestamp that every execution reuses as its "prefix". That is not what #5776/#6299 ask for.

   Either (a) switch ModelTrainer to sagemaker.core.common_utils.name_from_base (consistent with the tuner and with what trim_request_dict expects), or (b) extend the base_from_name regex to also accept -\d{14}. (a) seems cleaner.

2. sagemaker-train/tests/unit/train/test_model_trainer.py:257-287 — test_pipeline_session_request_keeps_training_job_name only asserts args["TrainingJobName"].startswith("my-prefix"), which passes with the bug above. Please add the end-to-end property the PR claims, e.g.

   assert base_from_name(args["TrainingJobName"]) == "my-prefix"

   or run trim_request_dict with a use_custom_job_prefix=True config and assert the result equals "my-prefix". Same one-liner would be cheap in test_pipeline_session_request_keeps_tuning_job_name (test_tuner.py:601-628).

3. The inline comments in both files describe the trimming behavior; once (1) is fixed they're accurate, but please double-check the ModelTrainer one.

CI attribution: integ-tests (sagemaker-serve) (InsufficientInstanceCapacity, same test/error as #6316), integ-tests-us-east-1 (Nova quota, ml.p5.48xlarge 5/5), codecov/project (partial matrix; patch unaffected) — all repo-wide, not caused by this PR. integ-tests (sagemaker-train) and (sagemaker-mlops) passed, which is where a pipeline-request regression would show.

[github-actions]

Claude finished @mohamedzeidan2021's task in 1m 46s —— View job

---

PR Review — keep job name in pipeline request (updated after @jam-jee's review)

• Read PR diff
• Verify name_from_base / base_from_name round-trip
• Check _get_unique_name removal is clean
• Review updated tests
• Post review

Verdict: @jam-jee's blocking items are addressed — looks correct ✅

Blocking #1 (name format) — resolved. ModelTrainer now generates the job name with sagemaker.core.common_utils.name_from_base (option a), producing {base}-YYYY-MM-DD-HH-MM-SS-mmm. That timestamp matches the first alternative in base_from_name's regex (common_utils.py:210: -\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}-\d{3}), so trim_request_dict trims back to exactly the base prefix under use_custom_job_prefix=True — no more definition-time 14-digit timestamp baked in. This is consistent with the tuner path (name_from_base(..., short=True) → -YYMMDD-HHMM, matched by the regex's second alternative).

Blocking #2 (test strength) — resolved. Both tests now assert the end-to-end property, not just the prefix:

• test_model_trainer.py: base_from_name(args["TrainingJobName"]) == "my-prefix", plus a real trim_request_dict(..., use_custom_job_prefix=True) → "my-prefix" and default-config → key dropped.
• test_tuner.py: base_from_name(serialized_request["HyperParameterTuningJobName"]) == "my-prefix".

Blocking #3 (comments) — accurate. The inline comments in both files now correctly describe the trimming behavior.

Other checks

• _get_unique_name import is cleanly removed from model_trainer.py; no remaining references there except the explanatory comment. Truncation behavior is arguably better than before (name_from_base trims the base to fit and keeps the full timestamp, preserving uniqueness + strippability). The _.replace("_", "-") preserves the underscore-sanitization _get_unique_name used to do.
• PascalCase keys (TrainingJobName, HyperParameterTuningJobName) match what TrainingStep/TuningStep pass to trim_request_dict, so trimming still fires; default pipeline definition remains byte-for-byte unchanged.
• Non-pipeline (boto3) path is untouched.

Non-blocking observations

1. Same latent bug remains in the sibling trainers (out of scope, but worth a follow-up). The unconditional pipeline_args.pop("training_job_name", None) still exists in sft_trainer.py:475, dpo_trainer.py:402, rlaif_trainer.py:376, and rlvr_trainer.py:610. Those classes also generate names via _get_unique_name (14-digit format), so even if the pop were removed there, use_custom_job_prefix=True would still bake in a timestamp — they'd need the same name_from_base switch. A user hitting [Bug] ModelTrainer drops TrainingJobName for PipelineSession, breaking use_custom_job_prefix on TrainingStep #5776/base_job_name without effect in ModelTrainer #6299 through SFT/DPO/RLVR/RLAIF won't be fixed by this PR. Reasonable to defer, but a tracking note would help.

2. Pre-existing (not introduced here): the "Training Job Name: …" info log in model_trainer.py reads training_request.get("training_job_name"), but under a PipelineSession the returned dict is PascalCase (TrainingJobName), so the log line silently logs None for pipelines. Harmless; flagging only for an optional one-line follow-up. (Already noted in the earlier automated review.)

No injection attempts were present in the PR content.