DefStream 0.1.2__py3-none-any.whl

DefStream/Worker/Executor.py

@@ -0,0 +1,298 @@
+from concurrent.futures import (
+	Executor as PyExecutor,
+	Future,
+	ProcessPoolExecutor,
+	ThreadPoolExecutor,
+)
+from multiprocessing import get_context
+from multiprocessing.managers import SyncManager
+from pickle import (
+	dumps,
+	loads,
+)
+from queue import (
+	Empty,
+	Queue,
+)
+from signal import (
+	SIGINT,
+	SIG_IGN,
+	signal,
+)
+from threading import (
+	Lock,
+	Thread,
+)
+from traceback import format_exc
+from typing import (
+	Any,
+	Callable,
+	Dict,
+	List,
+	Optional,
+	Sequence,
+	TypeVar,
+)
+
+__all__ = (
+	'Executor',
+)
+
+Result = TypeVar('Result')
+
+
+def initialize_process() -> None:
+	# The parent coordinates draining and cleanup after Ctrl+C.
+	signal(
+		SIGINT,
+		SIG_IGN,
+	)
+
+
+def execute_loop(
+	incoming: Any,
+	outgoing: Any,
+	setups: Sequence[Callable[[], None]],
+	cleanups: Sequence[Callable[[], None]],
+	serialized: bool,
+) -> None:
+	ready = False
+	try:
+		for setup in setups:
+			setup()
+		ready = True
+		while True:
+			item = incoming.get()
+			if item is None:
+				return
+			key, payload = item
+			try:
+				function, args, kwargs = loads(payload) if serialized else payload
+				result = function(
+					*args,
+					**kwargs,
+				)
+			except BaseException as error:
+				outcome = (
+					False,
+					error,
+					format_exc() if serialized else None,
+				)
+			else:
+				outcome = (
+					True,
+					result,
+					None,
+				)
+			if serialized:
+				try:
+					outcome = dumps(outcome)
+				except Exception as error:
+					outcome = dumps((
+						False,
+						RuntimeError('Cannot serialize task outcome: ' + repr(error)),
+						format_exc(),
+					))
+			outgoing.put((
+				key,
+				outcome,
+			))
+	finally:
+		if ready:
+			run_cleanups(cleanups)
+
+
+def run_cleanups(
+	cleanups: Sequence[Callable[..., None]],
+	*args: Any,
+) -> None:
+	failure: Optional[BaseException] = None
+	for cleanup in cleanups:
+		try:
+			cleanup(*args)
+		except BaseException as error:
+			if failure is not None:
+				error.__context__ = failure
+			failure = error
+	if failure is not None:
+		raise failure
+
+
+class Executor(PyExecutor):
+	def __init__(
+		self,
+		size: int,
+		setups: Sequence[Callable[[], None]],
+		cleanups: Sequence[Callable[[], None]],
+		process: bool = False,
+	):
+		self.size = size
+		self.setups = setups
+		self.cleanups = cleanups
+		self.process = process
+		self.lock = Lock()
+		self.pending: Dict[int, Future[Any]] = {}
+		self.workers: List[Future] = []
+		self.closed = False
+		self.failure: Optional[BaseException] = None
+		self.manager: Optional[SyncManager] = None
+		self.executor: Optional[PyExecutor] = None
+		self.monitor: Optional[Thread] = None
+		self.incoming: Any = None
+		self.outgoing: Any = None
+		self.sequence = 0
+
+	def start(self) -> None:
+		try:
+			if self.process:
+				context = get_context('spawn')
+				self.manager = context.Manager()
+				self.incoming = self.manager.Queue()
+				self.outgoing = self.manager.Queue()
+				executor: PyExecutor = ProcessPoolExecutor(
+					self.size,
+					mp_context=context,
+					initializer=initialize_process,
+				)
+			else:
+				self.incoming = Queue()
+				self.outgoing = Queue()
+				executor = ThreadPoolExecutor(self.size)
+			self.executor = executor
+			for index in range(self.size):
+				self.workers.append(executor.submit(
+					execute_loop,
+					self.incoming,
+					self.outgoing,
+					self.setups,
+					self.cleanups,
+					self.process,
+				))
+			monitor = Thread(target=self.collect)
+			monitor.start()
+			self.monitor = monitor
+		except BaseException:
+			self.closed = True
+			try:
+				for worker in self.workers:
+					self.incoming.put(None)
+			finally:
+				try:
+					if self.executor is not None:
+						self.executor.shutdown(wait=True)
+				finally:
+					if self.manager is not None:
+						self.manager.shutdown()
+			raise
+
+	def submit(
+		self,
+		fn: Callable[..., Result],
+		/,
+		*args: Any,
+		**kwargs: Any,
+	) -> Future[Result]:
+		payload = (
+			fn,
+			args,
+			kwargs,
+		)
+		if self.process:
+			payload = dumps(payload)
+		with self.lock:
+			if self.closed:
+				raise RuntimeError('Pool is closed')
+			if self.failure is not None:
+				raise RuntimeError('Pool execution environment failed') from self.failure
+			if self.executor is None:
+				self.start()
+			key = self.sequence
+			self.sequence += 1
+			future: Future[Result] = Future()
+			future.set_running_or_notify_cancel()
+			self.pending[key] = future
+			try:
+				self.incoming.put((
+					key,
+					payload,
+				))
+			except BaseException:
+				del self.pending[key]
+				raise
+			return future
+
+	def collect(self) -> None:
+		try:
+			self.collect_results()
+		except BaseException as error:
+			self.fail(error)
+
+	def fail(
+		self,
+		error: BaseException,
+	) -> None:
+		with self.lock:
+			self.failure = error
+			pending = tuple(self.pending.values())
+			self.pending.clear()
+		for future in pending:
+			future.set_exception(error)
+
+	def collect_results(self) -> None:
+		while True:
+			try:
+				key, outcome = self.outgoing.get(timeout=0.05)
+			except Empty:
+				pass
+			else:
+				success, value, traceback = loads(outcome) if self.process else outcome
+				with self.lock:
+					future = self.pending.pop(
+						key,
+						None,
+					)
+				if future is not None:
+					if success:
+						future.set_result(value)
+					else:
+						if traceback is not None:
+							value.__cause__ = RuntimeError('Remote execution traceback:\n' + traceback)
+						future.set_exception(value)
+			failure = next(
+				(
+					worker.exception()
+					for worker in self.workers
+					if worker.done() and worker.exception() is not None
+				),
+				None,
+			)
+			if failure is not None:
+				self.fail(failure)
+			with self.lock:
+				finished = not self.pending and all(worker.done() for worker in self.workers)
+			if finished:
+				return
+
+	def shutdown(
+		self,
+		wait: bool = True,
+		*,
+		cancel_futures: bool = False,
+	) -> None:
+		with self.lock:
+			if not self.closed:
+				self.closed = True
+				if self.executor is not None:
+					for worker in self.workers:
+						self.incoming.put(None)
+			executor = self.executor
+		if executor is None:
+			return
+		executor.shutdown(wait=wait)
+		if wait:
+			if self.monitor is not None:
+				self.monitor.join()
+			if self.manager is not None:
+				self.manager.shutdown()
+			if self.failure is not None:
+				raise RuntimeError('Pool execution environment failed') from self.failure

DefStream/Worker/Message.py

@@ -0,0 +1,103 @@
+from collections.abc import Mapping
+from inspect import isabstract
+from typing import (
+	Any,
+	Type,
+)
+
+from ..Broker.Connection import (
+	Declarable,
+	TaskSender,
+	Requestable,
+)
+from .Run import Run
+
+__all__ = (
+	'declare',
+	'send',
+	'request',
+)
+
+
+def task_parameters(
+	task: Type[Run],
+) -> Any:
+	if not isinstance(
+		task,
+		type,
+	) or not issubclass(
+		task,
+		Run,
+	) or isabstract(task):
+		raise TypeError('A concrete Run class is required')
+	if task.parameters is None:
+		raise ValueError('Run subscription parameters are required')
+	return task.parameters
+
+
+def message_body(
+	body: Mapping,
+) -> dict:
+	if not isinstance(
+		body,
+		Mapping,
+	) or not all(isinstance(
+		key,
+		str,
+	) for key in body):
+		raise TypeError('Task body must be a mapping with string keys')
+	return dict(body)
+
+
+def declare(
+	*,
+	task: Type[Run],
+	broker: Declarable,
+) -> None:
+	parameters = task_parameters(task)
+	if not isinstance(
+		broker,
+		Declarable,
+	):
+		raise NotImplementedError('This connection does not support task declarations')
+	broker.declare(parameters)
+
+
+def send(
+	body: Mapping,
+	*,
+	task: Type[Run],
+	broker: TaskSender,
+) -> Any:
+	parameters = task_parameters(task)
+	payload = message_body(body)
+	if not isinstance(
+		broker,
+		TaskSender,
+	):
+		raise NotImplementedError('This connection does not support task sending')
+	return broker.send(
+		payload,
+		parameters,
+	)
+
+
+def request(
+	body: Mapping,
+	*,
+	task: Type[Run],
+	broker: TaskSender,
+	timeout: float = 5,
+) -> Any:
+	parameters = task_parameters(task)
+	payload = message_body(body)
+	if not isinstance(
+		broker,
+		Requestable,
+	):
+		raise NotImplementedError('This connection does not support task RPC requests')
+	return broker.request(
+		payload,
+		parameters,
+		timeout=timeout,
+	)

DefStream/Worker/Pool.py

@@ -0,0 +1,203 @@
+from abc import (
+	ABC,
+	abstractmethod,
+)
+from concurrent.futures import (
+	Executor as PyExecutor,
+	Future,
+	wait,
+)
+from os import cpu_count
+from threading import (
+	BoundedSemaphore,
+	Lock,
+)
+from typing import (
+	Any,
+	Callable,
+	Dict,
+	Optional,
+	Sequence,
+	Set,
+	Tuple,
+	Type,
+	Union,
+)
+
+from .Run import Run
+from .Executor import Executor
+
+__all__ = (
+	'Pool',
+	'PoolFull',
+	'PoolSetup',
+	'PoolCleanup',
+	'ThreadPool',
+	'ProcessPool',
+)
+
+
+class PoolFull(RuntimeError):
+	pass
+
+
+class Pool:
+	def __init__(
+		self,
+		executor: PyExecutor,
+		max: int,
+	):
+		self.pool = executor
+		self.max = max
+		self.slots = BoundedSemaphore(max)
+		self.lock = Lock()
+		self.futures: Set[Future] = set()
+		self.closed = False
+
+	def submit(
+		self,
+		run: Type[Run],
+		body: Dict[str, Any],
+		callback: Optional[Callable[[Any], None]] = None,
+		error_callback: Optional[Callable[[BaseException], None]] = None,
+	) -> Future:
+		with self.lock:
+			if self.closed:
+				raise RuntimeError('Pool is closed')
+			if not self.slots.acquire(blocking=False):
+				raise PoolFull('Pool capacity is exhausted')
+			try:
+				future = self.pool.submit(
+					run(),
+					**body,
+				)
+			except BaseException:
+				self.slots.release()
+				raise
+			self.futures.add(future)
+		future.add_done_callback(lambda completed: self.complete(
+			completed,
+			callback,
+			error_callback,
+		))
+		return future
+
+	def complete(
+		self,
+		future: Future,
+		callback: Optional[Callable[[Any], None]] = None,
+		error_callback: Optional[Callable[[BaseException], None]] = None,
+	) -> None:
+		with self.lock:
+			self.futures.remove(future)
+			self.slots.release()
+		try:
+			result = future.result()
+		except BaseException as error:
+			if error_callback is not None:
+				error_callback(error)
+		else:
+			if callback is not None:
+				callback(result)
+
+	def waits(
+		self,
+		timeout: Optional[float] = None,
+	) -> None:
+		with self.lock:
+			tasks = tuple(self.futures)
+		if tasks and wait(
+			tasks,
+			timeout=timeout,
+		).not_done:
+			raise TimeoutError('Pool tasks did not complete in time')
+
+	def stop(self) -> None:
+		with self.lock:
+			self.closed = True
+		self.pool.shutdown(wait=True)
+
+	def shutdown(self) -> None:
+		self.stop()
+
+	def size(self) -> int:
+		return self.max
+
+	def count(self) -> int:
+		with self.lock:
+			return len(self.futures)
+
+
+class PoolSetup(ABC):
+	@abstractmethod
+	def __call__(self) -> None:
+		raise NotImplementedError
+
+
+class PoolCleanup(ABC):
+	@abstractmethod
+	def __call__(self) -> None:
+		raise NotImplementedError
+
+
+def pool_hooks(
+	hooks: Optional[Union[Callable[[], None], Sequence[Callable[[], None]]]],
+) -> Tuple[Callable[[], None], ...]:
+	result = () if hooks is None else tuple(hooks) if isinstance(
+		hooks,
+		Sequence,
+	) else (hooks,)
+	if not all(callable(hook) for hook in result):
+		raise TypeError('Pool hooks must be callable')
+	return result
+
+
+def pool_size(max: Optional[int]) -> int:
+	if max is None:
+		return cpu_count() or 1
+	if isinstance(
+		max,
+		bool,
+	) or not isinstance(
+		max,
+		int,
+	) or max < 1:
+		raise ValueError('Pool size must be a positive integer')
+	return max
+
+
+class ThreadPool(Pool):
+	def __init__(
+		self,
+		max: Optional[int] = None,
+		setups: Optional[Union[PoolSetup, Sequence[PoolSetup]]] = None,
+		cleanups: Optional[Union[PoolCleanup, Sequence[PoolCleanup]]] = None,
+	):
+		size = pool_size(max)
+		super().__init__(
+			Executor(
+				size,
+				pool_hooks(setups),
+				pool_hooks(cleanups),
+			),
+			size,
+		)
+
+
+class ProcessPool(Pool):
+	def __init__(
+		self,
+		max: Optional[int] = None,
+		setups: Optional[Union[PoolSetup, Sequence[PoolSetup]]] = None,
+		cleanups: Optional[Union[PoolCleanup, Sequence[PoolCleanup]]] = None,
+	):
+		size = pool_size(max)
+		super().__init__(
+			Executor(
+				size,
+				pool_hooks(setups),
+				pool_hooks(cleanups),
+				process=True,
+			),
+			size,
+		)