385 lines
13 KiB
Python
385 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
schedkit.py -- deterministic parallel scheduler for DAGs of agent tasks.
|
|
|
|
Lane: Codeforge (null).
|
|
|
|
A troupe schedules work as a graph: tasks with prerequisites. Given the same
|
|
graph, this module must always produce the SAME execution plan -- no threads,
|
|
no wall-clock, no set-iteration order leaking into the output. The plan is a
|
|
list of *waves*: the tasks that may start together, in lexicographic order.
|
|
|
|
Determinism contract (the reason this module exists):
|
|
plan(g) == plan(g) for every graph g
|
|
plan(g) == plan(relabel(g)) for every id permutation
|
|
|
|
Cycles are a *result*, not a crash: find_cycles() returns every distinct simple
|
|
cycle (canonically rotated, deduplicated, sorted), and plan() raises CycleError
|
|
carrying them.
|
|
|
|
Public API
|
|
----------
|
|
Graph graph container; add_task, add_dep, add_edge, copy, relabel
|
|
find_cycles(graph) -> tuple of cycles, each a tuple of task ids
|
|
plan(graph) -> tuple of waves (each wave a tuple of ids, sorted)
|
|
plan_graphviz(graph) -> Graphviz DOT source (str) with waves as ranks
|
|
plan_json(graph) -> JSON document (str) with waves, deps, validation
|
|
iter_plan(waves) -> flatten a plan in execution order
|
|
|
|
stdlib only. No I/O, no clock, no randomness.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Dict, Iterable, Iterator, List, Sequence, Set, Tuple
|
|
|
|
__all__ = [
|
|
"Graph",
|
|
"CycleError",
|
|
"find_cycles",
|
|
"plan",
|
|
"plan_graphviz",
|
|
"plan_json",
|
|
"iter_plan",
|
|
]
|
|
|
|
TASK = str
|
|
CYCLE = Tuple[TASK, ...]
|
|
TASKS = Tuple[TASK, ...]
|
|
WAVES = Tuple[TASKS, ...]
|
|
|
|
|
|
class CycleError(ValueError):
|
|
"""Raised by plan() when the dependency graph contains a cycle."""
|
|
|
|
def __init__(self, cycles: Sequence[CYCLE]) -> None:
|
|
self.cycles: WAVES = tuple(tuple(c) for c in cycles)
|
|
shown = ", ".join(" -> ".join(list(c) + [c[0]]) for c in self.cycles[:3])
|
|
more = "" if len(self.cycles) <= 3 else " (+%d more)" % (len(self.cycles) - 3)
|
|
super().__init__("dependency cycle detected: %s%s" % (shown, more))
|
|
|
|
|
|
class Graph:
|
|
"""A directed dependency graph over string task ids.
|
|
|
|
deps[t] is the set of tasks that must complete before t may start.
|
|
Edges point prerequisite -> dependent.
|
|
"""
|
|
|
|
__slots__ = ("_deps", "_dependents")
|
|
|
|
def __init__(self) -> None:
|
|
self._deps: Dict[TASK, Set[TASK]] = {}
|
|
self._dependents: Dict[TASK, Set[TASK]] = {}
|
|
|
|
# ------------------------------------------------------------------ build
|
|
def add_task(self, task: TASK, deps: Iterable[TASK] = ()) -> "Graph":
|
|
if not isinstance(task, str):
|
|
raise TypeError("task id must be str, got %s" % type(task).__name__)
|
|
if not task:
|
|
raise ValueError("task id must be non-empty")
|
|
if task in self._deps and deps:
|
|
raise ValueError(
|
|
"add_task(%r, deps=...) would silently drop existing deps" % task
|
|
)
|
|
self._deps.setdefault(task, set())
|
|
self._dependents.setdefault(task, set())
|
|
for d in deps:
|
|
self.add_dep(task, d)
|
|
return self
|
|
|
|
def add_dep(self, task: TASK, prerequisite: TASK) -> "Graph":
|
|
self.add_task(task)
|
|
self.add_task(prerequisite)
|
|
self._deps[task].add(prerequisite)
|
|
self._dependents[prerequisite].add(task)
|
|
return self
|
|
|
|
def add_edge(self, prerequisite: TASK, task: TASK) -> "Graph":
|
|
return self.add_dep(task, prerequisite)
|
|
|
|
def add_edges(self, pairs: Iterable[Tuple[TASK, TASK]]) -> "Graph":
|
|
for pre, post in pairs:
|
|
self.add_edge(pre, post)
|
|
return self
|
|
|
|
# --------------------------------------------------------------- inspect
|
|
def tasks(self) -> TASKS:
|
|
return tuple(sorted(self._deps))
|
|
|
|
def deps(self, task: TASK) -> TASKS:
|
|
return tuple(sorted(self._deps[task]))
|
|
|
|
def dependents(self, task: TASK) -> TASKS:
|
|
return tuple(sorted(self._dependents.get(task, ())))
|
|
|
|
def edges(self) -> Tuple[Tuple[TASK, TASK], ...]:
|
|
return tuple(
|
|
(pre, t) for t in sorted(self._deps) for pre in sorted(self._deps[t])
|
|
)
|
|
|
|
def __contains__(self, task: object) -> bool:
|
|
return task in self._deps
|
|
|
|
def __len__(self) -> int:
|
|
return len(self._deps)
|
|
|
|
def __bool__(self) -> bool:
|
|
return bool(self._deps)
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
return isinstance(other, Graph) and self.edges() == other.edges()
|
|
|
|
def __hash__(self) -> int:
|
|
return hash(self.edges())
|
|
|
|
def __repr__(self) -> str:
|
|
return "Graph(%r)" % (self.edges(),)
|
|
|
|
# ---------------------------------------------------------------- derive
|
|
def copy(self) -> "Graph":
|
|
g = Graph()
|
|
for t in sorted(self._deps):
|
|
g.add_task(t, sorted(self._deps[t]))
|
|
return g
|
|
|
|
def relabel(self, mapping: Dict[TASK, TASK]) -> "Graph":
|
|
"""Return an isomorphic graph under an injective id mapping.
|
|
|
|
Unmapped tasks keep their id. Raises ValueError if the mapping collides.
|
|
"""
|
|
values = list(mapping.values())
|
|
if len(set(values)) != len(values):
|
|
raise ValueError("relabel mapping is not injective")
|
|
out = Graph()
|
|
for t in sorted(self._deps):
|
|
out.add_task(
|
|
mapping.get(t, t),
|
|
sorted(mapping.get(d, d) for d in self._deps[t]),
|
|
)
|
|
return out
|
|
|
|
def topological_order(self) -> TASKS:
|
|
"""Kahn's algorithm, min-id first: a deterministic linear extension."""
|
|
order: List[TASK] = []
|
|
ready: List[TASK] = []
|
|
pending = {t: set(self._deps[t]) for t in self._deps}
|
|
for t in sorted(pending):
|
|
if not pending[t]:
|
|
ready.append(t)
|
|
while ready:
|
|
t = min(ready)
|
|
ready.remove(t)
|
|
order.append(t)
|
|
for d in sorted(self._dependents.get(t, ())):
|
|
pending[d].discard(t)
|
|
if not pending[d]:
|
|
ready.append(d)
|
|
if len(order) != len(self._deps):
|
|
raise CycleError(find_cycles(self))
|
|
return tuple(order)
|
|
|
|
|
|
# --------------------------------------------------------------------- cycles
|
|
def _strongly_connected(adj: Dict[TASK, List[TASK]], nodes: Sequence[TASK]) -> Set[TASK]:
|
|
"""Tarjan's SCC algorithm, iterative.
|
|
|
|
Returns the vertices of the first SCC with more than one vertex, or any
|
|
self-loop, in the subgraph induced by `nodes`. Empty set if none.
|
|
"""
|
|
index_of: Dict[TASK, int] = {}
|
|
low: Dict[TASK, int] = {}
|
|
on_stack: Set[TASK] = set()
|
|
stack: List[TASK] = []
|
|
counter = 0
|
|
|
|
for root in nodes:
|
|
if root in index_of:
|
|
continue
|
|
work: List[Tuple[TASK, int]] = [(root, 0)]
|
|
while work:
|
|
v, pi = work[-1]
|
|
if pi == 0:
|
|
index_of[v] = low[v] = counter
|
|
counter += 1
|
|
stack.append(v)
|
|
on_stack.add(v)
|
|
recurse = False
|
|
succ = adj[v]
|
|
for i in range(pi, len(succ)):
|
|
w = succ[i]
|
|
if w not in index_of:
|
|
work[-1] = (v, i + 1)
|
|
work.append((w, 0))
|
|
recurse = True
|
|
break
|
|
if w in on_stack:
|
|
low[v] = min(low[v], index_of[w])
|
|
if recurse:
|
|
continue
|
|
if low[v] == index_of[v]:
|
|
comp: List[TASK] = []
|
|
while True:
|
|
w = stack.pop()
|
|
on_stack.discard(w)
|
|
comp.append(w)
|
|
if w == v:
|
|
break
|
|
if len(comp) > 1 or (len(comp) == 1 and comp[0] in adj[comp[0]]):
|
|
return set(comp)
|
|
work.pop()
|
|
if work:
|
|
u, _ = work[-1]
|
|
low[u] = min(low[u], low[v])
|
|
return set()
|
|
|
|
|
|
def find_cycles(g: Graph) -> Tuple[CYCLE, ...]:
|
|
"""Return every distinct simple cycle: canonically rotated, deduped, sorted.
|
|
|
|
Johnson's algorithm -- enumerate elementary circuits without repetition.
|
|
Iterative throughout, so a 100k-node graph cannot exhaust the interpreter
|
|
stack. Deterministic: same graph -> same tuple of cycles in the same order.
|
|
"""
|
|
adj: Dict[TASK, List[TASK]] = {t: sorted(g.dependents(t)) for t in sorted(g._deps)}
|
|
ids = sorted(adj)
|
|
if not ids:
|
|
return ()
|
|
|
|
found: Set[CYCLE] = set()
|
|
|
|
def _unblock(node: TASK, blocked: Set[TASK], blocked_map: Dict[TASK, Set[TASK]]) -> None:
|
|
todo = [node]
|
|
while todo:
|
|
n = todo.pop()
|
|
if n in blocked:
|
|
blocked.discard(n)
|
|
todo.extend(sorted(blocked_map.get(n, ())))
|
|
blocked_map[n] = set()
|
|
|
|
def _circuit(start: TASK, adj_s: Dict[TASK, List[TASK]], nodes: Set[TASK]) -> bool:
|
|
blocked: Set[TASK] = set()
|
|
blocked_map: Dict[TASK, Set[TASK]] = {v: set() for v in nodes}
|
|
stack: List[TASK] = []
|
|
found_here: List[CYCLE] = []
|
|
|
|
def _walk(v: TASK) -> bool:
|
|
closed = False
|
|
stack.append(v)
|
|
blocked.add(v)
|
|
for w in adj_s[v]:
|
|
if w == start:
|
|
found_here.append(tuple(stack))
|
|
closed = True
|
|
elif w not in blocked and w in nodes:
|
|
if _walk(w):
|
|
closed = True
|
|
if closed:
|
|
_unblock(v, blocked, blocked_map)
|
|
else:
|
|
for w in adj_s[v]:
|
|
if w in nodes:
|
|
blocked_map.setdefault(w, set()).add(v)
|
|
stack.pop()
|
|
return closed
|
|
|
|
_walk(start)
|
|
found.update(_canon(c) for c in found_here)
|
|
return bool(found_here)
|
|
|
|
remaining: Dict[TASK, List[TASK]] = dict(adj)
|
|
for s in ids:
|
|
if s not in remaining:
|
|
continue
|
|
sub_nodes = set(_strongly_connected(remaining, [s]))
|
|
if sub_nodes:
|
|
_circuit(s, remaining, sub_nodes)
|
|
# Remove s from the graph for all subsequent iterations (Johnson 1975).
|
|
del remaining[s]
|
|
for v in remaining:
|
|
if s in remaining[v]:
|
|
remaining[v] = [w for w in remaining[v] if w != s]
|
|
|
|
return tuple(sorted(found))
|
|
|
|
|
|
def _canon(cycle: Sequence[TASK]) -> CYCLE:
|
|
"""Rotate a cycle so its lexicographically smallest node comes first."""
|
|
cyc = tuple(cycle)
|
|
if not cyc:
|
|
return cyc
|
|
i = min(range(len(cyc)), key=lambda k: (cyc[k], k))
|
|
return cyc[i:] + cyc[:i]
|
|
|
|
|
|
# ------------------------------------------------------------------------ plan
|
|
def plan(g: Graph) -> WAVES:
|
|
"""Return the deterministic parallel execution plan: a tuple of waves.
|
|
|
|
A task sits in wave k where k = 1 + max(wave of each prerequisite); tasks
|
|
with no prerequisites are in wave 0. Ids sort lexicographically within a
|
|
wave. Raises CycleError if the graph is cyclic.
|
|
"""
|
|
if not g:
|
|
return ()
|
|
cycles = find_cycles(g)
|
|
if cycles:
|
|
raise CycleError(cycles)
|
|
|
|
order = g.topological_order() # deterministic linear extension
|
|
wave_of: Dict[TASK, int] = {}
|
|
for t in order:
|
|
wave_of[t] = 1 + max((wave_of[d] for d in g.deps(t)), default=-1)
|
|
|
|
n = max(wave_of.values()) + 1
|
|
waves: List[List[TASK]] = [[] for _ in range(n)]
|
|
for t in order:
|
|
waves[wave_of[t]].append(t)
|
|
return tuple(tuple(w) for w in waves)
|
|
|
|
|
|
def plan_graphviz(g: Graph) -> str:
|
|
"""Graphviz DOT source for the plan; tasks in a wave share a rank."""
|
|
cycles = find_cycles(g)
|
|
if cycles:
|
|
raise CycleError(cycles)
|
|
waves = plan(g)
|
|
rank_of = {t: i for i, w in enumerate(waves) for t in w}
|
|
lines = ["digraph plan {", " rankdir=LR;", " node [shape=box];"]
|
|
for t in g.tasks():
|
|
lines.append(' "%s" [label="w%d: %s"];' % (t, rank_of[t], t))
|
|
for pre, post in g.edges():
|
|
lines.append(' "%s" -> "%s";' % (pre, post))
|
|
for w in waves:
|
|
lines.append(" { rank=same; %s }" % "; ".join('"%s"' % t for t in w))
|
|
lines.append("}")
|
|
return "\n".join(lines) + "\n"
|
|
|
|
|
|
def plan_json(g: Graph) -> str:
|
|
"""Machine-readable plan: waves, deps, validation, and the graph's edges."""
|
|
waves = plan(g)
|
|
rank_of = {t: i for i, w in enumerate(waves) for t in w}
|
|
doc = {
|
|
"tasks": [
|
|
{"id": t, "wave": rank_of[t], "deps": list(g.deps(t))}
|
|
for t in g.tasks()
|
|
],
|
|
"waves": [list(w) for w in waves],
|
|
"validation": {
|
|
"acyclic": True,
|
|
"n_tasks": len(g),
|
|
"n_waves": len(waves),
|
|
"edges": [list(e) for e in g.edges()],
|
|
},
|
|
}
|
|
return json.dumps(doc, indent=2) + "\n"
|
|
|
|
|
|
def iter_plan(waves: Iterable[Sequence[TASK]]) -> Iterator[TASK]:
|
|
"""Flatten a plan in execution order (wave by wave, id order inside)."""
|
|
for w in waves:
|
|
for t in w:
|
|
yield t
|