Add awaitable UDF server shutdown

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
 
 ```

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

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

singlestoredb/tests/test_function_url.py

@@ -1,16 +1,26 @@
 # 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
 from typing import Tuple
 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()