Improve Tenacity Verbosity and Limit Retries to Specific Errors (#167)

* Tenacity Patches

* Ensure We Reraise Exceptions

* Bump Version

* Create Common Tenacity Function That Reuses Similar Params

---------

Co-authored-by: srhinos <[email protected]>

Author: macwilk

hatchet_sdk/clients/admin.py

@@ -4,8 +4,8 @@
 
 import grpc
 from google.protobuf import timestamp_pb2
-from tenacity import retry, stop_after_attempt, wait_exponential_jitter
 
+from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.clients.run_event_listener import new_listener
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
 from hatchet_sdk.connection import new_conn
@@ -169,10 +169,7 @@ async def run(
             wrr.workflow_run_id, wrr.workflow_listener, wrr.workflow_run_event_listener
         )
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     async def run_workflow(
         self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
     ) -> WorkflowRunRef:
@@ -200,10 +197,7 @@ async def run_workflow(
 
             raise ValueError(f"gRPC error: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     async def put_workflow(
         self,
         name: str,
@@ -220,10 +214,7 @@ async def put_workflow(
         except grpc.RpcError as e:
             raise ValueError(f"Could not put workflow: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     async def put_rate_limit(
         self,
         key: str,
@@ -242,10 +233,7 @@ async def put_rate_limit(
         except grpc.RpcError as e:
             raise ValueError(f"Could not put rate limit: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     async def schedule_workflow(
         self,
         name: str,
@@ -279,10 +267,7 @@ def __init__(self, config: ClientConfig):
         self.listener_client = new_listener(config)
         self.namespace = config.namespace
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     def put_workflow(
         self,
         name: str,
@@ -301,10 +286,7 @@ def put_workflow(
         except grpc.RpcError as e:
             raise ValueError(f"Could not put workflow: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     def put_rate_limit(
         self,
         key: str,
@@ -323,10 +305,7 @@ def put_rate_limit(
         except grpc.RpcError as e:
             raise ValueError(f"Could not put rate limit: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     def schedule_workflow(
         self,
         name: str,
@@ -349,10 +328,7 @@ def schedule_workflow(
 
             raise ValueError(f"gRPC error: {e}")
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     def run_workflow(
         self, workflow_name: str, input: any, options: TriggerWorkflowOptions = None
     ) -> WorkflowRunRef:

hatchet_sdk/clients/events.py

@@ -4,8 +4,8 @@
 
 import grpc
 from google.protobuf import timestamp_pb2
-from tenacity import retry, stop_after_attempt, wait_exponential_jitter
 
+from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.contracts.events_pb2 import (
     Event,
     PushEventRequest,
@@ -43,10 +43,7 @@ def __init__(self, client: EventsServiceStub, config: ClientConfig):
         self.token = config.token
         self.namespace = config.namespace
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     def push(self, event_key, payload, options: PushEventOptions = None) -> Event:
 
         namespaced_event_key = self.namespace + event_key

hatchet_sdk/clients/rest/tenacity_utils.py

@@ -0,0 +1,34 @@
+import grpc
+import tenacity
+
+from hatchet_sdk.logger import logger
+
+
+def tenacity_retry(func):
+    return tenacity.retry(
+        reraise=True,
+        wait=tenacity.wait_exponential_jitter(),
+        stop=tenacity.stop_after_attempt(5),
+        before_sleep=tenacity_alert_retry,
+        retry=tenacity.retry_if_exception(tenacity_should_retry),
+    )(func)
+
+
+def tenacity_alert_retry(retry_state: tenacity.RetryCallState) -> None:
+    """Called between tenacity retries."""
+    logger.debug(
+        f"Retrying {retry_state.fn}: attempt "
+        f"{retry_state.attempt_number} ended with: {retry_state.outcome}",
+    )
+
+
+def tenacity_should_retry(ex: Exception) -> bool:
+    if isinstance(ex, grpc.aio.AioRpcError):
+        if ex.code in [
+            grpc.StatusCode.UNIMPLEMENTED,
+            grpc.StatusCode.NOT_FOUND,
+        ]:
+            return False
+        return True
+    else:
+        return False

hatchet_sdk/context/context.py

@@ -3,9 +3,8 @@
 import traceback
 from concurrent.futures import Future, ThreadPoolExecutor
 
-from tenacity import retry, stop_after_attempt, wait_exponential_jitter
-
 from hatchet_sdk.clients.events import EventClient
+from hatchet_sdk.clients.rest.tenacity_utils import tenacity_retry
 from hatchet_sdk.clients.rest_client import RestApi
 from hatchet_sdk.clients.run_event_listener import RunEventListenerClient
 from hatchet_sdk.clients.workflow_listener import PooledWorkflowRunListener
@@ -85,10 +84,7 @@ def __init__(
         self.spawn_index = -1
         self.worker = worker
 
-    @retry(
-        wait=wait_exponential_jitter(),
-        stop=stop_after_attempt(5),
-    )
+    @tenacity_retry
     async def spawn_workflow(
         self,
         workflow_name: str,

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "0.36.12"
+version = "0.36.13"
 description = ""
 authors = ["Alexander Belanger <[email protected]>"]
 readme = "README.md"
@@ -20,7 +20,7 @@ aiostream = "^0.5.2"
 nest-asyncio = "^1.6.0"
 aiohttp = "^3.10.5"
 aiohttp-retry = "^2.8.3"
-tenacity = ">=8.0.0"
+tenacity = ">=8.4.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.2"