scoped job execution

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<BTreeSet<InferenceCategory>>,
+  pub maybe_scope_by_model_type: Option<&'a BTreeSet<InferenceModelType>>,
+  pub maybe_scope_by_job_category: Option<&'a BTreeSet<InferenceCategory>>,
   pub mysql_pool: &'a MySqlPool,
 }
 
@@ -85,17 +87,18 @@ pub async fn list_available_generic_inference_jobs(
 )
   -> AnyhowResult<Vec<AvailableInferenceJob>>
 {
-  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<InferenceCategory>) -> Result<Vec<AvailableInferenceJobRawInternal>, sqlx::Error> {
+async fn list_sorted_by_id(args: ListAvailableGenericInferenceJobArgs<'_>, model_types: BTreeSet<InferenceModelType>, inference_categories: BTreeSet<InferenceCategory>) -> Result<Vec<AvailableInferenceJobRawInternal>, 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<InferenceCategory>) -> Result<Vec<AvailableInferenceJobRawInternal>, sqlx::Error> {
+async fn list_sorted_by_priority(args: ListAvailableGenericInferenceJobArgs<'_>, model_types: BTreeSet<InferenceModelType>, inference_categories: BTreeSet<InferenceCategory>) -> Result<Vec<AvailableInferenceJobRawInternal>, 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<T>, etc.
 /// Issue: https://github.com/launchbadge/sqlx/issues/875
-fn to_where_in_predicate(categories: &BTreeSet<InferenceCategory>) -> String {
+fn model_type_predicate(types: &BTreeSet<InferenceModelType>) -> String {
+  let mut vec = types.iter()
+      .map(|ty| ty.to_str())
+      .map(|ty| format!("\"{}\"", ty))
+      .collect::<Vec<String>>();
+  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<T>, etc.
+/// Issue: https://github.com/launchbadge/sqlx/issues/875
+fn inference_category_predicate(categories: &BTreeSet<InferenceCategory>) -> 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<InferenceCategory>) -> 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());
   }
 }

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;
 

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,
 

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(),

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;

crates/service/job/inference_job/src/util/scope_model_types.rs

@@ -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<BTreeSet<InferenceModelType>> {
-  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<BTreeSet<InferenceModelType>> {
-  let scoped_types = comma_separated_types.trim()
-      .split(",")
-      .map(|val| val.to_lowercase())
-      .collect::<Vec<String>>();
-
-  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]));
-  }
-}

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