From 01e4f687afd99f62a6731c390eec5c9ebf8f70ac Mon Sep 17 00:00:00 2001 From: Shubham Kumar Date: Fri, 16 Oct 2026 09:33:55 +0530 Subject: [PATCH] Add awaitable UDF server shutdown 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. --- ARCHITECTURE.md | 16 ++ singlestoredb/apps/__init__.py | 1 + singlestoredb/apps/_python_udfs.py | 58 +++++- singlestoredb/tests/test_function_url.py | 233 ++++++++++++++++++++++- 4 files changed, 301 insertions(+), 7 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 8720daa7a..accfb0b7e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -637,6 +637,22 @@ as SingleStore external functions. UDF servers can be deployed in three modes: a (remote), a memory-mapped collocated server (lowest latency), or as a plugin server using Unix sockets with a thread/process pool (CLI-driven). +### Managed UDF shutdown + +The notebook runtime can call `await singlestoredb.apps.shutdown_udf_app(timeout=20)` +on the same event loop that started `run_udf_app`. Uvicorn closes listeners and idle +connections, allows active requests to finish within the supplied drain budget, and +cancels requests exceeding it. The call waits for the serving task to exit and closes +any remaining connections. Repeated calls are safe; the first shutdown request sets +the budget. Passing `None` preserves the unbounded drain used when replacing an +existing server inside a kernel. + +The example's 20 seconds is not a platform default. The host must choose a budget +below its pod termination deadline, leave time for cleanup, and keep the kernel and +networking alive while awaiting shutdown. This SDK operation does not configure +Kubernetes, intercept Jupyter shutdown, or forcibly stop synchronous Python work in +worker threads; the notebook runtime must coordinate process termination afterward. + ### Architecture ``` diff --git a/singlestoredb/apps/__init__.py b/singlestoredb/apps/__init__.py index a07f228cb..b32bb9aa8 100644 --- a/singlestoredb/apps/__init__.py +++ b/singlestoredb/apps/__init__.py @@ -1,3 +1,4 @@ from ._cloud_functions import run_function_app # noqa: F401 from ._dashboards import run_dashboard_app # noqa: F401 from ._python_udfs import run_udf_app # noqa: F401 +from ._python_udfs import shutdown_udf_app # noqa: F401 diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index c061abc82..547c1ae99 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -1,4 +1,5 @@ import asyncio +import math import os import typing import warnings @@ -13,6 +14,7 @@ # Keep track of currently running server _running_server: 'typing.Optional[AwaitableUvicornServer]' = None +_running_server_task: 'typing.Optional[asyncio.Task[None]]' = None # Maximum number of UDFs allowed MAX_UDFS_LIMIT = 10 @@ -28,7 +30,7 @@ async def run_udf_app( log_level: str = 'error', kill_existing_app_server: bool = True, ) -> UdfConnectionInfo: - global _running_server + global _running_server, _running_server_task from ._uvicorn_util import AwaitableUvicornServer try: @@ -65,9 +67,7 @@ async def run_udf_app( # Shutdown the server gracefully if it was started by us. # Since the uvicorn server doesn't start a new subprocess # killing the process would result in kernel dying. - if _running_server is not None: - await _running_server.shutdown() - _running_server = None + await shutdown_udf_app(timeout=None) # Kill if any other process is occupying the port kill_process_by_port(app_config.listen_port) @@ -85,7 +85,7 @@ async def run_udf_app( app.register_interactive_functions() _running_server = AwaitableUvicornServer(config) - asyncio.create_task(_running_server.serve()) + _running_server_task = asyncio.create_task(_running_server.serve()) await _running_server.wait_for_startup() print(f'Python UDF registered at {base_url}') @@ -99,6 +99,54 @@ async def run_udf_app( return UdfConnectionInfo(base_url, app.get_function_info()) +async def shutdown_udf_app(timeout: typing.Optional[float]) -> None: + """ + Stop the UDF server started by ``run_udf_app`` in this Python kernel. + + Uvicorn closes listeners and idle connections, then allows active requests + to finish for up to ``timeout`` seconds before cancelling them. ``None`` + waits without a drain limit, preserving the existing server-replacement + behavior. A finite timeout must be non-negative. + + Call and await this operation on the same event loop as ``run_udf_app``, + before shutting down the kernel or its networking. The budget covers + request draining; allow additional time for server cleanup. Python cannot + forcibly stop arbitrary synchronous UDF code running in worker threads. + + It is safe to call when no server is running. Concurrent calls await the + same server task; the first shutdown request determines the drain budget. + """ + global _running_server, _running_server_task + + if timeout is not None and (not math.isfinite(timeout) or timeout < 0): + raise ValueError('timeout must be a finite non-negative number or None') + + server, task = _running_server, _running_server_task + if server is None: + return + if task is None: + raise RuntimeError('UDF server has no serving task') + + if not server.should_exit: + server.config.timeout_graceful_shutdown = timeout + server.should_exit = True + 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 + + def generate_base_url(app_config: AppConfig) -> str: if not app_config.is_gateway_enabled: raise RuntimeError('Python UDFs are not available if Nova Gateway is not enabled') diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py index 771f7f977..790e2603a 100644 --- a/singlestoredb/tests/test_function_url.py +++ b/singlestoredb/tests/test_function_url.py @@ -1,8 +1,14 @@ # type: ignore +import asyncio +import http.client +import json +import os import sys +import threading import unittest from collections.abc import Iterator from typing import Any +from typing import Callable from typing import List from typing import Optional from typing import Set @@ -10,7 +16,11 @@ from unittest.mock import MagicMock from unittest.mock import patch +from singlestoredb.apps import _python_udfs +from singlestoredb.apps import run_udf_app +from singlestoredb.apps import shutdown_udf_app from singlestoredb.apps._config import AppConfig +from singlestoredb.functions import udf from singlestoredb.functions.ext.function_url import classify_interactive_registration from singlestoredb.functions.ext.function_url import extract_service_url from singlestoredb.functions.ext.function_url import is_function_not_defined @@ -305,6 +315,7 @@ class TestRunUdfAppInteractiveOrder(unittest.IsolatedAsyncioTestCase): def tearDown(self) -> None: from singlestoredb.apps import _python_udfs _python_udfs._running_server = None + _python_udfs._running_server_task = None def _interactive_config(self) -> AppConfig: return AppConfig( @@ -391,12 +402,15 @@ async def test_preflight_then_shutdown_then_register(self) -> None: order: List[str] = [] existing = MagicMock() + existing.should_exit = False - async def shutdown() -> None: + async def serve() -> None: + while not existing.should_exit: + await asyncio.sleep(0) order.append('shutdown') - existing.shutdown.side_effect = shutdown _python_udfs._running_server = existing + _python_udfs._running_server_task = asyncio.create_task(serve()) app = self._mock_app(order, refuse=False) class FakeServer: @@ -431,5 +445,220 @@ async def wait_for_startup(self) -> None: ) +_request_started = threading.Event() +_release_request = threading.Event() +_request_finished = threading.Event() + + +@udf +async def _draining_udf(value: float) -> float: + _request_started.set() + try: + while not _release_request.is_set(): + await asyncio.sleep(0.01) + return value * 2 + finally: + _request_finished.set() + + +class TestUdfShutdown(unittest.IsolatedAsyncioTestCase): + + async def asyncSetUp(self) -> None: + self.env = { + 'SINGLESTOREDB_APP_LISTEN_PORT': '0', + 'SINGLESTOREDB_APP_BASE_URL': 'http://127.0.0.1', + 'SINGLESTOREDB_APP_BASE_PATH': '/', + 'SINGLESTOREDB_NOTEBOOK_SERVER_ID': 'shutdown-test', + 'SINGLESTOREDB_IS_LOCAL_DEV': 'true', + 'SINGLESTOREDB_WORKLOAD_TYPE': 'PythonUdf', + 'SINGLESTOREDB_NOVA_GATEWAY_ENDPOINT': 'http://127.0.0.1', + 'SINGLESTOREDB_EXT_FUNCTIONS': __name__, + } + self.previous_env = {name: os.environ.get(name) for name in self.env} + os.environ.update(self.env) + _request_started.clear() + _request_finished.clear() + _release_request.set() + self.connections: List[http.client.HTTPConnection] = [] + await run_udf_app(kill_existing_app_server=False) + self.server = _python_udfs._running_server + self.server_task = _python_udfs._running_server_task + assert self.server is not None + self.port = self.server.servers[0].sockets[0].getsockname()[1] + + async def asyncTearDown(self) -> None: + _release_request.set() + try: + await asyncio.wait_for(shutdown_udf_app(timeout=0.5), 3) + finally: + for connection in self.connections: + connection.close() + for name, value in self.previous_env.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + async def _wait_for(self, predicate: Callable[[], bool]) -> None: + async def wait() -> None: + while not predicate(): + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait(), 3) + + def _connect(self) -> http.client.HTTPConnection: + connection = http.client.HTTPConnection('127.0.0.1', self.port, timeout=3) + self.connections.append(connection) + return connection + + async def _invoke( + self, connection: http.client.HTTPConnection, + ) -> Tuple[int, bytes]: + def request() -> Tuple[int, bytes]: + connection.request( + 'POST', '/invoke', + json.dumps({'data': [[0, 2.125]]}), + { + 'Content-Type': 'application/json', + 'Accepts': 'application/json', + 's2-ef-name': '_draining_udf', + 's2-ef-version': '1.0', + }, + ) + response = connection.getresponse() + return response.status, response.read() + + return await asyncio.to_thread(request) + + async def test_closes_idle_connections_without_waiting_for_budget(self) -> None: + connection = self._connect() + status, body = await self._invoke(connection) + self.assertEqual(status, 200) + self.assertEqual(json.loads(body), {'data': [[0, 4.25]]}) + sock = connection.sock + assert sock is not None + + start = asyncio.get_running_loop().time() + await asyncio.wait_for(shutdown_udf_app(timeout=30), 3) + self.assertLess(asyncio.get_running_loop().time() - start, 3) + self.assertEqual(await asyncio.to_thread(sock.recv, 1), b'') + assert self.server_task is not None + self.assertTrue(self.server_task.done()) + self.assertIsNone(_python_udfs._running_server) + self.assertIsNone(_python_udfs._running_server_task) + await shutdown_udf_app(timeout=0) + + async def test_drains_active_request_and_refuses_new_connections(self) -> None: + idle = self._connect() + self.assertEqual((await self._invoke(idle))[0], 200) + idle_socket = idle.sock + assert idle_socket is not None + _request_started.clear() + _release_request.clear() + request = asyncio.create_task(self._invoke(self._connect())) + await self._wait_for(_request_started.is_set) + + shutdown = asyncio.create_task(shutdown_udf_app(timeout=2)) + server = self.server + assert server is not None + await self._wait_for( + lambda: all(not listener.is_serving() for listener in server.servers), + ) + self.assertFalse(shutdown.done()) + self.assertEqual(await asyncio.to_thread(idle_socket.recv, 1), b'') + with self.assertRaises(OSError): + await asyncio.open_connection('127.0.0.1', self.port) + + _release_request.set() + status, body = await asyncio.wait_for(request, 3) + self.assertEqual(status, 200) + self.assertEqual(json.loads(body), {'data': [[0, 4.25]]}) + await asyncio.wait_for(shutdown, 3) + self.assertTrue(_request_finished.is_set()) + + async def test_deadline_closes_over_budget_request(self) -> None: + _release_request.clear() + request = asyncio.create_task(self._invoke(self._connect())) + await self._wait_for(_request_started.is_set) + + start = asyncio.get_running_loop().time() + await asyncio.wait_for(shutdown_udf_app(timeout=0.1), 3) + elapsed = asyncio.get_running_loop().time() - start + self.assertGreaterEqual(elapsed, 0.1) + self.assertLess(elapsed, 1) + try: + status, _ = await asyncio.wait_for(request, 3) + except ( + http.client.RemoteDisconnected, + http.client.IncompleteRead, + ConnectionResetError, + ): + pass + else: + self.assertGreaterEqual(status, 500) + await self._wait_for(_request_finished.is_set) + server = self.server + assert server is not None + await self._wait_for(lambda: not server.server_state.connections) + + async def test_cancelled_waiter_does_not_cancel_drain(self) -> None: + _release_request.clear() + request = asyncio.create_task(self._invoke(self._connect())) + await self._wait_for(_request_started.is_set) + + first = asyncio.create_task(shutdown_udf_app(timeout=2)) + server = self.server + assert server is not None + await self._wait_for(lambda: server.should_exit) + first.cancel() + with self.assertRaises(asyncio.CancelledError): + await first + assert self.server_task is not None + self.assertFalse(self.server_task.done()) + + second = asyncio.create_task(shutdown_udf_app(timeout=0)) + await asyncio.sleep(0) + self.assertEqual(server.config.timeout_graceful_shutdown, 2) + _release_request.set() + status, _ = await asyncio.wait_for(request, 3) + self.assertEqual(status, 200) + await asyncio.wait_for(second, 3) + + async def test_can_start_again_after_shutdown(self) -> None: + await shutdown_udf_app(timeout=1) + await run_udf_app(kill_existing_app_server=False) + server = _python_udfs._running_server + assert server is not None + self.assertIsNot(server, self.server) + self.port = server.servers[0].sockets[0].getsockname()[1] + status, _ = await self._invoke(self._connect()) + self.assertEqual(status, 200) + + async def test_shutdown_while_starting_does_not_leave_listener_open(self) -> None: + await shutdown_udf_app(timeout=1) + startup = asyncio.create_task(run_udf_app(kill_existing_app_server=False)) + await asyncio.sleep(0) + server = _python_udfs._running_server + assert server is not None + self.assertFalse(server.started) + try: + await asyncio.wait_for(shutdown_udf_app(timeout=0), 3) + await startup + self.assertTrue(all(not listener.is_serving() for listener in server.servers)) + finally: + for listener in server.servers: + listener.close() + + async def test_rejects_invalid_budget_without_stopping_server(self) -> None: + for timeout in (-1, float('inf'), float('-inf'), float('nan')): + with self.subTest(timeout=timeout): + with self.assertRaises(ValueError): + await shutdown_udf_app(timeout=timeout) + assert self.server is not None + self.assertFalse(self.server.should_exit) + status, _ = await self._invoke(self._connect()) + self.assertEqual(status, 200) + + if __name__ == '__main__': unittest.main()