dag_ml_core/
graph.rs

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

dag_ml_core/
graph.rs
