diff --git a/crates/lib/database_queries/src/column_types/job_status.rs b/crates/lib/database_queries/src/column_types/job_status.rs index a9c03a9860..e3606e50ca 100644 --- a/crates/lib/database_queries/src/column_types/job_status.rs +++ b/crates/lib/database_queries/src/column_types/job_status.rs @@ -7,8 +7,8 @@ use container_common::anyhow_result::AnyhowResult; /// See: https://www.gitmemory.com/issue/launchbadge/sqlx/1241/847154375 /// eg. job_status as `job_status: crate::column_types::job_status::JobStatus` #[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] -#[serde(rename_all = "lowercase")] -#[sqlx(rename_all = "lowercase")] +#[serde(rename_all = "snake_case")] +#[sqlx(rename_all = "snake_case")] pub enum JobStatus { Pending, Started, diff --git a/crates/lib/database_queries/src/column_types/mod.rs b/crates/lib/database_queries/src/column_types/mod.rs index 4b0786fb98..1c108ef9c1 100644 --- a/crates/lib/database_queries/src/column_types/mod.rs +++ b/crates/lib/database_queries/src/column_types/mod.rs @@ -1,2 +1,3 @@ pub mod job_status; -pub mod record_visibility; \ No newline at end of file +pub mod record_visibility; +pub mod vocoder_type; \ No newline at end of file diff --git a/crates/lib/database_queries/src/column_types/vocoder_type.rs b/crates/lib/database_queries/src/column_types/vocoder_type.rs new file mode 100644 index 0000000000..ad92b2687e --- /dev/null +++ b/crates/lib/database_queries/src/column_types/vocoder_type.rs @@ -0,0 +1,31 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum VocoderType { + /// Current state of the art vocoder + #[serde(rename = "hifigan-superres")] + HifiGanSuperResolution, + + /// Older WaveGlow vocoder + WaveGlow, +} + +impl VocoderType { + pub fn to_str(&self) -> &'static str { + match self { + VocoderType::HifiGanSuperResolution => "hifigan-superres", + VocoderType::WaveGlow => "waveglow", + } + } + + pub fn from_str(vocoder_type: &str) -> AnyhowResult { + match vocoder_type { + "hifigan-superres" => Ok(VocoderType::HifiGanSuperResolution), + "waveglow" => Ok(VocoderType::WaveGlow), + _ => Err(anyhow!("invalid value: {:?}", vocoder_type)), + } + } +} diff --git a/crates/lib/database_queries/src/tts/mod.rs b/crates/lib/database_queries/src/tts/mod.rs index 39427a96cc..00c9ce9ee9 100644 --- a/crates/lib/database_queries/src/tts/mod.rs +++ b/crates/lib/database_queries/src/tts/mod.rs @@ -1,2 +1,3 @@ -pub mod insert_tts_inference_job; -pub mod list_available_tts_inference_jobs; \ No newline at end of file +pub mod tts_inference_jobs; +pub mod tts_models; +pub mod tts_results; \ No newline at end of file diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/_keys.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/_keys.rs new file mode 100644 index 0000000000..0a6b60dd8f --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/_keys.rs @@ -0,0 +1,5 @@ +//! Primary and foreign keys for TTS-related models + +/// Since jobs are internal, we'll use IDs as primary keys. +#[derive(Clone, Copy, PartialOrd, PartialEq, Eq, Ord, Debug)] +pub struct TtsInferenceJobId(pub i64); diff --git a/crates/lib/database_queries/src/tts/insert_tts_inference_job.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/insert_tts_inference_job.rs similarity index 100% rename from crates/lib/database_queries/src/tts/insert_tts_inference_job.rs rename to crates/lib/database_queries/src/tts/tts_inference_jobs/insert_tts_inference_job.rs diff --git a/crates/lib/database_queries/src/tts/list_available_tts_inference_jobs.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs similarity index 84% rename from crates/lib/database_queries/src/tts/list_available_tts_inference_jobs.rs rename to crates/lib/database_queries/src/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs index 1c077bdb62..718414db09 100644 --- a/crates/lib/database_queries/src/tts/list_available_tts_inference_jobs.rs +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs @@ -4,13 +4,14 @@ use container_common::anyhow_result::AnyhowResult; use crate::column_types::job_status::JobStatus; use crate::column_types::record_visibility::RecordVisibility; use crate::helpers::boolean_converters::i8_to_bool; +use crate::tts::tts_inference_jobs::_keys::TtsInferenceJobId; use log::{warn, info}; use sqlx::MySqlPool; /// table: tts_inference_jobs #[derive(Debug)] -pub struct TtsInferenceJob { - pub id: i64, +pub struct AvailableTtsInferenceJob { + pub id: TtsInferenceJobId, pub inference_job_token: String, pub uuid_idempotency_token: String, @@ -35,13 +36,13 @@ pub struct TtsInferenceJob { } /// Query jobs that are ready to run -pub async fn query_tts_inference_job_records( +pub async fn list_available_tts_inference_jobs( pool: &MySqlPool, num_records: u32 -) -> AnyhowResult> { +) -> AnyhowResult> { - let job_records : Vec = sqlx::query_as!( - TtsInferenceJobRawInternal, + let job_records : Vec = sqlx::query_as!( + AvailableTtsInferenceJobRawInternal, r#" SELECT id, @@ -85,9 +86,9 @@ WHERE .await?; let job_records = job_records.into_iter() - .map(|record : TtsInferenceJobRawInternal| { - TtsInferenceJob { - id: record.id, + .map(|record : AvailableTtsInferenceJobRawInternal| { + AvailableTtsInferenceJob { + id: TtsInferenceJobId(record.id), inference_job_token: record.inference_job_token, uuid_idempotency_token: record.uuid_idempotency_token, model_token: record.model_token, @@ -106,13 +107,13 @@ WHERE retry_at: record.retry_at, } }) - .collect::>(); + .collect::>(); Ok(job_records) } #[derive(Debug)] -pub struct TtsInferenceJobRawInternal { +struct AvailableTtsInferenceJobRawInternal { pub id: i64, pub inference_job_token: String, pub uuid_idempotency_token: String, diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_done.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_done.rs new file mode 100644 index 0000000000..5acb43700d --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_done.rs @@ -0,0 +1,33 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +pub async fn mark_tts_inference_job_done( + pool: &MySqlPool, + job_id: TtsInferenceJobId, + success: bool, + maybe_result_token: Option<&str> +) -> AnyhowResult<()> { + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = ?, + on_success_result_token = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_result_token, + job_id.0 + ) + .execute(pool) + .await?; + + Ok(()) +} + diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs new file mode 100644 index 0000000000..19824707be --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs @@ -0,0 +1,41 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; +use sqlx::MySqlPool; +use sqlx; + +// TODO: Rerun time shouldn't be 1-minute! + +/// Mark a single inference failure. The job may be re-run. +pub async fn mark_tts_inference_job_failure( + pool: &MySqlPool, + job: &AvailableTtsInferenceJob, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 1 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id.0, + ) + .execute(pool) + .await?; + + Ok(()) +} diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs new file mode 100644 index 0000000000..98708357bd --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs @@ -0,0 +1,85 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +// TODO(2022-02-09): It's unclear how stuck "pending" jobs don't become perma-dead +// I think this logic is flawed and the job system needs to pick up stuck "pending" jobs. + +pub struct TtsInferenceLockRecord { + id: i64, + status: String, +} + +pub async fn mark_tts_inference_job_pending_and_grab_lock( + pool: &MySqlPool, + job_id: TtsInferenceJobId +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + TtsInferenceLockRecord, + r#" +SELECT + id, + status +FROM tts_inference_jobs +WHERE id = ? +FOR UPDATE + "#, + job_id.0, + ) + .fetch_one(&mut transaction) + .await; + + let record : TtsInferenceLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job_id.0, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} + diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs new file mode 100644 index 0000000000..e1441afb38 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs @@ -0,0 +1,28 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +/// Mark a TTS job as dead ("do not ever retry"). +pub async fn mark_tts_inference_job_permanently_dead( + pool: &MySqlPool, + job_id: TtsInferenceJobId, + failure_reason: &str, +) -> AnyhowResult<()> { + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = "dead", + failure_reason = ?, + retry_at = NULL +WHERE id = ? + "#, + failure_reason.to_string(), + job_id.0, + ) + .execute(pool) + .await?; + + Ok(()) +} diff --git a/crates/lib/database_queries/src/tts/tts_inference_jobs/mod.rs b/crates/lib/database_queries/src/tts/tts_inference_jobs/mod.rs new file mode 100644 index 0000000000..7aaa57a6d6 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_inference_jobs/mod.rs @@ -0,0 +1,7 @@ +pub mod _keys; +pub mod insert_tts_inference_job; +pub mod list_available_tts_inference_jobs; +pub mod mark_tts_inference_job_done; +pub mod mark_tts_inference_job_failure; +pub mod mark_tts_inference_job_pending_and_grab_lock; +pub mod mark_tts_inference_job_permanently_dead; diff --git a/crates/lib/database_queries/src/tts/tts_models/get_tts_model_for_inference.rs b/crates/lib/database_queries/src/tts/tts_models/get_tts_model_for_inference.rs new file mode 100644 index 0000000000..a5d131a7c4 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_models/get_tts_model_for_inference.rs @@ -0,0 +1,103 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use log::warn; +use sqlx::MySqlPool; +use sqlx; + +// TODO: Can probably just reuse another query. + +pub struct TtsModelForInferenceRecord { + pub model_token: String, + pub tts_model_type: String, + + pub maybe_default_pretrained_vocoder: Option, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + + pub title: String, + pub private_bucket_hash: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub mod_deleted_at: Option>, + pub user_deleted_at: Option>, +} + + +#[derive(Clone, Debug)] +pub enum TtsModelForInferenceError { + ModelNotFound, + ModelDeleted, + DatabaseError { reason: String }, +} + +impl std::fmt::Display for TtsModelForInferenceError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TtsModelForInferenceError::ModelNotFound => write!(f, "ModelNotFound"), + TtsModelForInferenceError::ModelDeleted => write!(f, "ModelDeleted"), + TtsModelForInferenceError::DatabaseError { reason} => + write!(f, "Database error: {:?}", reason), + } + } +} + +impl std::error::Error for TtsModelForInferenceError {} + +pub async fn get_tts_model_for_inference( + pool: &MySqlPool, + model_token: &str +) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_model = sqlx::query_as!( + TtsModelForInferenceRecord, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.maybe_default_pretrained_vocoder, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + tts.title, + tts.private_bucket_hash, + tts.created_at, + tts.updated_at, + tts.user_deleted_at, + tts.mod_deleted_at +FROM tts_models as tts +JOIN users +ON users.token = tts.creator_user_token +WHERE tts.token = ? + "#, + &model_token + ) + .fetch_one(pool) + .await; // TODO: This will return error if it doesn't exist + + let model : TtsModelForInferenceRecord = match maybe_model { + Ok(model) => model, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(TtsModelForInferenceError::ModelNotFound); + }, + _ => { + warn!("tts model query error: {:?}", err); + return Err(TtsModelForInferenceError::DatabaseError { + reason: format!("Mysql error: {:?}", err) + }); + } + } + } + }; + + if model.mod_deleted_at.is_some() || model.user_deleted_at.is_some() { + return Err(TtsModelForInferenceError::ModelDeleted); + } + + Ok(model) +} diff --git a/crates/lib/database_queries/src/tts/tts_models/mod.rs b/crates/lib/database_queries/src/tts/tts_models/mod.rs new file mode 100644 index 0000000000..0bd84e0e68 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_models/mod.rs @@ -0,0 +1 @@ +pub mod get_tts_model_for_inference; \ No newline at end of file diff --git a/crates/lib/database_queries/src/tts/tts_results/insert_tts_result.rs b/crates/lib/database_queries/src/tts/tts_results/insert_tts_result.rs new file mode 100644 index 0000000000..be881d67f3 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_results/insert_tts_result.rs @@ -0,0 +1,165 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::vocoder_type::VocoderType; +use crate::tokens::Tokens; +use crate::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; +use log::warn; +use sqlx::MySqlPool; +use sqlx; +use std::path::Path; + +/// Used to give user-facing order to logged in user inference requests +pub struct SyntheticIdRecord { + pub next_id: i64, +} + +pub async fn insert_tts_result>( + pool: &MySqlPool, + job: &AvailableTtsInferenceJob, + text_hash: &str, + pretrained_vocoder_used: VocoderType, + bucket_audio_results_path: P, + bucket_spectrogram_results_path: P, + file_size_bytes: u64, + duration_millis: u64 +) -> AnyhowResult<(u64, String)> +{ + let inference_result_token = Tokens::new_tts_result()?; + + let bucket_audio_result_path = &bucket_audio_results_path + .as_ref() + .display() + .to_string(); + + let bucket_spectrogram_result_path = &bucket_spectrogram_results_path + .as_ref() + .display() + .to_string(); + + let normalized_inference_text = job.raw_inference_text.clone(); // TODO + + let maybe_creator_user_token = job.maybe_creator_user_token.clone(); + let mut maybe_creator_synthetic_id : Option = None; + + let mut transaction = pool.begin().await?; + + if let Some(creator_user_token) = maybe_creator_user_token.as_deref() { + let query_result = sqlx::query!( + r#" +INSERT INTO tts_result_synthetic_ids +SET + user_token = ?, + next_id = 1 +ON DUPLICATE KEY UPDATE + user_token = ?, + next_id = next_id + 1 + "#, + creator_user_token, + creator_user_token + ) + .execute(&mut transaction) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + //transaction.rollback().await?; + warn!("Transaction failure: {:?}", err); + } + } + + let query_result = sqlx::query_as!( + SyntheticIdRecord, + r#" +SELECT + next_id +FROM + tts_result_synthetic_ids +WHERE + user_token = ? +LIMIT 1 + "#, + creator_user_token, + ) + .fetch_one(&mut transaction) + .await; + + let record : SyntheticIdRecord = match query_result { + Ok(record) => record, + Err(err) => { + warn!("Transaction failure: {:?}", err); + transaction.rollback().await?; + return Err(anyhow!("Transaction failure: {:?}", err)); + } + }; + + let next_id = record.next_id as u64; + maybe_creator_synthetic_id = Some(next_id); + } + + let record_id = { + let query_result = sqlx::query!( + r#" +INSERT INTO tts_results +SET + token = ?, + + model_token = ?, + maybe_pretrained_vocoder_used = ?, + raw_inference_text = ?, + raw_inference_text_hash_sha2 = ?, + normalized_inference_text = ?, + + maybe_creator_user_token = ?, + maybe_creator_synthetic_id = ?, + + creator_ip_address = ?, + creator_set_visibility = ?, + + public_bucket_wav_audio_path = ?, + public_bucket_spectrogram_path = ?, + + file_size_bytes = ?, + duration_millis = ? + "#, + inference_result_token, + job.model_token.clone(), + pretrained_vocoder_used.to_str(), + job.raw_inference_text.clone(), + text_hash, + normalized_inference_text, + + maybe_creator_user_token, + maybe_creator_synthetic_id, + + job.creator_ip_address.clone(), + job.creator_set_visibility.to_str(), + + bucket_audio_result_path, + bucket_spectrogram_result_path, + + file_size_bytes, + duration_millis + ) + .execute(&mut transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((record_id, inference_result_token.clone())) +} + diff --git a/crates/lib/database_queries/src/tts/tts_results/mod.rs b/crates/lib/database_queries/src/tts/tts_results/mod.rs new file mode 100644 index 0000000000..933ac21575 --- /dev/null +++ b/crates/lib/database_queries/src/tts/tts_results/mod.rs @@ -0,0 +1 @@ +pub mod insert_tts_result; \ No newline at end of file diff --git a/crates/service/obs_gateway/src/threads/twitch_pubsub_user_subscriber_thread.rs b/crates/service/obs_gateway/src/threads/twitch_pubsub_user_subscriber_thread.rs index 2f18feac06..81bd4cac4e 100644 --- a/crates/service/obs_gateway/src/threads/twitch_pubsub_user_subscriber_thread.rs +++ b/crates/service/obs_gateway/src/threads/twitch_pubsub_user_subscriber_thread.rs @@ -9,7 +9,7 @@ use crate::twitch::pubsub::build_pubsub_topics_for_user::build_pubsub_topics_for use crate::twitch::twitch_user_id::TwitchUserId; use crate::twitch::websocket_client::TwitchWebsocketClient; use database_queries::tokens::Tokens; -use database_queries::tts::insert_tts_inference_job::TtsInferenceJobInsertBuilder; +use database_queries::tts::tts_inference_jobs::insert_tts_inference_job::TtsInferenceJobInsertBuilder; use database_queries::twitch_oauth::find::{TwitchOauthTokenRecord, TwitchOauthTokenFinder}; use database_queries::twitch_oauth::insert::TwitchOauthTokenInsertBuilder; use database_queries::twitch_pubsub::insert_bits::TwitchPubsubBitsInsertBuilder; diff --git a/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs b/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs index 149f956239..4e990e78b3 100644 --- a/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs +++ b/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs @@ -1,7 +1,7 @@ use actix_web::http::header; use anyhow::anyhow; use container_common::anyhow_result::AnyhowResult; -use crate::database::enums::vocoder_type::VocoderType; +use database_queries::column_types::vocoder_type::VocoderType; use hyper::client::Client; use hyper::{Body, Request}; use log::info; diff --git a/crates/service/storyteller_web/src/job_queries/mod.rs b/crates/service/storyteller_web/src/job_queries/mod.rs index fe1fc8930b..c25ac9cc81 100644 --- a/crates/service/storyteller_web/src/job_queries/mod.rs +++ b/crates/service/storyteller_web/src/job_queries/mod.rs @@ -1,4 +1,6 @@ +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + pub mod tts_download_job_queries; -pub mod tts_inference_job_queries; pub mod w2l_download_job_queries; pub mod w2l_inference_job_queries; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs b/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs index b3a4a85560..3108000ded 100644 --- a/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs +++ b/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs @@ -9,6 +9,9 @@ use database_queries::tokens::Tokens; use sqlx::MySqlPool; use std::path::Path; +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + /// table: tts_model_upload_jobs #[derive(Debug)] pub struct TtsUploadJobRecord { diff --git a/crates/service/storyteller_web/src/job_queries/tts_inference_job_queries.rs b/crates/service/storyteller_web/src/job_queries/tts_inference_job_queries.rs deleted file mode 100644 index 31cb75f245..0000000000 --- a/crates/service/storyteller_web/src/job_queries/tts_inference_job_queries.rs +++ /dev/null @@ -1,504 +0,0 @@ -//! NB: This seems required for sqlx to generate the cached queries. -//! Sqlx's prepare needs a *single* binary to work against, so we need to -//! include these in the main binary to generate all the queries. - -use anyhow::anyhow; -use chrono::{Utc, DateTime}; -use container_common::anyhow_result::AnyhowResult; -use crate::database::enums::record_visibility::RecordVisibility; -use crate::database::enums::vocoder_type::VocoderType; -use database_queries::tokens::Tokens; -use log::{warn, info}; -use sqlx::{MySqlPool, Transaction, MySql}; -use std::error; -use std::path::Path; - -/// table: tts_inference_jobs -#[derive(Debug)] -pub struct TtsInferenceJobRecord { - pub id: i64, - pub inference_job_token: String, - pub uuid_idempotency_token: String, - - pub model_token: String, - pub raw_inference_text: String, - - pub creator_ip_address: String, - pub maybe_creator_user_token: Option, - pub creator_set_visibility: RecordVisibility, - - pub is_from_api: i8, - pub is_for_twitch: i8, - - pub status: String, // TODO - pub priority_level: u8, - pub attempt_count: i32, - pub failure_reason: Option, - - pub created_at: chrono::DateTime, - pub updated_at: chrono::DateTime, - pub retry_at: Option>, -} - -pub async fn query_tts_inference_job_records( - pool: &MySqlPool, - num_records: u32 -) -> AnyhowResult> { - - let job_records = sqlx::query_as!( - TtsInferenceJobRecord, - r#" -SELECT - id, - token AS inference_job_token, - uuid_idempotency_token, - - model_token, - raw_inference_text, - - creator_ip_address, - maybe_creator_user_token, - creator_set_visibility as `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, - - is_from_api, - is_for_twitch, - - status, - priority_level, - attempt_count, - failure_reason, - created_at, - updated_at, - retry_at -FROM tts_inference_jobs -WHERE - ( - status IN ("pending", "attempt_failed") - ) - AND - ( - retry_at IS NULL - OR - retry_at < CURRENT_TIMESTAMP - ) - LIMIT ? - "#, - num_records, - ) - .fetch_all(pool) - .await?; - - Ok(job_records) -} - - -pub struct TtsInferenceLockRecord { - id: i64, - status: String, -} - -pub async fn grab_job_lock_and_mark_pending( - pool: &MySqlPool, - job: &TtsInferenceJobRecord -) -> AnyhowResult { - - // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. - let mut transaction = pool.begin().await?; - - let maybe_record = sqlx::query_as!( - TtsInferenceLockRecord, - r#" -SELECT - id, - status -FROM tts_inference_jobs -WHERE id = ? -FOR UPDATE - "#, - job.id, - ) - .fetch_one(&mut transaction) - .await; - - let record : TtsInferenceLockRecord = match maybe_record { - Ok(record) => record, - Err(err) => { - match err { - sqlx::Error::RowNotFound => { - return Err(anyhow!("could not job")); - }, - _ => { - return Err(anyhow!("query error")); - } - } - } - }; - - let can_transact = match record.status.as_ref() { - "pending" => true, // It's okay for us to take the lock. - "attempt_failed" => true, // We can retry. - "started" => false, // Job in progress (another job beat us, and we can't take the lock) - "complete_success" => false, // Job already complete - "complete_failure" => false, // Job already complete (permanently dead; no need to retry) - "dead" => false, // Job already complete (permanently dead; retries exhausted) - _ => false, // Future-proof - }; - - if !can_transact { - transaction.rollback().await?; - return Ok(false); - } - - let _acquire_lock = sqlx::query!( - r#" -UPDATE tts_inference_jobs -SET - status = 'started', - attempt_count = attempt_count + 1, - retry_at = NOW() + interval 2 minute -WHERE id = ? - "#, - job.id, - ) - .execute(&mut transaction) - .await?; - - transaction.commit().await?; - - Ok(true) -} - -pub async fn mark_tts_inference_job_failure( - pool: &MySqlPool, - job: &TtsInferenceJobRecord, - failure_reason: &str, - max_attempts: i32 -) -> AnyhowResult<()> { - - // statuses: "attempt_failed", "complete_failure", "dead" - let mut next_status = "attempt_failed"; - - if job.attempt_count >= max_attempts { - // NB: Job attempt count is incremented at start - next_status = "dead"; - } - - let query_result = sqlx::query!( - r#" -UPDATE tts_inference_jobs -SET - status = ?, - failure_reason = ?, - retry_at = NOW() + interval 1 minute -WHERE id = ? - "#, - next_status, - failure_reason.to_string(), - job.id, - ) - .execute(pool) - .await?; - - Ok(()) -} - -pub async fn mark_tts_inference_job_permanently_dead( - pool: &MySqlPool, - job: &TtsInferenceJobRecord, - failure_reason: &str, -) -> AnyhowResult<()> { - let query_result = sqlx::query!( - r#" -UPDATE tts_inference_jobs -SET - status = "dead", - failure_reason = ?, - retry_at = NULL -WHERE id = ? - "#, - failure_reason.to_string(), - job.id, - ) - .execute(pool) - .await?; - - Ok(()) -} - -pub async fn mark_tts_inference_job_done( - pool: &MySqlPool, - job: &TtsInferenceJobRecord, - success: bool, - maybe_result_token: Option<&str> -) -> AnyhowResult<()> { - let status = if success { "complete_success" } else { "complete_failure" }; - - let query_result = sqlx::query!( - r#" -UPDATE tts_inference_jobs -SET - status = ?, - on_success_result_token = ?, - failure_reason = NULL, - retry_at = NULL -WHERE id = ? - "#, - status, - maybe_result_token, - job.id - ) - .execute(pool) - .await?; - - Ok(()) -} - -pub struct SyntheticIdRecord { - pub next_id: i64, -} - -pub async fn insert_tts_result>( - pool: &MySqlPool, - job: &TtsInferenceJobRecord, - text_hash: &str, - pretrained_vocoder_used: VocoderType, - bucket_audio_results_path: P, - bucket_spectrogram_results_path: P, - file_size_bytes: u64, - duration_millis: u64 -) -> AnyhowResult<(u64, String)> -{ - let inference_result_token = Tokens::new_tts_result()?; - - let bucket_audio_result_path = &bucket_audio_results_path - .as_ref() - .display() - .to_string(); - - let bucket_spectrogram_result_path = &bucket_spectrogram_results_path - .as_ref() - .display() - .to_string(); - - let normalized_inference_text = job.raw_inference_text.clone(); // TODO - - let maybe_creator_user_token = job.maybe_creator_user_token.clone(); - let mut maybe_creator_synthetic_id : Option = None; - - let mut transaction = pool.begin().await?; - - if let Some(creator_user_token) = maybe_creator_user_token.as_deref() { - let query_result = sqlx::query!( - r#" -INSERT INTO tts_result_synthetic_ids -SET - user_token = ?, - next_id = 1 -ON DUPLICATE KEY UPDATE - user_token = ?, - next_id = next_id + 1 - "#, - creator_user_token, - creator_user_token - ) - .execute(&mut transaction) - .await; - - match query_result { - Ok(_) => {}, - Err(err) => { - //transaction.rollback().await?; - warn!("Transaction failure: {:?}", err); - } - } - - let query_result = sqlx::query_as!( - SyntheticIdRecord, - r#" -SELECT - next_id -FROM - tts_result_synthetic_ids -WHERE - user_token = ? -LIMIT 1 - "#, - creator_user_token, - ) - .fetch_one(&mut transaction) - .await; - - let record : SyntheticIdRecord = match query_result { - Ok(record) => record, - Err(err) => { - warn!("Transaction failure: {:?}", err); - transaction.rollback().await?; - return Err(anyhow!("Transaction failure: {:?}", err)); - } - }; - - let next_id = record.next_id as u64; - maybe_creator_synthetic_id = Some(next_id); - } - - let record_id = { - let query_result = sqlx::query!( - r#" -INSERT INTO tts_results -SET - token = ?, - - model_token = ?, - maybe_pretrained_vocoder_used = ?, - raw_inference_text = ?, - raw_inference_text_hash_sha2 = ?, - normalized_inference_text = ?, - - maybe_creator_user_token = ?, - maybe_creator_synthetic_id = ?, - - creator_ip_address = ?, - creator_set_visibility = ?, - - public_bucket_wav_audio_path = ?, - public_bucket_spectrogram_path = ?, - - file_size_bytes = ?, - duration_millis = ? - "#, - inference_result_token, - job.model_token.clone(), - pretrained_vocoder_used.to_str(), - job.raw_inference_text.clone(), - text_hash, - normalized_inference_text, - - maybe_creator_user_token, - maybe_creator_synthetic_id, - - job.creator_ip_address.clone(), - job.creator_set_visibility.to_str(), - - bucket_audio_result_path, - bucket_spectrogram_result_path, - - file_size_bytes, - duration_millis - ) - .execute(&mut transaction) - .await; - - let record_id = match query_result { - Ok(res) => { - res.last_insert_id() - }, - Err(err) => { - // TODO: handle better - //transaction.rollback().await?; - return Err(anyhow!("Mysql error: {:?}", err)); - } - }; - - record_id - }; - - transaction.commit().await?; - - Ok((record_id, inference_result_token.clone())) -} - -pub struct TtsModelForInferenceRecord { - pub model_token: String, - pub tts_model_type: String, - - pub maybe_default_pretrained_vocoder: Option, - - pub creator_user_token: String, - pub creator_username: String, - pub creator_display_name: String, - - pub title: String, - pub private_bucket_hash: String, - pub created_at: DateTime, - pub updated_at: DateTime, - pub mod_deleted_at: Option>, - pub user_deleted_at: Option>, -} - - -#[derive(Clone, Debug)] -pub enum TtsModelForInferenceError { - ModelNotFound, - ModelDeleted, - DatabaseError { reason: String }, -} - -impl std::fmt::Display for TtsModelForInferenceError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - TtsModelForInferenceError::ModelNotFound => write!(f, "ModelNotFound"), - TtsModelForInferenceError::ModelDeleted => write!(f, "ModelDeleted"), - TtsModelForInferenceError::DatabaseError { reason} => - write!(f, "Database error: {:?}", reason), - } - } -} - -impl std::error::Error for TtsModelForInferenceError {} - -pub async fn get_tts_model_by_token( - pool: &MySqlPool, - model_token: &str -) -> Result -{ - // NB: Lookup failure is Err(RowNotFound). - // NB: Since this is publicly exposed, we don't query sensitive data. - let maybe_model = sqlx::query_as!( - TtsModelForInferenceRecord, - r#" -SELECT - tts.token as model_token, - tts.tts_model_type, - tts.maybe_default_pretrained_vocoder, - tts.creator_user_token, - users.username as creator_username, - users.display_name as creator_display_name, - tts.title, - tts.private_bucket_hash, - tts.created_at, - tts.updated_at, - tts.user_deleted_at, - tts.mod_deleted_at -FROM tts_models as tts -JOIN users -ON users.token = tts.creator_user_token -WHERE tts.token = ? - "#, - &model_token - ) - .fetch_one(pool) - .await; // TODO: This will return error if it doesn't exist - - let model : TtsModelForInferenceRecord = match maybe_model { - Ok(model) => model, - Err(err) => { - match err { - sqlx::Error::RowNotFound => { - return Err(TtsModelForInferenceError::ModelNotFound); - }, - _ => { - warn!("tts model query error: {:?}", err); - return Err(TtsModelForInferenceError::DatabaseError { - reason: format!("Mysql error: {:?}", err) - }); - } - } - } - }; - - if model.mod_deleted_at.is_some() || model.user_deleted_at.is_some() { - return Err(TtsModelForInferenceError::ModelDeleted); - } - - Ok(model) -} diff --git a/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs b/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs index c49a5128d1..fd8c8347ad 100644 --- a/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs +++ b/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs @@ -8,6 +8,9 @@ use container_common::anyhow_result::AnyhowResult; use database_queries::tokens::Tokens; use sqlx::MySqlPool; +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + /// table: w2l_template_upload_jobs #[derive(Debug)] pub struct W2lTemplateUploadJobRecord { diff --git a/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs b/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs index 11f698fa1c..3e4b2878fb 100644 --- a/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs +++ b/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs @@ -11,6 +11,9 @@ use log::{warn, info}; use sqlx::MySqlPool; use std::path::Path; +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + /// table: w2l_inference_jobs #[derive(Debug)] pub struct W2lInferenceJobRecord { diff --git a/crates/service/storyteller_web/src/tts_inference_job.rs b/crates/service/storyteller_web/src/tts_inference_job.rs index 2875a043c0..68d92f37c4 100644 --- a/crates/service/storyteller_web/src/tts_inference_job.rs +++ b/crates/service/storyteller_web/src/tts_inference_job.rs @@ -24,19 +24,8 @@ use config::shared_constants::DEFAULT_RUST_LOG; use container_common::anyhow_result::AnyhowResult; use container_common::token::random_uuid::generate_random_uuid; use crate::common_env::CommonEnv; -use crate::database::enums::vocoder_type::VocoderType; use crate::database::mediators::firehose_publisher::FirehosePublisher; use crate::http_clients::tts_inference_sidecar_client::TtsInferenceSidecarClient; -use crate::job_queries::tts_inference_job_queries::TtsInferenceJobRecord; -use crate::job_queries::tts_inference_job_queries::TtsModelForInferenceError; -use crate::job_queries::tts_inference_job_queries::TtsModelForInferenceRecord; -use crate::job_queries::tts_inference_job_queries::get_tts_model_by_token; -use crate::job_queries::tts_inference_job_queries::grab_job_lock_and_mark_pending; -use crate::job_queries::tts_inference_job_queries::insert_tts_result; -use crate::job_queries::tts_inference_job_queries::mark_tts_inference_job_done; -use crate::job_queries::tts_inference_job_queries::mark_tts_inference_job_failure; -use crate::job_queries::tts_inference_job_queries::mark_tts_inference_job_permanently_dead; -use crate::job_queries::tts_inference_job_queries::query_tts_inference_job_records; use crate::script_execution::tacotron_inference_command::TacotronInferenceCommand; use crate::util::buckets::bucket_client::BucketClient; use crate::util::buckets::bucket_path_unifier::BucketPathUnifier; @@ -55,6 +44,14 @@ use crate::util::noop_logger::NoOpLogger; use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger; use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir; use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; +use database_queries::column_types::vocoder_type::VocoderType; +use database_queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs}; +use database_queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done; +use database_queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure; +use database_queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock; +use database_queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead; +use database_queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord}; +use database_queries::tts::tts_results::insert_tts_result::insert_tts_result; use log::{warn, info, error}; use newrelic_telemetry::Client as NewRelicClient; use newrelic_telemetry::ClientBuilder; @@ -332,12 +329,10 @@ async fn main_loop(inferencer: Inferencer) { loop { let num_records = inferencer.job_batch_size; - let query_result = query_tts_inference_job_records( - &inferencer.mysql_pool, - num_records) - .await; + let maybe_available_jobs + = list_available_tts_inference_jobs(&inferencer.mysql_pool, num_records).await; - let jobs = match query_result { + let jobs = match maybe_available_jobs { Ok(jobs) => jobs, Err(e) => { warn!("Error querying jobs: {:?}", e); @@ -385,7 +380,7 @@ async fn main_loop(inferencer: Inferencer) { /// Process a batch of jobs, returning the count of cold-cache skipped jobs. async fn process_jobs( inferencer: &Inferencer, - jobs: Vec, + jobs: Vec, ) -> AnyhowResult> { let mut batch_spans = Vec::new(); @@ -423,7 +418,7 @@ async fn process_jobs( let _r = mark_tts_inference_job_permanently_dead( &inferencer.mysql_pool, - &job, + job.id, failure_reason ).await; } else { @@ -552,26 +547,18 @@ impl std::error::Error for ModelStateError {} impl ModelState { /// Query the model details and see if the model file is on the filesystem in one go. pub async fn query_model_and_check_filesystem( - job: &TtsInferenceJobRecord, + job: &AvailableTtsInferenceJob, mysql_pool: &MySqlPool, semi_persistent_cache: &SemiPersistentCacheDir, virtual_cache: &SyncVirtualLfuCache, ) -> Result { info!("Looking up TTS model by token: {}", &job.model_token); - let tts_model = get_tts_model_by_token( + let tts_model = get_tts_model_for_inference( &mysql_pool, &job.model_token ).await?; - /*let model_record = match query_result { - Ok(model) => model, - Err(e) => { - error!("TTS model lookup error: {}, reason: {:?}", &job.model_token, &e); - return Err(e) - }, - };*/ - let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path( &tts_model.model_token); @@ -595,7 +582,7 @@ impl ModelState { async fn process_job( inferencer: &Inferencer, - job: &TtsInferenceJobRecord, + job: &AvailableTtsInferenceJob, model_record: &TtsModelForInferenceRecord, ) -> AnyhowResult<(Span, Span)> { @@ -641,10 +628,11 @@ async fn process_job( // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // - let lock_acquired = grab_job_lock_and_mark_pending(&inferencer.mysql_pool, job).await?; + let lock_acquired = + mark_tts_inference_job_pending_and_grab_lock(&inferencer.mysql_pool, job.id).await?; if !lock_acquired { - warn!("Could not acquire job lock for: {}", &job.id); + warn!("Could not acquire job lock for: {:?}", &job.id); let duration = start.elapsed(); since_creation_span.set_attribute("status", "failure"); @@ -773,7 +761,7 @@ async fn process_job( // ==================== WRITE TEXT TO FILE ==================== // - let temp_dir = format!("temp_tts_inference_{}", job.id); + let temp_dir = format!("temp_tts_inference_{}", job.id.0); let temp_dir = TempDir::new(&temp_dir)?; // NB: Exists until it goes out of scope. let text_input_fs_path = temp_dir.path().join("inference_input.txt"); @@ -905,7 +893,7 @@ async fn process_job( info!("Marking job complete..."); mark_tts_inference_job_done( &inferencer.mysql_pool, - job, + job.id, true, Some(&inference_result_token) ).await?; @@ -922,7 +910,7 @@ async fn process_job( redis_logger.log_status("done")?; - info!("Job {} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}", + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}", job.id, id); let duration = start.elapsed();