aws / aws/sagemaker-python-sdk

ModelTrainer (V3) does not support output_kms_key for source code uploads

Ouverte
#5,956 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
Langage dominant
Python
Étoiles
2.3k
Forks
1.3k
Merge moyen
1 j 22 h
PR mergées (30 j)
35

Description

## Describe the bug

In SageMaker SDK V2, the `Estimator` uses the `output_kms_key` when uploading user training scripts to S3 (see `_stage_user_code_in_s3`):
https://github.com/aws/sagemaker-python-sdk/blob/5b3b127a2d3d12a6ff0d877ddfd9ddf13527c27f/src/sagemaker/estimator.py#L1044-L1108

In SageMaker SDK V3, the new `ModelTrainer` does not apply any KMS key when uploading source code while creating the input data channel:
https://github.com/aws/sagemaker-python-sdk/blob/9101cef1a589cf2b44e73c5456dfa46c10e32691/sagemaker-train/src/sagemaker/train/model_trainer.py#L834-L953

Additionally, even when providing an S3 URI in the `SourceCode` object (instead of a local path), the `ModelTrainer` still uploads additional driver files whenever source code is specified:
https://github.com/aws/sagemaker-python-sdk/blob/9101cef1a589cf2b44e73c5456dfa46c10e32691/sagemaker-train/src/sagemaker/train/model_trainer.py#L684-L689

These uploads do not use a KMS key.

## Expected behavior

`ModelTrainer` should:
- either respect a user-provided KMS key (similar to `output_kms_key` in V2), OR
- allow configuration of a KMS key for all S3 uploads related to source code.

## Actual behavior

- Source code and driver files are uploaded to S3 without KMS encryption.
- There is no apparent way to configure a KMS key for these uploads.

## Impact

In restricted environments (like ours), S3 policies enforce server-side encryption with KMS.
As a result, `ModelTrainer` cannot be used with custom training scripts.

This blocks use cases that rely on custom code, such as MLflow serverless integration.

## Steps to reproduce

1. Create a `ModelTrainer` with a `SourceCode` object
2. Provide either:
- a local path, or
- an S3 URI
3. Observe that S3 uploads occur without KMS encryption

## Possible solution

Expose a parameter similar to `output_kms_key` in V2, or reuse existing encryption configuration mechanisms.

## Additional context

This is a regression compared to V2 `Estimator` behavior and impacts secure environments with strict S3 encryption policies.

Guide de contribution

Ouvrir le guide de contribution

Piste de recherche

Commencez dans sagemaker-train/src/sagemaker/train/model_trainer.py, autour des chemins d’upload de l’input-channel et du driver-file cités dans l’issue, puis comparez-les avec _stage_user_code_in_s3 dans src/sagemaker/estimator.py. Suivez la configuration de chiffrement existante et ajoutez une couverture ciblée pour les uploads de SourceCode locaux et S3. La tâche est terminée lorsque tous les uploads liés au code source peuvent utiliser la clé KMS configurée.

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é
3/5
Temps estimé
1-2 jours
Activité
Calme
Clarté
Plutôt claire
Accessibilité débutants
68/100

Recevez les nouvelles issues par e-mail

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