change: Do lazy initialization in predictor (#2206)

Co-authored-by: Ahsan Khan <[email protected]>
Co-authored-by: Dan <[email protected]>

Author: 3 people

src/sagemaker/predictor.py

@@ -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.
 
@@ -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 {}
@@ -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]
 
@@ -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,
@@ -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.
@@ -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
@@ -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):

tests/integ/test_mxnet.py

@@ -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"]

tests/unit/test_predictor.py

@@ -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,