MCPcopy Index your code
hub / github.com/pytorch/pytorch / tasks_by_node

Method tasks_by_node

caffe2/python/task.py:264–335  ·  view source on GitHub ↗
(self, node_remap=None)

Source from the content-addressed store, hash-verified

262 return self._report_nets[node][0]
263
264 def tasks_by_node(self, node_remap=None):
265 # tasks_by_node can't be called twice because the setup won't
266 # work properly a second time.
267 node_map = {}
268 for task in self.tasks():
269 node_map[task.node] =\
270 node_remap(task.node) if node_remap else task.node
271 if self._tasks_by_node is not None:
272 tasks_by_node, prev_node_map = self._tasks_by_node
273 assert prev_node_map == node_map, (
274 'Cannot call tasks_by_node multiple times.')
275 return tasks_by_node
276
277 # now we have report_steps. report_net is deprecated
278 for node, (net, interval) in self._report_nets.items():
279 self.report_step(net, node=node, interval_ms=interval * 1000)
280 self._report_nets = {}
281
282 tasks_by_node = defaultdict(list)
283 for task in self.tasks():
284 mapped_node = node_map[task.node]
285 tasks_by_node[mapped_node].append(task)
286
287 report_steps_by_node = defaultdict(list)
288 for original_node, step in self._report_steps:
289 report_steps_by_node[node_map[original_node]].append(step)
290
291 grouped_by_node = TaskGroup()
292 for node, tasks in tasks_by_node.items():
293 report_steps = report_steps_by_node[node]
294 node_inits, node_exits = get_setup_nets(
295 TaskGroup.LOCAL_SETUP,
296 [t.get_step() for t in tasks] + report_steps,
297 self)
298 # shortcut for single task with no queue
299 steps = report_steps
300 outputs = []
301 grouped_workspace_type = WorkspaceType.PRIVATE
302 for task in tasks:
303 step = task.get_step()
304 step.SetCreateWorkspace(
305 task.workspace_type() == WorkspaceType.PRIVATE)
306 if step is not None:
307 steps.append(step)
308 outputs += task.outputs()
309 # If any of the tasks in the node uses the global workspace,
310 # then set the grouped task to use the global workspace as well
311 if task.workspace_type() == WorkspaceType.GLOBAL:
312 grouped_workspace_type = WorkspaceType.GLOBAL
313 if len(steps) == 0:
314 steps.append(core.execution_step('empty', []))
315 if len(steps) == 1:
316 step = steps[0]
317 else:
318 step = core.execution_step(
319 '%s:body' % node, steps, concurrent_substeps=True)
320 if len(node_inits) > 0 or len(node_exits) > 0:
321 steps = []

Callers 3

to_taskMethod · 0.95
analyze_task_groupFunction · 0.80
print_task_groupFunction · 0.80

Calls 12

tasksMethod · 0.95
report_stepMethod · 0.95
TaskGroupClass · 0.85
get_setup_netsFunction · 0.85
get_stepMethod · 0.80
SetCreateWorkspaceMethod · 0.80
execution_stepMethod · 0.80
TaskClass · 0.70
itemsMethod · 0.45
appendMethod · 0.45
workspace_typeMethod · 0.45
outputsMethod · 0.45

Tested by

no test coverage detected