Skip to content

Flush buffered events before the async tracking client's worker exits - #1692

Open
AmirF194 wants to merge 1 commit into
apache:mainfrom
AmirF194:fix/1214-async-client-stop-atexit
Open

Flush buffered events before the async tracking client's worker exits#1692
AmirF194 wants to merge 1 commit into
apache:mainfrom
AmirF194:fix/1214-async-client-stop-atexit

Conversation

@AmirF194

Copy link
Copy Markdown

Fixes #1214.

Changes

BasicAsynchronousHamiltonClient (ui/sdk/src/hamilton_sdk/api/clients.py) had no stop()
and registered no atexit hook, unlike its synchronous sibling. 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 batch sat until the periodic flush timer happened to catch it, or was
lost outright if the process (or the event loop that owned the worker task) exited first.

  • Added async def stop() on BasicAsynchronousHamiltonClient: sends the sentinel and awaits
    the worker task so it flushes and exits. This is not wired to atexit, deliberately: the
    worker is a task on the caller's event loop, and by the time atexit callbacks run at
    interpreter shutdown that loop may already be closed, so callers must await stop()
    themselves before the loop stops running (mirroring the sync client isn't possible here for
    that reason; happy to add an atexit-scheduled best-effort variant too if maintainers want
    one, but that has its own shutdown-ordering edge cases worth discussing first).
  • Added the matching async def stop() on AsyncHamiltonTracker (adapters.py), delegating
    to the client.
  • Fixed a bug in worker()'s sentinel handling that this exercises for the first time: it
    appended the received item to batch before checking whether it was the None sentinel, so
    the first real stop() call flushed a batch containing None and raised. Reordered to check
    first, matching the synchronous client's worker.

How I tested this

Ran in a clean python:3.13-slim and python:3.10-slim Docker container (the SDK's CI matrix
runs 3.10-3.14; I could not run the 3.11/3.12/3.14 legs locally but the change has no
version-specific code paths):

  • New regression tests in tests/test_clients.py and tests/test_adapters.py: fail with
    AttributeError on unmodified main (confirmed via git stash), pass on this branch.
    test_async_client_stop_flushes_buffered_item_and_exits_worker enqueues an item, calls
    stop() before the flush interval elapses, and asserts the mocked backend request actually
    fired and the queue drained; it caught the sentinel-ordering bug above on first run.
  • Full pytest tests/ (minus the ray/pyspark-gated tests, which are pytest.importorskip
    guarded and unrelated to this change): 140 passed, 5 skipped, no failures.
  • ruff check and ruff format --diff on the four changed files: clean. This repo's CI does
    not run a lint/static-checks job for ui/sdk-only changes (checked gh pr checks on a
    recently merged ui/sdk-only PR, fix(ui/sdk): fix failing CI on master due to formatting issues #1626), so this was a courtesy check, not a required gate.

Not verified: behavior under a real ASGI/FastAPI shutdown sequence (e.g. the app in
examples/async/fastapi_example.py, which doesn't call the new stop() either); the fix is
scoped to the tracker itself, not to updating that example.

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 apache#1214
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Stop python process from exiting too early with async tracker still having data to send

1 participant