Skip to content

fix: use the provided kms key to upload code #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Christian-kam merged 1 commit into from
Nov 18, 2020
Merged

fix: use the provided kms key to upload code #1

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions src/sagemaker/processing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,14 +205,14 @@ def _normalize_args(
"""
self._current_job_name = self._generate_current_job_name(job_name=job_name)

inputs_with_code = self._include_code_in_inputs(inputs, code)
inputs_with_code = self._include_code_in_inputs(inputs, code, kms_key)
normalized_inputs = self._normalize_inputs(inputs_with_code, kms_key)
normalized_outputs = self._normalize_outputs(outputs)
self.arguments = arguments

return normalized_inputs, normalized_outputs

def _include_code_in_inputs(self, inputs, _code):
def _include_code_in_inputs(self, inputs, _code, kms_key=None):
"""A no op in the base class to include code in the processing job inputs.

Args:
Expand Down Expand Up @@ -475,7 +475,7 @@ def run(
if wait:
self.latest_job.wait(logs=logs)

def _include_code_in_inputs(self, inputs, code):
def _include_code_in_inputs(self, inputs, code, kms_key=None):
"""Converts code to appropriate input and includes in input list.

Side effects include:
Expand All @@ -493,7 +493,7 @@ def _include_code_in_inputs(self, inputs, code):
list[:class:`~sagemaker.processing.ProcessingInput`]: inputs together with the
code as `ProcessingInput`.
"""
user_code_s3_uri = self._handle_user_code_url(code)
user_code_s3_uri = self._handle_user_code_url(code,kms_key)
user_script_name = self._get_user_code_name(code)

inputs_with_code = self._convert_code_and_add_to_inputs(inputs, user_code_s3_uri)
Expand All @@ -515,7 +515,7 @@ def _get_user_code_name(self, code):
code_url = urlparse(code)
return os.path.basename(code_url.path)

def _handle_user_code_url(self, code):
def _handle_user_code_url(self, code, kms_key=None):
"""Gets the S3 URL containing the user's code.

Inspects the scheme the customer passed in ("s3://" for code in S3, "file://" or nothing
Expand Down Expand Up @@ -551,7 +551,7 @@ def _handle_user_code_url(self, code):
code
)
)
user_code_s3_uri = self._upload_code(code_path)
user_code_s3_uri = self._upload_code(code_path,kms_key)
else:
raise ValueError(
"code {} url scheme {} is not recognized. Please pass a file path or S3 url".format(
Expand All @@ -560,7 +560,7 @@ def _handle_user_code_url(self, code):
)
return user_code_s3_uri

def _upload_code(self, code):
def _upload_code(self, code, kms_key=None):
"""Uploads a code file or directory specified as a string
and returns the S3 URI.

Expand All @@ -579,7 +579,7 @@ def _upload_code(self, code):
self._CODE_CONTAINER_INPUT_NAME,
)
return s3.S3Uploader.upload(
local_path=code, desired_s3_uri=desired_s3_uri, sagemaker_session=self.sagemaker_session
local_path=code, desired_s3_uri=desired_s3_uri, sagemaker_session=self.sagemaker_session, kms_key=kms_key
)

def _convert_code_and_add_to_inputs(self, inputs, s3_uri):
Expand Down