Add warning log for failed tries.

PiperOrigin-RevId: 630604504
google · May 4, 2024 · c01069a · c01069a
1 parent 8f69eb2
commit c01069a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/edward2/maps.py b/edward2/maps.py
@@ -107,12 +107,18 @@ def robust_map(
     fn_with_backoff = tenacity.retry(
         retry=retry,
         wait=tenacity.wait_random_exponential(min=1, max=30),
+        before_sleep=tenacity.before_sleep_log(
+            logging.get_absl_logger(), logging.WARNING
+        ),
     )(fn)
   else:
     fn_with_backoff = tenacity.retry(
         retry=retry,
         wait=tenacity.wait_random_exponential(min=1, max=30),
         stop=tenacity.stop_after_attempt(max_retries + 1),
+        before_sleep=tenacity.before_sleep_log(
+            logging.get_absl_logger(), logging.WARNING
+        ),
     )(fn)
   if index_to_output is None:
     index_to_output = {}