aws / aws/sagemaker-python-sdk

[Bug] ModelTrainer drops TrainingJobName for PipelineSession, breaking use_custom_job_prefix on TrainingStep

Ouverte Adaptée aux débutants
#5,776 1 commentaire 0 réactions 0 personnes assignées Voir sur GitHub
component: pipelines type: bug
Langage dominant
Python
Étoiles
2.3k
Forks
1.3k
Merge moyen
1 j 22 h
PR mergées (30 j)
35

Description

**PySDK Version**
- [ ] PySDK V2 (2.x)
- [x] PySDK V3 (3.x)

**Describe the bug**
When a `ModelTrainer` is executed under a `PipelineSession` (i.e. produces a `TrainingStep`), `ModelTrainer._create_training_job_args` explicitly removes `training_job_name` from the request before serializing it to PascalCase:

```python
# sagemaker/train/model_trainer.py (sagemaker-train 1.8.0)
if boto3 or isinstance(self.sagemaker_session, PipelineSession):
if isinstance(self.sagemaker_session, PipelineSession):
training_request.pop("training_job_name", None)
# Convert snake_case to PascalCase for AWS API
pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
serialized_request = serialize(pipeline_request)
return serialized_request
```

Because the key is popped, the resulting request dict has no `TrainingJobName`. Downstream, `TrainingStep.arguments` (with `PipelineDefinitionConfig(use_custom_job_prefix=True)`) relies on `TrainingJobName` being present in the request so the prefix is preserved (and `trim_request_dict` removes it when `use_custom_job_prefix=False`).

The net effect is that `use_custom_job_prefix=True` is silently ignored for `TrainingStep` when the step is built from a `ModelTrainer`: every pipeline execution produces a random auto-generated training job name instead of the configured `base_job_name` prefix.

This is the same class of bug as #3991 and #4590 (which were about `TransformStep`), but for the new V3 `ModelTrainer` → `TrainingStep` path.

**To reproduce**

```python
from sagemaker.core.workflow.pipeline import Pipeline
from sagemaker.core.workflow.pipeline_context import PipelineSession
from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig
from sagemaker.train.model_trainer import ModelTrainer
# ... build a ModelTrainer `trainer` with base_job_name=\"my-prefix\" ...

pipeline_session = PipelineSession()
trainer.sagemaker_session = pipeline_session

step_args = trainer._create_training_job_args()
assert \"TrainingJobName\" in step_args, step_args # FAILS — key was popped

pipeline = Pipeline(
name=\"repro\",
steps=[...], # TrainingStep built from trainer
sagemaker_session=pipeline_session,
pipeline_definition_config=PipelineDefinitionConfig(use_custom_job_prefix=True),
)
# Pipeline executions will NOT use \"my-prefix-...\" as the training job name.
```

**Expected behavior**
`TrainingJobName` should remain in the request dict so that `PipelineDefinitionConfig(use_custom_job_prefix=True)` produces training jobs named with the configured prefix. When `use_custom_job_prefix=False`, `TrainingStep.arguments`/`trim_request_dict` will strip the key as usual.

A minimal fix is to stop popping the key:

```python
if boto3 or isinstance(self.sagemaker_session, PipelineSession):
pipeline_request = {to_pascal_case(k): v for k, v in training_request.items()}
serialized_request = serialize(pipeline_request)
return serialized_request
```

As a workaround we currently monkey-patch `_create_training_job_args` to re-insert `TrainingJobName = _get_unique_name(self.base_job_name)` when the session is a `PipelineSession`.

**Screenshots or logs**
N/A — silent misbehavior; the pipeline executes but job names use the default random name instead of the configured prefix.

**System information**
- **SageMaker Python SDK version**: sagemaker-train 1.8.0, sagemaker-core 2.8.0, sagemaker-mlops 1.8.0, sagemaker-serve 1.8.0 (also reproduces on 1.7.1 / 2.7.1)
- **Framework name or algorithm**: custom (source_code via ModelTrainer)
- **Framework version**: N/A
- **Python version**: 3.13
- **CPU or GPU**: CPU (irrelevant, bug is SDK-side)
- **Custom Docker image (Y/N)**: Y

**Additional context**
Related closed issues for other step types: #3991 (TransformStep), #4590 (TransformStep/ProcessingStep).

Guide de contribution

Ouvrir le guide de contribution

Piste de recherche

Commencez dans sagemaker/train/model_trainer.py, au niveau de ModelTrainer._create_training_job_args, et reproduisez l’assertion TrainingJobName montrée avec PipelineSession. Vérifiez TrainingStep.arguments et trim_request_dict ; c’est terminé lorsque la requête sérialisée conserve TrainingJobName pour l’utilisation avec use_custom_job_prefix=True, tandis que le cas false le supprime toujours.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Évaluation

Stack technique
aws, python
Domaine
cloud, machine-learning
Type d'issue
Bug
Difficulté
2/5
Temps estimé
1-3 heures
Activité
Calme
Clarté
Clairement spécifiée
Accessibilité débutants
72/100

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.