[backend:fix] fix same-kernel Python UDF registration and invocation MCDB-103581

singlestoredb/apps/_python_udfs.py

@@ -3,6 +3,7 @@
 import os
 import typing
 import warnings
+from concurrent.futures import Future
 
 from ..functions.ext.asgi import Application
 from ._config import AppConfig
@@ -14,7 +15,9 @@
 
 # Keep track of currently running server
 _running_server: 'typing.Optional[AwaitableUvicornServer]' = None
-_running_server_task: 'typing.Optional[asyncio.Task[None]]' = None
+_running_server_task: typing.Optional[
+    typing.Union[asyncio.Task[None], Future[None]]
+] = None
 
 # Maximum number of UDFs allowed
 MAX_UDFS_LIMIT = 10
@@ -32,6 +35,7 @@ async def run_udf_app(
 ) -> UdfConnectionInfo:
     global _running_server, _running_server_task
     from ._uvicorn_util import AwaitableUvicornServer
+    from ._uvicorn_util import ThreadedUvicornServer
 
     try:
         import uvicorn
@@ -84,9 +88,31 @@ async def run_udf_app(
     if app_config.running_interactively:
         app.register_interactive_functions()
 
-    _running_server = AwaitableUvicornServer(config)
-    _running_server_task = asyncio.create_task(_running_server.serve())
-    await _running_server.wait_for_startup()
+    server: AwaitableUvicornServer
+    task: typing.Union[asyncio.Task[None], Future[None]]
+    if app_config.running_interactively:
+        server = ThreadedUvicornServer(config)
+        task = server.start()
+    else:
+        server = AwaitableUvicornServer(config)
+        task = asyncio.create_task(server.serve())
+    _running_server = server
+    _running_server_task = task
+    try:
+        await server.wait_for_startup()
+    except BaseException:
+        if isinstance(server, ThreadedUvicornServer):
+            server.request_shutdown(0)
+            cleanup = asyncio.gather(
+                asyncio.wrap_future(task) if isinstance(task, Future) else task,
+                return_exceptions=True,
+            )
+            await asyncio.shield(cleanup)
+            server.join()
+            if _running_server is server:
+                _running_server = None
+                _running_server_task = None
+        raise
 
     print(f'Python UDF registered at {base_url}')
     if app_config.running_interactively:
@@ -108,8 +134,10 @@ async def shutdown_udf_app(timeout: typing.Optional[float]) -> 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
+    Call and await this operation before shutting down the kernel or its
+    networking. Published UDFs require the same event loop as ``run_udf_app``.
+    Interactive UDFs forward shutdown to their dedicated server loop, so the
+    caller may use another event loop in the same kernel. The budget covers
     request draining; allow additional time for server cleanup. Python cannot
     forcibly stop arbitrary synchronous UDF code running in worker threads.
 
@@ -127,21 +155,30 @@ async def shutdown_udf_app(timeout: typing.Optional[float]) -> None:
     if task is None:
         raise RuntimeError('UDF server has no serving task')
 
-    if not server.should_exit:
+    from ._uvicorn_util import ThreadedUvicornServer
+
+    if isinstance(server, ThreadedUvicornServer):
+        server.request_shutdown(timeout)
+    elif 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)
+        await asyncio.shield(
+            asyncio.wrap_future(task) if isinstance(task, Future) else 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 isinstance(server, ThreadedUvicornServer):
+                server.join()
+            else:
+                # 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

singlestoredb/apps/_uvicorn_util.py

@@ -1,5 +1,9 @@
 import asyncio
+import contextvars
+import logging
 import socket
+import threading
+from concurrent.futures import Future
 from typing import List
 from typing import Optional
 try:
@@ -17,20 +21,125 @@ class AwaitableUvicornServer(uvicorn.Server):
 
     def __init__(self, config: 'uvicorn.Config') -> None:
         super().__init__(config)
-        self._startup_future = asyncio.get_event_loop().create_future()
+        self._startup_future: Future[None] = Future()
 
     async def startup(self, sockets: Optional[List[socket.socket]] = None) -> None:
         try:
             result = await super().startup(sockets)
-            self._startup_future.set_result(True)
+            self._startup_future.set_result(None)
             return result
         except Exception as error:
             self._startup_future.set_exception(error)
             raise error
 
     async def wait_for_startup(self) -> None:
-        await self._startup_future
+        await asyncio.shield(asyncio.wrap_future(self._startup_future))
 
     async def shutdown(self, sockets: Optional[list[socket.socket]] = None) -> None:
         if self.started:
             await super().shutdown(sockets)
+
+
+class ThreadedUvicornServer(AwaitableUvicornServer):
+    """Serve interactive UDFs independently of a blocking notebook cell."""
+
+    should_exit: bool
+
+    def __init__(self, config: 'uvicorn.Config') -> None:
+        super().__init__(config)
+        self._finished: Future[None] = Future()
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._lock = threading.Lock()
+        self._shutdown_requested = False
+        self._shutdown_timeout: Optional[float] = None
+        self._thread = threading.Thread(
+            target=contextvars.copy_context().run,
+            args=(self._run_server_thread,),
+            name='singlestoredb-udf-server',
+            daemon=True,
+        )
+
+    def start(self) -> Future[None]:
+        try:
+            self._thread.start()
+        except Exception as exc:
+            self._startup_future.set_exception(exc)
+            self._finished.set_exception(exc)
+        return self._finished
+
+    def request_shutdown(self, timeout: Optional[float]) -> None:
+        with self._lock:
+            if not self._shutdown_requested:
+                self._shutdown_timeout = timeout
+                self._shutdown_requested = True
+            if self._loop is not None:
+                self._loop.call_soon_threadsafe(self._begin_shutdown)
+
+    def _begin_shutdown(self) -> None:
+        if not self.should_exit:
+            self.config.timeout_graceful_shutdown = self._shutdown_timeout
+            self.should_exit = True
+
+    async def shutdown(self, sockets: Optional[List[socket.socket]] = None) -> None:
+        await super().shutdown(sockets)
+        if not self.force_exit:
+            await asyncio.gather(*self.server_state.tasks, return_exceptions=True)
+
+    def join(self) -> None:
+        if self._thread.ident is not None:
+            self._thread.join()
+
+    def _run_server_thread(self) -> None:
+        try:
+            self._run_loop()
+        except BaseException as exc:
+            logging.getLogger('uvicorn.error').exception('UDF server thread failed')
+            error: BaseException = exc
+            if isinstance(exc, SystemExit):
+                error = RuntimeError(f'UDF server exited with code {exc.code}')
+                error.__cause__ = exc
+            if not self._startup_future.done():
+                self._startup_future.set_exception(error)
+            self._finished.set_exception(error)
+        else:
+            if not self._startup_future.done():
+                self._startup_future.set_exception(
+                    RuntimeError('UDF server stopped before startup'),
+                )
+            self._finished.set_result(None)
+
+    def _run_loop(self) -> None:
+        loop = asyncio.new_event_loop()
+        try:
+            asyncio.set_event_loop(loop)
+            with self._lock:
+                self._loop = loop
+                if self._shutdown_requested:
+                    self._begin_shutdown()
+            try:
+                loop.run_until_complete(self.serve())
+            finally:
+                with self._lock:
+                    self._loop = None
+                if self.started:
+                    for listener in self.servers:
+                        listener.close()
+                for connection in list(self.server_state.connections):
+                    connection.transport.abort()
+                pending = asyncio.all_tasks(loop)
+                for task in pending:
+                    task.cancel()
+                if pending:
+                    results = loop.run_until_complete(
+                        asyncio.gather(*pending, return_exceptions=True),
+                    )
+                    for result in results:
+                        if isinstance(result, Exception):
+                            loop.call_exception_handler({
+                                'message': 'Unexpected error during UDF server shutdown',
+                                'exception': result,
+                            })
+                loop.run_until_complete(loop.shutdown_asyncgens())
+        finally:
+            asyncio.set_event_loop(None)
+            loop.close()

singlestoredb/functions/ext/asgi.py

@@ -69,7 +69,9 @@
 from . import utils
 from ... import connection
 from ...config import get_option
+from ...exceptions import DatabaseError
 from ...management.stage import get_stage
+from ...mysql.constants import ER
 from ...mysql.constants import FIELD_TYPE as ft
 from ..signature import get_signature
 from ..signature import signature_to_sql
@@ -1892,7 +1894,41 @@ def _show_create_service_url(self, cur: Any, sql_name: str) -> Optional[str]:
             )
         else:
             qualified = escape_name(sql_name)
-        return self._service_url_for_qualified(cur, qualified)
+        try:
+            return self._service_url_for_qualified(cur, qualified)
+        except DatabaseError as exc:
+            if exc.errno != ER.WRONG_OBJECT:
+                raise
+
+        database = self.function_database or None
+        cur.execute(
+            'SELECT ROUTINE_TYPE FROM information_schema.ROUTINES '
+            'WHERE ROUTINE_SCHEMA = COALESCE(%s, DATABASE()) '
+            'AND LOWER(ROUTINE_NAME) = LOWER(%s) '
+            'UNION ALL SELECT TABLE_TYPE FROM information_schema.TABLES '
+            'WHERE TABLE_SCHEMA = COALESCE(%s, DATABASE()) '
+            'AND LOWER(TABLE_NAME) = LOWER(%s) '
+            'LIMIT 1',
+            (database, sql_name, database, sql_name),
+        )
+        conflicts = list(cur)
+        if not conflicts:
+            query = 'SHOW AGGREGATES'
+            if database:
+                query += f' IN {escape_name(database)}'
+            pattern = (
+                sql_name.replace('\\', '\\\\').replace('_', '\\_').replace('%', '\\%')
+            )
+            cur.execute(query + ' LIKE %s', (pattern,))
+            if list(cur):
+                conflicts = [('AGGREGATE',)]
+        if conflicts:
+            raise RuntimeError(
+                f'Cannot register SQL function `{sql_name}`: '
+                f'name is already used by a {conflicts[0][0]}. '
+                'Interactive registration will not replace an existing SQL object.',
+            )
+        return None
 
     def _classify_interactive_name(self, cur: Any, sql_name: str) -> str:
         existing = self._show_create_service_url(cur, sql_name)