scope download job type

crates/_types/enums/src/by_table/generic_download_jobs/generic_download_type.rs

@@ -8,7 +8,7 @@
 /// These types are present in the HTTP API and database columns as serialized here.
 ///
 /// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY.
-#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)]
+#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type, Hash, Ord, PartialOrd)]
 pub enum GenericDownloadType {
   /// NB: Note - this is hifigan for Tacotron2.
   /// Some work will be needed to unify this with other hifigan types.

crates/lib/mysql_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs

@@ -1,13 +1,11 @@
-use anyhow::anyhow;
 use chrono::Utc;
 use crate::column_types::job_status::JobStatus;
 use crate::queries::generic_download::job::_keys::GenericDownloadJobId;
-use crate::tokens::Tokens;
 use enums::by_table::generic_download_jobs::generic_download_type::GenericDownloadType;
 use enums::common::visibility::Visibility;
 use errors::AnyhowResult;
 use sqlx::MySqlPool;
-use std::path::Path;
+use std::collections::BTreeSet;
 use tokens::jobs::download::DownloadJobToken;
 
 /// table: generic_download_jobs
@@ -33,12 +31,11 @@ pub struct AvailableDownloadJob {
   pub retry_at: Option<chrono::DateTime<Utc>>,
 }
 
-pub async fn list_available_generic_download_jobs(pool: &MySqlPool, num_records: u32)
+pub async fn list_available_generic_download_jobs(pool: &MySqlPool, num_records: u32, download_types: &BTreeSet<GenericDownloadType>)
   -> AnyhowResult<Vec<AvailableDownloadJob>>
 {
-  let job_records = sqlx::query_as!(
-      AvailableDownloadJobRawInternal,
-        r#"
+
+  let mut query = String::from(r#"
 SELECT
   id,
   token AS `download_job_token: tokens::jobs::download::DownloadJobToken`,
@@ -69,11 +66,20 @@ WHERE
     OR
     retry_at < CURRENT_TIMESTAMP
   )
-  ORDER BY id ASC
-  LIMIT ?
-        "#,
-      num_records,
-    )
+  "#);
+
+  if let Some(clause) = download_type_clause(download_types) {
+    query.push_str(" AND ");
+    query.push_str(&clause);
+    query.push_str(" ");
+  }
+
+  query.push_str(&format!(r#"
+    ORDER BY id ASC
+    LIMIT {}
+  "#, num_records));
+
+  let mut job_records = sqlx::query_as::<_, AvailableDownloadJobRawInternal>(&query)
       .fetch_all(pool)
       .await?;
 
@@ -101,7 +107,21 @@ WHERE
   Ok(job_records)
 }
 
-#[derive(Debug)]
+fn download_type_clause(download_types: &BTreeSet<GenericDownloadType>) -> Option<String> {
+  if download_types.is_empty() {
+    return None;
+  }
+
+  let download_types = download_types.into_iter()
+      .map(|download_type| download_type.to_str())
+      .map(|download_type| format!("\"{}\"", download_type))
+      .collect::<Vec<_>>()
+      .join(", ");
+
+  Some(format!("download_type IN ( {} )", download_types))
+}
+
+#[derive(Debug, sqlx::FromRow)]
 struct AvailableDownloadJobRawInternal {
   pub id: i64,
   pub download_job_token: DownloadJobToken,
@@ -122,3 +142,23 @@ struct AvailableDownloadJobRawInternal {
   pub updated_at: chrono::DateTime<Utc>,
   pub retry_at: Option<chrono::DateTime<Utc>>,
 }
+
+
+#[cfg(test)]
+mod tests {
+  use std::collections::BTreeSet;
+  use enums::by_table::generic_download_jobs::generic_download_type::GenericDownloadType;
+  use crate::queries::generic_download::job::list_available_generic_download_jobs::download_type_clause;
+
+  #[test]
+  fn test_download_types_with_clause() {
+    let clause = download_type_clause(&BTreeSet::from([GenericDownloadType::HifiGan, GenericDownloadType::RvcV2]));
+    assert_eq!(clause.as_deref(), Some("download_type IN ( \"hifigan\", \"rvc_v2\" )"));
+  }
+
+  #[test]
+  fn test_download_types_with_clause_absent() {
+    let clause = download_type_clause(&BTreeSet::from([]));
+    assert_eq!(clause, None);
+  }
+}

crates/service/job/download_job/src/job_loop/main_loop.rs

@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use container_common::anyhow_result::AnyhowResult;
 use crate::JobState;
 use crate::job_loop::process_single_job::process_single_job;
@@ -6,7 +7,6 @@ use log::{error, warn};
 use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::{AvailableDownloadJob, list_available_generic_download_jobs};
 use mysql_queries::queries::generic_download::job::mark_generic_download_job_failure::mark_generic_download_job_failure;
 use std::time::Duration;
-use anyhow::anyhow;
 
 // Job runner timeouts (guards MySQL)
 const START_TIMEOUT_MILLIS : u64 = 500;
@@ -25,7 +25,11 @@ pub async fn main_loop(job_state: JobState) {
     }
 
     let num_records = 1;
-    let maybe_available_jobs = list_available_generic_download_jobs(&job_state.mysql_pool, num_records).await;
+
+    let maybe_available_jobs = list_available_generic_download_jobs(
+      &job_state.mysql_pool,
+      num_records,
+      &job_state.download_types).await;
 
     let jobs = match maybe_available_jobs {
       Ok(jobs) => jobs,

crates/service/job/download_job/src/job_state.rs

@@ -10,6 +10,7 @@ use crate::job_types::voice_conversion::rvc_v2::rvc_v2_model_check_command::RvcV
 use crate::job_types::voice_conversion::so_vits_svc::so_vits_svc_model_check_command::SoVitsSvcModelCheckCommand;
 use crate::job_types::voice_conversion::softvc::softvc_model_check_command::SoftVcModelCheckCommand;
 use crate::threads::nvidia_smi_checker::nvidia_smi_health_check_status::NvidiaSmiHealthCheckStatus;
+use enums::by_table::generic_download_jobs::generic_download_type::GenericDownloadType;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
 use mysql_queries::common_inputs::container_environment_arg::ContainerEnvironmentArg;
 use mysql_queries::mediators::badge_granter::BadgeGranter;
@@ -17,9 +18,13 @@ use mysql_queries::mediators::firehose_publisher::FirehosePublisher;
 use r2d2_redis::RedisConnectionManager;
 use r2d2_redis::r2d2;
 use sqlx::MySqlPool;
+use std::collections::BTreeSet;
 use std::path::PathBuf;
 
 pub struct JobState {
+  /// The types of downloads this instance of the job will handle.
+  pub download_types: BTreeSet<GenericDownloadType>,
+
   pub download_temp_directory: PathBuf,
   pub mysql_pool: MySqlPool,
 

crates/service/job/download_job/src/main.rs

@@ -23,6 +23,7 @@ pub mod job_types;
 pub mod threads;
 pub mod util;
 
+use std::collections::BTreeSet;
 use bootstrap::bootstrap::{bootstrap, BootstrapArgs};
 use cloud_storage::bucket_client::BucketClient;
 use cloud_storage::bucket_path_unifier::BucketPathUnifier;
@@ -55,6 +56,7 @@ use std::path::PathBuf;
 use std::time::Duration;
 use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
 use tokio::runtime::Runtime;
+use enums::by_table::generic_download_jobs::generic_download_type::GenericDownloadType;
 use crate::job_types::voice_conversion::rvc_v2::pretrained_hubert_model::PretrainedHubertModel;
 
 // Buckets
@@ -337,7 +339,17 @@ async fn main() -> AnyhowResult<()> {
     ).await;
   });
 
+  let mut download_types = BTreeSet::new();
+
+  download_types.insert(GenericDownloadType::RvcV2);
+  download_types.insert(GenericDownloadType::SoVitsSvc);
+  //download_types.insert(GenericDownloadType::Vits);
+  //download_types.insert(GenericDownloadType::HifiGan);
+
+  info!("Looking up download jobs of type {:?}", download_types);
+
   let job_state = JobState {
+    download_types,
     download_temp_directory: temp_directory,
     mysql_pool,
     redis_pool,