Source code for traits_futures.traits_executor

# (C) Copyright 2018-2024 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

"""
Executor to submit background tasks.
"""

import concurrent.futures
import logging
import threading
import warnings

from traits.api import (
    Any,
    Bool,
    Enum,
    HasStrictTraits,
    Instance,
    observe,
    Property,
    Set,
)

from traits_futures.background_call import submit_call
from traits_futures.background_iteration import submit_iteration
from traits_futures.background_progress import submit_progress
from traits_futures.ets_event_loop import ETSEventLoop
from traits_futures.executor_states import (
    ExecutorState,
    RUNNING,
    STOPPED,
    STOPPING,
)
from traits_futures.i_event_loop import IEventLoop
from traits_futures.i_parallel_context import IParallelContext
from traits_futures.multithreading_context import MultithreadingContext
from traits_futures.wrappers import FutureWrapper, run_background_task

logger = logging.getLogger(__name__)


# The TraitsExecutor class maintains an internal state that maps to the
# publicly visible state. The internal state keeps track of some extra
# details about the shutdown.

#: Internal state arising from a timeout on "shutdown": all tasks have been
#: cancelled and the background tasks have been unlinked from their
#: corresponding futures, but some background tasks may still be executing.
#: Maps to the STOPPING public state.
_TERMINATING = "terminating"

#: Mapping from each internal state to the corresponding user-visible state.
_INTERNAL_STATE_TO_EXECUTOR_STATE = {
    RUNNING: RUNNING,
    STOPPING: STOPPING,
    STOPPED: STOPPED,
    _TERMINATING: STOPPING,
}

#: Set of internal states that are considered to be "running" states.
_RUNNING_INTERNAL_STATES = {
    internal_state
    for internal_state, state in _INTERNAL_STATE_TO_EXECUTOR_STATE.items()
    if state == RUNNING
}

#: Set of internal states that are considered to be "stopped" states.
_STOPPED_INTERNAL_STATES = {
    internal_state
    for internal_state, state in _INTERNAL_STATE_TO_EXECUTOR_STATE.items()
    if state == STOPPED
}


class _StateTransitionError(Exception):
    """
    Exception used to indicate a bad state transition.

    Users should never see this exception. It always indicates an error in the
    executor logic.
    """


[docs] class TraitsExecutor(HasStrictTraits): """ Executor to initiate and manage background tasks. Parameters ---------- thread_pool : concurrent.futures.Executor, optional Deprecated alias for worker_pool. .. deprecated:: 0.2 Use ``worker_pool`` instead. worker_pool : concurrent.futures.Executor, optional If supplied, provides the underlying worker pool executor to use. In this case, the creator of the TraitsExecutor is responsible for shutting down the worker pool once it's no longer needed. If not supplied, a new private worker pool will be created, and this object's ``stop`` method will shut down that worker pool. max_workers : int or None, optional Maximum number of workers for the private worker pool. This parameter is mutually exclusive with ``worker_pool``. The default is ``None``, which delegates the choice of number of workers to Python's ``concurrent.futures`` module. context : IParallelContext, optional Parallelism context, providing appropriate concurrent primitives and worker pools for a given choice of parallelism (for example multithreading or multiprocessing). If not given, assumes multithreading. Note that if both ``context`` and ``worker_pool`` are given, they must be compatible. event_loop : IEventLoop, optional The event loop to use for message dispatch. If not given, uses an :class:`~.ETSEventLoop` instance, which determines the appropriate toolkit based on availability and the value of the ETS_TOOLKIT environment variable. """ #: Current state of this executor. state = Property(ExecutorState) #: Derived state: true if this executor is running; False if it's #: stopped or stopping. running = Property(Bool()) #: Derived state: true if this executor is stopped and it's safe #: to dispose of related resources (like the worker pool). stopped = Property(Bool()) def __init__( self, thread_pool=None, *, worker_pool=None, max_workers=None, context=None, event_loop=None, **traits, ): super().__init__(**traits) if thread_pool is not None: warnings.warn( ( "The thread_pool argument to TraitsExecutor is " "deprecated. Use worker_pool instead." ), category=DeprecationWarning, stacklevel=2, ) worker_pool = thread_pool if context is not None: self._context = context if event_loop is not None: self._event_loop = event_loop own_worker_pool = worker_pool is None if own_worker_pool: logger.debug(f"{self} creating worker pool") if max_workers is None: worker_pool = self._context.worker_pool() else: worker_pool = self._context.worker_pool( max_workers=max_workers ) elif max_workers is not None: raise TypeError( "at most one of 'worker_pool' and 'max_workers' " "should be supplied" ) self._worker_pool = worker_pool self._own_worker_pool = own_worker_pool logger.debug(f"{self} running")
[docs] def submit_call(self, callable, *args, **kwargs): """ Convenience function to submit a background call. .. deprecated:: 0.2 Use the :func:`~.submit_call` function instead. Parameters ---------- callable Function to execute in the background. *args Positional arguments to pass to that function. **kwargs Named arguments to pass to that function. Returns ------- future : CallFuture Object representing the state of the background call. """ warnings.warn( "The submit_call method is deprecated. Use the submit_call " "convenience function instead.", DeprecationWarning, stacklevel=2, ) return submit_call(self, callable, *args, **kwargs)
[docs] def submit_iteration(self, callable, *args, **kwargs): """ Convenience function to submit a background iteration. .. deprecated:: 0.2 Use the :func:`~.submit_iteration` function instead. Parameters ---------- callable Function executed in the background to provide the iterable. *args Positional arguments to pass to that function. **kwargs Named arguments to pass to that function. Returns ------- future : IterationFuture Object representing the state of the background iteration. """ warnings.warn( "The submit_iteration method is deprecated. Use the " "submit_iteration convenience function instead.", DeprecationWarning, stacklevel=2, ) return submit_iteration(self, callable, *args, **kwargs)
[docs] def submit_progress(self, callable, *args, **kwargs): """ Convenience function to submit a background progress call. .. deprecated:: 0.2 Use the :func:`~.submit_progress` function instead. Parameters ---------- callable Function executed in the background to provide the iterable. This should accept a "progress" named argument. The callable can then call the "progress" object to report progress. *args Positional arguments to pass to that function. **kwargs Named arguments to pass to that function. These should not include "progress". Returns ------- future : ProgressFuture Object representing the state of the background task. """ warnings.warn( "The submit_progress method is deprecated. Use the " "submit_progress convenience function instead.", DeprecationWarning, stacklevel=2, ) return submit_progress(self, callable, *args, **kwargs)
[docs] def submit(self, task): """ Submit a task to the executor, and return the corresponding future. This method is not thread-safe. It may only be used from the main thread. Parameters ---------- task : ITaskSpecification Returns ------- future : IFuture Future for this task. """ # We may relax this one day, but for now tasks may only be submitted # from the main thread. ref: enthought/traits-futures#302. if threading.current_thread() != threading.main_thread(): raise RuntimeError( "Tasks may only be sumitted on the main thread." ) if not self.running: raise RuntimeError("Can't submit task unless executor is running.") cancel_event = self._context.event() sender, receiver = self._message_router.pipe() runner = task.task() future = task.future(cancel_event.set) self._worker_pool.submit( run_background_task, runner, sender, cancel_event.is_set ) future_wrapper = FutureWrapper( future=future, receiver=receiver, ) self._wrappers.add(future_wrapper) logger.debug(f"{self} created future {future}") return future
[docs] def shutdown(self, *, timeout=None): """ Wait for all tasks to complete and then shut this executor down. All waiting or executing background tasks that are cancellable will be cancelled, and then this executor will wait for all tasks to complete. If a timeout is given and that timeout is reached before all tasks complete, then :exc:`RuntimeError` will be raised and the executor will remain in :data:`~.STOPPING` state. Otherwise, on return from this method the executor will be in :data:`~.STOPPED` state This method may be called at any time. If called on an executor that's already stopped, this method does nothing. Parameters ---------- timeout : float, optional Maximum time to wait for background tasks to complete, in seconds. If not given, this method will wait indefinitely. Raises ------ RuntimeError If a timeout is given, and the background tasks fail to complete within the given timeout. """ if self.stopped: return if self._internal_state == RUNNING: self._initiate_stop() if self._internal_state == STOPPING: self._initiate_manual_stop() assert self._internal_state == _TERMINATING if self._have_message_router: try: self._message_router.route_until( lambda: not self._wrappers, timeout=timeout, ) except RuntimeError as exc: # Re-raise with a more user-friendly error message. raise RuntimeError( "Shutdown timed out; " f"{len(self._wrappers)} tasks still running" ) from exc self._complete_stop()
[docs] def stop(self): """ Initiate stop: cancel existing jobs and prevent new ones. """ if not self.running: raise RuntimeError("Executor is not currently running.") self._initiate_stop() # If there are no tasks pending we can complete the stop immediately; # otherwise, we check as each task completes using the observer below. if not self._wrappers: self._complete_stop()
@observe("_wrappers:items:done") def _finalize_task_and_check_for_stop(self, event): wrapper = event.object self._message_router.close_pipe(wrapper.receiver) self._wrappers.remove(wrapper) logger.debug( f"{self} future {wrapper.future} done ({wrapper.future.state})" ) # If we're in STOPPING state and the last future has just exited, # clean up and stop. if self._internal_state == STOPPING: if not self._wrappers: self._complete_stop() # State transitions ####################################################### def _initiate_stop(self): """ Prevent new tasks from being submitted and cancel existing tasks. Internal state: RUNNING -> STOPPING """ if self._internal_state == RUNNING: self._cancel_tasks() self._internal_state = STOPPING else: raise _StateTransitionError( "Unexpected state transition in internal state {!r}".format( self._internal_state ) ) def _complete_stop(self): """ Move to stopped state when all remaining futures have completed. Internal state: * STOPPING -> STOPPED * _TERMINATING -> STOPPED """ if self._internal_state in {STOPPING, _TERMINATING}: # We should only get here once all futures have completed. assert not self._wrappers self._stop_router() self._close_context() self._shutdown_worker_pool() self._internal_state = STOPPED else: raise _StateTransitionError( "Unexpected state transition in internal state {!r}".format( self._internal_state ) ) def _initiate_manual_stop(self): """ Move into manual stopping mode (_TERMINATING internal state). This differs from the STOPPING internal state in its handling of completed futures: an executor in STOPPING state moves to STOPPED state automatically when the last task completes. An executor in _TERMINATING state must complete the stop manually. Internal state: STOPPING -> _TERMINATING """ if self._internal_state == STOPPING: self._internal_state = _TERMINATING else: raise _StateTransitionError( "Unexpected state transition in internal state {!r}".format( self._internal_state ) ) # Private methods ######################################################### def _cancel_tasks(self): """ Cancel all currently running tasks. """ logger.debug(f"{self} cancelling incomplete tasks") cancel_count = 0 for wrapper in self._wrappers: cancel_count += wrapper.future.cancel() logger.debug(f"{self} cancelled {cancel_count} tasks") def _stop_router(self): """ Stop the message router. """ if self._have_message_router: logger.debug(f"{self} stopping message router") self._message_router.stop() self._message_router = None self._have_message_router = False logger.debug(f"{self} message router stopped") def _close_context(self): """ Close the context, if we own it. """ if self._own_context: logger.debug(f"{self} closing context") self._context.close() logger.debug(f"{self} context closed") self._context = None def _shutdown_worker_pool(self): """ Shut down the worker pool, if we own it. """ if self._own_worker_pool: logger.debug(f"{self} shutting down owned worker pool") # The worker pool shutdown call is potentially blocking, but we # should only ever reach this line when all the background tasks # are complete, so in practice it should never block for long. self._worker_pool.shutdown() logger.debug(f"{self} worker pool is now shut down") self._worker_pool = None # Private traits ########################################################## #: Internal state of the executor. _internal_state = Enum(RUNNING, list(_INTERNAL_STATE_TO_EXECUTOR_STATE)) #: Wrappers for currently-executing futures. _wrappers = Set(Instance(FutureWrapper)) #: Parallelization context _context = Instance(IParallelContext) #: Event loop used for message dispatch _event_loop = Instance(IEventLoop) #: True if we own this context, else False. _own_context = Bool(False) #: concurrent.futures.Executor instance providing the worker pool. _worker_pool = Instance(concurrent.futures.Executor) #: True if we own this worker pool (and are therefore responsible #: for shutting it down), else False. _own_worker_pool = Bool() #: Router providing message connections between background tasks #: and foreground futures. _message_router = Any() #: True if we've created a message router, and need to shut it down. _have_message_router = Bool(False) # Private methods ######################################################### def _get_state(self): """Property getter for the "state" trait.""" return _INTERNAL_STATE_TO_EXECUTOR_STATE[self._internal_state] def _get_running(self): """Property getter for the "running" trait.""" return self._internal_state in _RUNNING_INTERNAL_STATES def _get_stopped(self): """Property getter for the "stopped" trait.""" return self._internal_state in _STOPPED_INTERNAL_STATES @observe("_internal_state") def _update_property_traits(self, event): """Trait change handler for the "_internal_state" trait.""" old_internal_state, new_internal_state = event.old, event.new logger.debug( f"{self} internal state changed " f"from {old_internal_state} to {new_internal_state}" ) old_state = _INTERNAL_STATE_TO_EXECUTOR_STATE[old_internal_state] new_state = _INTERNAL_STATE_TO_EXECUTOR_STATE[new_internal_state] if old_state != new_state: self.trait_property_changed("state", old_state, new_state) old_running = old_internal_state in _RUNNING_INTERNAL_STATES new_running = new_internal_state in _RUNNING_INTERNAL_STATES if old_running != new_running: self.trait_property_changed("running", old_running, new_running) old_stopped = old_internal_state in _STOPPED_INTERNAL_STATES new_stopped = new_internal_state in _STOPPED_INTERNAL_STATES if old_stopped != new_stopped: self.trait_property_changed("stopped", old_stopped, new_stopped) def __message_router_default(self): # Toolkit-specific message router. router = self._context.message_router(event_loop=self._event_loop) router.start() self._have_message_router = True return router def __event_loop_default(self): # By default we use the "ETS" event loop, which chooses which # event loop to use based on the ETS_TOOLKIT environment variable # and the available installed packages. return ETSEventLoop() def __context_default(self): # By default, we use multithreading. context = MultithreadingContext() self._own_context = True return context