arize 8.0.0a21__py3-none-any.whl → 8.0.0a23__py3-none-any.whl

arize/experiments/evaluators/executors.py

@@ -1,3 +1,5 @@
+"""Evaluator executors for managing concurrent evaluation tasks."""
+
 from __future__ import annotations
 
 import asyncio
@@ -10,18 +12,14 @@ import traceback
 from contextlib import contextmanager
 from enum import Enum
 from typing import (
+    TYPE_CHECKING,
     Any,
-    Callable,
-    Coroutine,
-    Generator,
-    List,
-    Optional,
     Protocol,
-    Sequence,
-    Tuple,
-    Union,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine, Generator, Sequence
+
 from tqdm.auto import tqdm
 
 from arize.experiments.evaluators.exceptions import ArizeException
@@ -30,13 +28,15 @@ logger = logging.getLogger(__name__)
 
 
 class Unset:
-    pass
+    """Sentinel class representing an unset or undefined value."""
 
 
 _unset = Unset()
 
 
 class ExecutionStatus(Enum):
+    """Enum representing execution status states for experiment tasks."""
+
     DID_NOT_RUN = "DID NOT RUN"
     COMPLETED = "COMPLETED"
     COMPLETED_WITH_RETRIES = "COMPLETED WITH RETRIES"
@@ -44,42 +44,52 @@ class ExecutionStatus(Enum):
 
 
 class ExecutionDetails:
+    """Container for tracking execution status, exceptions, and runtime metrics."""
+
     def __init__(self) -> None:
-        self.exceptions: List[Exception] = []
+        """Initialize execution details with default state."""
+        self.exceptions: list[Exception] = []
         self.status = ExecutionStatus.DID_NOT_RUN
         self.execution_seconds: float = 0
 
     def fail(self) -> None:
+        """Mark this execution as failed."""
         self.status = ExecutionStatus.FAILED
 
     def complete(self) -> None:
+        """Mark this execution as completed, with or without retries."""
         if self.exceptions:
             self.status = ExecutionStatus.COMPLETED_WITH_RETRIES
         else:
             self.status = ExecutionStatus.COMPLETED
 
     def log_exception(self, exc: Exception) -> None:
+        """Log an exception that occurred during execution."""
         self.exceptions.append(exc)
 
     def log_runtime(self, start_time: float) -> None:
+        """Log the runtime duration for this execution."""
         self.execution_seconds += time.time() - start_time
 
 
 class Executor(Protocol):
+    """Protocol defining the interface for experiment task executors."""
+
     def run(
         self, inputs: Sequence[Any]
-    ) -> Tuple[List[Any], List[ExecutionDetails]]: ...
+    ) -> tuple[list[object], list[ExecutionDetails]]:
+        """Execute the generation function on all inputs and return outputs with execution details."""
+        ...
 
 
 class AsyncExecutor(Executor):
-    """
-    A class that provides asynchronous execution of tasks using a producer-consumer pattern.
+    """A class that provides asynchronous execution of tasks using a producer-consumer pattern.
 
     An async interface is provided by the `execute` method, which returns a coroutine, and a sync
     interface is provided by the `run` method.
 
     Args:
-        generation_fn (Callable[[Any], Coroutine[Any, Any, Any]]): A coroutine function that
+        generation_fn (Callable[[object], Coroutine[Any, Any, Any]]): A coroutine function that
             generates tasks to be executed.
 
         concurrency (int, optional): The number of concurrent consumers. Defaults to 3.
@@ -102,14 +112,25 @@ class AsyncExecutor(Executor):
 
     def __init__(
         self,
-        generation_fn: Callable[[Any], Coroutine[Any, Any, Any]],
+        generation_fn: Callable[[object], Coroutine[Any, Any, Any]],
         concurrency: int = 3,
-        tqdm_bar_format: Optional[str] = None,
+        tqdm_bar_format: str | None = None,
         max_retries: int = 10,
         exit_on_error: bool = True,
-        fallback_return_value: Union[Unset, Any] = _unset,
+        fallback_return_value: Unset | object = _unset,
         termination_signal: signal.Signals = signal.SIGINT,
-    ):
+    ) -> None:
+        """Initialize the async executor with configuration parameters.
+
+        Args:
+            generation_fn: Async function to execute for each item.
+            concurrency: Maximum number of concurrent executions.
+            tqdm_bar_format: Custom format string for progress bar.
+            max_retries: Maximum retry attempts for failed executions.
+            exit_on_error: Whether to exit on first error.
+            fallback_return_value: Value to return when execution fails.
+            termination_signal: Signal to handle for graceful termination.
+        """
         self.generate = generation_fn
         self.fallback_return_value = fallback_return_value
         self.concurrency = concurrency
@@ -122,11 +143,12 @@ class AsyncExecutor(Executor):
     async def producer(
         self,
         inputs: Sequence[Any],
-        queue: asyncio.PriorityQueue[Tuple[int, Any]],
+        queue: asyncio.PriorityQueue[tuple[int, Any]],
         max_fill: int,
         done_producing: asyncio.Event,
         termination_signal: asyncio.Event,
     ) -> None:
+        """Produce tasks by adding inputs to the queue for consumers to process."""
         try:
             for index, input in enumerate(inputs):
                 if termination_signal.is_set():
@@ -140,13 +162,14 @@ class AsyncExecutor(Executor):
 
     async def consumer(
         self,
-        outputs: List[Any],
-        execution_details: List[ExecutionDetails],
-        queue: asyncio.PriorityQueue[Tuple[int, Any]],
+        outputs: list[object],
+        execution_details: list[ExecutionDetails],
+        queue: asyncio.PriorityQueue[tuple[int, Any]],
         done_producing: asyncio.Event,
         termination_event: asyncio.Event,
         progress_bar: tqdm[Any],
     ) -> None:
+        """Consume tasks from the queue and execute the generation function on each input."""
         termination_event_watcher = None
         while True:
             marked_done = False
@@ -170,7 +193,7 @@ class AsyncExecutor(Executor):
                 termination_event_watcher = asyncio.create_task(
                     termination_event.wait()
                 )
-                done, pending = await asyncio.wait(
+                done, _pending = await asyncio.wait(
                     [generate_task, termination_event_watcher],
                     timeout=120,
                     return_when=asyncio.FIRST_COMPLETED,
@@ -205,7 +228,7 @@ class AsyncExecutor(Executor):
                     retry_count := abs(priority)
                 ) < self.max_retries and not is_arize_exception:
                     tqdm.write(
-                        f"Exception in worker on attempt {retry_count + 1}: raised {repr(exc)}"
+                        f"Exception in worker on attempt {retry_count + 1}: raised {exc!r}"
                     )
                     tqdm.write("Requeuing...")
                     await queue.put((priority - 1, item))
@@ -229,10 +252,11 @@ class AsyncExecutor(Executor):
 
     async def execute(
         self, inputs: Sequence[Any]
-    ) -> Tuple[List[Any], List[ExecutionDetails]]:
+    ) -> tuple[list[object], list[ExecutionDetails]]:
+        """Execute all inputs asynchronously using producer-consumer pattern."""
         termination_event = asyncio.Event()
 
-        def termination_handler(signum: int, frame: Any) -> None:
+        def termination_handler(signum: int, frame: object) -> None:
             termination_event.set()
             tqdm.write(
                 "Process was interrupted. The return value will be incomplete..."
@@ -251,7 +275,7 @@ class AsyncExecutor(Executor):
         max_fill = max_queue_size - (
             2 * self.concurrency
         )  # ensure there is always room to requeue
-        queue: asyncio.PriorityQueue[Tuple[int, Any]] = asyncio.PriorityQueue(
+        queue: asyncio.PriorityQueue[tuple[int, Any]] = asyncio.PriorityQueue(
             maxsize=max_queue_size
         )
         done_producing = asyncio.Event()
@@ -280,7 +304,7 @@ class AsyncExecutor(Executor):
         termination_event_watcher = asyncio.create_task(
             termination_event.wait()
         )
-        done, pending = await asyncio.wait(
+        done, _pending = await asyncio.wait(
             [join_task, termination_event_watcher],
             return_when=asyncio.FIRST_COMPLETED,
         )
@@ -305,16 +329,16 @@ class AsyncExecutor(Executor):
 
     def run(
         self, inputs: Sequence[Any]
-    ) -> Tuple[List[Any], List[ExecutionDetails]]:
+    ) -> tuple[list[object], list[ExecutionDetails]]:
+        """Execute all inputs asynchronously and return outputs with execution details."""
         return asyncio.run(self.execute(inputs))
 
 
 class SyncExecutor(Executor):
-    """
-    Synchronous executor for generating outputs from inputs using a given generation function.
+    """Synchronous executor for generating outputs from inputs using a given generation function.
 
     Args:
-        generation_fn (Callable[[Any], Any]): The generation function that takes an input and
+        generation_fn (Callable[[object], Any]): The generation function that takes an input and
             returns an output.
 
         tqdm_bar_format (Optional[str], optional): The format string for the progress bar. Defaults
@@ -333,13 +357,23 @@ class SyncExecutor(Executor):
 
     def __init__(
         self,
-        generation_fn: Callable[[Any], Any],
-        tqdm_bar_format: Optional[str] = None,
+        generation_fn: Callable[[object], Any],
+        tqdm_bar_format: str | None = None,
         max_retries: int = 10,
         exit_on_error: bool = True,
-        fallback_return_value: Union[Unset, Any] = _unset,
-        termination_signal: Optional[signal.Signals] = signal.SIGINT,
-    ):
+        fallback_return_value: Unset | object = _unset,
+        termination_signal: signal.Signals | None = signal.SIGINT,
+    ) -> None:
+        """Initialize the sync executor with configuration parameters.
+
+        Args:
+            generation_fn: Synchronous function to execute for each item.
+            tqdm_bar_format: Custom format string for progress bar.
+            max_retries: Maximum retry attempts for failed executions.
+            exit_on_error: Whether to exit on first error.
+            fallback_return_value: Value to return when execution fails.
+            termination_signal: Signal to handle for graceful termination.
+        """
         self.generate = generation_fn
         self.fallback_return_value = fallback_return_value
         self.tqdm_bar_format = tqdm_bar_format
@@ -349,7 +383,7 @@ class SyncExecutor(Executor):
 
         self._TERMINATE = False
 
-    def _signal_handler(self, signum: int, frame: Any) -> None:
+    def _signal_handler(self, signum: int, frame: object) -> None:
         tqdm.write(
             "Process was interrupted. The return value will be incomplete..."
         )
@@ -357,7 +391,7 @@ class SyncExecutor(Executor):
 
     @contextmanager
     def _executor_signal_handling(
-        self, signum: Optional[int]
+        self, signum: int | None
     ) -> Generator[None, None, None]:
         original_handler = None
         if signum is not None:
@@ -369,10 +403,11 @@ class SyncExecutor(Executor):
         else:
             yield
 
-    def run(self, inputs: Sequence[Any]) -> Tuple[List[Any], List[Any]]:
+    def run(self, inputs: Sequence[Any]) -> tuple[list[object], list[object]]:
+        """Execute all inputs synchronously and return outputs with execution details."""
         with self._executor_signal_handling(self.termination_signal):
             outputs = [self.fallback_return_value] * len(inputs)
-            execution_details: List[ExecutionDetails] = [
+            execution_details: list[ExecutionDetails] = [
                 ExecutionDetails() for _ in range(len(inputs))
             ]
             progress_bar = tqdm(
@@ -398,12 +433,11 @@ class SyncExecutor(Executor):
                                 attempt >= self.max_retries
                                 or is_arize_exception
                             ):
-                                raise exc
-                            else:
-                                tqdm.write(
-                                    f"Exception in worker on attempt {attempt + 1}: {exc}"
-                                )
-                                tqdm.write("Retrying...")
+                                raise
+                            tqdm.write(
+                                f"Exception in worker on attempt {attempt + 1}: {exc}"
+                            )
+                            tqdm.write("Retrying...")
                 except Exception as exc:
                     execution_details[index].fail()
                     tqdm.write(
@@ -411,23 +445,40 @@ class SyncExecutor(Executor):
                     )
                     if self.exit_on_error:
                         return outputs, execution_details
-                    else:
-                        progress_bar.update()
+                    progress_bar.update()
                 finally:
                     execution_details[index].log_runtime(task_start_time)
         return outputs, execution_details
 
 
 def get_executor_on_sync_context(
-    sync_fn: Callable[[Any], Any],
-    async_fn: Callable[[Any], Coroutine[Any, Any, Any]],
+    sync_fn: Callable[[object], Any],
+    async_fn: Callable[[object], Coroutine[Any, Any, Any]],
     run_sync: bool = False,
     concurrency: int = 3,
-    tqdm_bar_format: Optional[str] = None,
+    tqdm_bar_format: str | None = None,
     max_retries: int = 10,
     exit_on_error: bool = True,
-    fallback_return_value: Union[Unset, Any] = _unset,
+    fallback_return_value: Unset | object = _unset,
 ) -> Executor:
+    """Get an appropriate executor based on the current threading context.
+
+    Automatically selects between sync and async execution based on thread context.
+    Falls back to sync execution when not in the main thread.
+
+    Args:
+        sync_fn: Synchronous function to execute.
+        async_fn: Asynchronous function to execute.
+        run_sync: Force synchronous execution. Defaults to False.
+        concurrency: Number of concurrent executions for async. Defaults to 3.
+        tqdm_bar_format: Format string for progress bar. Defaults to None.
+        max_retries: Maximum number of retry attempts. Defaults to 10.
+        exit_on_error: Whether to exit on first error. Defaults to True.
+        fallback_return_value: Value to return on failure. Defaults to unset.
+
+    Returns:
+        An Executor instance configured for the current context.
+    """
     if threading.current_thread() is not threading.main_thread():
         # run evals synchronously if not in the main thread
 
@@ -463,33 +514,30 @@ def get_executor_on_sync_context(
                 exit_on_error=exit_on_error,
                 fallback_return_value=fallback_return_value,
             )
-        else:
-            logger.warning(
-                "🐌!! If running inside a notebook, patching the event loop with "
-                "nest_asyncio will allow asynchronous eval submission, and is significantly "
-                "faster. To patch the event loop, run `nest_asyncio.apply()`."
-            )
-            return SyncExecutor(
-                sync_fn,
-                tqdm_bar_format=tqdm_bar_format,
-                max_retries=max_retries,
-                exit_on_error=exit_on_error,
-                fallback_return_value=fallback_return_value,
-            )
-    else:
-        return AsyncExecutor(
-            async_fn,
-            concurrency=concurrency,
+        logger.warning(
+            "🐌!! If running inside a notebook, patching the event loop with "
+            "nest_asyncio will allow asynchronous eval submission, and is significantly "
+            "faster. To patch the event loop, run `nest_asyncio.apply()`."
+        )
+        return SyncExecutor(
+            sync_fn,
             tqdm_bar_format=tqdm_bar_format,
             max_retries=max_retries,
             exit_on_error=exit_on_error,
             fallback_return_value=fallback_return_value,
         )
+    return AsyncExecutor(
+        async_fn,
+        concurrency=concurrency,
+        tqdm_bar_format=tqdm_bar_format,
+        max_retries=max_retries,
+        exit_on_error=exit_on_error,
+        fallback_return_value=fallback_return_value,
+    )
 
 
 def _running_event_loop_exists() -> bool:
-    """
-    Checks for a running event loop.
+    """Checks for a running event loop.
 
     Returns:
         bool: True if a running event loop exists, False otherwise.
@@ -497,6 +545,6 @@ def _running_event_loop_exists() -> bool:
     """
     try:
         asyncio.get_running_loop()
-        return True
     except RuntimeError:
         return False
+    return True