Add optional max concurrency support to Runner

Author: joy13975

src/ragas/evaluation.py

@@ -13,7 +13,7 @@
 from ragas.embeddings.base import BaseRagasEmbeddings, LangchainEmbeddingsWrapper, embedding_factory
 from ragas.llms import llm_factory
 from ragas.exceptions import ExceptionInRunner
-from ragas.executor import Executor
+from ragas.executor import Executor, DEFAULT_MAX_CONCURRENCY
 from ragas.llms.base import BaseRagasLLM, LangchainLLMWrapper
 from ragas.metrics._answer_correctness import AnswerCorrectness
 from ragas.metrics.base import Metric, MetricWithEmbeddings, MetricWithLLM
@@ -44,6 +44,7 @@ def evaluate(
     run_config: t.Optional[RunConfig] = None,
     raise_exceptions: bool = True,
     column_map: t.Dict[str, str] = {},
+    max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
 ) -> Result:
     """
     Run the evaluation on the dataset with different metrics
@@ -180,6 +181,7 @@ def evaluate(
         desc="Evaluating",
         keep_progress_bar=True,
         raise_exceptions=raise_exceptions,
+        max_concurrency=max_concurrency,
     )
     # new evaluation chain
     row_run_managers = []

src/ragas/executor.py

@@ -14,6 +14,8 @@
 logger = logging.getLogger(__name__)
 
 
+DEFAULT_MAX_CONCURRENCY = 16
+
 def runner_exception_hook(args: threading.ExceptHookArgs):
     print(args)
     raise args.exc_type
@@ -22,6 +24,16 @@ def runner_exception_hook(args: threading.ExceptHookArgs):
 # set a custom exception hook
 # threading.excepthook = runner_exception_hook
 
+def as_completed(loop, coros, max_concurrency):
+    if max_concurrency == -1:
+        return asyncio.as_completed(coros, loop=loop)
+    
+    semaphore = asyncio.Semaphore(max_concurrency, loop=loop)
+    async def sem_coro(coro):
+        async with semaphore:
+            return await coro
+    
+    return asyncio.as_completed([sem_coro(c) for c in coros], loop=loop)
 
 class Runner(threading.Thread):
     def __init__(
@@ -30,26 +42,29 @@ def __init__(
         desc: str,
         keep_progress_bar: bool = True,
         raise_exceptions: bool = True,
+        max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
     ):
         super().__init__()
         self.jobs = jobs
         self.desc = desc
         self.keep_progress_bar = keep_progress_bar
         self.raise_exceptions = raise_exceptions
-        self.futures = []
+        self.max_concurrency = max_concurrency
 
         # create task
         self.loop = asyncio.new_event_loop()
-        for job in self.jobs:
-            coroutine, name = job
-            self.futures.append(self.loop.create_task(coroutine, name=name))
+        self.futures = as_completed(
+            self.loop,
+            [coro for coro, _ in self.jobs],
+            self.max_concurrency)
 
     async def _aresults(self) -> t.List[t.Any]:
         results = []
+
         for future in tqdm(
-            asyncio.as_completed(self.futures),
+            self.futures,
             desc=self.desc,
-            total=len(self.futures),
+            total=len(self.jobs),
             # whether you want to keep the progress bar after completion
             leave=self.keep_progress_bar,
         ):
@@ -85,6 +100,7 @@ class Executor:
     keep_progress_bar: bool = True
     jobs: t.List[t.Any] = field(default_factory=list, repr=False)
     raise_exceptions: bool = False
+    max_concurrency: int = DEFAULT_MAX_CONCURRENCY
 
     def wrap_callable_with_index(self, callable: t.Callable, counter):
         async def wrapped_callable_async(*args, **kwargs):
@@ -104,6 +120,7 @@ def results(self) -> t.List[t.Any]:
             desc=self.desc,
             keep_progress_bar=self.keep_progress_bar,
             raise_exceptions=self.raise_exceptions,
+            max_concurrency=self.max_concurrency,
         )
         executor_job.start()
         try:

src/ragas/testset/docstore.py

@@ -16,7 +16,7 @@
 
 from ragas.embeddings.base import BaseRagasEmbeddings
 from ragas.exceptions import ExceptionInRunner
-from ragas.executor import Executor
+from ragas.executor import Executor, DEFAULT_MAX_CONCURRENCY
 from ragas.run_config import RunConfig
 from ragas.testset.utils import rng
 
@@ -210,7 +210,11 @@ def add_documents(self, docs: t.Sequence[Document], show_progress=True):
         self.add_nodes(nodes, show_progress=show_progress)
 
     def add_nodes(
-        self, nodes: t.Sequence[Node], show_progress=True, desc: str = "embedding nodes"
+        self,
+        nodes: t.Sequence[Node],
+        show_progress=True,
+        desc: str = "embedding nodes",
+        max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
     ):
         assert self.embeddings is not None, "Embeddings must be set"
         assert self.extractor is not None, "Extractor must be set"
@@ -224,6 +228,7 @@ def add_nodes(
             desc="embedding nodes",
             keep_progress_bar=False,
             raise_exceptions=True,
+            max_concurrency=max_concurrency,
         )
         result_idx = 0
         for i, n in enumerate(nodes):

src/ragas/testset/generator.py

@@ -13,7 +13,7 @@
 from ragas._analytics import TesetGenerationEvent, track
 from ragas.embeddings.base import BaseRagasEmbeddings, LangchainEmbeddingsWrapper
 from ragas.exceptions import ExceptionInRunner
-from ragas.executor import Executor
+from ragas.executor import Executor, DEFAULT_MAX_CONCURRENCY
 from ragas.llms import BaseRagasLLM, LangchainLLMWrapper
 from ragas.run_config import RunConfig
 from ragas.testset.docstore import Document, DocumentStore, InMemoryDocumentStore
@@ -183,6 +183,7 @@ def generate(
         is_async: bool = True,
         raise_exceptions: bool = True,
         run_config: t.Optional[RunConfig] = None,
+        max_concurrency: int = DEFAULT_MAX_CONCURRENCY,
     ):
         # validate distributions
         if not check_if_sum_is_close(list(distributions.values()), 1.0, 3):
@@ -213,6 +214,7 @@ def generate(
             desc="Generating",
             keep_progress_bar=True,
             raise_exceptions=raise_exceptions,
+            max_concurrency=max_concurrency,
         )
 
         current_nodes = [