dag_ml_core/
graph.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{DagMlError, Result};
6use crate::ids::NodeId;
7use crate::relation::EntityUnitLevel;
8
9pub const GRAPH_SPEC_SCHEMA_VERSION: u32 = 1;
10pub const GRAPH_SPEC_SCHEMA_ID: &str =
11    "https://github.com/GBeurier/dag-ml/schemas/graph_spec.v1.schema.json";
12
13#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15pub enum NodeKind {
16    Transform,
17    YTransform,
18    Split,
19    Model,
20    Fork,
21    Map,
22    FeatureJoin,
23    PredictionJoin,
24    MixedJoin,
25    SourceJoin,
26    Tag,
27    Exclude,
28    Augmentation,
29    Adapter,
30    Aggregator,
31    Generator,
32    Restructure,
33    Tuner,
34    Subgraph,
35    Chart,
36}
37
38#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
39#[serde(rename_all = "snake_case")]
40pub enum PortKind {
41    Data,
42    Target,
43    Prediction,
44    Artifact,
45    Metric,
46    Control,
47}
48
49#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
50#[serde(rename_all = "snake_case")]
51pub enum PortCardinality {
52    One,
53    Many,
54    Optional,
55}
56
57#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
58pub struct PortSpec {
59    pub name: String,
60    pub kind: PortKind,
61    pub representation: Option<String>,
62    pub cardinality: PortCardinality,
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub unit_level: Option<EntityUnitLevel>,
65    #[serde(default, skip_serializing_if = "Option::is_none")]
66    pub alignment_key: Option<String>,
67    #[serde(default, skip_serializing_if = "Option::is_none")]
68    pub target_level: Option<EntityUnitLevel>,
69    #[serde(default)]
70    pub description: String,
71}
72
73#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
74pub struct PortSchema {
75    #[serde(default)]
76    pub inputs: Vec<PortSpec>,
77    #[serde(default)]
78    pub outputs: Vec<PortSpec>,
79}
80
81#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
82pub struct PortRef {
83    pub node_id: NodeId,
84    pub port_name: String,
85}
86
87#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
88pub struct EdgeContract {
89    pub kind: PortKind,
90    pub representation: Option<String>,
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub unit_level: Option<EntityUnitLevel>,
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub alignment_key: Option<String>,
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub target_level: Option<EntityUnitLevel>,
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub relation_contract: Option<RelationContract>,
99    #[serde(default, skip_serializing_if = "is_false")]
100    pub allows_broadcast: bool,
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub missingness_policy: Option<MissingnessPolicy>,
103    #[serde(default)]
104    pub requires_oof: bool,
105    #[serde(default)]
106    pub requires_fold_alignment: bool,
107    #[serde(default = "default_true")]
108    pub propagates_lineage: bool,
109}
110
111#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
112pub struct RelationContract {
113    #[serde(default, skip_serializing_if = "Option::is_none")]
114    pub relation_fingerprint: Option<String>,
115    #[serde(default, skip_serializing_if = "is_false")]
116    pub required: bool,
117}
118
119#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
120#[serde(rename_all = "snake_case")]
121pub enum MissingnessPolicy {
122    Strict,
123    Warn,
124    ImputeDeclared,
125    Mask,
126    PartialModel,
127    PadRepresentation,
128}
129
130fn default_true() -> bool {
131    true
132}
133
134fn is_false(value: &bool) -> bool {
135    !*value
136}
137
138impl EdgeContract {
139    pub fn new(kind: PortKind, representation: Option<String>) -> Self {
140        Self {
141            kind,
142            representation,
143            unit_level: None,
144            alignment_key: None,
145            target_level: None,
146            relation_contract: None,
147            allows_broadcast: false,
148            missingness_policy: None,
149            requires_oof: false,
150            requires_fold_alignment: false,
151            propagates_lineage: true,
152        }
153    }
154}
155
156#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
157pub struct EdgeSpec {
158    pub source: PortRef,
159    pub target: PortRef,
160    pub contract: EdgeContract,
161}
162
163#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
164pub struct GraphInterface {
165    #[serde(default)]
166    pub inputs: Vec<PortSpec>,
167    #[serde(default)]
168    pub outputs: Vec<PortSpec>,
169}
170
171#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
172pub struct NodeSpec {
173    pub id: NodeId,
174    pub kind: NodeKind,
175    pub operator: Option<serde_json::Value>,
176    #[serde(default)]
177    pub params: BTreeMap<String, serde_json::Value>,
178    #[serde(default)]
179    pub ports: PortSchema,
180    #[serde(default)]
181    pub metadata: BTreeMap<String, serde_json::Value>,
182    #[serde(default)]
183    pub seed_label: Option<String>,
184}
185
186#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
187pub struct GraphSpec {
188    pub id: String,
189    #[serde(default)]
190    pub interface: GraphInterface,
191    #[serde(default)]
192    pub nodes: Vec<NodeSpec>,
193    #[serde(default)]
194    pub edges: Vec<EdgeSpec>,
195    #[serde(default)]
196    pub search_space_fingerprint: Option<String>,
197    #[serde(default)]
198    pub metadata: BTreeMap<String, serde_json::Value>,
199}
200
201impl GraphSpec {
202    pub fn validate(&self) -> Result<()> {
203        if self.id.trim().is_empty() {
204            return Err(DagMlError::GraphValidation(
205                "graph id must not be empty".to_string(),
206            ));
207        }
208        if self.nodes.is_empty() {
209            return Err(DagMlError::GraphValidation(
210                "graph must contain at least one node".to_string(),
211            ));
212        }
213        if let Some(fingerprint) = &self.search_space_fingerprint {
214            if fingerprint.trim().is_empty() {
215                return Err(DagMlError::GraphValidation(format!(
216                    "graph `{}` has empty search_space_fingerprint",
217                    self.id
218                )));
219            }
220        }
221
222        let mut nodes = BTreeMap::new();
223        validate_unique_ports(
224            &NodeId::new("graph:interface").expect("static identifier is valid"),
225            "interface input",
226            &self.interface.inputs,
227        )?;
228        validate_unique_ports(
229            &NodeId::new("graph:interface").expect("static identifier is valid"),
230            "interface output",
231            &self.interface.outputs,
232        )?;
233        for node in &self.nodes {
234            if nodes.insert(node.id.clone(), node).is_some() {
235                return Err(DagMlError::GraphValidation(format!(
236                    "duplicate node id `{}`",
237                    node.id
238                )));
239            }
240            crate::oof::StackingOofRefitContract::from_metadata(&node.metadata).map_err(
241                |error| {
242                    DagMlError::GraphValidation(format!(
243                        "node `{}` carries invalid `{}` metadata: {}",
244                        node.id,
245                        crate::oof::STACKING_OOF_REFIT_CONTRACT_METADATA_KEY,
246                        error
247                    ))
248                },
249            )?;
250            validate_unique_ports(&node.id, "input", &node.ports.inputs)?;
251            validate_unique_ports(&node.id, "output", &node.ports.outputs)?;
252        }
253
254        let mut adjacency: BTreeMap<NodeId, Vec<NodeId>> = nodes
255            .keys()
256            .cloned()
257            .map(|id| (id, Vec::new()))
258            .collect::<BTreeMap<_, _>>();
259        let mut indegree: BTreeMap<NodeId, usize> =
260            nodes.keys().cloned().map(|id| (id, 0)).collect();
261
262        for edge in &self.edges {
263            let source = nodes.get(&edge.source.node_id).ok_or_else(|| {
264                DagMlError::GraphValidation(format!(
265                    "edge source node `{}` does not exist",
266                    edge.source.node_id
267                ))
268            })?;
269            let target = nodes.get(&edge.target.node_id).ok_or_else(|| {
270                DagMlError::GraphValidation(format!(
271                    "edge target node `{}` does not exist",
272                    edge.target.node_id
273                ))
274            })?;
275
276            let source_port =
277                find_port(&source.ports.outputs, &edge.source.port_name).ok_or_else(|| {
278                    DagMlError::GraphValidation(format!(
279                        "source port `{}.{}` does not exist",
280                        edge.source.node_id, edge.source.port_name
281                    ))
282                })?;
283            let target_port =
284                find_port(&target.ports.inputs, &edge.target.port_name).ok_or_else(|| {
285                    DagMlError::GraphValidation(format!(
286                        "target port `{}.{}` does not exist",
287                        edge.target.node_id, edge.target.port_name
288                    ))
289                })?;
290
291            if source_port.kind != edge.contract.kind || target_port.kind != edge.contract.kind {
292                return Err(DagMlError::GraphValidation(format!(
293                    "edge `{}.{}` -> `{}.{}` has kind {:?}, but ports are {:?} and {:?}",
294                    edge.source.node_id,
295                    edge.source.port_name,
296                    edge.target.node_id,
297                    edge.target.port_name,
298                    edge.contract.kind,
299                    source_port.kind,
300                    target_port.kind
301                )));
302            }
303            validate_edge_contract(edge, source_port, target_port)?;
304            if edge.contract.requires_oof && edge.contract.kind != PortKind::Prediction {
305                return Err(DagMlError::GraphValidation(format!(
306                    "edge `{}.{}` -> `{}.{}` requires OOF but is not a prediction edge",
307                    edge.source.node_id,
308                    edge.source.port_name,
309                    edge.target.node_id,
310                    edge.target.port_name
311                )));
312            }
313
314            adjacency
315                .get_mut(&edge.source.node_id)
316                .expect("source exists")
317                .push(edge.target.node_id.clone());
318            *indegree
319                .get_mut(&edge.target.node_id)
320                .expect("target exists") += 1;
321        }
322
323        ensure_acyclic(adjacency, indegree)
324    }
325
326    pub fn topological_order(&self) -> Result<Vec<NodeId>> {
327        self.validate()?;
328        let nodes = self
329            .nodes
330            .iter()
331            .map(|node| node.id.clone())
332            .collect::<BTreeSet<_>>();
333        let mut adjacency = nodes
334            .iter()
335            .cloned()
336            .map(|id| (id, Vec::new()))
337            .collect::<BTreeMap<_, _>>();
338        let mut indegree: BTreeMap<NodeId, usize> =
339            nodes.iter().cloned().map(|id| (id, 0usize)).collect();
340        for edge in &self.edges {
341            adjacency
342                .get_mut(&edge.source.node_id)
343                .expect("source exists after validate")
344                .push(edge.target.node_id.clone());
345            *indegree
346                .get_mut(&edge.target.node_id)
347                .expect("target exists after validate") += 1;
348        }
349        topological_order(adjacency, indegree)
350    }
351
352    pub fn parallel_levels(&self) -> Result<Vec<Vec<NodeId>>> {
353        self.validate()?;
354        let nodes = self
355            .nodes
356            .iter()
357            .map(|node| node.id.clone())
358            .collect::<BTreeSet<_>>();
359        let mut adjacency = nodes
360            .iter()
361            .cloned()
362            .map(|id| (id, Vec::new()))
363            .collect::<BTreeMap<_, _>>();
364        let mut indegree: BTreeMap<NodeId, usize> =
365            nodes.iter().cloned().map(|id| (id, 0usize)).collect();
366        for edge in &self.edges {
367            adjacency
368                .get_mut(&edge.source.node_id)
369                .expect("source exists after validate")
370                .push(edge.target.node_id.clone());
371            *indegree
372                .get_mut(&edge.target.node_id)
373                .expect("target exists after validate") += 1;
374        }
375        topological_levels(adjacency, indegree)
376    }
377
378    pub fn upstream_nodes(&self, node_id: &NodeId) -> Vec<NodeId> {
379        let mut upstream = self
380            .edges
381            .iter()
382            .filter_map(|edge| {
383                (edge.target.node_id == *node_id).then_some(edge.source.node_id.clone())
384            })
385            .collect::<Vec<_>>();
386        upstream.sort();
387        upstream.dedup();
388        upstream
389    }
390
391    pub fn downstream_nodes(&self, node_id: &NodeId) -> Vec<NodeId> {
392        let mut downstream = self
393            .edges
394            .iter()
395            .filter_map(|edge| {
396                (edge.source.node_id == *node_id).then_some(edge.target.node_id.clone())
397            })
398            .collect::<Vec<_>>();
399        downstream.sort();
400        downstream.dedup();
401        downstream
402    }
403}
404
405fn validate_unique_ports(node_id: &NodeId, direction: &str, ports: &[PortSpec]) -> Result<()> {
406    let mut seen = BTreeSet::new();
407    for port in ports {
408        if port.name.trim().is_empty() {
409            return Err(DagMlError::GraphValidation(format!(
410                "{} port on node `{}` has an empty name",
411                direction, node_id
412            )));
413        }
414        if !seen.insert(port.name.as_str()) {
415            return Err(DagMlError::GraphValidation(format!(
416                "duplicate {} port `{}` on node `{}`",
417                direction, port.name, node_id
418            )));
419        }
420        validate_port_contract(node_id, direction, port)?;
421    }
422    Ok(())
423}
424
425fn find_port<'a>(ports: &'a [PortSpec], name: &str) -> Option<&'a PortSpec> {
426    ports.iter().find(|port| port.name == name)
427}
428
429fn validate_port_contract(node_id: &NodeId, direction: &str, port: &PortSpec) -> Result<()> {
430    validate_optional_non_empty(
431        &format!("{direction} port `{}` representation", port.name),
432        port.representation.as_deref(),
433    )?;
434    validate_optional_non_empty(
435        &format!("{direction} port `{}` alignment_key", port.name),
436        port.alignment_key.as_deref(),
437    )?;
438    if port
439        .alignment_key
440        .as_deref()
441        .is_some_and(|key| !is_identifier(key))
442    {
443        return Err(DagMlError::GraphValidation(format!(
444            "{direction} port `{}` on node `{node_id}` has invalid alignment_key",
445            port.name
446        )));
447    }
448    Ok(())
449}
450
451fn validate_edge_contract(
452    edge: &EdgeSpec,
453    source_port: &PortSpec,
454    target_port: &PortSpec,
455) -> Result<()> {
456    let label = format!(
457        "edge `{}.{}` -> `{}.{}`",
458        edge.source.node_id, edge.source.port_name, edge.target.node_id, edge.target.port_name
459    );
460    validate_optional_non_empty(
461        &format!("{label} representation"),
462        edge.contract.representation.as_deref(),
463    )?;
464    validate_optional_non_empty(
465        &format!("{label} alignment_key"),
466        edge.contract.alignment_key.as_deref(),
467    )?;
468    if edge
469        .contract
470        .alignment_key
471        .as_deref()
472        .is_some_and(|key| !is_identifier(key))
473    {
474        return Err(DagMlError::GraphValidation(format!(
475            "{label} has invalid alignment_key"
476        )));
477    }
478    if let Some(relation_contract) = &edge.contract.relation_contract {
479        validate_relation_contract(&label, relation_contract)?;
480    }
481
482    validate_edge_unit_alignment(&label, edge, source_port, target_port)?;
483
484    if relation_aware_edge(edge, source_port, target_port) {
485        let relation_fingerprint = edge
486            .contract
487            .relation_contract
488            .as_ref()
489            .and_then(|contract| contract.relation_fingerprint.as_deref());
490        if relation_fingerprint.is_none() {
491            return Err(DagMlError::GraphValidation(format!(
492                "{label} is relation-aware but has no relation_fingerprint"
493            )));
494        }
495        if !has_effective_unit_level(edge, source_port, target_port) {
496            return Err(DagMlError::GraphValidation(format!(
497                "{label} is relation-aware but has no unit_level metadata"
498            )));
499        }
500        if !has_effective_alignment_key(edge, source_port, target_port) {
501            return Err(DagMlError::GraphValidation(format!(
502                "{label} is relation-aware but has no alignment_key"
503            )));
504        }
505    }
506    Ok(())
507}
508
509fn validate_relation_contract(label: &str, contract: &RelationContract) -> Result<()> {
510    if let Some(fingerprint) = &contract.relation_fingerprint {
511        validate_sha256(label, "relation_fingerprint", fingerprint)?;
512    } else if contract.required {
513        return Err(DagMlError::GraphValidation(format!(
514            "{label} relation_contract is required but has no relation_fingerprint"
515        )));
516    }
517    Ok(())
518}
519
520fn validate_edge_unit_alignment(
521    label: &str,
522    edge: &EdgeSpec,
523    source_port: &PortSpec,
524    target_port: &PortSpec,
525) -> Result<()> {
526    if let Some(contract_unit) = edge.contract.unit_level {
527        for (endpoint, unit) in [
528            ("source", source_port.unit_level),
529            ("target", target_port.unit_level),
530        ] {
531            if let Some(unit) = unit {
532                if unit != contract_unit && !edge.contract.allows_broadcast {
533                    return Err(DagMlError::GraphValidation(format!(
534                        "{label} {endpoint} unit {:?} does not match edge unit {:?}",
535                        unit, contract_unit
536                    )));
537                }
538            }
539        }
540    }
541
542    if let (Some(source_unit), Some(target_unit)) = (source_port.unit_level, target_port.unit_level)
543    {
544        if source_unit != target_unit && !edge.contract.allows_broadcast {
545            return Err(DagMlError::GraphValidation(format!(
546                "{label} joins incompatible unit levels {:?} and {:?}",
547                source_unit, target_unit
548            )));
549        }
550    }
551
552    if let (Some(source_target), Some(target_target)) =
553        (source_port.target_level, target_port.target_level)
554    {
555        if source_target != target_target {
556            return Err(DagMlError::GraphValidation(format!(
557                "{label} joins incompatible target levels {:?} and {:?}",
558                source_target, target_target
559            )));
560        }
561    }
562    if let Some(contract_target) = edge.contract.target_level {
563        for (endpoint, target_level) in [
564            ("source", source_port.target_level),
565            ("target", target_port.target_level),
566        ] {
567            if let Some(target_level) = target_level {
568                if target_level != contract_target {
569                    return Err(DagMlError::GraphValidation(format!(
570                        "{label} {endpoint} target level {:?} does not match edge target_level {:?}",
571                        target_level, contract_target
572                    )));
573                }
574            }
575        }
576    }
577
578    if let (Some(source_alignment), Some(target_alignment)) = (
579        source_port.alignment_key.as_deref(),
580        target_port.alignment_key.as_deref(),
581    ) {
582        if source_alignment != target_alignment && !edge.contract.allows_broadcast {
583            return Err(DagMlError::GraphValidation(format!(
584                "{label} joins incompatible alignment keys `{source_alignment}` and `{target_alignment}`"
585            )));
586        }
587    }
588
589    if let Some(edge_alignment) = edge.contract.alignment_key.as_deref() {
590        for (endpoint, alignment) in [
591            ("source", source_port.alignment_key.as_deref()),
592            ("target", target_port.alignment_key.as_deref()),
593        ] {
594            if let Some(alignment) = alignment {
595                if alignment != edge_alignment && !edge.contract.allows_broadcast {
596                    return Err(DagMlError::GraphValidation(format!(
597                        "{label} {endpoint} alignment `{alignment}` does not match edge alignment `{edge_alignment}`"
598                    )));
599                }
600            }
601        }
602    }
603
604    if edge.contract.allows_broadcast
605        && edge.contract.alignment_key.is_none()
606        && source_port.alignment_key.is_none()
607        && target_port.alignment_key.is_none()
608    {
609        return Err(DagMlError::GraphValidation(format!(
610            "{label} allows broadcast but declares no alignment_key"
611        )));
612    }
613    Ok(())
614}
615
616fn relation_aware_edge(edge: &EdgeSpec, source_port: &PortSpec, target_port: &PortSpec) -> bool {
617    edge.contract.relation_contract.is_some()
618        || edge.contract.allows_broadcast
619        || edge.contract.alignment_key.is_some()
620        || non_physical(edge.contract.unit_level)
621        || non_physical(edge.contract.target_level)
622        || non_physical(source_port.unit_level)
623        || non_physical(source_port.target_level)
624        || non_physical(target_port.unit_level)
625        || non_physical(target_port.target_level)
626        || source_port.alignment_key.is_some()
627        || target_port.alignment_key.is_some()
628}
629
630fn has_effective_unit_level(
631    edge: &EdgeSpec,
632    source_port: &PortSpec,
633    target_port: &PortSpec,
634) -> bool {
635    edge.contract.unit_level.is_some()
636        || source_port.unit_level.is_some()
637        || target_port.unit_level.is_some()
638}
639
640fn has_effective_alignment_key(
641    edge: &EdgeSpec,
642    source_port: &PortSpec,
643    target_port: &PortSpec,
644) -> bool {
645    edge.contract.alignment_key.is_some()
646        || source_port.alignment_key.is_some()
647        || target_port.alignment_key.is_some()
648}
649
650fn non_physical(unit_level: Option<EntityUnitLevel>) -> bool {
651    unit_level.is_some_and(|level| level != EntityUnitLevel::PhysicalSample)
652}
653
654fn validate_optional_non_empty(label: &str, value: Option<&str>) -> Result<()> {
655    if value.is_some_and(|value| value.trim().is_empty()) {
656        return Err(DagMlError::GraphValidation(format!(
657            "{label} must not be empty"
658        )));
659    }
660    Ok(())
661}
662
663fn validate_sha256(owner: &str, field: &str, value: &str) -> Result<()> {
664    if value.len() == 64 && value.bytes().all(|byte| byte.is_ascii_hexdigit()) {
665        Ok(())
666    } else {
667        Err(DagMlError::GraphValidation(format!(
668            "{owner} has invalid {field}"
669        )))
670    }
671}
672
673fn is_identifier(value: &str) -> bool {
674    !value.is_empty()
675        && value.len() <= 128
676        && value
677            .bytes()
678            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'_' | b'-' | b'.' | b':'))
679}
680
681fn ensure_acyclic(
682    adjacency: BTreeMap<NodeId, Vec<NodeId>>,
683    indegree: BTreeMap<NodeId, usize>,
684) -> Result<()> {
685    topological_order(adjacency, indegree).map(|_| ())
686}
687
688fn topological_order(
689    adjacency: BTreeMap<NodeId, Vec<NodeId>>,
690    mut indegree: BTreeMap<NodeId, usize>,
691) -> Result<Vec<NodeId>> {
692    let mut queue = indegree
693        .iter()
694        .filter_map(|(id, degree)| (*degree == 0).then_some(id.clone()))
695        .collect::<BTreeSet<_>>();
696    let mut order = Vec::with_capacity(indegree.len());
697
698    while let Some(node) = queue.pop_first() {
699        order.push(node.clone());
700        if let Some(next_nodes) = adjacency.get(&node) {
701            for next in next_nodes {
702                let degree = indegree.get_mut(next).expect("node exists");
703                *degree -= 1;
704                if *degree == 0 {
705                    queue.insert(next.clone());
706                }
707            }
708        }
709    }
710
711    if order.len() == indegree.len() {
712        Ok(order)
713    } else {
714        Err(DagMlError::GraphValidation(
715            "graph contains at least one cycle".to_string(),
716        ))
717    }
718}
719
720fn topological_levels(
721    adjacency: BTreeMap<NodeId, Vec<NodeId>>,
722    mut indegree: BTreeMap<NodeId, usize>,
723) -> Result<Vec<Vec<NodeId>>> {
724    let mut queue = indegree
725        .iter()
726        .filter_map(|(id, degree)| (*degree == 0).then_some(id.clone()))
727        .collect::<BTreeSet<_>>();
728    let mut levels = Vec::new();
729    let mut visited = 0usize;
730
731    while !queue.is_empty() {
732        let level = queue.iter().cloned().collect::<Vec<_>>();
733        queue.clear();
734        for node in &level {
735            visited += 1;
736            if let Some(next_nodes) = adjacency.get(node) {
737                for next in next_nodes {
738                    let degree = indegree.get_mut(next).expect("node exists");
739                    *degree -= 1;
740                    if *degree == 0 {
741                        queue.insert(next.clone());
742                    }
743                }
744            }
745        }
746        levels.push(level);
747    }
748
749    if visited == indegree.len() {
750        Ok(levels)
751    } else {
752        Err(DagMlError::GraphValidation(
753            "graph contains at least one cycle".to_string(),
754        ))
755    }
756}
757
758#[cfg(test)]
759mod tests {
760    use super::*;
761
762    fn port(name: &str, kind: PortKind) -> PortSpec {
763        PortSpec {
764            name: name.to_string(),
765            kind,
766            representation: None,
767            cardinality: PortCardinality::One,
768            unit_level: None,
769            alignment_key: None,
770            target_level: None,
771            description: String::new(),
772        }
773    }
774
775    fn node(id: &str, inputs: Vec<PortSpec>, outputs: Vec<PortSpec>) -> NodeSpec {
776        NodeSpec {
777            id: NodeId::new(id).unwrap(),
778            kind: NodeKind::Model,
779            operator: None,
780            params: BTreeMap::new(),
781            ports: PortSchema { inputs, outputs },
782            metadata: BTreeMap::new(),
783            seed_label: None,
784        }
785    }
786
787    fn edge(source: &str, source_port: &str, target: &str, target_port: &str) -> EdgeSpec {
788        EdgeSpec {
789            source: PortRef {
790                node_id: NodeId::new(source).unwrap(),
791                port_name: source_port.to_string(),
792            },
793            target: PortRef {
794                node_id: NodeId::new(target).unwrap(),
795                port_name: target_port.to_string(),
796            },
797            contract: EdgeContract {
798                requires_oof: true,
799                requires_fold_alignment: true,
800                ..EdgeContract::new(PortKind::Prediction, None)
801            },
802        }
803    }
804
805    #[test]
806    fn validates_simple_graph() {
807        let graph = GraphSpec {
808            id: "g".to_string(),
809            interface: GraphInterface::default(),
810            nodes: vec![
811                node("model:a", vec![], vec![port("pred", PortKind::Prediction)]),
812                node("model:b", vec![port("pred", PortKind::Prediction)], vec![]),
813            ],
814            edges: vec![edge("model:a", "pred", "model:b", "pred")],
815            search_space_fingerprint: None,
816            metadata: BTreeMap::new(),
817        };
818
819        assert!(graph.validate().is_ok());
820    }
821
822    #[test]
823    fn computes_deterministic_parallel_levels() {
824        let graph = GraphSpec {
825            id: "g".to_string(),
826            interface: GraphInterface::default(),
827            nodes: vec![
828                node("model:a", vec![], vec![port("pred", PortKind::Prediction)]),
829                node(
830                    "model:b",
831                    vec![port("pred", PortKind::Prediction)],
832                    vec![port("pred", PortKind::Prediction)],
833                ),
834                node(
835                    "model:c",
836                    vec![port("pred", PortKind::Prediction)],
837                    vec![port("pred", PortKind::Prediction)],
838                ),
839                node("model:d", vec![port("pred", PortKind::Prediction)], vec![]),
840            ],
841            edges: vec![
842                edge("model:a", "pred", "model:b", "pred"),
843                edge("model:a", "pred", "model:c", "pred"),
844                edge("model:b", "pred", "model:d", "pred"),
845                edge("model:c", "pred", "model:d", "pred"),
846            ],
847            search_space_fingerprint: None,
848            metadata: BTreeMap::new(),
849        };
850
851        let levels = graph.parallel_levels().unwrap();
852
853        assert_eq!(
854            levels,
855            vec![
856                vec![NodeId::new("model:a").unwrap()],
857                vec![
858                    NodeId::new("model:b").unwrap(),
859                    NodeId::new("model:c").unwrap()
860                ],
861                vec![NodeId::new("model:d").unwrap()]
862            ]
863        );
864    }
865
866    #[test]
867    fn rejects_missing_edge_endpoint() {
868        let graph = GraphSpec {
869            id: "g".to_string(),
870            interface: GraphInterface::default(),
871            nodes: vec![node(
872                "model:a",
873                vec![],
874                vec![port("pred", PortKind::Prediction)],
875            )],
876            edges: vec![edge("model:a", "pred", "model:b", "pred")],
877            search_space_fingerprint: None,
878            metadata: BTreeMap::new(),
879        };
880
881        assert!(graph.validate().is_err());
882    }
883
884    #[test]
885    fn rejects_oof_contract_on_non_prediction_edge() {
886        let graph = GraphSpec {
887            id: "g".to_string(),
888            interface: GraphInterface::default(),
889            nodes: vec![
890                node("model:a", vec![], vec![port("x", PortKind::Data)]),
891                node("model:b", vec![port("x", PortKind::Data)], vec![]),
892            ],
893            edges: vec![EdgeSpec {
894                source: PortRef {
895                    node_id: NodeId::new("model:a").unwrap(),
896                    port_name: "x".to_string(),
897                },
898                target: PortRef {
899                    node_id: NodeId::new("model:b").unwrap(),
900                    port_name: "x".to_string(),
901                },
902                contract: EdgeContract {
903                    requires_oof: true,
904                    requires_fold_alignment: true,
905                    ..EdgeContract::new(PortKind::Data, None)
906                },
907            }],
908            search_space_fingerprint: None,
909            metadata: BTreeMap::new(),
910        };
911
912        let error = graph.validate().unwrap_err().to_string();
913
914        assert!(error.contains("requires OOF"));
915    }
916
917    fn unit_port(name: &str, kind: PortKind, unit_level: EntityUnitLevel) -> PortSpec {
918        let mut port = port(name, kind);
919        port.unit_level = Some(unit_level);
920        port.alignment_key = Some("sample_id".to_string());
921        port
922    }
923
924    fn data_edge_contract() -> EdgeContract {
925        EdgeContract::new(PortKind::Data, Some("tabular".to_string()))
926    }
927
928    fn relation_contract() -> RelationContract {
929        RelationContract {
930            relation_fingerprint: Some("a".repeat(64)),
931            required: true,
932        }
933    }
934
935    #[test]
936    fn rejects_unit_mismatch_without_explicit_broadcast() {
937        let graph = GraphSpec {
938            id: "g".to_string(),
939            interface: GraphInterface::default(),
940            nodes: vec![
941                node(
942                    "transform:obs",
943                    vec![],
944                    vec![unit_port("x", PortKind::Data, EntityUnitLevel::Observation)],
945                ),
946                node(
947                    "join:sample",
948                    vec![unit_port(
949                        "x",
950                        PortKind::Data,
951                        EntityUnitLevel::PhysicalSample,
952                    )],
953                    vec![],
954                ),
955            ],
956            edges: vec![EdgeSpec {
957                source: PortRef {
958                    node_id: NodeId::new("transform:obs").unwrap(),
959                    port_name: "x".to_string(),
960                },
961                target: PortRef {
962                    node_id: NodeId::new("join:sample").unwrap(),
963                    port_name: "x".to_string(),
964                },
965                contract: EdgeContract {
966                    relation_contract: Some(relation_contract()),
967                    ..data_edge_contract()
968                },
969            }],
970            search_space_fingerprint: None,
971            metadata: BTreeMap::new(),
972        };
973
974        let error = graph.validate().unwrap_err().to_string();
975
976        assert!(error.contains("incompatible unit levels"));
977    }
978
979    #[test]
980    fn relation_aware_edge_requires_relation_fingerprint() {
981        let graph = GraphSpec {
982            id: "g".to_string(),
983            interface: GraphInterface::default(),
984            nodes: vec![
985                node(
986                    "source:a",
987                    vec![],
988                    vec![unit_port("x", PortKind::Data, EntityUnitLevel::Observation)],
989                ),
990                node(
991                    "model:a",
992                    vec![unit_port("x", PortKind::Data, EntityUnitLevel::Observation)],
993                    vec![],
994                ),
995            ],
996            edges: vec![EdgeSpec {
997                source: PortRef {
998                    node_id: NodeId::new("source:a").unwrap(),
999                    port_name: "x".to_string(),
1000                },
1001                target: PortRef {
1002                    node_id: NodeId::new("model:a").unwrap(),
1003                    port_name: "x".to_string(),
1004                },
1005                contract: data_edge_contract(),
1006            }],
1007            search_space_fingerprint: None,
1008            metadata: BTreeMap::new(),
1009        };
1010
1011        let error = graph.validate().unwrap_err().to_string();
1012
1013        assert!(error.contains("relation-aware"));
1014    }
1015
1016    #[test]
1017    fn relation_aware_edge_requires_alignment_key() {
1018        let mut source_port = port("x", PortKind::Data);
1019        source_port.unit_level = Some(EntityUnitLevel::Observation);
1020        let mut target_port = port("x", PortKind::Data);
1021        target_port.unit_level = Some(EntityUnitLevel::Observation);
1022
1023        let graph = GraphSpec {
1024            id: "g".to_string(),
1025            interface: GraphInterface::default(),
1026            nodes: vec![
1027                node("source:a", vec![], vec![source_port]),
1028                node("model:a", vec![target_port], vec![]),
1029            ],
1030            edges: vec![EdgeSpec {
1031                source: PortRef {
1032                    node_id: NodeId::new("source:a").unwrap(),
1033                    port_name: "x".to_string(),
1034                },
1035                target: PortRef {
1036                    node_id: NodeId::new("model:a").unwrap(),
1037                    port_name: "x".to_string(),
1038                },
1039                contract: EdgeContract {
1040                    relation_contract: Some(relation_contract()),
1041                    ..data_edge_contract()
1042                },
1043            }],
1044            search_space_fingerprint: None,
1045            metadata: BTreeMap::new(),
1046        };
1047
1048        let error = graph.validate().unwrap_err().to_string();
1049
1050        assert!(error.contains("alignment_key"));
1051    }
1052
1053    #[test]
1054    fn explicit_broadcast_allows_sample_to_observation_edge() {
1055        let mut contract = data_edge_contract();
1056        contract.allows_broadcast = true;
1057        contract.alignment_key = Some("sample_id".to_string());
1058        contract.relation_contract = Some(relation_contract());
1059
1060        let graph = GraphSpec {
1061            id: "g".to_string(),
1062            interface: GraphInterface::default(),
1063            nodes: vec![
1064                node(
1065                    "source:sample",
1066                    vec![],
1067                    vec![unit_port(
1068                        "x",
1069                        PortKind::Data,
1070                        EntityUnitLevel::PhysicalSample,
1071                    )],
1072                ),
1073                node(
1074                    "adapter:broadcast",
1075                    vec![unit_port("x", PortKind::Data, EntityUnitLevel::Observation)],
1076                    vec![],
1077                ),
1078            ],
1079            edges: vec![EdgeSpec {
1080                source: PortRef {
1081                    node_id: NodeId::new("source:sample").unwrap(),
1082                    port_name: "x".to_string(),
1083                },
1084                target: PortRef {
1085                    node_id: NodeId::new("adapter:broadcast").unwrap(),
1086                    port_name: "x".to_string(),
1087                },
1088                contract,
1089            }],
1090            search_space_fingerprint: None,
1091            metadata: BTreeMap::new(),
1092        };
1093
1094        graph.validate().unwrap();
1095    }
1096
1097    #[test]
1098    fn rejects_cycles() {
1099        let graph = GraphSpec {
1100            id: "g".to_string(),
1101            interface: GraphInterface::default(),
1102            nodes: vec![
1103                node(
1104                    "model:a",
1105                    vec![port("pred", PortKind::Prediction)],
1106                    vec![port("pred", PortKind::Prediction)],
1107                ),
1108                node(
1109                    "model:b",
1110                    vec![port("pred", PortKind::Prediction)],
1111                    vec![port("pred", PortKind::Prediction)],
1112                ),
1113            ],
1114            edges: vec![
1115                edge("model:a", "pred", "model:b", "pred"),
1116                edge("model:b", "pred", "model:a", "pred"),
1117            ],
1118            search_space_fingerprint: None,
1119            metadata: BTreeMap::new(),
1120        };
1121
1122        assert!(graph.validate().is_err());
1123    }
1124
1125    #[test]
1126    fn published_graph_spec_schema_declares_current_contract() {
1127        let schema: serde_json::Value = serde_json::from_str(include_str!(
1128            "../../../docs/contracts/graph_spec.schema.json"
1129        ))
1130        .unwrap();
1131
1132        assert_eq!(schema["$id"], GRAPH_SPEC_SCHEMA_ID);
1133        assert!(schema["required"]
1134            .as_array()
1135            .unwrap()
1136            .iter()
1137            .any(|field| field.as_str() == Some("nodes")));
1138        assert_eq!(
1139            schema["$defs"]["node_kind"]["enum"]
1140                .as_array()
1141                .unwrap()
1142                .len(),
1143            20
1144        );
1145        assert!(schema["$defs"]["port_kind"]["enum"]
1146            .as_array()
1147            .unwrap()
1148            .iter()
1149            .any(|kind| kind.as_str() == Some("prediction")));
1150        assert!(schema["$defs"]["entity_unit_level"]["enum"]
1151            .as_array()
1152            .unwrap()
1153            .iter()
1154            .any(|level| level.as_str() == Some("combo")));
1155        assert!(schema["$defs"]["edge_contract"]["properties"]
1156            .as_object()
1157            .unwrap()
1158            .contains_key("relation_contract"));
1159    }
1160}
dag_ml_core/graph.rs

dag_ml_core/
graph.rs
