from asyncio import InvalidStateError from asyncio.tasks import FIRST_EXCEPTION from collections import OrderedDict from threading import Lock import time import asyncio import concurrent.futures import fnmatch import math import enum from types import GeneratorType from typing import Any, Awaitable, Callable, Dict, Generator, Iterable, List, Optional, Set, Union import inspect from ding.framework.context import Context from ding.framework.parallel import Parallel from ding.framework.event_loop import EventLoop from functools import wraps def enable_async(func: Callable) -> Callable: """ Overview: Empower the function with async ability. Arguments: - func (:obj:`Callable`): The original function. Returns: - runtime_handler (:obj:`Callable`): The wrap function. """ @wraps(func) def runtime_handler(task: "Task", *args, async_mode: Optional[bool] = None, **kwargs) -> "Task": """ Overview: If task's async mode is enabled, execute the step in current loop executor asyncly, or execute the task sync. Arguments: - task (:obj:`Task`): The task instance. - async_mode (:obj:`Optional[bool]`): Whether using async mode. Returns: - result (:obj:`Union[Any, Awaitable]`): The result or future object of middleware. """ if async_mode is None: async_mode = task.async_mode if async_mode: assert not kwargs, "Should not use kwargs in async_mode, use position parameters, kwargs: {}".format(kwargs) t = task._async_loop.run_in_executor(task._thread_pool, func, task, *args, **kwargs) task._async_stack.append(t) return task else: return func(task, *args, **kwargs) return runtime_handler class Role(str, enum.Enum): LEARNER = "learner" COLLECTOR = "collector" EVALUATOR = "evaluator" FETCHER = 'fetcher' class VoidMiddleware: def __call__(self, _): return class Task: """ Task will manage the execution order of the entire pipeline, register new middleware, and generate new context objects. """ role = Role def __init__(self) -> None: self.router = Parallel() self._finish = False def start( self, async_mode: bool = False, n_async_workers: int = 3, ctx: Optional[Context] = None, labels: Optional[Set[str]] = None ) -> "Task": # This flag can be modified by external or associated processes self._finish = False # This flag can only be modified inside the class, it will be set to False in the end of stop self._running = True self._middleware = [] self._wrappers = [] self.ctx = ctx or Context() self._backward_stack = OrderedDict() self._roles = set() # Bind event loop functions self._event_loop = EventLoop("task_{}".format(id(self))) # Async segment self.async_mode = async_mode self.n_async_workers = n_async_workers self._async_stack = [] self._async_loop = None self._thread_pool = None self._exception = None self._thread_lock = Lock() self.labels = labels or set() # Parallel segment if async_mode or self.router.is_active: self._activate_async() if self.router.is_active: def sync_finish(value): self._finish = value self.on("finish", sync_finish) self.init_labels() return self def add_role(self, role: Role): self._roles.add(role) def has_role(self, role: Role) -> bool: if len(self._roles) == 0: return True return role in self._roles @property def roles(self) -> Set[Role]: return self._roles def void(self): return VoidMiddleware() def init_labels(self): if self.async_mode: self.labels.add("async") if self.router.is_active: self.labels.add("distributed") self.labels.add("node.{}".format(self.router.node_id)) for label in self.router.labels: self.labels.add(label) else: self.labels.add("standalone") def use(self, fn: Callable, lock: Union[bool, Lock] = False) -> 'Task': """ Overview: Register middleware to task. The middleware will be executed by it's registry order. Arguments: - fn (:obj:`Callable`): A middleware is a function with only one argument: ctx. - lock (:obj:`Union[bool, Lock]`): There can only be one middleware execution under lock at any one time. Returns: - task (:obj:`Task`): The task. """ assert isinstance(fn, Callable), "Middleware function should be a callable object, current fn {}".format(fn) if isinstance(fn, VoidMiddleware): # Skip void function return self for wrapper in self._wrappers: fn = wrapper(fn) self._middleware.append(self.wrap(fn, lock=lock)) return self def use_wrapper(self, fn: Callable) -> 'Task': """ Overview: Register wrappers to task. A wrapper works like a decorator, but task will apply this \ decorator on top of each middleware. Arguments: - fn (:obj:`Callable`): A wrapper is a decorator, so the first argument is a callable function. Returns: - task (:obj:`Task`): The task. """ # Wrap exist middlewares for i, middleware in enumerate(self._middleware): self._middleware[i] = fn(middleware) self._wrappers.append(fn) return self def match_labels(self, patterns: Union[Iterable[str], str]) -> bool: """ Overview: A list of patterns to match labels. Arguments: - patterns (:obj:`Union[Iterable[str], str]`): Glob like pattern, e.g. node.1, node.*. """ if isinstance(patterns, str): patterns = [patterns] return any([fnmatch.filter(self.labels, p) for p in patterns]) def run(self, max_step: int = int(1e12)) -> None: """ Overview: Execute the iterations, when reach the max_step or task.finish is true, The loop will be break. Arguments: - max_step (:obj:`int`): Max step of iterations. """ assert self._running, "Please make sure the task is running before calling the this method, see the task.start" if len(self._middleware) == 0: return for i in range(max_step): for fn in self._middleware: self.forward(fn) # Sync should be called before backward, otherwise it is possible # that some generators have not been pushed to backward_stack. self.sync() self.backward() self.sync() if i == max_step - 1: self.finish = True if self.finish: break self.renew() def wrap(self, fn: Callable, lock: Union[bool, Lock] = False) -> Callable: """ Overview: Wrap the middleware, make it can be called directly in other middleware. Arguments: - fn (:obj:`Callable`): The middleware. - lock (:obj:`Union[bool, Lock]`): There can only be one middleware execution under lock at any one time. Returns: - fn_back (:obj:`Callable`): It will return a backward function, which will call the rest part of the middleware after yield. If this backward function was not called, the rest part of the middleware will be called in the global backward step. """ if lock is True: lock = self._thread_lock def forward(ctx: Context): if lock: with lock: g = self.forward(fn, ctx, async_mode=False) else: g = self.forward(fn, ctx, async_mode=False) def backward(): backward_stack = OrderedDict() key = id(g) backward_stack[key] = self._backward_stack.pop(key) if lock: with lock: self.backward(backward_stack, async_mode=False) else: self.backward(backward_stack, async_mode=False) return backward if hasattr(fn, "__name__"): forward = wraps(fn)(forward) else: forward = wraps(fn.__class__)(forward) return forward @enable_async def forward(self, fn: Callable, ctx: Optional[Context] = None) -> Optional[Generator]: """ Overview: This function will execute the middleware until the first yield statment, or the end of the middleware. Arguments: - fn (:obj:`Callable`): Function with contain the ctx argument in middleware. - ctx (:obj:`Optional[Context]`): Replace global ctx with a customized ctx. Returns: - g (:obj:`Optional[Generator]`): The generator if the return value of fn is a generator. """ assert self._running, "Please make sure the task is running before calling the this method, see the task.start" if not ctx: ctx = self.ctx g = fn(ctx) if isinstance(g, GeneratorType): try: next(g) self._backward_stack[id(g)] = g return g except StopIteration: pass @enable_async def backward(self, backward_stack: Optional[Dict[str, Generator]] = None) -> None: """ Overview: Execute the rest part of middleware, by the reversed order of registry. Arguments: - backward_stack (:obj:`Optional[Dict[str, Generator]]`): Replace global backward_stack with a customized \ stack. """ assert self._running, "Please make sure the task is running before calling the this method, see the task.start" if not backward_stack: backward_stack = self._backward_stack while backward_stack: # FILO _, g = backward_stack.popitem() try: next(g) except StopIteration: continue @property def running(self): return self._running def serial(self, *fns: List[Callable]) -> Callable: """ Overview: Wrap functions and keep them run in serial, Usually in order to avoid the confusion of dependencies in async mode. Arguments: - fn (:obj:`Callable`): Chain a serial of middleware, wrap them into one middleware function. """ def _serial(ctx): backward_keys = [] for fn in fns: g = self.forward(fn, ctx, async_mode=False) if isinstance(g, GeneratorType): backward_keys.append(id(g)) yield backward_stack = OrderedDict() for k in backward_keys: backward_stack[k] = self._backward_stack.pop(k) self.backward(backward_stack=backward_stack, async_mode=False) name = ",".join([fn.__name__ for fn in fns]) _serial.__name__ = "serial<{}>".format(name) return _serial def parallel(self, *fns: List[Callable]) -> Callable: """ Overview: Wrap functions and keep them run in parallel, should not use this funciton in async mode. Arguments: - fn (:obj:`Callable`): Parallelized middleware, wrap them into one middleware function. """ self._activate_async() def _parallel(ctx): backward_keys = [] for fn in fns: g = self.forward(fn, ctx, async_mode=True) if isinstance(g, GeneratorType): backward_keys.append(id(g)) self.sync() yield backward_stack = OrderedDict() for k in backward_keys: backward_stack[k] = self._backward_stack.pop(k) self.backward(backward_stack, async_mode=True) self.sync() name = ",".join([fn.__name__ for fn in fns]) _parallel.__name__ = "parallel<{}>".format(name) return _parallel def renew(self) -> 'Task': """ Overview: Renew the context instance, this function should be called after backward in the end of iteration. """ assert self._running, "Please make sure the task is running before calling the this method, see the task.start" self.ctx = self.ctx.renew() return self def __enter__(self) -> "Task": return self def __exit__(self, exc_type, exc_val, exc_tb): self.stop() def stop(self) -> None: """ Overview: Stop and cleanup every thing in the runtime of task. """ if self.router.is_active: self.emit("finish", True) if self._thread_pool: self._thread_pool.shutdown() self._event_loop.stop() self.router.off(self._wrap_event_name("*")) if self._async_loop: self._async_loop.stop() self._async_loop.close() # The middleware and listeners may contain some methods that reference to task, # If we do not clear them after the task exits, we may find that gc will not clean up the task object. self._middleware.clear() self._wrappers.clear() self._backward_stack.clear() self._async_stack.clear() self._running = False def sync(self) -> 'Task': if self._async_loop: self._async_loop.run_until_complete(self.sync_tasks()) return self async def sync_tasks(self) -> Awaitable[None]: if self._async_stack: await asyncio.wait(self._async_stack, return_when=FIRST_EXCEPTION) while self._async_stack: t = self._async_stack.pop(0) try: e = t.exception() if e: self._exception = e raise e except InvalidStateError: # Not finished. https://docs.python.org/3/library/asyncio-task.html#asyncio.Task.exception pass def async_executor(self, fn: Callable, *args, **kwargs) -> None: """ Overview: Execute task in background, then apppend the future instance in _async_stack. Arguments: - fn (:obj:`Callable`): Synchronization fuction. """ if not self._async_loop: raise Exception("Event loop was not initialized, please call this function in async or parallel mode") t = self._async_loop.run_in_executor(self._thread_pool, fn, *args, **kwargs) self._async_stack.append(t) def emit(self, event: str, *args, only_remote: bool = False, only_local: bool = False, **kwargs) -> None: """ Overview: Emit an event, call listeners. Arguments: - event (:obj:`str`): Event name. - only_remote (:obj:`bool`): Only broadcast the event to the connected nodes, default is False. - only_local (:obj:`bool`): Only emit local event, default is False. - args (:obj:`any`): Rest arguments for listeners. """ # Check if need to broadcast event to connected nodes, default is True assert self._running, "Please make sure the task is running before calling the this method, see the task.start" if only_local: self._event_loop.emit(event, *args, **kwargs) elif only_remote: if self.router.is_active: self.async_executor(self.router.emit, self._wrap_event_name(event), event, *args, **kwargs) else: if self.router.is_active: self.async_executor(self.router.emit, self._wrap_event_name(event), event, *args, **kwargs) self._event_loop.emit(event, *args, **kwargs) def on(self, event: str, fn: Callable) -> None: """ Overview: Subscribe to an event, execute this function every time the event is emitted. Arguments: - event (:obj:`str`): Event name. - fn (:obj:`Callable`): The function. """ self._event_loop.on(event, fn) if self.router.is_active: self.router.on(self._wrap_event_name(event), self._event_loop.emit) def once(self, event: str, fn: Callable) -> None: """ Overview: Subscribe to an event, execute this function only once when the event is emitted. Arguments: - event (:obj:`str`): Event name. - fn (:obj:`Callable`): The function. """ self._event_loop.once(event, fn) if self.router.is_active: self.router.on(self._wrap_event_name(event), self._event_loop.emit) def off(self, event: str, fn: Optional[Callable] = None) -> None: """ Overview: Unsubscribe an event Arguments: - event (:obj:`str`): Event name. - fn (:obj:`Callable`): The function. """ self._event_loop.off(event, fn) if self.router.is_active: self.router.off(self._wrap_event_name(event)) def wait_for(self, event: str, timeout: float = math.inf, ignore_timeout_exception: bool = True) -> Any: """ Overview: Wait for an event and block the thread. Arguments: - event (:obj:`str`): Event name. - timeout (:obj:`float`): Timeout in seconds. - ignore_timeout_exception (:obj:`bool`): If this is False, an exception will occur when meeting timeout. """ assert self._running, "Please make sure the task is running before calling the this method, see the task.start" received = False result = None def _receive_event(*args, **kwargs): nonlocal result, received result = (args, kwargs) received = True self.once(event, _receive_event) start = time.time() while time.time() - start < timeout: if received or self._exception: return result time.sleep(0.01) if ignore_timeout_exception: return result else: raise TimeoutError("Timeout when waiting for event: {}".format(event)) @property def finish(self): return self._finish @finish.setter def finish(self, value: bool): self._finish = value def _wrap_event_name(self, event: str) -> str: """ Overview: Wrap the event name sent to the router. Arguments: - event (:obj:`str`): Event name """ return "task.{}".format(event) def _activate_async(self): if not self._thread_pool: self._thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=self.n_async_workers) if not self._async_loop: self._async_loop = asyncio.new_event_loop() def get_attch_to_len(self) -> int: """ Overview: Get the length of the 'attach_to' list in Parallel._mq. Returns: int: the length of the Parallel._mq. """ if self.router.is_active: return self.router.get_attch_to_len() else: raise RuntimeError("The router is inactive, failed to be obtained the length of 'attch_to' list.") task = Task()