fix: RepackStep must use the same KMS key as the Model (#3692)

Co-authored-by: Guillem Orellana Trullols <[email protected]>

Author: Guillem96

src/sagemaker/workflow/model_step.py

@@ -268,6 +268,7 @@ def _append_repack_model_step(self):
                     depends_on=self.depends_on,
                     retry_policies=self._repack_model_retry_policies,
                     output_path=self._runtime_repack_output_prefix,
+                    output_kms_key=model.model_kms_key,
                 )
                 self.steps.append(repack_model_step)
 

src/sagemaker/workflow/step_collections.py

@@ -400,6 +400,7 @@ def __init__(
                 security_group_ids=estimator.security_group_ids,
                 description=description,
                 display_name=display_name,
+                output_kms_key=estimator.output_kms_key,
             )
             steps.append(repack_model_step)
             model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts