taps.executor.utils

taps/executor/utils.py

FutureDependencyExecutor

FutureDependencyExecutor(executor: Executor)

Bases: Executor

Executor wrapper that adds DAG-like features.

An Executor implementation that wraps another executor with logic for delaying task submission until all Future instances which are args or kwargs of a task have completed. In other words, child tasks will not be scheduled until the results of the child's parent tasks are available.

Parameters:

• executor (Executor) –

  Executor to wrap.

Source code in taps/executor/utils.py

def __init__(self, executor: Executor) -> None:
    self.executor = executor
    self._tasks: dict[Future[Any], _Task[Any]] = {}

submit

submit(
    function: Callable[P, T],
    /,
    *args: args,
    **kwargs: kwargs,
) -> Future[T]

Schedule the callable to be executed.

Parameters:

• function (Callable[P, T]) –

  Callable to execute.

• args (args, default: () ) –

  Positional arguments.

• kwargs (kwargs, default: {} ) –

  Keyword arguments.

Returns:

• Future[T] –

  Future object representing the result of the execution of the callable.

Source code in taps/executor/utils.py

def submit(
    self,
    function: Callable[P, T],
    /,
    *args: P.args,
    **kwargs: P.kwargs,
) -> Future[T]:
    """Schedule the callable to be executed.

    Args:
        function: Callable to execute.
        args: Positional arguments.
        kwargs: Keyword arguments.

    Returns:
        [`Future`][concurrent.futures.Future] object representing the \
        result of the execution of the callable.
    """
    client_future: Future[T] = Future()
    task = _Task(self.executor, function, args, kwargs, client_future)
    self._tasks[client_future] = task
    client_future.add_done_callback(self._task_future_callback)
    return client_future

map

map(
    function: Callable[..., T],
    *iterables: Iterable[Any],
    timeout: float | None = None,
    chunksize: int = 1
) -> Iterator[T]

Map a function onto iterables of arguments.

Parameters:

• function (Callable[..., T]) –

  A callable that will take as many arguments as there are passed iterables.

• iterables (Iterable[Any], default: () ) –

  Variable number of iterables.

• timeout (float | None, default: None ) –

  The maximum number of seconds to wait. If None, then there is no limit on the wait time.

• chunksize (int, default: 1 ) –

  If greater than one, the iterables will be chopped into chunks of size chunksize and submitted to the executor. If set to one, the items in the list will be sent one at a time.

Returns:

• Iterator[T] –

  An iterator equivalent to: map(func, *iterables) but the calls may be evaluated out-of-order.

Raises:

• ValueError –

  if chunksize is less than one.

Source code in taps/executor/utils.py

def map(
    self,
    function: Callable[..., T],
    *iterables: Iterable[Any],
    timeout: float | None = None,
    chunksize: int = 1,
) -> Iterator[T]:
    """Map a function onto iterables of arguments.

    Args:
        function: A callable that will take as many arguments as there are
            passed iterables.
        iterables: Variable number of iterables.
        timeout: The maximum number of seconds to wait. If None, then there
            is no limit on the wait time.
        chunksize: If greater than one, the iterables will be chopped into
            chunks of size chunksize and submitted to the executor. If set
            to one, the items in the list will be sent one at a time.

    Returns:
        An iterator equivalent to: `map(func, *iterables)` but the calls \
        may be evaluated out-of-order.

    Raises:
        ValueError: if chunksize is less than one.
    """
    # Based on concurrent.futures.ProcessPoolExecutor.map()
    # https://github.com/python/cpython/blob/37959e25cbbe1d207c660b5bc9583b9bd1403f1a/Lib/concurrent/futures/process.py
    if chunksize < 1:
        raise ValueError('chunksize must be >= 1.')

    results = super().map(
        functools.partial(_process_chunk, function),
        _get_chunks(*iterables, chunksize=chunksize),
        timeout=timeout,
    )

    def _result_iterator(
        iterable: Iterator[list[T]],
    ) -> Generator[T, None, None]:
        for element in iterable:
            element.reverse()
            while element:
                yield element.pop()

    return _result_iterator(results)

shutdown

shutdown(
    wait: bool = True, *, cancel_futures: bool = False
) -> None

Shutdown the executor.

Parameters:

• wait (bool, default: True ) –

  Wait on all pending futures to complete.

• cancel_futures (bool, default: False ) –

  Cancel all pending futures that the executor has not started running. Only used in Python 3.9 and later.

Source code in taps/executor/utils.py

def shutdown(
    self,
    wait: bool = True,
    *,
    cancel_futures: bool = False,
) -> None:
    """Shutdown the executor.

    Args:
        wait: Wait on all pending futures to complete.
        cancel_futures: Cancel all pending futures that the executor
            has not started running. Only used in Python 3.9 and later.
    """
    if sys.version_info >= (3, 9):  # pragma: >=3.9 cover
        self.executor.shutdown(wait=wait, cancel_futures=cancel_futures)
    else:  # pragma: <3.9 cover
        self.executor.shutdown(wait=wait)

warmup_executor

warmup_executor(
    executor: Executor,
    min_connected_nodes: int,
    batch_size: int,
    max_batches: int,
    batch_sleep: int,
) -> None

Warm up an executor until enough nodes are seen.

Submits a bag of tasks to the executor where each task returns the hostname the task was executed on. Unique hostnames are used to identify active nodes.

Parameters:

• executor (Executor) –

  Executor to warm up.

• min_connected_nodes (int) –

  Number of unique nodes necessary to consider the executor warm.

• batch_size (int) –

  Number of tasks to submit.

• max_batches (int) –

  Number of iterations to try to warm nodes.

• batch_sleep (int) –

  Seconds to sleep between batches.

Raises:

• RuntimeError –

  If min_connected_nodes nodes are not detected within max_batches warmup iterations.

Source code in taps/executor/utils.py

def warmup_executor(
    executor: Executor,
    min_connected_nodes: int,
    batch_size: int,
    max_batches: int,
    batch_sleep: int,
) -> None:
    """Warm up an executor until enough nodes are seen.

    Submits a bag of tasks to the executor where each task returns the
    hostname the task was executed on. Unique hostnames are used to identify
    active nodes.

    Args:
        executor: Executor to warm up.
        min_connected_nodes: Number of unique nodes necessary to consider
            the executor warm.
        batch_size: Number of tasks to submit.
        max_batches: Number of iterations to try to warm nodes.
        batch_sleep: Seconds to sleep between batches.

    Raises:
        RuntimeError: If `min_connected_nodes` nodes are not detected
            within `max_batches` warmup iterations.
    """
    hosts = set()
    for i in range(max_batches):
        futures = [executor.submit(_warmup_task) for _ in range(batch_size)]
        logger.info(
            f'Submitting warmup batch of tasks (batch={i + 1}/{max_batches}, '
            f'tasks={batch_size})',
        )

        for f in futures:
            hosts.add(f.result())

        if len(hosts) >= min_connected_nodes:
            logger.info(f'Executor connected to {len(hosts)} node(s)')
            logger.debug(f'Connected hosts: {hosts}')
            return

        time.sleep(batch_sleep)

    raise RuntimeError(
        f'Could not connect to {min_connected_nodes} nodes within '
        f'{max_batches} warmup batches of {batch_size} tasks.'
        f'Found {len(hosts)} unique nodes.',
    )