Class CollectivesConcurrencyControl

caffe2/python/data_parallel_model.py:1307–1344

Creates common worlds (up to max_concurrent_context) and manage the sequential execution of collectives that shares the same context with cyclic control inputs.

class CollectivesConcurrencyControl:
    """
    Creates common worlds (up to max_concurrent_context) and manage the
    sequential execution of collectives that shares the same context with
    cyclic control inputs.
    """
    def __init__(
        self,
        name,
        max_concurrent_context,
        param_init_net,
        rendezvous
    ):
        self.name = name
        self.param_init_net = param_init_net
        self.max_concurrent_context = max_concurrent_context
        self.counter = 0
        self.common_worlds = []
        self.control_inputs = []
        self.rendezvous = rendezvous

    def get_control_and_context(self, control_output_blob):
        common_world, control_input = [None, None]
        current_slot = self.counter % self.max_concurrent_context
        if len(self.common_worlds) < self.max_concurrent_context:
            common_world = _CreateOrCloneCommonWorld(
                self.param_init_net,
                "{}_{}_cw".format(self.name, current_slot),
                rendezvous=self.rendezvous,
            )
            self.common_worlds.append(common_world)
            self.control_inputs.append(control_output_blob)
        else:
            common_world = self.common_worlds[current_slot]
            control_input = self.control_inputs[current_slot]
            self.control_inputs[current_slot] = control_output_blob
        self.counter += 1
        return common_world, control_input


def _AllReduceBlobsDistributed(