951 lines
29 KiB
Rust
951 lines
29 KiB
Rust
//! Sharded weight management for large model files.
|
|
|
|
use crate::{HubError, HubResult, ModelId, StorageBackend};
|
|
use serde::{Deserialize, Serialize};
|
|
use sha2::{Digest, Sha256};
|
|
use std::path::Path;
|
|
use tokio::fs;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
|
|
/// Maximum retry attempts for shard downloads.
|
|
const MAX_RETRIES: usize = 3;
|
|
|
|
/// Default shard size (1GB).
|
|
const DEFAULT_SHARD_SIZE: u64 = 1024 * 1024 * 1024;
|
|
|
|
/// Represents a single weight shard.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct WeightShard {
|
|
/// Shard identifier
|
|
pub shard_id: usize,
|
|
/// Filename of the shard
|
|
pub filename: String,
|
|
/// Size in bytes
|
|
pub size: u64,
|
|
/// SHA256 checksum
|
|
pub checksum: String,
|
|
/// Byte range in the original file (start, end)
|
|
pub byte_range: (u64, u64),
|
|
}
|
|
|
|
impl WeightShard {
|
|
/// Create a new weight shard.
|
|
pub fn new(
|
|
shard_id: usize,
|
|
filename: String,
|
|
size: u64,
|
|
checksum: String,
|
|
byte_range: (u64, u64),
|
|
) -> Self {
|
|
Self {
|
|
shard_id,
|
|
filename,
|
|
size,
|
|
checksum,
|
|
byte_range,
|
|
}
|
|
}
|
|
|
|
/// Verify checksum of shard data.
|
|
pub fn verify_checksum(&self, data: &[u8]) -> bool {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(data);
|
|
let computed = hex::encode(hasher.finalize());
|
|
computed == self.checksum
|
|
}
|
|
|
|
/// Compute checksum from data.
|
|
pub fn compute_checksum(data: &[u8]) -> String {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(data);
|
|
hex::encode(hasher.finalize())
|
|
}
|
|
}
|
|
|
|
/// Sharded weights metadata.
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct ShardedWeights {
|
|
/// Model identifier
|
|
pub model_id: ModelId,
|
|
/// Total size of all shards
|
|
pub total_size: u64,
|
|
/// Number of shards
|
|
pub num_shards: usize,
|
|
/// List of shards
|
|
pub shards: Vec<WeightShard>,
|
|
/// Shard filename pattern
|
|
pub shard_pattern: String,
|
|
}
|
|
|
|
impl ShardedWeights {
|
|
/// Create new sharded weights metadata.
|
|
pub fn new(model_id: ModelId, total_size: u64, shard_pattern: String) -> Self {
|
|
Self {
|
|
model_id,
|
|
total_size,
|
|
num_shards: 0,
|
|
shards: Vec::new(),
|
|
shard_pattern,
|
|
}
|
|
}
|
|
|
|
/// Add a shard.
|
|
pub fn add_shard(&mut self, shard: WeightShard) {
|
|
self.shards.push(shard);
|
|
self.num_shards = self.shards.len();
|
|
}
|
|
|
|
/// Get shard by ID.
|
|
pub fn get_shard(&self, shard_id: usize) -> Option<&WeightShard> {
|
|
self.shards.iter().find(|s| s.shard_id == shard_id)
|
|
}
|
|
|
|
/// Verify all shards are present and valid.
|
|
pub fn verify_complete(&self) -> bool {
|
|
// Empty shard set is not complete
|
|
if self.num_shards == 0 || self.shards.is_empty() {
|
|
return false;
|
|
}
|
|
|
|
if self.shards.len() != self.num_shards {
|
|
return false;
|
|
}
|
|
|
|
for i in 0..self.num_shards {
|
|
if !self.shards.iter().any(|s| s.shard_id == i) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
/// Get total size from shards.
|
|
pub fn compute_total_size(&self) -> u64 {
|
|
self.shards.iter().map(|s| s.size).sum()
|
|
}
|
|
|
|
/// Format shard filename from pattern.
|
|
pub fn format_shard_filename(&self, shard_id: usize) -> String {
|
|
self.shard_pattern
|
|
.replace("{shard:05d}", &format!("{:05}", shard_id))
|
|
.replace("{total:05d}", &format!("{:05}", self.num_shards))
|
|
}
|
|
}
|
|
|
|
/// Weight manager for sharded operations.
|
|
pub struct WeightManager {
|
|
storage: Box<dyn StorageBackend>,
|
|
}
|
|
|
|
impl WeightManager {
|
|
/// Create a new weight manager.
|
|
pub fn new(storage: Box<dyn StorageBackend>) -> Self {
|
|
Self { storage }
|
|
}
|
|
|
|
/// Split a large weight file into shards.
|
|
pub async fn create_shards(
|
|
&self,
|
|
model_id: ModelId,
|
|
source_path: &Path,
|
|
shard_pattern: String,
|
|
shard_size: Option<u64>,
|
|
) -> HubResult<ShardedWeights> {
|
|
let shard_size = shard_size.unwrap_or(DEFAULT_SHARD_SIZE);
|
|
|
|
// Get file size
|
|
let metadata = fs::metadata(source_path).await?;
|
|
let total_size = metadata.len();
|
|
|
|
// Calculate number of shards
|
|
let num_shards = ((total_size + shard_size - 1) / shard_size) as usize;
|
|
|
|
let mut sharded_weights =
|
|
ShardedWeights::new(model_id.clone(), total_size, shard_pattern.clone());
|
|
|
|
// Open source file
|
|
let mut file = fs::File::open(source_path).await?;
|
|
let mut buffer = vec![0u8; shard_size as usize];
|
|
|
|
for shard_id in 0..num_shards {
|
|
let start_byte = shard_id as u64 * shard_size;
|
|
let end_byte = std::cmp::min(start_byte + shard_size, total_size);
|
|
let actual_size = end_byte - start_byte;
|
|
|
|
// Read shard data
|
|
let bytes_read = file.read(&mut buffer[..actual_size as usize]).await?;
|
|
let shard_data = &buffer[..bytes_read];
|
|
|
|
// Compute checksum
|
|
let checksum = WeightShard::compute_checksum(shard_data);
|
|
|
|
// Create filename
|
|
let filename = shard_pattern
|
|
.replace("{shard:05d}", &format!("{:05}", shard_id))
|
|
.replace("{total:05d}", &format!("{:05}", num_shards));
|
|
|
|
// Store shard
|
|
let storage_path = format!("models/{}/{}", model_id.to_string(), filename);
|
|
self.storage.store(&storage_path, shard_data).await?;
|
|
|
|
// Add to metadata
|
|
let shard = WeightShard::new(
|
|
shard_id,
|
|
filename,
|
|
bytes_read as u64,
|
|
checksum,
|
|
(start_byte, end_byte),
|
|
);
|
|
sharded_weights.add_shard(shard);
|
|
}
|
|
|
|
Ok(sharded_weights)
|
|
}
|
|
|
|
/// Download a single shard with retry logic.
|
|
pub async fn download_shard(
|
|
&self,
|
|
model_id: &ModelId,
|
|
shard: &WeightShard,
|
|
retry_count: usize,
|
|
) -> HubResult<Vec<u8>> {
|
|
let storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);
|
|
|
|
for attempt in 0..retry_count {
|
|
match self.storage.load(&storage_path).await {
|
|
Ok(data) => {
|
|
// Verify checksum
|
|
if shard.verify_checksum(&data) {
|
|
return Ok(data);
|
|
}
|
|
if attempt == retry_count - 1 {
|
|
return Err(HubError::ValidationFailed {
|
|
details: format!(
|
|
"Checksum mismatch for shard {} after {} retries",
|
|
shard.shard_id, retry_count
|
|
),
|
|
});
|
|
}
|
|
// Retry on checksum failure
|
|
tracing::warn!(
|
|
"Checksum mismatch for shard {}, attempt {}/{}",
|
|
shard.shard_id,
|
|
attempt + 1,
|
|
retry_count
|
|
);
|
|
}
|
|
Err(e) => {
|
|
if attempt == retry_count - 1 {
|
|
return Err(e);
|
|
}
|
|
tracing::warn!(
|
|
"Failed to download shard {}, attempt {}/{}: {}",
|
|
shard.shard_id,
|
|
attempt + 1,
|
|
retry_count,
|
|
e
|
|
);
|
|
// Exponential backoff
|
|
tokio::time::sleep(std::time::Duration::from_millis(100 * (1 << attempt)))
|
|
.await;
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(HubError::InvalidPackage {
|
|
reason: "Max retries exceeded".to_string(),
|
|
})
|
|
}
|
|
|
|
/// Download all shards in parallel.
|
|
pub async fn download_all(
|
|
&self,
|
|
sharded_weights: &ShardedWeights,
|
|
output_path: &Path,
|
|
max_parallel: usize,
|
|
) -> HubResult<()> {
|
|
use std::sync::Arc;
|
|
use tokio::sync::Semaphore;
|
|
|
|
let semaphore = Arc::new(Semaphore::new(max_parallel));
|
|
let mut tasks = Vec::new();
|
|
|
|
for shard in &sharded_weights.shards {
|
|
let shard = shard.clone();
|
|
let model_id = sharded_weights.model_id.clone();
|
|
let permit =
|
|
semaphore
|
|
.clone()
|
|
.acquire_owned()
|
|
.await
|
|
.map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Semaphore error: {}", e),
|
|
})?;
|
|
|
|
let _storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);
|
|
|
|
// Clone storage backend for parallel access
|
|
// Note: In practice, storage backend should be Arc-wrapped for true parallel access
|
|
// For this implementation, we'll download sequentially but structure for parallel
|
|
let task = tokio::spawn(async move {
|
|
let _permit = permit;
|
|
// Return shard info for sequential processing
|
|
Ok::<_, HubError>(shard)
|
|
});
|
|
|
|
tasks.push(task);
|
|
}
|
|
|
|
// Wait for all downloads
|
|
let mut downloaded_shards = Vec::new();
|
|
for task in tasks {
|
|
let shard = task.await.map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Task join error: {}", e),
|
|
})??;
|
|
downloaded_shards.push(shard);
|
|
}
|
|
|
|
// Sort shards by ID
|
|
downloaded_shards.sort_by_key(|s| s.shard_id);
|
|
|
|
// Download and merge shards sequentially (in order)
|
|
let mut output_file = fs::File::create(output_path).await?;
|
|
|
|
for shard in downloaded_shards {
|
|
let data = self
|
|
.download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
|
|
.await?;
|
|
output_file.write_all(&data).await?;
|
|
}
|
|
|
|
output_file.flush().await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Verify checksum of a shard.
|
|
pub async fn verify_checksum(
|
|
&self,
|
|
model_id: &ModelId,
|
|
shard: &WeightShard,
|
|
) -> HubResult<bool> {
|
|
let storage_path = format!("models/{}/{}", model_id.to_string(), shard.filename);
|
|
let data = self.storage.load(&storage_path).await?;
|
|
Ok(shard.verify_checksum(&data))
|
|
}
|
|
|
|
/// Merge shards into a single file.
|
|
pub async fn merge_shards(
|
|
&self,
|
|
sharded_weights: &ShardedWeights,
|
|
output_path: &Path,
|
|
) -> HubResult<()> {
|
|
if !sharded_weights.verify_complete() {
|
|
return Err(HubError::ValidationFailed {
|
|
details: "Incomplete shard set".to_string(),
|
|
});
|
|
}
|
|
|
|
let mut output_file = fs::File::create(output_path).await?;
|
|
|
|
// Sort shards by ID to ensure correct order
|
|
let mut sorted_shards = sharded_weights.shards.clone();
|
|
sorted_shards.sort_by_key(|s| s.shard_id);
|
|
|
|
for shard in sorted_shards {
|
|
let data = self
|
|
.download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
|
|
.await?;
|
|
output_file.write_all(&data).await?;
|
|
}
|
|
|
|
output_file.flush().await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Apply delta/adapter weights to base model.
|
|
pub async fn apply_delta(
|
|
&self,
|
|
base_path: &Path,
|
|
delta_path: &Path,
|
|
output_path: &Path,
|
|
) -> HubResult<()> {
|
|
// Read base weights
|
|
let base_data = fs::read(base_path).await?;
|
|
|
|
// Read delta weights
|
|
let delta_data = fs::read(delta_path).await?;
|
|
|
|
// Detect format and apply delta
|
|
let merged_data = if Self::is_safetensors(&base_data) && Self::is_safetensors(&delta_data) {
|
|
// SafeTensors format - parse and merge
|
|
self.merge_safetensors(&base_data, &delta_data)?
|
|
} else {
|
|
// For non-SafeTensors, just concatenate or use custom logic
|
|
// This is a simplified implementation
|
|
[base_data, delta_data].concat()
|
|
};
|
|
|
|
// Write merged weights
|
|
fs::write(output_path, merged_data).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if data is in SafeTensors format.
|
|
fn is_safetensors(data: &[u8]) -> bool {
|
|
// SafeTensors files start with an 8-byte little-endian header size
|
|
if data.len() < 8 {
|
|
return false;
|
|
}
|
|
|
|
// Read header size
|
|
let header_size = u64::from_le_bytes([
|
|
data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
|
|
]);
|
|
|
|
// Header size should be reasonable (< 100MB)
|
|
if header_size > 100_000_000 {
|
|
return false;
|
|
}
|
|
|
|
// Check if we have enough data for header
|
|
if data.len() < 8 + header_size as usize {
|
|
return false;
|
|
}
|
|
|
|
// Try to parse header as JSON
|
|
let header_bytes = &data[8..8 + header_size as usize];
|
|
serde_json::from_slice::<serde_json::Value>(header_bytes).is_ok()
|
|
}
|
|
|
|
/// Merge SafeTensors files (simplified implementation).
|
|
fn merge_safetensors(&self, base: &[u8], delta: &[u8]) -> HubResult<Vec<u8>> {
|
|
// This is a simplified implementation
|
|
// In production, you would properly parse SafeTensors format and merge tensors
|
|
|
|
// For now, we'll just concatenate the tensors section
|
|
// Real implementation would parse headers and merge tensor dictionaries
|
|
|
|
let base_header_size = u64::from_le_bytes([
|
|
base[0], base[1], base[2], base[3], base[4], base[5], base[6], base[7],
|
|
]);
|
|
|
|
let delta_header_size = u64::from_le_bytes([
|
|
delta[0], delta[1], delta[2], delta[3], delta[4], delta[5], delta[6], delta[7],
|
|
]);
|
|
|
|
// Parse headers
|
|
let base_header_bytes = &base[8..8 + base_header_size as usize];
|
|
let delta_header_bytes = &delta[8..8 + delta_header_size as usize];
|
|
|
|
let mut base_header: serde_json::Value = serde_json::from_slice(base_header_bytes)?;
|
|
let delta_header: serde_json::Value = serde_json::from_slice(delta_header_bytes)?;
|
|
|
|
// Merge headers (simplified - just add delta tensors)
|
|
if let (Some(base_obj), Some(delta_obj)) =
|
|
(base_header.as_object_mut(), delta_header.as_object())
|
|
{
|
|
for (key, value) in delta_obj {
|
|
if key != "__metadata__" {
|
|
base_obj.insert(key.clone(), value.clone());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Serialize merged header
|
|
let merged_header = serde_json::to_vec(&base_header)?;
|
|
let merged_header_size = merged_header.len() as u64;
|
|
|
|
// Build merged file
|
|
let mut result = Vec::new();
|
|
result.extend_from_slice(&merged_header_size.to_le_bytes());
|
|
result.extend_from_slice(&merged_header);
|
|
|
|
// Add base tensors
|
|
let base_tensors_start = 8 + base_header_size as usize;
|
|
result.extend_from_slice(&base[base_tensors_start..]);
|
|
|
|
// Add delta tensors
|
|
let delta_tensors_start = 8 + delta_header_size as usize;
|
|
result.extend_from_slice(&delta[delta_tensors_start..]);
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Resume interrupted download.
|
|
pub async fn resume_download(
|
|
&self,
|
|
sharded_weights: &ShardedWeights,
|
|
output_path: &Path,
|
|
) -> HubResult<()> {
|
|
// Check which shards are already downloaded
|
|
let mut downloaded_shards = Vec::new();
|
|
|
|
if output_path.exists() {
|
|
let current_size = fs::metadata(output_path).await?.len();
|
|
|
|
// Determine which shards are complete
|
|
let mut accumulated_size = 0u64;
|
|
for shard in &sharded_weights.shards {
|
|
if accumulated_size + shard.size <= current_size {
|
|
downloaded_shards.push(shard.shard_id);
|
|
accumulated_size += shard.size;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Open file for appending
|
|
let mut output_file = fs::OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.open(output_path)
|
|
.await?;
|
|
|
|
// Download remaining shards
|
|
let mut sorted_shards = sharded_weights.shards.clone();
|
|
sorted_shards.sort_by_key(|s| s.shard_id);
|
|
|
|
for shard in sorted_shards {
|
|
if !downloaded_shards.contains(&shard.shard_id) {
|
|
let data = self
|
|
.download_shard(&sharded_weights.model_id, &shard, MAX_RETRIES)
|
|
.await?;
|
|
output_file.write_all(&data).await?;
|
|
}
|
|
}
|
|
|
|
output_file.flush().await?;
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::storage::LocalStorageBackend;
|
|
use tempfile::TempDir;
|
|
|
|
fn test_model_id() -> ModelId {
|
|
ModelId::new("rustytorch", "test-model")
|
|
}
|
|
|
|
#[test]
|
|
fn test_weight_shard_creation() {
|
|
let shard = WeightShard::new(
|
|
0,
|
|
"model-00000-of-00003.safetensors".to_string(),
|
|
1024,
|
|
"abc123".to_string(),
|
|
(0, 1024),
|
|
);
|
|
|
|
assert_eq!(shard.shard_id, 0);
|
|
assert_eq!(shard.filename, "model-00000-of-00003.safetensors");
|
|
assert_eq!(shard.size, 1024);
|
|
assert_eq!(shard.checksum, "abc123");
|
|
assert_eq!(shard.byte_range, (0, 1024));
|
|
}
|
|
|
|
#[test]
|
|
fn test_weight_shard_checksum_verification() {
|
|
let data = b"test data";
|
|
let checksum = WeightShard::compute_checksum(data);
|
|
|
|
let shard = WeightShard::new(0, "test.bin".to_string(), 9, checksum.clone(), (0, 9));
|
|
|
|
assert!(shard.verify_checksum(data));
|
|
assert!(!shard.verify_checksum(b"wrong data"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_creation() {
|
|
let model_id = test_model_id();
|
|
let sharded = ShardedWeights::new(
|
|
model_id.clone(),
|
|
3072,
|
|
"model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
);
|
|
|
|
assert_eq!(sharded.model_id, model_id);
|
|
assert_eq!(sharded.total_size, 3072);
|
|
assert_eq!(sharded.num_shards, 0);
|
|
assert!(sharded.shards.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_add_shard() {
|
|
let mut sharded = ShardedWeights::new(
|
|
test_model_id(),
|
|
3072,
|
|
"model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
);
|
|
|
|
let shard = WeightShard::new(
|
|
0,
|
|
"test.bin".to_string(),
|
|
1024,
|
|
"abc".to_string(),
|
|
(0, 1024),
|
|
);
|
|
sharded.add_shard(shard);
|
|
|
|
assert_eq!(sharded.num_shards, 1);
|
|
assert_eq!(sharded.shards.len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_get_shard() {
|
|
let mut sharded = ShardedWeights::new(
|
|
test_model_id(),
|
|
2048,
|
|
"model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
);
|
|
|
|
let shard1 = WeightShard::new(
|
|
0,
|
|
"shard0.bin".to_string(),
|
|
1024,
|
|
"abc".to_string(),
|
|
(0, 1024),
|
|
);
|
|
let shard2 = WeightShard::new(
|
|
1,
|
|
"shard1.bin".to_string(),
|
|
1024,
|
|
"def".to_string(),
|
|
(1024, 2048),
|
|
);
|
|
|
|
sharded.add_shard(shard1.clone());
|
|
sharded.add_shard(shard2.clone());
|
|
|
|
assert_eq!(sharded.get_shard(0), Some(&shard1));
|
|
assert_eq!(sharded.get_shard(1), Some(&shard2));
|
|
assert_eq!(sharded.get_shard(2), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_verify_complete() {
|
|
let mut sharded = ShardedWeights::new(
|
|
test_model_id(),
|
|
3072,
|
|
"model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
);
|
|
|
|
// Empty is not complete
|
|
assert!(!sharded.verify_complete());
|
|
|
|
// Add shard 0
|
|
sharded.add_shard(WeightShard::new(
|
|
0,
|
|
"s0.bin".to_string(),
|
|
1024,
|
|
"a".to_string(),
|
|
(0, 1024),
|
|
));
|
|
assert!(sharded.verify_complete()); // 1 shard is complete
|
|
|
|
// Add shard 2 (missing shard 1)
|
|
sharded.add_shard(WeightShard::new(
|
|
2,
|
|
"s2.bin".to_string(),
|
|
1024,
|
|
"c".to_string(),
|
|
(2048, 3072),
|
|
));
|
|
assert!(!sharded.verify_complete()); // Missing shard 1
|
|
|
|
// Add shard 1
|
|
sharded.add_shard(WeightShard::new(
|
|
1,
|
|
"s1.bin".to_string(),
|
|
1024,
|
|
"b".to_string(),
|
|
(1024, 2048),
|
|
));
|
|
assert!(sharded.verify_complete()); // All shards present
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_compute_total_size() {
|
|
let mut sharded = ShardedWeights::new(
|
|
test_model_id(),
|
|
0,
|
|
"model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
);
|
|
|
|
sharded.add_shard(WeightShard::new(
|
|
0,
|
|
"s0.bin".to_string(),
|
|
1024,
|
|
"a".to_string(),
|
|
(0, 1024),
|
|
));
|
|
sharded.add_shard(WeightShard::new(
|
|
1,
|
|
"s1.bin".to_string(),
|
|
2048,
|
|
"b".to_string(),
|
|
(1024, 3072),
|
|
));
|
|
|
|
assert_eq!(sharded.compute_total_size(), 3072);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_format_shard_filename() {
|
|
let sharded = ShardedWeights {
|
|
model_id: test_model_id(),
|
|
total_size: 3072,
|
|
num_shards: 3,
|
|
shards: vec![],
|
|
shard_pattern: "model-{shard:05d}-of-{total:05d}.safetensors".to_string(),
|
|
};
|
|
|
|
assert_eq!(
|
|
sharded.format_shard_filename(0),
|
|
"model-00000-of-00003.safetensors"
|
|
);
|
|
assert_eq!(
|
|
sharded.format_shard_filename(1),
|
|
"model-00001-of-00003.safetensors"
|
|
);
|
|
assert_eq!(
|
|
sharded.format_shard_filename(2),
|
|
"model-00002-of-00003.safetensors"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_weight_manager_create_shards() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let manager = WeightManager::new(storage);
|
|
|
|
// Create test file
|
|
let test_file = temp_dir.path().join("test_weights.bin");
|
|
let test_data = vec![0u8; 2500]; // 2.5KB
|
|
fs::write(&test_file, &test_data).await.unwrap();
|
|
|
|
// Create shards with 1KB shard size
|
|
let sharded = manager
|
|
.create_shards(
|
|
test_model_id(),
|
|
&test_file,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
Some(1024),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(sharded.num_shards, 3); // 2500 bytes / 1024 = 3 shards
|
|
assert_eq!(sharded.total_size, 2500);
|
|
assert_eq!(sharded.shards.len(), 3);
|
|
|
|
// Verify shard sizes
|
|
assert_eq!(sharded.shards[0].size, 1024);
|
|
assert_eq!(sharded.shards[1].size, 1024);
|
|
assert_eq!(sharded.shards[2].size, 452); // Remaining bytes
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_weight_manager_download_shard() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let manager = WeightManager::new(storage);
|
|
|
|
// Create test file and shards
|
|
let test_file = temp_dir.path().join("test_weights.bin");
|
|
let test_data = vec![1u8; 1024];
|
|
fs::write(&test_file, &test_data).await.unwrap();
|
|
|
|
let sharded = manager
|
|
.create_shards(
|
|
test_model_id(),
|
|
&test_file,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
Some(1024),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Download first shard
|
|
let downloaded = manager
|
|
.download_shard(&sharded.model_id, &sharded.shards[0], MAX_RETRIES)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(downloaded.len(), 1024);
|
|
assert_eq!(downloaded, test_data);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_weight_manager_verify_checksum() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let manager = WeightManager::new(storage);
|
|
|
|
// Create test file
|
|
let test_file = temp_dir.path().join("test_weights.bin");
|
|
let test_data = vec![2u8; 512];
|
|
fs::write(&test_file, &test_data).await.unwrap();
|
|
|
|
let sharded = manager
|
|
.create_shards(
|
|
test_model_id(),
|
|
&test_file,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
Some(512),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Verify checksum
|
|
let is_valid = manager
|
|
.verify_checksum(&sharded.model_id, &sharded.shards[0])
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(is_valid);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_weight_manager_merge_shards() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let manager = WeightManager::new(storage);
|
|
|
|
// Create test file
|
|
let test_file = temp_dir.path().join("test_weights.bin");
|
|
let test_data = vec![3u8; 2048];
|
|
fs::write(&test_file, &test_data).await.unwrap();
|
|
|
|
// Create shards
|
|
let sharded = manager
|
|
.create_shards(
|
|
test_model_id(),
|
|
&test_file,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
Some(1024),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Merge shards
|
|
let output_file = temp_dir.path().join("merged.bin");
|
|
manager.merge_shards(&sharded, &output_file).await.unwrap();
|
|
|
|
// Verify merged file
|
|
let merged_data = fs::read(&output_file).await.unwrap();
|
|
assert_eq!(merged_data, test_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_safetensors_detection() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let _manager = WeightManager::new(storage);
|
|
|
|
// Valid SafeTensors format
|
|
let header = serde_json::json!({"test": "value"});
|
|
let header_bytes = serde_json::to_vec(&header).unwrap();
|
|
let header_size = header_bytes.len() as u64;
|
|
|
|
let mut safetensors_data = Vec::new();
|
|
safetensors_data.extend_from_slice(&header_size.to_le_bytes());
|
|
safetensors_data.extend_from_slice(&header_bytes);
|
|
safetensors_data.extend_from_slice(&[0u8; 100]); // Tensor data
|
|
|
|
assert!(WeightManager::is_safetensors(&safetensors_data));
|
|
|
|
// Invalid data
|
|
assert!(!WeightManager::is_safetensors(&[1, 2, 3]));
|
|
assert!(!WeightManager::is_safetensors(&[]));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_weight_manager_resume_download() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = Box::new(LocalStorageBackend::new(temp_dir.path().to_path_buf()));
|
|
let manager = WeightManager::new(storage);
|
|
|
|
// Create test file
|
|
let test_file = temp_dir.path().join("test_weights.bin");
|
|
let test_data = vec![4u8; 3072];
|
|
fs::write(&test_file, &test_data).await.unwrap();
|
|
|
|
// Create shards
|
|
let sharded = manager
|
|
.create_shards(
|
|
test_model_id(),
|
|
&test_file,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
Some(1024),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Simulate partial download (first shard only)
|
|
let output_file = temp_dir.path().join("partial.bin");
|
|
let first_shard_data = manager
|
|
.download_shard(&sharded.model_id, &sharded.shards[0], MAX_RETRIES)
|
|
.await
|
|
.unwrap();
|
|
fs::write(&output_file, &first_shard_data).await.unwrap();
|
|
|
|
// Resume download
|
|
manager
|
|
.resume_download(&sharded, &output_file)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Verify complete file
|
|
let complete_data = fs::read(&output_file).await.unwrap();
|
|
assert_eq!(complete_data, test_data);
|
|
}
|
|
|
|
#[test]
|
|
fn test_weight_shard_serialization() {
|
|
let shard = WeightShard::new(
|
|
0,
|
|
"test.bin".to_string(),
|
|
1024,
|
|
"checksum".to_string(),
|
|
(0, 1024),
|
|
);
|
|
|
|
let json = serde_json::to_string(&shard).unwrap();
|
|
let deserialized: WeightShard = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(deserialized, shard);
|
|
}
|
|
|
|
#[test]
|
|
fn test_sharded_weights_serialization() {
|
|
let mut sharded = ShardedWeights::new(
|
|
test_model_id(),
|
|
2048,
|
|
"model-{shard:05d}-of-{total:05d}.bin".to_string(),
|
|
);
|
|
|
|
sharded.add_shard(WeightShard::new(
|
|
0,
|
|
"shard0.bin".to_string(),
|
|
1024,
|
|
"abc".to_string(),
|
|
(0, 1024),
|
|
));
|
|
|
|
let json = serde_json::to_string(&sharded).unwrap();
|
|
let deserialized: ShardedWeights = serde_json::from_str(&json).unwrap();
|
|
|
|
assert_eq!(deserialized, sharded);
|
|
}
|
|
}
|