Migrated arg-based max_concurrency to RunConfig.

Also renamed max_concurrency to max_workers to be consistent with
convention.

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, DEFAULT_MAX_CONCURRENCY
+from ragas.executor import Executor
 from ragas.llms.base import BaseRagasLLM, LangchainLLMWrapper
 from ragas.metrics._answer_correctness import AnswerCorrectness
 from ragas.metrics.base import Metric, MetricWithEmbeddings, MetricWithLLM
@@ -40,7 +40,6 @@ def evaluate(
     embeddings: t.Optional[BaseRagasEmbeddings] = None,
     callbacks: Callbacks = [],
     is_async: bool = False,
-    max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
     run_config: t.Optional[RunConfig] = None,
     raise_exceptions: bool = True,
     column_map: t.Dict[str, str] = {},
@@ -73,9 +72,6 @@ def evaluate(
         evaluation is run by calling the `metric.ascore` method. In case the llm or
         embeddings does not support async then the evaluation can be run in sync mode
         with `is_async=False`. Default is False.
-    max_concurrency: int, optional
-        The number of workers to use for the evaluation. This is used by the
-        `ThreadpoolExecutor` to run the evaluation in sync mode.
     run_config: RunConfig, optional
         Configuration for runtime settings like timeout and retries. If not provided,
         default values are used.
@@ -124,8 +120,7 @@ def evaluate(
         raise ValueError("Provide dataset!")
 
     # default run_config
-    if run_config is None:
-        run_config = RunConfig()
+    run_config = run_config or RunConfig()
     # default metrics
     if metrics is None:
         from ragas.metrics import (
@@ -180,7 +175,7 @@ def evaluate(
         desc="Evaluating",
         keep_progress_bar=True,
         raise_exceptions=raise_exceptions,
-        max_concurrency=max_concurrency,
+        run_config=run_config,
     )
     # new evaluation chain
     row_run_managers = []

src/ragas/executor.py

@@ -11,12 +11,11 @@
 from tqdm.auto import tqdm
 
 from ragas.exceptions import MaxRetriesExceeded
+from ragas.run_config import RunConfig
 
 logger = logging.getLogger(__name__)
 
 
-DEFAULT_MAX_CONCURRENCY = 16
-
 def runner_exception_hook(args: threading.ExceptHookArgs):
     print(args)
     raise args.exc_type
@@ -25,13 +24,13 @@ 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:
+def as_completed(loop, coros, max_workers):
+    if max_workers == -1:
         return asyncio.as_completed(coros, loop=loop)
 
     # loop argument is removed since Python 3.10
     semaphore = asyncio.Semaphore(
-        max_concurrency,
+        max_workers,
         **({"loop": loop} if sys.version_info[:2] < (3, 10) else {})
     )
     async def sem_coro(coro):
@@ -47,25 +46,25 @@ def __init__(
         desc: str,
         keep_progress_bar: bool = True,
         raise_exceptions: bool = True,
-        max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
+        run_config: RunConfig = None,
     ):
         super().__init__()
         self.jobs = jobs
         self.desc = desc
         self.keep_progress_bar = keep_progress_bar
         self.raise_exceptions = raise_exceptions
-        self.max_concurrency = max_concurrency
+        self.run_config = run_config or RunConfig()
 
         # create task
         self.loop = asyncio.new_event_loop()
         self.futures = as_completed(
-            self.loop,
-            [coro for coro, _ in self.jobs],
-            self.max_concurrency)
+            loop=self.loop,
+            coros=[coro for coro, _ in self.jobs],
+            max_workers=self.run_config.max_workers
+        )
 
     async def _aresults(self) -> t.List[t.Any]:
         results = []
-
         for future in tqdm(
             self.futures,
             desc=self.desc,
@@ -95,7 +94,6 @@ def run(self):
             results = self.loop.run_until_complete(self._aresults())
         finally:
             self.results = results
-            [f.cancel() for f in self.futures]
             self.loop.stop()
 
 
@@ -105,7 +103,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: t.Optional[int] = DEFAULT_MAX_CONCURRENCY
+    run_config: RunConfig = None
 
     def wrap_callable_with_index(self, callable: t.Callable, counter):
         async def wrapped_callable_async(*args, **kwargs):
@@ -125,7 +123,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,
+            run_config=self.run_config,
         )
         executor_job.start()
         try:

src/ragas/run_config.py

@@ -20,6 +20,7 @@ class RunConfig:
     timeout: int = 60
     max_retries: int = 10
     max_wait: int = 60
+    max_workers: int = 16
     exception_types: t.Union[
         t.Type[BaseException],
         t.Tuple[t.Type[BaseException], ...],

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, DEFAULT_MAX_CONCURRENCY
+from ragas.executor import Executor
 from ragas.run_config import RunConfig
 from ragas.testset.utils import rng
 
@@ -83,27 +83,16 @@ class Direction(str, Enum):
     UP = "up"
     DOWN = "down"
 
-
 class DocumentStore(ABC):
     def __init__(self):
         self.documents = {}
 
     @abstractmethod
-    def add_documents(
-        self,
-        docs: t.Sequence[Document],
-        show_progress=True,
-        max_concurrency: t.Optional[int]=DEFAULT_MAX_CONCURRENCY
-    ):
+    def add_documents(self, docs: t.Sequence[Document], show_progress=True):
         ...
 
     @abstractmethod
-    def add_nodes(
-        self,
-        nodes: t.Sequence[Node],
-        show_progress=True,
-        max_concurrency: t.Optional[int]=DEFAULT_MAX_CONCURRENCY
-    ):
+    def add_nodes(self, nodes: t.Sequence[Node], show_progress=True):
         ...
 
     @abstractmethod
@@ -201,16 +190,12 @@ class InMemoryDocumentStore(DocumentStore):
     nodes: t.List[Node] = field(default_factory=list)
     node_embeddings_list: t.List[Embedding] = field(default_factory=list)
     node_map: t.Dict[str, Node] = field(default_factory=dict)
+    run_config: RunConfig = None
 
     def _embed_items(self, items: t.Union[t.Sequence[Document], t.Sequence[Node]]):
         ...
 
-    def add_documents(
-            self,
-            docs: t.Sequence[Document],
-            show_progress=True,
-            max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY
-        ):
+    def add_documents(self, docs: t.Sequence[Document], show_progress=True):
         """
         Add documents in batch mode.
         """
@@ -222,15 +207,9 @@ def add_documents(
             for d in self.splitter.transform_documents(docs)
         ]
 
-        self.add_nodes(nodes, show_progress=show_progress, max_concurrency=max_concurrency)
+        self.add_nodes(nodes, show_progress=show_progress)
 
-    def add_nodes(
-        self,
-        nodes: t.Sequence[Node],
-        show_progress=True,
-        desc: str = "embedding nodes",
-        max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
-    ):
+    def add_nodes(self, nodes: t.Sequence[Node], show_progress=True):
         assert self.embeddings is not None, "Embeddings must be set"
         assert self.extractor is not None, "Extractor must be set"
 
@@ -243,7 +222,7 @@ def add_nodes(
             desc="embedding nodes",
             keep_progress_bar=False,
             raise_exceptions=True,
-            max_concurrency=max_concurrency,
+            run_config=self.run_config,
         )
         result_idx = 0
         for i, n in enumerate(nodes):
@@ -341,3 +320,4 @@ def get_adjacent(
     def set_run_config(self, run_config: RunConfig):
         if self.embeddings:
             self.embeddings.set_run_config(run_config)
+        self.run_config = run_config

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, DEFAULT_MAX_CONCURRENCY
+from ragas.executor import Executor
 from ragas.llms import BaseRagasLLM, LangchainLLMWrapper
 from ragas.run_config import RunConfig
 from ragas.testset.docstore import Document, DocumentStore, InMemoryDocumentStore
@@ -78,6 +78,7 @@ def with_openai(
         embeddings: str = "text-embedding-ada-002",
         docstore: t.Optional[DocumentStore] = None,
         chunk_size: int = 512,
+        run_config: RunConfig = None,
     ) -> "TestsetGenerator":
         generator_llm_model = LangchainLLMWrapper(ChatOpenAI(model=generator_llm))
         critic_llm_model = LangchainLLMWrapper(ChatOpenAI(model=critic_llm))
@@ -93,6 +94,7 @@ def with_openai(
                 splitter=splitter,
                 embeddings=embeddings_model,
                 extractor=keyphrase_extractor,
+                run_config=run_config,
             )
             return cls(
                 generator_llm=generator_llm_model,
@@ -119,12 +121,10 @@ def generate_with_llamaindex_docs(
         is_async: bool = True,
         raise_exceptions: bool = True,
         run_config: t.Optional[RunConfig] = None,
-        max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
     ):
         # chunk documents and add to docstore
         self.docstore.add_documents(
-            [Document.from_llamaindex_document(doc) for doc in documents],
-            max_concurrency=max_concurrency,
+            [Document.from_llamaindex_document(doc) for doc in documents]
         )
 
         return self.generate(
@@ -134,7 +134,6 @@ def generate_with_llamaindex_docs(
             is_async=is_async,
             run_config=run_config,
             raise_exceptions=raise_exceptions,
-            max_concurrency=max_concurrency,
         )
 
     # if you add any arguments to this function, make sure to add them to
@@ -148,12 +147,10 @@ def generate_with_langchain_docs(
         is_async: bool = True,
         raise_exceptions: bool = True,
         run_config: t.Optional[RunConfig] = None,
-        max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
     ):
         # chunk documents and add to docstore
         self.docstore.add_documents(
-            [Document.from_langchain_document(doc) for doc in documents],
-            max_concurrency=max_concurrency,
+            [Document.from_langchain_document(doc) for doc in documents]
         )
 
         return self.generate(
@@ -163,7 +160,6 @@ def generate_with_langchain_docs(
             is_async=is_async,
             raise_exceptions=raise_exceptions,
             run_config=run_config,
-            max_concurrency=max_concurrency,
         )
 
     def init_evolution(self, evolution: Evolution) -> None:
@@ -189,7 +185,6 @@ def generate(
         is_async: bool = True,
         raise_exceptions: bool = True,
         run_config: t.Optional[RunConfig] = None,
-        max_concurrency: t.Optional[int] = DEFAULT_MAX_CONCURRENCY,
     ):
         # validate distributions
         if not check_if_sum_is_close(list(distributions.values()), 1.0, 3):
@@ -220,7 +215,7 @@ def generate(
             desc="Generating",
             keep_progress_bar=True,
             raise_exceptions=raise_exceptions,
-            max_concurrency=max_concurrency,
+            run_config=run_config,
         )
 
         current_nodes = [