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

dag_ml_core/runtime/
artifact.rs
