diff --git a/CHANGELOG.md b/CHANGELOG.md index 63b4541..2fcd89c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,17 @@ All notable changes to this project are documented in this file. crates: `[dependencies]` stays empty for this feature; the upstream `safetensors` crate, `serde_json`, and `corinth-canal` are not used. Does not include payload mmap or Hugging Face `config.json` policy. +- **Safetensors shard-index invariants ([RM-1360](https://linear.app/rpd-34/issue/RM-1360)):** + index shard paths resolve relative to the checkpoint root; absolute paths + and `..` traversal that would leave the root are rejected; existing paths + are canonicalized to reject symlink escape where the platform can follow + links. Each indexed tensor must map to exactly one existing shard. + Duplicate/conflicting ownership fails with + `ParserError::DuplicateTensorOwnership` (tensor name + shard list). Missing + referenced shards fail with `ParserError::MissingShard` before a manifest is + returned. Unreferenced on-disk shards are reported via + `index:unreferenced_shards` rather than rejected. Directory listings and + tensor records are sorted so repeated inspection is byte-identical. - **Optional `mmap` feature (#45 option 1):** `load_gguf_mmap` maps a GGUF with `memmap2` 0.9.11 instead of `fs::read` into a `Vec`. Default builds stay zero-dep (`default = []`). Packed CPU dequant for **Q8_0**, **Q5_K**, diff --git a/README.md b/README.md index 1fe1d85..f5ebd6f 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,8 @@ The off-by-default `safetensors` feature owns reusable **metadata-side** support - Safetensors header deserialization; - deterministic tensor manifests; - single-file, Hugging Face shard-index, and directory layouts; +- checkpoint-relative shard resolution with path-escape rejection; +- unique tensor ownership and missing-shard diagnostics; - tensor name/dtype/shape/offset/shard metadata; - MoE router/expert candidate discovery and grouping; - metadata-only layout-family inference where it is reusable outside Corinth. @@ -233,7 +235,8 @@ Safetensors surface (`--features safetensors`) includes: - `SafetensorsManifest`, `SafetensorsCheckpointSource`, `SafetensorsTensorRecord`; - `classify_tensor`, `discover_candidates`; - `SafetensorsCandidateSummary`, `SafetensorsRouterCandidate`, `SafetensorsExpertGroup`; -- `dtype_size_bytes`. +- `dtype_size_bytes`; +- `ParserError::DuplicateTensorOwnership` and `ParserError::MissingShard` for shard-index diagnostics. ## Ecosystem / promotion model diff --git a/src/error.rs b/src/error.rs index 1d24feb..644ea3c 100644 --- a/src/error.rs +++ b/src/error.rs @@ -52,6 +52,22 @@ pub enum ParserError { /// Number of experts actually available. available: usize, }, + /// A tensor name is claimed by more than one Safetensors shard. + DuplicateTensorOwnership { + /// Tensor name with conflicting owners. + name: String, + /// Checkpoint-relative shard paths that declare the tensor, sorted. + shards: Vec, + /// Path of the checkpoint or index being inspected. + path: String, + }, + /// An index-referenced Safetensors shard is missing on disk. + MissingShard { + /// Shard path as declared in the index (checkpoint-relative). + shard: String, + /// Path of the index or checkpoint that referenced the shard. + path: String, + }, } impl fmt::Display for ParserError { @@ -75,6 +91,16 @@ impl fmt::Display for ParserError { f, "expert index out of range: block={block}, expert={expert}, available={available}" ), + Self::DuplicateTensorOwnership { name, shards, path } => { + write!( + f, + "duplicate tensor ownership for '{name}' across shards {} in '{path}'", + shards.join(", ") + ) + } + Self::MissingShard { shard, path } => { + write!(f, "missing shard '{shard}' referenced by '{path}'") + } } } } @@ -83,7 +109,12 @@ impl std::error::Error for ParserError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { Self::Io { source, .. } => Some(source), - _ => None, + Self::UnsupportedFormat { .. } + | Self::MissingTensor { .. } + | Self::InvalidLayout { .. } + | Self::ExpertOutOfRange { .. } + | Self::DuplicateTensorOwnership { .. } + | Self::MissingShard { .. } => None, } } } diff --git a/src/safetensors/manifest.rs b/src/safetensors/manifest.rs index 2bec1e7..1d2715a 100644 --- a/src/safetensors/manifest.rs +++ b/src/safetensors/manifest.rs @@ -13,7 +13,7 @@ use super::relative_path; use super::validate::{ expected_tensor_byte_size, reject_output_checkpoint_conflict, reject_tensor_data_ranges, }; -use super::{io_error, model_load, unsupported}; +use super::{duplicate_tensor_ownership, io_error, model_load, unsupported}; use crate::error::Result; use std::collections::{BTreeMap, BTreeSet}; use std::fs::{self, File}; @@ -24,6 +24,9 @@ pub(super) const SAFETENSORS_EXTENSION: &str = "safetensors"; pub(super) const MAX_HEADER_BYTES: usize = 64 * 1024 * 1024; pub(super) const MAX_INDEX_BYTES: u64 = 64 * 1024 * 1024; /// Reserved index metadata key for shards present on disk but not in `weight_map`. +/// Unreferenced shards are reported here rather than rejected. Missing referenced +/// shards fail closed with [`crate::ParserError::MissingShard`] before a manifest +/// is returned. pub const INDEX_UNREFERENCED_SHARDS_KEY: &str = "index:unreferenced_shards"; /// Unambiguous boundary between shard relative path and logical metadata key in /// `shard:*` manifest keys (avoids ambiguity when the logical key contains `:`). @@ -385,13 +388,14 @@ fn expert_group_json(group: &super::SafetensorsExpertGroup) -> JsonValue { pub(super) fn inspect_single_file(path: &Path) -> Result { let root = parent_or_current(path); let shard = inspect_shard(path, root)?; - Ok(build_manifest( + build_manifest( "single_file", None, vec![path.to_path_buf()], shard.metadata, shard.tensors, - )) + path, + ) } pub(super) fn inspect_directory(root: &Path) -> Result { @@ -462,18 +466,31 @@ pub(super) fn inspect_index_shards( index_tensor_count: usize, unreferenced_shards_json: Option, ) -> Result { - let mut tensors = Vec::new(); + let tensor_owners = expected_tensor_owners(&expected_by_shard); + let mut inspections = Vec::new(); for shard_path in &shards { - let expected = expected_by_shard.get(shard_path).ok_or_else(|| { - model_load( + if !expected_by_shard.contains_key(shard_path) { + return Err(model_load( shard_path, "internal error: index shard has no expected tensor set".into(), + )); + } + inspections.push((shard_path.clone(), inspect_shard(shard_path, root)?)); + } + + reject_indexed_tensor_ownership(root, &tensor_owners, &inspections)?; + + let mut tensors = Vec::new(); + for (shard_path, shard) in inspections { + let expected = expected_by_shard.get(&shard_path).ok_or_else(|| { + model_load( + &shard_path, + "internal error: index shard has no expected tensor set".into(), ) })?; - let shard = inspect_shard(shard_path, root)?; merge_shard_metadata( &mut metadata, - &relative_path(shard_path, root), + &relative_path(&shard_path, root), shard.metadata, ); @@ -484,7 +501,7 @@ pub(super) fn inspect_index_shards( .collect::>(); if let Some(missing) = expected.difference(&found).next() { return Err(model_load( - shard_path, + &shard_path, format!( "index maps tensor '{missing}' to this shard, but the shard header does not contain it" ), @@ -505,9 +522,7 @@ pub(super) fn inspect_index_shards( metadata.insert(INDEX_UNREFERENCED_SHARDS_KEY.into(), encoded); } - Ok(build_manifest( - "hf_index", index_file, shards, metadata, tensors, - )) + build_manifest("hf_index", index_file, shards, metadata, tensors, root) } pub(super) fn inspect_shards( @@ -529,9 +544,64 @@ pub(super) fn inspect_shards( tensors.extend(shard.tensors); } - Ok(build_manifest( - input_kind, index_file, shards, metadata, tensors, - )) + build_manifest(input_kind, index_file, shards, metadata, tensors, root) +} + +fn expected_tensor_owners( + expected_by_shard: &BTreeMap>, +) -> BTreeMap { + let mut owners = BTreeMap::new(); + for (shard_path, names) in expected_by_shard { + for name in names { + owners.insert(name.clone(), shard_path.clone()); + } + } + owners +} + +fn reject_indexed_tensor_ownership( + root: &Path, + tensor_owners: &BTreeMap, + inspections: &[(PathBuf, ShardInspection)], +) -> Result<()> { + let mut found_in: BTreeMap> = BTreeMap::new(); + for (shard_path, shard) in inspections { + let relative = relative_path(shard_path, root); + for tensor in &shard.tensors { + if tensor_owners.contains_key(&tensor.name) { + found_in + .entry(tensor.name.clone()) + .or_default() + .insert(relative.clone()); + } + } + } + + for (name, expected_path) in tensor_owners { + let Some(owners) = found_in.get(name) else { + continue; + }; + if owners.len() > 1 { + return Err(duplicate_tensor_ownership( + root, + name.clone(), + owners.iter().cloned().collect(), + )); + } + let expected_rel = relative_path(expected_path, root); + if !owners.contains(&expected_rel) { + return Err(duplicate_tensor_ownership( + root, + name.clone(), + owners + .iter() + .cloned() + .chain(std::iter::once(expected_rel)) + .collect(), + )); + } + } + Ok(()) } pub(super) fn build_manifest( @@ -540,15 +610,17 @@ pub(super) fn build_manifest( shards: Vec, metadata: BTreeMap, mut tensors: Vec, -) -> SafetensorsManifest { + error_path: &Path, +) -> Result { tensors.sort_by(|left, right| { left.name .cmp(&right.name) .then(left.source_shard.cmp(&right.source_shard)) }); + reject_duplicate_tensor_ownership(error_path, &tensors)?; let candidates = discover_candidates(&tensors); - SafetensorsManifest { + Ok(SafetensorsManifest { manifest_version: 2, format: "safetensors", checkpoint: SafetensorsCheckpointSource { @@ -560,7 +632,28 @@ pub(super) fn build_manifest( }, tensors, candidates, + }) +} + +fn reject_duplicate_tensor_ownership( + path: &Path, + tensors: &[SafetensorsTensorRecord], +) -> Result<()> { + let mut owners: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new(); + for tensor in tensors { + owners + .entry(tensor.name.as_str()) + .or_default() + .insert(tensor.source_shard.as_str()); + } + if let Some((name, shards)) = owners.into_iter().find(|(_, shards)| shards.len() > 1) { + return Err(duplicate_tensor_ownership( + path, + name, + shards.into_iter().map(str::to_string).collect(), + )); } + Ok(()) } pub(super) fn inspect_shard(path: &Path, root: &Path) -> Result { @@ -805,6 +898,7 @@ pub(super) fn read_dir_paths(root: &Path) -> Result> { let entry = entry.map_err(|e| model_load(root, format!("read directory entry: {e}")))?; paths.push(entry.path()); } + paths.sort(); Ok(paths) } diff --git a/src/safetensors/mod.rs b/src/safetensors/mod.rs index c909af0..21ef75c 100644 --- a/src/safetensors/mod.rs +++ b/src/safetensors/mod.rs @@ -75,3 +75,24 @@ pub(super) fn unsupported(path: &Path, reason: String) -> ParserError { reason, } } + +pub(super) fn duplicate_tensor_ownership( + path: &Path, + name: impl Into, + mut shards: Vec, +) -> ParserError { + shards.sort(); + shards.dedup(); + ParserError::DuplicateTensorOwnership { + name: name.into(), + shards, + path: path.display().to_string(), + } +} + +pub(super) fn missing_shard(path: &Path, shard: impl Into) -> ParserError { + ParserError::MissingShard { + shard: shard.into(), + path: path.display().to_string(), + } +} diff --git a/src/safetensors/paths.rs b/src/safetensors/paths.rs index 68e6445..fdd4d33 100644 --- a/src/safetensors/paths.rs +++ b/src/safetensors/paths.rs @@ -1,9 +1,10 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 //! Path validation and same-file checks for Safetensors shards. -use super::{io_error, model_load}; +use super::{io_error, missing_shard, model_load}; use crate::error::Result; use std::fs; +use std::io::ErrorKind; #[cfg(unix)] use std::os::unix::fs::MetadataExt; use std::path::{Component, Path, PathBuf}; @@ -181,31 +182,68 @@ pub(super) fn is_safetensors_index(path: &Path) -> bool { .is_some_and(|name| name.ends_with(".safetensors.index.json")) } +/// Resolve an index `weight_map` shard reference against `root`. +/// +/// Absolute paths and `..` segments that would leave `root` are rejected +/// lexically. Remaining `..` / `.` segments are normalized so +/// `nested/../model.safetensors` stays portable and checkpoint-relative. +/// Missing targets fail with [`crate::ParserError::MissingShard`]. Existing +/// paths are canonicalized so symlink escape is rejected where the platform +/// can follow links. pub(super) fn index_shard_path(root: &Path, index_path: &Path, relative: &str) -> Result { let relative_path = Path::new(relative); - let mut normalized = PathBuf::new(); - let escapes_root = relative_path.is_absolute() - || relative_path.components().any(|component| match component { - Component::Normal(part) => { - normalized.push(part); - false - } - Component::CurDir => false, - Component::ParentDir | Component::RootDir | Component::Prefix(_) => true, - }); - if escapes_root { + if relative_path.is_absolute() { return Err(model_load( index_path, - format!("index shard path '{relative}' must stay within the checkpoint directory"), + format!("index shard path '{relative}' must be checkpoint-relative, not absolute"), )); } + + let mut normalized = PathBuf::new(); + for component in relative_path.components() { + match component { + Component::Normal(part) => normalized.push(part), + Component::CurDir => {} + Component::ParentDir => { + if !normalized.pop() { + return Err(model_load( + index_path, + format!( + "index shard path '{relative}' must stay within the checkpoint directory" + ), + )); + } + } + Component::RootDir | Component::Prefix(_) => { + return Err(model_load( + index_path, + format!( + "index shard path '{relative}' must be checkpoint-relative, not absolute" + ), + )); + } + } + } if normalized.as_os_str().is_empty() { return Err(model_load( index_path, format!("index shard path '{relative}' must name a Safetensors shard"), )); } - let candidate = root.join(normalized); + let candidate = root.join(&normalized); + match fs::metadata(&candidate) { + Err(err) if err.kind() == ErrorKind::NotFound => { + return Err(missing_shard(index_path, relative)); + } + Err(err) => return Err(io_error(&candidate, err)), + Ok(metadata) if !metadata.is_file() => { + return Err(model_load( + index_path, + format!("index shard path '{relative}' must name a regular Safetensors file"), + )); + } + Ok(_) => {} + } validate_path_stays_under_root(root, &candidate)?; Ok(candidate) } diff --git a/src/safetensors/tests.rs b/src/safetensors/tests.rs index 47d21d2..27659f0 100644 --- a/src/safetensors/tests.rs +++ b/src/safetensors/tests.rs @@ -6,6 +6,7 @@ use super::manifest::{ }; use super::paths::{canonical_existing_or_parent, parent_or_current}; use super::*; +use crate::ParserError; use std::collections::BTreeMap; use std::fs::{self, File}; use std::io::Write; @@ -231,6 +232,183 @@ fn rejects_index_shard_paths_that_escape_checkpoint_directory() { ); } +#[test] +fn rejects_absolute_index_shard_paths() { + let dir = temp_dir("absolute"); + let outside = temp_dir("absolute-outside"); + let outside_shard = outside.join("outside.safetensors"); + write_safetensors( + &outside_shard, + r#"{"a.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + let escaped = outside_shard.to_string_lossy().replace('\\', "\\\\"); + fs::write( + dir.join("model.safetensors.index.json"), + format!(r#"{{"weight_map": {{"a.weight": "{escaped}"}}}}"#), + ) + .unwrap(); + + let err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + assert!( + err.to_string() + .contains("must be checkpoint-relative, not absolute") + ); +} + +#[test] +fn accepts_parent_dir_index_paths_that_stay_within_checkpoint_root() { + let dir = temp_dir("nested-parent"); + write_safetensors( + &dir.join("model.safetensors"), + r#"{"a.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + fs::write( + dir.join("model.safetensors.index.json"), + r#"{ + "weight_map": { + "a.weight": "nested/../model.safetensors" + } + }"#, + ) + .unwrap(); + + let manifest = inspect_safetensors_checkpoint(&dir).unwrap(); + assert_eq!(manifest.checkpoint.shard_count, 1); + assert_eq!(manifest.tensors[0].source_shard, "model.safetensors"); +} + +#[test] +fn missing_referenced_shard_fails_before_manifest() { + let dir = temp_dir("missing-shard"); + fs::write( + dir.join("model.safetensors.index.json"), + r#"{ + "weight_map": { + "a.weight": "missing.safetensors" + } + }"#, + ) + .unwrap(); + + let err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + match err { + ParserError::MissingShard { shard, path } => { + assert_eq!(shard, "missing.safetensors"); + assert!(path.ends_with("model.safetensors.index.json")); + } + other => panic!("expected MissingShard, got {other}"), + } +} + +#[test] +fn duplicate_directory_tensor_ownership_is_typed() { + let dir = temp_dir("dup-dir"); + write_safetensors( + &dir.join("a.safetensors"), + r#"{"shared.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + write_safetensors( + &dir.join("b.safetensors"), + r#"{"shared.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + + let err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + match err { + ParserError::DuplicateTensorOwnership { name, shards, path } => { + assert_eq!(name, "shared.weight"); + assert_eq!( + shards, + vec!["a.safetensors".to_string(), "b.safetensors".to_string()] + ); + assert_eq!(Path::new(&path), dir.as_ref()); + } + other => panic!("expected DuplicateTensorOwnership, got {other}"), + } +} + +#[test] +fn duplicate_index_tensor_ownership_is_typed() { + let dir = temp_dir("dup-index"); + write_safetensors( + &dir.join("model-00001-of-00002.safetensors"), + r#"{"shared.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + write_safetensors( + &dir.join("model-00002-of-00002.safetensors"), + r#"{ + "shared.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}, + "other.weight": {"dtype": "F16", "shape": [1], "data_offsets": [2, 4]} + }"#, + 4, + ); + fs::write( + dir.join("model.safetensors.index.json"), + r#"{ + "weight_map": { + "shared.weight": "model-00001-of-00002.safetensors", + "other.weight": "model-00002-of-00002.safetensors" + } + }"#, + ) + .unwrap(); + + let err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + match err { + ParserError::DuplicateTensorOwnership { name, shards, .. } => { + assert_eq!(name, "shared.weight"); + assert_eq!( + shards, + vec![ + "model-00001-of-00002.safetensors".to_string(), + "model-00002-of-00002.safetensors".to_string() + ] + ); + } + other => panic!("expected DuplicateTensorOwnership, got {other}"), + } +} + +#[test] +fn shuffled_directory_creation_order_yields_identical_manifests() { + let first = temp_dir("shuffle-a"); + write_safetensors( + &first.join("z.safetensors"), + r#"{"z.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + write_safetensors( + &first.join("a.safetensors"), + r#"{"a.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + + let second = temp_dir("shuffle-b"); + write_safetensors( + &second.join("a.safetensors"), + r#"{"a.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + write_safetensors( + &second.join("z.safetensors"), + r#"{"z.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + + let manifest_a = inspect_safetensors_checkpoint(&first).unwrap(); + let manifest_b = inspect_safetensors_checkpoint(&second).unwrap(); + let json_a = manifest_a.to_pretty_json(); + let json_b = manifest_b.to_pretty_json(); + assert_eq!(json_a, json_b); + assert_eq!(json_a, manifest_a.to_pretty_json()); + assert_eq!(manifest_a.tensors[0].name, "a.weight"); + assert_eq!(manifest_a.tensors[1].name, "z.weight"); +} + #[cfg(unix)] #[test] fn rejects_index_shard_paths_that_escape_via_symlink() { @@ -283,6 +461,49 @@ fn rejects_directory_index_that_escapes_via_symlink() { ); } +/// Non-unix targets cannot portably create out-of-root symlinks (Windows needs +/// `SeCreateSymbolicLinkPrivilege`; other platforms vary). Lexical absolute and +/// `..` rejection still applies; unix symlink coverage lives in +/// `rejects_index_shard_paths_that_escape_via_symlink` and +/// `rejects_directory_index_that_escapes_via_symlink`. Existing paths are still +/// checked with `canonicalize` via `validate_path_stays_under_root`. +#[cfg(not(unix))] +#[test] +fn non_unix_symlink_escape_falls_back_to_lexical_path_guards() { + let dir = temp_dir("non-unix-symlink-fallback"); + fs::write( + dir.join("model.safetensors.index.json"), + r#"{"weight_map": {"a.weight": "../outside.safetensors"}}"#, + ) + .unwrap(); + let parent_err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + assert!( + parent_err + .to_string() + .contains("must stay within the checkpoint directory") + ); + + let outside = temp_dir("non-unix-absolute-outside"); + let outside_shard = outside.join("outside.safetensors"); + write_safetensors( + &outside_shard, + r#"{"a.weight": {"dtype": "F16", "shape": [1], "data_offsets": [0, 2]}}"#, + 2, + ); + let escaped = outside_shard.to_string_lossy().replace('\\', "\\\\"); + fs::write( + dir.join("model.safetensors.index.json"), + format!(r#"{{"weight_map": {{"a.weight": "{escaped}"}}}}"#), + ) + .unwrap(); + let absolute_err = inspect_safetensors_checkpoint(&dir).unwrap_err(); + assert!( + absolute_err + .to_string() + .contains("must be checkpoint-relative, not absolute") + ); +} + #[test] fn rejects_multiple_directory_indexes() { let dir = temp_dir("multiple-indexes"); @@ -762,7 +983,9 @@ fn discovery_does_not_double_list_experts_as_routers() { vec![PathBuf::from("experts.safetensors")], BTreeMap::new(), vec![tensor], - ); + Path::new("experts.safetensors"), + ) + .unwrap(); assert_eq!(manifest.candidates.router_tensors, Vec::::new()); assert_eq!(manifest.candidates.expert_tensors, vec![expert_name]); assert_eq!(manifest.tensors[0].labels, vec!["moe_expert_candidate"]); @@ -786,7 +1009,9 @@ fn named_family_requires_router_and_expert_evidence() { vec![PathBuf::from("model.safetensors")], BTreeMap::new(), vec![tensor], - ); + Path::new("model.safetensors"), + ) + .unwrap(); assert_eq!(manifest.candidates.detected_layout_family, None); assert!(manifest.candidates.router_tensors.is_empty()); assert!(manifest.candidates.expert_tensors.is_empty()); @@ -824,7 +1049,9 @@ fn unknown_expert_weight_kind_is_retained() { ], BTreeMap::new(), vec![router, expert], - ); + Path::new("checkpoint"), + ) + .unwrap(); assert_eq!( manifest.candidates.detected_layout_family, Some("generic_moe") @@ -859,7 +1086,9 @@ fn expert_groups_keep_parallel_layer_stacks_separate() { vec![PathBuf::from("experts.safetensors")], BTreeMap::new(), tensors, - ); + Path::new("experts.safetensors"), + ) + .unwrap(); let group_keys = manifest .candidates .expert_groups diff --git a/tests/safetensors_smoke.rs b/tests/safetensors_smoke.rs index c711985..c97dbe8 100644 --- a/tests/safetensors_smoke.rs +++ b/tests/safetensors_smoke.rs @@ -77,3 +77,15 @@ fn fixture_manifest_write_is_byte_stable() { assert_eq!(first, second); assert!(first.starts_with("{\n \"candidates\":")); } + +#[test] +fn sharded_fixture_manifest_is_byte_stable() { + let path = fixture_root().join("sharded"); + let first = inspect_safetensors_checkpoint(&path) + .expect("sharded fixture") + .to_pretty_json(); + let second = inspect_safetensors_checkpoint(&path) + .expect("sharded fixture") + .to_pretty_json(); + assert_eq!(first, second); +}