fix python aio put stream transport

Author: juanluisdb

sdks/python/hatchet_sdk/clients/events.py

@@ -5,6 +5,7 @@
 from datetime import timezone
 from typing import cast
 
+import grpc.aio
 from google.protobuf import timestamp_pb2
 from pydantic import BaseModel, ConfigDict
 
@@ -26,6 +27,7 @@
     PushEventRequest,
     PutLogRequest,
     PutStreamEventRequest,
+    PutStreamEventResponse,
 )
 from hatchet_sdk.contracts.events_pb2 import Event as EventProto
 from hatchet_sdk.contracts.events_pb2 import Events as EventsProto
@@ -117,6 +119,8 @@ def __init__(self, config: ClientConfig):
 
         conn = new_conn(config, False)
         self.events_service_client = EventsServiceStub(conn)
+        self.aio_events_service_client: EventsServiceStub | None = None
+        self.aio_events_service_channel: grpc.aio.Channel | None = None
 
         self.token = config.token
         self.namespace = config.namespace
@@ -127,6 +131,15 @@ def _wra(self, client: ApiClient) -> WorkflowRunsApi:
     def _ea(self, client: ApiClient) -> EventApi:
         return EventApi(client)
 
+    def _get_or_create_aio_events_service_client(self) -> EventsServiceStub:
+        if self.aio_events_service_client is None:
+            self.aio_events_service_channel = new_conn(self.client_config, True)
+            self.aio_events_service_client = EventsServiceStub(
+                self.aio_events_service_channel
+            )
+
+        return self.aio_events_service_client
+
     async def aio_push(
         self,
         event_key: str,
@@ -289,29 +302,57 @@ def log(
 
         put_log(request, metadata=create_authorization_header(self.token))
 
-    def stream(self, data: str | bytes, step_run_id: str, index: int) -> None:
-        put_stream_event = tenacity_retry(
-            self.events_service_client.PutStreamEvent, self.client_config.tenacity
-        )
+    def _create_put_stream_event_request(
+        self, data: str | bytes, step_run_id: str, index: int
+    ) -> PutStreamEventRequest:
         if isinstance(data, str):
             data_bytes = data.encode("utf-8")
         elif isinstance(data, bytes):
             data_bytes = data
         else:
             raise ValueError("Invalid data type. Expected str, bytes, or file.")
 
-        request = PutStreamEventRequest(
+        return PutStreamEventRequest(
             task_run_external_id=step_run_id,
             created_at=proto_timestamp_now(),
             message=data_bytes,
             event_index=index,
         )
 
+    def stream(self, data: str | bytes, step_run_id: str, index: int) -> None:
+        put_stream_event = tenacity_retry(
+            self.events_service_client.PutStreamEvent, self.client_config.tenacity
+        )
+        request = self._create_put_stream_event_request(data, step_run_id, index)
+
         try:
             put_stream_event(request, metadata=create_authorization_header(self.token))
         except Exception:
             raise
 
+    async def aio_stream(self, data: str | bytes, step_run_id: str, index: int) -> None:
+        aio_events_service_client = self._get_or_create_aio_events_service_client()
+        request = self._create_put_stream_event_request(data, step_run_id, index)
+
+        async def put_stream_event(
+            request: PutStreamEventRequest,
+            metadata: tuple[tuple[str, str]],
+        ) -> PutStreamEventResponse:
+            return cast(
+                PutStreamEventResponse,
+                await aio_events_service_client.PutStreamEvent(  # type: ignore[misc]
+                    request, metadata=metadata
+                ),
+            )
+
+        retrying_put_stream_event = tenacity_retry(
+            put_stream_event, self.client_config.tenacity
+        )
+
+        await retrying_put_stream_event(
+            request, metadata=create_authorization_header(self.token)
+        )
+
     async def aio_list(
         self,
         offset: int | None = None,

sdks/python/hatchet_sdk/context/context.py

@@ -544,7 +544,16 @@ async def aio_put_stream(self, data: str | bytes) -> None:
         :param data: The data to send to the Hatchet API. Can be a string or bytes.
         :return: None
         """
-        await asyncio.to_thread(self.put_stream, data)
+        try:
+            ix = self._increment_stream_index()
+
+            await self._event_client.aio_stream(
+                data=data,
+                step_run_id=self._step_run_id,
+                index=ix,
+            )
+        except Exception:
+            logger.exception("error putting stream event")
 
     def refresh_timeout(self, increment_by: Duration) -> None:
         """

sdks/python/tests/unit/test_aio_put_stream.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable
+from typing import cast
+
+import grpc
+import grpc.aio
+import pytest
+import tenacity
+
+from hatchet_sdk.clients.events import EventClient
+from hatchet_sdk.config import ClientConfig, TenacityConfig
+from hatchet_sdk.context.context import Context
+from hatchet_sdk.contracts.events_pb2 import (
+    PutStreamEventRequest,
+    PutStreamEventResponse,
+)
+from hatchet_sdk.contracts.events_pb2_grpc import EventsServiceStub
+
+
+def _make_grpc_error(code: grpc.StatusCode, details: str = "") -> grpc.aio.AioRpcError:
+    empty: grpc.aio.Metadata = grpc.aio.Metadata()
+    return grpc.aio.AioRpcError(code, empty, empty, details)
+
+
+class _GeneratedAioUnaryCall:
+    """Matches grpc.aio generated unary calls: sync call, awaitable result."""
+
+    def __init__(self, failures_before_success: int) -> None:
+        self.failures_before_success = failures_before_success
+        self.calls = 0
+        self.requests: list[PutStreamEventRequest] = []
+        self.metadata: list[tuple[tuple[str, str]]] = []
+
+    def __call__(
+        self,
+        request: PutStreamEventRequest,
+        *,
+        metadata: tuple[tuple[str, str]],
+    ) -> Awaitable[PutStreamEventResponse]:
+        self.calls += 1
+        self.requests.append(request)
+        self.metadata.append(metadata)
+
+        async def response() -> PutStreamEventResponse:
+            if self.calls <= self.failures_before_success:
+                raise _make_grpc_error(grpc.StatusCode.UNAVAILABLE, "transient")
+
+            return PutStreamEventResponse()
+
+        return response()
+
+
+class _FakeAioEventsServiceStub:
+    def __init__(self, put_stream_event: _GeneratedAioUnaryCall) -> None:
+        self.PutStreamEvent = put_stream_event
+
+
+def _event_client(aio_stub: _FakeAioEventsServiceStub) -> EventClient:
+    client = EventClient.__new__(EventClient)
+    client.client_config = ClientConfig.model_construct(
+        tenant_id="tenant",
+        token="token",
+        namespace="",
+        server_url="http://localhost",
+        host_port="localhost:7070",
+        tenacity=TenacityConfig(max_attempts=3, wait=tenacity.wait_none),
+    )
+    client.token = "token"
+    client.namespace = ""
+    client.aio_events_service_client = cast(EventsServiceStub, aio_stub)
+    client.aio_events_service_channel = None
+
+    return client
+
+
+@pytest.mark.parametrize(
+    ("data", "expected_message"),
+    [
+        ("hello", b"hello"),
+        (b"hello", b"hello"),
+    ],
+)
+async def test_aio_stream_retries_generated_aio_callable(
+    data: str | bytes, expected_message: bytes
+) -> None:
+    put_stream_event = _GeneratedAioUnaryCall(failures_before_success=2)
+    client = _event_client(_FakeAioEventsServiceStub(put_stream_event))
+
+    await client.aio_stream(data, step_run_id="step-run-id", index=7)
+
+    assert put_stream_event.calls == 3
+    assert [request.task_run_external_id for request in put_stream_event.requests] == [
+        "step-run-id",
+        "step-run-id",
+        "step-run-id",
+    ]
+    assert [request.message for request in put_stream_event.requests] == [
+        expected_message,
+        expected_message,
+        expected_message,
+    ]
+    assert [request.event_index for request in put_stream_event.requests] == [7, 7, 7]
+    assert put_stream_event.metadata == [
+        (("authorization", "bearer token"),),
+        (("authorization", "bearer token"),),
+        (("authorization", "bearer token"),),
+    ]
+
+
+class _RecordingEventClient:
+    def __init__(self) -> None:
+        self.calls: list[tuple[str | bytes, str, int]] = []
+        self.both_calls_started = asyncio.Event()
+        self.release_sends = asyncio.Event()
+
+    async def aio_stream(self, data: str | bytes, step_run_id: str, index: int) -> None:
+        self.calls.append((data, step_run_id, index))
+        if len(self.calls) == 2:
+            self.both_calls_started.set()
+        await self.release_sends.wait()
+
+
+def _context(event_client: _RecordingEventClient) -> Context:
+    ctx = Context.__new__(Context)
+    ctx._stream_index = 0
+    ctx._step_run_id = "step-run-id"
+    ctx._event_client = cast(EventClient, event_client)
+
+    return ctx
+
+
+async def test_aio_put_stream_assigns_index_before_async_send() -> None:
+    event_client = _RecordingEventClient()
+    ctx = _context(event_client)
+
+    tasks = [
+        asyncio.create_task(ctx.aio_put_stream("first")),
+        asyncio.create_task(ctx.aio_put_stream(b"second")),
+    ]
+
+    try:
+        await asyncio.wait_for(event_client.both_calls_started.wait(), timeout=1)
+
+        assert event_client.calls == [
+            ("first", "step-run-id", 0),
+            (b"second", "step-run-id", 1),
+        ]
+    finally:
+        event_client.release_sends.set()
+        await asyncio.gather(*tasks)