Retry 429s in REST Client (mlflow#1846)

* Added tests * Addressed review comments. * nit
zhaohc10 · Sep 16, 2019 · 7323cfd · 7323cfd
1 parent 18fca9e
commit 7323cfd
Showing 2 changed files with 50 additions and 6 deletions.
diff --git a/mlflow/utils/rest_utils.py b/mlflow/utils/rest_utils.py
@@ -18,12 +18,14 @@
 }
 
 
-def http_request(host_creds, endpoint, retries=3, retry_interval=3, **kwargs):
+def http_request(host_creds, endpoint, retries=3, retry_interval=3,
+                 max_rate_limit_interval=60, **kwargs):
     """
-    Makes an HTTP request with the specified method to the specified hostname/endpoint. Retries
-    up to `retries` times if a request fails with a server error (e.g. error code 500), waiting
-    `retry_interval` seconds between successive retries. Parses the API response (assumed to be
-    JSON) into a Python object and returns it.
+    Makes an HTTP request with the specified method to the specified hostname/endpoint. Ratelimit
+    error code (429) will be retried with an exponential back off (1, 2, 4, ... seconds) for at most
+    `max_rate_limit_interval` seconds.  Internal errors (500s) will be retried up to `retries` times
+    , waiting `retry_interval` seconds between successive retries. Parses the API response
+    (assumed to be JSON) into a Python object and returns it.
 
     :param host_creds: A :py:class:`mlflow.rest_utils.MlflowHostCreds` object containing
         hostname and optional authentication.
@@ -43,10 +45,22 @@ def http_request(host_creds, endpoint, retries=3, retry_interval=3, **kwargs):
 
     verify = not host_creds.ignore_tls_verification
 
+    def request_with_ratelimit_retries(max_rate_limit_interval, **kwargs):
+        response = requests.request(**kwargs)
+        time_left = max_rate_limit_interval
+        sleep = 1
+        while response.status_code == 429 and time_left > 0:
+            time.sleep(sleep)
+            time_left -= sleep
+            response = requests.request(**kwargs)
+            sleep = min(time_left, sleep*2)  # sleep for 1, 2, 4, ... seconds;
+        return response
+
     cleaned_hostname = strip_suffix(hostname, '/')
     url = "%s%s" % (cleaned_hostname, endpoint)
     for i in range(retries):
-        response = requests.request(url=url, headers=headers, verify=verify, **kwargs)
+        response = request_with_ratelimit_retries(max_rate_limit_interval,
+                                                  url=url, headers=headers, verify=verify, **kwargs)
         if response.status_code >= 200 and response.status_code < 500:
             return response
         else:

diff --git a/tests/utils/test_rest_utils.py b/tests/utils/test_rest_utils.py
@@ -85,6 +85,36 @@ def test_http_request_with_insecure(request):
     )
 
 
+@mock.patch('requests.request')
+def test_429_retries(request):
+    host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
+
+    class MockedResponse(object):
+        def __init__(self, status_code):
+            self.status_code = status_code
+            self.text = "mocked text"
+
+    request.side_effect = [MockedResponse(x) for x in (429, 200)]
+    assert http_request(host_only, '/my/endpoint', max_rate_limit_interval=0).status_code == 429
+    request.side_effect = [MockedResponse(x) for x in (429, 200)]
+    assert http_request(host_only, '/my/endpoint', max_rate_limit_interval=1).status_code == 200
+    request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
+    assert http_request(host_only, '/my/endpoint', max_rate_limit_interval=1).status_code == 429
+    request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
+    assert http_request(host_only, '/my/endpoint', max_rate_limit_interval=2).status_code == 200
+    request.side_effect = [MockedResponse(x) for x in (429, 429, 200)]
+    assert http_request(host_only, '/my/endpoint', max_rate_limit_interval=3).status_code == 200
+    # Test that any non 429 code is returned
+    request.side_effect = [MockedResponse(x) for x in (429, 404, 429, 200)]
+    assert http_request(host_only, '/my/endpoint').status_code == 404
+    # Test that retries work as expected
+    request.side_effect = [MockedResponse(x) for x in (429, 503, 429, 200)]
+    with pytest.raises(MlflowException, match="failed to return code 200"):
+        http_request(host_only, '/my/endpoint', retries=1)
+    request.side_effect = [MockedResponse(x) for x in (429, 503, 429, 200)]
+    assert http_request(host_only, '/my/endpoint', retries=2).status_code == 200
+
+
 @mock.patch('requests.request')
 def test_http_request_wrapper(request):
     host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)