Retry on 5XXX codes (#1537)

## Description

Retry on certain 5XX codes.

## Changes

- 5XX codes throw a different exception than 429, so adjusts `backoff`
arguments to retry on both
- Specify set of HTTP response codes to retry on in `_RETRY_CODES`
variable


- [ ] I have reviewed the [Guidelines for Contributing](CONTRIBUTING.md)
and the [Code of Conduct](CODE_OF_CONDUCT.md).

python/whylogs/api/writer/whylabs_client.py

@@ -31,7 +31,11 @@
 from whylabs_client.model.transaction_commit_request import TransactionCommitRequest
 from whylabs_client.model.transaction_log_request import TransactionLogRequest
 from whylabs_client.model.transaction_start_request import TransactionStartRequest
-from whylabs_client.rest import ApiException, ForbiddenException  # type: ignore
+from whylabs_client.rest import (  # type: ignore
+    ApiException,
+    ForbiddenException,
+    ServiceException,
+)
 
 from whylogs.api.logger.result_set import ResultSet, SegmentedResultSet
 from whylogs.api.whylabs.session.session_manager import INIT_DOCS, default_init
@@ -90,6 +94,14 @@
 }
 
 
+_RETRY_EXCEPTIONS = (ApiException, ServiceException)
+_RETRY_CODES = {408, 429, 502, 503, 504}
+
+
+def _giveup(e) -> bool:
+    return (e.status not in _RETRY_CODES,)  # type: ignore
+
+
 def _get_column_names(x: Union[DatasetProfile, DatasetProfileView, SegmentedDatasetProfileView, ResultSet]) -> Set[str]:
     if isinstance(x, DatasetProfile):
         return _get_column_names(x.view())
@@ -435,8 +447,8 @@ def tag_custom_performance_column(
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -495,8 +507,8 @@ def _set_column_schema(self, column_name: str, column_schema: ColumnSchema):
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -658,8 +670,8 @@ def get_transaction_id(self) -> str:
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -678,8 +690,8 @@ def commit_transaction(self, id: str) -> None:
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -703,8 +715,8 @@ def abort_transaction(self, id: str) -> None:
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -719,8 +731,8 @@ def transaction_status(self, id: str) -> Dict[str, Any]:
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -740,8 +752,8 @@ def get_upload_url_transaction(
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -800,8 +812,8 @@ def _get_column_weights(self):
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -828,8 +840,8 @@ def _put_feature_weights(self, file: FeatureWeights):
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -858,8 +870,8 @@ def do_request():
     def _get_existing_column_schema(self, model_api_instance, column_name) -> Optional[ColumnSchema]:
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -907,8 +919,8 @@ def _put_column_schema(self, column_name: str, value: str) -> Tuple[int, str]:
 
             @backoff.on_exception(
                 backoff.expo,
-                ApiException,
-                giveup=lambda e: e.status != 429,  # type: ignore
+                _RETRY_EXCEPTIONS,
+                giveup=_giveup,
                 max_time=MAX_REQUEST_TIME,
                 max_tries=MAX_REQUEST_TRIES,
                 jitter=backoff.full_jitter,
@@ -949,8 +961,8 @@ def _post_log_async(
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -1003,8 +1015,8 @@ def _post_log_segmented_reference(
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,
@@ -1034,8 +1046,8 @@ def _post_log_reference(self, request: LogAsyncRequest, dataset_timestamp: int)
 
         @backoff.on_exception(
             backoff.expo,
-            ApiException,
-            giveup=lambda e: e.status != 429,  # type: ignore
+            _RETRY_EXCEPTIONS,
+            giveup=_giveup,
             max_time=MAX_REQUEST_TIME,
             max_tries=MAX_REQUEST_TRIES,
             jitter=backoff.full_jitter,