[FSTORE-1656][4.2] Fetching feature vectors in the predictor file of a deployment fails with event loop is already running error (logicalclocks#498)

* add thredpoolexecutor for running loops

* removing prints and reverting get_running_loop changes

* adding back set event loop

* moving creation of thread pool to the init function

* adding connection options for timeout

* making the number of threads in the thread pool configurable

* creating tasks using a event loop

* spawning threads instead of a thread pool executor to not not block calls of get_feature_vector from differnt threads

* creating a single thread to run all async tasks

* updating comments

* moving the connection pool variable into the Async thread

---------

Co-authored-by: DhananjayMukhedkar <dhananjay1098@gmail.com>

Author: manu-sj

python/hopsworks_common/util.py

@@ -16,10 +16,12 @@
 
 from __future__ import annotations
 
+import asyncio
 import inspect
 import itertools
 import json
 import os
+import queue
 import re
 import shutil
 import sys
@@ -713,3 +715,191 @@ def feature_view_to_json(obj):
 
             return humps.camelize(json.loads(obj.json()))
     return None
+
+
+class AsyncTask:
+    """
+    Generic class to represent an async task.
+
+    Args:
+        func (Callable): The function to run asynchronously.
+        requires_connection_pool (bool): Whether the task requires a connection pool.
+        **kwargs: Key word arguments to be passed to the functions.
+
+    Properties:
+        result (Any): The result of the async task.
+        event (threading.Event): The event that will be set when the async task is finished.
+    """
+
+    def __init__(
+        self,
+        task_function: Callable,
+        task_args: Tuple = (),
+        requires_connection_pool=None,
+        **kwargs,
+    ):
+        self.task_function = task_function
+        self.task_args = task_args
+        self.task_kwargs = kwargs
+        self._event: threading.Event = threading.Event()
+        self._result: Any = None
+        self._requires_connection_pool = requires_connection_pool
+
+    @property
+    def result(self) -> Any:
+        """
+        The result of the async task.
+        """
+        return self._result
+
+    @result.setter
+    def result(self, value) -> None:
+        self._result = value
+
+    @property
+    def event(self) -> threading.Event:
+        """
+        The event that will be set when the async task is finished.
+        """
+        return self._event
+
+    @event.setter
+    def event(self, value) -> None:
+        self._event = value
+
+    @property
+    def requires_connection_pool(self) -> bool:
+        """
+        Whether the task requires a connection pool.
+        """
+        return self._requires_connection_pool
+
+
+class AsyncTaskThread(threading.Thread):
+    """
+    Generic thread class that can be used to run async tasks in a separate thread.
+    The thread will create its own event loop and run submitted tasks in that loop.
+
+    The thread also store and fetches a connection pool that can be used by the async tasks.
+
+    # Args:
+        connection_pool_initializer (Callable): A function that initializes a connection pool.
+        connection_pool_params (Tuple): The parameters to pass to the connection pool initializer.
+        *thread_args: Arguments to be passed to the thread.
+        **thread_kwargs: Key word arguments to be passed to the thread.
+
+    # Properties:
+        event_loop (asyncio.AbstractEventLoop): The event loop used by the thread.
+        task_queue (queue.Queue[AsyncTask]): The queue used to submit tasks to the thread.
+        connection_pool: The connection pool used
+    """
+
+    def __init__(
+        self,
+        connection_pool_initializer: Callable = None,
+        connection_pool_params: Tuple = (),
+        *thread_args,
+        **thread_kwargs,
+    ):
+        super().__init__(*thread_args, **thread_kwargs)
+        self._task_queue: queue.Queue[AsyncTask] = queue.Queue()
+        self._event_loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
+        self.stop_event = threading.Event()
+        self._connection_pool_initializer: Callable = connection_pool_initializer
+        self._connection_pool_params: Tuple = connection_pool_params
+        self._connection_pool = None
+        self.daemon = True  # Setting the thread as a daemon thread by default, so it will be terminated when the main thread is terminated.
+
+    async def execute_task(self):
+        """
+        Execute the async tasks for the queue.
+        """
+        asyncio.set_event_loop(self._event_loop)
+
+        if self._connection_pool_initializer:
+            self._connection_pool = await self._connection_pool_initializer(
+                *self._connection_pool_params
+            )
+
+        while not self.stop_event.is_set():
+            # Fetch a task from the queue.
+            task = self.task_queue.get()
+            # Run the task in the event loop and get the result
+            try:
+                # Check if the task requires a connection pool and pass it to the function if it does.
+                if task.requires_connection_pool:
+                    task.result = await task.task_function(
+                        *task.task_args,
+                        **task.task_kwargs,
+                        connection_pool=self.connection_pool,
+                    )
+                else:
+                    task.result = await task.task_function(
+                        *task.task_args, **task.task_kwargs
+                    )
+
+                # Unblock the task, so the submit function can return the result.
+                task.event.set()
+            except Exception as e:
+                task.result = e
+                task.event.set()
+                raise e
+
+    def stop(self):
+        """
+        Stop the thread and close the event loop.
+        """
+        self.stop_event.set()
+        self._event_loop.stop()
+        self._event_loop.close()
+
+    def run(self):
+        """
+        Execute the async tasks for the queue.
+        """
+        asyncio.set_event_loop(self._event_loop)
+        self._event_loop.create_task(self.execute_task())
+        try:
+            self._event_loop.run_forever()
+        except Exception as e:
+            self._event_loop.stop()
+            self._event_loop.close()
+            raise e
+        finally:
+            self._event_loop.close()
+
+    def submit(self, task: AsyncTask):
+        """
+        Submit a async task to the thread and block until the execution of the function is completed.
+        """
+        # Submit a task to the queue.
+        self.task_queue.put(task)
+        # Block the execution until the task is finished.
+        task.event.wait()
+
+        if isinstance(task.result, Exception):
+            raise task.result
+        else:
+            # Return the result of the task.
+            return task.result
+
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        """
+        The event loop used by the thread.
+        """
+        return self._event_loop
+
+    @property
+    def task_queue(self) -> queue.Queue[AsyncTask]:
+        """
+        The queue used to submit tasks to the thread.
+        """
+        return self._task_queue
+
+    @property
+    def connection_pool(self):
+        """
+        The connection pool used by the thread.
+        """
+        return self._connection_pool

python/hsfs/core/online_store_sql_engine.py

@@ -22,6 +22,7 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
 
 from hopsworks_common.core import variable_api
+from hopsworks_common.util import AsyncTask, AsyncTaskThread
 from hsfs import util
 from hsfs.core import (
     feature_view_api,
@@ -73,7 +74,6 @@ def __init__(
         self._prefix_by_serving_index = None
         self._pkname_by_serving_index = None
         self._serving_key_by_serving_index: Dict[str, ServingKey] = {}
-        self._connection_pool = None
         self._serving_keys: Set[ServingKey] = set(serving_keys or [])
 
         self._prepared_statements: Dict[str, List[ServingPreparedStatement]] = {}
@@ -89,6 +89,14 @@ def __init__(
         self._hostname = None
         self._connection_options = None
 
+        self._async_task_thread = None
+
+    def __del__(self):
+        # Safely stop the async task thread.
+        # The connection pool will be closed during garbage collection by aiomysql.
+        if self._async_task_thread.is_alive():
+            self._async_task_thread.stop()
+
     def fetch_prepared_statements(
         self,
         entity: Union[feature_view.FeatureView, training_dataset.TrainingDataset],
@@ -245,13 +253,15 @@ def init_async_mysql_connection(self, options=None):
             else None
         )
 
-        if util.is_runtime_notebook():
-            _logger.debug("Running in Jupyter notebook, applying nest_asyncio")
-            import nest_asyncio
-
-            nest_asyncio.apply()
-        else:
-            _logger.debug("Running in python script. Not applying nest_asyncio")
+        if not self._async_task_thread:
+            # Create the async event thread if it is not already running and start it.
+            self._async_task_thread = AsyncTaskThread(
+                connection_pool_initializer=self._get_connection_pool,
+                connection_pool_params=(
+                    len(self._prepared_statements[self.SINGLE_VECTOR_KEY]),
+                ),
+            )
+            self._async_task_thread.start()
 
     def get_single_feature_vector(self, entry: Dict[str, Any]) -> Dict[str, Any]:
         """Retrieve single vector with parallel queries using aiomysql engine."""
@@ -325,9 +335,15 @@ def _single_vector_result(
         _logger.debug(
             f"Executing prepared statements for serving vector with entries: {bind_entries}"
         )
-        loop = self._get_or_create_event_loop()
-        results_dict = loop.run_until_complete(
-            self._execute_prep_statements(prepared_statement_execution, bind_entries)
+        results_dict = self._async_task_thread.submit(
+            AsyncTask(
+                task_function=self._execute_prep_statements,
+                task_args=(
+                    prepared_statement_execution,
+                    bind_entries,
+                ),
+                requires_connection_pool=True,
+            )
         )
         _logger.debug(f"Retrieved feature vectors: {results_dict}")
         _logger.debug("Constructing serving vector from results")
@@ -393,9 +409,12 @@ def _batch_vector_results(
             f"Executing prepared statements for batch vector with entries: {entry_values}"
         )
         # run all the prepared statements in parallel using aiomysql engine
-        loop = self._get_or_create_event_loop()
-        parallel_results = loop.run_until_complete(
-            self._execute_prep_statements(prepared_stmts_to_execute, entry_values)
+        parallel_results = self._async_task_thread.submit(
+            AsyncTask(
+                task_function=self._execute_prep_statements,
+                task_args=(prepared_stmts_to_execute, entry_values),
+                requires_connection_pool=True,
+            )
         )
 
         _logger.debug(f"Retrieved feature vectors: {parallel_results}, stitching them.")
@@ -441,20 +460,6 @@ def _batch_vector_results(
                 )
         return batch_results, serving_keys_all_fg
 
-    def _get_or_create_event_loop(self):
-        try:
-            _logger.debug("Acquiring or starting event loop for async engine.")
-            loop = asyncio.get_event_loop()
-            asyncio.set_event_loop(loop)
-        except RuntimeError as ex:
-            if "There is no current event loop in thread" in str(ex):
-                _logger.debug(
-                    "No existing running event loop. Creating new event loop."
-                )
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-        return loop
-
     def refresh_mysql_connection(self):
         _logger.debug("Refreshing MySQL connection.")
         try:
@@ -547,22 +552,24 @@ def get_prepared_statement_labels(
         ]
 
     async def _get_connection_pool(self, default_min_size: int) -> None:
-        self._connection_pool = await util_sql.create_async_engine(
+        connection_pool = await util_sql.create_async_engine(
             self._online_connector,
             self._external,
             default_min_size,
             options=self._connection_options,
             hostname=self._hostname,
         )
+        return connection_pool
 
-    async def _query_async_sql(self, stmt, bind_params):
+    async def _query_async_sql(
+        self,
+        stmt,
+        bind_params,
+        connection_pool: aiomysql.utils._ConnectionContextManager,
+    ):
         """Query prepared statement together with bind params using aiomysql connection pool"""
         # create connection pool
-        await self._get_connection_pool(
-            len(self._prepared_statements[self.SINGLE_VECTOR_KEY])
-        )
-
-        async with self._connection_pool.acquire() as conn:
+        async with connection_pool.acquire() as conn:
             # Execute the prepared statement
             _logger.debug(
                 f"Executing prepared statement: {stmt} with bind params: {bind_params}"
@@ -580,6 +587,7 @@ async def _execute_prep_statements(
         self,
         prepared_statements: Dict[int, str],
         entries: Union[List[Dict[str, Any]], Dict[str, Any]],
+        connection_pool: aiomysql.utils._ConnectionContextManager,  # The connection pool required is passed as a parameter from the AsyncTaskThread.
     ):
         """Iterate over prepared statements to create async tasks
         and gather all tasks results for a given list of entries."""
@@ -595,7 +603,9 @@ async def _execute_prep_statements(
         try:
             tasks = [
                 asyncio.create_task(
-                    self._query_async_sql(prepared_statements[key], entries[key]),
+                    self._query_async_sql(
+                        prepared_statements[key], entries[key], connection_pool
+                    ),
                     name="query_prep_statement_key" + str(key),
                 )
                 for key in prepared_statements
@@ -749,7 +759,3 @@ def connection_options(self) -> Dict[str, Any]:
     @property
     def online_connector(self) -> storage_connector.StorageConnector:
         return self._online_connector
-
-    @property
-    def connection_pool(self) -> aiomysql.utils._ConnectionContextManager:
-        return self._connection_pool