Hotfix: Handle EOF Properly (#1557)

* fix: handle EOF properly

* chore: version

* fix: debug logs

* fix: rm eof type

Author: mrkaye97

sdks/python/hatchet_sdk/clients/dispatcher/action_listener.py

@@ -7,7 +7,6 @@
 
 import grpc
 import grpc.aio
-from grpc._cython import cygrpc  # type: ignore[attr-defined]
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
@@ -267,7 +266,6 @@ async def _generator(self) -> AsyncGenerator[Action | None, None]:
                     await self.interrupt.wait()
 
                     if not t.done():
-                        # print a warning
                         logger.warning(
                             "Interrupted read_with_interrupt task of action listener"
                         )
@@ -277,9 +275,10 @@ async def _generator(self) -> AsyncGenerator[Action | None, None]:
 
                         break
 
-                    assigned_action, _ = t.result()
+                    assigned_action, _, is_eof = t.result()
 
-                    if assigned_action is cygrpc.EOF:
+                    if is_eof:
+                        logger.debug("Handling EOF in Action Listener")
                         self.retries = self.retries + 1
                         break
 

sdks/python/hatchet_sdk/clients/event_ts.py

@@ -4,6 +4,8 @@
 import grpc.aio
 from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
+from hatchet_sdk.logger import logger
+
 
 class ThreadSafeEvent(asyncio.Event):
     """
@@ -32,30 +34,31 @@ async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: Callable[[TResponse], str],
-) -> tuple[TResponse, str]: ...
+) -> tuple[TResponse, str, bool]: ...
 
 
 @overload
 async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: None = None,
-) -> tuple[TResponse, None]: ...
+) -> tuple[TResponse, None, bool]: ...
 
 
 async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse],
     interrupt: ThreadSafeEvent,
     key_generator: Callable[[TResponse], str] | None = None,
-) -> tuple[TResponse, str | None]:
+) -> tuple[TResponse, str | None, bool]:
     try:
         result = cast(TResponse, await listener.read())
 
         if result is cygrpc.EOF:
-            raise ValueError("Unexpected EOF")
+            logger.warning("Received EOF from engine")
+            return cast(TResponse, None), None, True
 
         key = key_generator(result) if key_generator else None
 
-        return result, key
+        return result, key, False
     finally:
         interrupt.set()

sdks/python/hatchet_sdk/clients/listeners/pooled_listener.py

@@ -5,7 +5,6 @@
 
 import grpc
 import grpc.aio
-from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
 from hatchet_sdk.clients.event_ts import ThreadSafeEvent, read_with_interrupt
 from hatchet_sdk.config import ClientConfig
@@ -131,9 +130,12 @@ async def _init_producer(self) -> None:
                                 await asyncio.sleep(DEFAULT_LISTENER_RETRY_INTERVAL)
                                 break
 
-                            event, key = t.result()
+                            event, key, is_eof = t.result()
 
-                            if event is cygrpc.EOF:
+                            if is_eof:
+                                logger.debug(
+                                    f"Handling EOF in Pooled Listener {self.__class__.__name__}"
+                                )
                                 break
 
                             subscriptions = self.to_subscriptions.get(key, [])

sdks/python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "hatchet-sdk"
-version = "1.6.0"
+version = "1.6.1"
 description = ""
 authors = ["Alexander Belanger <[email protected]>"]
 readme = "README.md"