Fix a bug that sometimes TraceGraph does not follow ordering.

microsoft · Jan 14, 2025 · 6c4c52d · 6c4c52d
1 parent 3636d7c
commit 6c4c52d
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/opto/trace/propagators/graph_propagator.py b/opto/trace/propagators/graph_propagator.py
@@ -103,6 +103,7 @@ def init_feedback(self, node, feedback: Any):
 
     def _propagate(self, child: MessageNode):
         graph = [(p.level, p) for p in child.parents]  # add the parents
+        graph = sorted(graph, key=lambda x: x[0])  # sort by level, heapq in TraceGraph requires this
         feedback = self.aggregate(child.feedback) + TraceGraph(graph=graph, user_feedback=None)
         assert isinstance(feedback, TraceGraph)