Training Pipeline Orchestration Mock DAG in Python

from collections import deque
from dataclasses import dataclass, field


@dataclass
class DAGNode:
    name: str
    task: callable
    dependencies: list[str] = field(default_factory=list)


class MockDAG:
    def __init__(self, nodes: list[DAGNode]):
        self.nodes = {n.name: n for n in nodes}
        self.executed: list[str] = []

    def run(self):
        # Build adjacency and in-degree from dependencies
        adjacency = {name: [] for name in self.nodes}
        in_degree = {name: 0 for name in self.nodes}
        for name, node in self.nodes.items():
            for dep in node.dependencies:
                adjacency[dep].append(name)
                in_degree[name] += 1

        # Topological sort via Kahn's algorithm
        queue = deque([n for n, d in in_degree.items() if d == 0])
        while queue:
            current = queue.popleft()
            node = self.nodes[current]
            node.task()
            self.executed.append(node.name)
            for neighbor in adjacency[current]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        if len(self.executed) != len(self.nodes):
            raise RuntimeError("DAG contains cycles")


def main():
    # Example: training pipeline with data loading → preprocessing → training → evaluation
    dag = MockDAG(nodes=[
        DAGNode(name="load_data", task=lambda: print("Loading data")),
        DAGNode(name="preprocess", task=lambda: print("Preprocessing"),
                dependencies=["load_data"]),
        DAGNode(name="train", task=lambda: print("Training model"),
                dependencies=["preprocess"]),
        DAGNode(name="evaluate", task=lambda: print("Evaluating"),
                dependencies=["train"]),
    ])
    dag.run()
    print("Execution order:", dag.executed)


if __name__ == "__main__":
    main()