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())