Skip to content

Commit 2c1899a

Browse files
mufaddal-rohawalaknikure
authored andcommitted
fix: support kms key in processor pack local code (aws#3333)
1 parent 0a2e9dd commit 2c1899a

File tree

2 files changed

+6
-3
lines changed

2 files changed

+6
-3
lines changed

src/sagemaker/processing.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1631,7 +1631,7 @@ def run( # type: ignore[override]
16311631
:class:`~sagemaker.workflow.pipeline_context.PipelineSession`
16321632
"""
16331633
s3_runproc_sh, inputs, job_name = self._pack_and_upload_code(
1634-
code, source_dir, dependencies, git_config, job_name, inputs
1634+
code, source_dir, dependencies, git_config, job_name, inputs, kms_key
16351635
)
16361636

16371637
# Submit a processing job.
@@ -1647,7 +1647,9 @@ def run( # type: ignore[override]
16471647
kms_key=kms_key,
16481648
)
16491649

1650-
def _pack_and_upload_code(self, code, source_dir, dependencies, git_config, job_name, inputs):
1650+
def _pack_and_upload_code(
1651+
self, code, source_dir, dependencies, git_config, job_name, inputs, kms_key=None
1652+
):
16511653
"""Pack local code bundle and upload to Amazon S3."""
16521654
if code.startswith("s3://"):
16531655
return code, inputs, job_name
@@ -1685,6 +1687,7 @@ def _pack_and_upload_code(self, code, source_dir, dependencies, git_config, job_
16851687
s3_runproc_sh = S3Uploader.upload_string_as_file_body(
16861688
self._generate_framework_script(script),
16871689
desired_s3_uri=entrypoint_s3_uri,
1690+
kms_key=kms_key,
16881691
sagemaker_session=self.sagemaker_session,
16891692
)
16901693
logger.info("runproc.sh uploaded to %s", s3_runproc_sh)

tests/unit/sagemaker/workflow/test_pipeline_session.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ def test_pipeline_session_init(sagemaker_client_config, boto_session):
5050
sagemaker_client=sagemaker_client,
5151
)
5252
assert sess.sagemaker_client is not None
53-
assert sess.default_bucket() is not None
53+
assert sess.default_bucket is not None
5454
assert sess.context is None
5555

5656

0 commit comments

Comments
 (0)