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 f16a2448b0..a03daed518 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 @@ -48,7 +48,8 @@ pub async fn main_loop(job_dependencies: JobDependencies) { sort_by_priority = false; } - let maybe_scoped_model_types = job_dependencies.job.system.scoped_execution.get_scoped_model_types(); + let maybe_scoped_job_types = job_dependencies.job.system.scoped_job_type_execution.get_scoped_job_types(); + let maybe_scoped_model_types = job_dependencies.job.system.scoped_model_type_execution.get_scoped_model_types(); let batch_query_start_time = Instant::now(); @@ -56,7 +57,7 @@ pub async fn main_loop(job_dependencies: JobDependencies) { num_records: job_dependencies.job.system.job_batch_size, is_debug_worker: false, // TODO sort_by_priority, - maybe_scope_by_job_type: None, + maybe_scope_by_job_type: maybe_scoped_job_types, maybe_scope_by_model_type: maybe_scoped_model_types, maybe_scope_by_job_category: None, mysql_pool: &job_dependencies.db.mysql_pool, diff --git a/crates/service/job/inference_job/src/job_dependencies.rs b/crates/service/job/inference_job/src/job_dependencies.rs index 30d31e2058..7b3f1ddd06 100644 --- a/crates/service/job/inference_job/src/job_dependencies.rs +++ b/crates/service/job/inference_job/src/job_dependencies.rs @@ -26,7 +26,8 @@ use mysql_queries::mediators::firehose_publisher::FirehosePublisher; use crate::job_specific_dependencies::JobSpecificDependencies; use crate::util::instrumentation::JobInstruments; use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory; -use crate::util::scoped_execution::ScopedExecution; +use crate::util::scoped_job_type_execution::ScopedJobTypeExecution; +use crate::util::scoped_model_type_execution::ScopedModelTypeExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct JobDependencies { @@ -60,9 +61,13 @@ pub struct JobSystemDependencies { } pub struct JobSystemControls { - /// The job should only run on these types of models. + /// The application should only run jobs associated with these types of models. /// This is provided at job start. - pub scoped_execution: ScopedExecution, + pub scoped_model_type_execution: ScopedModelTypeExecution, + + /// The application should only run on these types of jobs. + /// This is provided at job start. + pub scoped_job_type_execution: ScopedJobTypeExecution, // Allow for models not to exist on the filesystem. All jobs will execute when first tried // regardless of whether their models were previously downloaded. diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index c7ac093e77..0d34456972 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -1,3 +1,4 @@ +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use errors::AnyhowResult; @@ -14,7 +15,8 @@ use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; use crate::job::job_types::videofilter::rerender_a_video::rerender_dependencies::RerenderDependencies; use crate::job::job_types::workflow::comfy_ui_dependencies::ComfyDependencies; -use crate::util::scoped_execution::ScopedExecution; +use crate::util::scoped_job_type_execution::ScopedJobTypeExecution; +use crate::util::scoped_model_type_execution::ScopedModelTypeExecution; pub struct JobSpecificDependencies { pub maybe_rvc_v2_dependencies: Option, @@ -34,7 +36,10 @@ pub struct JobSpecificDependencies { impl JobSpecificDependencies { - pub async fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult { + pub async fn setup_for_jobs( + scoped_job_type_execution: &ScopedJobTypeExecution, + scoped_model_type_execution: &ScopedModelTypeExecution + ) -> AnyhowResult { let mut maybe_rvc_v2_dependencies = None; let mut maybe_sad_talker_dependencies = None; let mut maybe_svc_dependencies = None; @@ -49,67 +54,71 @@ impl JobSpecificDependencies { let mut maybe_convert_fbx_to_gltf_dependencies = None; let mut maybe_convert_bvh_to_workflow_dependencies = None; - if scoped_execution.can_run_job(InferenceModelType::RvcV2) { - print_with_space("Setting RVCv2 dependencies..."); - maybe_rvc_v2_dependencies = Some(RvcV2Dependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::SadTalker) { - print_with_space("Setting SadTalker dependencies..."); - maybe_sad_talker_dependencies = Some(SadTalkerDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::SoVitsSvc) { - print_with_space("Setting SVC dependencies..."); - maybe_svc_dependencies = Some(SvcDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::Tacotron2) { - print_with_space("Setting Tacotron2 dependencies..."); - maybe_tacotron2_dependencies = Some(Tacotron2Dependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::VallEX) { - print_with_space("Setting VALL-E-X dependencies..."); - maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::StyleTTS2) { - print_with_space("Setting StyleTTS2 dependencies..."); - maybe_styletts2_dependencies = Some(StyleTTS2Dependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::Vits) { - print_with_space("Setting Vits dependencies..."); - maybe_vits_dependencies = Some(VitsDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::RerenderAVideo) { - print_with_space("Setting Rerender dependencies..."); - maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::StableDiffusion) { - print_with_space("Setting Stable Diffusion dependencies..."); - maybe_stable_diffusion_dependencies = Some(StableDiffusionDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::MocapNet) { - print_with_space("Setting MocapNet dependencies..."); - maybe_mocapnet_dependencies = Some(MocapNetDependencies::setup()?); - } - - if scoped_execution.can_run_job(InferenceModelType::ComfyUi) { + if scoped_model_type_execution.can_run_job(InferenceModelType::ComfyUi) + || scoped_job_type_execution.can_run_job(InferenceJobType::LivePortrait) + || scoped_job_type_execution.can_run_job(InferenceJobType::VideoRender) + || scoped_job_type_execution.can_run_job(InferenceJobType::ComfyUi) + { print_with_space("Setting ComfyUI dependencies..."); maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup().await?); } - if scoped_execution.can_run_job(InferenceModelType::ConvertFbxToGltf) { + if scoped_model_type_execution.can_run_job(InferenceModelType::RvcV2) { + print_with_space("Setting RVCv2 dependencies..."); + maybe_rvc_v2_dependencies = Some(RvcV2Dependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::SadTalker) { + print_with_space("Setting SadTalker dependencies..."); + maybe_sad_talker_dependencies = Some(SadTalkerDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::SoVitsSvc) { + print_with_space("Setting SVC dependencies..."); + maybe_svc_dependencies = Some(SvcDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::Tacotron2) { + print_with_space("Setting Tacotron2 dependencies..."); + maybe_tacotron2_dependencies = Some(Tacotron2Dependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::VallEX) { + print_with_space("Setting VALL-E-X dependencies..."); + maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::StyleTTS2) { + print_with_space("Setting StyleTTS2 dependencies..."); + maybe_styletts2_dependencies = Some(StyleTTS2Dependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::Vits) { + print_with_space("Setting Vits dependencies..."); + maybe_vits_dependencies = Some(VitsDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::RerenderAVideo) { + print_with_space("Setting Rerender dependencies..."); + maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::StableDiffusion) { + print_with_space("Setting Stable Diffusion dependencies..."); + maybe_stable_diffusion_dependencies = Some(StableDiffusionDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::MocapNet) { + print_with_space("Setting MocapNet dependencies..."); + maybe_mocapnet_dependencies = Some(MocapNetDependencies::setup()?); + } + + if scoped_model_type_execution.can_run_job(InferenceModelType::ConvertFbxToGltf) { print_with_space("Setting ConvertFbxToGltf dependencies..."); maybe_convert_fbx_to_gltf_dependencies = Some(FbxToGltfDependencies::setup()?); } - if scoped_execution.can_run_job(InferenceModelType::BvhToWorkflow) { + if scoped_model_type_execution.can_run_job(InferenceModelType::BvhToWorkflow) { print_with_space("Setting ConvertBvhToWorkflow dependencies..."); maybe_convert_bvh_to_workflow_dependencies = Some(RenderEngineSceneToVideoDependencies::setup()?); } diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 9c887262be..775f82e9ce 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -54,7 +54,8 @@ use crate::job_specific_dependencies::JobSpecificDependencies; use crate::util::instrumentation::{init_otel_metrics_pipeline, JobInstrumentLabels}; use crate::util::instrumentation::JobInstruments; use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory; -use crate::util::scoped_execution::ScopedExecution; +use crate::util::scoped_job_type_execution::ScopedJobTypeExecution; +use crate::util::scoped_model_type_execution::ScopedModelTypeExecution; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub mod http_server; @@ -216,7 +217,9 @@ async fn main() -> AnyhowResult<()> { job_stats: job_stats.clone(), }; - let scoped_execution = ScopedExecution::new_from_env()?; + // TODO(bt,2024-07-16): Phase out model type scoping in favor of job type scoping + let scoped_job_type_execution = ScopedJobTypeExecution::new_from_env()?; + let scoped_model_type_execution = ScopedModelTypeExecution::new_from_env()?; let build_sha = std::fs::read_to_string("/GIT_SHA") .unwrap_or(String::from("unknown")) @@ -237,7 +240,8 @@ async fn main() -> AnyhowResult<()> { let meter = opentelemetry::global::meter("inference-job"); - let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution).await?; + let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs( + &scoped_job_type_execution, &scoped_model_type_execution).await?; let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory( easyenv::get_env_pathbuf_or_default( @@ -275,7 +279,8 @@ async fn main() -> AnyhowResult<()> { }, job: JobSystemDependencies { system: JobSystemControls { - scoped_execution, + scoped_model_type_execution, + scoped_job_type_execution, always_allow_cold_filesystem_cache: easyenv::get_env_bool_or_default("ALWAYS_ALLOW_COLD_FILESYSTEM_CACHE", false), cold_filesystem_cache_starvation_threshold: easyenv::get_env_num("COLD_FILESYSTEM_CACHE_STARVATION_THRESHOLD", 3)?, job_batch_wait_millis: common_env.job_batch_wait_millis, diff --git a/crates/service/job/inference_job/src/util/mod.rs b/crates/service/job/inference_job/src/util/mod.rs index ed1595a457..440d9ab368 100644 --- a/crates/service/job/inference_job/src/util/mod.rs +++ b/crates/service/job/inference_job/src/util/mod.rs @@ -5,7 +5,8 @@ pub mod get_polymorphic_args_from_job; pub mod instrumentation; pub mod maybe_download_file_from_bucket; pub mod model_weights_cache; -pub mod scoped_execution; +pub mod scoped_job_type_execution; +pub mod scoped_model_type_execution; pub mod scoped_temp_dir_creator; #[macro_use] diff --git a/crates/service/job/inference_job/src/util/scoped_job_type_execution.rs b/crates/service/job/inference_job/src/util/scoped_job_type_execution.rs new file mode 100644 index 0000000000..dcafdedee2 --- /dev/null +++ b/crates/service/job/inference_job/src/util/scoped_job_type_execution.rs @@ -0,0 +1,101 @@ +use std::collections::BTreeSet; + +use anyhow::anyhow; +use log::info; + +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; +use errors::AnyhowResult; + +/// Execution can be scoped down to run on only certain model types or inference categories. +#[derive(Clone)] +pub struct ScopedJobTypeExecution { + /// If set, only these types of model type will be inferred. + /// None means no scoping, so all types can execute. + scoped_types: Option>, +} + +impl ScopedJobTypeExecution { + + 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_JOB_TYPES") { + Some(comma_separated_types) => Some(parse_job_types(&comma_separated_types)?), + None => None, + }; + + if let Some(types) = scoped_types.as_ref() { + info!("Scoping execution to job types: {:?}", types); + } + + Ok(Self { + scoped_types, + }) + } + + pub fn can_run_job(&self, job_type: InferenceJobType) -> bool { + match self.scoped_types { + None => true, + Some(ref types) => types.contains(&job_type), + } + } + + pub fn get_scoped_job_types(&self) -> Option<&BTreeSet> { + self.scoped_types.as_ref() + } +} + +pub fn parse_job_types(comma_separated_types: &str) -> AnyhowResult> { + let scoped_types = comma_separated_types.trim() + .split(",") + .map(|val| val.to_lowercase()) + .collect::>(); + + let mut job_types = BTreeSet::new(); + + for t in scoped_types.into_iter() { + let model_type = InferenceJobType::from_str(&t) + .map_err(|_err| anyhow!( + "Invalid job type: {:?}; should include only items from: {:?}", + t, + InferenceJobType::all_variants()))?; + + job_types.insert(model_type); + } + + Ok(job_types) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; + + use crate::util::scoped_job_type_execution::{parse_job_types, ScopedJobTypeExecution}; + + #[test] + fn test_parse() { + assert_eq!(parse_job_types("rvc_v2,live_portrait").unwrap(), + BTreeSet::from([InferenceJobType::RvcV2, InferenceJobType::LivePortrait])); + } + + #[test] + fn test_can_execute() { + let scoping = ScopedJobTypeExecution::new_from_set(BTreeSet::from([ + InferenceJobType::RvcV2, + InferenceJobType::SadTalker, + ])); + + assert_eq!(true, scoping.can_run_job(InferenceJobType::RvcV2)); + assert_eq!(true, scoping.can_run_job(InferenceJobType::SadTalker)); + + assert_eq!(false, scoping.can_run_job(InferenceJobType::SoVitsSvc)); + assert_eq!(false, scoping.can_run_job(InferenceJobType::LivePortrait)); + } +} diff --git a/crates/service/job/inference_job/src/util/scoped_execution.rs b/crates/service/job/inference_job/src/util/scoped_model_type_execution.rs similarity index 83% rename from crates/service/job/inference_job/src/util/scoped_execution.rs rename to crates/service/job/inference_job/src/util/scoped_model_type_execution.rs index 4e0546883e..5ec5058c11 100644 --- a/crates/service/job/inference_job/src/util/scoped_execution.rs +++ b/crates/service/job/inference_job/src/util/scoped_model_type_execution.rs @@ -8,20 +8,23 @@ use errors::AnyhowResult; /// Execution can be scoped down to run on only certain model types or inference categories. #[derive(Clone)] -pub struct ScopedExecution { +#[deprecated(note="Use ScopedJobTypeExecution instead.")] +pub struct ScopedModelTypeExecution { /// 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 { +impl ScopedModelTypeExecution { + #[deprecated(note="Use ScopedJobTypeExecution instead.")] pub fn new_from_set(scoped_types: BTreeSet) -> Self { Self { scoped_types: Some(scoped_types), } } + #[deprecated(note="Use ScopedJobTypeExecution instead.")] pub fn new_from_env() -> AnyhowResult { let scoped_types = match easyenv::get_env_string_optional("SCOPED_EXECUTION_MODEL_TYPES") { @@ -38,6 +41,7 @@ impl ScopedExecution { }) } + #[deprecated(note="Use ScopedJobTypeExecution instead.")] pub fn can_run_job(&self, job_model_type: InferenceModelType) -> bool { match self.scoped_types { None => true, @@ -45,6 +49,7 @@ impl ScopedExecution { } } + #[deprecated(note="Use ScopedJobTypeExecution instead.")] pub fn get_scoped_model_types(&self) -> Option<&BTreeSet> { self.scoped_types.as_ref() } @@ -77,7 +82,7 @@ mod tests { use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; - use crate::util::scoped_execution::{parse_model_types, ScopedExecution}; + use crate::util::scoped_model_type_execution::{parse_model_types, ScopedModelTypeExecution}; #[test] fn test_parse() { @@ -87,7 +92,7 @@ mod tests { #[test] fn test_can_execute() { - let scoping = ScopedExecution::new_from_set(BTreeSet::from([ + let scoping = ScopedModelTypeExecution::new_from_set(BTreeSet::from([ InferenceModelType::RvcV2, InferenceModelType::Vits ])); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs index b4195f7a8b..2056c24e2c 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_live_portrait_workflow_handler.rs @@ -215,7 +215,7 @@ pub async fn enqueue_live_portrait_workflow_handler( uuid_idempotency_token: &request.uuid_idempotency_token, job_type: InferenceJobType::LivePortrait, inference_category: InferenceCategory::DeprecatedField, - maybe_model_type: Some(InferenceModelType::ComfyUi), // TODO(bt): Remove! + maybe_model_type: None, maybe_model_token: None, maybe_input_source_token: None, maybe_input_source_token_type: None,