From 3d9a0cfcc1922b90628691b366c6fccc5845bc36 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 23 Jul 2023 21:32:39 -0400 Subject: [PATCH] scoped job execution --- .../list_available_generic_inference_jobs.rs | 89 +++++++++++++---- .../src/job/job_loop/main_loop.rs | 3 +- .../job/inference_job/src/job_dependencies.rs | 5 + crates/service/job/inference_job/src/main.rs | 4 +- .../service/job/inference_job/src/util/mod.rs | 2 +- .../src/util/scope_model_types.rs | 44 --------- .../src/util/scoped_execution.rs | 97 +++++++++++++++++++ 7 files changed, 177 insertions(+), 67 deletions(-) delete mode 100644 crates/service/job/inference_job/src/util/scope_model_types.rs create mode 100644 crates/service/job/inference_job/src/util/scoped_execution.rs diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs b/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs index c9ce668d99..3feaa53738 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs @@ -11,6 +11,7 @@ use sqlx::{MySql, MySqlPool}; use std::collections::BTreeSet; use std::future::Future; use std::path::Path; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use tokens::jobs::inference::InferenceJobToken; use crate::payloads::generic_inference_args::GenericInferenceArgs; @@ -76,7 +77,8 @@ pub struct ListAvailableGenericInferenceJobArgs<'a> { pub num_records: u32, pub is_debug_worker: bool, pub sort_by_priority: bool, - pub maybe_scope_by_job_category: Option>, + pub maybe_scope_by_model_type: Option<&'a BTreeSet>, + pub maybe_scope_by_job_category: Option<&'a BTreeSet>, pub mysql_pool: &'a MySqlPool, } @@ -85,17 +87,18 @@ pub async fn list_available_generic_inference_jobs( ) -> AnyhowResult> { - let ALL_TYPES = InferenceCategory::all_variants(); - - let inference_types = args.maybe_scope_by_job_category - .as_ref() + let model_types = args.maybe_scope_by_model_type .map(|types| types.clone()) - .unwrap_or(ALL_TYPES); + .unwrap_or(InferenceModelType::all_variants()); // NB: All model types + + let inference_categories = args.maybe_scope_by_job_category + .map(|types| types.clone()) + .unwrap_or(InferenceCategory::all_variants()); // NB: All categories let query = if args.sort_by_priority { - list_sorted_by_priority(args, inference_types).await + list_sorted_by_priority(args, model_types, inference_categories).await } else { - list_sorted_by_id(args, inference_types).await + list_sorted_by_id(args, model_types, inference_categories).await }; let job_records = query?; @@ -146,7 +149,7 @@ pub async fn list_available_generic_inference_jobs( Ok(job_records) } -async fn list_sorted_by_id(args: ListAvailableGenericInferenceJobArgs<'_>, inference_categories: BTreeSet) -> Result, sqlx::Error> { +async fn list_sorted_by_id(args: ListAvailableGenericInferenceJobArgs<'_>, model_types: BTreeSet, inference_categories: BTreeSet) -> Result, sqlx::Error> { // NB: Can't be type checked because of WHERE IN clause with dynamic contents // Also had to remove the following typing: @@ -199,11 +202,18 @@ SELECT FROM generic_inference_jobs"#.to_string(); query.push_str(&format!(r#" -WHERE + WHERE + ( + maybe_model_type IN ({}) + ) + "#, model_type_predicate(&model_types))); + + query.push_str(&format!(r#" + and ( inference_category IN ({}) ) - "#, to_where_in_predicate(&inference_categories))); + "#, inference_category_predicate(&inference_categories))); query.push_str(r#" AND @@ -232,7 +242,7 @@ WHERE .await } -async fn list_sorted_by_priority(args: ListAvailableGenericInferenceJobArgs<'_>, inference_categories: BTreeSet) -> Result, sqlx::Error> { +async fn list_sorted_by_priority(args: ListAvailableGenericInferenceJobArgs<'_>, model_types: BTreeSet, inference_categories: BTreeSet) -> Result, sqlx::Error> { // NB: Can't be type checked because of WHERE IN clause with dynamic contents // Also had to remove the following typing: @@ -284,13 +294,19 @@ SELECT FROM generic_inference_jobs"#.to_string(); + query.push_str(&format!(r#" + WHERE + ( + maybe_model_type IN ({}) + ) + "#, model_type_predicate(&model_types))); query.push_str(&format!(r#" -WHERE + and ( inference_category IN ({}) ) - "#, to_where_in_predicate(&inference_categories))); + "#, inference_category_predicate(&inference_categories))); query.push_str(r#" AND @@ -374,7 +390,18 @@ struct AvailableInferenceJobRawInternal { /// Return a comma-separated predicate, since SQLx does not yet support WHERE IN(?) for Vec, etc. /// Issue: https://github.com/launchbadge/sqlx/issues/875 -fn to_where_in_predicate(categories: &BTreeSet) -> String { +fn model_type_predicate(types: &BTreeSet) -> String { + let mut vec = types.iter() + .map(|ty| ty.to_str()) + .map(|ty| format!("\"{}\"", ty)) + .collect::>(); + vec.sort(); // NB: For the benefit of tests. + vec.join(", ") +} + +/// Return a comma-separated predicate, since SQLx does not yet support WHERE IN(?) for Vec, etc. +/// Issue: https://github.com/launchbadge/sqlx/issues/875 +fn inference_category_predicate(categories: &BTreeSet) -> String { let mut vec = categories.iter() .map(|ty| ty.to_str()) .map(|ty| format!("\"{}\"", ty)) @@ -385,27 +412,49 @@ fn to_where_in_predicate(categories: &BTreeSet) -> String { #[cfg(test)] mod tests { - use crate::queries::generic_inference::job::list_available_generic_inference_jobs::to_where_in_predicate; + use crate::queries::generic_inference::job::list_available_generic_inference_jobs::{inference_category_predicate, model_type_predicate}; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; + use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use std::collections::BTreeSet; #[test] - fn test_where_in_clause() { + fn test_model_type_predicate() { // None let types = BTreeSet::from([]); - assert_eq!(to_where_in_predicate(&types), "".to_string()); + assert_eq!(model_type_predicate(&types), "".to_string()); + + // One + let types = BTreeSet::from([ + InferenceModelType::RvcV2, + ]); + + assert_eq!(model_type_predicate(&types), "\"rvc_v2\"".to_string()); + + // Multiple + let types = BTreeSet::from([ + InferenceModelType::RvcV2, + InferenceModelType::SoVitsSvc, + ]); + assert_eq!(model_type_predicate(&types), "\"rvc_v2\", \"so_vits_svc\"".to_string()); + } + + #[test] + fn test_inference_category_predicate() { + // None + let types = BTreeSet::from([]); + assert_eq!(inference_category_predicate(&types), "".to_string()); // Some let types = BTreeSet::from([ InferenceCategory::VoiceConversion, ]); - assert_eq!(to_where_in_predicate(&types), "\"voice_conversion\"".to_string()); + assert_eq!(inference_category_predicate(&types), "\"voice_conversion\"".to_string()); // All let types = BTreeSet::from([ InferenceCategory::TextToSpeech, InferenceCategory::VoiceConversion, ]); - assert_eq!(to_where_in_predicate(&types), "\"text_to_speech\", \"voice_conversion\"".to_string()); + assert_eq!(inference_category_predicate(&types), "\"text_to_speech\", \"voice_conversion\"".to_string()); } } diff --git a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs index 0d6e0325e1..00365bfaa9 100644 --- a/crates/service/job/inference_job/src/job/job_loop/main_loop.rs +++ b/crates/service/job/inference_job/src/job/job_loop/main_loop.rs @@ -44,7 +44,8 @@ pub async fn main_loop(job_dependencies: JobDependencies) { num_records: job_dependencies.job_batch_size, is_debug_worker: false, // TODO sort_by_priority, - maybe_scope_by_job_category: None, // TODO + maybe_scope_by_model_type: job_dependencies.scoped_execution.get_scoped_model_types(), + maybe_scope_by_job_category: None, mysql_pool: &job_dependencies.mysql_pool, }).await; diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index bb9e8d8fb8..8c7e4b1f61 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/crates/service/job/inference_job/src/job_dependencies.rs @@ -7,6 +7,7 @@ use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceComma use crate::job::job_types::vc::rvc_v2::pretrained_hubert_model::PretrainedHubertModel; use crate::job::job_types::vc::rvc_v2::rvc_v2_inference_command::RvcV2InferenceCommand; use crate::job::job_types::vc::so_vits_svc::so_vits_svc_inference_command::SoVitsSvcInferenceCommand; +use crate::util::scoped_execution::ScopedExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporterBuilder; use jobs_common::job_stats::JobStats; @@ -24,6 +25,10 @@ use sqlx::MySqlPool; use std::path::PathBuf; pub struct JobDependencies { + /// The job should only run on these types of models. + /// This is provided at job start. + pub scoped_execution: ScopedExecution, + /// Filesystem info and utils pub fs: FileSystemDetails, diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 6063951b9b..4f810c7649 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -36,9 +36,11 @@ use crate::http_server::run_http_server::launch_http_server; use crate::job::job_loop::main_loop::main_loop; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::Tacotron2InferenceCommand; use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand; +use crate::job::job_types::vc::rvc_v2::pretrained_hubert_model::PretrainedHubertModel; use crate::job::job_types::vc::rvc_v2::rvc_v2_inference_command::RvcV2InferenceCommand; use crate::job::job_types::vc::so_vits_svc::so_vits_svc_inference_command::SoVitsSvcInferenceCommand; use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, PretrainedModels, RvcV2Details, SoVitsSvcDetails, Tacotron2VocodesDetails, VitsDetails}; +use crate::util::scoped_execution::ScopedExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; use filesys::create_dir_all_if_missing::create_dir_all_if_missing; use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporterBuilder; @@ -58,7 +60,6 @@ use sqlx::mysql::MySqlPoolOptions; use std::path::PathBuf; use std::time::Duration; use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions}; -use crate::job::job_types::vc::rvc_v2::pretrained_hubert_model::PretrainedHubertModel; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -241,6 +242,7 @@ async fn main() -> AnyhowResult<()> { }; let job_dependencies = JobDependencies { + scoped_execution: ScopedExecution::new_from_env()?, fs: FileSystemDetails { temp_directory_downloads: temp_directory_downloads.clone(), temp_directory_work: temp_directory_work.clone(), diff --git a/crates/service/job/inference_job/src/util/mod.rs b/crates/service/job/inference_job/src/util/mod.rs index 0b59161195..63f98be1ca 100644 --- a/crates/service/job/inference_job/src/util/mod.rs +++ b/crates/service/job/inference_job/src/util/mod.rs @@ -1,4 +1,4 @@ pub mod audiowmark; -pub mod scope_model_types; +pub mod scoped_execution; pub mod maybe_download_file_from_bucket; pub mod scoped_temp_dir_creator; diff --git a/crates/service/job/inference_job/src/util/scope_model_types.rs b/crates/service/job/inference_job/src/util/scope_model_types.rs deleted file mode 100644 index b28e3990d7..0000000000 --- a/crates/service/job/inference_job/src/util/scope_model_types.rs +++ /dev/null @@ -1,44 +0,0 @@ -use anyhow::anyhow; -use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; -use errors::AnyhowResult; -use std::collections::BTreeSet; - -pub fn scope_model_types() -> AnyhowResult> { - let comma_separated_types = match easyenv::get_env_string_optional("SCOPE_MODEL_TYPES") { - None => return Ok(InferenceModelType::all_variants()), - Some(val) => val, - }; - - parse_model_types(&comma_separated_types) -} - -pub fn parse_model_types(comma_separated_types: &str) -> AnyhowResult> { - let scoped_types = comma_separated_types.trim() - .split(",") - .map(|val| val.to_lowercase()) - .collect::>(); - - let mut model_types = BTreeSet::new(); - - for t in scoped_types.into_iter() { - let model_type = InferenceModelType::from_str(&t) - .map_err(|err| anyhow!("Invalid model type: {:?}", t))?; - - model_types.insert(model_type); - } - - Ok(model_types) -} - -#[cfg(test)] -mod tests { - use crate::util::scope_model_types::parse_model_types; - use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; - use std::collections::BTreeSet; - - #[test] - fn test_parse() { - assert_eq!(parse_model_types("rvc_v2,so_vits_svc").unwrap(), - BTreeSet::from([InferenceModelType::RvcV2, InferenceModelType::SoVitsSvc])); - } -} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/util/scoped_execution.rs b/crates/service/job/inference_job/src/util/scoped_execution.rs new file mode 100644 index 0000000000..e5d3163cc6 --- /dev/null +++ b/crates/service/job/inference_job/src/util/scoped_execution.rs @@ -0,0 +1,97 @@ +use anyhow::anyhow; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use errors::AnyhowResult; +use std::collections::BTreeSet; +use log::info; + +/// Execution can be scoped down to run on only certain model types or inference categories. +#[derive(Clone)] +pub struct ScopedExecution { + /// If set, only these types of model type will be inferred. + /// None means no scoping, so all types can execute. + scoped_types: Option>, +} + +impl ScopedExecution { + + pub fn new_from_set(scoped_types: BTreeSet) -> Self { + Self { + scoped_types: Some(scoped_types), + } + } + + pub fn new_from_env() -> AnyhowResult { + let scoped_types = + match easyenv::get_env_string_optional("SCOPED_EXECUTION_MODEL_TYPES") { + Some(comma_separated_types) => Some(parse_model_types(&comma_separated_types)?), + None => None, + }; + + if let Some(types) = scoped_types.as_ref() { + info!("Scoping execution to model types: {:?}", types); + } + + Ok(Self { + scoped_types, + }) + } + + pub fn can_run_job(&self, job_model_type: InferenceModelType) -> bool { + match self.scoped_types { + None => true, + Some(ref types) => types.contains(&job_model_type), + } + } + + pub fn get_scoped_model_types(&self) -> Option<&BTreeSet> { + self.scoped_types.as_ref() + } +} + +pub fn parse_model_types(comma_separated_types: &str) -> AnyhowResult> { + let scoped_types = comma_separated_types.trim() + .split(",") + .map(|val| val.to_lowercase()) + .collect::>(); + + let mut model_types = BTreeSet::new(); + + for t in scoped_types.into_iter() { + let model_type = InferenceModelType::from_str(&t) + .map_err(|err| anyhow!( + "Invalid model type: {:?}; should include only items from: {:?}", + t, + InferenceModelType::all_variants()))?; + + model_types.insert(model_type); + } + + Ok(model_types) +} + +#[cfg(test)] +mod tests { + use crate::util::scoped_execution::{parse_model_types, ScopedExecution}; + use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; + use std::collections::BTreeSet; + + #[test] + fn test_parse() { + assert_eq!(parse_model_types("rvc_v2,so_vits_svc").unwrap(), + BTreeSet::from([InferenceModelType::RvcV2, InferenceModelType::SoVitsSvc])); + } + + #[test] + fn test_can_execute() { + let scoping = ScopedExecution::new_from_set(BTreeSet::from([ + InferenceModelType::RvcV2, + InferenceModelType::Vits + ])); + + assert_eq!(true, scoping.can_run_job(InferenceModelType::RvcV2)); + assert_eq!(true, scoping.can_run_job(InferenceModelType::Vits)); + + assert_eq!(false, scoping.can_run_job(InferenceModelType::SoVitsSvc)); + assert_eq!(false, scoping.can_run_job(InferenceModelType::Tacotron2)); + } +}