Add support to delete model within Predictor and Pipeline class. (#647)

Author: chuyang-deng

CHANGELOG.rst

@@ -6,6 +6,8 @@ CHANGELOG
 ==========
 
 * doc-fix: Remove incorrect parameter for EI TFS Python README
+* feature: ``Predictor``: delete SageMaker model
+* feature: ``Pipeline``: delete SageMaker model
 
 1.18.3.post1
 ============

README.rst

@@ -192,6 +192,8 @@ Here is an end to end example of how to use a SageMaker Estimator:
     # Tears down the SageMaker endpoint and endpoint configuration
     mxnet_predictor.delete_endpoint()
 
+    # Deletes the SageMaker model
+    mxnet_predictor.delete_model()
 
 The example above will eventually delete both the SageMaker endpoint and endpoint configuration through `delete_endpoint()`. If you want to keep your SageMaker endpoint configuration, use the value False for the `delete_endpoint_config` parameter, as shown below.
 
@@ -230,6 +232,9 @@ For more `information <https://boto3.amazonaws.com/v1/documentation/api/latest/r
     # Tears down the SageMaker endpoint and endpoint configuration
     mxnet_predictor.delete_endpoint()
 
+    # Deletes the SageMaker model
+    mxnet_predictor.delete_model()
+
 Training Metrics
 ~~~~~~~~~~~~~~~~
 The SageMaker Python SDK allows you to specify a name and a regular expression for metrics you want to track for training.
@@ -284,6 +289,9 @@ We can take the example in  `Using Estimators <#using-estimators>`__ , and use e
     # Tears down the endpoint container and deletes the corresponding endpoint configuration
     mxnet_predictor.delete_endpoint()
 
+    # Deletes the model
+    mxnet_predictor.delete_model()
+
 
 If you have an existing model and want to deploy it locally, don't specify a sagemaker_session argument to the ``MXNetModel`` constructor.
 The correct session is generated when you call ``model.deploy()``.
@@ -307,6 +315,9 @@ Here is an end-to-end example:
     # Tear down the endpoint container and delete the corresponding endpoint configuration
     predictor.delete_endpoint()
 
+    # Deletes the model
+    predictor.delete_model()
+
 
 If you don't want to deploy your model locally, you can also choose to perform a Local Batch Transform Job. This is
 useful if you want to test your container before creating a Sagemaker Batch Transform Job. Note that the performance

src/sagemaker/pipeline.py

@@ -103,3 +103,14 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None, tags
         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)
+
+    def delete_model(self):
+        """Delete the SageMaker model backing this pipeline model. This does not delete the list of SageMaker models used
+        in multiple containers to build the inference pipeline.
+
+        """
+
+        if self.name is None:
+            raise ValueError('The SageMaker model must be created before attempting to delete.')
+
+        self.sagemaker_session.delete_model(self.name)

src/sagemaker/predictor.py

@@ -56,6 +56,8 @@ def __init__(self, endpoint, sagemaker_session=None, serializer=None, deserializ
         self.deserializer = deserializer
         self.content_type = content_type or getattr(serializer, 'content_type', None)
         self.accept = accept or getattr(deserializer, 'accept', None)
+        self._endpoint_config_name = self._get_endpoint_config_name()
+        self._model_names = self._get_model_names()
 
     def predict(self, data, initial_args=None):
         """Return the inference from the specified endpoint.
@@ -109,23 +111,51 @@ def _delete_endpoint_config(self):
         """Delete the Amazon SageMaker endpoint configuration
 
         """
-        endpoint_description = self.sagemaker_session.sagemaker_client.describe_endpoint(EndpointName=self.endpoint)
-        endpoint_config_name = endpoint_description['EndpointConfigName']
-        self.sagemaker_session.delete_endpoint_config(endpoint_config_name)
+        self.sagemaker_session.delete_endpoint_config(self._endpoint_config_name)
 
     def delete_endpoint(self, delete_endpoint_config=True):
-        """Delete the Amazon SageMaker endpoint and endpoint configuration backing this predictor.
+        """Delete the Amazon SageMaker endpoint backing this predictor. Also delete the endpoint configuration attached
+        to it if delete_endpoint_config is True.
 
         Args:
-            delete_endpoint_config (bool): Flag to indicate whether to delete the corresponding SageMaker endpoint
-                configuration tied to the endpoint. If False, only the endpoint will be deleted. (default: True)
+            delete_endpoint_config (bool, optional): Flag to indicate whether to delete endpoint configuration together
+                with endpoint. Defaults to True. If True, both endpoint and endpoint configuration will be deleted. If
+                False, only endpoint will be deleted.
 
         """
         if delete_endpoint_config:
             self._delete_endpoint_config()
 
         self.sagemaker_session.delete_endpoint(self.endpoint)
 
+    def delete_model(self):
+        """Deletes the Amazon SageMaker models backing this predictor.
+
+        """
+        request_failed = False
+        failed_models = []
+        for model_name in self._model_names:
+            try:
+                self.sagemaker_session.delete_model(model_name)
+            except Exception:  # pylint: disable=broad-except
+                request_failed = True
+                failed_models.append(model_name)
+
+        if request_failed:
+            raise Exception('One or more models cannot be deleted, please retry. \n'
+                            'Failed models: {}'.format(', '.join(failed_models)))
+
+    def _get_endpoint_config_name(self):
+        endpoint_desc = self.sagemaker_session.sagemaker_client.describe_endpoint(EndpointName=self.endpoint)
+        endpoint_config_name = endpoint_desc['EndpointConfigName']
+        return endpoint_config_name
+
+    def _get_model_names(self):
+        endpoint_config = self.sagemaker_session.sagemaker_client.describe_endpoint_config(
+            EndpointConfigName=self._endpoint_config_name)
+        production_variants = endpoint_config['ProductionVariants']
+        return map(lambda d: d['ModelName'], production_variants)
+
 
 class _CsvSerializer(object):
     def __init__(self):

tests/integ/test_inference_pipeline.py

@@ -92,3 +92,8 @@ def test_inference_pipeline_model_deploy(sagemaker_session):
 
         invalid_data = "1.0,28.0,C,38.0,71.5,1.0"
         assert (predictor.predict(invalid_data) is None)
+
+    model.delete_model()
+    with pytest.raises(Exception) as exception:
+        sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
+        assert 'Could not find model' in str(exception.value)

tests/integ/test_kmeans.py

@@ -75,6 +75,11 @@ def test_kmeans(sagemaker_session):
             assert record.label["closest_cluster"] is not None
             assert record.label["distance_to_cluster"] is not None
 
+    predictor.delete_model()
+    with pytest.raises(Exception) as exception:
+        sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
+        assert 'Could not find model' in str(exception.value)
+
 
 def test_async_kmeans(sagemaker_session):
     training_job_name = ""

tests/integ/test_mxnet_train.py

@@ -72,6 +72,11 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
         data = numpy.zeros(shape=(1, 1, 28, 28))
         predictor.predict(data)
 
+    predictor.delete_model()
+    with pytest.raises(Exception) as exception:
+        sagemaker_session.sagemaker_client.describe_model(ModelName=model.name)
+        assert 'Could not find model' in str(exception.value)
+
 
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
     endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())

tests/unit/test_chainer.py

@@ -45,6 +45,15 @@
 GPU = 'ml.p2.xlarge'
 CPU = 'ml.c4.xlarge'
 
+ENDPOINT_DESC = {
+    'EndpointConfigName': 'test-endpoint'
+}
+
+ENDPOINT_CONFIG_DESC = {
+    'ProductionVariants': [{'ModelName': 'model-1'},
+                           {'ModelName': 'model-2'}]
+}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -54,6 +63,8 @@ def sagemaker_session():
 
     describe = {'ModelArtifacts': {'S3ModelArtifacts': 's3://m/m.tar.gz'}}
     session.sagemaker_client.describe_training_job = Mock(return_value=describe)
+    session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
+    session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
     session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     session.expand_role = Mock(name="expand_role", return_value=ROLE)
     return session

tests/unit/test_estimator.py

@@ -102,6 +102,15 @@
     'ModelDataUrl': MODEL_DATA,
 }
 
+ENDPOINT_DESC = {
+    'EndpointConfigName': 'test-endpoint'
+}
+
+ENDPOINT_CONFIG_DESC = {
+    'ProductionVariants': [{'ModelName': 'model-1'},
+                           {'ModelName': 'model-2'}]
+}
+
 
 class DummyFramework(Framework):
     __framework_name__ = 'dummy'
@@ -146,6 +155,8 @@ def sagemaker_session():
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
                                                       return_value=DESCRIBE_TRAINING_JOB_RESULT)
+    sms.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
+    sms.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
     return sms
 
 

tests/unit/test_fm.py

@@ -37,6 +37,15 @@
     }
 }
 
+ENDPOINT_DESC = {
+    'EndpointConfigName': 'test-endpoint'
+}
+
+ENDPOINT_CONFIG_DESC = {
+    'ProductionVariants': [{'ModelName': 'model-1'},
+                           {'ModelName': 'model-2'}]
+}
+
 
 @pytest.fixture()
 def sagemaker_session():
@@ -47,6 +56,8 @@ def sagemaker_session():
     sms.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
     sms.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
                                                       return_value=DESCRIBE_TRAINING_JOB_RESULT)
+    sms.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
+    sms.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
     return sms