DBFS API 429 retries: Fix bug on computing the time until the next retry (#326)

In #319 we added 429 retries for the DBFS API.

Computing the time until the next retry as `retry_state.idle_for - time_for_last_retry` is not going to work as expected in the case where more than 1 call to functions wrapped with `retry_429` happens (e.g. in the case of `put_file`). The problem is that `time_for_last_retry` is only initialised once per databricks-cli command, while it should be initialised every time a function wrapped with `retry_429` is called, in the first attempt.

We fix this by initialising `time_for_last_retry` if `retry_state.attempt_number == 1`. This only works if the code is single-threaded, which is the case here.

Additional improvements:
* Remove unnecessary retry log, since we already have one.
* Increase max number of retries to 8, so that we have one additional retry with long delay.
* Use click.echo instead of logger, to maintain consistency with the rest of the code.
* Add clarifying note on `retry_state.idle_for`.
* Display seconds with 2 decimal places only.

Author: bogdanghita-db

databricks_cli/dbfs/api.py

@@ -28,13 +28,9 @@
 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
@@ -45,13 +41,9 @@
 BUFFER_SIZE_BYTES = 2**20
 EXPONENTIAL_BACKOFF_MULTIPLIER = 1
 MAX_SECONDS_WAIT = 60
-MAX_RETRY_ATTEMPTS = 7
+MAX_RETRY_ATTEMPTS = 8
 time_for_last_retry = 0
 
-logging.basicConfig(stream=sys.stderr, level=logging.WARNING)
-
-logger = logging.getLogger(__name__)
-
 
 class ParseException(Exception):
     pass
@@ -94,13 +86,21 @@ class DbfsErrorCodes(object):
 def before_sleep_on_429(retry_state):
     global time_for_last_retry
     if retry_state.attempt_number < 1:
-        loglevel = logging.INFO
+        click.echo("Warning: Unexpected retry_state.attempt_number={}.".format(
+            retry_state.attempt_number))
+        click.echo("Received 429 REQUEST_LIMIT_EXCEEDED. Retrying with exponential backoff.")
     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
+        # Initialize time_for_last_retry on the first attempt.
+        if retry_state.attempt_number == 1:
+            time_for_last_retry = 0
+        # Note: Here idle_for represents the total time spent sleeping in all retries so far +
+        # the time that we will sleep until the next retry. We determined this empirically,
+        # as it is not clearly stated in the Tenacity docs.
+        time_until_next_retry = retry_state.idle_for - time_for_last_retry
+        click.echo(("Received 429 REQUEST_LIMIT_EXCEEDED for attempt {}. "
+                    "Retrying in {:.2f} seconds.").format(retry_state.attempt_number,
+                                                          time_until_next_retry))
+        time_for_last_retry = retry_state.idle_for
 
 
 def retry_429(func):
@@ -113,7 +113,6 @@ def wrapped_function(*args, **kwargs):
             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

tests/dbfs/test_api.py

@@ -40,6 +40,7 @@
     'file_size': 1
 }
 TEST_FILE_INFO = api.FileInfo(TEST_DBFS_PATH, False, 1)
+MAX_RETRY_ATTEMPTS = 8
 
 
 def get_resource_does_not_exist_exception():
@@ -140,7 +141,7 @@ def test_mkdirs_stop_retrying(self, dbfs_api):
         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
+        assert dbfs_api.client.mkdirs.call_count == MAX_RETRY_ATTEMPTS
 
     def test_file_exists_true(self, dbfs_api):
         dbfs_api.client.get_status.return_value = TEST_FILE_JSON