scope jobs by type

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,

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.

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<RvcV2Dependencies>,
@@ -34,7 +36,10 @@ pub struct JobSpecificDependencies {
 
 impl JobSpecificDependencies {
 
-  pub async fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult<Self> {
+  pub async fn setup_for_jobs(
+    scoped_job_type_execution: &ScopedJobTypeExecution,
+    scoped_model_type_execution: &ScopedModelTypeExecution
+  ) -> AnyhowResult<Self> {
     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()?);
     }

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,

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]

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<BTreeSet<InferenceJobType>>,
+}
+
+impl ScopedJobTypeExecution {
+
+  pub fn new_from_set(scoped_types: BTreeSet<InferenceJobType>) -> Self {
+    Self {
+      scoped_types: Some(scoped_types),
+    }
+  }
+
+  pub fn new_from_env() -> AnyhowResult<Self> {
+    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<InferenceJobType>> {
+    self.scoped_types.as_ref()
+  }
+}
+
+pub fn parse_job_types(comma_separated_types: &str) -> AnyhowResult<BTreeSet<InferenceJobType>> {
+  let scoped_types = comma_separated_types.trim()
+      .split(",")
+      .map(|val| val.to_lowercase())
+      .collect::<Vec<String>>();
+
+  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));
+  }
+}

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<BTreeSet<InferenceModelType>>,
 }
 
-impl ScopedExecution {
+impl ScopedModelTypeExecution {
 
+  #[deprecated(note="Use ScopedJobTypeExecution instead.")]
   pub fn new_from_set(scoped_types: BTreeSet<InferenceModelType>) -> Self {
     Self {
       scoped_types: Some(scoped_types),
     }
   }
 
+  #[deprecated(note="Use ScopedJobTypeExecution instead.")]
   pub fn new_from_env() -> AnyhowResult<Self> {
     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<InferenceModelType>> {
     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
     ]));

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,