diff --git a/crates/lib/api_clients/cloud_storage/Cargo.toml b/crates/lib/api_clients/cloud_storage/Cargo.toml index 3da8b27c71..86a3de927e 100644 --- a/crates/lib/api_clients/cloud_storage/Cargo.toml +++ b/crates/lib/api_clients/cloud_storage/Cargo.toml @@ -30,6 +30,8 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } anyhow = "1.0.75" log = "0.4.14" env_logger = "0.10.1" +async-trait = "0.1.58" + # The "never-encode-slash" feature seems to have been removed: # https://github.com/durch/rust-s3/pull/254 rust-s3 = { version = "0.33.0" } diff --git a/crates/lib/api_clients/cloud_storage/src/lib.rs b/crates/lib/api_clients/cloud_storage/src/lib.rs index da2cb76964..47b8b2ac76 100644 --- a/crates/lib/api_clients/cloud_storage/src/lib.rs +++ b/crates/lib/api_clients/cloud_storage/src/lib.rs @@ -1,3 +1,4 @@ pub mod bucket_client; pub mod bucket_path_unifier; -pub mod remote_file_manager; \ No newline at end of file +pub mod remote_file_manager; +pub mod remote_file_manager_mock; \ No newline at end of file diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs index e12c746e77..4d071769f0 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs @@ -2,6 +2,27 @@ use errors::AnyhowResult; use crate::bucket_client::BucketClient; use tokio::time::Duration; use log::info; +use crate::remote_file_manager_mock::bucket_orchestration_mock; +use async_trait::async_trait; +#[async_trait] +pub trait BucketOrchestrationDownload { + async fn download_file_to_disk( + &self, + object_path: String, + filesystem_path: String, + is_public:bool + ) -> AnyhowResult<()>; +} + +#[async_trait] +pub trait BucketOrchestrationUpload { + 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, @@ -11,6 +32,35 @@ pub struct BucketOrchestration { private_bucket_name: String, } +impl BucketOrchestrationDownload for BucketOrchestration { + #[async_trait] + 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 + + } +} + +impl BucketOrchestrationUpload for BucketOrchestration { + #[async_trait] + 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 { @@ -48,32 +98,9 @@ impl BucketOrchestration { } } - pub 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 - - } - - pub 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 - - } async fn get_bucket_with_visbility(&self, public:bool) -> AnyhowResult { - + let bucket_timeout = easyenv::get_env_duration_seconds_or_default( "BUCKET_TIMEOUT_SECONDS", Duration::from_secs(60 * 10)); let bucket_client:BucketClient; diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs index 39b9ada6b9..e63d4e2270 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs @@ -30,8 +30,17 @@ impl RemoteCloudBucketDetails { pub fn file_descriptor_from_bucket_details(&self) -> Box { match self.prefix.as_str() { "loRA" => Box::new(weights_descriptor::WeightsLoRADescriptor {}), - "SD15" => Box::new(weights_descriptor::WeightsSD15Descriptor {}), - "SDXL" => Box::new(weights_descriptor::WeightsSDXLDescriptor {}), + "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) } } diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs index bccc6aee13..2f6c98d19e 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs @@ -73,4 +73,100 @@ impl FileDescriptor for WeightsSDXLDescriptor { 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::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 "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 + } } \ No newline at end of file diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/bucket_orchestration_mock.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/bucket_orchestration_mock.rs new file mode 100644 index 0000000000..d7b13b677a --- /dev/null +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/bucket_orchestration_mock.rs @@ -0,0 +1,5 @@ + + + +struct BucketOrchestrationMock { +} \ No newline at end of file diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/mod.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/mod.rs new file mode 100644 index 0000000000..90ad55263e --- /dev/null +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager_mock/mod.rs @@ -0,0 +1 @@ +pub mod bucket_orchestration_mock; \ No newline at end of file