Support for Predictor to delete SageMaker model.

Author: Chuyang Deng

CHANGELOG.rst

@@ -8,8 +8,8 @@ CHANGELOG
 * doc-fix: update information about saving models in the MXNet README
 * doc-fix: change ReadTheDocs links from latest to stable
 * doc-fix: add ``transform_fn`` information and fix ``input_fn`` signature in the MXNet README
-* feature: Support for ``Predictor`` to delete endpoint configuration by default when calling ``delete_endpoint()``
-* feature: Support for ``model`` to delete SageMaker model
+* feature: Support for ``Predictor`` to delete model, and delete endpoint configuration by default when calling ``delete_endpoint()``
+* feature: Support for ``Model`` to delete SageMaker model
 * feature: Support for ``Transformer`` to delete SageMaker model
 * bug-fix: Default account for SKLearnModel fixed
 

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 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.
 
@@ -283,6 +285,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()``.
@@ -306,6 +311,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/predictor.py

@@ -127,6 +127,22 @@ def delete_endpoint(self, delete_endpoint_config=True):
 
         self.sagemaker_session.delete_endpoint(self.endpoint)
 
+    def delete_model(self):
+        """Deletes the Amazon SageMaker models backing this predictor.
+
+        """
+        model_names = self._get_model_names()
+        for model_name in model_names:
+            self.sagemaker_session.delete_model(model_name)
+
+    def _get_model_names(self):
+        endpoint_desc = self.sagemaker_session.sagemaker_client.describe_endpoint(EndpointName=self.endpoint)
+        endpoint_config_name = endpoint_desc['EndpointConfigName']
+        endpoint_config = self.sagemaker_session.sagemaker_client.describe_endpoint_config(EndpointConfigName=
+                                                                                               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_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_predictor.py

@@ -466,3 +466,18 @@ def test_delete_endpoint_only():
 
     sagemaker_session.delete_endpoint.assert_called_with(ENDPOINT)
     sagemaker_session.delete_endpoint_config.assert_not_called()
+
+
+def test_delete_model():
+    endpoint_desc = {
+        'ProductionVariants': [{
+            'VariantName': 'my-model'
+        }]
+    }
+    sagemaker_session = empty_sagemaker_session()
+    sagemaker_session.sagemaker_client.describe_endpoint = Mock(return_value=endpoint_desc)
+    predictor = RealTimePredictor(ENDPOINT, sagemaker_session=sagemaker_session)
+
+    predictor.delete_model()
+    sagemaker_session.delete_model.assert_called_with('my-model')
+