aws / aws/sagemaker-python-sdk
ModelTrainer (V3) does not support output_kms_key for source code uploads
- 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