Source code for wolframclient.evaluation.pool

# -*- coding: utf-8 -*-

from __future__ import absolute_import, print_function, unicode_literals

import itertools
import logging
from asyncio import CancelledError

from wolframclient.evaluation.base import WolframAsyncEvaluator
from wolframclient.evaluation.kernel.asyncsession import WolframLanguageAsyncSession
from wolframclient.exception import WolframKernelException
from wolframclient.utils import six
from wolframclient.utils.api import asyncio
from wolframclient.utils.functional import is_iterable

logger = logging.getLogger(__name__)

__all__ = ["WolframEvaluatorPool", "parallel_evaluate"]


[docs]class WolframEvaluatorPool(WolframAsyncEvaluator): """ A pool of kernels to dispatch one-shot evaluations asynchronously. Evaluators can be specified in various ways: as a string representing the path to a local kernel, a :class:`~wolframclient.evaluation.WolframCloudAsyncSession` or an instance of a :class:`~wolframclient.evaluation.WolframLanguageAsyncSession`. More than one evaluator specification can be provided in the form of an iterable object, yielding the abovementioned specification. If the number of evaluators is less than the requested pool size (`poolsize`), elements are duplicated until the requested number of evaluators is reached. Create a pool from a Wolfram kernel default location:: async with WolframEvaluatorPool() as pool: await pool.evaluate('$InstallationDirectory') Create a pool from a specific Wolfram kernel:: async with WolframEvaluatorPool('/path/to/local/kernel') as pool: await pool.evaluate('1+1') Create a pool from a cloud evaluator:: cloud_session = WolframCloudAsyncSession(credentials=myCredentials) async with WolframEvaluatorPool(cloud_session) as pool: await pool.evaluate('$MachineName') Create a pool from a list of specifications:: evaluators = [ WolframCloudAsyncSession(credentials=myCredentials), '/path/to/local/kernel' ] async with WolframEvaluatorPool(evaluators) as pool: await pool.evaluate('$MachineName') Set `poolsize` to the number of kernel instances. The requested size may not be reached due to licencing restrictions. Set `load_factor` to specify how many workloads are queued per kernel before a new evaluation becomes a blocking operation. Values below or equal to 0 mean an infinite queue size. Set `loop` to the event loop to use. `kwargs` are passed to :class:`~wolframclient.evaluation.WolframLanguageAsyncSession` during initialization. """ def __init__( self, async_evaluators=None, poolsize=4, load_factor=0, loop=None, async_language_session_class=WolframLanguageAsyncSession, **kwargs ): super().__init__(loop) if poolsize <= 0: raise ValueError("Invalid pool size value %i. Expecting a positive integer." % i) self._queue = asyncio.Queue(load_factor * poolsize, loop=self._loop) self.async_language_session_class = async_language_session_class self._evaluators = set() if async_evaluators is None or isinstance(async_evaluators, six.string_types): for _ in range(poolsize): self._add_evaluator(async_evaluators, **kwargs) else: if not is_iterable(async_evaluators): async_evaluators = itertools.repeat(async_evaluators) for evaluator in itertools.cycle(async_evaluators): if len(self._evaluators) >= poolsize: break self._add_evaluator(evaluator) self._started_tasks = [] self._pending_init_tasks = () self.last = 0 self.eval_count = 0 self.requestedsize = poolsize def _add_evaluator(self, evaluator, **kwargs): if evaluator is None or isinstance(evaluator, six.string_types): self._evaluators.add( self.async_language_session_class(kernel=evaluator, loop=self._loop, **kwargs) ) elif isinstance(evaluator, WolframAsyncEvaluator): if evaluator in self._evaluators: self._evaluators.add(evaluator.duplicate()) else: self._evaluators.add(evaluator) else: raise ValueError( "Invalid asynchronous evaluator specifications. %s is neither a string nor a WolframAsyncEvaluator instance." % evaluator ) async def _kernel_loop(self, kernel): while True: try: future = None task = None logger.debug("Wait for a new queue entry.") task = await self._queue.get() if task is None: logger.info("Termination requested for kernel: %s." % kernel) break # func is one of the evaluate* methods from WolframAsyncEvaluator. future, func, args, kwargs = task # those method can't be cancelled since the kernel is evaluating anyway. try: func = getattr(kernel, func) result = await asyncio.shield(func(*args, **kwargs)) future.set_result(result) except Exception as e: future.set_exception(e) # First exceptions are those we can't recover from. except KeyboardInterrupt as interrupt: logger.error("Loop associated to kernel %s interrupted by user.", kernel) raise interrupt except CancelledError as cancel: logger.warning("Loop associated to kernel %s cancelled.", kernel) raise cancel except RuntimeError as runtime: logger.error("Unexpected runtime error: {}", runtime) raise runtime except Exception as e: if future: logger.warning( "Exception raised in loop returned in future object. Exception was: %s" % e ) future.set_exception(e) else: logger.warning("No future object. Exception raised in loop was: %s" % e) raise e finally: if task: self._queue.task_done() async def _async_start_kernel(self, kernel): kernel_started = False try: # start the kernel await kernel.start() kernel_started = True except asyncio.CancelledError: logger.info("Cancelled signal during kernel start.") except Exception as e: try: if logger.isEnabledFor(logging.INFO): logger.info( "A kernel from pool failed to start: %s. Reason is %s", kernel, e ) await kernel.stop() except asyncio.CancelledError: logger.info("Cancelled signal.") except Exception as e2: logger.info("Exception raised during clean-up after failed start: %s", e2) if kernel_started: # schedule the infinite evaluation loop task = asyncio.create_task(self._kernel_loop(kernel)) if logger.isEnabledFor(logging.INFO): logger.info("New kernel started in pool: %s.", kernel) # register the task. The loop is not always started at this point. self._started_tasks.append(task) @property def started(self): return len(self._started_tasks) > 0
[docs] async def start(self): """ Start a pool of kernels and wait for at least one of them to be ready for evaluation. This method is a coroutine. If not all the kernels were able to start, it fails and terminates the pool. """ self.stopped = False # keep track of the init tasks. We have to wait before terminating. self._pending_init_tasks = { (asyncio.ensure_future(self._async_start_kernel(kernel), loop=self._loop)) for kernel in self._evaluators } # uninitialized kernels are removed if they failed to start # if they do start the task (the loop) is added to _started_tasks. # we need at least one working kernel. # we also need to keep track of start kernel tasks in case of early termination. while len(self._started_tasks) == 0: if len(self._pending_init_tasks) == 0: raise WolframKernelException("Failed to start any kernel.") _, self._pending_init_tasks = await asyncio.wait( self._pending_init_tasks, return_when=asyncio.FIRST_COMPLETED ) logger.info("Pool initialized with %i running kernels", len(self._started_tasks))
[docs] async def stop(self): if self.stopped: return self.stopped = True # make sure all init tasks are finished. if len(self._pending_init_tasks) > 0: for task in self._pending_init_tasks: task.cancel() await asyncio.wait(self._pending_init_tasks) if len(self._started_tasks) > 0: try: # request for loop termination. for _ in range(len(self._started_tasks)): await self._queue.put(None) # wait for loop to finish before terminating the kernels await asyncio.wait(self._started_tasks, loop=self._loop) except CancelledError: pass except Exception as e: logger.warning("Exception raised while terminating loop: %s", e) # terminate the kernel instances, if any started. tasks = {asyncio.create_task(kernel.stop()) for kernel in self._evaluators} # `wait` raises the first exception, but wait for all tasks to finish. await asyncio.wait(tasks, loop=self._loop)
[docs] async def terminate(self): await self.stop()
[docs] async def ensure_started(self): if not self.started: await self.start() if self.stopped: await self.restart()
async def _put_evaluation_task(self, future, func, expr, **kwargs): await self.ensure_started() await self._queue.put((future, func, (expr,), kwargs)) self.eval_count += 1
[docs] async def evaluate(self, expr, **kwargs): future = asyncio.Future(loop=self._loop) await self._put_evaluation_task(future, "evaluate", expr, **kwargs) return await future
[docs] async def evaluate_wxf(self, expr, **kwargs): future = asyncio.Future(loop=self._loop) await self._put_evaluation_task(future, "evaluate_wxf", expr, **kwargs) return await future
[docs] async def evaluate_wrap(self, expr, **kwargs): future = asyncio.Future(loop=self._loop) await self._put_evaluation_task(future, "evaluate_wrap", expr, **kwargs) return await future
[docs] def evaluate_all(self, iterable): return self._loop.run_until_complete(self._evaluate_all(iterable))
async def _evaluate_all(self, iterable): tasks = [asyncio.create_task(self.evaluate(expr)) for expr in iterable] return await asyncio.gather(*tasks) def __repr__(self): return "<%s %i/%i started evaluators, cumulating %i evaluations>" % ( self.__class__.__name__, len(self._started_tasks), self.requestedsize, self.eval_count, ) def __len__(self): return len(self._started_tasks)
[docs]def parallel_evaluate(expressions, evaluator_spec=None, max_evaluators=4, loop=None): """ Start a kernel pool and evaluate the expressions in parallel. The pool is created with the value of `evaluator_spec`. The pool is automatically stopped when it is no longer needed. The expressions are evaluated and returned in order. Note that each evaluation should be independent and not rely on any previous one. There is no guarantee that two given expressions evaluate on the same kernel. """ loop = loop or asyncio.get_event_loop() pool = None try: pool = WolframEvaluatorPool(evaluator_spec, poolsize=max_evaluators, loop=loop) loop.run_until_complete(pool.start()) return pool.evaluate_all(expressions) finally: if pool: loop.run_until_complete(pool.terminate())