Adding retries for DBFS API on Error Code 429 (#319)

This PR attempts to use the well received [tenacity library ](https://github.com/jd/tenacity) to add retries for the DBFS cli commands. This is needed since we are imposing request limiting for these APIs. 

The retries are done with an exponential backoff only for error code 429. The exponential multiplier is 1 sec and the max wait is 30 seconds. We will randomly wait up to 2^x * 1 seconds between each retry until the range reaches 60 seconds, then randomly up to 60 seconds afterwards where x is the number of retry. We will try a maximum for 7 times.

Manually testing by using the following script produced this error. 
```
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 1. Retrying in 0.2678150505033847 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 2. Retrying in 1.4529390049015964 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 3. Retrying in 1.4825864280042993 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 4. Retrying in 3.79336630796632 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 5. Retrying in 19.54595718573239 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 6. Retrying in 32.52860286202574 seconds.
DbfsTestingRateLimitSuite-c5458501-c08b-4add-b782-d3198fee31a2
bogdan-ghita
databricks-results
test-dir-1
tmp
DbfsTestingRateLimitSuite-c5458501-c08b-4add-b782-d3198fee31a2
bogdan-ghita
databricks-results
test-dir-1
tmp
DbfsTestingRateLimitSuite-c5458501-c08b-4add-b782-d3198fee31a2
bogdan-ghita
databricks-results
test-dir-1
tmp
```

In case of an error, we see the following output.
```
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 1. Retrying in 0.25618165976582064 seconds.
Rate limit exceeded. Retrying with exponential backoff.
WARNING:databricks_cli.dbfs.api: Received 429 REQUEST_LIMIT_EXCEEDED for attempt 2. Retrying in 1.3924883604006733 seconds.
Rate limit exceeded. Retrying with exponential backoff.
Error: RateLimitException: 429 Too Many Requests
```

Script: 
```
#!/bin/bash

CONCURRENT_REQ_LIMIT=30
NUM_REQUESTS_TO_EXCEED=1

NUM_REQUESTS=$((CONCURRENT_REQ_LIMIT + NUM_REQUESTS_TO_EXCEED))

for (( i=1; i<=$NUM_REQUESTS; i++ )); do
	echo "$(date) Sending request $i"
	dbfs ls dbfs:/ &
done
wait
```
Unit tests were added to test the changes.

Author: gotibhai

databricks_cli/dbfs/api.py

@@ -28,16 +28,29 @@
 import tempfile
 
 import re
+import logging
+import sys
 import click
 
+
+from tenacity import retry, wait_random_exponential, retry_if_exception_type, stop_after_attempt
+
 from requests.exceptions import HTTPError
 
 from databricks_cli.sdk import DbfsService
 from databricks_cli.utils import error_and_quit
 from databricks_cli.dbfs.dbfs_path import DbfsPath
-from databricks_cli.dbfs.exceptions import LocalFileExistsException
+from databricks_cli.dbfs.exceptions import LocalFileExistsException, RateLimitException
 
 BUFFER_SIZE_BYTES = 2**20
+EXPONENTIAL_BACKOFF_MULTIPLIER = 1
+MAX_SECONDS_WAIT = 60
+MAX_RETRY_ATTEMPTS = 7
+time_for_last_retry = 0
+
+logging.basicConfig(stream=sys.stderr, level=logging.WARNING)
+
+logger = logging.getLogger(__name__)
 
 
 class ParseException(Exception):
@@ -75,12 +88,42 @@ class DbfsErrorCodes(object):
     RESOURCE_DOES_NOT_EXIST = 'RESOURCE_DOES_NOT_EXIST'
     RESOURCE_ALREADY_EXISTS = 'RESOURCE_ALREADY_EXISTS'
     PARTIAL_DELETE = 'PARTIAL_DELETE'
+    TOO_MANY_REQUESTS = 'TOO_MANY_REQUESTS'
+
+
+def before_sleep_on_429(retry_state):
+    global time_for_last_retry
+    if retry_state.attempt_number < 1:
+        loglevel = logging.INFO
+    else:
+        loglevel = logging.WARNING   
+    logger.log(
+        loglevel, ' Received 429 REQUEST_LIMIT_EXCEEDED for attempt %s. Retrying in %s seconds.',
+        retry_state.attempt_number, retry_state.idle_for - time_for_last_retry)
+    time_for_last_retry = retry_state.idle_for
+
+
+def retry_429(func):
+    @retry(wait=wait_random_exponential(multiplier=EXPONENTIAL_BACKOFF_MULTIPLIER, 
+           max=MAX_SECONDS_WAIT), retry=retry_if_exception_type(RateLimitException), 
+           stop=stop_after_attempt(MAX_RETRY_ATTEMPTS), reraise=True, 
+           before_sleep=before_sleep_on_429)
+    def wrapped_function(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except HTTPError as e:
+            if e.response.status_code == 429:
+                click.echo("Rate limit exceeded. Retrying with exponential backoff.")
+                raise RateLimitException("429 Too Many Requests")
+            raise e
+    return wrapped_function
 
 
 class DbfsApi(object):
     def __init__(self, api_client):
         self.client = DbfsService(api_client)
 
+    @retry_429
     def list_files(self, dbfs_path, headers=None):
         list_response = self.client.list(dbfs_path.absolute_path, headers=headers)
         if 'files' in list_response:
@@ -101,20 +144,37 @@ def file_exists(self, dbfs_path, headers=None):
             raise e
         return True
 
+    @retry_429
     def get_status(self, dbfs_path, headers=None):
         json = self.client.get_status(dbfs_path.absolute_path, headers=headers)
         return FileInfo.from_json(json)
 
+    @retry_429
+    def create(self, dbfs_path, overwrite, headers):
+        return self.client.create(dbfs_path.absolute_path, overwrite, headers=headers)
+
+    @retry_429
+    def add_block(self, handle, contents, headers):
+        self.client.add_block(handle, contents, headers=headers)
+
+    @retry_429
+    def close(self, handle, headers):
+        self.client.close(handle, headers=headers)
+
     def put_file(self, src_path, dbfs_path, overwrite, headers=None):
-        handle = self.client.create(dbfs_path.absolute_path, overwrite, headers=headers)['handle']
+        handle = self.create(dbfs_path, overwrite, headers=headers)['handle']
         with open(src_path, 'rb') as local_file:
             while True:
                 contents = local_file.read(BUFFER_SIZE_BYTES)
                 if len(contents) == 0:
                     break
                 # add_block should not take a bytes object.
-                self.client.add_block(handle, b64encode(contents).decode(), headers=headers)
-            self.client.close(handle, headers=headers)
+                self.add_block(handle, b64encode(contents).decode(), headers=headers)
+            self.close(handle, headers=headers)
+
+    @retry_429
+    def read(self, dbfs_path, offset, headers):
+        return self.client.read(dbfs_path.absolute_path, offset, BUFFER_SIZE_BYTES, headers=headers)       
 
     def get_file(self, dbfs_path, dst_path, overwrite, headers=None):
         if os.path.exists(dst_path) and not overwrite:
@@ -126,8 +186,7 @@ def get_file(self, dbfs_path, dst_path, overwrite, headers=None):
         offset = 0
         with open(dst_path, 'wb') as local_file:
             while offset < length:
-                response = self.client.read(dbfs_path.absolute_path, offset, BUFFER_SIZE_BYTES,
-                                            headers=headers)
+                response = self.read(dbfs_path, offset, headers=headers)
                 bytes_read = response['bytes_read']
                 data = response['data']
                 offset += bytes_read
@@ -146,11 +205,13 @@ def get_num_files_deleted(partial_delete_error):
                     message))
         return int(m.group(1))
 
+    @retry_429
     def delete(self, dbfs_path, recursive, headers=None):
         num_files_deleted = 0
         while True:
             try:
-                self.client.delete(dbfs_path.absolute_path, recursive=recursive, headers=headers)
+                self.client.delete(dbfs_path.absolute_path, 
+                                   recursive=recursive, headers=headers)
             except HTTPError as e:
                 if e.response.status_code == 503:
                     try:
@@ -172,9 +233,11 @@ def delete(self, dbfs_path, recursive, headers=None):
             break
         click.echo("\rDelete finished successfully.\033[K")
 
+    @retry_429
     def mkdirs(self, dbfs_path, headers=None):
         self.client.mkdirs(dbfs_path.absolute_path, headers=headers)
 
+    @retry_429
     def move(self, dbfs_src, dbfs_dst, headers=None):
         self.client.move(dbfs_src.absolute_path, dbfs_dst.absolute_path, headers=headers)
 

databricks_cli/dbfs/exceptions.py

@@ -24,3 +24,7 @@
 
 class LocalFileExistsException(Exception):
     pass
+
+
+class RateLimitException(Exception):
+    pass

setup.py

@@ -39,6 +39,7 @@
         'tabulate>=0.7.7',
         'six>=1.10.0',
         'configparser>=0.3.5;python_version < "3.6"',
+        'tenacity>=6.2.0'
     ],
     entry_points='''
         [console_scripts]

tests/dbfs/test_api.py

@@ -31,7 +31,7 @@
 
 import databricks_cli.dbfs.api as api
 from databricks_cli.dbfs.dbfs_path import DbfsPath
-from databricks_cli.dbfs.exceptions import LocalFileExistsException
+from databricks_cli.dbfs.exceptions import LocalFileExistsException, RateLimitException
 
 TEST_DBFS_PATH = DbfsPath('dbfs:/test')
 TEST_FILE_JSON = {
@@ -55,6 +55,13 @@ def get_partial_delete_exception(message="[...] operation has deleted 10 files [
     return requests.exceptions.HTTPError(response=response)
 
 
+def get_rate_limit_exception():
+    response = requests.Response()
+    response.status_code = 429
+    response._content = ('{{"error_code": "{}"}}'.format(api.DbfsErrorCodes.TOO_MANY_REQUESTS)).encode() #  NOQA
+    return requests.exceptions.HTTPError(response=response)
+
+
 class TestFileInfo(object):
     def test_to_row_not_long_form_not_absolute(self):
         file_info = api.FileInfo(TEST_DBFS_PATH, False, 1)
@@ -103,6 +110,38 @@ def test_list_files_does_not_exist(self, dbfs_api):
 
         assert len(files) == 0
 
+    def test_list_files_rate_limited(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        # Simulate 2 rate limit exceptions followed by a full successful operation
+        exception_sequence = [rate_limit_exception, rate_limit_exception, None]
+        dbfs_api.client.list_files = mock.Mock(side_effect=exception_sequence)
+        # Should succeed
+        files = dbfs_api.list_files(TEST_DBFS_PATH)
+
+        assert len(files) == 0
+
+    def test_mkdirs_rate_limited(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        # Simulate 2 rate limit exceptions followed by a full successful operation
+        exception_sequence = [rate_limit_exception, rate_limit_exception, None]
+        dbfs_api.client.mkdirs = mock.Mock(side_effect=exception_sequence)
+        # Should succeed
+        dbfs_api.mkdirs(DbfsPath('dbfs:/test/mkdir'))
+        files = dbfs_api.client.list(DbfsPath('dbfs:/test/mkdir'))
+
+        assert len(files) == 0
+
+    def test_mkdirs_stop_retrying(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        # Simulate 9 rate limit exceptions which will fail eventually
+        exception_sequence = [rate_limit_exception, rate_limit_exception, rate_limit_exception, 
+                              rate_limit_exception, rate_limit_exception, rate_limit_exception, 
+                              rate_limit_exception, rate_limit_exception, rate_limit_exception]
+        dbfs_api.client.mkdirs = mock.Mock(side_effect=exception_sequence)
+        with pytest.raises(RateLimitException):
+            dbfs_api.mkdirs(DbfsPath('dbfs:/test/mkdir'))
+        assert dbfs_api.client.mkdirs.call_count == 7
+
     def test_file_exists_true(self, dbfs_api):
         dbfs_api.client.get_status.return_value = TEST_FILE_JSON
         assert dbfs_api.file_exists(TEST_DBFS_PATH)
@@ -122,6 +161,16 @@ def test_get_status_fail(self, dbfs_api):
         with pytest.raises(exception.__class__):
             dbfs_api.get_status(TEST_DBFS_PATH)
 
+    def test_get_status_rate_limited(self, dbfs_api):
+        exception = get_rate_limit_exception()
+        dbfs_api.client.get_status = mock.Mock(side_effect=[exception, {
+            'path': '/test',
+            'is_dir': False,
+            'file_size': 1
+        }])
+        # Should succeed
+        assert dbfs_api.get_status(TEST_DBFS_PATH) is not None
+
     def test_put_file(self, dbfs_api, tmpdir):
         test_file_path = os.path.join(tmpdir.strpath, 'test')
         with open(test_file_path, 'wt') as f:
@@ -138,6 +187,26 @@ def test_put_file(self, dbfs_api, tmpdir):
         assert api_mock.close.call_count == 1
         assert test_handle == api_mock.close.call_args[0][0]
 
+    def test_put_file_rate_limited(self, dbfs_api, tmpdir):
+        test_file_path = os.path.join(tmpdir.strpath, 'test')
+        with open(test_file_path, 'wt') as f:
+            f.write('test')
+
+        exception = get_rate_limit_exception()
+        api_mock = dbfs_api.client
+        test_handle = 0
+        api_mock.create = mock.Mock(side_effect=[exception, {'handle': test_handle}])
+        api_mock.add_block = mock.Mock(side_effect=[exception, None])
+        api_mock.close = mock.Mock(side_effect=[exception, None])
+        dbfs_api.put_file(test_file_path, TEST_DBFS_PATH, True)
+
+        assert api_mock.create.call_count == 2
+        assert api_mock.add_block.call_count == 2
+        assert test_handle == api_mock.add_block.call_args[0][0]
+        assert b64encode(b'test').decode() == api_mock.add_block.call_args[0][1]
+        assert api_mock.close.call_count == 2
+        assert test_handle == api_mock.close.call_args[0][0]
+
     def test_get_file_check_overwrite(self, dbfs_api, tmpdir):
         test_file_path = os.path.join(tmpdir.strpath, 'test')
         with open(test_file_path, 'w') as f:
@@ -159,6 +228,22 @@ def test_get_file(self, dbfs_api, tmpdir):
         with open(test_file_path, 'r') as f:
             assert f.read() == 'x'
 
+    def test_get_file_rate_limited(self, dbfs_api, tmpdir):
+        api_mock = dbfs_api.client
+        api_mock.get_status.return_value = TEST_FILE_JSON
+        rate_limit_exception = get_rate_limit_exception()
+
+        api_mock.read = mock.Mock(side_effect=[rate_limit_exception, {
+            'bytes_read': 1,
+            'data': b64encode(b'x'),
+        }])
+
+        test_file_path = os.path.join(tmpdir.strpath, 'test')
+        dbfs_api.get_file(TEST_DBFS_PATH, test_file_path, True)
+
+        with open(test_file_path, 'r') as f:
+            assert f.read() == 'x'
+
     def test_cat(self, dbfs_api):
         dbfs_api.client.get_status.return_value = {
             'path': '/test',
@@ -173,15 +258,58 @@ def test_cat(self, dbfs_api):
             dbfs_api.cat('dbfs:/whatever-doesnt-matter')
             click_mock.echo.assert_called_with('a', nl=False)
 
+    def test_cat_rate_limited(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        # Simulate 2 rate limit exceptions followed by a full successful operation
+        get_status_exception_sequence = [rate_limit_exception, rate_limit_exception, {
+            'path': '/test',
+            'is_dir': False,
+            'file_size': 1
+        }]
+        read_exception_sequence = [rate_limit_exception, rate_limit_exception, {
+            'bytes_read': 1,
+            'data': b64encode(b'a'),
+        }]
+
+        dbfs_api.client.get_status = mock.Mock(side_effect=get_status_exception_sequence)
+        dbfs_api.client.read = mock.Mock(side_effect=read_exception_sequence)
+
+        with mock.patch('databricks_cli.dbfs.api.click') as click_mock:
+            dbfs_api.cat('dbfs:/whatever-doesnt-matter')
+            click_mock.echo.assert_called_with('a', nl=False)
+
     def test_partial_delete(self, dbfs_api):
         e_partial_delete = get_partial_delete_exception()
-        # Simulate 3 partial deletes followed by a full successful delete
+        # Simulate 3 partial deletes followed by a full successful operation
         exception_sequence = [e_partial_delete, e_partial_delete, e_partial_delete, None]
         dbfs_api.client.delete = mock.Mock(side_effect=exception_sequence)
         dbfs_api.delete_retry_delay_millis = 1
         # Should succeed
         dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
 
+    def test_partial_delete_with_rate_limit(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        e_partial_delete = get_partial_delete_exception()
+        # Simulate 3 partial deletes interspersed with rate limiting exceptions 
+        # followed by a full successful delete
+        exception_sequence = [e_partial_delete, rate_limit_exception, 
+                              e_partial_delete, rate_limit_exception,
+                              e_partial_delete, None]
+
+        dbfs_api.client.delete = mock.Mock(side_effect=exception_sequence)
+        dbfs_api.delete_retry_delay_millis = 1
+        # Should succeed
+        dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
+
+    def test_delete_with_rate_limit(self, dbfs_api):
+        rate_limit_exception = get_rate_limit_exception()
+        # Simulate a rate limit exception followed by a full successful delete
+        exception_sequence = [rate_limit_exception, None]
+        dbfs_api.client.delete = mock.Mock(side_effect=exception_sequence)
+        dbfs_api.delete_retry_delay_millis = 1
+        # Should succeed
+        dbfs_api.delete(DbfsPath('dbfs:/whatever-doesnt-matter'), recursive=True)
+
     def test_partial_delete_exception_message_parse_error(self, dbfs_api):
         message = "unexpected partial delete exception message"
         e_partial_delete = get_partial_delete_exception(message)