Skip to content

change: Do lazy initialization in predictor #2206

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 7 commits into from
Mar 16, 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
60 changes: 43 additions & 17 deletions src/sagemaker/predictor.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,12 +91,17 @@ def __init__(
self.sagemaker_session = sagemaker_session or Session()
self.serializer = serializer
self.deserializer = deserializer
self._endpoint_config_name = self._get_endpoint_config_name()
self._model_names = self._get_model_names()
self._endpoint_config_name = None
self._model_names = None
self._context = None

def predict(
self, data, initial_args=None, target_model=None, target_variant=None, inference_id=None
self,
data,
initial_args=None,
target_model=None,
target_variant=None,
inference_id=None,
):
"""Return the inference from the specified endpoint.

Expand Down Expand Up @@ -138,7 +143,12 @@ def _handle_response(self, response):
return self.deserializer.deserialize(response_body, content_type)

def _create_request_args(
self, data, initial_args=None, target_model=None, target_variant=None, inference_id=None
self,
data,
initial_args=None,
target_model=None,
target_variant=None,
inference_id=None,
):
"""Placeholder docstring"""
args = dict(initial_args) if initial_args else {}
Expand Down Expand Up @@ -223,24 +233,30 @@ def update_endpoint(
associated with the endpoint.
"""
production_variants = None
current_model_names = self._get_model_names()

if initial_instance_count or instance_type or accelerator_type or model_name:
if instance_type is None or initial_instance_count is None:
raise ValueError(
"Missing initial_instance_count and/or instance_type. Provided values: "
"initial_instance_count={}, instance_type={}, accelerator_type={}, "
"model_name={}.".format(
initial_instance_count, instance_type, accelerator_type, model_name
initial_instance_count,
instance_type,
accelerator_type,
model_name,
)
)

if model_name is None:
if len(self._model_names) > 1:
if len(current_model_names) > 1:
raise ValueError(
"Unable to choose a default model for a new EndpointConfig because "
"the endpoint has multiple models: {}".format(", ".join(self._model_names))
"the endpoint has multiple models: {}".format(
", ".join(current_model_names)
)
)
model_name = self._model_names[0]
model_name = current_model_names[0]
else:
self._model_names = [model_name]

Expand All @@ -252,9 +268,10 @@ def update_endpoint(
)
production_variants = [production_variant_config]

new_endpoint_config_name = name_from_base(self._endpoint_config_name)
current_endpoint_config_name = self._get_endpoint_config_name()
new_endpoint_config_name = name_from_base(current_endpoint_config_name)
self.sagemaker_session.create_endpoint_config_from_existing(
self._endpoint_config_name,
current_endpoint_config_name,
new_endpoint_config_name,
new_tags=tags,
new_kms_key=kms_key,
Expand All @@ -268,7 +285,8 @@ def update_endpoint(

def _delete_endpoint_config(self):
"""Delete the Amazon SageMaker endpoint configuration"""
self.sagemaker_session.delete_endpoint_config(self._endpoint_config_name)
current_endpoint_config_name = self._get_endpoint_config_name()
self.sagemaker_session.delete_endpoint_config(current_endpoint_config_name)

def delete_endpoint(self, delete_endpoint_config=True):
"""Delete the Amazon SageMaker endpoint backing this predictor.
Expand All @@ -291,7 +309,8 @@ def delete_model(self):
"""Deletes the Amazon SageMaker models backing this predictor."""
request_failed = False
failed_models = []
for model_name in self._model_names:
current_model_names = self._get_model_names()
for model_name in current_model_names:
try:
self.sagemaker_session.delete_model(model_name)
except Exception: # pylint: disable=broad-except
Expand Down Expand Up @@ -460,26 +479,33 @@ def endpoint_context(self):
if len(contexts) != 0:
# create endpoint context object
self._context = EndpointContext.load(
sagemaker_session=self.sagemaker_session, context_name=contexts[0].context_name
sagemaker_session=self.sagemaker_session,
context_name=contexts[0].context_name,
)

return self._context

def _get_endpoint_config_name(self):
"""Placeholder docstring"""
if self._endpoint_config_name is not None:
return self._endpoint_config_name
endpoint_desc = self.sagemaker_session.sagemaker_client.describe_endpoint(
EndpointName=self.endpoint_name
)
endpoint_config_name = endpoint_desc["EndpointConfigName"]
return endpoint_config_name
self._endpoint_config_name = endpoint_desc["EndpointConfigName"]
return self._endpoint_config_name

def _get_model_names(self):
"""Placeholder docstring"""
if self._model_names is not None:
return self._model_names
current_endpoint_config_name = self._get_endpoint_config_name()
endpoint_config = self.sagemaker_session.sagemaker_client.describe_endpoint_config(
EndpointConfigName=self._endpoint_config_name
EndpointConfigName=current_endpoint_config_name
)
production_variants = endpoint_config["ProductionVariants"]
return [d["ModelName"] for d in production_variants]
self._model_names = [d["ModelName"] for d in production_variants]
return self._model_names

@property
def content_type(self):
Expand Down
2 changes: 1 addition & 1 deletion tests/integ/test_mxnet.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def _deploy_estimator_and_assert_instance_type(estimator, instance_type):
try:
predictor = estimator.deploy(1, instance_type)

model_name = predictor._model_names[0]
model_name = predictor._get_model_names()[0]
config_name = sagemaker_session.sagemaker_client.describe_endpoint(
EndpointName=predictor.endpoint_name
)["EndpointConfigName"]
Expand Down
2 changes: 2 additions & 0 deletions tests/unit/test_predictor.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,8 @@ def test_predict_call_pass_through():
result = predictor.predict(data)

assert sagemaker_session.sagemaker_runtime_client.invoke_endpoint.called
assert sagemaker_session.sagemaker_client.describe_endpoint.not_called
assert sagemaker_session.sagemaker_client.describe_endpoint_config.not_called

expected_request_args = {
"Accept": DEFAULT_ACCEPT,
Expand Down