1053 lines
34 KiB
Rust
1053 lines
34 KiB
Rust
//! Model packaging system for deployment-ready model bundles.
|
|
|
|
use crate::{HubError, HubResult, ModelId, ModelMetadata, StorageBackend};
|
|
use base64::Engine;
|
|
use chrono::{DateTime, Utc};
|
|
use serde::{Deserialize, Serialize};
|
|
use sha2::{Digest, Sha256};
|
|
use std::collections::HashMap;
|
|
use std::future::Future;
|
|
use std::io::{Read, Write};
|
|
use std::path::Path;
|
|
use std::pin::Pin;
|
|
use tokio::fs;
|
|
|
|
/// Model package format version.
|
|
pub const PACKAGE_FORMAT_VERSION: &str = "1.0";
|
|
|
|
/// Model package containing all assets and metadata.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelPackage {
|
|
/// Package metadata
|
|
pub metadata: PackageMetadata,
|
|
/// Model assets
|
|
pub assets: Vec<ModelAsset>,
|
|
/// Dependencies
|
|
pub dependencies: Vec<PackageDependency>,
|
|
/// Package manifest
|
|
pub manifest: PackageManifest,
|
|
}
|
|
|
|
/// Package-level metadata.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PackageMetadata {
|
|
/// Package format version
|
|
pub format_version: String,
|
|
/// Model metadata
|
|
pub model: ModelMetadata,
|
|
/// Package creation timestamp
|
|
pub created_at: DateTime<Utc>,
|
|
/// Package creator
|
|
pub created_by: String,
|
|
/// Package size in bytes
|
|
pub size: u64,
|
|
/// Content hash for integrity verification
|
|
pub content_hash: String,
|
|
/// Compression settings
|
|
pub compression: CompressionSettings,
|
|
/// Package signature (for security)
|
|
pub signature: Option<PackageSignature>,
|
|
}
|
|
|
|
/// Package compression settings.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct CompressionSettings {
|
|
/// Compression algorithm used
|
|
pub algorithm: CompressionAlgorithm,
|
|
/// Compression level (1-9)
|
|
pub level: u8,
|
|
/// Original size before compression
|
|
pub original_size: u64,
|
|
/// Compressed size
|
|
pub compressed_size: u64,
|
|
/// Compression ratio
|
|
pub ratio: f64,
|
|
}
|
|
|
|
/// Supported compression algorithms.
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
|
pub enum CompressionAlgorithm {
|
|
None,
|
|
Gzip,
|
|
Zstd,
|
|
Lz4,
|
|
}
|
|
|
|
/// Package signature for verification.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PackageSignature {
|
|
/// Signature algorithm
|
|
pub algorithm: String,
|
|
/// Signature bytes (base64 encoded)
|
|
pub signature: String,
|
|
/// Public key identifier
|
|
pub key_id: String,
|
|
/// Signing timestamp
|
|
pub signed_at: DateTime<Utc>,
|
|
}
|
|
|
|
/// Model asset within a package.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ModelAsset {
|
|
/// Asset path within package
|
|
pub path: String,
|
|
/// Asset type
|
|
pub asset_type: AssetType,
|
|
/// Asset size in bytes
|
|
pub size: u64,
|
|
/// Asset hash
|
|
pub hash: String,
|
|
/// MIME type
|
|
pub mime_type: String,
|
|
/// Asset metadata
|
|
pub metadata: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// Types of model assets.
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
|
pub enum AssetType {
|
|
/// Model weights/parameters
|
|
Weights,
|
|
/// Model configuration
|
|
Config,
|
|
/// Tokenizer data
|
|
Tokenizer,
|
|
/// Vocabulary file
|
|
Vocabulary,
|
|
/// Model code/implementation
|
|
Code,
|
|
/// Documentation
|
|
Documentation,
|
|
/// Example inputs/outputs
|
|
Examples,
|
|
/// License file
|
|
License,
|
|
/// Custom asset type
|
|
Custom(String),
|
|
}
|
|
|
|
/// Package dependency information.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PackageDependency {
|
|
/// Dependency model ID
|
|
pub model_id: ModelId,
|
|
/// Version constraint
|
|
pub version_constraint: String,
|
|
/// Dependency type
|
|
pub dependency_type: String,
|
|
/// Whether dependency is optional
|
|
pub optional: bool,
|
|
/// Content hash of dependency
|
|
pub content_hash: Option<String>,
|
|
}
|
|
|
|
/// Package manifest with file listings.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PackageManifest {
|
|
/// List of all files in package
|
|
pub files: Vec<ManifestEntry>,
|
|
/// Directory structure
|
|
pub directories: Vec<String>,
|
|
/// Total number of files
|
|
pub file_count: usize,
|
|
/// Total uncompressed size
|
|
pub total_size: u64,
|
|
}
|
|
|
|
/// Manifest entry for a single file.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ManifestEntry {
|
|
/// File path
|
|
pub path: String,
|
|
/// File size
|
|
pub size: u64,
|
|
/// File hash
|
|
pub hash: String,
|
|
/// File permissions (Unix-style)
|
|
pub permissions: u32,
|
|
/// Last modified timestamp
|
|
pub modified: DateTime<Utc>,
|
|
}
|
|
|
|
/// Options for model packaging.
|
|
#[derive(Debug, Clone)]
|
|
pub struct PackagingOptions {
|
|
/// Compression settings
|
|
pub compression: CompressionAlgorithm,
|
|
/// Compression level
|
|
pub compression_level: u8,
|
|
/// Include source code
|
|
pub include_code: bool,
|
|
/// Include documentation
|
|
pub include_docs: bool,
|
|
/// Include examples
|
|
pub include_examples: bool,
|
|
/// Validate package integrity
|
|
pub validate: bool,
|
|
/// Sign the package
|
|
pub sign: Option<SigningOptions>,
|
|
/// Exclude patterns (glob patterns)
|
|
pub exclude_patterns: Vec<String>,
|
|
/// Include patterns (glob patterns)
|
|
pub include_patterns: Vec<String>,
|
|
}
|
|
|
|
/// Package signing options.
|
|
#[derive(Debug, Clone)]
|
|
pub struct SigningOptions {
|
|
/// Private key for signing
|
|
pub private_key: Vec<u8>,
|
|
/// Key identifier
|
|
pub key_id: String,
|
|
/// Signing algorithm
|
|
pub algorithm: String,
|
|
}
|
|
|
|
impl Default for PackagingOptions {
|
|
fn default() -> Self {
|
|
Self {
|
|
compression: CompressionAlgorithm::Zstd,
|
|
compression_level: 6,
|
|
include_code: true,
|
|
include_docs: true,
|
|
include_examples: false,
|
|
validate: true,
|
|
sign: None,
|
|
exclude_patterns: vec![
|
|
"*.log".to_string(),
|
|
"*.tmp".to_string(),
|
|
".git/**".to_string(),
|
|
"__pycache__/**".to_string(),
|
|
],
|
|
include_patterns: vec!["**/*".to_string()],
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model packager for creating deployment bundles.
|
|
pub struct ModelPackager {
|
|
/// Storage backend
|
|
storage: Box<dyn StorageBackend>,
|
|
/// Packaging options
|
|
options: PackagingOptions,
|
|
}
|
|
|
|
impl ModelPackager {
|
|
/// Create a new model packager.
|
|
pub fn new(storage: Box<dyn StorageBackend>, options: PackagingOptions) -> Self {
|
|
Self { storage, options }
|
|
}
|
|
|
|
/// Package a model from a directory.
|
|
pub async fn package_from_directory(
|
|
&self,
|
|
model_dir: &Path,
|
|
metadata: ModelMetadata,
|
|
) -> HubResult<ModelPackage> {
|
|
// Validate input directory
|
|
if !model_dir.exists() || !model_dir.is_dir() {
|
|
return Err(HubError::InvalidPackage {
|
|
reason: format!("Invalid model directory: {}", model_dir.display()),
|
|
});
|
|
}
|
|
|
|
// Scan directory for assets
|
|
let assets = self.scan_assets(model_dir).await?;
|
|
|
|
// Create manifest
|
|
let manifest = self.create_manifest(model_dir).await?;
|
|
|
|
// Calculate package hash
|
|
let content_hash = self.calculate_package_hash(&assets, &manifest)?;
|
|
|
|
// Create package metadata
|
|
let package_metadata = PackageMetadata {
|
|
format_version: PACKAGE_FORMAT_VERSION.to_string(),
|
|
model: metadata,
|
|
created_at: Utc::now(),
|
|
created_by: whoami::username(),
|
|
size: manifest.total_size,
|
|
content_hash,
|
|
compression: CompressionSettings {
|
|
algorithm: self.options.compression.clone(),
|
|
level: self.options.compression_level,
|
|
original_size: manifest.total_size,
|
|
compressed_size: 0, // Will be updated after compression
|
|
ratio: 0.0, // Will be calculated after compression
|
|
},
|
|
signature: None, // Will be added if signing is enabled
|
|
};
|
|
|
|
// Create package
|
|
let mut package = ModelPackage {
|
|
metadata: package_metadata,
|
|
assets,
|
|
dependencies: vec![], // TODO: Extract from model metadata
|
|
manifest,
|
|
};
|
|
|
|
// Sign package if requested
|
|
if let Some(signing_options) = &self.options.sign {
|
|
package.metadata.signature = Some(self.sign_package(&package, signing_options)?);
|
|
}
|
|
|
|
// Validate package if requested
|
|
if self.options.validate {
|
|
self.validate_package(&package)?;
|
|
}
|
|
|
|
Ok(package)
|
|
}
|
|
|
|
/// Save a package to storage.
|
|
pub async fn save_package(&self, package: &ModelPackage, package_path: &str) -> HubResult<()> {
|
|
// Serialize package to bytes
|
|
let package_bytes = self.serialize_package(package)?;
|
|
|
|
// Compress if needed
|
|
let (compressed_bytes, _compression_stats) = self.compress_package(&package_bytes)?;
|
|
|
|
// Update compression statistics in package metadata
|
|
// Note: This would require making the package mutable, which we can't do here
|
|
// In a real implementation, you'd pass compression stats back to the caller
|
|
|
|
// Store package
|
|
self.storage.store(package_path, &compressed_bytes).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Load a package from storage.
|
|
pub async fn load_package(&self, package_path: &str) -> HubResult<ModelPackage> {
|
|
// Load compressed package data
|
|
let compressed_bytes = self.storage.load(package_path).await?;
|
|
|
|
// Decompress
|
|
let package_bytes = self.decompress_package(&compressed_bytes)?;
|
|
|
|
// Deserialize package
|
|
let package = self.deserialize_package(&package_bytes)?;
|
|
|
|
// Validate package integrity
|
|
if self.options.validate {
|
|
self.validate_package(&package)?;
|
|
}
|
|
|
|
Ok(package)
|
|
}
|
|
|
|
/// Extract a package to a directory.
|
|
pub async fn extract_package(
|
|
&self,
|
|
package: &ModelPackage,
|
|
target_dir: &Path,
|
|
) -> HubResult<()> {
|
|
// Create target directory
|
|
fs::create_dir_all(target_dir).await?;
|
|
|
|
// Extract each asset
|
|
for asset in &package.assets {
|
|
let asset_path = target_dir.join(&asset.path);
|
|
|
|
// Create parent directories
|
|
if let Some(parent) = asset_path.parent() {
|
|
fs::create_dir_all(parent).await?;
|
|
}
|
|
|
|
// Load asset data from storage
|
|
let asset_data = self.load_asset_data(asset).await?;
|
|
|
|
// Write asset to file
|
|
fs::write(&asset_path, asset_data).await?;
|
|
|
|
// Verify asset integrity
|
|
self.verify_asset_integrity(asset, &asset_path).await?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn scan_assets<'a>(
|
|
&'a self,
|
|
model_dir: &'a Path,
|
|
) -> Pin<Box<dyn Future<Output = HubResult<Vec<ModelAsset>>> + Send + 'a>> {
|
|
Box::pin(async move {
|
|
let mut assets = Vec::new();
|
|
let mut entries = fs::read_dir(model_dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
|
|
if path.is_file() {
|
|
// Skip files matching exclude patterns
|
|
if self.should_exclude(&path) {
|
|
continue;
|
|
}
|
|
|
|
let asset = self.create_asset(&path, model_dir).await?;
|
|
assets.push(asset);
|
|
} else if path.is_dir() {
|
|
// Recursively scan subdirectories
|
|
let sub_assets = self.scan_assets(&path).await?;
|
|
assets.extend(sub_assets);
|
|
}
|
|
}
|
|
|
|
Ok(assets)
|
|
})
|
|
}
|
|
|
|
async fn create_asset(&self, file_path: &Path, base_dir: &Path) -> HubResult<ModelAsset> {
|
|
let metadata = fs::metadata(file_path).await?;
|
|
let size = metadata.len();
|
|
|
|
// Calculate relative path
|
|
let relative_path = file_path
|
|
.strip_prefix(base_dir)
|
|
.map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Invalid path structure: {e}"),
|
|
})?
|
|
.to_string_lossy()
|
|
.to_string();
|
|
|
|
// Read file and calculate hash
|
|
let content = fs::read(file_path).await?;
|
|
let hash = format!("sha256:{}", hex::encode(Sha256::digest(&content)));
|
|
|
|
// Determine asset type from file extension and content
|
|
let asset_type = self.determine_asset_type(&relative_path, &content);
|
|
|
|
// Determine MIME type
|
|
let mime_type = self.determine_mime_type(&relative_path);
|
|
|
|
Ok(ModelAsset {
|
|
path: relative_path,
|
|
asset_type,
|
|
size,
|
|
hash,
|
|
mime_type,
|
|
metadata: HashMap::new(),
|
|
})
|
|
}
|
|
|
|
async fn create_manifest(&self, model_dir: &Path) -> HubResult<PackageManifest> {
|
|
let mut files = Vec::new();
|
|
let mut directories = Vec::new();
|
|
let mut total_size = 0;
|
|
|
|
self.scan_directory_recursive(
|
|
model_dir,
|
|
model_dir,
|
|
&mut files,
|
|
&mut directories,
|
|
&mut total_size,
|
|
)
|
|
.await?;
|
|
|
|
let file_count = files.len();
|
|
Ok(PackageManifest {
|
|
files,
|
|
directories,
|
|
file_count,
|
|
total_size,
|
|
})
|
|
}
|
|
|
|
fn scan_directory_recursive<'a>(
|
|
&'a self,
|
|
current_dir: &'a Path,
|
|
base_dir: &'a Path,
|
|
files: &'a mut Vec<ManifestEntry>,
|
|
directories: &'a mut Vec<String>,
|
|
total_size: &'a mut u64,
|
|
) -> Pin<Box<dyn Future<Output = HubResult<()>> + Send + 'a>> {
|
|
Box::pin(async move {
|
|
let mut entries = fs::read_dir(current_dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
let metadata = entry.metadata().await?;
|
|
|
|
let relative_path = path
|
|
.strip_prefix(base_dir)
|
|
.map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Invalid path structure: {e}"),
|
|
})?
|
|
.to_string_lossy()
|
|
.to_string();
|
|
|
|
if metadata.is_file() {
|
|
if !self.should_exclude(&path) {
|
|
let content = fs::read(&path).await?;
|
|
let hash = format!("sha256:{}", hex::encode(Sha256::digest(&content)));
|
|
|
|
files.push(ManifestEntry {
|
|
path: relative_path,
|
|
size: metadata.len(),
|
|
hash,
|
|
permissions: 0o644, // Default permissions
|
|
modified: DateTime::<Utc>::from(metadata.modified()?),
|
|
});
|
|
|
|
*total_size += metadata.len();
|
|
}
|
|
} else if metadata.is_dir() {
|
|
directories.push(relative_path);
|
|
self.scan_directory_recursive(&path, base_dir, files, directories, total_size)
|
|
.await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
fn calculate_package_hash(
|
|
&self,
|
|
assets: &[ModelAsset],
|
|
manifest: &PackageManifest,
|
|
) -> HubResult<String> {
|
|
let mut hasher = Sha256::new();
|
|
|
|
// Hash all asset hashes
|
|
for asset in assets {
|
|
hasher.update(asset.hash.as_bytes());
|
|
}
|
|
|
|
// Hash manifest
|
|
let manifest_json = serde_json::to_string(manifest)?;
|
|
hasher.update(manifest_json.as_bytes());
|
|
|
|
Ok(format!("sha256:{}", hex::encode(hasher.finalize())))
|
|
}
|
|
|
|
fn serialize_package(&self, package: &ModelPackage) -> HubResult<Vec<u8>> {
|
|
match self.options.compression {
|
|
CompressionAlgorithm::None => {
|
|
bincode::serialize(package).map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Failed to serialize package: {e}"),
|
|
})
|
|
}
|
|
_ => {
|
|
// For compressed formats, use JSON for better compression
|
|
serde_json::to_vec(package).map_err(|e| HubError::SerializationError { source: e })
|
|
}
|
|
}
|
|
}
|
|
|
|
fn deserialize_package(&self, bytes: &[u8]) -> HubResult<ModelPackage> {
|
|
match self.options.compression {
|
|
CompressionAlgorithm::None => {
|
|
bincode::deserialize(bytes).map_err(|e| HubError::InvalidPackage {
|
|
reason: format!("Failed to deserialize package: {e}"),
|
|
})
|
|
}
|
|
_ => serde_json::from_slice(bytes)
|
|
.map_err(|e| HubError::SerializationError { source: e }),
|
|
}
|
|
}
|
|
|
|
fn compress_package(&self, bytes: &[u8]) -> HubResult<(Vec<u8>, CompressionSettings)> {
|
|
let original_size = bytes.len() as u64;
|
|
|
|
let (compressed_bytes, algorithm) = match self.options.compression {
|
|
CompressionAlgorithm::None => (bytes.to_vec(), CompressionAlgorithm::None),
|
|
CompressionAlgorithm::Gzip => {
|
|
let mut encoder = flate2::write::GzEncoder::new(
|
|
Vec::new(),
|
|
flate2::Compression::new(self.options.compression_level as u32),
|
|
);
|
|
encoder.write_all(bytes)?;
|
|
(encoder.finish()?, CompressionAlgorithm::Gzip)
|
|
}
|
|
CompressionAlgorithm::Zstd => {
|
|
let compressed = zstd::bulk::compress(bytes, self.options.compression_level as i32)
|
|
.map_err(|e| HubError::CompressionError {
|
|
details: e.to_string(),
|
|
})?;
|
|
(compressed, CompressionAlgorithm::Zstd)
|
|
}
|
|
CompressionAlgorithm::Lz4 => {
|
|
// For now, fall back to no compression for LZ4
|
|
(bytes.to_vec(), CompressionAlgorithm::None)
|
|
}
|
|
};
|
|
|
|
let compressed_size = compressed_bytes.len() as u64;
|
|
let ratio = if original_size > 0 {
|
|
compressed_size as f64 / original_size as f64
|
|
} else {
|
|
1.0
|
|
};
|
|
|
|
let settings = CompressionSettings {
|
|
algorithm,
|
|
level: self.options.compression_level,
|
|
original_size,
|
|
compressed_size,
|
|
ratio,
|
|
};
|
|
|
|
Ok((compressed_bytes, settings))
|
|
}
|
|
|
|
fn decompress_package(&self, bytes: &[u8]) -> HubResult<Vec<u8>> {
|
|
match self.options.compression {
|
|
CompressionAlgorithm::None => Ok(bytes.to_vec()),
|
|
CompressionAlgorithm::Gzip => {
|
|
let mut decoder = flate2::read::GzDecoder::new(bytes);
|
|
let mut decompressed = Vec::new();
|
|
decoder.read_to_end(&mut decompressed)?;
|
|
Ok(decompressed)
|
|
}
|
|
CompressionAlgorithm::Zstd => {
|
|
zstd::bulk::decompress(bytes, 1024 * 1024 * 100) // 100MB limit
|
|
.map_err(|e| HubError::CompressionError {
|
|
details: e.to_string(),
|
|
})
|
|
}
|
|
CompressionAlgorithm::Lz4 => Ok(bytes.to_vec()),
|
|
}
|
|
}
|
|
|
|
fn sign_package(
|
|
&self,
|
|
package: &ModelPackage,
|
|
signing_options: &SigningOptions,
|
|
) -> HubResult<PackageSignature> {
|
|
// For now, create a placeholder signature
|
|
// In a real implementation, you'd use proper cryptographic signing
|
|
let package_json = serde_json::to_string(package)?;
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(package_json.as_bytes());
|
|
hasher.update(&signing_options.private_key);
|
|
let signature_hash = hasher.finalize();
|
|
|
|
Ok(PackageSignature {
|
|
algorithm: signing_options.algorithm.clone(),
|
|
signature: base64::engine::general_purpose::STANDARD.encode(signature_hash),
|
|
key_id: signing_options.key_id.clone(),
|
|
signed_at: Utc::now(),
|
|
})
|
|
}
|
|
|
|
fn validate_package(&self, package: &ModelPackage) -> HubResult<()> {
|
|
// Validate package format version
|
|
if package.metadata.format_version != PACKAGE_FORMAT_VERSION {
|
|
return Err(HubError::ValidationFailed {
|
|
details: format!(
|
|
"Unsupported package format version: {}",
|
|
package.metadata.format_version
|
|
),
|
|
});
|
|
}
|
|
|
|
// Validate asset count matches manifest
|
|
if package.assets.len() != package.manifest.file_count {
|
|
return Err(HubError::ValidationFailed {
|
|
details: "Asset count mismatch with manifest".to_string(),
|
|
});
|
|
}
|
|
|
|
// Validate required assets are present
|
|
let has_weights = package
|
|
.assets
|
|
.iter()
|
|
.any(|a| a.asset_type == AssetType::Weights);
|
|
let has_config = package
|
|
.assets
|
|
.iter()
|
|
.any(|a| a.asset_type == AssetType::Config);
|
|
|
|
if !has_weights {
|
|
return Err(HubError::ValidationFailed {
|
|
details: "Package must contain model weights".to_string(),
|
|
});
|
|
}
|
|
|
|
if !has_config {
|
|
return Err(HubError::ValidationFailed {
|
|
details: "Package must contain model configuration".to_string(),
|
|
});
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn load_asset_data(&self, _asset: &ModelAsset) -> HubResult<Vec<u8>> {
|
|
// In a real implementation, you'd load asset data from storage
|
|
// For now, return empty data
|
|
Ok(Vec::new())
|
|
}
|
|
|
|
async fn verify_asset_integrity(&self, asset: &ModelAsset, file_path: &Path) -> HubResult<()> {
|
|
let content = fs::read(file_path).await?;
|
|
let hash = format!("sha256:{}", hex::encode(Sha256::digest(&content)));
|
|
|
|
if hash != asset.hash {
|
|
return Err(HubError::ValidationFailed {
|
|
details: format!("Asset integrity check failed for {}", asset.path),
|
|
});
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn should_exclude(&self, path: &Path) -> bool {
|
|
let path_str = path.to_string_lossy();
|
|
// Also get just the filename for simple patterns like *.log
|
|
let filename = path
|
|
.file_name()
|
|
.map(|n| n.to_string_lossy().to_string())
|
|
.unwrap_or_default();
|
|
|
|
// Check exclude patterns against both full path and filename
|
|
for pattern in &self.options.exclude_patterns {
|
|
if glob_match::glob_match(pattern, &path_str)
|
|
|| glob_match::glob_match(pattern, &filename)
|
|
{
|
|
return true;
|
|
}
|
|
}
|
|
|
|
// Check include patterns
|
|
if !self.options.include_patterns.is_empty() {
|
|
let included = self.options.include_patterns.iter().any(|pattern| {
|
|
glob_match::glob_match(pattern, &path_str)
|
|
|| glob_match::glob_match(pattern, &filename)
|
|
});
|
|
return !included;
|
|
}
|
|
|
|
false
|
|
}
|
|
|
|
fn determine_asset_type(&self, path: &str, _content: &[u8]) -> AssetType {
|
|
let path_lower = path.to_lowercase();
|
|
|
|
// Check specific patterns before generic ones
|
|
if path_lower.contains("tokenizer") {
|
|
AssetType::Tokenizer
|
|
} else if path_lower.contains("vocab") {
|
|
AssetType::Vocabulary
|
|
} else if path_lower.contains("weight")
|
|
|| path_lower.ends_with(".pth")
|
|
|| path_lower.ends_with(".safetensors")
|
|
{
|
|
AssetType::Weights
|
|
} else if path_lower.contains("config")
|
|
|| (path_lower.ends_with(".json") && !path_lower.contains("tokenizer"))
|
|
{
|
|
AssetType::Config
|
|
} else if path_lower.ends_with(".py") || path_lower.ends_with(".rs") {
|
|
AssetType::Code
|
|
} else if path_lower.ends_with(".md") || path_lower.ends_with(".txt") {
|
|
AssetType::Documentation
|
|
} else if path_lower.contains("example") {
|
|
AssetType::Examples
|
|
} else if path_lower.contains("license") {
|
|
AssetType::License
|
|
} else {
|
|
AssetType::Custom(path.split('.').next_back().unwrap_or("unknown").to_string())
|
|
}
|
|
}
|
|
|
|
fn determine_mime_type(&self, path: &str) -> String {
|
|
let extension = path.split('.').next_back().unwrap_or("").to_lowercase();
|
|
|
|
match extension.as_str() {
|
|
"json" => "application/json",
|
|
"txt" | "md" => "text/plain",
|
|
"py" => "text/x-python",
|
|
"rs" => "text/x-rust",
|
|
"pth" => "application/octet-stream",
|
|
"safetensors" => "application/octet-stream",
|
|
_ => "application/octet-stream",
|
|
}
|
|
.to_string()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::model::{ModelId, ModelMetadata, ModelSchema, ModelStatus};
|
|
use crate::storage::{LocalStorageBackend, StorageBackend};
|
|
use crate::versioning::ModelVersion;
|
|
use chrono::Utc;
|
|
use semver::Version;
|
|
use std::collections::HashMap;
|
|
use tempfile::TempDir;
|
|
use tokio::fs;
|
|
|
|
async fn create_test_model_dir() -> HubResult<TempDir> {
|
|
let temp_dir = TempDir::new()?;
|
|
|
|
// Create test files
|
|
let model_path = temp_dir.path();
|
|
|
|
// Create config.json
|
|
let config = serde_json::json!({
|
|
"model_type": "transformer",
|
|
"hidden_size": 768,
|
|
"num_layers": 12
|
|
});
|
|
fs::write(model_path.join("config.json"), config.to_string()).await?;
|
|
|
|
// Create model weights (dummy file)
|
|
fs::write(model_path.join("model.safetensors"), b"dummy weights data").await?;
|
|
|
|
// Create tokenizer data
|
|
fs::write(model_path.join("tokenizer.json"), "{}").await?;
|
|
|
|
// Create README
|
|
fs::write(
|
|
model_path.join("README.md"),
|
|
"# Test Model\n\nThis is a test model.",
|
|
)
|
|
.await?;
|
|
|
|
Ok(temp_dir)
|
|
}
|
|
|
|
fn create_test_metadata() -> ModelMetadata {
|
|
ModelMetadata {
|
|
id: ModelId::new("test", "model"),
|
|
version: ModelVersion::new(Version::parse("1.0.0").unwrap()),
|
|
title: "Test Model".to_string(),
|
|
description: "A test model for packaging".to_string(),
|
|
architecture: "transformer".to_string(),
|
|
framework: "rustytorch".to_string(),
|
|
framework_version: "1.0.0".to_string(),
|
|
tags: vec!["test".to_string()],
|
|
author: "Test Author".to_string(),
|
|
license: Some("MIT".to_string()),
|
|
created_at: Utc::now(),
|
|
updated_at: Utc::now(),
|
|
status: ModelStatus::Available,
|
|
size: 1024,
|
|
content_hash: "test".to_string(),
|
|
dependencies: vec![],
|
|
schema: ModelSchema {
|
|
inputs: vec![],
|
|
outputs: vec![],
|
|
config: None,
|
|
},
|
|
metrics: HashMap::new(),
|
|
metadata: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_package_creation() {
|
|
let model_dir = create_test_model_dir().await.unwrap();
|
|
let metadata = create_test_metadata();
|
|
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
let options = PackagingOptions::default();
|
|
let packager = ModelPackager::new(storage, options);
|
|
|
|
let package = packager
|
|
.package_from_directory(model_dir.path(), metadata)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Verify package structure
|
|
assert_eq!(package.metadata.format_version, PACKAGE_FORMAT_VERSION);
|
|
assert!(!package.assets.is_empty());
|
|
assert!(package.assets.len() >= 4); // config, weights, tokenizer, README
|
|
|
|
// Check for required assets
|
|
let has_config = package
|
|
.assets
|
|
.iter()
|
|
.any(|a| a.asset_type == AssetType::Config);
|
|
let has_weights = package
|
|
.assets
|
|
.iter()
|
|
.any(|a| a.asset_type == AssetType::Weights);
|
|
assert!(has_config);
|
|
assert!(has_weights);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_asset_type_detection() {
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
let options = PackagingOptions::default();
|
|
let packager = ModelPackager::new(storage, options);
|
|
|
|
// Test different asset types
|
|
assert_eq!(
|
|
packager.determine_asset_type("config.json", b"{}"),
|
|
AssetType::Config
|
|
);
|
|
assert_eq!(
|
|
packager.determine_asset_type("model.safetensors", b""),
|
|
AssetType::Weights
|
|
);
|
|
assert_eq!(
|
|
packager.determine_asset_type("tokenizer.json", b""),
|
|
AssetType::Tokenizer
|
|
);
|
|
assert_eq!(
|
|
packager.determine_asset_type("README.md", b""),
|
|
AssetType::Documentation
|
|
);
|
|
|
|
if let AssetType::Custom(ext) = packager.determine_asset_type("test.xyz", b"") {
|
|
assert_eq!(ext, "xyz");
|
|
} else {
|
|
panic!("Expected Custom asset type");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_compression_algorithms() {
|
|
let test_data =
|
|
b"This is test data that should be compressed well with repetition. ".repeat(100);
|
|
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
// Test different compression algorithms
|
|
let algorithms = vec![
|
|
CompressionAlgorithm::None,
|
|
CompressionAlgorithm::Gzip,
|
|
CompressionAlgorithm::Zstd,
|
|
];
|
|
|
|
for algorithm in algorithms {
|
|
let mut options = PackagingOptions::default();
|
|
options.compression = algorithm.clone();
|
|
|
|
let packager = ModelPackager::new(
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf())),
|
|
options,
|
|
);
|
|
|
|
let (compressed, stats) = packager.compress_package(&test_data).unwrap();
|
|
|
|
match algorithm {
|
|
CompressionAlgorithm::None => {
|
|
assert_eq!(compressed.len(), test_data.len());
|
|
assert_eq!(stats.ratio, 1.0);
|
|
}
|
|
_ => {
|
|
// Compression should reduce size for repetitive data
|
|
assert!(compressed.len() < test_data.len());
|
|
assert!(stats.ratio < 1.0);
|
|
}
|
|
}
|
|
|
|
// Test decompression
|
|
let decompressed = packager.decompress_package(&compressed).unwrap();
|
|
assert_eq!(decompressed, test_data);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_package_validation() {
|
|
let model_dir = create_test_model_dir().await.unwrap();
|
|
let metadata = create_test_metadata();
|
|
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
let options = PackagingOptions::default();
|
|
let packager = ModelPackager::new(storage, options);
|
|
|
|
let package = packager
|
|
.package_from_directory(model_dir.path(), metadata)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Valid package should pass validation
|
|
assert!(packager.validate_package(&package).is_ok());
|
|
|
|
// Test invalid format version
|
|
let mut invalid_package = package.clone();
|
|
invalid_package.metadata.format_version = "0.1".to_string();
|
|
assert!(packager.validate_package(&invalid_package).is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_exclude_patterns() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let model_path = temp_dir.path();
|
|
|
|
// Create files that should be excluded
|
|
fs::write(model_path.join("model.log"), "log data")
|
|
.await
|
|
.unwrap();
|
|
fs::write(model_path.join("temp.tmp"), "temp data")
|
|
.await
|
|
.unwrap();
|
|
fs::write(model_path.join("config.json"), "{}")
|
|
.await
|
|
.unwrap();
|
|
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
let options = PackagingOptions::default();
|
|
let packager = ModelPackager::new(storage, options);
|
|
|
|
let assets = packager.scan_assets(model_path).await.unwrap();
|
|
|
|
// Only config.json should be included
|
|
assert_eq!(assets.len(), 1);
|
|
assert_eq!(assets[0].path, "config.json");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_package_serialization() {
|
|
let model_dir = create_test_model_dir().await.unwrap();
|
|
let metadata = create_test_metadata();
|
|
|
|
let storage_dir = TempDir::new().unwrap();
|
|
let storage: Box<dyn StorageBackend> =
|
|
Box::new(LocalStorageBackend::new(storage_dir.path().to_path_buf()));
|
|
|
|
let options = PackagingOptions::default();
|
|
let packager = ModelPackager::new(storage, options);
|
|
|
|
let original_package = packager
|
|
.package_from_directory(model_dir.path(), metadata)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Test serialization/deserialization
|
|
let serialized = packager.serialize_package(&original_package).unwrap();
|
|
let deserialized = packager.deserialize_package(&serialized).unwrap();
|
|
|
|
assert_eq!(
|
|
original_package.metadata.model.id,
|
|
deserialized.metadata.model.id
|
|
);
|
|
assert_eq!(original_package.assets.len(), deserialized.assets.len());
|
|
}
|
|
|
|
#[test]
|
|
fn test_packaging_options_default() {
|
|
let options = PackagingOptions::default();
|
|
|
|
assert_eq!(options.compression, CompressionAlgorithm::Zstd);
|
|
assert_eq!(options.compression_level, 6);
|
|
assert!(options.include_code);
|
|
assert!(options.include_docs);
|
|
assert!(!options.include_examples);
|
|
assert!(options.validate);
|
|
assert!(options.sign.is_none());
|
|
assert!(!options.exclude_patterns.is_empty());
|
|
}
|
|
}
|