Add concurrency limit option

Author: kesmit13

singlestoredb/functions/decorator.py

@@ -103,6 +103,7 @@ def _func(
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
     timeout: Optional[int] = None,
+    concurrency_limit: Optional[int] = None,
 ) -> UDFType:
     """Generic wrapper for UDF and TVF decorators."""
 
@@ -112,6 +113,7 @@ def _func(
             args=expand_types(args),
             returns=expand_types(returns),
             timeout=timeout,
+            concurrency_limit=concurrency_limit,
         ).items() if v is not None
     }
 
@@ -155,6 +157,7 @@ def udf(
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
     timeout: Optional[int] = None,
+    concurrency_limit: Optional[int] = None,
 ) -> UDFType:
     """
     Define a user-defined function (UDF).
@@ -185,6 +188,10 @@ def udf(
     timeout : int, optional
         The timeout in seconds for the UDF execution. If not specified,
         the global default timeout is used.
+    concurrency_limit : int, optional
+        The maximum number of concurrent subsets of rows that will be
+        processed simultaneously by the UDF. If not specified,
+        the global default concurrency limit is used.
 
     Returns
     -------
@@ -197,4 +204,5 @@ def udf(
         args=args,
         returns=returns,
         timeout=timeout,
+        concurrency_limit=concurrency_limit,
     )

singlestoredb/functions/ext/asgi.py

@@ -301,7 +301,6 @@ async def run_in_parallel(
     func: Callable[..., Any],
     params_list: Sequence[Sequence[Any]],
     cancel_event: threading.Event,
-    limit: int = get_option('external_function.concurrency_limit'),
     transformer: Callable[[Any], Any] = identity,
 ) -> List[Any]:
     """"
@@ -315,8 +314,6 @@ async def run_in_parallel(
         The parameters to pass to the function
     cancel_event : threading.Event
         The event to check for cancellation
-    limit : int
-        The maximum number of concurrent tasks to run
     transformer : Callable[[Any], Any]
         A function to transform the results
 
@@ -326,6 +323,7 @@ async def run_in_parallel(
         The results of the function calls
 
     """
+    limit = get_concurrency_limit(func)
     is_async = asyncio.iscoroutinefunction(func)
 
     async def call(batch: Sequence[Any]) -> Any:
@@ -352,6 +350,29 @@ async def thread_call(batch: Sequence[Any]) -> Any:
     return list(itertools.chain.from_iterable(results))
 
 
+def get_concurrency_limit(func: Callable[..., Any]) -> int:
+    """
+    Get the concurrency limit for a function.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to get the concurrency limit for
+
+    Returns
+    -------
+    int
+        The concurrency limit for the function
+
+    """
+    return max(
+        1, func._singlestoredb_attrs.get(  # type: ignore
+            'concurrency_limit',
+            get_option('external_function.concurrency_limit'),
+        ),
+    )
+
+
 def build_udf_endpoint(
     func: Callable[..., Any],
     returns_data_format: str,