Controller class initializer. Args: item: The metagraph to place wrapped in a cluster. cluster: A cluster of devices on which to place the item.
(self, item, cluster)
| 25 | """Controller class.""" |
| 26 | |
| 27 | def __init__(self, item, cluster): |
| 28 | """Controller class initializer. |
| 29 | |
| 30 | Args: |
| 31 | item: The metagraph to place wrapped in a cluster. |
| 32 | cluster: A cluster of devices on which to place the item. |
| 33 | """ |
| 34 | self.item = item |
| 35 | |
| 36 | self._node = {} |
| 37 | for node in item.metagraph.graph_def.node: |
| 38 | self._node[node.name] = node |
| 39 | |
| 40 | self._fanout = defaultdict(lambda: []) |
| 41 | for node in item.metagraph.graph_def.node: |
| 42 | for fanin in self._get_node_fanin(node): |
| 43 | self._fanout[fanin.name].append(node) |
| 44 | |
| 45 | important_op_names = item.IdentifyImportantOps(sort_topologically=True) |
| 46 | |
| 47 | # List of important ops (these are the ops to place) sorted in topological |
| 48 | # order. The order of this collection is deterministic. |
| 49 | self.important_ops = [] |
| 50 | for name in important_op_names: |
| 51 | self.important_ops.append(self._node[name]) |
| 52 | |
| 53 | self.node_properties = item.GetOpProperties() |
| 54 | |
| 55 | self.cluster = cluster |
| 56 | self.devices = cluster.ListDevices() |
| 57 | |
| 58 | self.colocation_constraints = item.GetColocationGroups() |
| 59 | |
| 60 | self.placement_constraints = cluster.GetSupportedDevices(item) |
| 61 | for node_name, dev in self.placement_constraints.items(): |
| 62 | if len(dev) == 1: |
| 63 | # Place the node on the supported device |
| 64 | node = self._node[node_name] |
| 65 | node.device = dev[0] |
| 66 | fanout = self.get_node_fanout(node) |
| 67 | # Update the fanout of the fanin to bypass the node |
| 68 | for fanin in self._get_node_fanin(node): |
| 69 | fanout_of_fanin = self.get_node_fanout(fanin) |
| 70 | fanout_of_fanin += fanout |
| 71 | fanout_of_fanin.remove(node) |
| 72 | # Remove node from the list of important ops since we don't need to |
| 73 | # place the node. |
| 74 | if node in self.important_ops: |
| 75 | self.important_ops.remove(node) |
| 76 | important_op_names.remove(node.name) |
| 77 | |
| 78 | # List of important op names, in non deterministic order. |
| 79 | self.important_op_names = frozenset(important_op_names) |
| 80 | |
| 81 | @property |
| 82 | def input_graph_def(self): |
nothing calls this directly
no test coverage detected