Deleting endpoint for PyTorch test, retrying endpoint deletion 3 times (#297)

Author: mvsusp

tests/integ/test_pytorch_train.py

@@ -41,7 +41,7 @@ def fixture_training_job(sagemaker_session, pytorch_full_version):
 def test_sync_fit_deploy(pytorch_training_job, sagemaker_session):
     # TODO: add tests against local mode when it's ready to be used
     endpoint_name = 'test-pytorch-sync-fit-attach-deploy{}'.format(sagemaker_timestamp())
-    with timeout(minutes=20):
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, minutes=20):
         estimator = PyTorch.attach(pytorch_training_job, sagemaker_session=sagemaker_session)
         predictor = estimator.deploy(1, 'ml.c4.xlarge', endpoint_name=endpoint_name)
         data = numpy.zeros(shape=(1, 1, 28, 28), dtype=numpy.float32)

tests/integ/timeout.py

@@ -15,6 +15,8 @@
 import signal
 from contextlib import contextmanager
 import logging
+from time import sleep
+
 from awslogs.core import AWSLogs
 from botocore.exceptions import ClientError
 
@@ -65,17 +67,24 @@ def timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, second
             yield [t]
             no_errors = True
         finally:
-            try:
-                sagemaker_session.delete_endpoint(endpoint_name)
-                LOGGER.info('deleted endpoint {}'.format(endpoint_name))
-
-                _show_endpoint_logs(endpoint_name, sagemaker_session)
-                if no_errors:
-                    _cleanup_endpoint_logs(endpoint_name, sagemaker_session)
-            except ClientError as ce:
-                if ce.response['Error']['Code'] == 'ValidationException':
-                    # avoids the inner exception to be overwritten
-                    pass
+            attempts = 3
+
+            while attempts > 0:
+                attempts -= 1
+                try:
+                    sagemaker_session.delete_endpoint(endpoint_name)
+                    LOGGER.info('deleted endpoint {}'.format(endpoint_name))
+
+                    _show_endpoint_logs(endpoint_name, sagemaker_session)
+                    if no_errors:
+                        _cleanup_endpoint_logs(endpoint_name, sagemaker_session)
+                    return
+                except ClientError as ce:
+                    if ce.response['Error']['Code'] == 'ValidationException':
+                        # avoids the inner exception to be overwritten
+                        pass
+                # trying to delete the resource again in 10 seconds
+                sleep(10)
 
 
 def _show_endpoint_logs(endpoint_name, sagemaker_session):