fix: fix params in collect_ops

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -42,15 +42,21 @@ def __init__(
         self.agg_mode = agg_mode
 
 
-def op(name: str, deps=None):
+def op(name: str, deps=None, agg_mode: AggMode = AggMode.ALL_REDUCE):
     deps = deps or []
 
     def decorator(func):
         @wraps(func)
         def _wrapper(*args, **kwargs):
             return func(*args, **kwargs)
 
-        _wrapper.op_node = OpNode(name, deps, lambda self, ctx: func(self, **ctx))
+        _wrapper.op_node = OpNode(
+            name=name,
+            deps=deps,
+            compute_func=lambda self, ctx: func(self),
+            callback_func=lambda self, ctx, results: None,
+            agg_mode=agg_mode,
+        )
         return _wrapper
 
     return decorator
@@ -59,48 +65,45 @@ def _wrapper(*args, **kwargs):
 class Engine:
     def __init__(self, max_workers: int = 4):
         self.max_workers = max_workers
+        self.bucket_mgr = BucketManager()
 
     def run(self, ops: List[OpNode], ctx: Context):
-        name2op = {operation.name: operation for operation in ops}
-
-        # topological sort
-        graph = {n: set(name2op[n].deps) for n in name2op}
-        topo = []
-        q = [n for n, d in graph.items() if not d]
-        while q:
-            cur = q.pop(0)
-            topo.append(cur)
-            for child in [c for c, d in graph.items() if cur in d]:
-                graph[child].remove(cur)
-                if not graph[child]:
-                    q.append(child)
+        name2op = {op.name: op for op in ops}
+        topo_names = [op.name for op in self._topo_sort(ops)]
 
-        if len(topo) != len(ops):
-            raise ValueError(
-                "Cyclic dependencies detected among operations."
-                "Please check your configuration."
-            )
-
-        # semaphore for max_workers
         sem = threading.Semaphore(self.max_workers)
         done = {n: threading.Event() for n in name2op}
         exc = {}
 
+        for node in ops:
+            bucket_size = ctx.get(f"_bucket_size_{node.name}", 1)
+            self.bucket_mgr.register(
+                node.name,
+                bucket_size,
+                node.agg_mode,
+                lambda results, n=node: self._callback_wrapper(n, ctx, results),
+            )
+
         def _exec(n: str):
             with sem:
                 for d in name2op[n].deps:
                     done[d].wait()
                 if any(d in exc for d in name2op[n].deps):
-                    exc[n] = Exception("Skipped due to failed dependencies")
+                    exc[n] = "Skipped due to failed dependencies"
                     done[n].set()
                     return
+
                 try:
-                    name2op[n].func(name2op[n], ctx)
+                    name2op[n].compute_func(name2op[n], ctx)
                 except Exception:  # pylint: disable=broad-except
                     exc[n] = traceback.format_exc()
-                done[n].set()
+                finally:
+                    done[n].set()
 
-        ts = [threading.Thread(target=_exec, args=(n,), daemon=True) for n in topo]
+        ts = [
+            threading.Thread(target=_exec, args=(name,), daemon=True)
+            for name in topo_names
+        ]
         for t in ts:
             t.start()
         for t in ts:
@@ -111,6 +114,34 @@ def _exec(n: str):
                 + "\n".join(f"---- {op} ----\n{tb}" for op, tb in exc.items())
             )
 
+    @staticmethod
+    def _callback_wrapper(node: OpNode, ctx: Context, results: List[Any]):
+        try:
+            node.callback_func(node, ctx, results)
+        except Exception:  # pylint: disable=broad-except
+            traceback.print_exc()
+
+    @staticmethod
+    def _topo_sort(ops: List[OpNode]) -> List[OpNode]:
+        name2op = {operation.name: operation for operation in ops}
+        graph = {n: set(name2op[n].deps) for n in name2op}
+        topo = []
+        q = [n for n, d in graph.items() if not d]
+        while q:
+            cur = q.pop(0)
+            topo.append(name2op[cur])
+            for child in [c for c, d in graph.items() if cur in d]:
+                graph[child].remove(cur)
+                if not graph[child]:
+                    q.append(child)
+
+        if len(topo) != len(ops):
+            raise ValueError(
+                "Cyclic dependencies detected among operations."
+                "Please check your configuration."
+            )
+        return topo
+
 
 class Bucket:
     """
@@ -190,8 +221,9 @@ def collect_ops(config: dict, graph_gen) -> List[OpNode]:
         op_node.deps = runtime_deps
 
         if "params" in stage:
-            op_node.func = lambda self, ctx, m=method, sc=stage: m(sc.get("params", {}))
+            params = stage["params"]
+            op_node.compute_func = lambda self, ctx, m=method, p=params: m(p)
         else:
-            op_node.func = lambda self, ctx, m=method: m()
+            op_node.compute_func = lambda self, ctx, m=method: m()
         ops.append(op_node)
     return ops