691 lines
20 KiB
Rust
691 lines
20 KiB
Rust
//! Storage backends for model registry data and packages.
|
|
|
|
use crate::{HubError, HubResult};
|
|
use async_trait::async_trait;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use tokio::fs;
|
|
use tokio::io::AsyncRead;
|
|
|
|
/// Storage configuration for different backends.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum StorageConfig {
|
|
/// Local file system storage
|
|
Local { base_path: PathBuf },
|
|
/// S3-compatible storage
|
|
S3 {
|
|
bucket: String,
|
|
region: String,
|
|
access_key_id: String,
|
|
secret_access_key: String,
|
|
endpoint: Option<String>,
|
|
},
|
|
/// Google Cloud Storage
|
|
Gcs {
|
|
bucket: String,
|
|
project_id: String,
|
|
service_account_key: Option<String>,
|
|
},
|
|
/// Azure Blob Storage
|
|
Azure {
|
|
account_name: String,
|
|
container: String,
|
|
access_key: String,
|
|
},
|
|
/// Redis-based caching layer
|
|
Redis {
|
|
url: String,
|
|
db: u8,
|
|
ttl: Option<u64>,
|
|
},
|
|
}
|
|
|
|
/// Storage backend trait for model registry operations.
|
|
#[async_trait]
|
|
pub trait StorageBackend: Send + Sync {
|
|
/// Store data at the specified path.
|
|
async fn store(&self, path: &str, data: &[u8]) -> HubResult<()>;
|
|
|
|
/// Load data from the specified path.
|
|
async fn load(&self, path: &str) -> HubResult<Vec<u8>>;
|
|
|
|
/// Check if data exists at the specified path.
|
|
async fn exists(&self, path: &str) -> HubResult<bool>;
|
|
|
|
/// Delete data at the specified path.
|
|
async fn delete(&self, path: &str) -> HubResult<()>;
|
|
|
|
/// List all paths with the given prefix.
|
|
async fn list(&self, prefix: &str) -> HubResult<Vec<String>>;
|
|
|
|
/// Get metadata for a stored object.
|
|
async fn metadata(&self, path: &str) -> HubResult<StorageMetadata>;
|
|
|
|
/// Copy data from one path to another.
|
|
async fn copy(&self, from: &str, to: &str) -> HubResult<()>;
|
|
|
|
/// Get the size of stored data.
|
|
async fn size(&self, path: &str) -> HubResult<u64>;
|
|
|
|
/// Stream data for large files.
|
|
async fn stream(&self, path: &str) -> HubResult<Box<dyn AsyncRead + Send + Unpin>>;
|
|
}
|
|
|
|
/// Metadata for stored objects.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StorageMetadata {
|
|
/// Object size in bytes
|
|
pub size: u64,
|
|
/// Last modified timestamp
|
|
pub modified: chrono::DateTime<chrono::Utc>,
|
|
/// Content type/MIME type
|
|
pub content_type: Option<String>,
|
|
/// ETag or content hash
|
|
pub etag: Option<String>,
|
|
/// Custom metadata
|
|
pub metadata: HashMap<String, String>,
|
|
}
|
|
|
|
/// Local file system storage backend.
|
|
pub struct LocalStorageBackend {
|
|
base_path: PathBuf,
|
|
}
|
|
|
|
impl LocalStorageBackend {
|
|
/// Create a new local storage backend.
|
|
pub fn new(base_path: PathBuf) -> Self {
|
|
Self { base_path }
|
|
}
|
|
|
|
/// Get the full path for a storage key.
|
|
fn full_path(&self, path: &str) -> PathBuf {
|
|
self.base_path.join(path)
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl StorageBackend for LocalStorageBackend {
|
|
async fn store(&self, path: &str, data: &[u8]) -> HubResult<()> {
|
|
let full_path = self.full_path(path);
|
|
|
|
// Create parent directories if they don't exist
|
|
if let Some(parent) = full_path.parent() {
|
|
fs::create_dir_all(parent).await?;
|
|
}
|
|
|
|
// Write data to file
|
|
fs::write(&full_path, data).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn load(&self, path: &str) -> HubResult<Vec<u8>> {
|
|
let full_path = self.full_path(path);
|
|
|
|
if !full_path.exists() {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: path.to_string(),
|
|
});
|
|
}
|
|
|
|
let data = fs::read(&full_path).await?;
|
|
Ok(data)
|
|
}
|
|
|
|
async fn exists(&self, path: &str) -> HubResult<bool> {
|
|
let full_path = self.full_path(path);
|
|
Ok(full_path.exists())
|
|
}
|
|
|
|
async fn delete(&self, path: &str) -> HubResult<()> {
|
|
let full_path = self.full_path(path);
|
|
|
|
if full_path.exists() {
|
|
if full_path.is_dir() {
|
|
fs::remove_dir_all(&full_path).await?;
|
|
} else {
|
|
fs::remove_file(&full_path).await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn list(&self, prefix: &str) -> HubResult<Vec<String>> {
|
|
let search_path = self.full_path(prefix);
|
|
let mut results = Vec::new();
|
|
|
|
if search_path.is_dir() {
|
|
let mut entries = fs::read_dir(&search_path).await?;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
let relative_path = path
|
|
.strip_prefix(&self.base_path)
|
|
.map_err(|e| HubError::InvalidPackage {
|
|
reason: e.to_string(),
|
|
})?
|
|
.to_string_lossy()
|
|
.to_string();
|
|
|
|
results.push(relative_path);
|
|
}
|
|
}
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
async fn metadata(&self, path: &str) -> HubResult<StorageMetadata> {
|
|
let full_path = self.full_path(path);
|
|
|
|
if !full_path.exists() {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: path.to_string(),
|
|
});
|
|
}
|
|
|
|
let metadata = fs::metadata(&full_path).await?;
|
|
|
|
Ok(StorageMetadata {
|
|
size: metadata.len(),
|
|
modified: chrono::DateTime::<chrono::Utc>::from(metadata.modified()?),
|
|
content_type: None,
|
|
etag: None,
|
|
metadata: HashMap::new(),
|
|
})
|
|
}
|
|
|
|
async fn copy(&self, from: &str, to: &str) -> HubResult<()> {
|
|
let from_path = self.full_path(from);
|
|
let to_path = self.full_path(to);
|
|
|
|
if !from_path.exists() {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: from.to_string(),
|
|
});
|
|
}
|
|
|
|
// Create parent directories for destination
|
|
if let Some(parent) = to_path.parent() {
|
|
fs::create_dir_all(parent).await?;
|
|
}
|
|
|
|
fs::copy(&from_path, &to_path).await?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn size(&self, path: &str) -> HubResult<u64> {
|
|
let full_path = self.full_path(path);
|
|
|
|
if !full_path.exists() {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: path.to_string(),
|
|
});
|
|
}
|
|
|
|
let metadata = fs::metadata(&full_path).await?;
|
|
Ok(metadata.len())
|
|
}
|
|
|
|
async fn stream(&self, path: &str) -> HubResult<Box<dyn AsyncRead + Send + Unpin>> {
|
|
let full_path = self.full_path(path);
|
|
|
|
if !full_path.exists() {
|
|
return Err(HubError::ModelNotFound {
|
|
model_id: path.to_string(),
|
|
});
|
|
}
|
|
|
|
let file = tokio::fs::File::open(&full_path).await?;
|
|
Ok(Box::new(file))
|
|
}
|
|
}
|
|
|
|
/// S3-compatible storage backend.
|
|
pub struct S3StorageBackend {
|
|
bucket: String,
|
|
client: Option<reqwest::Client>, // Placeholder for actual S3 client
|
|
}
|
|
|
|
impl S3StorageBackend {
|
|
/// Create a new S3 storage backend.
|
|
pub fn new(
|
|
bucket: String,
|
|
_region: String,
|
|
_access_key_id: String,
|
|
_secret_access_key: String,
|
|
_endpoint: Option<String>,
|
|
) -> Self {
|
|
Self {
|
|
bucket,
|
|
client: Some(reqwest::Client::new()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl StorageBackend for S3StorageBackend {
|
|
async fn store(&self, path: &str, data: &[u8]) -> HubResult<()> {
|
|
// Placeholder implementation - would use AWS SDK in production
|
|
tracing::info!(
|
|
"Storing {} bytes to S3 at {}/{}",
|
|
data.len(),
|
|
self.bucket,
|
|
path
|
|
);
|
|
|
|
// In a real implementation, this would upload to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn load(&self, path: &str) -> HubResult<Vec<u8>> {
|
|
// Placeholder implementation
|
|
tracing::info!("Loading from S3 at {}/{}", self.bucket, path);
|
|
|
|
// In a real implementation, this would download from S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
|
|
Ok(vec![]) // Return empty data for now
|
|
}
|
|
|
|
async fn exists(&self, path: &str) -> HubResult<bool> {
|
|
tracing::info!("Checking existence on S3 at {}/{}", self.bucket, path);
|
|
|
|
// Placeholder - would use HEAD request to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
|
|
Ok(true)
|
|
}
|
|
|
|
async fn delete(&self, path: &str) -> HubResult<()> {
|
|
tracing::info!("Deleting from S3 at {}/{}", self.bucket, path);
|
|
|
|
// Placeholder - would use DELETE request to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn list(&self, prefix: &str) -> HubResult<Vec<String>> {
|
|
tracing::info!("Listing S3 objects with prefix: {}", prefix);
|
|
|
|
// Placeholder - would use LIST request to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
|
|
|
Ok(vec![])
|
|
}
|
|
|
|
async fn metadata(&self, path: &str) -> HubResult<StorageMetadata> {
|
|
tracing::info!("Getting S3 metadata for {}/{}", self.bucket, path);
|
|
|
|
// Placeholder - would use HEAD request to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
|
|
|
Ok(StorageMetadata {
|
|
size: 0,
|
|
modified: chrono::Utc::now(),
|
|
content_type: Some("application/octet-stream".to_string()),
|
|
etag: Some("placeholder-etag".to_string()),
|
|
metadata: HashMap::new(),
|
|
})
|
|
}
|
|
|
|
async fn copy(&self, from: &str, to: &str) -> HubResult<()> {
|
|
tracing::info!("Copying S3 object from {} to {}", from, to);
|
|
|
|
// Placeholder - would use COPY request to S3
|
|
tokio::time::sleep(std::time::Duration::from_millis(15)).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn size(&self, path: &str) -> HubResult<u64> {
|
|
let metadata = self.metadata(path).await?;
|
|
Ok(metadata.size)
|
|
}
|
|
|
|
async fn stream(&self, path: &str) -> HubResult<Box<dyn AsyncRead + Send + Unpin>> {
|
|
tracing::info!("Streaming from S3 at {}/{}", self.bucket, path);
|
|
|
|
// Placeholder - would create streaming reader from S3
|
|
let cursor = std::io::Cursor::new(vec![]);
|
|
Ok(Box::new(cursor))
|
|
}
|
|
}
|
|
|
|
/// Caching storage backend that wraps another backend.
|
|
pub struct CachingStorageBackend {
|
|
/// Primary storage backend
|
|
primary: Box<dyn StorageBackend>,
|
|
/// Cache storage backend
|
|
cache: Box<dyn StorageBackend>,
|
|
/// Cache TTL in seconds
|
|
ttl: u64,
|
|
}
|
|
|
|
impl CachingStorageBackend {
|
|
/// Create a new caching storage backend.
|
|
pub fn new(primary: Box<dyn StorageBackend>, cache: Box<dyn StorageBackend>, ttl: u64) -> Self {
|
|
Self {
|
|
primary,
|
|
cache,
|
|
ttl,
|
|
}
|
|
}
|
|
|
|
async fn is_cache_valid(&self, path: &str) -> bool {
|
|
if let Ok(metadata) = self.cache.metadata(path).await {
|
|
let age = chrono::Utc::now()
|
|
.signed_duration_since(metadata.modified)
|
|
.num_seconds() as u64;
|
|
age < self.ttl
|
|
} else {
|
|
false
|
|
}
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl StorageBackend for CachingStorageBackend {
|
|
async fn store(&self, path: &str, data: &[u8]) -> HubResult<()> {
|
|
// Store in primary
|
|
self.primary.store(path, data).await?;
|
|
|
|
// Store in cache (ignore errors)
|
|
let _ = self.cache.store(path, data).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn load(&self, path: &str) -> HubResult<Vec<u8>> {
|
|
// Check cache first
|
|
if self.is_cache_valid(path).await
|
|
&& let Ok(data) = self.cache.load(path).await
|
|
{
|
|
return Ok(data);
|
|
}
|
|
|
|
// Load from primary
|
|
let data = self.primary.load(path).await?;
|
|
|
|
// Update cache (ignore errors)
|
|
let _ = self.cache.store(path, &data).await;
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
async fn exists(&self, path: &str) -> HubResult<bool> {
|
|
// Check cache first
|
|
if self.is_cache_valid(path).await && self.cache.exists(path).await.unwrap_or(false) {
|
|
return Ok(true);
|
|
}
|
|
|
|
// Check primary
|
|
self.primary.exists(path).await
|
|
}
|
|
|
|
async fn delete(&self, path: &str) -> HubResult<()> {
|
|
// Delete from both primary and cache
|
|
self.primary.delete(path).await?;
|
|
let _ = self.cache.delete(path).await; // Ignore cache errors
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn list(&self, prefix: &str) -> HubResult<Vec<String>> {
|
|
// Always use primary for listing
|
|
self.primary.list(prefix).await
|
|
}
|
|
|
|
async fn metadata(&self, path: &str) -> HubResult<StorageMetadata> {
|
|
// Check cache first
|
|
if self.is_cache_valid(path).await
|
|
&& let Ok(metadata) = self.cache.metadata(path).await
|
|
{
|
|
return Ok(metadata);
|
|
}
|
|
|
|
// Get from primary
|
|
self.primary.metadata(path).await
|
|
}
|
|
|
|
async fn copy(&self, from: &str, to: &str) -> HubResult<()> {
|
|
// Copy in primary
|
|
self.primary.copy(from, to).await?;
|
|
|
|
// Invalidate cache for destination
|
|
let _ = self.cache.delete(to).await;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
async fn size(&self, path: &str) -> HubResult<u64> {
|
|
// Check cache first
|
|
if self.is_cache_valid(path).await
|
|
&& let Ok(size) = self.cache.size(path).await
|
|
{
|
|
return Ok(size);
|
|
}
|
|
|
|
// Get from primary
|
|
self.primary.size(path).await
|
|
}
|
|
|
|
async fn stream(&self, path: &str) -> HubResult<Box<dyn AsyncRead + Send + Unpin>> {
|
|
// For streaming, always use primary (caching streams is complex)
|
|
self.primary.stream(path).await
|
|
}
|
|
}
|
|
|
|
/// Create a storage backend from configuration.
|
|
pub fn create_storage_backend(config: StorageConfig) -> HubResult<Box<dyn StorageBackend>> {
|
|
match config {
|
|
StorageConfig::Local { base_path } => Ok(Box::new(LocalStorageBackend::new(base_path))),
|
|
StorageConfig::S3 {
|
|
bucket,
|
|
region,
|
|
access_key_id,
|
|
secret_access_key,
|
|
endpoint,
|
|
} => Ok(Box::new(S3StorageBackend::new(
|
|
bucket,
|
|
region,
|
|
access_key_id,
|
|
secret_access_key,
|
|
endpoint,
|
|
))),
|
|
StorageConfig::Gcs { .. } => {
|
|
// Placeholder for GCS implementation
|
|
Err(HubError::ConfigError {
|
|
details: "GCS storage not yet implemented".to_string(),
|
|
})
|
|
}
|
|
StorageConfig::Azure { .. } => {
|
|
// Placeholder for Azure implementation
|
|
Err(HubError::ConfigError {
|
|
details: "Azure storage not yet implemented".to_string(),
|
|
})
|
|
}
|
|
StorageConfig::Redis { .. } => {
|
|
// Placeholder for Redis implementation
|
|
Err(HubError::ConfigError {
|
|
details: "Redis storage not yet implemented".to_string(),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tempfile::TempDir;
|
|
|
|
#[tokio::test]
|
|
async fn test_local_storage_basic_operations() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = LocalStorageBackend::new(temp_dir.path().to_path_buf());
|
|
|
|
let test_data = b"Hello, world!";
|
|
let test_path = "test/file.txt";
|
|
|
|
// Test store
|
|
storage.store(test_path, test_data).await.unwrap();
|
|
|
|
// Test exists
|
|
assert!(storage.exists(test_path).await.unwrap());
|
|
|
|
// Test load
|
|
let loaded_data = storage.load(test_path).await.unwrap();
|
|
assert_eq!(loaded_data, test_data);
|
|
|
|
// Test size
|
|
let size = storage.size(test_path).await.unwrap();
|
|
assert_eq!(size, test_data.len() as u64);
|
|
|
|
// Test metadata
|
|
let metadata = storage.metadata(test_path).await.unwrap();
|
|
assert_eq!(metadata.size, test_data.len() as u64);
|
|
|
|
// Test copy
|
|
let copy_path = "test/copy.txt";
|
|
storage.copy(test_path, copy_path).await.unwrap();
|
|
assert!(storage.exists(copy_path).await.unwrap());
|
|
|
|
let copied_data = storage.load(copy_path).await.unwrap();
|
|
assert_eq!(copied_data, test_data);
|
|
|
|
// Test list
|
|
let files = storage.list("test").await.unwrap();
|
|
assert_eq!(files.len(), 2); // Original and copy
|
|
|
|
// Test delete
|
|
storage.delete(test_path).await.unwrap();
|
|
assert!(!storage.exists(test_path).await.unwrap());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_storage_nested_paths() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = LocalStorageBackend::new(temp_dir.path().to_path_buf());
|
|
|
|
let test_data = b"Nested file content";
|
|
let nested_path = "models/nlp/bert/config.json";
|
|
|
|
// Store in nested path
|
|
storage.store(nested_path, test_data).await.unwrap();
|
|
|
|
// Verify file exists and content is correct
|
|
assert!(storage.exists(nested_path).await.unwrap());
|
|
let loaded_data = storage.load(nested_path).await.unwrap();
|
|
assert_eq!(loaded_data, test_data);
|
|
|
|
// Check that parent directories were created
|
|
let full_path = storage.full_path(nested_path);
|
|
assert!(full_path.parent().unwrap().exists());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_local_storage_error_handling() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let storage = LocalStorageBackend::new(temp_dir.path().to_path_buf());
|
|
|
|
let nonexistent_path = "nonexistent/file.txt";
|
|
|
|
// Test load nonexistent file
|
|
assert!(storage.load(nonexistent_path).await.is_err());
|
|
|
|
// Test metadata for nonexistent file
|
|
assert!(storage.metadata(nonexistent_path).await.is_err());
|
|
|
|
// Test size for nonexistent file
|
|
assert!(storage.size(nonexistent_path).await.is_err());
|
|
|
|
// Test copy from nonexistent file
|
|
assert!(
|
|
storage
|
|
.copy(nonexistent_path, "destination.txt")
|
|
.await
|
|
.is_err()
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_s3_storage_placeholder() {
|
|
let storage = S3StorageBackend::new(
|
|
"test-bucket".to_string(),
|
|
"us-east-1".to_string(),
|
|
"access-key".to_string(),
|
|
"secret-key".to_string(),
|
|
None,
|
|
);
|
|
|
|
// Test that placeholder implementation doesn't crash
|
|
assert!(storage.exists("test-path").await.is_ok());
|
|
assert!(storage.store("test-path", b"test-data").await.is_ok());
|
|
assert!(storage.load("test-path").await.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_caching_storage_backend() {
|
|
let temp_dir_primary = TempDir::new().unwrap();
|
|
let temp_dir_cache = TempDir::new().unwrap();
|
|
|
|
let primary = Box::new(LocalStorageBackend::new(
|
|
temp_dir_primary.path().to_path_buf(),
|
|
));
|
|
let cache = Box::new(LocalStorageBackend::new(
|
|
temp_dir_cache.path().to_path_buf(),
|
|
));
|
|
|
|
let caching_storage = CachingStorageBackend::new(primary, cache, 60); // 60 second TTL
|
|
|
|
let test_data = b"Cached data";
|
|
let test_path = "cached/file.txt";
|
|
|
|
// Store data
|
|
caching_storage.store(test_path, test_data).await.unwrap();
|
|
|
|
// Load data (should populate cache)
|
|
let loaded_data = caching_storage.load(test_path).await.unwrap();
|
|
assert_eq!(loaded_data, test_data);
|
|
|
|
// Verify data exists in both primary and cache
|
|
assert!(caching_storage.exists(test_path).await.unwrap());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_create_storage_backend() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
|
|
let config = StorageConfig::Local {
|
|
base_path: temp_dir.path().to_path_buf(),
|
|
};
|
|
|
|
let storage = create_storage_backend(config).unwrap();
|
|
|
|
// Test that created backend works
|
|
storage.store("test.txt", b"test").await.unwrap();
|
|
assert!(storage.exists("test.txt").await.unwrap());
|
|
}
|
|
|
|
#[test]
|
|
fn test_storage_config_serialization() {
|
|
let config = StorageConfig::Local {
|
|
base_path: PathBuf::from("/tmp/models"),
|
|
};
|
|
|
|
let json = serde_json::to_string(&config).unwrap();
|
|
let deserialized: StorageConfig = serde_json::from_str(&json).unwrap();
|
|
|
|
match deserialized {
|
|
StorageConfig::Local { base_path } => {
|
|
assert_eq!(base_path, PathBuf::from("/tmp/models"));
|
|
}
|
|
_ => panic!("Expected Local config"),
|
|
}
|
|
}
|
|
}
|