1use std::collections::BTreeMap;
2use std::error::Error;
3use std::fmt;
4
5use crate::{Artifact, PackageManifest, TensorIntegrity, TensorStorage, ValidationErrors};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct MaterializationTensor<'a> {
9 pub tensor_id: &'a str,
10 pub native_name: &'a str,
11 pub ggml_type: u32,
12 pub dimensions: &'a [u64],
13 pub artifact: &'a Artifact,
14 pub data_offset: u64,
15 pub stored_length: u64,
16 pub alignment: u64,
17 pub integrity: &'a TensorIntegrity,
18}
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum MaterializationBindingError {
22 InvalidManifest(ValidationErrors),
23 EmptyTensorId {
24 index: usize,
25 },
26 TensorIdsNotStrictlySorted {
27 index: usize,
28 previous: String,
29 current: String,
30 },
31 UnknownTensorId {
32 tensor_id: String,
33 },
34 AliasNotMaterializable {
35 tensor_id: String,
36 native_name: String,
37 target_tensor_id: String,
38 },
39 MissingArtifact {
40 tensor_id: String,
41 artifact_id: String,
42 },
43}
44
45impl fmt::Display for MaterializationBindingError {
46 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
47 match self {
48 Self::InvalidManifest(error) => write!(formatter, "{error}"),
49 Self::EmptyTensorId { index } => {
50 write!(
51 formatter,
52 "materialization tensor id at index {index} is empty"
53 )
54 }
55 Self::TensorIdsNotStrictlySorted {
56 index,
57 previous,
58 current,
59 } => write!(
60 formatter,
61 "materialization tensor ids are not strictly sorted at index {index}: {previous:?} then {current:?}"
62 ),
63 Self::UnknownTensorId { tensor_id } => {
64 write!(formatter, "unknown materialization tensor id {tensor_id:?}")
65 }
66 Self::AliasNotMaterializable {
67 tensor_id,
68 native_name,
69 target_tensor_id,
70 } => write!(
71 formatter,
72 "selected tensor {tensor_id:?} ({native_name:?}) aliases {target_tensor_id:?}, but alias materialization is unsupported"
73 ),
74 Self::MissingArtifact {
75 tensor_id,
76 artifact_id,
77 } => write!(
78 formatter,
79 "selected tensor {tensor_id:?} references missing artifact {artifact_id:?}"
80 ),
81 }
82 }
83}
84
85impl Error for MaterializationBindingError {
86 fn source(&self) -> Option<&(dyn Error + 'static)> {
87 match self {
88 Self::InvalidManifest(error) => Some(error),
89 _ => None,
90 }
91 }
92}
93
94impl PackageManifest {
95 pub fn materialization_tensors<'a>(
101 &'a self,
102 exact_tensor_ids: &[String],
103 ) -> Result<Vec<MaterializationTensor<'a>>, MaterializationBindingError> {
104 self.validate()
105 .map_err(MaterializationBindingError::InvalidManifest)?;
106 validate_requested_ids(exact_tensor_ids)?;
107
108 let tensors = self
109 .tensor_catalog
110 .entries
111 .iter()
112 .map(|tensor| (tensor.id.as_str(), tensor))
113 .collect::<BTreeMap<_, _>>();
114 let artifacts = self
115 .artifact_catalog
116 .entries
117 .iter()
118 .map(|artifact| (artifact.id.as_str(), artifact))
119 .collect::<BTreeMap<_, _>>();
120
121 exact_tensor_ids
122 .iter()
123 .map(|tensor_id| {
124 let tensor = tensors.get(tensor_id.as_str()).ok_or_else(|| {
125 MaterializationBindingError::UnknownTensorId {
126 tensor_id: tensor_id.clone(),
127 }
128 })?;
129 let TensorStorage::Owned {
130 artifact_id,
131 data_offset,
132 stored_length,
133 alignment,
134 integrity,
135 } = &tensor.storage
136 else {
137 let TensorStorage::Alias { target_tensor_id } = &tensor.storage else {
138 unreachable!()
139 };
140 return Err(MaterializationBindingError::AliasNotMaterializable {
141 tensor_id: tensor.id.clone(),
142 native_name: tensor.name.clone(),
143 target_tensor_id: target_tensor_id.clone(),
144 });
145 };
146 let artifact = artifacts.get(artifact_id.as_str()).ok_or_else(|| {
147 MaterializationBindingError::MissingArtifact {
148 tensor_id: tensor.id.clone(),
149 artifact_id: artifact_id.clone(),
150 }
151 })?;
152 Ok(MaterializationTensor {
153 tensor_id: &tensor.id,
154 native_name: &tensor.name,
155 ggml_type: tensor.ggml_type,
156 dimensions: &tensor.dimensions,
157 artifact,
158 data_offset: *data_offset,
159 stored_length: *stored_length,
160 alignment: *alignment,
161 integrity,
162 })
163 })
164 .collect()
165 }
166}
167
168fn validate_requested_ids(exact_tensor_ids: &[String]) -> Result<(), MaterializationBindingError> {
169 for (index, tensor_id) in exact_tensor_ids.iter().enumerate() {
170 if tensor_id.is_empty() {
171 return Err(MaterializationBindingError::EmptyTensorId { index });
172 }
173 }
174 for (index, pair) in exact_tensor_ids.windows(2).enumerate() {
175 if pair[0] >= pair[1] {
176 return Err(MaterializationBindingError::TensorIdsNotStrictlySorted {
177 index: index + 1,
178 previous: pair[0].clone(),
179 current: pair[1].clone(),
180 });
181 }
182 }
183 Ok(())
184}
185
186#[cfg(test)]
187mod tests {
188 use std::collections::BTreeMap;
189
190 use serde_json::Value;
191
192 use super::*;
193 use crate::{
194 ArtifactCatalog, PACKAGE_SCHEMA_VERSION, SourceFile, SourceModel, Tensor, TensorCatalog,
195 };
196
197 const DIGEST: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";
198
199 #[test]
200 fn binds_exact_id_to_distinct_native_name() {
201 let manifest = fixture();
202 let bindings = manifest
203 .materialization_tensors(&["a".to_string()])
204 .unwrap();
205 assert_eq!(bindings.len(), 1);
206 assert_eq!(bindings[0].tensor_id, "a");
207 assert_eq!(bindings[0].native_name, "native-only");
208 assert_eq!(bindings[0].artifact.id, "weights");
209 assert_eq!(bindings[0].artifact.path, "artifacts/source-00000.gguf");
210 assert_eq!(bindings[0].artifact.sha256, DIGEST);
211 assert_eq!(bindings[0].data_offset, 0);
212 assert_eq!(bindings[0].stored_length, 32);
213 assert_eq!(bindings[0].alignment, 32);
214 assert_eq!(bindings[0].integrity, &TensorIntegrity::ArtifactSha256);
215
216 assert!(matches!(
217 manifest.materialization_tensors(&["native-only".to_string()]),
218 Err(MaterializationBindingError::UnknownTensorId { tensor_id })
219 if tensor_id == "native-only"
220 ));
221 assert!(matches!(
222 manifest.materialization_tensors(&["b".to_string()]),
223 Ok(bindings) if bindings[0].native_name == "a"
224 ));
225 }
226
227 #[test]
228 fn rejects_non_normalized_or_missing_ids() {
229 let manifest = fixture();
230 assert!(matches!(
231 manifest.materialization_tensors(&["b".to_string(), "a".to_string()]),
232 Err(MaterializationBindingError::TensorIdsNotStrictlySorted { .. })
233 ));
234 assert!(matches!(
235 manifest.materialization_tensors(&["a".to_string(), "a".to_string()]),
236 Err(MaterializationBindingError::TensorIdsNotStrictlySorted { .. })
237 ));
238 assert!(matches!(
239 manifest.materialization_tensors(&["missing".to_string()]),
240 Err(MaterializationBindingError::UnknownTensorId { .. })
241 ));
242
243 let bindings = manifest
244 .materialization_tensors(&["a".to_string(), "b".to_string()])
245 .unwrap();
246 assert_eq!(
247 bindings
248 .iter()
249 .map(|binding| binding.tensor_id)
250 .collect::<Vec<_>>(),
251 ["a", "b"]
252 );
253 }
254
255 #[test]
256 fn rejects_selected_alias_without_substitution() {
257 let mut manifest = fixture();
258 manifest.tensor_catalog.entries.push(Tensor {
259 id: "c".to_string(),
260 name: "alias-native".to_string(),
261 ggml_type: 1,
262 dimensions: vec![4, 4],
263 layer_ordinal: None,
264 storage: TensorStorage::Alias {
265 target_tensor_id: "a".to_string(),
266 },
267 });
268 manifest.package_id = manifest.computed_package_id().unwrap();
269 manifest.validate().unwrap();
270
271 assert!(manifest.materialization_tensors(&["a".to_string()]).is_ok());
272 assert!(matches!(
273 manifest.materialization_tensors(&["c".to_string()]),
274 Err(MaterializationBindingError::AliasNotMaterializable {
275 tensor_id,
276 native_name,
277 target_tensor_id,
278 }) if tensor_id == "c" && native_name == "alias-native" && target_tensor_id == "a"
279 ));
280 assert!(matches!(
281 manifest.materialization_tensors(&["a".to_string(), "c".to_string()]),
282 Err(MaterializationBindingError::AliasNotMaterializable { .. })
283 ));
284 }
285
286 #[test]
287 fn rejects_invalid_manifest_before_binding() {
288 let mut manifest = fixture();
289 manifest.package_id = "sha256:invalid".to_string();
290
291 assert!(matches!(
292 manifest.materialization_tensors(&["a".to_string()]),
293 Err(MaterializationBindingError::InvalidManifest(_))
294 ));
295 }
296
297 fn fixture() -> PackageManifest {
298 let mut manifest = PackageManifest {
299 schema_version: PACKAGE_SCHEMA_VERSION,
300 package_id: String::new(),
301 model_id: "model".to_string(),
302 source_model: SourceModel {
303 sha256: DIGEST.to_string(),
304 metadata_artifact_id: "weights".to_string(),
305 repo: None,
306 revision: None,
307 primary_file: Some("model.gguf".to_string()),
308 canonical_ref: None,
309 distribution_id: None,
310 files: vec![SourceFile {
311 path: "model.gguf".to_string(),
312 byte_size: 128,
313 sha256: DIGEST.to_string(),
314 }],
315 },
316 format: "gguf-v2".to_string(),
317 layer_count: 2,
318 model_metadata: BTreeMap::from([(
319 "general.architecture".to_string(),
320 Value::String("llama".to_string()),
321 )]),
322 artifact_catalog: ArtifactCatalog {
323 entries: vec![Artifact {
324 id: "weights".to_string(),
325 path: "artifacts/source-00000.gguf".to_string(),
326 byte_size: 128,
327 sha256: DIGEST.to_string(),
328 }],
329 },
330 tensor_catalog: TensorCatalog {
331 entries: vec![
332 Tensor {
333 id: "a".to_string(),
334 name: "native-only".to_string(),
335 ggml_type: 1,
336 dimensions: vec![4, 4],
337 layer_ordinal: None,
338 storage: TensorStorage::Owned {
339 artifact_id: "weights".to_string(),
340 data_offset: 0,
341 stored_length: 32,
342 alignment: 32,
343 integrity: TensorIntegrity::ArtifactSha256,
344 },
345 },
346 Tensor {
347 id: "b".to_string(),
348 name: "a".to_string(),
349 ggml_type: 1,
350 dimensions: vec![4, 4],
351 layer_ordinal: Some(1),
352 storage: TensorStorage::Owned {
353 artifact_id: "weights".to_string(),
354 data_offset: 64,
355 stored_length: 32,
356 alignment: 32,
357 integrity: TensorIntegrity::ArtifactSha256,
358 },
359 },
360 ],
361 },
362 sidecars: Vec::new(),
363 publisher_metadata: Vec::new(),
364 publisher_defaults: None,
365 generation: None,
366 native_abi_version: "0.1.49".to_string(),
367 generator_version: "test".to_string(),
368 created_at_unix_secs: 1,
369 };
370 manifest.package_id = manifest.computed_package_id().unwrap();
371 manifest
372 }
373}