Shortcuts

Source code for ding.framework.parallel

import atexit
import os
import random
import time
import traceback
import pickle
from mpire.pool import WorkerPool
from ditk import logging
import tempfile
import socket
from os import path
from typing import Callable, Dict, List, Optional, Tuple, Union, Set
from threading import Thread
from ding.framework.event_loop import EventLoop
from ding.utils.design_helper import SingletonMetaclass
from ding.framework.message_queue import *
from ding.utils.registry_factory import MQ_REGISTRY

# Avoid ipc address conflict, random should always use random seed
random = random.Random()


[docs]class Parallel(metaclass=SingletonMetaclass): def __init__(self) -> None: # Init will only be called once in a process self._listener = None self.is_active = False self.node_id = None self.local_id = None self.labels = set() self._event_loop = EventLoop("parallel_{}".format(id(self))) self._retries = 0 # Retries in auto recovery def _run( self, node_id: int, local_id: int, n_parallel_workers: int, labels: Optional[Set[str]] = None, auto_recover: bool = False, max_retries: int = float("inf"), mq_type: str = "nng", startup_interval: int = 1, **kwargs ) -> None: self.node_id = node_id self.local_id = local_id self.startup_interval = startup_interval self.n_parallel_workers = n_parallel_workers self.labels = labels or set() self.auto_recover = auto_recover self.max_retries = max_retries self._mq = MQ_REGISTRY.get(mq_type)(**kwargs) time.sleep(self.local_id * self.startup_interval) self._listener = Thread(target=self.listen, name="mq_listener", daemon=True) self._listener.start() self.mq_type = mq_type self.barrier_runtime = Parallel.get_barrier_runtime()(self.node_id)
[docs] @classmethod def runner( cls, n_parallel_workers: int, mq_type: str = "nng", attach_to: Optional[List[str]] = None, protocol: str = "ipc", address: Optional[str] = None, ports: Optional[Union[List[int], int]] = None, topology: str = "mesh", labels: Optional[Set[str]] = None, node_ids: Optional[Union[List[int], int]] = None, auto_recover: bool = False, max_retries: int = float("inf"), redis_host: Optional[str] = None, redis_port: Optional[int] = None, startup_interval: int = 1 ) -> Callable: """ Overview: This method allows you to configure parallel parameters, and now you are still in the parent process. Arguments: - n_parallel_workers (:obj:`int`): Workers to spawn. - mq_type (:obj:`str`): Embedded message queue type, i.e. nng, redis. - attach_to (:obj:`Optional[List[str]]`): The node's addresses you want to attach to. - protocol (:obj:`str`): Network protocol. - address (:obj:`Optional[str]`): Bind address, ip or file path. - ports (:obj:`Optional[List[int]]`): Candidate ports. - topology (:obj:`str`): Network topology, includes: `mesh` (default): fully connected between each other; `star`: only connect to the first node; `alone`: do not connect to any node, except the node attached to; - labels (:obj:`Optional[Set[str]]`): Labels. - node_ids (:obj:`Optional[List[int]]`): Candidate node ids. - auto_recover (:obj:`bool`): Auto recover from uncaught exceptions from main. - max_retries (:obj:`int`): Max retries for auto recover. - redis_host (:obj:`str`): Redis server host. - redis_port (:obj:`int`): Redis server port. - startup_interval (:obj:`int`): Start up interval between each task. Returns: - _runner (:obj:`Callable`): The wrapper function for main. """ all_args = locals() del all_args["cls"] args_parsers = {"nng": cls._nng_args_parser, "redis": cls._redis_args_parser} assert n_parallel_workers > 0, "Parallel worker number should bigger than 0" def _runner(main_process: Callable, *args, **kwargs) -> None: """ Overview: Prepare to run in subprocess. Arguments: - main_process (:obj:`Callable`): The main function, your program start from here. """ runner_params = args_parsers[mq_type](**all_args) params_group = [] for i, runner_kwargs in enumerate(runner_params): runner_kwargs["local_id"] = i params_group.append([runner_kwargs, (main_process, args, kwargs)]) if n_parallel_workers == 1: cls._subprocess_runner(*params_group[0]) else: with WorkerPool(n_jobs=n_parallel_workers, start_method="spawn", daemon=False) as pool: # Cleanup the pool just in case the program crashes. atexit.register(pool.__exit__) pool.map(cls._subprocess_runner, params_group) return _runner
@classmethod def _nng_args_parser( cls, n_parallel_workers: int, attach_to: Optional[List[str]] = None, protocol: str = "ipc", address: Optional[str] = None, ports: Optional[Union[List[int], int]] = None, topology: str = "mesh", node_ids: Optional[Union[List[int], int]] = None, **kwargs ) -> Dict[str, dict]: attach_to = attach_to or [] nodes = cls.get_node_addrs(n_parallel_workers, protocol=protocol, address=address, ports=ports) def cleanup_nodes(): for node in nodes: protocol, file_path = node.split("://") if protocol == "ipc" and path.exists(file_path): os.remove(file_path) atexit.register(cleanup_nodes) def topology_network(i: int) -> List[str]: if topology == "mesh": return nodes[:i] + attach_to elif topology == "star": return nodes[:min(1, i)] + attach_to elif topology == "alone": return attach_to else: raise ValueError("Unknown topology: {}".format(topology)) runner_params = [] candidate_node_ids = cls.padding_param(node_ids, n_parallel_workers, 0) for i in range(n_parallel_workers): runner_kwargs = { **kwargs, "node_id": candidate_node_ids[i], "listen_to": nodes[i], "attach_to": topology_network(i), "n_parallel_workers": n_parallel_workers, } runner_params.append(runner_kwargs) return runner_params @classmethod def _redis_args_parser(cls, n_parallel_workers: int, node_ids: Optional[Union[List[int], int]] = None, **kwargs): runner_params = [] candidate_node_ids = cls.padding_param(node_ids, n_parallel_workers, 0) for i in range(n_parallel_workers): runner_kwargs = {**kwargs, "n_parallel_workers": n_parallel_workers, "node_id": candidate_node_ids[i]} runner_params.append(runner_kwargs) return runner_params @classmethod def _subprocess_runner(cls, runner_kwargs: dict, main_params: Tuple[Union[List, Dict]]) -> None: """ Overview: Really run in subprocess. Arguments: - runner_params (:obj:`Tuple[Union[List, Dict]]`): Args and kwargs for runner. - main_params (:obj:`Tuple[Union[List, Dict]]`): Args and kwargs for main function. """ logging.getLogger().setLevel(logging.INFO) main_process, args, kwargs = main_params with Parallel() as router: router.is_active = True router._run(**runner_kwargs) time.sleep(0.3) # Waiting for network pairing router._supervised_runner(main_process, *args, **kwargs) def _supervised_runner(self, main: Callable, *args, **kwargs) -> None: """ Overview: Run in supervised mode. Arguments: - main (:obj:`Callable`): Main function. """ if self.auto_recover: while True: try: main(*args, **kwargs) break except Exception as e: if self._retries < self.max_retries: logging.warning( "Auto recover from exception: {}, node: {}, retries: {}".format( e, self.node_id, self._retries ) ) logging.warning(traceback.format_exc()) self._retries += 1 else: logging.warning( "Exceed the max retries, node: {}, retries: {}, max_retries: {}".format( self.node_id, self._retries, self.max_retries ) ) raise e else: main(*args, **kwargs) @classmethod def get_node_addrs( cls, n_workers: int, protocol: str = "ipc", address: Optional[str] = None, ports: Optional[Union[List[int], int]] = None ) -> None: if protocol == "ipc": node_name = "".join(random.choices("abcdefghijklmnopqrstuvwxyz0123456789", k=4)) tmp_dir = tempfile.gettempdir() nodes = ["ipc://{}/ditask_{}_{}.ipc".format(tmp_dir, node_name, i) for i in range(n_workers)] elif protocol == "tcp": address = address or cls.get_ip() ports = cls.padding_param(ports, n_workers, 50515) assert len(ports) == n_workers, "The number of ports must be the same as the number of workers, \ now there are {} ports and {} workers".format(len(ports), n_workers) nodes = ["tcp://{}:{}".format(address, port) for port in ports] else: raise Exception("Unknown protocol {}".format(protocol)) return nodes @classmethod def padding_param(cls, int_or_list: Optional[Union[List[int], int]], n_max: int, start_value: int) -> List[int]: """ Overview: Padding int or list param to the length of n_max. Arguments: - int_or_list (:obj:`Optional[Union[List[int], int]]`): Int or list typed value. - n_max (:obj:`int`): Max length. - start_value (:obj:`int`): Start from value. """ param = int_or_list if isinstance(param, List) and len(param) == 1: param = param[0] # List with only 1 element is equal to int if isinstance(param, int): param = range(param, param + n_max) else: param = param or range(start_value, start_value + n_max) return param def listen(self): self._mq.listen() while True: if not self._mq: break msg = self._mq.recv() # msg is none means that the message queue is no longer being listened to, # especially if the message queue is already closed if not msg: break topic, msg = msg self._handle_message(topic, msg) def on(self, event: str, fn: Callable) -> None: """ Overview: Register an remote event on parallel instance, this function will be executed \ when a remote process emit this event via network. Arguments: - event (:obj:`str`): Event name. - fn (:obj:`Callable`): Function body. """ if self.is_active: self._mq.subscribe(event) self._event_loop.on(event, fn) def once(self, event: str, fn: Callable) -> None: """ Overview: Register an remote event which will only call once on parallel instance, this function will be executed when a remote process emit this event via network. Arguments: - event (:obj:`str`): Event name. - fn (:obj:`Callable`): Function body. """ if self.is_active: self._mq.subscribe(event) self._event_loop.once(event, fn) def off(self, event: str) -> None: """ Overview: Unregister an event. Arguments: - event (:obj:`str`): Event name. """ if self.is_active: self._mq.unsubscribe(event) self._event_loop.off(event) def emit(self, event: str, *args, **kwargs) -> None: """ Overview: Send an remote event via network to subscribed processes. Arguments: - event (:obj:`str`): Event name. """ if self.is_active: payload = {"a": args, "k": kwargs} try: data = pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL) except AttributeError as e: logging.error("Arguments are not pickable! Event: {}, Args: {}".format(event, args)) raise e self._mq.publish(event, data) def _handle_message(self, topic: str, msg: bytes) -> None: """ Overview: Recv and parse payload from other processes, and call local functions. Arguments: - topic (:obj:`str`): Recevied topic. - msg (:obj:`bytes`): Recevied message. """ event = topic if not self._event_loop.listened(event): logging.debug("Event {} was not listened in parallel {}".format(event, self.node_id)) return try: payload = pickle.loads(msg) except Exception as e: logging.error("Error when unpacking message on node {}, msg: {}".format(self.node_id, e)) return self._event_loop.emit(event, *payload["a"], **payload["k"]) @classmethod def get_ip(cls): s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: # doesn't even have to be reachable s.connect(('10.255.255.255', 1)) ip = s.getsockname()[0] except Exception: ip = '127.0.0.1' finally: s.close() return ip def get_attch_to_len(self) -> int: """ Overview: Get the length of the 'attach_to' list of message queue. Returns: int: the length of the self._mq.attach_to. Returns 0 if self._mq is not initialized """ if self._mq: if hasattr(self._mq, 'attach_to'): return len(self._mq.attach_to) return 0 def __enter__(self) -> "Parallel": return self def __exit__(self, exc_type, exc_val, exc_tb): self.stop() def stop(self): logging.info("Stopping parallel worker on node: {}".format(self.node_id)) self.is_active = False time.sleep(0.03) if self._mq: self._mq.stop() self._mq = None if self._listener: self._listener.join(timeout=1) self._listener = None self._event_loop.stop() @classmethod def get_barrier_runtime(cls): # We get the BarrierRuntime object in the closure to avoid circular import. from ding.framework.middleware.barrier import BarrierRuntime return BarrierRuntime