def _create_subcomm(self): """ Internal function to create the sub-communicator that will be used by each independent worker when running the task function Notes ----- This function updates the global ``nbodykit`` communicator via ``set_nbkit_comm`` so all plugins will be initialized with this sub-communicator """ self.subcomm = None chain_ranks = [] color = 0 total_ranks = 0 # split the ranks for i, ranks in split_ranks(self.size, self.cpus_per_worker, include_all=self.use_all_cpus): chain_ranks.append(ranks[0]) if self.rank in ranks: color = i + 1 total_ranks += len(ranks) self.workers = i + 1 # store the total number of workers leftover = (self.size - 1) - total_ranks if leftover and self.rank == 0: args = (self.cpus_per_worker, self.size - 1, leftover) self.logger.warning( "with `cpus_per_worker` = %d and %d available ranks, %d ranks will do no work" % args) # crash if we only have one process or one worker if self.size <= self.workers: args = (self.size, self.workers + 1, self.workers) raise ValueError( "only have %d ranks; need at least %d to use the desired %d workers" % args) # ranks that will do work have a nonzero color now self._valid_worker = color > 0 # split the comm between the workers self.subcomm = self.comm.Split(color, 0) # set the global extension point comm GlobalComm.set(self.subcomm)
def _create_subcomm(self): """ Internal function to create the sub-communicator that will be used by each independent worker when running the task function Notes ----- This function updates the global ``nbodykit`` communicator via ``set_nbkit_comm`` so all plugins will be initialized with this sub-communicator """ self.subcomm = None chain_ranks = [] color = 0 total_ranks = 0 # split the ranks for i, ranks in split_ranks(self.size, self.cpus_per_worker, include_all=self.use_all_cpus): chain_ranks.append(ranks[0]) if self.rank in ranks: color = i+1 total_ranks += len(ranks) self.workers = i+1 # store the total number of workers leftover = (self.size - 1) - total_ranks if leftover and self.rank == 0: args = (self.cpus_per_worker, self.size-1, leftover) self.logger.warning("with `cpus_per_worker` = %d and %d available ranks, %d ranks will do no work" %args) # crash if we only have one process or one worker if self.size <= self.workers: args = (self.size, self.workers+1, self.workers) raise ValueError("only have %d ranks; need at least %d to use the desired %d workers" %args) # ranks that will do work have a nonzero color now self._valid_worker = color > 0 # split the comm between the workers self.subcomm = self.comm.Split(color, 0) # set the global extension point comm GlobalComm.set(self.subcomm)
def wrapper(self, *args, **kwargs): from nbodykit import GlobalComm setattr(self, 'comm', kwargs.pop('comm', GlobalComm.get())) return init(self, *args, **kwargs)