wehub-resource-sync
137 lines
5.5 KiB
Python
137 lines
5.5 KiB
Python
import collections
|
|
from dataclasses import dataclass, field
|
|
from typing import Dict, Iterator, List, Optional, Set, Tuple, Type
|
|
|
|
from ray.data._internal.logical.interfaces import Rule
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
|
|
@DeveloperAPI
|
|
class Ruleset:
|
|
"""A collection of rules to apply to a plan.
|
|
|
|
This is a utility class to ensure that, if rules depend on each other, they're
|
|
applied in a correct order.
|
|
"""
|
|
|
|
@dataclass(frozen=True)
|
|
class _Node:
|
|
rule: Type[Rule]
|
|
dependents: List["Ruleset._Node"] = field(default_factory=list)
|
|
|
|
def __init__(self, rules: Optional[List[Type[Rule]]] = None):
|
|
if rules is None:
|
|
rules = []
|
|
|
|
self._rules = list(rules)
|
|
|
|
def add(self, rule: Type[Rule]):
|
|
if rule in self._rules:
|
|
raise ValueError(f"Rule {rule} already in ruleset")
|
|
|
|
self._rules.append(rule)
|
|
|
|
if self._contains_cycle():
|
|
raise ValueError("Cannot add rule that would create a cycle")
|
|
|
|
def remove(self, rule: Type[Rule]):
|
|
if rule not in self._rules:
|
|
raise ValueError(f"Rule {rule} not found in ruleset")
|
|
|
|
self._rules.remove(rule)
|
|
|
|
def __iter__(self) -> Iterator[Type[Rule]]:
|
|
"""Iterate over the rules in this ruleset.
|
|
|
|
This method yields rules in dependency order. For example, if B depends on A,
|
|
then this method yields A before B. Each rule is yielded exactly once, and a
|
|
rule is only yielded once *all* of its dependencies have been yielded (so a
|
|
rule that several others must precede is not emitted early or duplicated).
|
|
Insertion order breaks ties among rules that are ready at the same time.
|
|
"""
|
|
order, _ = self._topological_order()
|
|
for node in order:
|
|
yield node.rule
|
|
|
|
def _topological_order(self) -> Tuple[List["Ruleset._Node"], int]:
|
|
"""Order the nodes by dependency using Kahn's algorithm.
|
|
|
|
Returns the topologically-ordered nodes and the total node count.
|
|
A node is enqueued the moment its in-degree (count of not-yet-emitted
|
|
dependencies) hits zero; since an in-degree only decreases and we
|
|
enqueue solely on the zero-crossing, each node is emitted exactly once.
|
|
Insertion order breaks ties among nodes that are ready together.
|
|
|
|
Nodes that participate in a cycle never reach in-degree zero, so they
|
|
are absent from the result -- i.e. ``len(order) < total`` exactly when
|
|
the graph contains a cycle.
|
|
"""
|
|
nodes, indegree = self._build_graph()
|
|
queue = collections.deque(n for n in nodes if indegree[id(n)] == 0)
|
|
order: List["Ruleset._Node"] = []
|
|
while queue:
|
|
node = queue.popleft()
|
|
order.append(node)
|
|
for dep in node.dependents:
|
|
indegree[id(dep)] -= 1
|
|
if indegree[id(dep)] == 0:
|
|
queue.append(dep)
|
|
return order, len(nodes)
|
|
|
|
def _build_graph(
|
|
self,
|
|
) -> Tuple[List["Ruleset._Node"], Dict[int, int]]:
|
|
"""Build the dependency DAG.
|
|
|
|
Returns the nodes (one per rule, in insertion order) and their
|
|
in-degrees -- the number of rules that must be applied before each.
|
|
The in-degree map is keyed by node identity (``id``) rather than rule
|
|
type so that distinct nodes never share a counter, and is computed as
|
|
the edges are added (every incoming edge bumps the target's in-degree)
|
|
rather than re-derived by a second traversal. A node whose in-degree
|
|
is zero is a root.
|
|
"""
|
|
rule_to_node: Dict[Type[Rule], "Ruleset._Node"] = {
|
|
rule: Ruleset._Node(rule) for rule in self._rules
|
|
}
|
|
indegree: Dict[int, int] = {id(node): 0 for node in rule_to_node.values()}
|
|
|
|
# De-duplicate edges. The same ordering can be declared from both ends
|
|
# -- rule A lists B in ``dependencies()`` while B lists A in
|
|
# ``dependents()`` -- which would otherwise add the edge twice,
|
|
# double-counting the in-degree and duplicating ``dependents`` entries.
|
|
seen_edges: Set[Tuple[int, int]] = set()
|
|
|
|
def add_edge(before: "Ruleset._Node", after: "Ruleset._Node") -> None:
|
|
"""Record that ``before`` must be applied before ``after``."""
|
|
edge = (id(before), id(after))
|
|
if edge in seen_edges:
|
|
return
|
|
seen_edges.add(edge)
|
|
before.dependents.append(after)
|
|
indegree[id(after)] += 1
|
|
|
|
for rule in self._rules:
|
|
node = rule_to_node[rule]
|
|
|
|
# Rules that must be applied *before* this rule: dependency -> node.
|
|
for dependency in rule.dependencies():
|
|
if dependency in rule_to_node:
|
|
add_edge(rule_to_node[dependency], node)
|
|
|
|
# Rules that must be applied *after* this rule: node -> dependent.
|
|
for dependent in rule.dependents():
|
|
if dependent in rule_to_node:
|
|
add_edge(node, rule_to_node[dependent])
|
|
|
|
return list(rule_to_node.values()), indegree
|
|
|
|
def _contains_cycle(self) -> bool:
|
|
# Kahn's traversal drops any node stuck in a cycle (its in-degree never
|
|
# reaches zero), so a shortfall between the ordered nodes and the total
|
|
# means a cycle exists. This correctly flags a graph that mixes an
|
|
# acyclic component with a disjoint cycle -- a plain "does any root
|
|
# exist?" check would be fooled by the acyclic root.
|
|
order, total = self._topological_order()
|
|
return len(order) != total
|