Skip to content

fix: null checks for uploaded_code and entry_point #2535

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.

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
merged 6 commits into from
Jul 28, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 11 additions & 4 deletions src/sagemaker/estimator.py
Original file line number Diff line number Diff line change
Expand Up @@ -2319,9 +2319,13 @@ def _model_source_dir(self):
str: Either a local or an S3 path pointing to the ``source_dir`` to be
used for code by the model to be deployed
"""
return (
self.source_dir if self.sagemaker_session.local_mode else self.uploaded_code.s3_prefix
)
if self.sagemaker_session.local_mode:
return self.source_dir

if self.uploaded_code is not None:
return self.uploaded_code.s3_prefix

return None

def _model_entry_point(self):
"""Get the appropriate value to pass as ``entry_point`` to a model constructor.
Expand All @@ -2333,7 +2337,10 @@ def _model_entry_point(self):
if self.sagemaker_session.local_mode or (self._model_source_dir() is None):
return self.entry_point

return self.uploaded_code.script_name
if self.uploaded_code is not None:
return self.uploaded_code.script_name

return None

def hyperparameters(self):
"""Return the hyperparameters as a dictionary to use for training.
Expand Down
3 changes: 3 additions & 0 deletions src/sagemaker/workflow/_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,9 @@ def arguments(self) -> RequestType:
model = self.estimator.create_model(**self.kwargs)
self.image_uri = model.image_uri

if self.model_data is None:
self.model_data = model.model_data

# reset placeholder
self.estimator.output_path = output_path

Expand Down
3 changes: 2 additions & 1 deletion src/sagemaker/workflow/step_collections.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ def __init__(
else:
sagemaker_session = model_entity.sagemaker_session
role = model_entity.role
if hasattr(model_entity, "entry_point"):
if hasattr(model_entity, "entry_point") and model_entity.entry_point is not None:
repack_model = True
entry_point = model_entity.entry_point
source_dir = model_entity.source_dir
Expand All @@ -169,6 +169,7 @@ def __init__(
model_entity.model_data = (
repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
)

# remove kwargs consumed by model repacking step
kwargs.pop("output_kms_key", None)

Expand Down