`ding.framework.middleware.distributer`¶

`ding.framework.middleware.distributer` ¶

`ContextExchanger` ¶

`init(skip_n_iter=1, storage_loader=None)` ¶

Overview

Exchange context between processes, support properties: trajectories, episodes, env_step, env_episode, train_iter

Arguments: - skip_n_iter (:obj:int): For collectors, it may be necessary to skip waiting for the first n iterations to collect data for the learner to learn. This parameter will not work on learner. - storage_loader (:obj:Optional[StorageLoader]): Turn data into storage class to reduce the network overhead.

`put(payload)` ¶

Overview

Get attributes from ctx on the callback of event. Each attribute should have a standalone put handler, which named _put_{key}

`fetch(ctx)` ¶

Overview

Fetch attributes from ctx before emit them to the event bus. Each attribute should have a standalone fetch handler, which named _fetch_{key}

`ModelExchanger` ¶

`init(model, model_loader=None)` ¶

Overview

Exchange model between processes, only the learner will send the model, otherwise the model will only be received. If you are using a shared model on a single host, there is no need to use this middleware.

Arguments: - model (:obj:torch.nn.Module): Pytorch module. - model_loader (:obj:ModelLoader): Encode model in subprocess.

`PeriodicalModelExchanger` ¶

`init(model, mode, period=1, delay_toleration=np.inf, stale_toleration=1, event_name='model_exchanger', model_loader=None)` ¶

Overview

Exchange model between processes, set the mode to "send" or "receive" to specify the role of the process. If you are using a shared model on a single host, there is no need to use this middleware.

Arguments: - model (:obj:torch.nn.Module): Pytorch module. - mode (:obj:str): "send" or "receive". - period (:obj:int): The period of model exchange. - delay_toleration (:obj:float): The permitted time interval for receiving model after being sent. - stale_toleration (:obj:int): The permitted number of iterations for receiving model after being sent. - event_name (:obj:str): The event name for model exchange. - model_loader (:obj:ModelLoader): ModelLoader for this PeriodicalModelExchanger to use.

Full Source Code

../ding/framework/middleware/distributer.py

import numpy as npfrom time import sleep, timefrom dataclasses import fieldsfrom typing import TYPE_CHECKING, List, Dict, Any, Optional, Unionfrom ditk import loggingfrom ding.framework import taskfrom ding.data import StorageLoader, Storage, ModelLoaderif TYPE_CHECKING:    from ding.framework.context import Context    from torch.nn import Moduleclass ContextExchanger:    def __init__(self, skip_n_iter: int = 1, storage_loader: Optional[StorageLoader] = None) -> None:        """        Overview:            Exchange context between processes,            support properties: trajectories, episodes, env_step, env_episode, train_iter        Arguments:            - skip_n_iter (:obj:`int`): For collectors, it may be necessary to skip waiting \                for the first n iterations to collect data for the learner to learn. This parameter \                will not work on learner.            - storage_loader (:obj:`Optional[StorageLoader]`): Turn data into storage class to reduce \                the network overhead.        """        if not task.router.is_active:            raise RuntimeError("ContextHandler should be used in parallel mode!")        self._state = {}        self._local_state = {}  # just save local state, not send to remote node        if task.has_role(task.role.COLLECTOR):            self._local_state['env_step'] = 0            self._local_state['env_episode'] = 0        self._event_name = "context_exchanger_{role}"        self._skip_n_iter = skip_n_iter        self._storage_loader = storage_loader        for role in task.role:  # Only subscribe to other roles            if not task.has_role(role):                task.on(self._event_name.format(role=role), self.put)        if storage_loader:            task.once("finish", lambda _: storage_loader.shutdown())    def __new__(cls, *args, **kwargs):        if not task.router.is_active:            return task.void()        if len(task.roles) == 0:            logging.warning("The task does not have any roles defined, the ContextExchanger will not work.")            return task.void()        if len(task.roles) > 1:            logging.warning(                "Use multiple roles in one exchanger may lead to unexpected result, please check your code."            )        return super(ContextExchanger, cls).__new__(cls)    def __call__(self, ctx: "Context"):        self.merge(ctx)        yield        payload = self.fetch(ctx)        if payload:            if self._storage_loader and task.has_role(task.role.COLLECTOR):                payload = self._storage_loader.save(payload)            for role in task.roles:                task.emit(self._event_name.format(role=role), payload, only_remote=True)    def __del__(self):        if self._storage_loader:            self._storage_loader.shutdown()    def put(self, payload: Union[Dict, Storage]):        """        Overview:            Get attributes from ctx on the callback of event.            Each attribute should have a standalone put handler, which named `_put_{key}`        """        def callback(payload: Dict):            for key, item in payload.items():                fn_name = "_put_{}".format(key)                if hasattr(self, fn_name):                    getattr(self, fn_name)(item)                else:                    logging.warning("Receive unexpected key ({}) in context exchanger".format(key))        if isinstance(payload, Storage):            assert self._storage_loader is not None, "Storage loader is not defined when data is a storage object."            self._storage_loader.load(payload, callback)        else:            callback(payload)    def fetch(self, ctx: "Context") -> Dict[str, Any]:        """        Overview:            Fetch attributes from ctx before emit them to the event bus.            Each attribute should have a standalone fetch handler, which named `_fetch_{key}`        """        payload = {}        for field in fields(ctx):            key, item = field.name, getattr(ctx, field.name)            fn_name = "_fetch_{}".format(key)            if hasattr(self, fn_name):                value = getattr(self, fn_name)(item)                if value is not None:                    payload[key] = value        return payload    def merge(self, ctx: "Context"):        if task.has_role(task.role.LEARNER):            # Learner should always wait for trajs.            # TODO: Automaticlly wait based on properties, not roles.            while len(self._state) == 0:                sleep(0.01)        elif ctx.total_step >= self._skip_n_iter:            start = time()            while len(self._state) == 0:                if time() - start > 60:                    logging.warning("Timeout when waiting for new context! Node id: {}".format(task.router.node_id))                    break                sleep(0.01)        for k, v in self._state.items():            if not task.has_role(task.role.COLLECTOR) and k.startswith('increment_'):                pure_k = k.split('increment_')[-1]                setattr(ctx, pure_k, getattr(ctx, pure_k) + v)            else:                setattr(ctx, k, v)        self._state = {}    # Handle each attibute of context    def _put_trajectories(self, traj: List[Any]):        if not task.has_role(task.role.LEARNER):            return        if "trajectories" not in self._state:            self._state["trajectories"] = []        self._state["trajectories"].extend(traj)    def _fetch_trajectories(self, traj: List[Any]):        if task.has_role(task.role.COLLECTOR):            return traj    def _put_episodes(self, episodes: List[Any]):        if not task.has_role(task.role.LEARNER):            return        if "episodes" not in self._state:            self._state["episodes"] = []        self._state["episodes"].extend(episodes)    def _fetch_episodes(self, episodes: List[Any]):        if task.has_role(task.role.COLLECTOR):            return episodes    def _put_trajectory_end_idx(self, trajectory_end_idx: List[str]):        if not task.has_role(task.role.LEARNER):            return        if "trajectory_end_idx" not in self._state:            self._state["trajectory_end_idx"] = []        self._state["trajectory_end_idx"].extend(trajectory_end_idx)    def _fetch_trajectory_end_idx(self, trajectory_end_idx: List[str]):        if task.has_role(task.role.COLLECTOR):            return trajectory_end_idx    def _put_env_step(self, increment_env_step: int):        if not task.has_role(task.role.COLLECTOR):            if 'increment_env_step' not in self._state:                self._state['increment_env_step'] = 0            self._state["increment_env_step"] += increment_env_step    def _fetch_env_step(self, env_step: int):        if task.has_role(task.role.COLLECTOR):            increment_env_step = env_step - self._local_state['env_step']            self._local_state['env_step'] = env_step            return increment_env_step    def _put_env_episode(self, increment_env_episode: int):        if not task.has_role(task.role.COLLECTOR):            if 'increment_env_episode' not in self._state:                self._state['increment_env_episode'] = 0            self._state["increment_env_episode"] += increment_env_episode    def _fetch_env_episode(self, env_episode: int):        if task.has_role(task.role.COLLECTOR):            increment_env_episode = env_episode - self._local_state['env_episode']            self._local_state['env_episode'] = env_episode            return increment_env_episode    def _put_train_iter(self, train_iter: int):        if not task.has_role(task.role.LEARNER):            self._state["train_iter"] = train_iter    def _fetch_train_iter(self, train_iter: int):        if task.has_role(task.role.LEARNER):            return train_iterclass ModelExchanger:    def __init__(self, model: "Module", model_loader: Optional[ModelLoader] = None) -> None:        """        Overview:            Exchange model between processes, only the learner will send the model,            otherwise the model will only be received.            If you are using a shared model on a single host, there is no need to use this middleware.        Arguments:            - model (:obj:`torch.nn.Module`): Pytorch module.            - model_loader (:obj:`ModelLoader`): Encode model in subprocess.        """        self._model = model        self._model_loader = model_loader        self._event_name = "model_exchanger"        self._state_dict_cache: Optional[Union[object, Storage]] = None        self._is_learner = task.has_role(task.role.LEARNER)        if not self._is_learner:            task.on(self._event_name, self._cache_state_dict)        if model_loader:            task.once("finish", lambda _: model_loader.shutdown())    def _cache_state_dict(self, state_dict: Union[object, Storage]):        self._state_dict_cache = state_dict    def __new__(cls, *args, **kwargs):        if not task.router.is_active:            return task.void()        if len(task.roles) == 0:            logging.warning("The task does not have any roles defined, the ModelExchanger will not work.")            return task.void()        if len(task.roles) > 1:            logging.warning(                "Use multiple roles in one exchanger may lead to unexpected result, please check your code."            )        return super(ModelExchanger, cls).__new__(cls)    def __call__(self, ctx: "Context") -> Any:        if self._model_loader:            self._model_loader.start()        if not self._is_learner:            if ctx.total_step != 0:  # Skip first iteration                self._update_model()        else:            yield            self._send_model()    def _update_model(self):        start = time()        while True:            if task.finish:                return            if time() - start > 60:                logging.warning("Timeout when waiting for new model! Node id: {}".format(task.router.node_id))                break            if self._state_dict_cache is None:                sleep(0.01)            else:                if isinstance(self._state_dict_cache, Storage) and self._model_loader is not None:                    try:                        self._model.load_state_dict(self._model_loader.load(self._state_dict_cache))                        self._state_dict_cache = None                        break                    except FileNotFoundError as e:                        logging.warning(                            "Model file has been deleted on node {}, maybe you can increase the ttl.".format(                                task.router.node_id                            )                        )                        self._state_dict_cache = None                        continue                else:                    self._model.load_state_dict(self._state_dict_cache)                    self._state_dict_cache = None                    break    def _send_model(self):        if self._model_loader:            self._model_loader.save(self._send_callback)        else:            task.emit(self._event_name, self._model.state_dict(), only_remote=True)    def _send_callback(self, storage: Storage):        if task.running:            task.emit(self._event_name, storage, only_remote=True)    def __del__(self):        if self._model_loader:            self._model_loader.shutdown()class PeriodicalModelExchanger:    def __init__(            self,            model: "Module",            mode: str,            period: int = 1,            delay_toleration: float = np.inf,            stale_toleration: int = 1,            event_name: str = "model_exchanger",            model_loader: Optional[ModelLoader] = None    ) -> None:        """        Overview:            Exchange model between processes, set the mode to "send" or "receive" to specify the role of the process.            If you are using a shared model on a single host, there is no need to use this middleware.        Arguments:            - model (:obj:`torch.nn.Module`): Pytorch module.            - mode (:obj:`str`): "send" or "receive".            - period (:obj:`int`): The period of model exchange.            - delay_toleration (:obj:`float`): The permitted time interval for receiving model after being sent.            - stale_toleration (:obj:`int`): The permitted number of iterations for receiving model after being sent.            - event_name (:obj:`str`): The event name for model exchange.            - model_loader (:obj:`ModelLoader`): ModelLoader for this PeriodicalModelExchanger to use.        """        self._model = model        self._model_loader = model_loader        self._event_name = event_name        self._period = period        self._mode = mode        if self._mode == "receive":            self._id_counter = -1            self._model_id = -1        else:            self._id_counter = 0        self._stale_toleration = stale_toleration        self._model_stale = stale_toleration        self._delay_toleration = delay_toleration        self._state_dict_cache: Optional[Union[object, Storage]] = None        if self._mode == "receive":            task.on(self._event_name, self._cache_state_dict)        if model_loader:            task.once("finish", lambda _: model_loader.shutdown())    def _cache_state_dict(self, msg: Dict[str, Any]):        if msg['id'] % self._period == 0:            self._state_dict_cache = msg['model']            self._id_counter = msg['id']            self._time = msg['time']    def __new__(cls, *args, **kwargs):        return super(PeriodicalModelExchanger, cls).__new__(cls)    def __call__(self, ctx: "Context") -> Any:        if self._model_loader:            self._model_loader.start()        if self._mode == "receive":            if ctx.total_step != 0:  # Skip first iteration                self._update_model()        elif self._mode == "send":            yield            if self._id_counter % self._period == 0:                self._send_model(id=self._id_counter)            self._id_counter += 1        else:            raise NotImplementedError    def _update_model(self):        start = time()        while True:            if task.finish:                return            if time() - start > 60:                logging.warning("Timeout when waiting for new model! Node id: {}".format(task.router.node_id))                self._model_stale += 1                break            if self._state_dict_cache is None:                if self._model_stale < self._stale_toleration and time() - self._time < self._delay_toleration:                    self._model_stale += 1                    break                else:                    sleep(0.01)            else:                if self._id_counter > self._model_id and time() - self._time < self._delay_toleration:                    if isinstance(self._state_dict_cache, Storage) and self._model_loader is not None:                        try:                            self._model.load_state_dict(self._model_loader.load(self._state_dict_cache))                            self._state_dict_cache = None                            self._model_id = self._id_counter                            self._model_stale = 1                            break                        except FileNotFoundError as e:                            logging.warning(                                "Model file has been deleted on node {}, maybe you can increase the ttl.".format(                                    task.router.node_id                                )                            )                            self._state_dict_cache = None                            continue                    else:                        self._model.load_state_dict(self._state_dict_cache)                        self._state_dict_cache = None                        self._model_id = self._id_counter                        self._model_stale = 1                        break                else:                    self._model_stale += 1    def _send_model(self, id: int):        if self._model_loader:            self._model_loader.save(self._send_callback)        else:            task.emit(self._event_name, {'id': id, 'model': self._model.state_dict(), 'time': time()}, only_remote=True)    def _send_callback(self, storage: Storage):        if task.running:            task.emit(self._event_name, storage, only_remote=True)    def __del__(self):        if self._model_loader:            self._model_loader.shutdown()

ding.framework.middleware.distributer¶