572 lines
17 KiB
Rust
572 lines
17 KiB
Rust
//! # RTX Hub - Model Registry and Versioning System
|
|
//!
|
|
//! RTX Hub provides a comprehensive model registry and versioning system for
|
|
//! RustyTorch production deployments. It handles model packaging, versioning,
|
|
//! metadata management, and dependency tracking.
|
|
//!
|
|
//! ## Core Features
|
|
//!
|
|
//! - Content-addressable model storage
|
|
//! - Semantic versioning with dependency resolution
|
|
//! - Model metadata and configuration management
|
|
//! - Multi-backend storage support (local, S3, GCS, Azure)
|
|
//! - Model validation and integrity checking
|
|
//! - Registry synchronization and replication
|
|
//! - HuggingFace-style pretrained model loading API
|
|
//! - Comprehensive model search and discovery
|
|
//! - SafeTensors file format support
|
|
//! - HuggingFace Hub compatibility
|
|
//!
|
|
//! ## Quick Start
|
|
//!
|
|
//! Load a model with a single line:
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_hub::load;
|
|
//!
|
|
//! // Load a model from RustyTorch Hub
|
|
//! let model = load("rustytorch/llama-3-8b-q4").await?;
|
|
//!
|
|
//! // Or from HuggingFace Hub
|
|
//! let model = load("meta-llama/Llama-3-8B").await?;
|
|
//! ```
|
|
//!
|
|
//! ## Loading Options
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_hub::{load_with_config, LoadConfig};
|
|
//!
|
|
//! let config = LoadConfig::default()
|
|
//! .with_revision("v1.0.0")
|
|
//! .with_dtype(DType::BF16)
|
|
//! .with_device("cuda:0");
|
|
//!
|
|
//! let model = load_with_config("rustytorch/llama-3-8b", config).await?;
|
|
//! ```
|
|
|
|
#![deny(clippy::unwrap_used)]
|
|
#![cfg_attr(test, allow(clippy::unwrap_used))]
|
|
|
|
pub mod discovery;
|
|
pub mod error;
|
|
pub mod huggingface;
|
|
pub mod model;
|
|
pub mod model_card;
|
|
pub mod packaging;
|
|
pub mod pretrained;
|
|
pub mod registry;
|
|
pub mod safetensors;
|
|
pub mod storage;
|
|
pub mod versioning;
|
|
pub mod weights;
|
|
|
|
pub use discovery::{
|
|
ModelDiscovery, ModelMetrics, ModelSummary, SearchQuery, SearchResult, SortBy,
|
|
};
|
|
pub use error::{HubError, HubResult};
|
|
pub use huggingface::{
|
|
HFGenerationConfig, HFHubClient, HFModelConfig, HFTokenizerConfig, RTXActivation,
|
|
RTXArchitecture, RTXDType, RTXModelConfig, convert_hf_config,
|
|
};
|
|
pub use model::{ModelId, ModelInfo, ModelMetadata, ModelStatus};
|
|
pub use model_card::{
|
|
CO2Emissions, MetricResult, ModelCardBuilder, ModelCardDoc, ModelIndex, Widget,
|
|
};
|
|
pub use packaging::{ModelPackage, ModelPackager, PackagingOptions};
|
|
pub use pretrained::{FromPretrained, LocalCache, PretrainedConfig, PretrainedLoader};
|
|
pub use registry::{ModelRegistry, QueryOptions, Registry, RegistryConfig, SortField};
|
|
pub use safetensors::{
|
|
SafeTensors, SafeTensorsBuilder, SafeTensorsDType, ShardedSafeTensors, TensorInfo,
|
|
};
|
|
pub use storage::{StorageBackend, StorageConfig};
|
|
pub use versioning::{
|
|
DependencySpec, ModelVersion, ResolvedDependency, VersionConstraint, VersionResolver,
|
|
};
|
|
pub use weights::{ShardedWeights, WeightManager, WeightShard};
|
|
|
|
use std::path::PathBuf;
|
|
use tracing::info;
|
|
|
|
// ============================================================================
|
|
// One-Command Model Loading API
|
|
// ============================================================================
|
|
|
|
/// Configuration for model loading.
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct LoadConfig {
|
|
/// Git-like revision (e.g., "main", "v1.0.0")
|
|
pub revision: Option<String>,
|
|
/// Target data type for weights
|
|
pub dtype: Option<RTXDType>,
|
|
/// Target device (e.g., "cuda:0", "metal", "cpu")
|
|
pub device: Option<String>,
|
|
/// Custom cache directory
|
|
pub cache_dir: Option<PathBuf>,
|
|
/// Force re-download even if cached
|
|
pub force_download: bool,
|
|
/// Offline mode (local files only)
|
|
pub offline: bool,
|
|
/// Authentication token for private models
|
|
pub auth_token: Option<String>,
|
|
/// Use HuggingFace Hub instead of RustyTorch Hub
|
|
pub use_hf_hub: bool,
|
|
/// Trust remote code (for custom model classes)
|
|
pub trust_remote_code: bool,
|
|
}
|
|
|
|
impl LoadConfig {
|
|
/// Create a new load configuration with defaults.
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Set the revision/tag.
|
|
pub fn with_revision(mut self, revision: impl Into<String>) -> Self {
|
|
self.revision = Some(revision.into());
|
|
self
|
|
}
|
|
|
|
/// Set the target data type.
|
|
pub fn with_dtype(mut self, dtype: RTXDType) -> Self {
|
|
self.dtype = Some(dtype);
|
|
self
|
|
}
|
|
|
|
/// Set the target device.
|
|
pub fn with_device(mut self, device: impl Into<String>) -> Self {
|
|
self.device = Some(device.into());
|
|
self
|
|
}
|
|
|
|
/// Set the cache directory.
|
|
pub fn with_cache_dir(mut self, cache_dir: impl Into<PathBuf>) -> Self {
|
|
self.cache_dir = Some(cache_dir.into());
|
|
self
|
|
}
|
|
|
|
/// Enable force download.
|
|
pub fn with_force_download(mut self, force: bool) -> Self {
|
|
self.force_download = force;
|
|
self
|
|
}
|
|
|
|
/// Enable offline mode.
|
|
pub fn with_offline(mut self, offline: bool) -> Self {
|
|
self.offline = offline;
|
|
self
|
|
}
|
|
|
|
/// Set authentication token.
|
|
pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
|
|
self.auth_token = Some(token.into());
|
|
self
|
|
}
|
|
|
|
/// Use HuggingFace Hub.
|
|
pub fn with_hf_hub(mut self) -> Self {
|
|
self.use_hf_hub = true;
|
|
self
|
|
}
|
|
|
|
/// Trust remote code.
|
|
pub fn with_trust_remote_code(mut self, trust: bool) -> Self {
|
|
self.trust_remote_code = trust;
|
|
self
|
|
}
|
|
|
|
/// Get the cache directory, using default if not set.
|
|
pub fn get_cache_dir(&self) -> PathBuf {
|
|
self.cache_dir.clone().unwrap_or_else(|| {
|
|
dirs::home_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join(".cache")
|
|
.join("rtx-hub")
|
|
})
|
|
}
|
|
|
|
/// Convert to PretrainedConfig.
|
|
pub fn to_pretrained_config(&self) -> PretrainedConfig {
|
|
PretrainedConfig {
|
|
revision: self.revision.clone(),
|
|
cache_dir: self.cache_dir.clone(),
|
|
force_download: self.force_download,
|
|
resume_download: true,
|
|
local_files_only: self.offline,
|
|
use_auth_token: self.auth_token.clone(),
|
|
mirror: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Loaded model with configuration and weights path.
|
|
#[derive(Debug)]
|
|
pub struct LoadedModel {
|
|
/// Model ID
|
|
pub model_id: String,
|
|
/// Path to cached model files
|
|
pub path: PathBuf,
|
|
/// Model configuration (converted from HF format)
|
|
pub config: RTXModelConfig,
|
|
/// Original HuggingFace configuration
|
|
pub hf_config: HFModelConfig,
|
|
/// Tokenizer configuration (if available)
|
|
pub tokenizer_config: Option<HFTokenizerConfig>,
|
|
/// Generation configuration (if available)
|
|
pub generation_config: Option<HFGenerationConfig>,
|
|
/// Load configuration used
|
|
pub load_config: LoadConfig,
|
|
}
|
|
|
|
impl LoadedModel {
|
|
/// Get the path to the weights file.
|
|
pub fn weights_path(&self) -> PathBuf {
|
|
// Check for sharded model first
|
|
let index_path = self.path.join("model.safetensors.index.json");
|
|
if index_path.exists() {
|
|
return self.path.clone();
|
|
}
|
|
|
|
// Single file
|
|
let safetensors_path = self.path.join("model.safetensors");
|
|
if safetensors_path.exists() {
|
|
return safetensors_path;
|
|
}
|
|
|
|
// Fallback to PyTorch format
|
|
self.path.join("pytorch_model.bin")
|
|
}
|
|
|
|
/// Check if model uses sharded weights.
|
|
pub fn is_sharded(&self) -> bool {
|
|
self.path.join("model.safetensors.index.json").exists()
|
|
}
|
|
|
|
/// Load SafeTensors weights.
|
|
pub async fn load_safetensors(&self) -> HubResult<SafeTensors> {
|
|
let weights_path = self.weights_path();
|
|
|
|
if weights_path.is_dir() {
|
|
// Sharded model - return error, use load_sharded_safetensors instead
|
|
return Err(HubError::InvalidPackage {
|
|
reason: "Model uses sharded weights. Use load_sharded_safetensors() instead."
|
|
.to_string(),
|
|
});
|
|
}
|
|
|
|
SafeTensors::load(&weights_path).await
|
|
}
|
|
|
|
/// Load sharded SafeTensors weights.
|
|
pub async fn load_sharded_safetensors(&self) -> HubResult<ShardedSafeTensors> {
|
|
ShardedSafeTensors::load_from_dir(&self.path).await
|
|
}
|
|
|
|
/// Get model architecture.
|
|
pub fn architecture(&self) -> RTXArchitecture {
|
|
self.config.architecture
|
|
}
|
|
|
|
/// Get vocabulary size.
|
|
pub fn vocab_size(&self) -> usize {
|
|
self.config.vocab_size
|
|
}
|
|
|
|
/// Get hidden size.
|
|
pub fn hidden_size(&self) -> usize {
|
|
self.config.hidden_size
|
|
}
|
|
|
|
/// Get number of layers.
|
|
pub fn num_layers(&self) -> usize {
|
|
self.config.num_layers
|
|
}
|
|
|
|
/// Get number of attention heads.
|
|
pub fn num_heads(&self) -> usize {
|
|
self.config.num_heads
|
|
}
|
|
}
|
|
|
|
/// Load a model with default configuration.
|
|
///
|
|
/// This is the simplest way to load a model:
|
|
///
|
|
/// ```rust,ignore
|
|
/// let model = rtx_hub::load("rustytorch/llama-3-8b-q4").await?;
|
|
/// ```
|
|
///
|
|
/// Model IDs can be:
|
|
/// - RustyTorch Hub: `rustytorch/model-name`
|
|
/// - HuggingFace Hub: `meta-llama/Llama-3-8B`
|
|
/// - Local path: `/path/to/model`
|
|
pub async fn load(model_id: &str) -> HubResult<LoadedModel> {
|
|
load_with_config(model_id, LoadConfig::default()).await
|
|
}
|
|
|
|
/// Load a model with custom configuration.
|
|
///
|
|
/// ```rust,ignore
|
|
/// let config = LoadConfig::default()
|
|
/// .with_revision("v1.0.0")
|
|
/// .with_dtype(RTXDType::BF16);
|
|
///
|
|
/// let model = rtx_hub::load_with_config("rustytorch/llama-3-8b", config).await?;
|
|
/// ```
|
|
pub async fn load_with_config(model_id: &str, config: LoadConfig) -> HubResult<LoadedModel> {
|
|
info!("Loading model: {}", model_id);
|
|
|
|
// Check if it's a local path
|
|
let local_path = PathBuf::from(model_id);
|
|
if local_path.exists() && local_path.is_dir() {
|
|
return load_from_local(model_id, local_path, config).await;
|
|
}
|
|
|
|
// Determine which hub to use
|
|
let use_hf = config.use_hf_hub || is_hf_model_id(model_id);
|
|
|
|
if use_hf {
|
|
load_from_hf_hub(model_id, config).await
|
|
} else {
|
|
load_from_rtx_hub(model_id, config).await
|
|
}
|
|
}
|
|
|
|
/// Check if model ID looks like a HuggingFace model.
|
|
fn is_hf_model_id(model_id: &str) -> bool {
|
|
// HuggingFace models typically have format "org/model" or "user/model"
|
|
// RustyTorch hub uses "rustytorch/model" prefix
|
|
if model_id.starts_with("rustytorch/") {
|
|
return false;
|
|
}
|
|
|
|
// Known HuggingFace organizations
|
|
let hf_orgs = [
|
|
"meta-llama/",
|
|
"mistralai/",
|
|
"openai/",
|
|
"google/",
|
|
"facebook/",
|
|
"microsoft/",
|
|
"EleutherAI/",
|
|
"huggingface/",
|
|
"bigscience/",
|
|
"stabilityai/",
|
|
"Qwen/",
|
|
];
|
|
|
|
hf_orgs.iter().any(|org| model_id.starts_with(org))
|
|
}
|
|
|
|
/// Load model from local directory.
|
|
async fn load_from_local(
|
|
model_id: &str,
|
|
path: PathBuf,
|
|
config: LoadConfig,
|
|
) -> HubResult<LoadedModel> {
|
|
info!("Loading local model from: {}", path.display());
|
|
|
|
// Load HuggingFace config
|
|
let config_path = path.join("config.json");
|
|
if !config_path.exists() {
|
|
return Err(HubError::InvalidPackage {
|
|
reason: format!("config.json not found in {}", path.display()),
|
|
});
|
|
}
|
|
|
|
let config_content = tokio::fs::read_to_string(&config_path).await?;
|
|
let hf_config: HFModelConfig = serde_json::from_str(&config_content)?;
|
|
|
|
// Convert to RTX config
|
|
let rtx_config = convert_hf_config(&hf_config)?;
|
|
|
|
// Load optional configs
|
|
let tokenizer_config =
|
|
load_optional_json::<HFTokenizerConfig>(&path.join("tokenizer_config.json")).await;
|
|
let generation_config =
|
|
load_optional_json::<HFGenerationConfig>(&path.join("generation_config.json")).await;
|
|
|
|
Ok(LoadedModel {
|
|
model_id: model_id.to_string(),
|
|
path,
|
|
config: rtx_config,
|
|
hf_config,
|
|
tokenizer_config,
|
|
generation_config,
|
|
load_config: config,
|
|
})
|
|
}
|
|
|
|
/// Load model from HuggingFace Hub.
|
|
async fn load_from_hf_hub(model_id: &str, config: LoadConfig) -> HubResult<LoadedModel> {
|
|
info!("Loading from HuggingFace Hub: {}", model_id);
|
|
|
|
let cache_dir = config.get_cache_dir();
|
|
let mut client = HFHubClient::new(cache_dir);
|
|
|
|
if let Some(ref token) = config.auth_token {
|
|
client = client.with_token(token);
|
|
}
|
|
|
|
let pretrained_config = config.to_pretrained_config();
|
|
let model_path = client.download_model(model_id, &pretrained_config).await?;
|
|
|
|
// Load configs
|
|
let hf_config = client.load_config(&model_path).await?;
|
|
let rtx_config = convert_hf_config(&hf_config)?;
|
|
|
|
let tokenizer_config = client.load_tokenizer_config(&model_path).await.ok();
|
|
let generation_config = client.load_generation_config(&model_path).await.ok();
|
|
|
|
Ok(LoadedModel {
|
|
model_id: model_id.to_string(),
|
|
path: model_path,
|
|
config: rtx_config,
|
|
hf_config,
|
|
tokenizer_config,
|
|
generation_config,
|
|
load_config: config,
|
|
})
|
|
}
|
|
|
|
/// Load model from RustyTorch Hub.
|
|
async fn load_from_rtx_hub(model_id: &str, config: LoadConfig) -> HubResult<LoadedModel> {
|
|
info!("Loading from RustyTorch Hub: {}", model_id);
|
|
|
|
// For now, fall back to HuggingFace Hub loading
|
|
// In the future, this will connect to the RustyTorch model registry
|
|
load_from_hf_hub(model_id, config.clone())
|
|
.await
|
|
.map_err(|_| HubError::ModelNotFound {
|
|
model_id: format!(
|
|
"{} (RustyTorch Hub not available, try with .with_hf_hub())",
|
|
model_id
|
|
),
|
|
})
|
|
}
|
|
|
|
/// Load optional JSON file.
|
|
async fn load_optional_json<T: serde::de::DeserializeOwned>(path: &PathBuf) -> Option<T> {
|
|
if !path.exists() {
|
|
return None;
|
|
}
|
|
|
|
match tokio::fs::read_to_string(path).await {
|
|
Ok(content) => serde_json::from_str(&content).ok(),
|
|
Err(_) => None,
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Model Loading Utilities
|
|
// ============================================================================
|
|
|
|
/// List available models in cache.
|
|
pub async fn list_cached_models(cache_dir: Option<PathBuf>) -> HubResult<Vec<String>> {
|
|
let cache_dir = cache_dir.unwrap_or_else(|| {
|
|
dirs::home_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join(".cache")
|
|
.join("rtx-hub")
|
|
.join("models")
|
|
});
|
|
|
|
if !cache_dir.exists() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let mut models = Vec::new();
|
|
let mut entries = tokio::fs::read_dir(&cache_dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let name = entry.file_name().to_string_lossy().to_string();
|
|
// Convert back from cache format (namespace--model) to model ID
|
|
let model_id = name.replace("--", "/");
|
|
models.push(model_id);
|
|
}
|
|
|
|
Ok(models)
|
|
}
|
|
|
|
/// Clear cached models.
|
|
pub async fn clear_cache(cache_dir: Option<PathBuf>) -> HubResult<()> {
|
|
let cache_dir = cache_dir.unwrap_or_else(|| {
|
|
dirs::home_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join(".cache")
|
|
.join("rtx-hub")
|
|
});
|
|
|
|
if cache_dir.exists() {
|
|
tokio::fs::remove_dir_all(&cache_dir).await?;
|
|
}
|
|
|
|
info!("Cache cleared: {}", cache_dir.display());
|
|
Ok(())
|
|
}
|
|
|
|
/// Get cache size in bytes.
|
|
pub async fn cache_size(cache_dir: Option<PathBuf>) -> HubResult<u64> {
|
|
let cache_dir = cache_dir.unwrap_or_else(|| {
|
|
dirs::home_dir()
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
.join(".cache")
|
|
.join("rtx-hub")
|
|
});
|
|
|
|
if !cache_dir.exists() {
|
|
return Ok(0);
|
|
}
|
|
|
|
calculate_dir_size(&cache_dir).await
|
|
}
|
|
|
|
/// Calculate directory size recursively.
|
|
async fn calculate_dir_size(dir: &PathBuf) -> HubResult<u64> {
|
|
let mut total_size = 0u64;
|
|
|
|
let mut entries = tokio::fs::read_dir(dir).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
let metadata = tokio::fs::metadata(&path).await?;
|
|
|
|
if metadata.is_file() {
|
|
total_size += metadata.len();
|
|
} else if metadata.is_dir() {
|
|
total_size += Box::pin(calculate_dir_size(&path)).await?;
|
|
}
|
|
}
|
|
|
|
Ok(total_size)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tempfile::TempDir;
|
|
|
|
/// Integration test setup helper
|
|
pub struct TestRegistry {
|
|
pub registry: ModelRegistry,
|
|
pub temp_dir: TempDir,
|
|
}
|
|
|
|
impl TestRegistry {
|
|
pub async fn new() -> HubResult<Self> {
|
|
let temp_dir = TempDir::new()?;
|
|
let config = RegistryConfig {
|
|
storage: StorageConfig::Local {
|
|
base_path: temp_dir.path().to_path_buf(),
|
|
},
|
|
database_url: format!("sqlite://{}/registry.db", temp_dir.path().display()),
|
|
enable_validation: true,
|
|
enable_compression: true,
|
|
..Default::default()
|
|
};
|
|
|
|
let registry = ModelRegistry::new(config).await?;
|
|
Ok(Self { registry, temp_dir })
|
|
}
|
|
}
|
|
}
|