Skip to main content

cubecl_environment/bundle/
manifest.rs

1use alloc::string::{String, ToString};
2use alloc::vec::Vec;
3
4/// The `meta` key the manifest is stored under.
5#[cfg(native_cache)]
6pub(super) const MANIFEST_KEY: &str = "manifest";
7
8/// The manifest schema version this build reads and writes.
9pub const MANIFEST_SCHEMA: u32 = 1;
10
11/// Error opening or creating a bundle.
12#[derive(Debug)]
13pub enum BundleError {
14    /// The bundle file couldn't be read or written.
15    #[cfg(native_cache)]
16    Io(std::io::Error),
17    /// The environment database couldn't be opened or queried.
18    #[cfg(native_cache)]
19    Storage(String),
20    /// The file opened but carries no manifest, so it isn't a bundle.
21    NotABundle,
22    /// The manifest is not valid for the expected schema.
23    InvalidManifest(String),
24    /// The manifest declares a schema this build doesn't understand.
25    UnsupportedSchema(u32),
26    /// The bundle database is at a schema this build doesn't read.
27    #[cfg(native_cache)]
28    UnsupportedDatabase(String),
29    /// The bundle exceeds what the format can address.
30    TooLarge,
31    /// The blob isn't a readable flat bundle.
32    Flat(super::EmbeddedBundleError),
33}
34
35impl core::fmt::Display for BundleError {
36    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
37        match self {
38            #[cfg(native_cache)]
39            BundleError::Io(err) => write!(f, "bundle io error: {err}"),
40            #[cfg(native_cache)]
41            BundleError::Storage(err) => write!(f, "bundle storage error: {err}"),
42            BundleError::NotABundle => write!(f, "the file carries no bundle manifest"),
43            BundleError::InvalidManifest(err) => write!(f, "invalid bundle manifest: {err}"),
44            BundleError::UnsupportedSchema(schema) => {
45                write!(
46                    f,
47                    "unsupported bundle schema {schema} (this build supports {MANIFEST_SCHEMA})"
48                )
49            }
50            #[cfg(native_cache)]
51            BundleError::UnsupportedDatabase(schema) => {
52                write!(
53                    f,
54                    "unsupported bundle database schema {schema} (this build reads {})",
55                    crate::persistence::turso::SCHEMA_VERSION
56                )
57            }
58            BundleError::TooLarge => write!(
59                f,
60                "the flat bundle format addresses at most {} bytes; export fewer namespaces",
61                u32::MAX
62            ),
63            // Already a complete sentence, and prefixing it would read as two
64            // diagnoses of the same failure.
65            BundleError::Flat(err) => write!(f, "{err}"),
66        }
67    }
68}
69
70impl core::error::Error for BundleError {}
71
72impl From<super::EmbeddedBundleError> for BundleError {
73    fn from(err: super::EmbeddedBundleError) -> Self {
74        Self::Flat(err)
75    }
76}
77
78#[cfg(native_cache)]
79impl From<std::io::Error> for BundleError {
80    fn from(err: std::io::Error) -> Self {
81        Self::Io(err)
82    }
83}
84
85/// The manifest of an environment bundle, stored as a row of the bundle
86/// database in [`BundleFormat::Sqlite`](super::BundleFormat::Sqlite) and as the
87/// metadata blob in [`BundleFormat::Flat`](super::BundleFormat::Flat).
88#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
89pub struct BundleManifest {
90    /// Manifest schema version.
91    pub schema: u32,
92    /// Human-chosen bundle name, e.g. "H100 Linux".
93    pub name: String,
94    /// The cubecl version the bundle was exported with. Entries are only
95    /// visible to the same version, exactly like local caches.
96    pub cubecl_version: String,
97    /// Creation time as seconds since the unix epoch, informational.
98    #[serde(default, skip_serializing_if = "Option::is_none")]
99    pub created_unix_secs: Option<u64>,
100    /// The environments the bundle was captured on, informational in v1.
101    #[serde(default, rename = "environments")]
102    pub environments: Vec<EnvironmentInfo>,
103}
104
105/// Description of one environment a bundle was captured on. Informational:
106/// correctness never depends on these fields.
107#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
108pub struct EnvironmentInfo {
109    /// Short machine-friendly label, e.g. "h100-linux".
110    #[serde(default)]
111    pub label: String,
112    /// Operating system, e.g. "linux".
113    #[serde(default)]
114    pub os: String,
115    /// CPU architecture, e.g. `x86_64`.
116    #[serde(default)]
117    pub arch: String,
118    /// Free-form device fingerprints, e.g. `cuda-0: NVIDIA H100 PCIe (sm_90)`.
119    #[serde(default)]
120    pub devices: Vec<String>,
121}
122
123impl BundleManifest {
124    /// Parses and validates a serialized manifest, the JSON both formats
125    /// store.
126    ///
127    /// This is the guard [`SqliteBundle`](super::SqliteBundle) applies at open,
128    /// available to the flat format too through
129    /// [`EmbeddedBundle::manifest`](super::EmbeddedBundle::manifest).
130    pub fn parse(content: &[u8]) -> Result<Self, BundleError> {
131        if content.is_empty() {
132            return Err(BundleError::NotABundle);
133        }
134
135        let manifest: Self = serde_json::from_slice(content)
136            .map_err(|err| BundleError::InvalidManifest(err.to_string()))?;
137
138        if manifest.schema != MANIFEST_SCHEMA {
139            return Err(BundleError::UnsupportedSchema(manifest.schema));
140        }
141
142        Ok(manifest)
143    }
144
145    /// Warns when the bundle was built for another cubecl version.
146    ///
147    /// Not an error: the bundle still installs, its entries are simply never
148    /// looked up, because the cubecl version is part of every namespace. A
149    /// clear warning beats silent emptiness.
150    pub fn warn_on_version_mismatch(&self) {
151        if self.cubecl_version != env!("CARGO_PKG_VERSION") {
152            log::warn!(
153                "Bundle '{}' was built for cubecl {}, running {}; its entries will be ignored.",
154                self.name,
155                self.cubecl_version,
156                env!("CARGO_PKG_VERSION"),
157            );
158        }
159    }
160
161    /// The manifest of the database `database` opened.
162    ///
163    /// # Errors
164    ///
165    /// [`NotABundle`](BundleError::NotABundle) for an environment no export
166    /// wrote, whatever [`parse`](Self::parse) refuses in the one it holds,
167    /// and a storage error when the database can't be read.
168    #[cfg(native_cache)]
169    pub fn new(database: &crate::persistence::Database) -> Result<Self, BundleError> {
170        database.with_connection(Self::read)
171    }
172
173    /// Writes this manifest into the database `database` opened, making it a
174    /// bundle [`SqliteBundle`](super::SqliteBundle) reads — or replacing the
175    /// manifest it had.
176    ///
177    /// # Errors
178    ///
179    /// When the database refuses the write, a read-only one among them.
180    #[cfg(native_cache)]
181    pub fn write_to(&self, database: &crate::persistence::Database) -> Result<(), BundleError> {
182        database.with_connection(|connection| self.write(connection))
183    }
184
185    /// Reads and validates the manifest of a bundle database.
186    #[cfg(native_cache)]
187    pub(super) fn read(connection: &turso::Connection) -> Result<Self, BundleError> {
188        let content = crate::persistence::turso::meta_get(connection, MANIFEST_KEY)
189            .map_err(super::sqlite::missing_meta)?
190            .ok_or(BundleError::NotABundle)?;
191
192        Self::parse(content.as_bytes())
193    }
194
195    /// Writes the manifest into a bundle database.
196    #[cfg(native_cache)]
197    pub(super) fn write(&self, connection: &turso::Connection) -> Result<(), BundleError> {
198        let content = serde_json::to_string_pretty(self)
199            .map_err(|err| BundleError::InvalidManifest(err.to_string()))?;
200
201        crate::persistence::turso::meta_set(connection, MANIFEST_KEY, &content)
202            .map_err(super::export::storage_error)
203    }
204}