Class represented as a GraphScope dataflow dag. A :class:`Dag` is always belongs to a session and contains a set of :class:`Operation` object, which performs computations on tensors.
| 27 | |
| 28 | |
| 29 | class Dag(object): |
| 30 | """Class represented as a GraphScope dataflow dag. |
| 31 | |
| 32 | A :class:`Dag` is always belongs to a session and contains a set of |
| 33 | :class:`Operation` object, which performs computations on tensors. |
| 34 | """ |
| 35 | |
| 36 | def __init__(self): |
| 37 | # the order in which op joins the dag, starting by 1. |
| 38 | self._seq = 1 |
| 39 | # mapping from op's key to op |
| 40 | self._ops_by_key = dict() |
| 41 | self._ops_seq_by_key = dict() |
| 42 | |
| 43 | def __str__(self): |
| 44 | return str(self.as_dag_def()) |
| 45 | |
| 46 | def __repr__(self): |
| 47 | return self.__str__() |
| 48 | |
| 49 | def exists(self, op): |
| 50 | if not isinstance(op, Operation): |
| 51 | raise TypeError("op must be an Operation: {0}".format(op)) |
| 52 | return op.key in self._ops_by_key |
| 53 | |
| 54 | def add_op(self, op): |
| 55 | if not isinstance(op, Operation): |
| 56 | raise TypeError("op must be an Operation: {0}".format(op)) |
| 57 | if not op.evaluated and op.key in self._ops_by_key: |
| 58 | raise ValueError("op named {0} already exist in dag".format(op.key)) |
| 59 | self._ops_by_key[op.key] = op |
| 60 | self._ops_seq_by_key[op.key] = self._seq |
| 61 | self._seq += 1 |
| 62 | |
| 63 | def as_dag_def(self): |
| 64 | """Return :class:`Dag` as a :class:`DagDef` proto buffer.""" |
| 65 | dag_def = op_def_pb2.DagDef() |
| 66 | for _, op in self._ops_by_key.items(): |
| 67 | dag_def.op.extend([op.as_op_def()]) |
| 68 | return dag_def |
| 69 | |
| 70 | def to_json(self): |
| 71 | return dict({k: op.to_json() for k, op in self._ops_by_key.items()}) |
| 72 | |
| 73 | def extract_subdag_for(self, ops): |
| 74 | """Extract all nodes included the path that can reach the target ops.""" |
| 75 | out = op_def_pb2.DagDef() |
| 76 | # leaf op handle |
| 77 | # there are two kinds of leaf op: |
| 78 | # 1) unload graph / app |
| 79 | # 2) networkx related op |
| 80 | if len(ops) == 1 and ops[0].is_leaf_op(): |
| 81 | out.op.extend([ops[0].as_op_def()]) |
| 82 | return out |
| 83 | op_keys = list() |
| 84 | # assert op is not present in current dag |
| 85 | for op in ops: |
| 86 | assert op.key in self._ops_by_key, "%s is not in the dag" % op.key |