1use crate::checkpoint::{TensorCatalog, TensorDescriptor, TensorDtype, TensorStorage};
7use eredu_checkpoint::safetensors::SafetensorsShards;
8use eredu_gguf::{Checkpoint as GgufCheckpoint, GgmlType, MetadataValue};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11use std::{
12 collections::{BTreeMap, BTreeSet},
13 fs::File,
14 io::Read,
15 path::{Path, PathBuf},
16};
17
18#[derive(Debug, Clone, Copy, Eq, Hash, PartialEq, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24#[non_exhaustive]
25pub enum LoadingProtocol {
26 Model,
28 Realtime,
30}
31
32#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
34#[serde(rename_all = "snake_case")]
35#[non_exhaustive]
36pub enum ArtifactFormat {
37 SafeTensors,
39 Gguf,
41}
42
43#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
45pub struct ModelConfiguration {
46 declared_model_type: String,
48 effective_model_type: String,
50 family: String,
52 loading_protocol: LoadingProtocol,
54 #[serde(skip_serializing_if = "Option::is_none")]
56 json: Option<Value>,
57}
58
59impl ModelConfiguration {
60 pub fn new(
62 declared_model_type: impl Into<String>,
63 effective_model_type: impl Into<String>,
64 family: impl Into<String>,
65 loading_protocol: LoadingProtocol,
66 json: Option<Value>,
67 ) -> Result<Self, ArtifactError> {
68 let configuration = Self {
69 declared_model_type: declared_model_type.into(),
70 effective_model_type: effective_model_type.into(),
71 family: family.into(),
72 loading_protocol,
73 json,
74 };
75 if [
76 &configuration.declared_model_type,
77 &configuration.effective_model_type,
78 &configuration.family,
79 ]
80 .into_iter()
81 .any(|value| value.trim().is_empty())
82 {
83 return Err(ArtifactError::InvalidArtifact(
84 "model configuration identities must be non-empty".into(),
85 ));
86 }
87 Ok(configuration)
88 }
89
90 pub fn declared_model_type(&self) -> &str {
92 &self.declared_model_type
93 }
94
95 pub fn effective_model_type(&self) -> &str {
97 &self.effective_model_type
98 }
99
100 pub fn family(&self) -> &str {
102 &self.family
103 }
104
105 pub const fn loading_protocol(&self) -> LoadingProtocol {
107 self.loading_protocol
108 }
109
110 pub const fn json(&self) -> Option<&Value> {
112 self.json.as_ref()
113 }
114}
115
116#[derive(Debug, Clone)]
118pub struct ResolvedModelConfiguration<P> {
119 configuration: ModelConfiguration,
121 architecture_plan: P,
123}
124
125impl<P> ResolvedModelConfiguration<P> {
126 pub fn new(configuration: ModelConfiguration, architecture_plan: P) -> Self {
128 Self {
129 configuration,
130 architecture_plan,
131 }
132 }
133
134 pub const fn configuration(&self) -> &ModelConfiguration {
136 &self.configuration
137 }
138
139 pub const fn architecture_plan(&self) -> &P {
141 &self.architecture_plan
142 }
143
144 pub fn into_parts(self) -> (ModelConfiguration, P) {
146 (self.configuration, self.architecture_plan)
147 }
148}
149
150pub trait ModelConfigurationResolver {
155 type ArtifactPlan: Clone + std::fmt::Debug;
157
158 fn resolve_safetensors(
160 &self,
161 json: &Value,
162 ) -> Result<ResolvedModelConfiguration<Self::ArtifactPlan>, ArtifactError>;
163
164 fn resolve_gguf(
166 &self,
167 architecture: &str,
168 checkpoint: &GgufCheckpoint,
169 ) -> Result<ResolvedModelConfiguration<Self::ArtifactPlan>, ArtifactError>;
170
171 fn gguf_companion_requirements(
173 &self,
174 architecture: &str,
175 checkpoint: &GgufCheckpoint,
176 ) -> Result<Vec<GgufCompanionRequirement>, ArtifactError>;
177
178 fn artifact_plan(
181 &self,
182 _path: &Path,
183 _format: ArtifactFormat,
184 _configuration: &ModelConfiguration,
185 _tensors: &TensorCatalog,
186 _validated_gguf: Option<&ValidatedGguf>,
187 resolved_plan: Self::ArtifactPlan,
188 ) -> Result<Self::ArtifactPlan, ArtifactError> {
189 Ok(resolved_plan)
190 }
191}
192
193#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
195#[non_exhaustive]
196pub enum GgufCompanionRole {
197 MediaProjector,
199 Named(String),
201}
202
203#[derive(Debug, Clone, Copy, Eq, PartialEq)]
205#[non_exhaustive]
206pub enum GgufCompanionEncoding {
207 DenseRequired,
209 DensePreferred,
211}
212
213#[derive(Debug, Clone, Eq, PartialEq)]
215pub struct GgufCompanionRequirement {
216 role: GgufCompanionRole,
217 required: bool,
218 filename_prefix: String,
219 parent_search_depth: usize,
220 encoding: GgufCompanionEncoding,
221}
222
223impl GgufCompanionRequirement {
224 pub fn new(
226 role: GgufCompanionRole,
227 required: bool,
228 filename_prefix: impl Into<String>,
229 parent_search_depth: usize,
230 encoding: GgufCompanionEncoding,
231 ) -> Result<Self, ArtifactError> {
232 let filename_prefix = filename_prefix.into();
233 if filename_prefix.trim().is_empty()
234 || matches!(&role, GgufCompanionRole::Named(name) if name.trim().is_empty())
235 {
236 return Err(ArtifactError::InvalidArtifact(
237 "GGUF companion roles and filename prefixes must be non-empty".into(),
238 ));
239 }
240 Ok(Self {
241 role,
242 required,
243 filename_prefix,
244 parent_search_depth,
245 encoding,
246 })
247 }
248
249 pub fn role(&self) -> &GgufCompanionRole {
251 &self.role
252 }
253}
254
255pub fn gguf_u32_metadata_values(
257 key: &str,
258 value: Option<&MetadataValue>,
259) -> Result<Vec<u32>, ArtifactError> {
260 let Some(value) = value else {
261 return Ok(Vec::new());
262 };
263 value.to_u32_vec().ok_or_else(|| {
264 ArtifactError::InvalidArtifact(format!(
265 "GGUF metadata key {key:?} must contain an integer or integer array whose values fit in u32"
266 ))
267 })
268}
269
270#[derive(Debug, Clone)]
272pub struct ArtifactInspection<P = ()> {
273 path: PathBuf,
274 format: ArtifactFormat,
275 configuration: ModelConfiguration,
276 tensors: TensorCatalog,
277 safetensors_shards: Option<SafetensorsShards>,
278 validated_gguf: Option<ValidatedGguf>,
279 architecture_plan: P,
280}
281
282#[derive(Debug, Clone)]
287pub struct ValidatedGguf {
288 checkpoint: GgufCheckpoint,
289 companions: BTreeMap<GgufCompanionRole, ValidatedGgufCompanion>,
290}
291
292#[derive(Debug, Clone)]
294pub struct ValidatedGgufCompanion {
295 path: PathBuf,
296 checkpoint: GgufCheckpoint,
297}
298
299impl ValidatedGgufCompanion {
300 pub fn path(&self) -> &Path {
302 &self.path
303 }
304
305 pub fn checkpoint(&self) -> &GgufCheckpoint {
307 &self.checkpoint
308 }
309}
310
311impl ValidatedGguf {
312 pub fn checkpoint(&self) -> &GgufCheckpoint {
314 &self.checkpoint
315 }
316
317 pub fn companion(&self, role: &GgufCompanionRole) -> Option<&ValidatedGgufCompanion> {
319 self.companions.get(role)
320 }
321
322 pub fn companions(
324 &self,
325 ) -> impl Iterator<Item = (&GgufCompanionRole, &ValidatedGgufCompanion)> {
326 self.companions.iter()
327 }
328
329 pub fn into_parts(
331 self,
332 ) -> (
333 GgufCheckpoint,
334 BTreeMap<GgufCompanionRole, ValidatedGgufCompanion>,
335 ) {
336 (self.checkpoint, self.companions)
337 }
338}
339
340impl<P> ArtifactInspection<P> {
341 pub fn path(&self) -> &Path {
343 &self.path
344 }
345 pub const fn format(&self) -> ArtifactFormat {
347 self.format
348 }
349 pub fn configuration(&self) -> &ModelConfiguration {
351 &self.configuration
352 }
353 pub fn tensors(&self) -> &TensorCatalog {
355 &self.tensors
356 }
357 pub fn safetensors_shards(&self) -> Option<&SafetensorsShards> {
359 self.safetensors_shards.as_ref()
360 }
361 pub fn validated_gguf(&self) -> Option<&ValidatedGguf> {
363 self.validated_gguf.as_ref()
364 }
365 pub fn gguf_checkpoint(&self) -> Option<&GgufCheckpoint> {
367 self.validated_gguf().map(ValidatedGguf::checkpoint)
368 }
369 pub fn architecture_plan(&self) -> &P {
371 &self.architecture_plan
372 }
373 pub fn architecture_plan_mut(&mut self) -> &mut P {
375 &mut self.architecture_plan
376 }
377}
378
379#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
381#[serde(tag = "kind", rename_all = "snake_case")]
382#[non_exhaustive]
383pub enum QuantizationRequest {
384 Affine {
386 group_size: u32,
388 bits: u8,
390 },
391 MxFp4,
393}
394
395#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
397#[serde(rename_all = "snake_case")]
398#[non_exhaustive]
399pub enum ResidencyRequest {
400 #[default]
402 FullyResident,
403 LayerwiseHost,
405 DenseDiskStream,
407 AddressableParameterBanks,
409}
410
411#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
413pub struct PreparationPolicy {
414 quantization: Option<QuantizationRequest>,
416 residency: ResidencyRequest,
418 topology: Option<crate::topology::ParallelTopology>,
420 required_session_capabilities: crate::backend::SessionCapabilities,
422}
423
424impl PreparationPolicy {
425 pub const fn new(
427 quantization: Option<QuantizationRequest>,
428 residency: ResidencyRequest,
429 ) -> Self {
430 Self {
431 quantization,
432 residency,
433 topology: None,
434 required_session_capabilities: crate::backend::SessionCapabilities::new(
435 false, false, false,
436 ),
437 }
438 }
439
440 pub const fn quantization(self) -> Option<QuantizationRequest> {
442 self.quantization
443 }
444 pub const fn residency(self) -> ResidencyRequest {
446 self.residency
447 }
448 pub const fn topology(self) -> Option<crate::topology::ParallelTopology> {
450 self.topology
451 }
452 pub const fn required_session_capabilities(self) -> crate::backend::SessionCapabilities {
454 self.required_session_capabilities
455 }
456 pub const fn with_topology(mut self, topology: crate::topology::ParallelTopology) -> Self {
458 self.topology = Some(topology);
459 self
460 }
461 pub const fn with_required_session_capabilities(
463 mut self,
464 capabilities: crate::backend::SessionCapabilities,
465 ) -> Self {
466 self.required_session_capabilities = capabilities;
467 self
468 }
469
470 pub fn validate_session_capabilities(
472 &self,
473 available: &crate::backend::SessionCapabilities,
474 ) -> Result<(), crate::backend::SessionCapabilityError> {
475 self.required_session_capabilities.validate(available)
476 }
477}
478
479#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
481#[serde(rename_all = "snake_case")]
482#[non_exhaustive]
483pub enum MaterializationRoute {
484 Resident,
486 Layerwise,
488 AddressableParameterBanks,
490}
491
492#[derive(Debug, Clone)]
494pub struct ModelPreparationPlan<P = ()> {
495 inspection: ArtifactInspection<P>,
496 policy: PreparationPolicy,
497 route: MaterializationRoute,
498 admitted_session_capabilities: crate::backend::SessionCapabilities,
499}
500
501impl<P> ModelPreparationPlan<P> {
502 pub fn inspection(&self) -> &ArtifactInspection<P> {
504 &self.inspection
505 }
506 pub const fn policy(&self) -> PreparationPolicy {
508 self.policy
509 }
510 pub const fn route(&self) -> MaterializationRoute {
512 self.route
513 }
514 pub const fn admitted_session_capabilities(&self) -> crate::backend::SessionCapabilities {
516 self.admitted_session_capabilities
517 }
518 pub fn into_artifact(self) -> ModelArtifact {
525 match self.inspection.validated_gguf {
526 Some(validated) => ModelArtifact::Gguf {
527 path: self.inspection.path,
528 configuration: self.inspection.configuration,
529 tensors: self.inspection.tensors,
530 validated,
531 },
532 None => ModelArtifact::SafeTensors {
533 path: self.inspection.path,
534 configuration: self.inspection.configuration,
535 tensors: self.inspection.tensors,
536 shards: self
537 .inspection
538 .safetensors_shards
539 .expect("SafeTensors inspection retains its admitted shard set"),
540 },
541 }
542 }
543}
544
545#[derive(Debug, Clone)]
547#[non_exhaustive]
548pub enum ModelArtifact {
549 SafeTensors {
551 path: PathBuf,
553 configuration: ModelConfiguration,
555 tensors: TensorCatalog,
557 shards: SafetensorsShards,
559 },
560 Gguf {
562 path: PathBuf,
564 configuration: ModelConfiguration,
566 tensors: TensorCatalog,
568 validated: ValidatedGguf,
570 },
571}
572
573pub fn inspect_artifact<R: ModelConfigurationResolver>(
575 path: impl AsRef<Path>,
576 resolver: &R,
577) -> Result<ArtifactInspection<R::ArtifactPlan>, ArtifactError> {
578 let path = path.as_ref();
579 if is_gguf(path) {
580 inspect_gguf(path, resolver)
581 } else if path.is_dir() {
582 inspect_safetensors(path, resolver)
583 } else if !path.exists() {
584 Err(ArtifactError::MissingArtifact(path.to_path_buf()))
585 } else {
586 Err(ArtifactError::UnsupportedContainer(path.to_path_buf()))
587 }
588}
589
590pub fn plan_model_preparation<P>(
592 inspection: ArtifactInspection<P>,
593 policy: PreparationPolicy,
594 admitted_session_capabilities: crate::backend::SessionCapabilities,
595) -> Result<ModelPreparationPlan<P>, ArtifactError> {
596 let route = validate_preparation_policy(inspection.configuration.loading_protocol, policy)?;
597 Ok(ModelPreparationPlan {
598 inspection,
599 policy,
600 route,
601 admitted_session_capabilities,
602 })
603}
604
605pub fn validate_preparation_policy(
607 protocol: LoadingProtocol,
608 policy: PreparationPolicy,
609) -> Result<MaterializationRoute, ArtifactError> {
610 if protocol != LoadingProtocol::Model {
611 return Err(ArtifactError::UnsupportedLoadingProtocol(protocol));
612 }
613 let route = match policy.residency {
614 ResidencyRequest::FullyResident => MaterializationRoute::Resident,
615 ResidencyRequest::LayerwiseHost | ResidencyRequest::DenseDiskStream => {
616 MaterializationRoute::Layerwise
617 }
618 ResidencyRequest::AddressableParameterBanks => {
619 MaterializationRoute::AddressableParameterBanks
620 }
621 };
622 Ok(route)
623}
624
625fn inspect_gguf<R: ModelConfigurationResolver>(
626 path: &Path,
627 resolver: &R,
628) -> Result<ArtifactInspection<R::ArtifactPlan>, ArtifactError> {
629 let checkpoint = GgufCheckpoint::open(path)?;
630 let architecture_name = checkpoint
631 .metadata()
632 .get("general.architecture")
633 .and_then(MetadataValue::as_str)
634 .ok_or(ArtifactError::MissingGgufArchitecture)?;
635 let (configuration, resolved_plan) = resolver
636 .resolve_gguf(architecture_name, &checkpoint)?
637 .into_parts();
638 let requirements = resolver.gguf_companion_requirements(architecture_name, &checkpoint)?;
639 let companions = resolve_gguf_companions(path, &requirements)?;
640 validate_gguf_container(&checkpoint)?;
641 let tensors = checkpoint
642 .tensors()
643 .map(|tensor| {
644 let descriptor = tensor.descriptor();
645 let shape = descriptor
646 .dimensions
647 .iter()
648 .map(|&dimension| {
649 usize::try_from(dimension).map_err(|_| {
650 ArtifactError::InvalidArtifact(format!(
651 "GGUF tensor {:?} dimension {dimension} exceeds the host address space",
652 descriptor.name
653 ))
654 })
655 })
656 .collect::<Result<Vec<_>, _>>()?;
657 Ok(TensorDescriptor {
658 name: descriptor.name.clone(),
659 shape,
660 dtype: gguf_dtype(descriptor.ggml_type),
661 storage: None,
662 })
663 })
664 .collect::<Result<Vec<_>, ArtifactError>>()?;
665 let tensors = TensorCatalog::new(tensors)?;
666 let validated_gguf = ValidatedGguf {
667 checkpoint,
668 companions,
669 };
670 let architecture_plan = resolver.artifact_plan(
671 path,
672 ArtifactFormat::Gguf,
673 &configuration,
674 &tensors,
675 Some(&validated_gguf),
676 resolved_plan,
677 )?;
678 Ok(ArtifactInspection {
679 path: path.to_path_buf(),
680 format: ArtifactFormat::Gguf,
681 configuration,
682 tensors,
683 safetensors_shards: None,
684 validated_gguf: Some(validated_gguf),
685 architecture_plan,
686 })
687}
688
689pub fn resolve_gguf_companions(
691 primary: &Path,
692 requirements: &[GgufCompanionRequirement],
693) -> Result<BTreeMap<GgufCompanionRole, ValidatedGgufCompanion>, ArtifactError> {
694 let mut resolved = BTreeMap::new();
695 let mut declared_roles = BTreeSet::new();
696 for requirement in requirements {
697 if !declared_roles.insert(requirement.role.clone()) {
698 return Err(ArtifactError::InvalidArtifact(format!(
699 "GGUF companion role {:?} was declared more than once",
700 requirement.role
701 )));
702 }
703 let mut directories = Vec::new();
704 let mut directory = primary.parent().unwrap_or_else(|| Path::new("."));
705 directories.push(directory.to_path_buf());
706 for _ in 0..requirement.parent_search_depth {
707 let Some(parent) = directory.parent() else {
708 break;
709 };
710 if parent == directory {
711 break;
712 }
713 directories.push(parent.to_path_buf());
714 directory = parent;
715 }
716 let mut candidates = Vec::new();
717 for directory in &directories {
718 let candidate_start = candidates.len();
719 for entry in std::fs::read_dir(directory)? {
720 let path = entry?.path();
721 let name = path
722 .file_name()
723 .and_then(|name| name.to_str())
724 .unwrap_or_default();
725 if path != primary
726 && path.is_file()
727 && name
728 .get(..requirement.filename_prefix.len())
729 .is_some_and(|prefix| {
730 prefix.eq_ignore_ascii_case(&requirement.filename_prefix)
731 })
732 && is_gguf(&path)
733 {
734 let checkpoint = GgufCheckpoint::open(&path)?;
735 if checkpoint.physical_tensor_count() == 0 {
736 return Err(ArtifactError::InvalidArtifact(format!(
737 "GGUF companion {} contains no tensors",
738 path.display()
739 )));
740 }
741 let dense = checkpoint.tensors().all(|tensor| {
742 matches!(
743 tensor.descriptor().ggml_type,
744 eredu_gguf::GgmlType::F32
745 | eredu_gguf::GgmlType::F16
746 | eredu_gguf::GgmlType::Bf16
747 )
748 });
749 candidates.push((path, checkpoint, dense));
750 }
751 }
752 if candidates.len() != candidate_start {
753 break;
754 }
755 }
756 candidates.sort_by(|left, right| left.0.cmp(&right.0));
757 candidates.dedup_by(|left, right| left.0 == right.0);
758 let dense = candidates
759 .iter()
760 .filter(|candidate| candidate.2)
761 .collect::<Vec<_>>();
762 let selected = match requirement.encoding {
763 GgufCompanionEncoding::DenseRequired => match dense.as_slice() {
764 [candidate] => Some(*candidate),
765 [] if candidates.is_empty() => None,
766 [] => {
767 return Err(ArtifactError::InvalidArtifact(format!(
768 "GGUF companion {:?} requires dense F32, F16, or BF16 tensors, but all {} matching candidates are quantized",
769 requirement.role,
770 candidates.len()
771 )))
772 }
773 _ => return Err(ambiguous_companion(requirement, &directories, dense.len())),
774 },
775 GgufCompanionEncoding::DensePreferred => match dense.as_slice() {
776 [candidate] => Some(*candidate),
777 [] => match candidates.as_slice() {
778 [candidate] => Some(candidate),
779 [] => None,
780 _ => {
781 return Err(ambiguous_companion(
782 requirement,
783 &directories,
784 candidates.len(),
785 ))
786 }
787 },
788 _ => return Err(ambiguous_companion(requirement, &directories, dense.len())),
789 },
790 };
791 match selected {
792 Some((path, checkpoint, _)) => {
793 resolved.insert(
794 requirement.role.clone(),
795 ValidatedGgufCompanion {
796 path: path.clone(),
797 checkpoint: checkpoint.clone(),
798 },
799 );
800 }
801 None if requirement.required => {
802 return Err(ArtifactError::MissingRequiredGgufCompanion {
803 role: requirement.role.clone(),
804 filename_prefix: requirement.filename_prefix.clone(),
805 searched_directories: directories,
806 })
807 }
808 None => {}
809 }
810 }
811 Ok(resolved)
812}
813
814fn ambiguous_companion(
815 requirement: &GgufCompanionRequirement,
816 directories: &[PathBuf],
817 candidates: usize,
818) -> ArtifactError {
819 ArtifactError::InvalidArtifact(format!(
820 "GGUF companion {:?} is ambiguous: found {candidates} preferred candidates in {}",
821 requirement.role,
822 display_directories(directories)
823 ))
824}
825
826fn display_directories(directories: &[PathBuf]) -> String {
827 directories
828 .iter()
829 .map(|directory| directory.display().to_string())
830 .collect::<Vec<_>>()
831 .join(", ")
832}
833
834fn validate_gguf_container(checkpoint: &GgufCheckpoint) -> Result<(), ArtifactError> {
835 if checkpoint.physical_tensor_count() == 0 {
836 return Err(ArtifactError::InvalidArtifact(
837 "GGUF model checkpoint contains no tensors".into(),
838 ));
839 }
840 Ok(())
841}
842
843fn inspect_safetensors<R: ModelConfigurationResolver>(
844 path: &Path,
845 resolver: &R,
846) -> Result<ArtifactInspection<R::ArtifactPlan>, ArtifactError> {
847 let config_path = path.join("config.json");
848 let json: Value = serde_json::from_reader(File::open(&config_path)?)?;
849 let (configuration, resolved_plan) = resolver.resolve_safetensors(&json)?.into_parts();
850 let shards = SafetensorsShards::discover(path)?;
851 let mut descriptors = Vec::new();
852 let mut names = BTreeSet::new();
853 for shard in shards.payload_paths() {
854 for descriptor in inspect_safetensors_header(shard)? {
855 if !names.insert(descriptor.name.clone()) {
856 return Err(ArtifactError::DuplicateTensor(descriptor.name));
857 }
858 descriptors.push(descriptor);
859 }
860 }
861 let tensors = TensorCatalog::new(descriptors)?;
862 if tensors.is_empty() {
863 return Err(ArtifactError::InvalidArtifact(
864 "SafeTensors checkpoint contains no tensors".into(),
865 ));
866 }
867 let architecture_plan = resolver.artifact_plan(
868 path,
869 ArtifactFormat::SafeTensors,
870 &configuration,
871 &tensors,
872 None,
873 resolved_plan,
874 )?;
875 Ok(ArtifactInspection {
876 path: path.to_path_buf(),
877 format: ArtifactFormat::SafeTensors,
878 configuration,
879 tensors,
880 safetensors_shards: Some(shards),
881 validated_gguf: None,
882 architecture_plan,
883 })
884}
885
886#[derive(Deserialize)]
887struct RawSafetensorInfo {
888 dtype: String,
889 shape: Vec<usize>,
890 data_offsets: [u64; 2],
891}
892
893fn inspect_safetensors_header(path: &Path) -> Result<Vec<TensorDescriptor>, ArtifactError> {
894 const MAX_HEADER_BYTES: u64 = 100_000_000;
895 let mut file = File::open(path)?;
896 let file_len = file.metadata()?.len();
897 let mut length = [0_u8; 8];
898 file.read_exact(&mut length)?;
899 let header_len = u64::from_le_bytes(length);
900 if header_len > MAX_HEADER_BYTES {
901 return Err(ArtifactError::InvalidArtifact(format!(
902 "SafeTensors header in {} exceeds {MAX_HEADER_BYTES} bytes",
903 path.display()
904 )));
905 }
906 let mut header = vec![
907 0_u8;
908 usize::try_from(header_len).map_err(|_| {
909 ArtifactError::InvalidArtifact("SafeTensors header length overflows usize".into())
910 })?
911 ];
912 file.read_exact(&mut header)?;
913 let raw: BTreeMap<String, Value> = serde_json::from_slice(&header)?;
914 let payload_start = 8_u64
915 .checked_add(header_len)
916 .ok_or_else(|| ArtifactError::InvalidArtifact("SafeTensors offset overflow".into()))?;
917 let mut entries = raw
918 .into_iter()
919 .filter(|(name, _)| name != "__metadata__")
920 .map(|(name, value)| {
921 serde_json::from_value::<RawSafetensorInfo>(value).map(|info| (name, info))
922 })
923 .collect::<Result<Vec<_>, _>>()?;
924 entries.sort_by_key(|(_, info)| info.data_offsets[0]);
925 let mut output = Vec::with_capacity(entries.len());
926 let mut expected_offset = 0_u64;
927 for (name, info) in entries {
928 if info.shape.contains(&0) {
931 return Err(ArtifactError::InvalidArtifact(format!(
932 "SafeTensors tensor {name:?} has an invalid shape"
933 )));
934 }
935 let [start, end] = info.data_offsets;
936 if start != expected_offset || end < start {
937 return Err(ArtifactError::InvalidArtifact(format!(
938 "SafeTensors tensor {name:?} has non-contiguous data offsets"
939 )));
940 }
941 expected_offset = end;
942 let absolute = payload_start
943 .checked_add(start)
944 .ok_or_else(|| ArtifactError::InvalidArtifact("SafeTensors offset overflow".into()))?;
945 output.push(TensorDescriptor {
946 name,
947 shape: info.shape,
948 dtype: safetensors_dtype(&info.dtype),
949 storage: Some(TensorStorage {
950 member: path.display().to_string(),
951 offset: absolute,
952 length: end - start,
953 }),
954 });
955 }
956 if payload_start
957 .checked_add(expected_offset)
958 .ok_or_else(|| ArtifactError::InvalidArtifact("SafeTensors length overflow".into()))?
959 != file_len
960 {
961 return Err(ArtifactError::InvalidArtifact(format!(
962 "SafeTensors payload length does not match header in {}",
963 path.display()
964 )));
965 }
966 Ok(output)
967}
968
969fn safetensors_dtype(dtype: &str) -> TensorDtype {
970 match dtype {
971 "BOOL" => TensorDtype::Bool,
972 "U8" => TensorDtype::U8,
973 "I8" => TensorDtype::I8,
974 "I16" => TensorDtype::I16,
975 "U16" => TensorDtype::U16,
976 "F16" => TensorDtype::F16,
977 "BF16" => TensorDtype::Bf16,
978 "I32" => TensorDtype::I32,
979 "U32" => TensorDtype::U32,
980 "F32" => TensorDtype::F32,
981 "F64" => TensorDtype::F64,
982 "I64" => TensorDtype::I64,
983 "U64" => TensorDtype::U64,
984 "C64" => TensorDtype::Complex64,
985 other => TensorDtype::Encoded(other.into()),
986 }
987}
988
989fn gguf_dtype(dtype: GgmlType) -> TensorDtype {
990 match dtype {
991 GgmlType::F32 => TensorDtype::F32,
992 GgmlType::F16 => TensorDtype::F16,
993 GgmlType::Bf16 => TensorDtype::Bf16,
994 GgmlType::I8 => TensorDtype::I8,
995 GgmlType::I16 => TensorDtype::I16,
996 GgmlType::I32 => TensorDtype::I32,
997 GgmlType::I64 => TensorDtype::I64,
998 GgmlType::F64 => TensorDtype::F64,
999 encoded => TensorDtype::Encoded(format!("{encoded:?}")),
1000 }
1001}
1002
1003fn is_gguf(path: &Path) -> bool {
1004 path.extension()
1005 .and_then(|extension| extension.to_str())
1006 .is_some_and(|extension| extension.eq_ignore_ascii_case("gguf"))
1007}
1008
1009#[derive(Debug, thiserror::Error)]
1011#[non_exhaustive]
1012pub enum ArtifactError {
1013 #[error("model artifact does not exist: {0}")]
1015 MissingArtifact(PathBuf),
1016 #[error("model artifact must be a SafeTensors directory or .gguf file: {0}")]
1018 UnsupportedContainer(PathBuf),
1019 #[error("unsupported model type: {0}")]
1021 UnsupportedModelType(String),
1022 #[error("unsupported GGUF architecture: {0}")]
1024 UnsupportedGgufArchitecture(String),
1025 #[error("GGUF metadata is missing string key \"general.architecture\"")]
1027 MissingGgufArchitecture,
1028 #[error(
1030 "required GGUF companion {role:?} matching {filename_prefix:?} was not found in {searched}",
1031 searched = display_directories(.searched_directories)
1032 )]
1033 MissingRequiredGgufCompanion {
1034 role: GgufCompanionRole,
1036 filename_prefix: String,
1038 searched_directories: Vec<PathBuf>,
1040 },
1041 #[error("invalid model artifact: {0}")]
1043 InvalidArtifact(String),
1044 #[error("invalid architecture artifact plan: {0}")]
1046 InvalidArchitecturePlan(String),
1047 #[error("duplicate checkpoint tensor {0:?}")]
1049 DuplicateTensor(String),
1050 #[error(transparent)]
1052 SafetensorsShards(#[from] eredu_checkpoint::safetensors::SafetensorsShardError),
1053 #[error("unsupported model quantization policy: {0}")]
1055 UnsupportedQuantizationPolicy(String),
1056 #[error("unsupported model residency policy: {0}")]
1058 UnsupportedResidencyPolicy(String),
1059 #[error("model artifact requires the {0:?} loading protocol")]
1061 UnsupportedLoadingProtocol(LoadingProtocol),
1062 #[error(transparent)]
1064 Io(#[from] std::io::Error),
1065 #[error(transparent)]
1067 Json(#[from] serde_json::Error),
1068 #[error(transparent)]
1070 Gguf(#[from] eredu_gguf::Error),
1071 #[error(transparent)]
1073 Catalog(#[from] crate::checkpoint::CatalogError),
1074}
1075
1076#[cfg(test)]
1077mod tests {
1078 use super::*;
1079 use eredu_gguf::{GgmlType, MetadataArray, TensorInput, Writer};
1080 use std::io::Write;
1081
1082 struct FixtureResolver;
1083
1084 #[derive(Debug, Clone, Default, Eq, PartialEq)]
1085 struct FixtureArtifactPlan {
1086 format: Option<ArtifactFormat>,
1087 }
1088
1089 impl ModelConfigurationResolver for FixtureResolver {
1090 type ArtifactPlan = FixtureArtifactPlan;
1091
1092 fn resolve_safetensors(
1093 &self,
1094 json: &Value,
1095 ) -> Result<ResolvedModelConfiguration<Self::ArtifactPlan>, ArtifactError> {
1096 let model_type = json
1097 .get("model_type")
1098 .and_then(Value::as_str)
1099 .ok_or_else(|| ArtifactError::InvalidArtifact("missing model_type".into()))?;
1100 let family = match model_type {
1101 "llama" => "llama",
1102 "gemma4" => "gemma4",
1103 "future" => "future_family",
1104 other => return Err(ArtifactError::UnsupportedModelType(other.into())),
1105 };
1106 Ok(ResolvedModelConfiguration::new(
1107 ModelConfiguration {
1108 declared_model_type: model_type.into(),
1109 effective_model_type: model_type.into(),
1110 family: family.into(),
1111 loading_protocol: LoadingProtocol::Model,
1112 json: Some(json.clone()),
1113 },
1114 FixtureArtifactPlan::default(),
1115 ))
1116 }
1117
1118 fn resolve_gguf(
1119 &self,
1120 architecture: &str,
1121 _checkpoint: &GgufCheckpoint,
1122 ) -> Result<ResolvedModelConfiguration<Self::ArtifactPlan>, ArtifactError> {
1123 let family = match architecture {
1124 "llama" => "llama",
1125 "future" => "future_family",
1126 other => return Err(ArtifactError::UnsupportedGgufArchitecture(other.into())),
1127 };
1128 Ok(ResolvedModelConfiguration::new(
1129 ModelConfiguration {
1130 declared_model_type: architecture.into(),
1131 effective_model_type: architecture.into(),
1132 family: family.into(),
1133 loading_protocol: LoadingProtocol::Model,
1134 json: None,
1135 },
1136 FixtureArtifactPlan::default(),
1137 ))
1138 }
1139
1140 fn gguf_companion_requirements(
1141 &self,
1142 architecture: &str,
1143 _checkpoint: &GgufCheckpoint,
1144 ) -> Result<Vec<GgufCompanionRequirement>, ArtifactError> {
1145 if architecture == "future" {
1146 return Ok(vec![GgufCompanionRequirement::new(
1147 GgufCompanionRole::MediaProjector,
1148 false,
1149 "mmproj",
1150 0,
1151 GgufCompanionEncoding::DensePreferred,
1152 )?]);
1153 }
1154 Ok(Vec::new())
1155 }
1156
1157 fn artifact_plan(
1158 &self,
1159 _path: &Path,
1160 format: ArtifactFormat,
1161 _configuration: &ModelConfiguration,
1162 _tensors: &TensorCatalog,
1163 _validated_gguf: Option<&ValidatedGguf>,
1164 _resolved_plan: Self::ArtifactPlan,
1165 ) -> Result<Self::ArtifactPlan, ArtifactError> {
1166 Ok(FixtureArtifactPlan {
1167 format: Some(format),
1168 })
1169 }
1170 }
1171
1172 fn write_safetensors_fixture(root: &Path, model_type: &str) {
1173 std::fs::write(
1174 root.join("config.json"),
1175 format!(r#"{{"model_type":"{model_type}"}}"#),
1176 )
1177 .unwrap();
1178 let header =
1179 br#"{"token_embd.weight":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#;
1180 let mut file = File::create(root.join("model.safetensors")).unwrap();
1181 file.write_all(&(header.len() as u64).to_le_bytes())
1182 .unwrap();
1183 file.write_all(header).unwrap();
1184 file.write_all(&[0_u8; 16]).unwrap();
1185 }
1186
1187 fn write_gguf_fixture(path: &Path, ggml_type: GgmlType) {
1188 let metadata = BTreeMap::from([(
1189 "general.architecture".into(),
1190 MetadataValue::String("clip".into()),
1191 )]);
1192 let (dimensions, data) = match ggml_type {
1193 GgmlType::F32 => (vec![1], 1.0_f32.to_le_bytes().to_vec()),
1194 GgmlType::Q8_0 => (vec![32], vec![0_u8; 34]),
1195 other => panic!("unsupported fixture encoding {other:?}"),
1196 };
1197 Writer::default()
1198 .write(
1199 File::create(path).unwrap(),
1200 &metadata,
1201 &[TensorInput {
1202 name: "projector.weight",
1203 dimensions: &dimensions,
1204 ggml_type,
1205 data: &data,
1206 }],
1207 )
1208 .unwrap();
1209 }
1210
1211 #[test]
1212 fn companion_planning_selects_by_catalog_encoding_not_filename() {
1213 let root = tempfile::tempdir().unwrap();
1214 let primary = root.path().join("model.gguf");
1215 write_gguf_fixture(&primary, GgmlType::F32);
1216 let quantized_name = root.path().join("mmproj-f16.gguf");
1217 let dense_name = root.path().join("mmproj-q4_k.gguf");
1218 write_gguf_fixture(&quantized_name, GgmlType::Q8_0);
1219 write_gguf_fixture(&dense_name, GgmlType::F32);
1220 let requirement = GgufCompanionRequirement::new(
1221 GgufCompanionRole::MediaProjector,
1222 true,
1223 "mmproj",
1224 1,
1225 GgufCompanionEncoding::DensePreferred,
1226 )
1227 .unwrap();
1228
1229 let companions = resolve_gguf_companions(&primary, &[requirement]).unwrap();
1230
1231 assert_eq!(
1232 companions
1233 .get(&GgufCompanionRole::MediaProjector)
1234 .unwrap()
1235 .path(),
1236 dense_name
1237 );
1238 }
1239
1240 #[test]
1241 fn dense_only_and_required_companion_policies_fail_closed() {
1242 let root = tempfile::tempdir().unwrap();
1243 let primary = root.path().join("model.gguf");
1244 write_gguf_fixture(&primary, GgmlType::F32);
1245 write_gguf_fixture(&root.path().join("mmproj.gguf"), GgmlType::Q8_0);
1246 let optional = GgufCompanionRequirement::new(
1247 GgufCompanionRole::MediaProjector,
1248 false,
1249 "mmproj",
1250 0,
1251 GgufCompanionEncoding::DenseRequired,
1252 )
1253 .unwrap();
1254 assert!(resolve_gguf_companions(&primary, &[optional]).is_err());
1255 let required = GgufCompanionRequirement::new(
1256 GgufCompanionRole::MediaProjector,
1257 true,
1258 "mmproj",
1259 0,
1260 GgufCompanionEncoding::DenseRequired,
1261 )
1262 .unwrap();
1263 assert!(resolve_gguf_companions(&primary, &[required]).is_err());
1264 }
1265
1266 #[test]
1267 fn missing_required_companion_preserves_its_semantic_role() {
1268 let root = tempfile::tempdir().unwrap();
1269 let primary = root.path().join("model.gguf");
1270 write_gguf_fixture(&primary, GgmlType::F32);
1271 let requirement = GgufCompanionRequirement::new(
1272 GgufCompanionRole::MediaProjector,
1273 true,
1274 "mmproj",
1275 0,
1276 GgufCompanionEncoding::DensePreferred,
1277 )
1278 .unwrap();
1279
1280 let error = resolve_gguf_companions(&primary, &[requirement]).unwrap_err();
1281
1282 assert!(matches!(
1283 error,
1284 ArtifactError::MissingRequiredGgufCompanion {
1285 role: GgufCompanionRole::MediaProjector,
1286 filename_prefix,
1287 searched_directories,
1288 } if filename_prefix == "mmproj" && searched_directories == [root.path()]
1289 ));
1290 }
1291
1292 #[test]
1293 fn loading_protocol_is_family_agnostic() {
1294 assert!(matches!(
1295 validate_preparation_policy(LoadingProtocol::Realtime, PreparationPolicy::default()),
1296 Err(ArtifactError::UnsupportedLoadingProtocol(
1297 LoadingProtocol::Realtime
1298 ))
1299 ));
1300 }
1301
1302 #[test]
1303 fn gguf_u32_metadata_is_lossless_and_fail_closed() {
1304 let values = MetadataValue::Array(MetadataArray::Uint64(vec![0, u32::MAX.into()]));
1305 assert_eq!(
1306 gguf_u32_metadata_values("tokenizer.ids", Some(&values)).unwrap(),
1307 vec![0, u32::MAX]
1308 );
1309 assert!(gguf_u32_metadata_values(
1310 "tokenizer.ids",
1311 Some(&MetadataValue::Uint64(u64::from(u32::MAX) + 1))
1312 )
1313 .is_err());
1314 assert!(
1315 gguf_u32_metadata_values("tokenizer.ids", Some(&MetadataValue::Int32(-1))).is_err()
1316 );
1317 assert!(gguf_u32_metadata_values(
1318 "tokenizer.ids",
1319 Some(&MetadataValue::String("1".into()))
1320 )
1321 .is_err());
1322 assert!(gguf_u32_metadata_values("tokenizer.ids", None)
1323 .unwrap()
1324 .is_empty());
1325 }
1326
1327 #[test]
1328 fn safetensors_inspection_and_planning_are_backend_neutral() {
1329 let root = tempfile::tempdir().unwrap();
1330 write_safetensors_fixture(root.path(), "llama");
1331 let inspection = inspect_artifact(root.path(), &FixtureResolver).unwrap();
1332 assert_eq!(inspection.configuration().family, "llama");
1333 assert_eq!(inspection.tensors().len(), 1);
1334 assert_eq!(
1335 inspection
1336 .safetensors_shards()
1337 .unwrap()
1338 .payload_paths()
1339 .len(),
1340 1
1341 );
1342 let plan = plan_model_preparation(
1343 inspection,
1344 PreparationPolicy::default(),
1345 crate::backend::SessionCapabilities::default(),
1346 )
1347 .unwrap();
1348 assert_eq!(plan.route(), MaterializationRoute::Resident);
1349 let architecture_plan = plan.inspection().architecture_plan().clone();
1350 let artifact = plan.into_artifact();
1351 let ModelArtifact::SafeTensors { shards, .. } = artifact else {
1352 panic!("expected SafeTensors artifact");
1353 };
1354 assert_eq!(shards.payload_paths().len(), 1);
1355 assert_eq!(architecture_plan.format, Some(ArtifactFormat::SafeTensors));
1356 }
1357
1358 #[test]
1359 fn safetensors_inspection_rejects_index_entries_missing_from_their_shard() {
1360 let root = tempfile::tempdir().unwrap();
1361 write_safetensors_fixture(root.path(), "llama");
1362 std::fs::rename(
1363 root.path().join("model.safetensors"),
1364 root.path().join("model-00001.safetensors"),
1365 )
1366 .unwrap();
1367 std::fs::write(
1368 root.path().join("model.safetensors.index.json"),
1369 r#"{"weight_map":{"missing.weight":"model-00001.safetensors"}}"#,
1370 )
1371 .unwrap();
1372
1373 assert!(matches!(
1374 inspect_artifact(root.path(), &FixtureResolver),
1375 Err(ArtifactError::SafetensorsShards(
1376 eredu_checkpoint::safetensors::SafetensorsShardError::MalformedIndex { .. }
1377 ))
1378 ));
1379 }
1380
1381 #[cfg(unix)]
1382 #[test]
1383 fn safetensors_inspection_rejects_indexed_symlinks_outside_the_access_root() {
1384 use std::os::unix::fs::symlink;
1385
1386 let parent = tempfile::tempdir().unwrap();
1387 let outside = parent.path().join("outside");
1388 std::fs::create_dir(&outside).unwrap();
1389 write_safetensors_fixture(&outside, "llama");
1390
1391 let checkpoint = parent.path().join("checkpoint");
1392 std::fs::create_dir(&checkpoint).unwrap();
1393 std::fs::write(checkpoint.join("config.json"), r#"{"model_type":"llama"}"#).unwrap();
1394 symlink(
1395 outside.join("model.safetensors"),
1396 checkpoint.join("model-00001.safetensors"),
1397 )
1398 .unwrap();
1399 std::fs::write(
1400 checkpoint.join("model.safetensors.index.json"),
1401 r#"{"weight_map":{"token_embd.weight":"model-00001.safetensors"}}"#,
1402 )
1403 .unwrap();
1404
1405 assert!(matches!(
1406 inspect_artifact(&checkpoint, &FixtureResolver),
1407 Err(ArtifactError::SafetensorsShards(
1408 eredu_checkpoint::safetensors::SafetensorsShardError::UnsafeShardPath { .. }
1409 ))
1410 ));
1411 }
1412
1413 #[test]
1414 fn safetensors_native_dtypes_remain_typed_in_the_portable_catalog() {
1415 assert_eq!(safetensors_dtype("BOOL"), TensorDtype::Bool);
1416 assert_eq!(safetensors_dtype("I64"), TensorDtype::I64);
1417 assert_eq!(safetensors_dtype("U32"), TensorDtype::U32);
1418 assert_eq!(safetensors_dtype("F64"), TensorDtype::F64);
1419 assert_eq!(safetensors_dtype("C64"), TensorDtype::Complex64);
1420 assert_eq!(
1421 safetensors_dtype("F8_E4M3"),
1422 TensorDtype::Encoded("F8_E4M3".into())
1423 );
1424 }
1425
1426 #[test]
1427 fn gguf_dense_dtypes_remain_typed_in_the_portable_catalog() {
1428 assert_eq!(gguf_dtype(GgmlType::F16), TensorDtype::F16);
1429 assert_eq!(gguf_dtype(GgmlType::Bf16), TensorDtype::Bf16);
1430 assert_eq!(gguf_dtype(GgmlType::F32), TensorDtype::F32);
1431 assert_eq!(
1432 gguf_dtype(GgmlType::Q4K),
1433 TensorDtype::Encoded("Q4K".into())
1434 );
1435 }
1436
1437 #[test]
1438 fn core_accepts_families_defined_only_by_the_resolver() {
1439 let root = tempfile::tempdir().unwrap();
1440 write_safetensors_fixture(root.path(), "future");
1441 let inspection = inspect_artifact(root.path(), &FixtureResolver).unwrap();
1442 assert_eq!(inspection.configuration().family, "future_family");
1443 assert_eq!(
1444 inspection.configuration().loading_protocol,
1445 LoadingProtocol::Model
1446 );
1447 assert!(plan_model_preparation(
1448 inspection,
1449 PreparationPolicy::default(),
1450 crate::backend::SessionCapabilities::default(),
1451 )
1452 .is_ok());
1453 }
1454
1455 #[test]
1456 fn safetensors_inspection_accepts_rank_zero_scalar_parameters() {
1457 let root = tempfile::tempdir().unwrap();
1458 std::fs::write(
1459 root.path().join("config.json"),
1460 r#"{"model_type":"gemma4"}"#,
1461 )
1462 .unwrap();
1463 let header = br#"{"clip.output_max":{"dtype":"F32","shape":[],"data_offsets":[0,4]}}"#;
1464 let mut file = File::create(root.path().join("model.safetensors")).unwrap();
1465 file.write_all(&(header.len() as u64).to_le_bytes())
1466 .unwrap();
1467 file.write_all(header).unwrap();
1468 file.write_all(&0.0_f32.to_le_bytes()).unwrap();
1469
1470 let inspection = inspect_artifact(root.path(), &FixtureResolver).unwrap();
1471 assert_eq!(
1472 inspection.tensors().get("clip.output_max").unwrap().shape,
1473 Vec::<usize>::new()
1474 );
1475 }
1476
1477 #[test]
1478 fn parallel_policy_binds_the_exact_neutral_topology() {
1479 let root = tempfile::tempdir().unwrap();
1480 write_safetensors_fixture(root.path(), "llama");
1481 let topology = crate::topology::ParallelTopology::new(2, 3, 4, 1).unwrap();
1482 let policy = PreparationPolicy {
1483 topology: Some(topology),
1484 ..PreparationPolicy::default()
1485 };
1486
1487 let plan = plan_model_preparation(
1488 inspect_artifact(root.path(), &FixtureResolver).unwrap(),
1489 policy,
1490 crate::backend::SessionCapabilities::default(),
1491 )
1492 .unwrap();
1493
1494 assert_eq!(plan.policy(), policy);
1495 assert_eq!(plan.policy().topology, Some(topology));
1496 assert_eq!(plan.route(), MaterializationRoute::Resident);
1497 }
1498
1499 #[test]
1500 fn policy_leaves_expert_cache_capability_to_architecture_and_backend() {
1501 let root = tempfile::tempdir().unwrap();
1502 write_safetensors_fixture(root.path(), "llama");
1503 let plan = plan_model_preparation(
1504 inspect_artifact(root.path(), &FixtureResolver).unwrap(),
1505 PreparationPolicy {
1506 residency: ResidencyRequest::AddressableParameterBanks,
1507 ..PreparationPolicy::default()
1508 },
1509 crate::backend::SessionCapabilities::default(),
1510 )
1511 .unwrap();
1512 assert_eq!(
1513 plan.route(),
1514 MaterializationRoute::AddressableParameterBanks
1515 );
1516 }
1517
1518 #[test]
1519 fn policy_leaves_nonresident_quantization_capability_to_architecture_and_backend() {
1520 let root = tempfile::tempdir().unwrap();
1521 write_safetensors_fixture(root.path(), "llama");
1522 let policy = PreparationPolicy {
1523 quantization: Some(QuantizationRequest::MxFp4),
1524 residency: ResidencyRequest::LayerwiseHost,
1525 ..PreparationPolicy::default()
1526 };
1527 let plan = plan_model_preparation(
1528 inspect_artifact(root.path(), &FixtureResolver).unwrap(),
1529 policy,
1530 crate::backend::SessionCapabilities::default(),
1531 )
1532 .unwrap();
1533 assert_eq!(plan.policy(), policy);
1534 assert_eq!(plan.route(), MaterializationRoute::Layerwise);
1535 }
1536
1537 #[test]
1538 fn gguf_plan_owns_the_portable_checkpoint_for_later_materialization() {
1539 let root = tempfile::tempdir().unwrap();
1540 let path = root.path().join("model.gguf");
1541 let data = [0_u8; 2];
1542 let metadata = BTreeMap::from([
1543 (
1544 "general.architecture".into(),
1545 MetadataValue::String("llama".into()),
1546 ),
1547 ("llama.block_count".into(), MetadataValue::Uint32(1)),
1548 ("llama.embedding_length".into(), MetadataValue::Uint32(1)),
1549 ]);
1550 Writer::default()
1551 .write(
1552 File::create(&path).unwrap(),
1553 &metadata,
1554 &[TensorInput {
1555 name: "token_embd.weight",
1556 dimensions: &[1],
1557 ggml_type: GgmlType::F16,
1558 data: &data,
1559 }],
1560 )
1561 .unwrap();
1562
1563 let inspection = inspect_artifact(&path, &FixtureResolver).unwrap();
1564 let validated = inspection.validated_gguf().unwrap();
1565 assert_eq!(validated.checkpoint().physical_tensor_count(), 1);
1566 assert_eq!(inspection.configuration().declared_model_type, "llama");
1567 assert_eq!(
1568 inspection.tensors().get("token_embd.weight").unwrap().dtype,
1569 TensorDtype::F16
1570 );
1571
1572 let plan = plan_model_preparation(
1573 inspection,
1574 PreparationPolicy::default(),
1575 crate::backend::SessionCapabilities::default(),
1576 )
1577 .unwrap();
1578 let architecture_plan = plan.inspection().architecture_plan().clone();
1579 let artifact = plan.into_artifact();
1580 let ModelArtifact::Gguf {
1581 configuration,
1582 validated,
1583 ..
1584 } = artifact
1585 else {
1586 panic!("expected GGUF artifact");
1587 };
1588 assert_eq!(architecture_plan.format, Some(ArtifactFormat::Gguf));
1589 assert_eq!(configuration.family, "llama");
1590 assert_eq!(validated.checkpoint().physical_tensor_count(), 1);
1591 }
1592
1593 #[test]
1594 fn core_accepts_architecture_owned_gguf_schema() {
1595 let root = tempfile::tempdir().unwrap();
1596 let path = root.path().join("model.gguf");
1597 let data = 1.0_f32.to_le_bytes();
1598 let metadata = BTreeMap::from([
1599 (
1600 "general.architecture".into(),
1601 MetadataValue::String("future".into()),
1602 ),
1603 ("future.state_width".into(), MetadataValue::Uint32(1)),
1604 ]);
1605 Writer::default()
1606 .write(
1607 File::create(&path).unwrap(),
1608 &metadata,
1609 &[TensorInput {
1610 name: "state.in_proj",
1611 dimensions: &[1],
1612 ggml_type: GgmlType::F32,
1613 data: &data,
1614 }],
1615 )
1616 .unwrap();
1617
1618 let inspection = inspect_artifact(&path, &FixtureResolver).unwrap();
1619
1620 assert_eq!(inspection.configuration().family, "future_family");
1621 assert!(inspection.tensors().get("state.in_proj").is_some());
1622 }
1623
1624 #[test]
1625 fn preparation_plan_carries_the_exact_inspected_companion() {
1626 let root = tempfile::tempdir().unwrap();
1627 let primary = root.path().join("model.gguf");
1628 let scalar = 1.0_f32.to_le_bytes();
1629 Writer::default()
1630 .write(
1631 File::create(&primary).unwrap(),
1632 &BTreeMap::from([(
1633 "general.architecture".into(),
1634 MetadataValue::String("future".into()),
1635 )]),
1636 &[TensorInput {
1637 name: "state.in_proj",
1638 dimensions: &[1],
1639 ggml_type: GgmlType::F32,
1640 data: &scalar,
1641 }],
1642 )
1643 .unwrap();
1644 let projector = root.path().join("mmproj.gguf");
1645 write_gguf_fixture(&projector, GgmlType::F32);
1646
1647 let inspection = inspect_artifact(&primary, &FixtureResolver).unwrap();
1648 assert_eq!(
1649 inspection
1650 .validated_gguf()
1651 .unwrap()
1652 .companion(&GgufCompanionRole::MediaProjector)
1653 .unwrap()
1654 .path(),
1655 projector
1656 );
1657 let ModelArtifact::Gguf { validated, .. } = plan_model_preparation(
1658 inspection,
1659 PreparationPolicy::default(),
1660 crate::backend::SessionCapabilities::default(),
1661 )
1662 .unwrap()
1663 .into_artifact() else {
1664 panic!("expected GGUF preparation artifact");
1665 };
1666 assert_eq!(
1667 validated
1668 .companion(&GgufCompanionRole::MediaProjector)
1669 .unwrap()
1670 .path(),
1671 projector
1672 );
1673 }
1674}