spatialrust_distribute/
graph.rs

1//! Partition graphs for edge/distributed placement.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use crate::{DistributeError, DistributeResult};
6
7/// One execution node on a device/host.
8#[derive(Clone, Debug, PartialEq, Eq, Hash)]
9pub struct ExecutionNode {
10    /// Node id.
11    pub id: String,
12    /// Device or host label.
13    pub device: String,
14}
15
16/// Named partition of execution nodes.
17#[derive(Clone, Debug, PartialEq, Eq)]
18pub struct ExecutionPartition {
19    /// Partition id.
20    pub id: String,
21    /// Member node ids.
22    pub nodes: Vec<String>,
23}
24
25impl ExecutionPartition {
26    /// Creates a validated non-empty partition.
27    pub fn try_new(id: impl Into<String>, nodes: Vec<String>) -> DistributeResult<Self> {
28        let id = id.into();
29        if id.is_empty() {
30            return Err(DistributeError::InvalidConfiguration(
31                "partition id must be non-empty".into(),
32            ));
33        }
34        if nodes.is_empty() {
35            return Err(DistributeError::InvalidConfiguration(
36                "partition must contain at least one node".into(),
37            ));
38        }
39        if nodes.iter().any(|n| n.is_empty()) {
40            return Err(DistributeError::InvalidConfiguration("node ids must be non-empty".into()));
41        }
42        Ok(Self { id, nodes })
43    }
44
45    /// Returns true when `node_id` is a member.
46    #[must_use]
47    pub fn contains(&self, node_id: &str) -> bool {
48        self.nodes.iter().any(|n| n == node_id)
49    }
50}
51
52/// Directed partition graph with adjacency between partitions.
53#[derive(Clone, Debug, Default)]
54pub struct PartitionGraph {
55    partitions: HashMap<String, ExecutionPartition>,
56    edges: Vec<(String, String)>,
57}
58
59impl PartitionGraph {
60    /// Creates an empty graph.
61    #[must_use]
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    /// Inserts a partition.
67    pub fn insert_partition(&mut self, partition: ExecutionPartition) -> DistributeResult<()> {
68        if partition.id.is_empty() {
69            return Err(DistributeError::InvalidConfiguration(
70                "partition id must be non-empty".into(),
71            ));
72        }
73        if partition.nodes.is_empty() {
74            return Err(DistributeError::InvalidConfiguration(
75                "partition must contain at least one node".into(),
76            ));
77        }
78        self.partitions.insert(partition.id.clone(), partition);
79        Ok(())
80    }
81
82    /// Connects two partitions with a directed edge.
83    pub fn connect(
84        &mut self,
85        from: impl Into<String>,
86        to: impl Into<String>,
87    ) -> DistributeResult<()> {
88        let from = from.into();
89        let to = to.into();
90        if !self.partitions.contains_key(&from) || !self.partitions.contains_key(&to) {
91            return Err(DistributeError::Missing("partition endpoint".into()));
92        }
93        if from == to {
94            return Err(DistributeError::InvalidConfiguration("self-edges are not allowed".into()));
95        }
96        if self.edges.iter().any(|(a, b)| a == &from && b == &to) {
97            return Ok(());
98        }
99        self.edges.push((from, to));
100        Ok(())
101    }
102
103    /// Returns partitions.
104    #[must_use]
105    pub fn partitions(&self) -> &HashMap<String, ExecutionPartition> {
106        &self.partitions
107    }
108
109    /// Returns edges.
110    #[must_use]
111    pub fn edges(&self) -> &[(String, String)] {
112        &self.edges
113    }
114
115    /// Looks up a partition by id.
116    #[must_use]
117    pub fn partition(&self, id: &str) -> Option<&ExecutionPartition> {
118        self.partitions.get(id)
119    }
120
121    /// Finds which partition owns a node id.
122    #[must_use]
123    pub fn partition_of_node(&self, node_id: &str) -> Option<&ExecutionPartition> {
124        self.partitions.values().find(|partition| partition.contains(node_id))
125    }
126
127    /// Returns outgoing neighbors of a partition.
128    #[must_use]
129    pub fn successors(&self, id: &str) -> Vec<&str> {
130        self.edges.iter().filter_map(|(from, to)| (from == id).then_some(to.as_str())).collect()
131    }
132
133    /// Returns a topological order of partitions, or errors on cycles / missing nodes.
134    pub fn topological_order(&self) -> DistributeResult<Vec<String>> {
135        let mut indegree: HashMap<&str, usize> =
136            self.partitions.keys().map(|id| (id.as_str(), 0usize)).collect();
137        for (from, to) in &self.edges {
138            if !indegree.contains_key(from.as_str()) || !indegree.contains_key(to.as_str()) {
139                return Err(DistributeError::Missing("partition endpoint".into()));
140            }
141            *indegree.get_mut(to.as_str()).unwrap() += 1;
142        }
143
144        let mut queue: VecDeque<String> = indegree
145            .iter()
146            .filter(|(_, deg)| **deg == 0)
147            .map(|(id, _)| (*id).to_string())
148            .collect();
149        // Stable order for deterministic schedules.
150        let mut queued: HashSet<String> = queue.iter().cloned().collect();
151        let mut order = Vec::with_capacity(self.partitions.len());
152
153        while let Some(id) = {
154            // Pop lexicographically smallest ready id for determinism.
155            let next = queue.iter().min().cloned();
156            if let Some(ref n) = next {
157                queue.retain(|x| x != n);
158                queued.remove(n);
159            }
160            next
161        } {
162            order.push(id.clone());
163            for succ in self.successors(&id) {
164                let deg = indegree.get_mut(succ).unwrap();
165                *deg = deg.saturating_sub(1);
166                if *deg == 0 && !queued.contains(succ) {
167                    queued.insert(succ.to_string());
168                    queue.push_back(succ.to_string());
169                }
170            }
171        }
172
173        if order.len() != self.partitions.len() {
174            return Err(DistributeError::CycleDetected);
175        }
176        Ok(order)
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::{ExecutionPartition, PartitionGraph};
183
184    #[test]
185    fn connects_partitions_and_orders() {
186        let mut graph = PartitionGraph::new();
187        graph
188            .insert_partition(ExecutionPartition::try_new("edge", vec!["cam".into()]).unwrap())
189            .unwrap();
190        graph
191            .insert_partition(ExecutionPartition::try_new("cloud", vec!["infer".into()]).unwrap())
192            .unwrap();
193        graph.connect("edge", "cloud").unwrap();
194        assert_eq!(graph.edges().len(), 1);
195        assert_eq!(graph.topological_order().unwrap(), vec!["edge", "cloud"]);
196        assert_eq!(graph.partition_of_node("cam").unwrap().id, "edge");
197    }
198
199    #[test]
200    fn detects_cycle() {
201        let mut graph = PartitionGraph::new();
202        graph
203            .insert_partition(ExecutionPartition::try_new("a", vec!["n0".into()]).unwrap())
204            .unwrap();
205        graph
206            .insert_partition(ExecutionPartition::try_new("b", vec!["n1".into()]).unwrap())
207            .unwrap();
208        graph.connect("a", "b").unwrap();
209        graph.connect("b", "a").unwrap();
210        assert!(matches!(graph.topological_order(), Err(crate::DistributeError::CycleDetected)));
211    }
212}
spatialrust_distribute/graph.rs

spatialrust_distribute/
graph.rs
