514 lines
17 KiB
Python
514 lines
17 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.
|
|
|
|
Cost (measured, CPython 3.14, see spec_schedkit.py section G):
|
|
acyclic graph O(V+E) -- one SCC pass, no circuit search at all
|
|
one SCC of k O(k^2 + k*e_k + circuits) (Johnson's own bound)
|
|
|
|
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 heapq
|
|
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.
|
|
|
|
Precisely: the *lexicographically smallest* topological order of the
|
|
graph. A min-heap over the ready set is what buys that -- popping the
|
|
largest ready id instead yields a valid order that is not the smallest,
|
|
so the heap is part of the contract, not an optimisation detail.
|
|
|
|
A heap, not a list scan. The previous `min(ready)` + `list.remove` made
|
|
a *wide* graph O(V^2): 20,000 tasks ready at once cost ~10s for a single
|
|
plan, and 5,000-wide layers cost ~12s.
|
|
"""
|
|
order: List[TASK] = []
|
|
pending = {t: set(self._deps[t]) for t in self._deps}
|
|
ready: List[TASK] = [t for t in self._deps if not pending[t]]
|
|
heapq.heapify(ready)
|
|
while ready:
|
|
t = heapq.heappop(ready)
|
|
order.append(t)
|
|
for d in sorted(self._dependents.get(t, ())):
|
|
pending[d].discard(t)
|
|
if not pending[d]:
|
|
heapq.heappush(ready, d)
|
|
if len(order) != len(self._deps):
|
|
raise CycleError(find_cycles(self))
|
|
return tuple(order)
|
|
|
|
|
|
# --------------------------------------------------------------------- cycles
|
|
def _sccs(adj, roots):
|
|
"""Iterative Tarjan. Every SCC of the graph, as sorted tuples of indices.
|
|
|
|
`adj[i]` is node i's successor list. Deterministic: roots are taken in the
|
|
given order and every adjacency list is sorted, so the returned list is
|
|
reproducible. Iterative throughout -- no recursion to blow up on a deep
|
|
graph.
|
|
"""
|
|
index_of: Dict[int, int] = {}
|
|
low: Dict[int, int] = {}
|
|
on_stack: Set[int] = set()
|
|
stack: List[int] = []
|
|
counter = 0
|
|
out: List[Tuple[int, ...]] = []
|
|
|
|
for root in roots:
|
|
if root in index_of:
|
|
continue
|
|
work: List[List[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][1] = 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[int] = []
|
|
while True:
|
|
w = stack.pop()
|
|
on_stack.discard(w)
|
|
comp.append(w)
|
|
if w == v:
|
|
break
|
|
out.append(tuple(sorted(comp)))
|
|
work.pop()
|
|
if work:
|
|
u = work[-1][0]
|
|
low[u] = min(low[u], low[v])
|
|
return out
|
|
|
|
|
|
def _scc_from(adj, root, floor):
|
|
"""The SCC containing `root` in the subgraph induced by nodes >= `floor`.
|
|
|
|
Neighbours below `floor` are invisible: that is Johnson's outer-loop vertex
|
|
removal done by index comparison -- O(1) per step -- instead of rewriting
|
|
every adjacency list, which cost O(E) per step and made the previous build
|
|
quadratic on graphs that contained no cycle at all.
|
|
"""
|
|
index_of: Dict[int, int] = {}
|
|
low: Dict[int, int] = {}
|
|
on_stack: Set[int] = set()
|
|
stack: List[int] = []
|
|
counter = 0
|
|
comp: List[int] = [root]
|
|
work: List[List[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 < floor:
|
|
continue
|
|
if w not in index_of:
|
|
work[-1][1] = 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]:
|
|
got: List[int] = []
|
|
while True:
|
|
w = stack.pop()
|
|
on_stack.discard(w)
|
|
got.append(w)
|
|
if w == v:
|
|
break
|
|
if v == root:
|
|
comp = got
|
|
work.pop()
|
|
if work:
|
|
u = work[-1][0]
|
|
low[u] = min(low[u], low[v])
|
|
return comp
|
|
|
|
|
|
def _unblock(node, blocked, blocked_map) -> None:
|
|
todo = [node]
|
|
while todo:
|
|
n = todo.pop()
|
|
if n in blocked:
|
|
blocked.discard(n)
|
|
todo.extend(blocked_map.get(n, ()))
|
|
blocked_map[n] = set()
|
|
|
|
|
|
def _circuit(start, adj, nodes, floor, emit) -> None:
|
|
"""Johnson's CIRCUIT: every elementary circuit through `start` that stays
|
|
inside `nodes` and at or above `floor`.
|
|
|
|
Iterative: an explicit frame stack replaces the recursion, so a 50,000-node
|
|
ring cannot exhaust the interpreter stack.
|
|
"""
|
|
blocked: Set[int] = {start}
|
|
blocked_map: Dict[int, Set[int]] = {}
|
|
path: List[int] = [start]
|
|
frames: List[List[int]] = [[start, 0, False]] # vertex, scan position, closed
|
|
while frames:
|
|
fr = frames[-1]
|
|
v = fr[0]
|
|
succ = adj[v]
|
|
i = fr[1]
|
|
closed = fr[2]
|
|
descended = False
|
|
while i < len(succ):
|
|
w = succ[i]
|
|
i += 1
|
|
if w < floor or w not in nodes:
|
|
continue
|
|
if w == start:
|
|
emit(tuple(path))
|
|
closed = True
|
|
elif w not in blocked:
|
|
fr[1] = i
|
|
fr[2] = closed
|
|
path.append(w)
|
|
blocked.add(w)
|
|
frames.append([w, 0, False])
|
|
descended = True
|
|
break
|
|
if descended:
|
|
continue
|
|
fr[1] = i
|
|
fr[2] = closed
|
|
if closed:
|
|
_unblock(v, blocked, blocked_map)
|
|
else:
|
|
for w in succ:
|
|
if w >= floor and w in nodes:
|
|
blocked_map.setdefault(w, set()).add(v)
|
|
frames.pop()
|
|
path.pop()
|
|
if frames and closed:
|
|
frames[-1][2] = True
|
|
|
|
|
|
def _johnson(comp, adj, ids, found) -> None:
|
|
"""Every elementary circuit inside one SCC (`comp` = sorted indices).
|
|
|
|
Faithful Johnson: walk the local indices upward; at each step take the SCC
|
|
holding the least still-visible vertex and run CIRCUIT from it.
|
|
|
|
One addition to Johnson's own bookkeeping: a visible vertex with no visible
|
|
predecessor cannot lie on a cycle, so it is skipped without an SCC pass.
|
|
Without it a k-node ring costs O(k^2) (a Tarjan walk per node); with it the
|
|
ring costs O(k), because removing each node leaves its successor a source.
|
|
"""
|
|
local = {g: i for i, g in enumerate(comp)}
|
|
ladj = [[local[w] for w in adj[g] if w in local] for g in comp]
|
|
k = len(comp)
|
|
|
|
indeg = [0] * k
|
|
for row in ladj:
|
|
for w in row:
|
|
indeg[w] += 1
|
|
|
|
s = 0
|
|
while s < k:
|
|
if indeg[s]: # a self-loop counts: it is an incoming edge
|
|
sub = _scc_from(ladj, s, s)
|
|
if len(sub) > 1 or (len(sub) == 1 and s in ladj[s]):
|
|
nodes = set(sub)
|
|
_circuit(
|
|
s,
|
|
ladj,
|
|
nodes,
|
|
s,
|
|
lambda c: found.add(_canon(tuple(ids[comp[i]] for i in c))),
|
|
)
|
|
for w in ladj[s]: # remove s from the visible subgraph
|
|
indeg[w] -= 1
|
|
s += 1
|
|
|
|
|
|
def find_cycles(g: Graph) -> Tuple[CYCLE, ...]:
|
|
"""Return every distinct simple cycle: canonically rotated, deduped, sorted.
|
|
|
|
Johnson's algorithm, fronted by a single SCC pass: a vertex that cannot lie
|
|
on a cycle is never handed to the circuit search, so an acyclic graph costs
|
|
O(V+E). Cost inside one SCC of k vertices is O(k^2 + k*e_k + circuits),
|
|
Johnson's own bound.
|
|
|
|
Deterministic: same graph -> same tuple of cycles in the same order.
|
|
"""
|
|
ids = sorted(g._deps)
|
|
if not ids:
|
|
return ()
|
|
idx = {t: i for i, t in enumerate(ids)}
|
|
adj = [[idx[d] for d in g.dependents(t)] for t in ids]
|
|
|
|
found: Set[CYCLE] = set()
|
|
for comp in _sccs(adj, range(len(ids))):
|
|
if len(comp) == 1:
|
|
v = comp[0]
|
|
if v in adj[v]:
|
|
found.add((ids[v],))
|
|
continue
|
|
_johnson(comp, adj, ids, found)
|
|
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 -- *sorted*, not merely "in topological pop order": wave membership is
|
|
the plan, and the order inside a wave must therefore depend on the ids
|
|
alone. Without that, plan(relabel(g)) was not the relabelled plan(g) and
|
|
the isomorphism contract in this module's header was false.
|
|
|
|
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(sorted(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
|