Repository navigation
Add awaitable UDF server shutdown - #143
ShubMishra wants to merge 1 commit into
Conversation
Expose a caller-budgeted shutdown operation for notebook runtime integration. Await the Uvicorn serving task, preserve existing replacement behavior, and keep shutdown coverage with the existing lifecycle tests.
6727e97 to
01e4f68
Compare
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit 01e4f68. Configure here.
| server.should_exit = True | ||
| try: | ||
| # Cancelling a caller must not interrupt the server's own drain. | ||
| await asyncio.shield(task) |
There was a problem hiding this comment.
Failed serve task breaks shutdown
Medium Severity
shutdown_udf_app awaits the serving task with asyncio.shield, so a prior serve failure is reraised. After run_udf_app fails during startup, it still leaves _running_server and _running_server_task set, so the next shutdown or replacement start surfaces that leftover error and does not bring a new server up.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 01e4f68. Configure here.
There was a problem hiding this comment.
🟡 Changes recommended
Cancelling the sole shutdown waiter can permanently skip cleanup and leave startup listeners open.
1 open issue
What changed in this PR
Adds awaitable shutdown for managed notebook UDF servers so hosts can drain requests before terminating the kernel.
Changes:
- Exports
shutdown_udf_app(timeout=...)and tracks the serving task. - Uses the shutdown API when replacing an existing server.
- Documents host responsibilities and adds lifecycle tests.
| File | Description |
|---|---|
| singlestoredb/tests/test_function_url.py | Adds shutdown and connection-draining tests. |
| singlestoredb/apps/_python_udfs.py | Implements awaitable shutdown and serving-task tracking. |
| singlestoredb/apps/__init__.py | Exports the shutdown API. |
| ARCHITECTURE.md | Documents drain budgets and host coordination. |
Review effort: Balanced
| try: | ||
| # Cancelling a caller must not interrupt the server's own drain. | ||
| await asyncio.shield(task) | ||
| finally: | ||
| if task.done(): | ||
| # Uvicorn can skip shutdown when exit was requested during startup. | ||
| if server.started: | ||
| for listener in server.servers: | ||
| listener.close() | ||
| # Uvicorn can finish after cancelling requests but before they close. | ||
| for connection in list(server.server_state.connections): | ||
| connection.transport.abort() | ||
| if _running_server is server: | ||
| _running_server = None | ||
| _running_server_task = None |
kesmit13
left a comment
kesmit13
left a comment
There was a problem hiding this comment.
I see any issues, but it might be better for @mgiannakopoulos to approve before we merge.
Summary
Add
await singlestoredb.apps.shutdown_udf_app(timeout=...)so the notebook runtime can stop the UDF HTTP server before terminating its Python kernel.This uses Uvicorn's shutdown path to close listeners/idle connections and drain active requests, waits for the serving task, and closes remaining sockets after the drain deadline. The caller supplies the budget;
Nonepreserves the existing unbounded in-kernel replacement behavior. No pod grace period or automatic signal hook is introduced.Test Plan
.venv/bin/python -m pytest --noconftest -o addopts="" -q singlestoredb/tests/test_function_url.py singlestoredb/tests/test_udf_event_loop.py— passed: 37 tests and 4 subtests, including real HTTP connections for idle closure, active-call completion, deadline expiry, cancelled waiters, restart, and shutdown during startup..venv/bin/pre-commit run --files singlestoredb/apps/_python_udfs.py singlestoredb/apps/__init__.py singlestoredb/tests/test_function_url.py ARCHITECTURE.md— passed, including lint and type checks.Helios follow-up
After this SDK change is released, a separate Helios PR will update the SDK pin/notebook image, make Jupyter await draining before kernel shutdown, and coordinate Nova's pod grace period and network cleanup. Gateway handling of abrupt failures remains separate. This provides the SDK entry point without choosing the production drain budget.
Tracked under MCDB-100416.
Note
Medium Risk
Changes notebook UDF server lifecycle and connection draining at kernel shutdown; behavior is well-tested but affects in-flight UDF HTTP traffic during restarts and pod teardown.
Overview
Adds
await singlestoredb.apps.shutdown_udf_app(timeout=...)so notebook runtimes can stop the managed UDF HTTP server on the same asyncio loop asrun_udf_app, before tearing down the kernel.Shutdown drives Uvicorn’s graceful path: listeners and idle connections close immediately, in-flight requests get up to
timeoutseconds (or unbounded whentimeout=None, matching in-kernel server replacement), then remaining sockets are aborted and globals are cleared.run_udf_appnow keeps the servingasyncio.Taskand delegates existing-server teardown to this API instead of callingserver.shutdown()directly.ARCHITECTURE.md documents host responsibilities (drain budget vs pod deadline; no K8s/Jupyter hooks). Integration tests cover idle shutdown, active-request draining, deadline expiry, cancelled waiters, restart, startup races, and invalid timeouts.
Reviewed by Cursor Bugbot for commit 01e4f68. Bugbot is set up for automated code reviews on this repo. Configure here.