From ffcf460291b596e08aaecaf7f2c04163ecda8f4d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Jul 2023 14:40:28 -0400 Subject: [PATCH] scope download job type --- .../generic_download_type.rs | 2 +- .../list_available_generic_download_jobs.rs | 66 +++++++++++++++---- .../download_job/src/job_loop/main_loop.rs | 8 ++- .../service/job/download_job/src/job_state.rs | 5 ++ crates/service/job/download_job/src/main.rs | 12 ++++ 5 files changed, 77 insertions(+), 16 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_download_jobs/generic_download_type.rs b/crates/_types/enums/src/by_table/generic_download_jobs/generic_download_type.rs index be825611f4..d007f13eef 100644 --- a/crates/_types/enums/src/by_table/generic_download_jobs/generic_download_type.rs +++ b/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. diff --git a/crates/lib/mysql_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs b/crates/lib/mysql_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs index 8be01c0eb7..e6273feb42 100644 --- a/crates/lib/mysql_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs +++ b/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>, } -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) -> AnyhowResult> { - 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) -> Option { + 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::>() + .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, pub retry_at: Option>, } + + +#[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); + } +} diff --git a/crates/service/job/download_job/src/job_loop/main_loop.rs b/crates/service/job/download_job/src/job_loop/main_loop.rs index 4008d6254e..1a9ebc065f 100644 --- a/crates/service/job/download_job/src/job_loop/main_loop.rs +++ b/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, diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index d48008f2b7..529a25da01 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/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, + pub download_temp_directory: PathBuf, pub mysql_pool: MySqlPool, diff --git a/crates/service/job/download_job/src/main.rs b/crates/service/job/download_job/src/main.rs index 435d9b92e6..3151b970d1 100644 --- a/crates/service/job/download_job/src/main.rs +++ b/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,