dag_ml_core/runtime/
artifact.rs

1// Auto-split from the former monolithic `runtime.rs` (pure refactor).
2use super::*;
3
4#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum HandleKind {
7    Data,
8    DataView,
9    Model,
10    Artifact,
11    Prediction,
12    Relation,
13}
14
15#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
16#[serde(deny_unknown_fields)]
17pub struct HandleRef {
18    pub handle: u64,
19    pub kind: HandleKind,
20    pub owner_controller: ControllerId,
21}
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum ArtifactBackend {
26    Joblib,
27    Torch,
28    Tensorflow,
29    Onnx,
30    Safetensors,
31    Json,
32    Raw,
33}
34
35#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
36pub struct ArtifactRef {
37    pub id: ArtifactId,
38    pub kind: String,
39    pub controller_id: ControllerId,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    pub backend: Option<ArtifactBackend>,
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub uri: Option<String>,
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub content_fingerprint: Option<String>,
46    pub size_bytes: Option<u64>,
47    #[serde(default, skip_serializing_if = "Option::is_none")]
48    pub plugin: Option<String>,
49    #[serde(default, skip_serializing_if = "Option::is_none")]
50    pub plugin_version: Option<String>,
51}
52
53impl ArtifactRef {
54    pub fn validate(&self) -> Result<()> {
55        if self.kind.trim().is_empty() {
56            return Err(DagMlError::RuntimeValidation(format!(
57                "artifact `{}` has empty kind",
58                self.id
59            )));
60        }
61        validate_artifact_optional_text("uri", &self.uri, &self.id)?;
62        validate_artifact_optional_text("plugin", &self.plugin, &self.id)?;
63        validate_artifact_optional_text("plugin_version", &self.plugin_version, &self.id)?;
64        if self.plugin_version.is_some() && self.plugin.is_none() {
65            return Err(DagMlError::RuntimeValidation(format!(
66                "artifact `{}` has plugin_version without plugin",
67                self.id
68            )));
69        }
70        if let Some(content_fingerprint) = &self.content_fingerprint {
71            validate_runtime_fingerprint("artifact content", content_fingerprint)?;
72        }
73        if self.uri.is_some() && self.backend.is_none() {
74            return Err(DagMlError::RuntimeValidation(format!(
75                "artifact `{}` has uri without backend",
76                self.id
77            )));
78        }
79        if self.uri.is_some() && self.content_fingerprint.is_none() {
80            return Err(DagMlError::RuntimeValidation(format!(
81                "artifact `{}` has uri without content_fingerprint",
82                self.id
83            )));
84        }
85        Ok(())
86    }
87
88    /// Validate that the artifact carries portable metadata: a backend, a safe
89    /// relative URI and a content fingerprint. Legacy artifacts that only carry
90    /// inline metadata stay readable through [`ArtifactRef::validate`] but are
91    /// refused here so persisted manifests can be moved with their payloads.
92    pub fn validate_portable(&self) -> Result<()> {
93        self.validate()?;
94        let Some(uri) = self.uri.as_deref() else {
95            return Err(DagMlError::RuntimeValidation(format!(
96                "artifact `{}` is not portable: requires backend, uri and content_fingerprint",
97                self.id
98            )));
99        };
100        // `validate` already guarantees that a present URI implies a backend and
101        // a 64-hex content fingerprint, so confirming the URI is enough here.
102        validate_relative_artifact_uri(&self.id, uri)
103    }
104}
105
106pub fn refit_artifact_input_key(artifact_id: &ArtifactId) -> String {
107    format!("artifact:{artifact_id}")
108}
109
110#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
111pub struct ArtifactMaterializationRequest {
112    pub run_id: RunId,
113    pub bundle_id: BundleId,
114    pub node_id: NodeId,
115    pub phase: Phase,
116    pub variant_id: Option<VariantId>,
117    pub controller_id: ControllerId,
118    pub artifact: ArtifactRef,
119    pub params_fingerprint: String,
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub training_loss_fingerprint: Option<String>,
122}
123
124#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
125pub struct ArtifactHandleRecord {
126    pub handle: HandleRef,
127    pub node_id: NodeId,
128    pub controller_id: ControllerId,
129    pub artifact: ArtifactRef,
130    pub params_fingerprint: String,
131    #[serde(default, skip_serializing_if = "Option::is_none")]
132    pub training_loss_fingerprint: Option<String>,
133}
134
135impl ArtifactHandleRecord {
136    pub fn validate(&self) -> Result<()> {
137        self.artifact.validate()?;
138        if !matches!(self.handle.kind, HandleKind::Model | HandleKind::Artifact) {
139            return Err(DagMlError::RuntimeValidation(format!(
140                "artifact `{}` is registered with non-artifact/model handle kind {:?}",
141                self.artifact.id, self.handle.kind
142            )));
143        }
144        if self.handle.owner_controller != self.controller_id {
145            return Err(DagMlError::RuntimeValidation(format!(
146                "artifact `{}` handle owner `{}` does not match controller `{}`",
147                self.artifact.id, self.handle.owner_controller, self.controller_id
148            )));
149        }
150        if self.artifact.controller_id != self.controller_id {
151            return Err(DagMlError::RuntimeValidation(format!(
152                "artifact `{}` controller `{}` does not match record controller `{}`",
153                self.artifact.id, self.artifact.controller_id, self.controller_id
154            )));
155        }
156        if self.params_fingerprint.trim().is_empty() {
157            return Err(DagMlError::RuntimeValidation(format!(
158                "artifact `{}` has empty params fingerprint",
159                self.artifact.id
160            )));
161        }
162        if let Some(fingerprint) = &self.training_loss_fingerprint {
163            validate_runtime_fingerprint("artifact training loss", fingerprint)?;
164        }
165        Ok(())
166    }
167}
168
169pub trait RuntimeArtifactStore {
170    fn materialize(&self, request: &ArtifactMaterializationRequest) -> Result<HandleRef>;
171}
172
173#[derive(Clone, Debug, Default)]
174pub struct InMemoryArtifactStore {
175    records: BTreeMap<ArtifactId, ArtifactHandleRecord>,
176    refit_artifacts: BTreeMap<ArtifactId, RefitArtifactRecord>,
177}
178
179impl InMemoryArtifactStore {
180    pub fn new() -> Self {
181        Self::default()
182    }
183
184    pub fn register(&mut self, artifact: &RefitArtifactRecord, handle: HandleRef) -> Result<()> {
185        artifact.validate()?;
186        let record = ArtifactHandleRecord {
187            handle,
188            node_id: artifact.node_id.clone(),
189            controller_id: artifact.controller_id.clone(),
190            artifact: artifact.artifact.clone(),
191            params_fingerprint: artifact.params_fingerprint.clone(),
192            training_loss_fingerprint: artifact.training_loss_fingerprint.clone(),
193        };
194        record.validate()?;
195        if self.records.contains_key(&record.artifact.id)
196            || self.refit_artifacts.contains_key(&record.artifact.id)
197        {
198            return Err(DagMlError::RuntimeValidation(format!(
199                "duplicate artifact handle for `{}`",
200                artifact.artifact.id
201            )));
202        }
203        let previous_record = self.records.insert(record.artifact.id.clone(), record);
204        debug_assert!(previous_record.is_none());
205        let previous_artifact = self
206            .refit_artifacts
207            .insert(artifact.artifact.id.clone(), artifact.clone());
208        debug_assert!(previous_artifact.is_none());
209        Ok(())
210    }
211
212    pub fn capture_refit_artifacts(
213        &mut self,
214        task: &NodeTask,
215        result: &NodeResult,
216    ) -> Result<Vec<RefitArtifactRecord>> {
217        if task.phase != Phase::Refit {
218            return Err(DagMlError::RuntimeValidation(format!(
219                "cannot capture refit artifacts from phase {:?}",
220                task.phase
221            )));
222        }
223        let mut records = Vec::new();
224        for artifact in &result.artifacts {
225            let handle = result.artifact_handles.get(&artifact.id).ok_or_else(|| {
226                DagMlError::RuntimeValidation(format!(
227                    "node `{}` emitted artifact `{}` without artifact handle",
228                    task.node_plan.node_id, artifact.id
229                ))
230            })?;
231            let record = RefitArtifactRecord {
232                node_id: task.node_plan.node_id.clone(),
233                controller_id: task.node_plan.controller_id.clone(),
234                artifact: artifact.clone(),
235                params_fingerprint: task.node_plan.params_fingerprint.clone(),
236                training_loss_fingerprint: task
237                    .node_plan
238                    .training_loss_fingerprint(Phase::Refit)?,
239                data_requirement_keys: task
240                    .node_plan
241                    .data_bindings
242                    .iter()
243                    .map(|binding| {
244                        data_binding_requirement_key(&binding.node_id, &binding.input_name)
245                    })
246                    .collect(),
247                // Only the Validation-OOF meta-feature inputs become replay
248                // prediction-cache requirements. The off-fold (REFIT/PREDICT)
249                // test/predict base predictions are recomputed each phase, not
250                // replayed from cache, and they share the same producer/port as the
251                // Validation OOF input — so including them would duplicate the
252                // requirement key. They are excluded here (partition != Validation).
253                prediction_requirement_keys: task
254                    .prediction_inputs
255                    .values()
256                    .filter(|spec| spec.partition == PredictionPartition::Validation)
257                    .map(|spec| {
258                        bundle_prediction_requirement_key(
259                            &spec.producer_node,
260                            &spec.source_port,
261                            &task.node_plan.node_id,
262                            &spec.target_port,
263                        )
264                    })
265                    .collect(),
266            };
267            self.register(&record, handle.clone())?;
268            records.push(record);
269        }
270        Ok(records)
271    }
272
273    pub fn get(&self, artifact_id: &ArtifactId) -> Option<&ArtifactHandleRecord> {
274        self.records.get(artifact_id)
275    }
276
277    pub fn len(&self) -> usize {
278        self.records.len()
279    }
280
281    pub fn is_empty(&self) -> bool {
282        self.records.is_empty()
283    }
284
285    pub fn refit_artifacts(&self) -> Vec<RefitArtifactRecord> {
286        self.refit_artifacts.values().cloned().collect()
287    }
288}
289
290impl RuntimeArtifactStore for InMemoryArtifactStore {
291    fn materialize(&self, request: &ArtifactMaterializationRequest) -> Result<HandleRef> {
292        let record = self.records.get(&request.artifact.id).ok_or_else(|| {
293            DagMlError::RuntimeValidation(format!(
294                "artifact store is missing refit artifact `{}` for bundle `{}`",
295                request.artifact.id, request.bundle_id
296            ))
297        })?;
298        if record.node_id != request.node_id {
299            return Err(DagMlError::RuntimeValidation(format!(
300                "artifact `{}` is registered for node `{}` but requested for `{}`",
301                request.artifact.id, record.node_id, request.node_id
302            )));
303        }
304        if record.controller_id != request.controller_id {
305            return Err(DagMlError::RuntimeValidation(format!(
306                "artifact `{}` is registered for controller `{}` but requested for `{}`",
307                request.artifact.id, record.controller_id, request.controller_id
308            )));
309        }
310        if record.artifact != request.artifact {
311            return Err(DagMlError::RuntimeValidation(format!(
312                "artifact `{}` metadata does not match bundle record",
313                request.artifact.id
314            )));
315        }
316        if record.params_fingerprint != request.params_fingerprint {
317            return Err(DagMlError::RuntimeValidation(format!(
318                "artifact `{}` params fingerprint does not match bundle record",
319                request.artifact.id
320            )));
321        }
322        if record.training_loss_fingerprint != request.training_loss_fingerprint {
323            return Err(DagMlError::RuntimeValidation(format!(
324                "artifact `{}` training loss fingerprint does not match bundle record",
325                request.artifact.id
326            )));
327        }
328        record.validate()?;
329        Ok(record.handle.clone())
330    }
331}
332
333pub const FILE_ARTIFACT_MANIFEST_SCHEMA_VERSION: u32 = 1;
334pub const FILE_ARTIFACT_MANIFEST_FILE: &str = "artifact_manifest.json";
335
336pub(crate) fn default_file_artifact_manifest_schema_version() -> u32 {
337    FILE_ARTIFACT_MANIFEST_SCHEMA_VERSION
338}
339
340/// One persisted artifact entry. Mirrors the bundle [`RefitArtifactRecord`]
341/// identity (node, controller, artifact, parameters and training loss) while requiring
342/// the [`ArtifactRef`] to be portable so the manifest stays movable with its
343/// payloads.
344#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
345pub struct FileArtifactManifestEntry {
346    pub node_id: NodeId,
347    pub controller_id: ControllerId,
348    pub artifact: ArtifactRef,
349    pub params_fingerprint: String,
350    #[serde(default, skip_serializing_if = "Option::is_none")]
351    pub training_loss_fingerprint: Option<String>,
352}
353
354impl FileArtifactManifestEntry {
355    fn from_refit_record(record: &RefitArtifactRecord) -> Result<Self> {
356        let entry = Self {
357            node_id: record.node_id.clone(),
358            controller_id: record.controller_id.clone(),
359            artifact: record.artifact.clone(),
360            params_fingerprint: record.params_fingerprint.clone(),
361            training_loss_fingerprint: record.training_loss_fingerprint.clone(),
362        };
363        entry.validate()?;
364        Ok(entry)
365    }
366
367    pub fn validate(&self) -> Result<()> {
368        self.artifact.validate_portable()?;
369        if self.artifact.controller_id != self.controller_id {
370            return Err(DagMlError::RuntimeValidation(format!(
371                "artifact manifest entry `{}` controller `{}` does not match artifact controller `{}`",
372                self.artifact.id, self.controller_id, self.artifact.controller_id
373            )));
374        }
375        validate_runtime_fingerprint("artifact manifest params", &self.params_fingerprint)?;
376        if let Some(fingerprint) = &self.training_loss_fingerprint {
377            validate_runtime_fingerprint("artifact manifest training loss", fingerprint)?;
378        }
379        Ok(())
380    }
381
382    fn matches_refit_record(&self, record: &RefitArtifactRecord) -> bool {
383        self.node_id == record.node_id
384            && self.controller_id == record.controller_id
385            && self.artifact == record.artifact
386            && self.params_fingerprint == record.params_fingerprint
387            && self.training_loss_fingerprint == record.training_loss_fingerprint
388    }
389}
390
391/// Versioned, file-backed artifact manifest. This is a manifest/portability
392/// layer only: it records portable [`ArtifactRef`] metadata for a bundle's
393/// refit artifacts. It does not deserialize ML objects or materialize artifact
394/// payloads; payload stores remain future work.
395#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
396pub struct FileArtifactManifest {
397    pub bundle_id: BundleId,
398    #[serde(default = "default_file_artifact_manifest_schema_version")]
399    pub schema_version: u32,
400    #[serde(default)]
401    pub artifacts: Vec<FileArtifactManifestEntry>,
402}
403
404impl FileArtifactManifest {
405    pub fn validate(&self) -> Result<()> {
406        if self.schema_version != FILE_ARTIFACT_MANIFEST_SCHEMA_VERSION {
407            return Err(DagMlError::RuntimeValidation(format!(
408                "file artifact manifest for bundle `{}` uses unsupported schema_version {}, expected {}",
409                self.bundle_id, self.schema_version, FILE_ARTIFACT_MANIFEST_SCHEMA_VERSION
410            )));
411        }
412        let mut artifact_ids = BTreeSet::new();
413        let mut uris = BTreeSet::new();
414        for entry in &self.artifacts {
415            entry.validate()?;
416            if !artifact_ids.insert(entry.artifact.id.as_str()) {
417                return Err(DagMlError::RuntimeValidation(format!(
418                    "file artifact manifest for bundle `{}` has duplicate artifact id `{}`",
419                    self.bundle_id, entry.artifact.id
420                )));
421            }
422            // `entry.validate` guarantees a portable URI is present.
423            if let Some(uri) = entry.artifact.uri.as_deref() {
424                if !uris.insert(uri) {
425                    return Err(DagMlError::RuntimeValidation(format!(
426                        "file artifact manifest for bundle `{}` has duplicate artifact uri `{}`",
427                        self.bundle_id, uri
428                    )));
429                }
430            }
431        }
432        Ok(())
433    }
434
435    pub fn validate_against_bundle(&self, bundle: &ExecutionBundle) -> Result<()> {
436        self.validate()?;
437        bundle.validate()?;
438        if self.bundle_id != bundle.bundle_id {
439            return Err(DagMlError::RuntimeValidation(format!(
440                "file artifact manifest bundle `{}` does not match bundle `{}`",
441                self.bundle_id, bundle.bundle_id
442            )));
443        }
444        if self.artifacts.len() != bundle.refit_artifacts.len() {
445            return Err(DagMlError::RuntimeValidation(format!(
446                "file artifact manifest for bundle `{}` has {} artifact(s) for {} bundle refit artifact(s)",
447                self.bundle_id,
448                self.artifacts.len(),
449                bundle.refit_artifacts.len()
450            )));
451        }
452        let entries_by_id = self
453            .artifacts
454            .iter()
455            .map(|entry| (entry.artifact.id.as_str(), entry))
456            .collect::<BTreeMap<_, _>>();
457        for record in &bundle.refit_artifacts {
458            let entry = entries_by_id
459                .get(record.artifact.id.as_str())
460                .ok_or_else(|| {
461                    DagMlError::RuntimeValidation(format!(
462                        "file artifact manifest for bundle `{}` is missing refit artifact `{}`",
463                        self.bundle_id, record.artifact.id
464                    ))
465                })?;
466            if !entry.matches_refit_record(record) {
467                return Err(DagMlError::RuntimeValidation(format!(
468                    "file artifact manifest entry `{}` does not match bundle refit artifact",
469                    entry.artifact.id
470                )));
471            }
472        }
473        Ok(())
474    }
475}
476
477/// File-backed artifact manifest store rooted at a directory.
478///
479/// This is a portability/manifest layer: [`FileArtifactManifestStore::write`]
480/// serializes portable artifact references from a validated bundle and
481/// [`FileArtifactManifestStore::open`] reloads and revalidates them against the
482/// bundle. It never reads, writes or deserializes artifact payloads.
483#[derive(Clone, Debug)]
484pub struct FileArtifactManifestStore {
485    root: PathBuf,
486    manifest: FileArtifactManifest,
487}
488
489impl FileArtifactManifestStore {
490    pub fn write(root: impl AsRef<Path>, bundle: &ExecutionBundle) -> Result<FileArtifactManifest> {
491        bundle.validate()?;
492        let root = root.as_ref();
493        fs::create_dir_all(root).map_err(|err| {
494            DagMlError::RuntimeValidation(format!(
495                "failed to create artifact manifest store `{}`: {err}",
496                root.display()
497            ))
498        })?;
499        let mut entries = Vec::with_capacity(bundle.refit_artifacts.len());
500        for record in &bundle.refit_artifacts {
501            entries.push(FileArtifactManifestEntry::from_refit_record(record)?);
502        }
503        entries.sort_by(|left, right| left.artifact.id.cmp(&right.artifact.id));
504        let manifest = FileArtifactManifest {
505            bundle_id: bundle.bundle_id.clone(),
506            schema_version: FILE_ARTIFACT_MANIFEST_SCHEMA_VERSION,
507            artifacts: entries,
508        };
509        manifest.validate_against_bundle(bundle)?;
510        write_runtime_json(
511            &root.join(FILE_ARTIFACT_MANIFEST_FILE),
512            &manifest,
513            "artifact manifest",
514        )?;
515        Ok(manifest)
516    }
517
518    pub fn open(root: impl Into<PathBuf>, bundle: &ExecutionBundle) -> Result<Self> {
519        bundle.validate()?;
520        let root = root.into();
521        let manifest: FileArtifactManifest =
522            read_runtime_json(&root.join(FILE_ARTIFACT_MANIFEST_FILE), "artifact manifest")?;
523        manifest.validate_against_bundle(bundle)?;
524        Ok(Self { root, manifest })
525    }
526
527    pub fn root(&self) -> &Path {
528        &self.root
529    }
530
531    pub fn manifest(&self) -> &FileArtifactManifest {
532        &self.manifest
533    }
534}
535
536#[derive(Clone, Debug, Eq, PartialEq)]
537pub struct ArtifactPayloadMaterializationRecord {
538    pub run_id: RunId,
539    pub bundle_id: BundleId,
540    pub node_id: NodeId,
541    pub phase: Phase,
542    pub variant_id: Option<VariantId>,
543    pub artifact_id: ArtifactId,
544    pub training_loss_fingerprint: Option<String>,
545    pub payload_uri: String,
546    pub content_fingerprint: String,
547    pub size_bytes: u64,
548    pub handle: HandleRef,
549}
550
551#[derive(Clone, Debug, Eq, PartialEq)]
552pub(crate) struct ArtifactPayloadMetadata {
553    pub(crate) uri: String,
554    pub(crate) content_fingerprint: String,
555    pub(crate) size_bytes: u64,
556}
557
558#[derive(Clone, Debug)]
559pub struct FileArtifactPayloadStore {
560    root: PathBuf,
561    manifest: FileArtifactManifest,
562    records_by_artifact_id: BTreeMap<ArtifactId, RefitArtifactRecord>,
563    materialization_records: RefCell<Vec<ArtifactPayloadMaterializationRecord>>,
564}
565
566impl FileArtifactPayloadStore {
567    pub fn write_from_source(
568        output_root: impl AsRef<Path>,
569        source_root: impl AsRef<Path>,
570        bundle: &ExecutionBundle,
571    ) -> Result<Self> {
572        bundle.validate()?;
573        let output_root = output_root.as_ref();
574        let source_root = source_root.as_ref();
575        fs::create_dir_all(output_root).map_err(|err| {
576            DagMlError::RuntimeValidation(format!(
577                "failed to create artifact payload store `{}`: {err}",
578                output_root.display()
579            ))
580        })?;
581        for record in &bundle.refit_artifacts {
582            record.artifact.validate_portable()?;
583            validate_artifact_payload_file(source_root, &record.artifact)?;
584            let source_path = artifact_payload_path(source_root, &record.artifact)?;
585            let output_path = artifact_payload_path(output_root, &record.artifact)?;
586            if let Some(parent) = output_path.parent() {
587                fs::create_dir_all(parent).map_err(|err| {
588                    DagMlError::RuntimeValidation(format!(
589                        "failed to create artifact payload directory `{}`: {err}",
590                        parent.display()
591                    ))
592                })?;
593            }
594            if source_path != output_path {
595                fs::copy(&source_path, &output_path).map_err(|err| {
596                    DagMlError::RuntimeValidation(format!(
597                        "failed to copy artifact payload `{}` from {} to {}: {err}",
598                        record.artifact.id,
599                        source_path.display(),
600                        output_path.display()
601                    ))
602                })?;
603            }
604        }
605        FileArtifactManifestStore::write(output_root, bundle)?;
606        Self::open(output_root.to_path_buf(), bundle)
607    }
608
609    pub fn open(root: impl Into<PathBuf>, bundle: &ExecutionBundle) -> Result<Self> {
610        bundle.validate()?;
611        let root = root.into();
612        let manifest_store = FileArtifactManifestStore::open(root.clone(), bundle)?;
613        let records_by_artifact_id = bundle
614            .refit_artifacts
615            .iter()
616            .cloned()
617            .map(|record| (record.artifact.id.clone(), record))
618            .collect::<BTreeMap<_, _>>();
619        let store = Self {
620            root,
621            manifest: manifest_store.manifest().clone(),
622            records_by_artifact_id,
623            materialization_records: RefCell::new(Vec::new()),
624        };
625        store.validate_payloads()?;
626        Ok(store)
627    }
628
629    pub fn root(&self) -> &Path {
630        &self.root
631    }
632
633    pub fn manifest(&self) -> &FileArtifactManifest {
634        &self.manifest
635    }
636
637    pub fn payload_count(&self) -> usize {
638        self.manifest.artifacts.len()
639    }
640
641    pub fn materialization_records(&self) -> Vec<ArtifactPayloadMaterializationRecord> {
642        self.materialization_records.borrow().clone()
643    }
644
645    pub fn validate_payloads(&self) -> Result<()> {
646        self.manifest.validate()?;
647        for entry in &self.manifest.artifacts {
648            let record = self
649                .records_by_artifact_id
650                .get(&entry.artifact.id)
651                .ok_or_else(|| {
652                    DagMlError::RuntimeValidation(format!(
653                        "artifact payload store for bundle `{}` has no bundle record for `{}`",
654                        self.manifest.bundle_id, entry.artifact.id
655                    ))
656                })?;
657            if !entry.matches_refit_record(record) {
658                return Err(DagMlError::RuntimeValidation(format!(
659                    "artifact payload store entry `{}` does not match bundle refit artifact",
660                    entry.artifact.id
661                )));
662            }
663            validate_artifact_payload_file(&self.root, &entry.artifact)?;
664        }
665        Ok(())
666    }
667}
668
669impl RuntimeArtifactStore for FileArtifactPayloadStore {
670    fn materialize(&self, request: &ArtifactMaterializationRequest) -> Result<HandleRef> {
671        request.artifact.validate_portable()?;
672        let record = self
673            .records_by_artifact_id
674            .get(&request.artifact.id)
675            .ok_or_else(|| {
676                DagMlError::RuntimeValidation(format!(
677                    "artifact payload store is missing refit artifact `{}` for bundle `{}`",
678                    request.artifact.id, request.bundle_id
679                ))
680            })?;
681        if record.node_id != request.node_id {
682            return Err(DagMlError::RuntimeValidation(format!(
683                "artifact `{}` is registered for node `{}` but requested for `{}`",
684                request.artifact.id, record.node_id, request.node_id
685            )));
686        }
687        if record.controller_id != request.controller_id {
688            return Err(DagMlError::RuntimeValidation(format!(
689                "artifact `{}` is registered for controller `{}` but requested for `{}`",
690                request.artifact.id, record.controller_id, request.controller_id
691            )));
692        }
693        if record.artifact != request.artifact {
694            return Err(DagMlError::RuntimeValidation(format!(
695                "artifact `{}` metadata does not match bundle record",
696                request.artifact.id
697            )));
698        }
699        if record.params_fingerprint != request.params_fingerprint {
700            return Err(DagMlError::RuntimeValidation(format!(
701                "artifact `{}` params fingerprint does not match bundle record",
702                request.artifact.id
703            )));
704        }
705        if record.training_loss_fingerprint != request.training_loss_fingerprint {
706            return Err(DagMlError::RuntimeValidation(format!(
707                "artifact `{}` training loss fingerprint does not match bundle record",
708                request.artifact.id
709            )));
710        }
711        let metadata = validate_artifact_payload_file(&self.root, &request.artifact)?;
712        let fingerprint = stable_json_fingerprint(&(
713            &request.run_id,
714            &request.bundle_id,
715            &request.node_id,
716            request.phase,
717            &request.variant_id,
718            &request.artifact.id,
719            &metadata.content_fingerprint,
720            &request.params_fingerprint,
721            &request.training_loss_fingerprint,
722        ))?;
723        let handle = HandleRef {
724            handle: u64::from_str_radix(&fingerprint[..16], 16)
725                .expect("sha256 hex prefix should fit into u64"),
726            kind: HandleKind::Artifact,
727            owner_controller: request.controller_id.clone(),
728        };
729        self.materialization_records
730            .borrow_mut()
731            .push(ArtifactPayloadMaterializationRecord {
732                run_id: request.run_id.clone(),
733                bundle_id: request.bundle_id.clone(),
734                node_id: request.node_id.clone(),
735                phase: request.phase,
736                variant_id: request.variant_id.clone(),
737                artifact_id: request.artifact.id.clone(),
738                training_loss_fingerprint: request.training_loss_fingerprint.clone(),
739                payload_uri: metadata.uri,
740                content_fingerprint: metadata.content_fingerprint,
741                size_bytes: metadata.size_bytes,
742                handle: handle.clone(),
743            });
744        Ok(handle)
745    }
746}
747
748#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
749#[serde(deny_unknown_fields)]
750pub struct LineageRecord {
751    pub record_id: LineageId,
752    pub run_id: RunId,
753    pub node_id: NodeId,
754    pub phase: Phase,
755    pub controller_id: ControllerId,
756    pub controller_version: String,
757    pub variant_id: Option<VariantId>,
758    pub fold_id: Option<FoldId>,
759    #[serde(default)]
760    pub branch_path: Vec<BranchId>,
761    #[serde(default)]
762    pub input_lineage: Vec<LineageId>,
763    #[serde(default)]
764    pub artifact_refs: Vec<ArtifactRef>,
765    pub params_fingerprint: String,
766    pub data_model_shape_fingerprint: Option<String>,
767    pub aggregation_policy_fingerprint: Option<String>,
768    pub seed: Option<u64>,
769    #[serde(default)]
770    pub unsafe_flags: BTreeSet<String>,
771    #[serde(default)]
772    pub metrics: BTreeMap<String, f64>,
773    #[serde(default, skip_serializing_if = "Vec::is_empty")]
774    pub loss_attestations: Vec<LossExecutionAttestation>,
775    #[serde(default, skip_serializing_if = "Vec::is_empty")]
776    pub early_stopping_records: Vec<EarlyStoppingRecord>,
777}
778
779impl LineageRecord {
780    pub fn validate(&self) -> Result<()> {
781        if self.params_fingerprint.trim().is_empty() {
782            return Err(DagMlError::RuntimeValidation(format!(
783                "lineage `{}` has empty params fingerprint",
784                self.record_id
785            )));
786        }
787        for artifact in &self.artifact_refs {
788            artifact.validate()?;
789        }
790        for attestation in &self.loss_attestations {
791            attestation.validate()?;
792            if attestation.node_id != self.node_id || attestation.phase != self.phase {
793                return Err(DagMlError::RuntimeValidation(format!(
794                    "lineage `{}` contains a loss attestation outside its node/phase scope",
795                    self.record_id
796                )));
797            }
798        }
799        let mut early_stopping_roles = BTreeSet::new();
800        for record in &self.early_stopping_records {
801            record.validate_against(&self.node_id, self.phase, self.fold_id.as_ref())?;
802            if !early_stopping_roles.insert(record.metric_role.role_id.as_str()) {
803                return Err(DagMlError::RuntimeValidation(format!(
804                    "lineage `{}` contains duplicate early-stopping role `{}`",
805                    self.record_id, record.metric_role.role_id
806                )));
807            }
808        }
809        Ok(())
810    }
811}
812
813#[derive(Clone, Debug, Default)]
814pub struct InMemoryLineageRecorder {
815    records: BTreeMap<LineageId, LineageRecord>,
816}
817
818impl InMemoryLineageRecorder {
819    pub fn new() -> Self {
820        Self::default()
821    }
822
823    pub fn record(&mut self, record: LineageRecord) -> Result<()> {
824        record.validate()?;
825        if self
826            .records
827            .insert(record.record_id.clone(), record)
828            .is_some()
829        {
830            return Err(DagMlError::RuntimeValidation(
831                "duplicate lineage record id".to_string(),
832            ));
833        }
834        Ok(())
835    }
836
837    pub fn get(&self, id: &LineageId) -> Option<&LineageRecord> {
838        self.records.get(id)
839    }
840
841    pub fn len(&self) -> usize {
842        self.records.len()
843    }
844
845    pub fn is_empty(&self) -> bool {
846        self.records.is_empty()
847    }
848
849    pub fn records(&self) -> impl Iterator<Item = &LineageRecord> {
850        self.records.values()
851    }
852}
dag_ml_core/runtime/artifact.rs

dag_ml_core/runtime/
artifact.rs
