in chatlearn/runtime/model_flow.py [0:0]
def topological_sort(self):
result = []
level_map = defaultdict(list)
in_degree = defaultdict(int)
# Calculate the in-degree of each vertex
for u in self.model_nodes:
for v in u.output_nodes:
in_degree[v] += 1
for v in u.dependent_output_nodes:
in_degree[v] += 1
# Enqueue all the vertices with an in-degree of 0
queue = deque([u for u in self.model_nodes if in_degree[u] == 0])
# Perform topological sorting
while queue:
current_level = []
for _ in range(len(queue)):
current = queue.popleft()
current_level.append(current)
result.append(current)
# Decrement the in-degree of adjacent vertices
for v in current.output_nodes + current.dependent_output_nodes:
in_degree[v] -= 1
if in_degree[v] == 0:
queue.append(v)
level_map[len(result)].extend(current_level)
# Check if the graph contains a cycle
if len(result) != len(self.model_nodes):
raise RuntimeError("Please check if the graph contains a cycle")
return [v[1] for v in sorted(level_map.items())]