fix: honor 'wait' flag when updating endpoint (#1222)

Co-authored-by: Nadia Yakimakha <[email protected]>

Author: laurenyu

src/sagemaker/model.py

@@ -473,7 +473,9 @@ def deploy(
                 kms_key=kms_key,
                 data_capture_config_dict=data_capture_config_dict,
             )
-            self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
+            self.sagemaker_session.update_endpoint(
+                self.endpoint_name, endpoint_config_name, wait=wait
+            )
         else:
             self.sagemaker_session.endpoint_from_production_variants(
                 name=self.endpoint_name,

src/sagemaker/multidatamodel.py

@@ -243,7 +243,9 @@ def deploy(
                 kms_key=kms_key,
                 data_capture_config_dict=data_capture_config_dict,
             )
-            self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
+            self.sagemaker_session.update_endpoint(
+                self.endpoint_name, endpoint_config_name, wait=wait
+            )
         else:
             self.sagemaker_session.endpoint_from_production_variants(
                 name=self.endpoint_name,

src/sagemaker/pipeline.py

@@ -159,7 +159,9 @@ def deploy(
                 tags=tags,
                 data_capture_config_dict=data_capture_config_dict,
             )
-            self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
+            self.sagemaker_session.update_endpoint(
+                self.endpoint_name, endpoint_config_name, wait=wait
+            )
         else:
             self.sagemaker_session.endpoint_from_production_variants(
                 name=self.endpoint_name,

src/sagemaker/session.py

@@ -2363,8 +2363,8 @@ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
             self.wait_for_endpoint(endpoint_name)
         return endpoint_name
 
-    def update_endpoint(self, endpoint_name, endpoint_config_name):
-        """ Update an Amazon SageMaker ``Endpoint`` according to the endpoint configuration
+    def update_endpoint(self, endpoint_name, endpoint_config_name, wait=True):
+        """Update an Amazon SageMaker ``Endpoint`` according to the endpoint configuration
         specified in the request
 
         Raise an error if endpoint with endpoint_name does not exist.
@@ -2373,10 +2373,14 @@ def update_endpoint(self, endpoint_name, endpoint_config_name):
             endpoint_name (str): Name of the Amazon SageMaker ``Endpoint`` to update.
             endpoint_config_name (str): Name of the Amazon SageMaker endpoint configuration to
                 deploy.
+            wait (bool): Whether to wait for the endpoint deployment to complete before returning
+                (default: True).
 
         Returns:
             str: Name of the Amazon SageMaker ``Endpoint`` being updated.
 
+        Raises:
+            ValueError: if the endpoint does not already exist
         """
         if not _deployment_entity_exists(
             lambda: self.sagemaker_client.describe_endpoint(EndpointName=endpoint_name)
@@ -2389,6 +2393,9 @@ def update_endpoint(self, endpoint_name, endpoint_config_name):
         self.sagemaker_client.update_endpoint(
             EndpointName=endpoint_name, EndpointConfigName=endpoint_config_name
         )
+
+        if wait:
+            self.wait_for_endpoint(endpoint_name)
         return endpoint_name
 
     def delete_endpoint(self, endpoint_name):

tests/unit/test_model.py

@@ -346,32 +346,66 @@ def test_deploy_creates_correct_session(local_session, session, tmpdir):
 @patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
 def test_deploy_update_endpoint(sagemaker_session, tmpdir):
     model = DummyFrameworkModel(sagemaker_session, source_dir=tmpdir)
+    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1, update_endpoint=True)
+    sagemaker_session.create_endpoint_config.assert_called_with(
+        name=model.name,
+        model_name=model.name,
+        initial_instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        accelerator_type=None,
+        tags=None,
+        kms_key=None,
+        data_capture_config_dict=None,
+    )
+    config_name = sagemaker_session.create_endpoint_config(
+        name=model.name,
+        model_name=model.name,
+        initial_instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        accelerator_type=ACCELERATOR_TYPE,
+    )
+    sagemaker_session.update_endpoint.assert_called_with(model.name, config_name, wait=True)
+    sagemaker_session.create_endpoint.assert_not_called()
+
+
+@patch("sagemaker.fw_utils.tar_and_upload_dir", MagicMock())
+def test_deploy_update_endpoint_optional_args(sagemaker_session, tmpdir):
     endpoint_name = "endpoint-name"
+    tags = [{"Key": "Value"}]
+    kms_key = "foo"
+    data_capture_config = MagicMock()
+
+    model = DummyFrameworkModel(sagemaker_session, source_dir=tmpdir)
     model.deploy(
         instance_type=INSTANCE_TYPE,
         initial_instance_count=1,
-        endpoint_name=endpoint_name,
         update_endpoint=True,
+        endpoint_name=endpoint_name,
         accelerator_type=ACCELERATOR_TYPE,
+        tags=tags,
+        kms_key=kms_key,
+        wait=False,
+        data_capture_config=data_capture_config,
     )
     sagemaker_session.create_endpoint_config.assert_called_with(
         name=model.name,
         model_name=model.name,
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         accelerator_type=ACCELERATOR_TYPE,
-        tags=None,
-        kms_key=None,
-        data_capture_config_dict=None,
+        tags=tags,
+        kms_key=kms_key,
+        data_capture_config_dict=data_capture_config._to_request_dict(),
     )
     config_name = sagemaker_session.create_endpoint_config(
         name=model.name,
         model_name=model.name,
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
         accelerator_type=ACCELERATOR_TYPE,
+        wait=False,
     )
-    sagemaker_session.update_endpoint.assert_called_with(endpoint_name, config_name)
+    sagemaker_session.update_endpoint.assert_called_with(endpoint_name, config_name, wait=False)
     sagemaker_session.create_endpoint.assert_not_called()
 
 

tests/unit/test_multidatamodel.py

@@ -295,7 +295,9 @@ def test_deploy_model_update(sagemaker_session):
         instance_type=INSTANCE_TYPE,
         accelerator_type=None,
     )
-    sagemaker_session.update_endpoint.assert_called_with(MULTI_MODEL_ENDPOINT_NAME, config_name)
+    sagemaker_session.update_endpoint.assert_called_with(
+        MULTI_MODEL_ENDPOINT_NAME, config_name, wait=True
+    )
     sagemaker_session.create_endpoint.assert_not_called()
 
 

tests/unit/test_pipeline_model.py

@@ -193,7 +193,7 @@ def test_deploy_update_endpoint(tfo, time, sagemaker_session):
         initial_instance_count=INSTANCE_COUNT,
         instance_type=INSTANCE_TYPE,
     )
-    sagemaker_session.update_endpoint.assert_called_with(endpoint_name, config_name)
+    sagemaker_session.update_endpoint.assert_called_with(endpoint_name, config_name, wait=True)
     sagemaker_session.create_endpoint.assert_not_called()
 
 

tests/unit/test_session.py

@@ -1785,6 +1785,18 @@ def test_update_endpoint_succeed(sagemaker_session):
     assert returned_endpoint_name == endpoint_name
 
 
+def test_update_endpoint_no_wait(sagemaker_session):
+    sagemaker_session.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "Updating"}
+    )
+    endpoint_name = "some-endpoint"
+    endpoint_config = "some-endpoint-config"
+    returned_endpoint_name = sagemaker_session.update_endpoint(
+        endpoint_name, endpoint_config, wait=False
+    )
+    assert returned_endpoint_name == endpoint_name
+
+
 def test_update_endpoint_non_existing_endpoint(sagemaker_session):
     error = ClientError(
         {"Error": {"Code": "ValidationException", "Message": "Could not find entity"}}, "foo"