Class TaskGroup

caffe2/python/task.py:158–351

Context that gathers tasks which will run concurrently, potentially on multiple nodes. All tasks in the same node will share the same workspace and thus can share blobs, while tasks running in different nodes won't be able to directly share data. All tasks of the task group wil

class TaskGroup(context.Managed):
    """
    Context that gathers tasks which will run concurrently, potentially on
    multiple nodes. All tasks in the same node will share the same workspace
    and thus can share blobs, while tasks running in different nodes won't
    be able to directly share data.

    All tasks of the task group will start concurrently, and the task group
    will finish execution when the last task of the group finishes.

    Example:
        # suppose that s1 ... s5 are execution steps or nets.
        with TaskGroup() as tg:
            # these tasks go to default node 'local'
            Task(step=s1)
            Task(step=s2)

            with Node('n2'):
                Task(step=s3)
            with Node('n1'):
                Task(step=s4)
            with Node('n2'):
                Task(step=s5)

        # this will run all steps in parallel.
        # s1 and s2 will run at default node 'local'
        # s3 and s5 will run at node 'n2'
        # s4 will run at node 'n1'
        session.run(tg)
    """
    LOCAL_SETUP = 'local_setup'

    def __init__(self, workspace_type=None):
        self._plan_cache = None
        self._tasks = []
        self._already_used = False
        self._prev_active = None
        self._tasks_to_add = []
        self._report_nets = {}
        self._report_steps = []
        self._workspace_type = workspace_type
        self._tasks_by_node = None
        self._remote_nets = []

    def add_remote_net(self, net):
        self._remote_nets.append(net)

    def remote_nets(self):
        return self._remote_nets

    def add(self, task):
        assert not self._already_used, (
            'Cannot add Task to an already used TaskGroup.')
        assert (
            self._workspace_type is None or
            task._workspace_type is None or
            self._workspace_type == task._workspace_type)
        if task._workspace_type is None: