Compare commits

...
Author SHA1 Message Date
ArEnSc 08f59a1989 Update .env 2023-12-21 19:53:07 -05:00
Michael Chung 15c218286f Bucket Orchestration work complete. 2023-12-21 19:51:32 -05:00
Michael Chung 0efe2d8177 Working Tests 2023-12-21 12:38:04 -05:00
Michael Chung ed97cdbb0d Ok builds again going to run unit tests. 2023-12-21 02:25:20 -05:00
Michael Chung 6b2d4ffdce Adding the mocks 2023-12-21 01:58:43 -05:00
Michael Chung 86c80fdf19 Things build again editor works well. 2023-12-21 01:01:35 -05:00
Michael Chung eca8b01f14 Missed 1 file. 2023-12-21 00:42:22 -05:00
Michael Chung 1bddb21e9f Missed 1 file. 2023-12-21 00:41:26 -05:00
Michael Chung 20b950794c Ok it builds need to test it using a mock. 2023-12-21 00:40:17 -05:00
Michael Chung 0d58176f0f Minor bug fixes then testing. 2023-12-20 02:23:45 -05:00
Michael Chung 57f383612f Now testing time! 2023-12-20 02:20:44 -05:00
Michael Chung ab3a312898 Ok so it actually compiles right now. 2023-12-20 00:13:02 -05:00
Michael Chung a2aa71e8a6 builds properly now working on the download aspect. 2023-12-19 23:25:18 -05:00
Michael Chung 3e02d791af Merge branch 'master' into feature/weights_files 2023-12-19 17:44:39 -05:00
Michael Chung 70694ce330 Need to replace the hashing function with a new one and merge master. 2023-12-19 17:43:24 -05:00
Brandon Thomas 5093621995 Merge pull request #227 from storytold/brandon-migration-fixes
migration fixes (tested locally)
2023-12-19 05:53:43 -05:00
Michael Chung 775d867add Ok fixed compilation issues. 2023-12-19 03:05:39 -05:00
Michael Chung da62443522 Fixed some issues using a box 2023-12-19 02:48:41 -05:00
Michael Chung 1ae2df8fa7 working on the file manager for uploads and downloads 2023-12-19 02:40:42 -05:00
Michael Chung 2e6d7d96e4 Merge branch 'master' into feature/weights_files 2023-12-18 21:18:01 -05:00
Michael Chung 39bae02259 Merge branch 'feature/weights_files' of https://github.com/storytold/storyteller-rust into feature/weights_files 2023-12-18 09:26:42 -05:00
Michael Chung 7755a5683b Working on file descriptors though gotta learn some rust things, quickly for boiler plate. 2023-12-18 09:26:12 -05:00
23 changed files with 768 additions and 59 deletions
Generated
+11 -2
View File
@@ -620,9 +620,9 @@ dependencies = [
[[package]]
name = "async-trait"
version = "0.1.74"
version = "0.1.75"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a66537f1bb974b254c98ed142ff995236e81b9d0fe4db0575f46612cb15eb0f9"
checksum = "fdf6721fb0140e4f897002dd086c06f6c27775df19cfe1fccb21181a48fd2c98"
dependencies = [
"proc-macro2",
"quote",
@@ -1124,10 +1124,19 @@ name = "cloud_storage"
version = "0.0.1"
dependencies = [
"anyhow",
"async-trait",
"buckets",
"crockford",
"easyenv",
"env_logger",
"errors",
"filesys",
"hashing",
"log",
"mimetypes",
"my-workspace-hack",
"rust-s3",
"storyteller_root",
"tokio",
]
+1
View File
@@ -18,6 +18,7 @@ path = "src/lib.rs"
crockford = { path = "../../lib/crockford" }
errors = { path = "../../lib/errors" }
# Workspace hack (faster compile times)
my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" }
+1 -1
View File
@@ -10,7 +10,7 @@
// Okay to toggle
#![forbid(unreachable_patterns)]
#![forbid(unused_imports)]
// #![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
@@ -1,3 +0,0 @@
pub trait descriptor {
fn my_method(&self);
}
@@ -1,4 +0,0 @@
pub mod bucket_directory;
pub mod bucket_file_path;
pub mod web_file_manager;
@@ -1,32 +0,0 @@
pub struct FileBucketPath {
pub bucket_clients: BucketClients
}
impl web_file_manager {
pub fn new(my_field: i32) -> Self {
web_file_manager { my_field }
}
pub fn increment(&mut self) {
self.my_field += 1;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new() {
let my_class = web_file_manager::new(10);
assert_eq!(my_class.my_field, 10);
}
#[test]
fn test_increment() {
let mut my_class = web_file_manager::new(10);
my_class.increment();
assert_eq!(my_class.my_field, 11);
}
}
+1 -1
View File
@@ -3,4 +3,4 @@ pub mod media_uploads;
pub mod public_path;
pub mod voice_conversion_results;
pub mod weight_uploads;
pub mod weight_files;
pub mod weight_files;
@@ -5,8 +5,10 @@ use sqlx::{MySql, Pool};
use buckets::public::media_files::bucket_file_path::MediaFileBucketPath;
use enums::by_table::media_files::media_file_type::MediaFileType;
use enums::common::visibility::Visibility;
use errors::AnyhowResult;
use filesys::file_read_bytes::file_read_bytes;
use filesys::file_size::file_size;
use filesys::path_to_string::path_to_string;
@@ -106,7 +108,7 @@ async fn seed_model(
let mut bucket_hash = "NOT_UPLOADED_BY_SEED_TOOL".to_string();
let mut maybe_bucket_prefix = None;
let mut maybe_bucket_extension = None;
if media_file_exists(mysql_pool, media_file_token).await? {
info!("Media file already seeded: {:?}", media_file_token);
return Ok(())
@@ -181,16 +183,16 @@ async fn seed_file_to_bucket(
.extension()
.map(|extension| extension.to_str())
.flatten();
// get the new bucket path ...
let bucket_location = MediaFileBucketPath::generate_new(maybe_bucket_prefix, maybe_bucket_extension);
let bucket_path = path_to_string(bucket_location.to_full_object_pathbuf());
info!("Reading media file: {:?}", media_file_path);
// get meta data
let bytes = file_read_bytes(media_file_path)?;
let mimetype = get_mimetype_for_bytes(&bytes).unwrap_or("application/octet-stream");
info!("Uploading media file to bucket path: {:?}", bucket_path);
let _r = bucket_clients.public.upload_file_with_content_type(
@@ -25,7 +25,7 @@ use tokens::tokens::users::UserToken;
use crate::bucket_clients::BucketClients;
use crate::seeding::users::HANASHI_USER_TOKEN;
use std::result::Result::Ok;
pub async fn seed_weights_files(
mysql_pool: &Pool<MySql>,
maybe_bucket_clients: Option<&BucketClients>
@@ -202,7 +202,7 @@ async fn seed_file_to_bucket(
.map(|extension| extension.to_str())
.flatten();
// we should just have this be file bucket path and then we can use a file descriptor to turn it into a specific type of path
let bucket_location = WeightFileBucketPath::generate_new(
let bucket_location: WeightFileBucketPath = WeightFileBucketPath::generate_new(
maybe_bucket_prefix,
maybe_bucket_extension
);
@@ -16,14 +16,21 @@ path = "src/lib.rs"
# Internal
buckets = { path = "../../../_types/buckets" }
crockford = { path = "../../../lib/crockford" }
errors = { path = "../../../lib/errors" }
easyenv = { path = "../../../lib/easyenv" }
filesys = { path = "../../../lib/files/filesys" }
hashing = { path = "../../../lib/files/hashing" }
mimetypes = { path = "../../../lib/files/mimetypes" }
storyteller_root = { path = "../../../lib/storyteller_root" }
# Workspace hack (faster compile times)
my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" }
# External
anyhow = "1.0.75"
log = "0.4.14"
env_logger = "0.10.1"
async-trait = "0.1.75"
# The "never-encode-slash" feature seems to have been removed:
# https://github.com/durch/rust-s3/pull/254
rust-s3 = { version = "0.33.0" }
@@ -31,3 +38,4 @@ rust-s3 = { version = "0.33.0" }
# Runtime
# NB: We switched to Tokio from `async-std` for broader ecosystem support.
tokio = { version = "1.34.0", features = ["macros"] }
@@ -1,6 +1,6 @@
use std::path::Path;
use std::time::Duration;
use errors::AnyhowResult;
use anyhow::anyhow;
use anyhow::bail;
use log::info;
@@ -10,15 +10,31 @@ use s3::creds::Credentials;
use s3::region::Region;
use tokio::fs::File;
use tokio::io::AsyncReadExt;
use std::error::Error;
#[derive(Clone)]
pub struct BucketClient {
bucket: Bucket,
/// If set, put all files under this root path.
optional_bucket_root: Option<String>,
}
#[derive(Debug)]
pub enum BucketClientError {
ErrorWithCodeAndMessage { code: u16, message: String },
}
impl std::fmt::Display for BucketClientError {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match self {
BucketClientError::ErrorWithCodeAndMessage { code, message } => write!(f, "Error {}: {}", code, message),
}
}
}
impl Error for BucketClientError {}
impl BucketClient {
pub fn create(
@@ -87,6 +103,25 @@ impl BucketClient {
Ok(())
}
pub async fn upload_file_with_content_type_process(&self, object_name: &str, bytes: &[u8], content_type: &str) -> AnyhowResult<()> {
info!("Filename for bucket: {}", object_name);
let object_name = self.get_rooted_object_name(object_name);
info!("Rooted filename for bucket: {}", object_name);
let response = self.bucket.put_object_with_content_type(&object_name, bytes, content_type).await?;
let body_bytes = response.bytes();
let code = response.status_code();
info!("upload code: {}", code);
if code != 200 {
let body = String::from_utf8_lossy(body_bytes);
warn!("upload body: {}", body);
return Err(anyhow!("upload failed: {}", code));
} else {
info!("upload success: {}", code);
Ok(())
}
}
#[deprecated = "Use upload_file instead above it returns an error we can surface and act on. upload_file_with_content_type_process"]
pub async fn upload_file_with_content_type(&self, object_name: &str, bytes: &[u8], content_type: &str) -> anyhow::Result<()> {
info!("Filename for bucket: {}", object_name);
@@ -191,8 +226,7 @@ impl BucketClient {
&self,
object_path: P,
filesystem_path: Q,
) -> anyhow::Result<()> {
// TODO I think we might want to return an error here? but there has to a reason why we are just returning ok?
) -> AnyhowResult<()> {
let object_path_str = object_path.as_ref()
.to_str()
.map(|s| s.to_string())
@@ -208,9 +242,7 @@ impl BucketClient {
404 => bail!("File not found in bucket: {}", &object_path_str),
_ => {},
}
info!("download code: {}", status_code);
Ok(())
}
@@ -1,2 +1,4 @@
pub mod bucket_client;
pub mod bucket_path_unifier;
pub mod remote_file_manager;
pub mod remote_file_manager_mock;
@@ -0,0 +1,131 @@
use log::info;
use tokio::time::Duration;
use async_trait::async_trait;
use errors::AnyhowResult;
use crate::bucket_client::BucketClient;
use crate::remote_file_manager::file_meta_data::FileMetaData;
#[async_trait]
pub trait BucketOrchestrationCore {
async fn download_file_to_disk(
&self,
object_path: String,
filesystem_path: String,
is_public: bool,
) -> AnyhowResult<()>;
async fn upload_file_with_content_type_process(&self, object_name: &str,
bytes: &[u8],
content_type: &str,
is_public: bool) -> AnyhowResult<()>;
}
pub struct BucketOrchestration {
access_key: String,
secret_key: String,
region_name: String,
public_bucket_name: String,
private_bucket_name: String,
}
#[async_trait]
impl BucketOrchestrationCore for BucketOrchestration {
async fn download_file_to_disk(
&self,
object_path: String,
filesystem_path: String,
is_public: bool,
) -> AnyhowResult<()> {
let bucket_client = self.get_bucket_with_visbility(is_public).await?;
bucket_client.download_file_to_disk(object_path, filesystem_path).await
}
async fn upload_file_with_content_type_process(&self, object_name: &str,
bytes: &[u8],
content_type: &str,
is_public: bool) -> AnyhowResult<()> {
let bucket_client = self.get_bucket_with_visbility(is_public).await?;
bucket_client.upload_file_with_content_type_process(
object_name,
bytes,
content_type,
).await
}
}
impl BucketOrchestration {
pub fn new_bucket_client_from_existing_env() -> AnyhowResult<Self> {
let access_key = easyenv::get_env_string_required("ACCESS_KEY")?;
let secret_key = easyenv::get_env_string_required("SECRET_KEY")?;
let region_name = easyenv::get_env_string_required("REGION_NAME")?;
let public_bucket_name = easyenv::get_env_string_required("PUBLIC_BUCKET_NAME")?;
let private_bucket_name = easyenv::get_env_string_required("PRIVATE_BUCKET_NAME")?;
let bucket_orchestration = BucketOrchestration::new(
access_key,
secret_key,
region_name,
public_bucket_name,
private_bucket_name,
);
Ok(bucket_orchestration)
}
pub fn new(
access_key: String,
secret_key: String,
region_name: String,
public_bucket_name: String,
private_bucket_name: String) -> Self {
Self {
access_key: access_key,
secret_key: secret_key,
region_name: region_name,
public_bucket_name: public_bucket_name,
private_bucket_name: private_bucket_name,
}
}
async fn get_bucket_with_visbility(&self, public: bool) -> AnyhowResult<BucketClient> {
let bucket_timeout = easyenv::get_env_duration_seconds_or_default(
"BUCKET_TIMEOUT_SECONDS", Duration::from_secs(60 * 10));
let bucket_client: BucketClient;
if public {
// use public bucket client
info!("Configuring public GCS bucket...");
bucket_client = BucketClient::create(
&self.access_key,
&self.secret_key,
&self.region_name,
&self.public_bucket_name,
None,
Some(bucket_timeout),
)?;
} else {
info!("Configuring private GCS bucket...");
bucket_client = BucketClient::create(
&self.access_key,
&self.secret_key,
&self.region_name,
&self.private_bucket_name,
None,
Some(bucket_timeout),
)?;
}
Ok(bucket_client)
}
}
#[cfg(test)]
mod tests {
#[test]
pub fn test() {}
}
@@ -0,0 +1,35 @@
// Describes the file and how it will be saved to GCBucket
// where the file should be stored based off the descriptor
// DEFAULT IMPLEMENTATION
const REMOTE_FILE_DIRECTORY: &str = "/implement_google_cloud_storage_bucket_name";
// File Descriptor Steers the Bucket Directory
pub trait FileDescriptor {
// By default a file belongs in the public bucket will help us figureout which bucket to use.
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file
// e.g requires! a period .safetensors .bin .jpg
fn get_suffix(&self)->String {
return "implement".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, loRA, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "implement".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
#[cfg(test)]
mod tests {
#[test]
pub fn test() {}
}
@@ -0,0 +1,83 @@
use std::path::PathBuf;
use crockford::crockford_entropy_lower;
use buckets::util::hashed_directory_path_long_string;
use crate::remote_file_manager::file_descriptor::FileDescriptor;
use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
pub struct FileBucketDirectory {
file_object_hash: String,
remote_cloud_base_directory: String,
full_remote_cloud_file_path: String,
file_name: String,
file_descriptor: Box<dyn FileDescriptor>
}
impl FileBucketDirectory {
pub fn generate_new(file_descriptor: Box<dyn FileDescriptor>) -> Self {
Self::new_from_file_descriptor(file_descriptor)
}
// TODO refactor this out of here
pub fn from_existing_bucket_details(bucket_details:RemoteCloudBucketDetails) -> Self {
let file_descriptor = bucket_details.file_descriptor_from_bucket_details();
let entropy = bucket_details.get_object_hash();
let middle = hashed_directory_path_long_string::hashed_directory_path_long_string(entropy.as_ref());
// gets you cloud bucket path e.g weights/a/b/c/d/clould_path_entropy
let remote_cloud_base_directory = format!("{}/{}{}", file_descriptor.remote_directory_path(), middle, entropy);
// gets you name of the file with suffix and prefix and entropy in the centre
let file_name = format!("{}_{}.{}", bucket_details.get_prefix(), entropy, bucket_details.get_suffix());
// note: no longer optional because it's easy to know what it would be in the db explicit is better than implcit.
// This is the full path you download from
let full_remote_cloud_file_path = format!("{}/{}", remote_cloud_base_directory , file_name);
Self {
file_object_hash: entropy.to_string(),
remote_cloud_base_directory: remote_cloud_base_directory,
full_remote_cloud_file_path: full_remote_cloud_file_path,
file_name:file_name,
file_descriptor:file_descriptor
}
}
pub fn new_from_file_descriptor(file_descriptor: Box<dyn FileDescriptor>) -> Self {
let entropy = crockford_entropy_lower(32);
// gets you wiki /a/b/c/d folder structure
let middle = hashed_directory_path_long_string::hashed_directory_path_long_string(entropy.as_ref());
// gets you cloud bucket path e.g weights/a/b/c/d/clould_path_entropy
let remote_cloud_base_directory = format!("{}/{}{}", file_descriptor.remote_directory_path(), middle, entropy);
// gets you name of the file with suffix and prefix and entropy in the centre
let file_name = format!("{}_{}.{}", file_descriptor.get_prefix(), entropy, file_descriptor.get_suffix());
// note: no longer optional because it's easy to know what it would be in the db explicit is better than implcit.
// This is the full path you upload to.
let full_remote_cloud_file_path = format!("{}/{}", remote_cloud_base_directory , file_name);
Self {
file_object_hash: entropy,
remote_cloud_base_directory: remote_cloud_base_directory,
full_remote_cloud_file_path: full_remote_cloud_file_path,
file_name:file_name,
file_descriptor:file_descriptor
}
}
pub fn get_file_object_hash(&self) -> &str {
&self.file_object_hash
}
pub fn get_remote_cloud_base_directory(&self) -> &str {
&self.remote_cloud_base_directory
}
pub fn get_full_remote_cloud_file_path(&self) -> &str {
&self.full_remote_cloud_file_path
}
pub fn to_full_remote_cloud_file_path_pathbuf(&self) -> PathBuf {
PathBuf::from(&self.full_remote_cloud_file_path)
}
pub fn get_file_name(&self) -> &str {
&self.file_name
}
}
@@ -0,0 +1,6 @@
#[derive(Debug, Clone)]
pub struct FileMetaData {
pub file_size_bytes: u64,
pub sha256_checksum: String,
pub mimetype: String
}
@@ -0,0 +1,12 @@
pub mod remote_cloud_file_manager;
pub mod remote_cloud_bucket_details;
pub mod bucket_orchestration;
pub mod file_directory;
pub mod file_meta_data;
pub mod file_descriptor;
pub mod media_descriptor;
pub mod weights_descriptor;
@@ -0,0 +1,58 @@
use super::file_descriptor::FileDescriptor;
use crate::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsSDXLDescriptor, self};
pub struct RemoteCloudBucketDetails {
pub object_hash: String,
pub prefix: String,
pub suffix: String,
}
impl RemoteCloudBucketDetails {
pub fn new(object_hash: String, prefix: String, suffix: String) -> Self {
Self {
object_hash: object_hash,
prefix: prefix,
suffix: suffix
}
}
pub fn get_object_hash(&self) -> &str {
&self.object_hash
}
pub fn get_prefix(&self) -> &str {
&self.prefix
}
pub fn get_suffix(&self) -> &str {
&self.suffix
}
pub fn file_descriptor_from_bucket_details(&self) -> Box<dyn FileDescriptor> {
match self.prefix.as_str() {
"loRA" => Box::new(weights_descriptor::WeightsLoRADescriptor {}),
"sd15" => Box::new(weights_descriptor::WeightsSD15Descriptor {}),
"sdxl" => Box::new(weights_descriptor::WeightsSDXLDescriptor {}),
"valle_prompt" => Box::new(weights_descriptor::WeightsVallePromptDescriptor {}),
"rvc" => {
match self.suffix.as_str() {
"safetensors" => Box::new(weights_descriptor::WeightsRVCDescriptor {}),
"index" => Box::new(weights_descriptor::WeightsRVCIndexDescriptor {}),
_ => panic!("Unknown suffix: {}",self.suffix)
}
},
"svc" => Box::new(weights_descriptor::WeightsSVCDescriptor {}),
_ => panic!("Unknown prefix: {}", self.prefix)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
pub fn test() {
}
}
@@ -0,0 +1,196 @@
use crate::remote_file_manager::file_directory::FileBucketDirectory;
use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails;
use super::bucket_orchestration;
use super::file_descriptor::{FileDescriptor};
use super::file_meta_data::FileMetaData;
use errors:: AnyhowResult;
use filesys::file_read_bytes::file_read_bytes;
use filesys::file_size::file_size;
use hashing::sha256::sha256_hash_file::sha256_hash_file;
use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes;
use crate::remote_file_manager::bucket_orchestration::{BucketOrchestrationCore};
struct RemoteCloudFileClient {
bucket_orchestration_client: Box<dyn BucketOrchestrationCore>
}
impl RemoteCloudFileClient {
pub fn new(bucket_orchestration_client: Box<dyn BucketOrchestrationCore>) -> Self {
Self {
bucket_orchestration_client: bucket_orchestration_client
}
}
pub async fn download_file(&self, remote_cloud_bucket_details:RemoteCloudBucketDetails, to_system_file_path:String) -> AnyhowResult<()> {
let file_descriptor = remote_cloud_bucket_details.file_descriptor_from_bucket_details();
let file_bucket_directory = FileBucketDirectory::from_existing_bucket_details(remote_cloud_bucket_details);
let full_remote_cloud_file_path = file_bucket_directory.get_full_remote_cloud_file_path().to_string();
let is_public = file_descriptor.is_public().clone();
self.bucket_orchestration_client.download_file_to_disk(full_remote_cloud_file_path, to_system_file_path, is_public).await?;
Ok(())
}
// return error or success with meta data.
pub async fn upload_file(&self, file_descriptor:Box<dyn FileDescriptor>, from_system_file_path:&str) -> AnyhowResult<FileMetaData> {
// get file meta data
println!("Reading media file: {:?}", from_system_file_path);
// get meta data
let bytes = file_read_bytes(from_system_file_path)?;
let result = Self::get_file_meta_data(from_system_file_path)?;
let is_public = file_descriptor.is_public();
let directory = FileBucketDirectory::generate_new(
file_descriptor
);
println!("Uploading media file to bucket path: {:?}",directory.get_full_remote_cloud_file_path());
self.bucket_orchestration_client.upload_file_with_content_type_process(
&directory.get_remote_cloud_base_directory(),
bytes.as_ref(),
result.mimetype.as_ref(),
is_public
).await?;
Ok(result)
}
// Retrieve the metadata from the file
fn get_file_meta_data(system_file_path:&str) -> AnyhowResult<FileMetaData> {
let file_size_bytes = file_size(system_file_path.clone())?;
let sha256_checksum = sha256_hash_file(system_file_path.clone())?;
let bytes = file_read_bytes(system_file_path)?;
let mimetype: &str = get_mimetype_for_bytes(&bytes).unwrap_or("application/octet-stream");
Ok(FileMetaData {
file_size_bytes: file_size_bytes,
sha256_checksum: sha256_checksum,
mimetype: mimetype.to_string()
})
}
}
#[cfg(test)]
mod tests {
use async_trait::async_trait;
use env_logger;
use errors::AnyhowResult;
use crate::remote_file_manager::bucket_orchestration::BucketOrchestrationCore;
use crate::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsSDXLDescriptor};
struct BucketOrchestrationMock {}
#[async_trait]
impl BucketOrchestrationCore for BucketOrchestrationMock {
async fn download_file_to_disk(
&self,
object_path: String,
filesystem_path: String,
is_public: bool,
) -> AnyhowResult<()>{
println!("Download File to Disk");
println!("{}",object_path);
println!("{}",filesystem_path);
println!("{}",is_public);
println!("Download Done Downloading");
assert_eq!(String::from("/weights/1/2/123/loRA_123.safetensors"),object_path);
assert_eq!(String::from("./file_path_here"),filesystem_path);
assert_eq!(is_public,true);
Ok(())
}
async fn upload_file_with_content_type_process(&self, object_name: &str,
bytes: &[u8],
content_type: &str,
is_public: bool) -> AnyhowResult<()> {
println!("Upload File to Disk");
println!("{}",object_name);
assert_eq!(object_name,String::from("/weights/2/y/q/m/2/2yqm2f1bamh88seyd690h9v24apgezhr/loRA_2yqm2f1bamh88seyd690h9v24apgezhr.safetensors"));
println!("ContentType:{}",content_type);
println!("{}",is_public);
assert_eq!(is_public,true);
println!("Is Done Uploading");
Ok(())
}
}
#[tokio::test]
async fn remote_file_manager_descriptor_test() {
use super::*;
let remote_cloud_bucket_details = RemoteCloudBucketDetails::new("object_hash".to_string(), "loRA".to_string(), "safetensors".to_string());
let file_descriptor = remote_cloud_bucket_details.file_descriptor_from_bucket_details();
assert_eq!(file_descriptor.get_prefix(), "loRA");
assert_eq!(file_descriptor.get_suffix(), "safetensors");
assert_eq!(file_descriptor.is_public(), true);
}
#[tokio::test]
async fn remote_file_manager_download_existing_file() {
use super::*;
env_logger::init();
let file_path:&str= "./file_path_here";
let remote_cloud_file_manager = RemoteCloudFileClient {
bucket_orchestration_client: Box::new(BucketOrchestrationMock {})
};
println!("begin upload from file_path: {:?}", file_path);
let details = RemoteCloudBucketDetails {
object_hash: String::from("123"),
prefix: String::from("loRA"),
suffix: String::from("safetensors")
};
let result = remote_cloud_file_manager.download_file(details,file_path.to_string()).await;
match result {
Ok(file_meta_data) => {
println!("file_meta_data: {:?}", file_meta_data);
},
Err(e) => {
println!("error: {:?}", e);
}
}
}
#[tokio::test]
async fn remote_file_manager_upload_existing_file() {
use super::*;
env_logger::init();
// Replace with a test file.
let weight_path:&str= "/home/tensor/code/storyteller/storyteller-rust/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs";
let remote_cloud_file_manager = RemoteCloudFileClient {
bucket_orchestration_client: Box::new(BucketOrchestrationMock {})
};
println!("begin upload from weight_path: {:?}", weight_path);
let result = remote_cloud_file_manager.upload_file(Box::new(WeightsLoRADescriptor {}),weight_path).await;
match result {
Ok(file_meta_data) => {
println!("file_meta_data: {:?}", file_meta_data);
},
Err(e) => {
println!("error: {:?}", e);
}
}
}
}
@@ -0,0 +1,172 @@
use super::file_descriptor::FileDescriptor;
pub struct WeightsLoRADescriptor {}
const REMOTE_FILE_DIRECTORY: &str = "/weights";
impl FileDescriptor for WeightsLoRADescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "safetensors".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, loRA, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "loRA".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsSD15Descriptor {}
impl FileDescriptor for WeightsSD15Descriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "safetensors".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "sd15".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsSDXLDescriptor {}
impl FileDescriptor for WeightsSDXLDescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "safetensors".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "sdxl".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsSVCDescriptor {}
impl FileDescriptor for crate::remote_file_manager::weights_descriptor::WeightsSVCDescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "safetensors".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "svc".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsRVCDescriptor {}
impl FileDescriptor for crate::remote_file_manager::weights_descriptor::WeightsRVCDescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "pth".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "rvc".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsRVCIndexDescriptor {}
impl FileDescriptor for crate::remote_file_manager::weights_descriptor::WeightsRVCIndexDescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "index".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "rvc".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
pub struct WeightsVallePromptDescriptor {}
impl FileDescriptor for crate::remote_file_manager::weights_descriptor::WeightsVallePromptDescriptor {
fn remote_directory_path(&self) -> &str {
return REMOTE_FILE_DIRECTORY;
}
// this will be the type of file peroid is handled by the file formatter
// e.g safetensors bin jpg
fn get_suffix(&self)->String {
return "safetensors".to_string();
}
// This will be the prefix of the media type or the weights type.
// name of the weights or the name of the media type
// vall-e_prompt, SD15, sd15, sdxl when implmenting add to the end
fn get_prefix(&self)->String {
return "valle_prompt".to_string();
}
// This will be ensure that the right bucket is picked
fn is_public(&self) -> bool {
true
}
}
@@ -0,0 +1 @@
pub mod bucket_orchestration_mock;