feat: add orchestration engine for GraphGen and tests

Author: ChenZiHong-Gavin

graphgen/engine.py

@@ -0,0 +1,85 @@
+"""
+orchestration engine for GraphGen
+"""
+
+import threading
+from typing import Any, Callable, List
+
+
+class Context(dict):
+    _lock = threading.Lock()
+
+    def set(self, k, v):
+        with self._lock:
+            self[k] = v
+
+    def get(self, k, default=None):
+        with self._lock:
+            return super().get(k, default)
+
+
+class OpNode:
+    def __init__(
+        self, name: str, deps: List[str], func: Callable[["OpNode", Context], Any]
+    ):
+        self.name, self.deps, self.func = name, deps, func
+
+
+def op(name: str, deps: List[str] = None):
+    def decorator(f: Callable[["OpNode", Context], Any]):
+        return OpNode(name, deps or [], f)
+
+    return decorator
+
+
+class Engine:
+    def __init__(self, max_workers: int = 4):
+        self.max_workers = max_workers
+
+    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)
+
+        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 = {}
+
+        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")
+                    done[n].set()
+                    return
+                try:
+                    name2op[n].func(name2op[n], ctx)
+                except Exception as e:  # pylint: disable=broad-except
+                    exc[n] = e
+                done[n].set()
+
+        ts = [threading.Thread(target=_exec, args=(n,), daemon=True) for n in topo]
+        for t in ts:
+            t.start()
+        for t in ts:
+            t.join()
+        if exc:
+            raise RuntimeError(f"Some operations failed: {exc}")

graphgen/generate.py graphgen/run.pygraphgen/generate.py renamed to graphgen/run.py

@@ -0,0 +1,78 @@
+import pytest
+
+from graphgen.engine import Context, Engine, op
+
+engine = Engine(max_workers=2)
+
+
+def test_simple_dag(capsys):
+    """Verify the DAG A->B/C->D execution results and print order."""
+    ctx = Context()
+
+    @op("A")
+    def op_a(self, ctx):
+        print("Running A")
+        ctx.set("A", 1)
+
+    @op("B", deps=["A"])
+    def op_b(self, ctx):
+        print("Running B")
+        ctx.set("B", ctx.get("A") + 1)
+
+    @op("C", deps=["A"])
+    def op_c(self, ctx):
+        print("Running C")
+        ctx.set("C", ctx.get("A") + 2)
+
+    @op("D", deps=["B", "C"])
+    def op_d(self, ctx):
+        print("Running D")
+        ctx.set("D", ctx.get("B") + ctx.get("C"))
+
+    # Explicitly list the nodes to run; avoid relying on globals().
+    ops = [op_a, op_b, op_c, op_d]
+    engine.run(ops, ctx)
+
+    # Assert final results.
+    assert ctx["A"] == 1
+    assert ctx["B"] == 2
+    assert ctx["C"] == 3
+    assert ctx["D"] == 5
+
+    # Assert print order: A must run before B and C; D must run after B and C.
+    captured = capsys.readouterr().out.strip().splitlines()
+    assert "Running A" in captured
+    assert "Running B" in captured
+    assert "Running C" in captured
+    assert "Running D" in captured
+
+    a_idx = next(i for i, line in enumerate(captured) if "Running A" in line)
+    b_idx = next(i for i, line in enumerate(captured) if "Running B" in line)
+    c_idx = next(i for i, line in enumerate(captured) if "Running C" in line)
+    d_idx = next(i for i, line in enumerate(captured) if "Running D" in line)
+
+    assert a_idx < b_idx
+    assert a_idx < c_idx
+    assert d_idx > b_idx
+    assert d_idx > c_idx
+
+
+def test_cyclic_detection():
+    """A cyclic dependency should raise ValueError."""
+    ctx = Context()
+
+    @op("X", deps=["Y"])
+    def op_x(self, ctx):
+        pass
+
+    @op("Y", deps=["X"])
+    def op_y(self, ctx):
+        pass
+
+    ops = [op_x, op_y]
+    with pytest.raises(ValueError, match="Cyclic dependencies"):
+        engine.run(ops, ctx)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])