From e62a70f7b580239de367ffc81d6156c4e8e928e2 Mon Sep 17 00:00:00 2001 From: Amir Fathi Date: Sat, 22 Aug 2026 00:25:37 +0000 Subject: [PATCH] Flush buffered events before the async tracking client's worker exits BasicAsynchronousHamiltonClient had no stop(), and nothing ever put the None sentinel worker() waits for, so a buffered batch sat until the periodic flush timer happened to catch it, or was lost outright if the process exited first. Add stop(), track the worker task, and delegate to it from AsyncHamiltonTracker.stop(). The sentinel branch in worker() had never actually run before this change: it appended the item to the batch before checking whether it was the None sentinel, so the first real stop() call flushed a batch containing None and crashed. Fixed the check order to match the synchronous client, which already gets this right. Fixes #1214 --- ui/sdk/src/hamilton_sdk/adapters.py | 8 +++++ ui/sdk/src/hamilton_sdk/api/clients.py | 19 +++++++++-- ui/sdk/tests/test_adapters.py | 27 ++++++++++++++++ ui/sdk/tests/test_clients.py | 45 ++++++++++++++++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/ui/sdk/src/hamilton_sdk/adapters.py b/ui/sdk/src/hamilton_sdk/adapters.py index f63248ee7..dae3044b0 100644 --- a/ui/sdk/src/hamilton_sdk/adapters.py +++ b/ui/sdk/src/hamilton_sdk/adapters.py @@ -469,6 +469,14 @@ async def ainit(self): self.initialized = True return self + async def stop(self): + """Flush any buffered tracking data and stop the background worker. + + Call this before your event loop shuts down; nothing else flushes + the tracker's queue for you. + """ + await self.client.stop() + async def post_graph_construct( self, graph: h_graph.FunctionGraph, modules: list[ModuleType], config: dict[str, Any] ): diff --git a/ui/sdk/src/hamilton_sdk/api/clients.py b/ui/sdk/src/hamilton_sdk/api/clients.py index 90d84d801..3a0664f2e 100644 --- a/ui/sdk/src/hamilton_sdk/api/clients.py +++ b/ui/sdk/src/hamilton_sdk/api/clients.py @@ -585,9 +585,24 @@ def __init__( self.data_queue = asyncio.Queue() self.running = True self.max_batch_size = 100 + self.worker_task: asyncio.Task | None = None async def ainit(self): - asyncio.create_task(self.worker()) + self.worker_task = asyncio.create_task(self.worker()) + + async def stop(self): + """Signal the worker to flush whatever is buffered and exit. + + Unlike BasicSynchronousHamiltonClient, this is not wired to atexit: + the worker is a task on this event loop, and by the time atexit + callbacks run at interpreter shutdown the loop may already be + closed, so callers must await this explicitly before the loop + that owns it stops running. + """ + if self.worker_task is None or self.worker_task.done(): + return + await self.data_queue.put(None) + await self.worker_task async def flush(self, batch): """Flush the batch (send it to the backend or process it).""" @@ -624,7 +639,6 @@ async def worker(self): ) try: item = await asyncio.wait_for(self.data_queue.get(), timeout=self.flush_interval) - batch.append(item) except asyncio.TimeoutError: # This is fine, we just keep waiting pass @@ -632,6 +646,7 @@ async def worker(self): if item is None: await self.flush(batch) return + batch.append(item) # Check if batch is full or flush interval has passed if ( diff --git a/ui/sdk/tests/test_adapters.py b/ui/sdk/tests/test_adapters.py index 754cec68d..d6d8a0230 100644 --- a/ui/sdk/tests/test_adapters.py +++ b/ui/sdk/tests/test_adapters.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import asyncio import os.path import pytest @@ -162,6 +163,32 @@ def test_parallel_ray_sample_error(): shutdown() +def test_async_hamilton_tracker_stop_delegates_to_client(): + """Regression for #1214: AsyncHamiltonTracker had no way for a caller to + flush its client's buffered tracking data before the event loop shuts + down. stop() must delegate to the underlying client's stop().""" + + class _FakeAsyncClient: + def __init__(self, *args, **kwargs): + self.stop_calls = 0 + + async def stop(self): + self.stop_calls += 1 + + async def run(): + tracker = adapters.AsyncHamiltonTracker( + project_id=1, + username="test-user", + dag_name="test-dag", + client_factory=_FakeAsyncClient, + api_key="test-key", + ) + await tracker.stop() + assert tracker.client.stop_calls == 1 + + asyncio.run(run()) + + if __name__ == "__main__": # test_adapters() diff --git a/ui/sdk/tests/test_clients.py b/ui/sdk/tests/test_clients.py index 105e98dd8..62ee04fa7 100644 --- a/ui/sdk/tests/test_clients.py +++ b/ui/sdk/tests/test_clients.py @@ -81,3 +81,48 @@ async def run(): await client.flush(batch) asyncio.run(run()) + + +def test_async_client_stop_flushes_buffered_item_and_exits_worker(): + """Regression for #1214: worker()'s only exit path waits for a `None` + sentinel on data_queue, and nothing ever put one there for the async + client, so a buffered item sat unflushed until the flush_interval timer + happened to fire, or was lost outright if the process exited first. + stop() must send that sentinel and wait for the worker to flush and + exit, well before the flush_interval timeout.""" + + async def run(): + client = _make_client() + session_cm = _mock_session_with_status(200) + with patch("aiohttp.ClientSession", return_value=session_cm): + await client.ainit() + await client.update_tasks(dag_run_id=1, attributes=[], task_updates=[]) + + assert not client.worker_task.done() + + await client.stop() + + assert client.worker_task.done() + assert client.data_queue.empty() + session_cm.__aenter__.return_value.put.assert_called_once() + + asyncio.run(run()) + + +def test_async_client_stop_is_idempotent(): + async def run(): + client = _make_client() + with patch("aiohttp.ClientSession", return_value=_mock_session_with_status(200)): + await client.ainit() + await client.stop() + await client.stop() # must not hang or raise + + asyncio.run(run()) + + +def test_async_client_stop_before_ainit_is_a_noop(): + async def run(): + client = _make_client() + await client.stop() # worker_task is None, must not raise + + asyncio.run(run())