cubecl_environment/bundle/
import.rs

1use alloc::string::String;
2use alloc::vec::Vec;
3
4use crate::bytes::Bytes;
5use crate::persistence::{InsertSummary, Origin, storage};
6
7use super::Bundle;
8
9/// How many entries one storage transaction carries.
10///
11/// A bundle is imported in batches rather than entry by entry, because each
12/// batch is one exclusive lock on the cache root and a 10k-entry bundle would
13/// otherwise block every other writer 10k times. The batch is buffered in
14/// memory, so it is bounded rather than "the whole bundle".
15const BATCH: usize = 1024;
16
17/// What an [`import`] copied.
18#[derive(Debug, Clone, Default, PartialEq, Eq)]
19pub struct ImportReport {
20    /// Namespaces the bundle held entries for.
21    pub namespaces: Vec<String>,
22    /// Entries written into the local storage.
23    pub imported: usize,
24    /// Entries skipped because the local storage already had the key. A local
25    /// value is never overwritten by a bundle.
26    pub skipped: usize,
27    /// Entries the storage refused to write: a full disk, an unwritable cache
28    /// root. They are simply absent, so the application recomputes them.
29    pub failed: usize,
30}
31
32/// Copies every entry of `bundle` into the local storage.
33///
34/// This is the only thing a bundle is for. Once imported, entries are ordinary
35/// storage rows and the bundle can be deleted: nothing consults it at runtime,
36/// so a lookup never depends on a file staying installed.
37///
38/// Importing is insert-only and therefore idempotent: a key the storage
39/// already holds keeps its value, whether it was computed here or imported
40/// earlier. Entries land with [`Origin::Imported`], which lets a locally
41/// computed value replace them later if the bundle turns out to be stale.
42///
43/// Fills the *active* environment; switch with
44/// [`environment::activate`](crate::environment::activate) beforehand to
45/// target another one.
46pub fn import(bundle: &dyn Bundle) -> ImportReport {
47    let mut report = ImportReport::default();
48
49    for namespace in bundle.namespaces() {
50        let target = storage::open(&namespace);
51
52        let mut batch: Vec<(Bytes, Bytes)> = Vec::with_capacity(BATCH);
53        let mut total = InsertSummary::default();
54        let commit = |batch: &mut Vec<(Bytes, Bytes)>, total: &mut InsertSummary| {
55            if batch.is_empty() {
56                return;
57            }
58
59            let summary = target.insert_many(&mut batch.drain(..), Origin::Imported);
60            total.stored += summary.stored;
61            total.conflict += summary.conflict;
62            total.failed += summary.failed;
63        };
64
65        bundle.scan(&namespace, &mut |key, value| {
66            batch.push((
67                Bytes::from_bytes_vec(key.to_vec()),
68                Bytes::from_bytes_vec(value.to_vec()),
69            ));
70            if batch.len() == BATCH {
71                commit(&mut batch, &mut total);
72            }
73        });
74        commit(&mut batch, &mut total);
75
76        let (imported, skipped) = (total.stored, total.conflict);
77        log::debug!("Imported {imported} entries into {namespace} ({skipped} already present)");
78        if total.failed > 0 {
79            log::warn!("Failed to import {} entries into {namespace}", total.failed);
80        }
81
82        report.namespaces.push(namespace);
83        report.imported += imported;
84        report.skipped += skipped;
85        report.failed += total.failed;
86    }
87
88    log::info!(
89        "Imported {} entries from {} into {} namespaces ({} already present)",
90        report.imported,
91        bundle.describe(),
92        report.namespaces.len(),
93        report.skipped,
94    );
95
96    report
97}
cubecl_environment/bundle/import.rs

cubecl_environment/bundle/
import.rs
