Skip to main content

runmat_execution/placement/
graph.rs

1use std::collections::BTreeSet;
2
3use serde::{Deserialize, Serialize};
4
5use super::{ExecutionCandidateDescriptor, PlacementResourceSnapshot, PlacementSignature};
6
7#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
8#[serde(deny_unknown_fields)]
9pub struct CandidateResourceDemand {
10    pub cpu_millicores: u32,
11    pub retained_bytes: u64,
12    pub scratch_bytes: u64,
13    pub queue_slots: u32,
14}
15
16#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
17#[serde(deny_unknown_fields)]
18pub struct PlacementGraphCandidate {
19    pub descriptor: ExecutionCandidateDescriptor,
20    pub resources: CandidateResourceDemand,
21    /// Provider failure may fall back only when outputs remain uncommitted.
22    pub transactional_results: bool,
23}
24
25#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
26#[serde(deny_unknown_fields)]
27pub struct PlacementGraphNode {
28    pub identity: String,
29    pub candidates: Vec<PlacementGraphCandidate>,
30}
31
32#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
33#[serde(deny_unknown_fields)]
34pub struct PlacementGraphEdge {
35    pub from: u32,
36    pub to: u32,
37    pub bytes: u64,
38    pub host_to_provider_ns: u64,
39    pub provider_to_host_ns: u64,
40    pub cross_provider_ns: u64,
41}
42
43#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
44#[serde(deny_unknown_fields)]
45pub struct PlacementGraph {
46    /// Nodes are topologically ordered. Edges must point from a lower index to
47    /// a higher index.
48    pub nodes: Vec<PlacementGraphNode>,
49    pub edges: Vec<PlacementGraphEdge>,
50}
51
52#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
53#[serde(deny_unknown_fields)]
54pub struct PlacementGraphLimits {
55    pub max_nodes: u32,
56    pub max_candidates_per_node: u32,
57    pub max_edges: u32,
58    pub max_providers: u32,
59    pub max_frontier_states: u32,
60    pub max_expansions: u32,
61}
62
63impl Default for PlacementGraphLimits {
64    fn default() -> Self {
65        Self {
66            max_nodes: 64,
67            max_candidates_per_node: 8,
68            max_edges: 256,
69            max_providers: 16,
70            max_frontier_states: 256,
71            max_expansions: 16_384,
72        }
73    }
74}
75
76#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
77#[serde(deny_unknown_fields)]
78pub struct PlacementPlanRequest {
79    pub signature: PlacementSignature,
80    pub graph: PlacementGraph,
81    pub limits: PlacementGraphLimits,
82    pub resources: PlacementResourceSnapshot,
83    pub deterministic: bool,
84    /// When true, provider candidates must guarantee staged results so a
85    /// pre-commit failure can select the CPU/runtime fallback without replay.
86    pub require_transactional_results: bool,
87}
88
89impl PlacementPlanRequest {
90    pub fn validate(&self) -> Result<(), &'static str> {
91        self.signature.validate()?;
92        if self.graph.nodes.is_empty() {
93            return Err("placement graph must contain at least one node");
94        }
95        if self.graph.nodes.len() > self.limits.max_nodes as usize {
96            return Err("placement graph exceeds node limit");
97        }
98        if self.graph.edges.len() > self.limits.max_edges as usize {
99            return Err("placement graph exceeds edge limit");
100        }
101        if self.resources.providers.len() > self.limits.max_providers as usize {
102            return Err("placement resources exceed provider limit");
103        }
104        let mut node_identities = BTreeSet::new();
105        for node in &self.graph.nodes {
106            if node.identity.is_empty()
107                || node.identity.len() > 128
108                || node.identity.chars().any(char::is_control)
109            {
110                return Err("placement node identity is invalid");
111            }
112            if !node_identities.insert(node.identity.as_str()) {
113                return Err("placement node identities must be unique");
114            }
115            if node.candidates.is_empty()
116                || node.candidates.len() > self.limits.max_candidates_per_node as usize
117            {
118                return Err("placement node candidate count is invalid");
119            }
120            if node
121                .candidates
122                .iter()
123                .all(|candidate| candidate.descriptor.kind.is_provider())
124            {
125                return Err("every placement node requires a CPU/runtime fallback");
126            }
127            let mut candidates = BTreeSet::new();
128            for candidate in &node.candidates {
129                candidate.descriptor.validate()?;
130                if !candidates.insert(candidate.descriptor.identity.as_str()) {
131                    return Err("candidate identities must be unique within a node");
132                }
133            }
134        }
135        let mut edges = BTreeSet::new();
136        for edge in &self.graph.edges {
137            if edge.from >= edge.to || edge.to as usize >= self.graph.nodes.len() {
138                return Err("placement edges must follow topological node order");
139            }
140            if !edges.insert((edge.from, edge.to)) {
141                return Err("placement graph edges must be unique");
142            }
143        }
144        if self.limits.max_frontier_states == 0
145            || self.limits.max_expansions == 0
146            || self.limits.max_edges == 0
147            || self.limits.max_providers == 0
148        {
149            return Err("placement search limits must be non-zero");
150        }
151        let mut provider_ids = BTreeSet::new();
152        for provider in &self.resources.providers {
153            if !provider_ids.insert(provider.device_id) {
154                return Err("provider resource device identities must be unique");
155            }
156            if provider
157                .capacity_bytes
158                .is_some_and(|capacity| provider.live_bytes > capacity)
159                || provider
160                    .scratch_available_bytes
161                    .zip(provider.capacity_bytes)
162                    .is_some_and(|(scratch, capacity)| scratch > capacity)
163                || provider.reclaimable_bytes > provider.live_bytes
164                || provider.queue_depth.is_some() != provider.queue_limit.is_some()
165                || provider
166                    .queue_depth
167                    .zip(provider.queue_limit)
168                    .is_some_and(|(depth, limit)| depth > limit)
169            {
170                return Err("provider resource snapshot is inconsistent");
171            }
172        }
173        Ok(())
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use crate::{
180        CandidateExecutionLocation, CandidateOutputResidency, CandidatePreparationState, Digest,
181        EstimateConfidence, EstimateSource, ExecutionCandidateKind, ExecutionCostComponents,
182        ExecutionCostEstimate, PlacementRevision, ProviderResourceSnapshot,
183    };
184
185    use super::*;
186
187    fn candidate(identity: &str, provider: bool) -> PlacementGraphCandidate {
188        let kind = if provider {
189            ExecutionCandidateKind::ProviderFusion
190        } else {
191            ExecutionCandidateKind::SharedRuntime
192        };
193        PlacementGraphCandidate {
194            descriptor: ExecutionCandidateDescriptor {
195                identity: identity.into(),
196                region: None,
197                kind,
198                execution_location: if provider {
199                    CandidateExecutionLocation::Provider { device_id: 7 }
200                } else {
201                    CandidateExecutionLocation::Host
202                },
203                preparation: CandidatePreparationState::Warm,
204                cost: ExecutionCostEstimate {
205                    components: ExecutionCostComponents {
206                        execution_ns: 1,
207                        ..ExecutionCostComponents::default()
208                    },
209                    scratch_bytes: 0,
210                    confidence: EstimateConfidence::Exact,
211                    source: EstimateSource::Synthetic,
212                },
213                // Execution and resulting residency are intentionally
214                // independent: provider work may materialize a host result.
215                output_residency: CandidateOutputResidency::Host,
216                guards: Vec::new(),
217            },
218            resources: CandidateResourceDemand::default(),
219            transactional_results: true,
220        }
221    }
222
223    fn request() -> PlacementPlanRequest {
224        PlacementPlanRequest {
225            signature: PlacementSignature {
226                region: None,
227                operation: "test.graph".into(),
228                runtime_facts: Digest::sha256(b"facts"),
229                revision: PlacementRevision {
230                    program: None,
231                    catalog: Digest::sha256(b"catalog"),
232                    compiler: Digest::sha256(b"compiler"),
233                    provider: Digest::sha256(b"provider"),
234                    policy: Digest::sha256(b"policy"),
235                },
236            },
237            graph: PlacementGraph {
238                nodes: vec![PlacementGraphNode {
239                    identity: "node".into(),
240                    candidates: vec![candidate("cpu", false), candidate("gpu-host", true)],
241                }],
242                edges: Vec::new(),
243            },
244            limits: PlacementGraphLimits::default(),
245            resources: PlacementResourceSnapshot {
246                cpu_millicores_available: 1_000,
247                memory_available_bytes: Some(1_024),
248                cancellation_requested: false,
249                providers: vec![ProviderResourceSnapshot {
250                    device_id: 7,
251                    capacity_bytes: Some(1_024),
252                    live_bytes: 0,
253                    reclaimable_bytes: 0,
254                    scratch_available_bytes: Some(1_024),
255                    queue_depth: Some(0),
256                    queue_limit: Some(1),
257                    lost: false,
258                    epoch: 1,
259                }],
260                epoch: 1,
261            },
262            deterministic: true,
263            require_transactional_results: true,
264        }
265    }
266
267    #[test]
268    fn graph_contract_round_trips_with_independent_execution_and_output_locations() {
269        let request = request();
270        request.validate().unwrap();
271        let encoded = serde_json::to_vec(&request).unwrap();
272        let decoded: PlacementPlanRequest = serde_json::from_slice(&encoded).unwrap();
273        assert_eq!(decoded, request);
274    }
275
276    #[test]
277    fn graph_contract_rejects_unbounded_or_duplicate_edges() {
278        let mut request = request();
279        request.graph.nodes.push(PlacementGraphNode {
280            identity: "second".into(),
281            candidates: vec![candidate("cpu.second", false)],
282        });
283        let edge = PlacementGraphEdge {
284            from: 0,
285            to: 1,
286            bytes: 1,
287            host_to_provider_ns: 1,
288            provider_to_host_ns: 1,
289            cross_provider_ns: 1,
290        };
291        request.graph.edges = vec![edge, edge];
292        assert_eq!(
293            request.validate().unwrap_err(),
294            "placement graph edges must be unique"
295        );
296
297        request.graph.edges = vec![edge];
298        request.limits.max_edges = 0;
299        assert!(request.validate().is_err());
300    }
301}