Melvin Strobl
8be7f9770e
The submodule collapse was only half applied: .gitmodules was deleted but backend/ and frontend/ were still recorded as gitlinks, so none of their files were tracked. Replace the gitlinks with the real trees. Also untrack .env (it carried placeholder secrets) in favour of a tracked .env.example, drop the committed __pycache__, and narrow the blanket *.png ignore that would have swallowed design assets.
219 lines
7.0 KiB
Python
219 lines
7.0 KiB
Python
"""
|
|
DAG generation utilities for testing pipelines.
|
|
|
|
This module provides functions to generate random directed acyclic graphs (DAGs)
|
|
with guaranteed properties like connectivity and no cycles.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import networkx as nx
|
|
import numpy as np
|
|
from typing import Callable
|
|
|
|
from util import Message, NodeParams
|
|
from nodes import Node, MLPNode
|
|
|
|
|
|
def generate_dag(
|
|
num_nodes: int,
|
|
edge_probability: float = 0.3,
|
|
seed: int | None = None,
|
|
) -> nx.DiGraph:
|
|
"""
|
|
Generate a random DAG using the Erdős-Rényi model with topological ordering.
|
|
|
|
Creates edges only from lower-indexed to higher-indexed nodes to guarantee
|
|
acyclicity, then removes isolated nodes to ensure connectivity.
|
|
|
|
:param num_nodes: Number of nodes in the graph.
|
|
:type num_nodes: int
|
|
:param edge_probability: Probability of edge between any two nodes.
|
|
:type edge_probability: float
|
|
:param seed: Random seed for reproducibility.
|
|
:type seed: int | None
|
|
:returns: A random DAG.
|
|
:rtype: nx.DiGraph
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
G = nx.DiGraph()
|
|
G.add_nodes_from(range(num_nodes))
|
|
|
|
# Add edges only from lower to higher index (guarantees DAG)
|
|
for i in range(num_nodes):
|
|
for j in range(i + 1, num_nodes):
|
|
if rng.random() < edge_probability:
|
|
G.add_edge(i, j)
|
|
|
|
# Ensure connectivity: connect isolated nodes
|
|
_ensure_connected(G, rng)
|
|
|
|
return G
|
|
|
|
|
|
def _ensure_connected(G: nx.DiGraph, rng: np.random.Generator) -> None:
|
|
"""
|
|
Ensure the DAG is weakly connected by adding edges.
|
|
|
|
:param G: The graph to modify in place.
|
|
:type G: nx.DiGraph
|
|
:param rng: Random number generator.
|
|
:type rng: np.random.Generator
|
|
"""
|
|
# Get weakly connected components
|
|
components = list(nx.weakly_connected_components(G))
|
|
|
|
if len(components) <= 1:
|
|
return
|
|
|
|
# Sort nodes in each component by index for DAG-safe edge addition
|
|
sorted_components = [sorted(c) for c in components]
|
|
|
|
# Connect components by adding edge from max of one to min of next
|
|
for i in range(len(sorted_components) - 1):
|
|
src = sorted_components[i][-1] # Last (highest) node in component
|
|
dst = sorted_components[i + 1][0] # First (lowest) node in next component
|
|
|
|
# Ensure edge direction maintains DAG property
|
|
if src < dst:
|
|
G.add_edge(src, dst)
|
|
else:
|
|
G.add_edge(dst, src)
|
|
|
|
|
|
def dag_to_pipeline_nodes(
|
|
G: nx.DiGraph,
|
|
node_factory: Callable[..., Node] | None = None,
|
|
params: dict | None = None,
|
|
seed: int | None = None,
|
|
) -> tuple[list[Node], list[Message], list[Node]]:
|
|
"""
|
|
Convert a networkx DAG to pipeline nodes with messages.
|
|
|
|
:param G: The DAG to convert.
|
|
:type G: nx.DiGraph
|
|
:param node_factory: Factory function to create nodes (default: MLPNode).
|
|
:type node_factory: Callable[..., Node] | None
|
|
:param params: Parameters to pass to node factory.
|
|
:type params: dict | None
|
|
:param seed: Random seed for reproducibility.
|
|
:type seed: int | None
|
|
:returns: Tuple of (nodes, all_messages, trigger_nodes).
|
|
:rtype: tuple[list[Node], list[Message], list[Node]]
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
params = params or {}
|
|
|
|
# Create a message for each edge
|
|
edge_messages: dict[tuple[int, int], Message] = {}
|
|
for src, dst in G.edges():
|
|
edge_messages[(src, dst)] = Message(name=f"msg_{src}_{dst}")
|
|
|
|
# Create nodes in topological order
|
|
nodes: list[Node] = []
|
|
trigger_nodes: list[Node] = []
|
|
|
|
for node_id in nx.topological_sort(G):
|
|
# Inputs: messages from incoming edges
|
|
requires = [edge_messages[(src, node_id)] for src in G.predecessors(node_id)]
|
|
|
|
# Outputs: messages for outgoing edges
|
|
provides = [edge_messages[(node_id, dst)] for dst in G.successors(node_id)]
|
|
|
|
# Skip nodes with no outputs (sink nodes produce no messages)
|
|
# But we still need to create them to consume inputs
|
|
if node_factory:
|
|
node = node_factory(
|
|
requires=requires,
|
|
provides=provides,
|
|
params=params,
|
|
name=f"{node_id}",
|
|
)
|
|
else:
|
|
node = MLPNode(
|
|
requires=requires,
|
|
provides=provides,
|
|
params=params,
|
|
name=f"{node_id}",
|
|
)
|
|
|
|
nodes.append(node)
|
|
|
|
# Trigger nodes are sources (no inputs)
|
|
if len(requires) == 0:
|
|
trigger_nodes.append(node)
|
|
|
|
all_messages = list(edge_messages.values())
|
|
return nodes, all_messages, trigger_nodes
|
|
|
|
|
|
def generate_multi_pipeline_dag(
|
|
num_pipelines: int,
|
|
nodes_per_pipeline: int,
|
|
cross_pipeline_edges: int = 2,
|
|
edge_probability: float = 0.3,
|
|
seed: int | None = None,
|
|
) -> tuple[list[nx.DiGraph], nx.DiGraph]:
|
|
"""
|
|
Generate multiple DAGs that can be composed into a parent pipeline.
|
|
|
|
Creates separate DAGs for each pipeline and adds cross-pipeline edges
|
|
to create dependencies between them.
|
|
|
|
:param num_pipelines: Number of child pipelines.
|
|
:type num_pipelines: int
|
|
:param nodes_per_pipeline: Nodes in each pipeline.
|
|
:type nodes_per_pipeline: int
|
|
:param cross_pipeline_edges: Number of edges connecting pipelines.
|
|
:type cross_pipeline_edges: int
|
|
:param edge_probability: Edge probability within each pipeline.
|
|
:type edge_probability: float
|
|
:param seed: Random seed for reproducibility.
|
|
:type seed: int | None
|
|
:returns: Tuple of (list of pipeline DAGs, combined DAG).
|
|
:rtype: tuple[list[nx.DiGraph], nx.DiGraph]
|
|
"""
|
|
rng = np.random.default_rng(seed)
|
|
|
|
pipeline_graphs: list[nx.DiGraph] = []
|
|
combined = nx.DiGraph()
|
|
|
|
# Generate each pipeline's DAG
|
|
for p in range(num_pipelines):
|
|
offset = p * nodes_per_pipeline
|
|
G = generate_dag(
|
|
nodes_per_pipeline, edge_probability, seed=rng.integers(0, 10000)
|
|
)
|
|
|
|
# Relabel nodes with offset
|
|
mapping = {n: n + offset for n in G.nodes()}
|
|
G = nx.relabel_nodes(G, mapping)
|
|
|
|
# Add pipeline attribute
|
|
for n in G.nodes():
|
|
G.nodes[n]["pipeline"] = p
|
|
|
|
pipeline_graphs.append(G)
|
|
combined = nx.compose(combined, G)
|
|
|
|
# Add cross-pipeline edges (from later nodes to earlier nodes of next pipeline)
|
|
for _ in range(cross_pipeline_edges):
|
|
p1, p2 = rng.choice(num_pipelines, size=2, replace=False)
|
|
if p1 > p2:
|
|
p1, p2 = p2, p1
|
|
|
|
# Get sink nodes from p1 (nodes with no outgoing edges within pipeline)
|
|
p1_nodes = [n for n in pipeline_graphs[p1].nodes()]
|
|
p1_sinks = [n for n in p1_nodes if pipeline_graphs[p1].out_degree(n) == 0]
|
|
|
|
# Get source-ish nodes from p2 (nodes with few incoming edges)
|
|
p2_nodes = [n for n in pipeline_graphs[p2].nodes()]
|
|
p2_sources = [n for n in p2_nodes if combined.in_degree(n) <= 1]
|
|
|
|
if p1_sinks and p2_sources:
|
|
src = rng.choice(p1_sinks)
|
|
dst = rng.choice(p2_sources)
|
|
combined.add_edge(src, dst)
|
|
|
|
return pipeline_graphs, combined
|