Add update_endpoint support to session class and add optional flag to… (#606)

Author: chuyang-deng

CHANGELOG.rst

@@ -4,8 +4,10 @@ CHANGELOG
 
 
 1.18.2.dev
-======
+==========
+
 * enhancement: Include SageMaker Notebook Instance version number in boto3 user agent, if available.
+* feature: Support for updating existing endpoint
 
 1.18.1
 ======

README.rst

@@ -192,6 +192,37 @@ Here is an end to end example of how to use a SageMaker Estimator:
     # Tears down the SageMaker endpoint
     mxnet_estimator.delete_endpoint()
 
+
+Additionally, it is possible to deploy a different endpoint configuration, which links to your model, to an already existing SageMaker endpoint.
+This can be done by specifying the existing endpoint name for the ``endpoint_name`` parameter along with the ``update_endpoint`` parameter as ``True`` within your ``deploy()`` call.
+For more `information <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.update_endpoint>`__.
+
+.. code:: python
+
+    from sagemaker.mxnet import MXNet
+
+    # Configure an MXNet Estimator (no training happens yet)
+    mxnet_estimator = MXNet('train.py',
+                            role='SageMakerRole',
+                            train_instance_type='ml.p2.xlarge',
+                            train_instance_count=1,
+                            framework_version='1.2.1')
+
+    # Starts a SageMaker training job and waits until completion.
+    mxnet_estimator.fit('s3://my_bucket/my_training_data/')
+
+    # Deploys the model that was generated by fit() to an existing SageMaker endpoint
+    mxnet_predictor = mxnet_estimator.deploy(initial_instance_count=1,
+                                             instance_type='ml.p2.xlarge',
+                                             update_endpoint=True,
+                                             endpoint_name='existing-endpoint')
+
+    # Serializes data and makes a prediction request to the SageMaker endpoint
+    response = mxnet_predictor.predict(data)
+
+    # Tears down the SageMaker endpoint
+    mxnet_estimator.delete_endpoint()
+
 Training Metrics
 ~~~~~~~~~~~~~~~~
 The SageMaker Python SDK allows you to specify a name and a regular expression for metrics you want to track for training.

src/sagemaker/estimator.py

@@ -330,7 +330,7 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name='m
         return estimator
 
     def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,
-               use_compiled_model=False, **kwargs):
+               use_compiled_model=False, update_endpoint=False, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
 
         More information:
@@ -347,6 +347,9 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             endpoint_name (str): Name to use for creating an Amazon SageMaker endpoint. If not specified, the name of
                 the training job is used.
             use_compiled_model (bool): Flag to select whether to use compiled (optimized) model. Default: False.
+            update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint.
+                If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources
+                corresponding to the previous EndpointConfig. Default: False
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy.
                 For more, see the implementation docs.
@@ -370,7 +373,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             instance_type=instance_type,
             initial_instance_count=initial_instance_count,
             accelerator_type=accelerator_type,
-            endpoint_name=endpoint_name)
+            endpoint_name=endpoint_name,
+            update_endpoint=update_endpoint)
 
     @property
     def model_data(self):

src/sagemaker/local/local_session.py

@@ -143,6 +143,9 @@ def create_endpoint(self, EndpointName, EndpointConfigName):
         LocalSagemakerClient._endpoints[EndpointName] = endpoint
         endpoint.serve()
 
+    def update_endpoint(self, EndpointName, EndpointConfigName):  # pylint: disable=unused-argument
+        raise NotImplementedError('Update endpoint name is not supported in local session.')
+
     def delete_endpoint(self, EndpointName):
         if EndpointName in LocalSagemakerClient._endpoints:
             LocalSagemakerClient._endpoints[EndpointName].stop()

src/sagemaker/model.py

@@ -195,7 +195,8 @@ def compile(self, target_instance_family, input_shape, output_path, role,
         self._is_compiled_model = True
         return self
 
-    def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None, tags=None):
+    def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,
+               update_endpoint=False, tags=None):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
 
         Create a SageMaker ``Model`` and ``EndpointConfig``, and deploy an ``Endpoint`` from this ``Model``.
@@ -217,6 +218,9 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
                 For more information: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): The name of the endpoint to create (default: None).
                 If not specified, a unique endpoint name will be created.
+            update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint.
+                If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources
+                corresponding to the previous EndpointConfig. If False, a new endpoint will be created. Default: False
             tags(List[dict[str, str]]): The list of tags to attach to this specific endpoint.
 
         Returns:
@@ -245,7 +249,18 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             self.endpoint_name = self.name
             if self._is_compiled_model and not self.endpoint_name.endswith(compiled_model_suffix):
                 self.endpoint_name += compiled_model_suffix
-        self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)
+
+        if update_endpoint:
+            endpoint_config_name = self.sagemaker_session.create_endpoint_config(
+                name=self.name,
+                model_name=self.name,
+                initial_instance_count=initial_instance_count,
+                instance_type=instance_type,
+                accelerator_type=accelerator_type)
+            self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
+        else:
+            self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)
+
         if self.predictor_cls:
             return self.predictor_cls(self.endpoint_name, self.sagemaker_session)
 

src/sagemaker/session.py

@@ -750,6 +750,26 @@ def create_endpoint(self, endpoint_name, config_name, 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 specified in the request
+
+        Raise an error if endpoint with endpoint_name does not exist.
+
+        Args:
+            endpoint_name (str): Name of the Amazon SageMaker ``Endpoint`` to update.
+            endpoint_config_name (str): Name of the Amazon SageMaker endpoint configuration to deploy.
+
+        Returns:
+            str: Name of the Amazon SageMaker ``Endpoint`` being updated.
+        """
+        if not _deployment_entity_exists(lambda: self.sagemaker_client.describe_endpoint(EndpointName=endpoint_name)):
+            raise ValueError('Endpoint with name "{}" does not exist; please use an existing endpoint name'
+                             .format(endpoint_name))
+
+        self.sagemaker_client.update_endpoint(EndpointName=endpoint_name,
+                                              EndpointConfigName=endpoint_config_name)
+        return endpoint_name
+
     def delete_endpoint(self, endpoint_name):
         """Delete an Amazon SageMaker ``Endpoint``.
 

tests/integ/test_mxnet_train.py

@@ -72,6 +72,48 @@ def test_deploy_model(mxnet_training_job, sagemaker_session):
         predictor.predict(data)
 
 
+def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session):
+    endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+        model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+        script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+        model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+                           py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session)
+        model.deploy(1, 'ml.t2.medium', endpoint_name=endpoint_name)
+        old_endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+        old_config_name = old_endpoint['EndpointConfigName']
+
+        model.deploy(1, 'ml.m4.xlarge', update_endpoint=True, endpoint_name=endpoint_name)
+        new_endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)['ProductionVariants']
+        new_production_variants = new_endpoint['ProductionVariants']
+        new_config_name = new_endpoint['EndpointConfigName']
+
+        assert old_config_name != new_config_name
+        assert new_production_variants['InstanceType'] == 'ml.m4.xlarge'
+        assert new_production_variants['InitialInstanceCount'] == 1
+        assert new_production_variants['AcceleratorType'] is None
+
+
+def test_deploy_model_with_update_non_existing_endpoint(mxnet_training_job, sagemaker_session):
+    endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+    expected_error_message = 'Endpoint with name "{}" does not exist; ' \
+                             'please use an existing endpoint name'.format(endpoint_name)
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+            desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+            model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+            script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+            model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+                               py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session)
+            model.deploy(1, 'ml.t2.medium', endpoint_name=endpoint_name)
+            sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+
+            with pytest.raises(ValueError, message=expected_error_message):
+                model.deploy(1, 'ml.m4.xlarge', update_endpoint=True, endpoint_name='non-existing-endpoint')
+
+
 @pytest.mark.continuous_testing
 @pytest.mark.regional_testing
 @pytest.mark.skipif(tests.integ.test_region() not in tests.integ.EI_SUPPORTED_REGIONS,

tests/unit/test_estimator.py

@@ -990,6 +990,21 @@ def test_generic_deploy_accelerator_type(sagemaker_session):
     assert args[1][0]['InstanceType'] == INSTANCE_TYPE
 
 
+def test_deploy_with_update_endpoint(sagemaker_session):
+    estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
+                          sagemaker_session=sagemaker_session)
+    estimator.set_hyperparameters(**HYPERPARAMS)
+    estimator.fit({'train': 's3://bucket/training-prefix'})
+    endpoint_name = 'endpoint-name'
+    estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_name=endpoint_name, update_endpoint=True)
+
+    update_endpoint_args = sagemaker_session.update_endpoint.call_args[0]
+    assert update_endpoint_args[0] == endpoint_name
+    assert update_endpoint_args[1].startWith(IMAGE_NAME)
+
+    sagemaker_session.create_endpoint.assert_not_called()
+
+
 @patch('sagemaker.estimator.LocalSession')
 @patch('sagemaker.estimator.Session')
 def test_local_mode(session_class, local_session_class):

tests/unit/test_local_session.py

@@ -310,6 +310,16 @@ def test_create_endpoint(describe_model, describe_endpoint_config, request, *arg
     assert 'my-endpoint' in sagemaker.local.local_session.LocalSagemakerClient._endpoints
 
 
+@patch('sagemaker.local.local_session.LocalSession')
+def test_update_endpoint(LocalSession):
+    local_sagemaker_client = sagemaker.local.local_session.LocalSagemakerClient()
+    endpoint_name = 'my-endpoint'
+    endpoint_config = 'my-endpoint-config'
+    expected_error_message = 'Update endpoint name is not supported in local session.'
+    with pytest.raises(NotImplementedError, message=expected_error_message):
+        local_sagemaker_client.update_endpoint(endpoint_name, endpoint_config)
+
+
 @patch('sagemaker.local.image._SageMakerContainer.serve')
 @patch('urllib3.PoolManager.request')
 def test_serve_endpoint_with_correct_accelerator(request, *args):

tests/unit/test_model.py

@@ -224,6 +224,33 @@ def test_deploy_creates_correct_session(local_session, session, tmpdir):
     assert model.sagemaker_session == session.return_value
 
 
+@patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
+def test_deploy_update_endpoint(sagemaker_session, tmpdir):
+    model = DummyFrameworkModel(sagemaker_session, source_dir=tmpdir)
+    endpoint_name = 'endpoint-name'
+    model.deploy(instance_type=INSTANCE_TYPE,
+                 initial_instance_count=1,
+                 endpoint_name=endpoint_name,
+                 update_endpoint=True,
+                 accelerator_type=ACCELERATOR_TYPE)
+    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
+    )
+    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(endpoint_name, config_name)
+    sagemaker_session.create_endpoint.assert_not_called()
+
+
 def test_model_enable_network_isolation(sagemaker_session):
     model = DummyFrameworkModel(sagemaker_session=sagemaker_session)
     assert model.enable_network_isolation() is False

tests/unit/test_session.py

@@ -920,6 +920,23 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
         Tags=tags)
 
 
+def test_update_endpoint_succeed(sagemaker_session):
+    sagemaker_session.sagemaker_client.describe_endpoint = Mock(return_value={'EndpointStatus': 'InService'})
+    endpoint_name = "some-endpoint"
+    endpoint_config = "some-endpoint-config"
+    returned_endpoint_name = sagemaker_session.update_endpoint(endpoint_name, endpoint_config)
+    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')
+    expected_error_message = 'Endpoint with name "non-existing-endpoint" does not exist; ' \
+                             'please use an existing endpoint name'
+    sagemaker_session.sagemaker_client.describe_endpoint = Mock(side_effect=error)
+    with pytest.raises(ValueError, message=expected_error_message):
+        sagemaker_session.update_endpoint("non-existing-endpoint", "non-existing-config")
+
+
 @patch('time.sleep')
 def test_wait_for_tuning_job(sleep, sagemaker_session):
     hyperparameter_tuning_job_desc = {'HyperParameterTuningJobStatus': 'Completed'}