daniel543x/artcraft
1
0 Fork 0
mirror of https://github.com/storytold/artcraft.git synced 2026-11-03 17:08:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

generic inference job queries

Browse Source
This commit is contained in:
Brandon Thomas committed 2023-01-11 00:33:40 -05:00
1 parent 9f873c33e7
commit 8755e7a533
14 files changed
+1218 -811
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+3
View File
@@ -2541,6 +2541,8 @@ dependencies = [
"database_queries",
"dotenv",
"easyenv",
"enums",
"errors",
"hostname",
"hyper 0.14.14",
"jobs_common",
@@ -2555,6 +2557,7 @@ dependencies = [
"sqlx",
"subprocess",
"tempdir",
"tokens",
"tokio 1.16.1",
"tts_common",
]
Dockerfile
+1 -1
View File
@@ -9,7 +9,7 @@
FROM ubuntu:jammy as rust-base
# NB: This can be "stable" or another version.
ARG RUST_TOOLCHAIN="1.65.0"
ARG RUST_TOOLCHAIN="1.66.0"
WORKDIR /tmp
crates/_types/enums/src/workers/generic_inference_type.rs
+21 -1
View File
@@ -1,12 +1,14 @@
// TODO: Use macro-derived impls
use std::collections::BTreeSet;
/// Our "generic inference" pipeline supports a wide variety of ML models and other media.
/// Each type of inference is identified by the following enum variants.
/// 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, Hash, Ord, PartialOrd, Deserialize, Serialize, sqlx::Type)]
pub enum GenericInferenceType {
#[serde(rename = "text_to_speech")]
#[sqlx(rename = "text_to_speech")]
@@ -33,6 +35,15 @@ impl GenericInferenceType {
_ => Err(format!("invalid value: {:?}", value)),
}
}
pub fn all_variants() -> BTreeSet<GenericInferenceType> {
// NB: BTreeSet is sorted
// NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc.
BTreeSet::from([
GenericInferenceType::TextToSpeech,
GenericInferenceType::VoiceConversion,
])
}
}
#[cfg(test)]
@@ -57,4 +68,13 @@ mod tests {
assert_eq!(GenericInferenceType::from_str("text_to_speech").unwrap(), GenericInferenceType::TextToSpeech);
assert_eq!(GenericInferenceType::from_str("voice_conversion").unwrap(), GenericInferenceType::VoiceConversion);
}
#[test]
fn all_variants() {
let mut variants = GenericInferenceType::all_variants();
assert_eq!(variants.len(), 2);
assert_eq!(variants.pop_first(), Some(GenericInferenceType::TextToSpeech));
assert_eq!(variants.pop_first(), Some(GenericInferenceType::VoiceConversion));
assert_eq!(variants.pop_first(), None);
}
}
crates/lib/database_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs
+181 -46
View File
@@ -6,7 +6,10 @@ use crate::helpers::boolean_converters::i8_to_bool;
use crate::queries::generic_inference::job::_keys::GenericInferenceJobId;
use enums::core::visibility::Visibility;
use enums::workers::generic_inference_type::GenericInferenceType;
use sqlx::MySqlPool;
use sqlx::mysql::MySqlArguments;
use sqlx::{MySql, MySqlPool};
use std::collections::BTreeSet;
use std::future::Future;
use std::path::Path;
use tokens::jobs::inference::InferenceJobToken;
@@ -41,56 +44,33 @@ pub struct AvailableInferenceJob {
pub retry_at: Option<chrono::DateTime<Utc>>,
}
pub async fn list_available_generic_inference_jobs(pool: &MySqlPool, num_records: u32)
pub struct ListAvailableGenericInferenceJobArgs<'a> {
pub num_records: u32,
pub is_debug_worker: bool,
pub sort_by_priority: bool,
pub maybe_scope_by_job_type: Option<BTreeSet<GenericInferenceType>>,
pub mysql_pool: &'a MySqlPool,
}
pub async fn list_available_generic_inference_jobs(
args: ListAvailableGenericInferenceJobArgs<'_>,
)
-> AnyhowResult<Vec<AvailableInferenceJob>>
{
let job_records = sqlx::query_as!(
AvailableInferenceJobRawInternal,
r#"
SELECT
id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`,
token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`,
let ALL_TYPES = GenericInferenceType::all_variants();
inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`,
maybe_inference_args,
maybe_raw_inference_text,
maybe_model_token,
let inference_types = args.maybe_scope_by_job_type
.as_ref()
.map(|types| types.clone())
.unwrap_or(ALL_TYPES);
maybe_creator_user_token,
creator_ip_address,
creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`,
let query = if args.sort_by_priority {
list_sorted_by_priority(args, inference_types).await
} else {
list_sorted_by_id(args, inference_types).await
};
status as `status: crate::column_types::job_status::JobStatus`,
attempt_count,
priority_level,
is_from_premium_user,
is_from_api_user,
is_for_twitch,
is_debug_request,
created_at,
updated_at,
retry_at
FROM generic_inference_jobs
WHERE
(
status IN ("pending", "attempt_failed")
)
AND
(
retry_at IS NULL
OR
retry_at < CURRENT_TIMESTAMP
)
ORDER BY id ASC
LIMIT ?
"#,
num_records,
)
.fetch_all(pool)
.await?;
let job_records = query?;
let job_records = job_records.into_iter()
.map(|record : AvailableInferenceJobRawInternal| {
@@ -121,6 +101,124 @@ WHERE
Ok(job_records)
}
async fn list_sorted_by_id(args: ListAvailableGenericInferenceJobArgs<'_>, inference_types: BTreeSet<GenericInferenceType>) -> Result<Vec<AvailableInferenceJobRawInternal>, sqlx::Error> {
sqlx::query_as!(
AvailableInferenceJobRawInternal,
r#"
SELECT
id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`,
token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`,
inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`,
maybe_inference_args,
maybe_raw_inference_text,
maybe_model_token,
maybe_creator_user_token,
creator_ip_address,
creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`,
status as `status: crate::column_types::job_status::JobStatus`,
attempt_count,
priority_level,
is_from_premium_user,
is_from_api_user,
is_for_twitch,
is_debug_request,
created_at,
updated_at,
retry_at
FROM generic_inference_jobs
WHERE
(
inference_type IN (?)
)
AND
(
status IN ("pending", "attempt_failed")
)
AND
(
retry_at IS NULL
OR
retry_at < CURRENT_TIMESTAMP
)
AND
(
is_debug_request = ?
)
ORDER BY id ASC
LIMIT ?
"#,
to_where_in_predicate(&inference_types),
args.is_debug_worker,
args.num_records,
)
.fetch_all(args.mysql_pool)
.await
}
async fn list_sorted_by_priority(args: ListAvailableGenericInferenceJobArgs<'_>, inference_types: BTreeSet<GenericInferenceType>) -> Result<Vec<AvailableInferenceJobRawInternal>, sqlx::Error> {
sqlx::query_as!(
AvailableInferenceJobRawInternal,
r#"
SELECT
id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`,
token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`,
inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`,
maybe_inference_args,
maybe_raw_inference_text,
maybe_model_token,
maybe_creator_user_token,
creator_ip_address,
creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`,
status as `status: crate::column_types::job_status::JobStatus`,
attempt_count,
priority_level,
is_from_premium_user,
is_from_api_user,
is_for_twitch,
is_debug_request,
created_at,
updated_at,
retry_at
FROM generic_inference_jobs
WHERE
(
inference_type IN (?)
)
AND
(
status IN ("pending", "attempt_failed")
)
AND
(
retry_at IS NULL
OR
retry_at < CURRENT_TIMESTAMP
)
AND
(
is_debug_request = ?
)
ORDER BY priority_level DESC, id ASC
LIMIT ?
"#,
to_where_in_predicate(&inference_types),
args.is_debug_worker,
args.num_records,
)
.fetch_all(args.mysql_pool)
.await
}
#[derive(Debug)]
struct AvailableInferenceJobRawInternal {
pub id: GenericInferenceJobId,
@@ -150,3 +248,40 @@ struct AvailableInferenceJobRawInternal {
pub updated_at: chrono::DateTime<Utc>,
pub retry_at: Option<chrono::DateTime<Utc>>,
}
/// 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(types: &BTreeSet<GenericInferenceType>) -> String {
let mut vec = types.iter()
.map(|ty| ty.to_str()) // TODO: These strings might need to be quoted.
.collect::<Vec<&'static str>>();
vec.sort(); // NB: For the benefit of tests.
vec.join(", ")
}
#[cfg(test)]
mod tests {
use crate::queries::generic_inference::job::list_available_generic_inference_jobs::to_where_in_predicate;
use enums::workers::generic_inference_type::GenericInferenceType;
use std::collections::BTreeSet;
#[test]
fn test_where_in_clause() {
// None
let types = BTreeSet::from([]);
assert_eq!(to_where_in_predicate(&types), "".to_string());
// Some
let types = BTreeSet::from([
GenericInferenceType::VoiceConversion,
]);
assert_eq!(to_where_in_predicate(&types), "voice_conversion".to_string());
// All
let types = BTreeSet::from([
GenericInferenceType::TextToSpeech,
GenericInferenceType::VoiceConversion,
]);
assert_eq!(to_where_in_predicate(&types), "text_to_speech, voice_conversion".to_string());
}
}
crates/lib/database_queries/src/queries/generic_inference/job/mod.rs
+2 -2
View File
@@ -1,6 +1,6 @@
pub mod _keys;
pub mod list_available_generic_inference_jobs;
pub mod mark_generic_inference_job_completely_failed;
pub mod mark_generic_inference_job_successfully_done;
pub mod mark_generic_inference_job_failure;
pub mod mark_generic_inference_job_pending_and_grab_lock;
pub mod mark_generic_inference_job_successfully_done;
pub mod list_available_generic_inference_jobs;
crates/service/job/download_job/src/job_steps/main_loop.rs
+2 -1
View File
@@ -7,6 +7,7 @@ use jobs_common::noop_logger::NoOpLogger;
use log::warn;
use std::time::Duration;
// Job runner timeouts (guards MySQL)
const START_TIMEOUT_MILLIS : u64 = 500;
const INCREASE_TIMEOUT_MILLIS : u64 = 1000;
@@ -41,7 +42,7 @@ pub async fn main_loop(job_state: JobState) {
match result {
Ok(_) => {},
Err(e) => {
warn!("Error querying jobs: {:?}", e);
warn!("Error processing jobs: {:?}", e);
std::thread::sleep(Duration::from_millis(error_timeout_millis));
error_timeout_millis += INCREASE_TIMEOUT_MILLIS;
continue;
crates/service/job/inference_job/Cargo.toml
+3
View File
@@ -19,8 +19,11 @@ cloud_storage = { path = "../../../lib/cloud_storage" }
config = { path = "../../../lib/config" }
container_common = { path = "../../../lib/container_common" }
database_queries = { path = "../../../lib/database_queries" }
enums = { path = "../../../_types/enums" }
errors = { path = "../../../lib/errors" }
jobs_common = { path = "../../../lib/jobs_common" }
memory_caching = { path = "../../../lib/memory_caching" }
tokens = { path = "../../../_types/tokens" }
tts_common = { path = "../../../lib/tts_common" }
# Vendored code
crates/service/job/inference_job/src/job_steps/main_loop.rs
+34 -357
View File
@@ -1,67 +1,38 @@
use container_common::anyhow_result::AnyhowResult;
use container_common::collections::multiple_random_from_vec::multiple_random_from_vec;
use crate::caching::cache_miss_strategizer::CacheMissStrategy;
use crate::caching::virtual_lfu_cache::SyncVirtualLfuCache;
use crate::job_steps::job_dependencies::JobDependencies;
use crate::job_steps::job_stats::JobStats;
use crate::job_steps::process_single_job::process_single_job;
use crate::job_steps::process_single_job_error::ProcessSingleJobError;
use crate::job_steps::sidecar_health_check_trap::maybe_block_on_sidecar_health_check;
use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs, list_available_tts_inference_jobs_with_minimum_priority};
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead;
use database_queries::queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord};
use crate::job_steps::process_single_job_old::process_single_job_old;
use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::{AvailableInferenceJob, list_available_generic_inference_jobs, ListAvailableGenericInferenceJobArgs};
use database_queries::queries::generic_inference::job::mark_generic_inference_job_failure::mark_generic_inference_job_failure;
use errors::AnyhowResult;
use jobs_common::noop_logger::NoOpLogger;
use jobs_common::semi_persistent_cache_dir::SemiPersistentCacheDir;
use log::{error, info, warn};
use memory_caching::multi_item_ttl_cache::MultiItemTtlCache;
use newrelic_telemetry::Span;
use sqlx::MySqlPool;
use std::time::Duration;
use crate::job_steps::process_single_job::process_single_job;
// Job runner timeouts (guards MySQL)
const START_TIMEOUT_MILLIS : u64 = 500;
const INCREASE_TIMEOUT_MILLIS : u64 = 1000;
pub async fn main_loop(job_args: JobDependencies) {
pub async fn main_loop(job_dependencies: JobDependencies) {
let mut noop_logger = NoOpLogger::new(job_dependencies.no_op_logger_millis as i64);
let mut error_timeout_millis = START_TIMEOUT_MILLIS;
let mut noop_logger = NoOpLogger::new(job_args.no_op_logger_millis as i64);
let mut span_batch = Vec::new();
let mut sort_by_priority = true;
let mut sort_by_priority_count = 0;
let mut needs_health_check_at_start = true; // Run health check at startup.
loop {
let num_records = job_args.job_batch_size;
// Don't completely starve low-priority jobs
if sort_by_priority_count >= job_args.low_priority_starvation_prevention_every_nth {
if sort_by_priority_count >= job_dependencies.low_priority_starvation_prevention_every_nth {
sort_by_priority_count = 0;
sort_by_priority = false;
}
let maybe_available_jobs =
if let Some(minimum_priority) = job_args.maybe_minimum_priority {
// Special high-priority workers
list_available_tts_inference_jobs_with_minimum_priority(
&job_args.mysql_pool,
minimum_priority,
num_records,
job_args.worker_details.is_debug_worker
).await
} else {
// Normal path
list_available_tts_inference_jobs(
&job_args.mysql_pool,
sort_by_priority,
num_records,
job_args.worker_details.is_debug_worker
).await
};
let maybe_available_jobs = list_available_generic_inference_jobs(ListAvailableGenericInferenceJobArgs {
num_records: job_dependencies.job_batch_size,
is_debug_worker: false, // TODO
sort_by_priority,
maybe_scope_by_job_type: None, // TODO
mysql_pool: &job_dependencies.mysql_pool,
}).await;
sort_by_priority = true;
sort_by_priority_count += 1;
@@ -69,7 +40,7 @@ pub async fn main_loop(job_args: JobDependencies) {
let jobs = match maybe_available_jobs {
Ok(jobs) => jobs,
Err(e) => {
warn!("Error querying jobs: {:?}", e);
error!("Error querying jobs: {:?}", e);
std::thread::sleep(Duration::from_millis(error_timeout_millis));
error_timeout_millis += INCREASE_TIMEOUT_MILLIS;
continue;
@@ -77,346 +48,52 @@ pub async fn main_loop(job_args: JobDependencies) {
};
if jobs.is_empty() {
noop_logger.log_message_after_awhile("No TTS jobs picked up from database!");
std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis));
noop_logger.log_message_after_awhile("No jobs picked up from database!");
std::thread::sleep(Duration::from_millis(job_dependencies.job_batch_wait_millis));
error_timeout_millis = START_TIMEOUT_MILLIS; // reset
continue;
}
info!("Queried {} jobs from database", jobs.len());
let batch_result = process_jobs(
&job_args,
jobs,
needs_health_check_at_start,
).await;
let batch_result = process_job_batch(&job_dependencies, jobs).await;
if needs_health_check_at_start {
needs_health_check_at_start = false;
}
let mut spans = match batch_result {
Ok(spans) => spans,
match batch_result {
Ok(_) => {},
Err(e) => {
warn!("Error running job batch: {:?}", e);
error!("Error processing jobs: {:?}", e);
std::thread::sleep(Duration::from_millis(error_timeout_millis));
error_timeout_millis += INCREASE_TIMEOUT_MILLIS;
continue;
}
};
if !job_args.newrelic_disabled {
span_batch.append(&mut spans);
if span_batch.len() > 50 {
let spans_to_send = span_batch.split_off(0).into();
job_args.newrelic_client.send_spans(spans_to_send).await;
}
}
error_timeout_millis = START_TIMEOUT_MILLIS; // reset
std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis));
std::thread::sleep(Duration::from_millis(job_dependencies.job_batch_wait_millis));
}
}
/// Process a batch of jobs, returning the count of cold-cache skipped jobs.
async fn process_jobs(
inferencer: &JobDependencies,
jobs: Vec<AvailableTtsInferenceJob>,
needs_health_check_at_start: bool,
) -> AnyhowResult<Vec<Span>> {
if needs_health_check_at_start {
maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await;
}
let mut maybe_sidecar_health_issue = false;
let mut batch_spans = Vec::new();
// TODO: A common interface/trait for each submodule (tts, webvc) to declare how to determine if the job is "ready".
// This probably returns a struct or enum with some measure of how many GB need to be downloaded.
async fn process_job_batch(job_state: &JobState, jobs: Vec<AvailableInferenceJob>) -> AnyhowResult<()> {
for job in jobs.into_iter() {
let model_state_result = ModelState::query_model_and_check_filesystem(
&job,
&inferencer.mysql_pool,
&inferencer.caches.tts_model_record_cache,
&inferencer.semi_persistent_cache,
&inferencer.virtual_model_lfu,
).await;
let model_state = match model_state_result {
Ok(model_state) => model_state,
Err(e) => {
error!("TTS model fetch and state check error: {}, reason: {:?}", &job.model_token, &e);
let (failure_reason, permanent_failure) = match e {
ModelStateError::ModelNotFound => ("model was not found", true),
ModelStateError::ModelDeleted => ("model has been deleted", true),
ModelStateError::CacheError { .. } => ("internal cache error", false),
ModelStateError::DatabaseError { .. } => ("unknown database error", false),
};
let internal_debugging_failure_reason = format!("model error: {:?}", e);
let mut job_progress_reporter = inferencer
.job_progress_reporter
.new_tts_inference(&job.inference_job_token)?;
job_progress_reporter.log_status(failure_reason)?;
if permanent_failure {
warn!("Marking job permanently dead: {} because: {:?}", job.inference_job_token, &e);
let _r = mark_tts_inference_job_permanently_dead(
&inferencer.mysql_pool,
job.id,
failure_reason,
&internal_debugging_failure_reason,
&inferencer.get_worker_name(),
).await;
} else {
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
).await;
}
continue;
}
};
if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache {
warn!("Model isn't ready: {} (downloaded = {}), (in memory = {})",
&model_state.model_record.model_token,
model_state.is_downloaded_to_filesystem,
model_state.is_in_memory_cache);
let maybe_strategy = if !model_state.is_downloaded_to_filesystem {
inferencer.cache_miss_strategizers.disk_cache_miss(&model_state.model_record.model_token)
} else {
inferencer.cache_miss_strategizers.memory_cache_miss(&model_state.model_record.model_token)
};
match maybe_strategy {
Err(err) => {
warn!("Unable to process job: {:?}", err);
let failure_reason = "cache error";
let internal_debugging_failure_reason = format!("cache error: {:?}", err);
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
).await;
continue;
},
Ok(CacheMissStrategy::WaitOrSkip) => {
// We're going to skip this for now.
// Maybe another worker has a warm cache and can continue.
warn!("Skipping TTS due to cold cache: {} ({})",
model_state.model_record.model_token,
model_state.model_record.title);
continue;
},
Ok(CacheMissStrategy::Proceed) => {}, // We're going to go ahead...
}
}
if maybe_sidecar_health_issue {
// Since we'll have a signal of the sidecar's health potentially being an issue, we don't
// need to background health check it from another thread. Instead we can react to the
// "potentially down" signal and block until it alleviates.
maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await;
maybe_sidecar_health_issue = false;
}
let result = process_single_job(inferencer, &job, &model_state.model_record).await;
let result = process_single_job(job_state, &job).await;
match result {
Ok((span1, span2)) => {
batch_spans.push(span1);
batch_spans.push(span2);
},
Ok(_) => {},
Err(e) => {
warn!("Failure to process job: {:?}", e);
record_failure_and_maybe_slow_down(&inferencer.job_stats);
maybe_sidecar_health_issue = true;
let failure_reason = "failure processing job";
let internal_debugging_failure_reason = format!("job error: {:?}", e);
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
let failure_reason = "";
let _r = mark_generic_inference_job_failure(
&job_state.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
job_state.job_max_attempts
).await;
match e {
ProcessSingleJobError::Other(_) => {} // No-op
ProcessSingleJobError::FilesystemFull => {
// TODO: Refactor - we should stop processing all of these jobs as we'll lose out
// on this entire batch by attempting to clear the filesystem. This should be handled
// in the calling code.
delete_tts_synthesizers_from_cache(&inferencer.semi_persistent_cache)?;
}
}
}
}
}
Ok(batch_spans)
}
fn record_failure_and_maybe_slow_down(job_stats: &JobStats) {
let stats = match job_stats.increment_failure_count() {
Ok(stats) => stats,
Err(e) => {
warn!("Error recording stats and reacting to repeated failures: {:?}", e);
return; // Can't really do anything.
}
};
let seconds_timeout = match stats.consecutive_failure_count {
t if t > 100 => 180,
t if t > 50 => 60,
t if t > 20 => 30,
t if t > 10 => 10,
t if t > 5 => 5,
_ => return, // No timeout
};
info!("Slowing down {} seconds due to significant repeated failures: {:?}",
seconds_timeout,
stats);
std::thread::sleep(Duration::from_secs(seconds_timeout));
}
/// Hack to delete locally cached TTS synthesizers to free up space from a full filesystem.
/// This is not intelligent and doesn't use any LRU/LFU mechanic.
/// This also relies on files not being read or written by concurrent workers while deleting.
fn delete_tts_synthesizers_from_cache(cache_dir: &SemiPersistentCacheDir) -> AnyhowResult<()> {
warn!("Deleting cached TTS synthesizers to free up disk space.");
let tts_synthesizer_dir = cache_dir.tts_synthesizer_model_directory();
// TODO: When this is no longer sufficient, delete other types of locally-cached data.
let paths = std::fs::read_dir(tts_synthesizer_dir)?
.map(|res| res.map(|e| e.path()))
.collect::<Result<Vec<_>, std::io::Error>>()?;
let models_to_delete = multiple_random_from_vec(&paths, 35);
for model_to_delete in models_to_delete {
warn!("Deleting cached model file: {:?}", model_to_delete);
let full_model_path = cache_dir.tts_synthesizer_model_path(model_to_delete);
std::fs::remove_file(full_model_path)?;
}
Ok(())
}
/// We check both of these in one go so that we can reuse the ModelRecord later
/// without another DB query.
struct ModelState {
pub model_record: TtsModelForInferenceRecord,
pub is_downloaded_to_filesystem: bool,
pub is_in_memory_cache: bool,
}
#[derive(Debug, Clone)]
enum ModelStateError {
ModelNotFound,
ModelDeleted,
CacheError { reason: String },
DatabaseError { reason: String },
}
impl std::fmt::Display for ModelStateError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ModelStateError::ModelNotFound => write!(f, "ModelNotFound"),
ModelStateError::ModelDeleted => write!(f, "ModelDeleted"),
ModelStateError::CacheError { reason} => write!(f, "Cache error: {:?}", reason),
ModelStateError::DatabaseError { reason} => write!(f, "Database error: {:?}", reason),
}
}
}
impl From<TtsModelForInferenceError> for ModelStateError {
fn from(error: TtsModelForInferenceError) -> Self {
match error {
TtsModelForInferenceError::ModelNotFound => ModelStateError::ModelNotFound,
TtsModelForInferenceError::ModelDeleted => ModelStateError::ModelDeleted,
TtsModelForInferenceError::DatabaseError { reason } => ModelStateError::DatabaseError { reason }
}
}
}
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: &AvailableTtsInferenceJob,
mysql_pool: &MySqlPool,
tts_model_record_cache: &MultiItemTtlCache<String, TtsModelForInferenceRecord>,
semi_persistent_cache: &SemiPersistentCacheDir,
virtual_cache: &SyncVirtualLfuCache,
) -> Result<Self, ModelStateError> {
// Many workers will be querying models constantly (n-many per batch).
// We can save on a lot of DB query volume by caching model records.
let maybe_cached_tts_model =
tts_model_record_cache.copy_without_bump_if_unexpired(job.model_token.clone())
.ok()
.flatten();
let tts_model = match maybe_cached_tts_model {
Some(tts_model) => tts_model,
None => {
info!("Looking up TTS model record by token: {}", &job.model_token);
let tts_model = get_tts_model_for_inference(
&mysql_pool,
&job.model_token
).await?;
tts_model_record_cache.store_copy(&job.model_token, &tts_model).ok();
tts_model
}
};
let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path(
&tts_model.model_token);
let is_downloaded_to_filesystem = tts_synthesizer_fs_path.exists();
let path = tts_synthesizer_fs_path
.to_str()
.ok_or(ModelStateError::CacheError { reason: "could not make path".to_string() })?
.to_string();
let is_in_memory_cache = virtual_cache.in_cache(&path)
.map_err(|e| ModelStateError::CacheError { reason: format!("Model cache error: {:?}", e) })?;
Ok(Self {
model_record: tts_model,
is_downloaded_to_filesystem,
is_in_memory_cache,
})
}
}
crates/service/job/inference_job/src/job_steps/main_loop_old.rs
+422
View File
@@ -0,0 +1,422 @@
use container_common::anyhow_result::AnyhowResult;
use container_common::collections::multiple_random_from_vec::multiple_random_from_vec;
use crate::caching::cache_miss_strategizer::CacheMissStrategy;
use crate::caching::virtual_lfu_cache::SyncVirtualLfuCache;
use crate::job_steps::job_dependencies::JobDependencies;
use crate::job_steps::job_stats::JobStats;
use crate::job_steps::process_single_job_old::process_single_job_old;
use crate::job_steps::process_single_job_error::ProcessSingleJobError;
use crate::job_steps::sidecar_health_check_trap::maybe_block_on_sidecar_health_check;
use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs, list_available_tts_inference_jobs_with_minimum_priority};
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead;
use database_queries::queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord};
use jobs_common::noop_logger::NoOpLogger;
use jobs_common::semi_persistent_cache_dir::SemiPersistentCacheDir;
use log::{error, info, warn};
use memory_caching::multi_item_ttl_cache::MultiItemTtlCache;
use newrelic_telemetry::Span;
use sqlx::MySqlPool;
use std::time::Duration;
// Job runner timeouts (guards MySQL)
const START_TIMEOUT_MILLIS : u64 = 500;
const INCREASE_TIMEOUT_MILLIS : u64 = 1000;
pub async fn main_loop_old(job_args: JobDependencies) {
let mut error_timeout_millis = START_TIMEOUT_MILLIS;
let mut noop_logger = NoOpLogger::new(job_args.no_op_logger_millis as i64);
let mut span_batch = Vec::new();
let mut sort_by_priority = true;
let mut sort_by_priority_count = 0;
let mut needs_health_check_at_start = true; // Run health check at startup.
loop {
let num_records = job_args.job_batch_size;
// Don't completely starve low-priority jobs
if sort_by_priority_count >= job_args.low_priority_starvation_prevention_every_nth {
sort_by_priority_count = 0;
sort_by_priority = false;
}
let maybe_available_jobs =
if let Some(minimum_priority) = job_args.maybe_minimum_priority {
// Special high-priority workers
list_available_tts_inference_jobs_with_minimum_priority(
&job_args.mysql_pool,
minimum_priority,
num_records,
job_args.worker_details.is_debug_worker
).await
} else {
// Normal path
list_available_tts_inference_jobs(
&job_args.mysql_pool,
sort_by_priority,
num_records,
job_args.worker_details.is_debug_worker
).await
};
sort_by_priority = true;
sort_by_priority_count += 1;
let jobs = match maybe_available_jobs {
Ok(jobs) => jobs,
Err(e) => {
warn!("Error querying jobs: {:?}", e);
std::thread::sleep(Duration::from_millis(error_timeout_millis));
error_timeout_millis += INCREASE_TIMEOUT_MILLIS;
continue;
}
};
if jobs.is_empty() {
noop_logger.log_message_after_awhile("No inference jobs picked up from database!");
std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis));
continue;
}
info!("Queried {} jobs from database", jobs.len());
let batch_result = process_jobs(
&job_args,
jobs,
needs_health_check_at_start,
).await;
if needs_health_check_at_start {
needs_health_check_at_start = false;
}
let mut spans = match batch_result {
Ok(spans) => spans,
Err(e) => {
warn!("Error running job batch: {:?}", e);
std::thread::sleep(Duration::from_millis(error_timeout_millis));
error_timeout_millis += INCREASE_TIMEOUT_MILLIS;
continue;
}
};
if !job_args.newrelic_disabled {
span_batch.append(&mut spans);
if span_batch.len() > 50 {
let spans_to_send = span_batch.split_off(0).into();
job_args.newrelic_client.send_spans(spans_to_send).await;
}
}
error_timeout_millis = START_TIMEOUT_MILLIS; // reset
std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis));
}
}
/// Process a batch of jobs, returning the count of cold-cache skipped jobs.
async fn process_jobs(
inferencer: &JobDependencies,
jobs: Vec<AvailableTtsInferenceJob>,
needs_health_check_at_start: bool,
) -> AnyhowResult<Vec<Span>> {
if needs_health_check_at_start {
maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await;
}
let mut maybe_sidecar_health_issue = false;
let mut batch_spans = Vec::new();
for job in jobs.into_iter() {
let model_state_result = ModelState::query_model_and_check_filesystem(
&job,
&inferencer.mysql_pool,
&inferencer.caches.tts_model_record_cache,
&inferencer.semi_persistent_cache,
&inferencer.virtual_model_lfu,
).await;
let model_state = match model_state_result {
Ok(model_state) => model_state,
Err(e) => {
error!("TTS model fetch and state check error: {}, reason: {:?}", &job.model_token, &e);
let (failure_reason, permanent_failure) = match e {
ModelStateError::ModelNotFound => ("model was not found", true),
ModelStateError::ModelDeleted => ("model has been deleted", true),
ModelStateError::CacheError { .. } => ("internal cache error", false),
ModelStateError::DatabaseError { .. } => ("unknown database error", false),
};
let internal_debugging_failure_reason = format!("model error: {:?}", e);
let mut job_progress_reporter = inferencer
.job_progress_reporter
.new_tts_inference(&job.inference_job_token)?;
job_progress_reporter.log_status(failure_reason)?;
if permanent_failure {
warn!("Marking job permanently dead: {} because: {:?}", job.inference_job_token, &e);
let _r = mark_tts_inference_job_permanently_dead(
&inferencer.mysql_pool,
job.id,
failure_reason,
&internal_debugging_failure_reason,
&inferencer.get_worker_name(),
).await;
} else {
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
).await;
}
continue;
}
};
if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache {
warn!("Model isn't ready: {} (downloaded = {}), (in memory = {})",
&model_state.model_record.model_token,
model_state.is_downloaded_to_filesystem,
model_state.is_in_memory_cache);
let maybe_strategy = if !model_state.is_downloaded_to_filesystem {
inferencer.cache_miss_strategizers.disk_cache_miss(&model_state.model_record.model_token)
} else {
inferencer.cache_miss_strategizers.memory_cache_miss(&model_state.model_record.model_token)
};
match maybe_strategy {
Err(err) => {
warn!("Unable to process job: {:?}", err);
let failure_reason = "cache error";
let internal_debugging_failure_reason = format!("cache error: {:?}", err);
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
).await;
continue;
},
Ok(CacheMissStrategy::WaitOrSkip) => {
// We're going to skip this for now.
// Maybe another worker has a warm cache and can continue.
warn!("Skipping TTS due to cold cache: {} ({})",
model_state.model_record.model_token,
model_state.model_record.title);
continue;
},
Ok(CacheMissStrategy::Proceed) => {}, // We're going to go ahead...
}
}
if maybe_sidecar_health_issue {
// Since we'll have a signal of the sidecar's health potentially being an issue, we don't
// need to background health check it from another thread. Instead we can react to the
// "potentially down" signal and block until it alleviates.
maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await;
maybe_sidecar_health_issue = false;
}
let result = process_single_job_old(inferencer, &job, &model_state.model_record).await;
match result {
Ok((span1, span2)) => {
batch_spans.push(span1);
batch_spans.push(span2);
},
Err(e) => {
warn!("Failure to process job: {:?}", e);
record_failure_and_maybe_slow_down(&inferencer.job_stats);
maybe_sidecar_health_issue = true;
let failure_reason = "failure processing job";
let internal_debugging_failure_reason = format!("job error: {:?}", e);
let _r = mark_tts_inference_job_failure(
&inferencer.mysql_pool,
&job,
failure_reason,
&internal_debugging_failure_reason,
inferencer.job_max_attempts,
&inferencer.get_worker_name(),
).await;
match e {
ProcessSingleJobError::Other(_) => {} // No-op
ProcessSingleJobError::FilesystemFull => {
// TODO: Refactor - we should stop processing all of these jobs as we'll lose out
// on this entire batch by attempting to clear the filesystem. This should be handled
// in the calling code.
delete_tts_synthesizers_from_cache(&inferencer.semi_persistent_cache)?;
}
}
}
}
}
Ok(batch_spans)
}
fn record_failure_and_maybe_slow_down(job_stats: &JobStats) {
let stats = match job_stats.increment_failure_count() {
Ok(stats) => stats,
Err(e) => {
warn!("Error recording stats and reacting to repeated failures: {:?}", e);
return; // Can't really do anything.
}
};
let seconds_timeout = match stats.consecutive_failure_count {
t if t > 100 => 180,
t if t > 50 => 60,
t if t > 20 => 30,
t if t > 10 => 10,
t if t > 5 => 5,
_ => return, // No timeout
};
info!("Slowing down {} seconds due to significant repeated failures: {:?}",
seconds_timeout,
stats);
std::thread::sleep(Duration::from_secs(seconds_timeout));
}
/// Hack to delete locally cached TTS synthesizers to free up space from a full filesystem.
/// This is not intelligent and doesn't use any LRU/LFU mechanic.
/// This also relies on files not being read or written by concurrent workers while deleting.
fn delete_tts_synthesizers_from_cache(cache_dir: &SemiPersistentCacheDir) -> AnyhowResult<()> {
warn!("Deleting cached TTS synthesizers to free up disk space.");
let tts_synthesizer_dir = cache_dir.tts_synthesizer_model_directory();
// TODO: When this is no longer sufficient, delete other types of locally-cached data.
let paths = std::fs::read_dir(tts_synthesizer_dir)?
.map(|res| res.map(|e| e.path()))
.collect::<Result<Vec<_>, std::io::Error>>()?;
let models_to_delete = multiple_random_from_vec(&paths, 35);
for model_to_delete in models_to_delete {
warn!("Deleting cached model file: {:?}", model_to_delete);
let full_model_path = cache_dir.tts_synthesizer_model_path(model_to_delete);
std::fs::remove_file(full_model_path)?;
}
Ok(())
}
/// We check both of these in one go so that we can reuse the ModelRecord later
/// without another DB query.
struct ModelState {
pub model_record: TtsModelForInferenceRecord,
pub is_downloaded_to_filesystem: bool,
pub is_in_memory_cache: bool,
}
#[derive(Debug, Clone)]
enum ModelStateError {
ModelNotFound,
ModelDeleted,
CacheError { reason: String },
DatabaseError { reason: String },
}
impl std::fmt::Display for ModelStateError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ModelStateError::ModelNotFound => write!(f, "ModelNotFound"),
ModelStateError::ModelDeleted => write!(f, "ModelDeleted"),
ModelStateError::CacheError { reason} => write!(f, "Cache error: {:?}", reason),
ModelStateError::DatabaseError { reason} => write!(f, "Database error: {:?}", reason),
}
}
}
impl From<TtsModelForInferenceError> for ModelStateError {
fn from(error: TtsModelForInferenceError) -> Self {
match error {
TtsModelForInferenceError::ModelNotFound => ModelStateError::ModelNotFound,
TtsModelForInferenceError::ModelDeleted => ModelStateError::ModelDeleted,
TtsModelForInferenceError::DatabaseError { reason } => ModelStateError::DatabaseError { reason }
}
}
}
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: &AvailableTtsInferenceJob,
mysql_pool: &MySqlPool,
tts_model_record_cache: &MultiItemTtlCache<String, TtsModelForInferenceRecord>,
semi_persistent_cache: &SemiPersistentCacheDir,
virtual_cache: &SyncVirtualLfuCache,
) -> Result<Self, ModelStateError> {
// Many workers will be querying models constantly (n-many per batch).
// We can save on a lot of DB query volume by caching model records.
let maybe_cached_tts_model =
tts_model_record_cache.copy_without_bump_if_unexpired(job.model_token.clone())
.ok()
.flatten();
let tts_model = match maybe_cached_tts_model {
Some(tts_model) => tts_model,
None => {
info!("Looking up TTS model record by token: {}", &job.model_token);
let tts_model = get_tts_model_for_inference(
&mysql_pool,
&job.model_token
).await?;
tts_model_record_cache.store_copy(&job.model_token, &tts_model).ok();
tts_model
}
};
let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path(
&tts_model.model_token);
let is_downloaded_to_filesystem = tts_synthesizer_fs_path.exists();
let path = tts_synthesizer_fs_path
.to_str()
.ok_or(ModelStateError::CacheError { reason: "could not make path".to_string() })?
.to_string();
let is_in_memory_cache = virtual_cache.in_cache(&path)
.map_err(|e| ModelStateError::CacheError { reason: format!("Model cache error: {:?}", e) })?;
Ok(Self {
model_record: tts_model,
is_downloaded_to_filesystem,
is_in_memory_cache,
})
}
}
crates/service/job/inference_job/src/job_steps/mod.rs
+3
View File
@@ -1,6 +1,9 @@
pub mod job_dependencies;
pub mod job_stats;
pub mod main_loop;
pub mod main_loop_old;
pub mod process_single_job;
pub mod process_single_job_error;
pub mod process_single_job_old;
pub mod select_job_batch;
pub mod sidecar_health_check_trap;
crates/service/job/inference_job/src/job_steps/process_single_job.rs
+88 -401
View File
@@ -1,433 +1,120 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use anyhow::anyhow;
use container_common::anyhow_result::AnyhowResult;
use container_common::filesystem::check_file_exists::check_file_exists;
use config::is_bad_download_url::is_bad_download_url;
use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
use container_common::hashing::hash_string_sha2::hash_string_sha2;
use container_common::token::random_uuid::generate_random_uuid;
use crate::job_steps::job_dependencies::JobDependencies;
use crate::job_steps::process_single_job_error::ProcessSingleJobError;
use crate::job_types::tts::seconds_to_decoder_steps::seconds_to_decoder_steps;
use crate::util::download_file_from_bucket::maybe_download_file_from_bucket;
use database_queries::column_types::vocoder_type::VocoderType;
use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock;
use database_queries::queries::tts::tts_models::get_tts_model_for_inference::TtsModelForInferenceRecord;
use database_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result;
use log::{warn, info, error};
use newrelic_telemetry::Span;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH, Instant};
use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
use database_queries::queries::generic_inference::job::mark_generic_inference_job_pending_and_grab_lock::mark_generic_inference_job_pending_and_grab_lock;
use database_queries::queries::generic_inference::job::mark_generic_inference_job_successfully_done::mark_generic_inference_job_successfully_done;
use enums::workers::generic_inference_type::GenericInferenceType;
use errors::AnyhowResult;
use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
use log::{info, warn};
use tempdir::TempDir;
use tts_common::clean_symbols::clean_symbols;
use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic;
use tts_common::text_pipelines::text_pipeline_type::TextPipelineType;
#[derive(Deserialize, Default)]
struct FileMetadata {
pub duration_millis: Option<u64>,
pub mimetype: Option<String>,
pub file_size_bytes: u64,
}
fn read_metadata_file(filename: &PathBuf) -> AnyhowResult<FileMetadata> {
let mut file = File::open(filename)?;
let mut buffer = String::new();
file.read_to_string(&mut buffer)?;
Ok(serde_json::from_str(&buffer)?)
}
pub async fn process_single_job(
job_args: &JobDependencies,
job: &AvailableTtsInferenceJob,
model_record: &TtsModelForInferenceRecord,
) -> Result<(Span, Span), ProcessSingleJobError> {
let start = Instant::now();
let span_id = generate_random_uuid();
let trace_id = generate_random_uuid();
let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or("");
let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis())
.name("single job execution")
.attribute("user_token", maybe_user_token)
.service_name("tts-inference-job");
let span_id = generate_random_uuid();
let trace_id = generate_random_uuid();
let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000;
let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp)
.name("job since creation")
.attribute("user_token", maybe_user_token)
.service_name("tts-inference-job");
let mut job_progress_reporter = job_args
.job_progress_reporter
.new_tts_inference(&job.inference_job_token)
.map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
pub async fn process_single_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> AnyhowResult<()> {
let mut redis = job_dependencies.redis_pool.get()?;
let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, job.download_job_token.as_str());
// ==================== ATTEMPT TO GRAB JOB LOCK ==================== //
info!("Attempting to grab lock for job: {}", job.inference_job_token);
let lock_acquired =
mark_tts_inference_job_pending_and_grab_lock(&job_args.mysql_pool, job.id)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
let lock_acquired = mark_generic_inference_job_pending_and_grab_lock(&job_dependencies.mysql_pool, job.id).await?;
if !lock_acquired {
warn!("Could not acquire job lock for: {:?}", &job.id);
let duration = start.elapsed();
since_creation_span.set_attribute("status", "failure");
since_creation_span.set_duration(duration);
job_iteration_span.set_attribute("status", "failure");
job_iteration_span.set_duration(duration);
return Ok((since_creation_span, job_iteration_span));
warn!("Could not acquire job lock for: {}", &job.id.0);
return Ok(())
}
info!("Lock acquired for job: {}", job.inference_job_token);
info!("Beginning work on {:?} = {} | {}", job.download_type, job.download_job_token, job.download_url);
// ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== //
// ==================== SETUP TEMP DIRS ==================== //
let waveglow_vocoder_model_fs_path = {
let waveglow_vocoder_model_filename = job_args.waveglow_vocoder_model_filename.clone();
let waveglow_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename);
let waveglow_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename);
// TODO: Scoped temp dir
maybe_download_file_from_bucket(
"waveglow vocoder model",
&waveglow_vocoder_model_fs_path,
&waveglow_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (1 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
let temp_dir = format!("temp_{}", job.id.0);
let temp_dir = TempDir::new(&temp_dir)?;
waveglow_vocoder_model_fs_path
};
// ==================== HANDLE DIFFERENT INFERENCE TYPES ==================== //
// ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== //
let mut entity_token : Option<String> = None;
let mut entity_type : Option<String> = None;
let pretrained_hifigan_vocoder_model_fs_path = {
let hifigan_vocoder_model_filename = job_args.hifigan_vocoder_model_filename.clone();
let hifigan_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename);
let hifigan_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename);
maybe_download_file_from_bucket(
"hifigan vocoder model",
&hifigan_vocoder_model_fs_path,
&hifigan_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (2 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
hifigan_vocoder_model_fs_path
};
// ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== //
let hifigan_superres_vocoder_model_fs_path = {
let hifigan_superres_vocoder_model_filename = job_args.hifigan_superres_vocoder_model_filename.clone();
let hifigan_superres_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename);
let hifigan_superres_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename);
maybe_download_file_from_bucket(
"hifigan superres vocoder model",
&hifigan_superres_vocoder_model_fs_path,
&hifigan_superres_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (3 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
hifigan_superres_vocoder_model_fs_path
};
// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== //
let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder {
None => None,
Some(vocoder) => {
let custom_vocoder_fs_path = job_args.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token);
let custom_vocoder_object_path = job_args.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash);
maybe_download_file_from_bucket(
"custom vocoder",
&custom_vocoder_fs_path,
&custom_vocoder_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading user vocoder",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
Some(custom_vocoder_fs_path)
match job.inference_type {
GenericInferenceType::TextToSpeech => {
// TODO
}
};
// ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== //
let tts_synthesizer_fs_path = {
let tts_synthesizer_fs_path = job_args.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token);
let tts_synthesizer_object_path = job_args.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash);
maybe_download_file_from_bucket(
"synthesizer",
&tts_synthesizer_fs_path,
&tts_synthesizer_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading synthesizer",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
tts_synthesizer_fs_path
};
// ==================== Preprocess text ==================== //
let cleaned_inference_text = clean_symbols(&job.raw_inference_text);
// ==================== WRITE TEXT TO FILE ==================== //
info!("Creating tempdir for inference results.");
let temp_dir = format!("temp_tts_inference_{}", job.id.0);
// NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
let temp_dir = TempDir::new(&temp_dir)
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
let text_input_fs_path = temp_dir.path().join("inference_input.txt");
std::fs::write(&text_input_fs_path, &cleaned_inference_text)
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
// ==================== RUN INFERENCE ==================== //
job_progress_reporter.log_status("running inference")
.map_err(|e| ProcessSingleJobError::Other(e))?;
// TODO: Fix this.
let maybe_unload_model_path = job_args
.virtual_model_lfu
.insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or(""))
.map_err(|e| ProcessSingleJobError::Other(e))?;
if let Some(model_path) = maybe_unload_model_path.as_deref() {
warn!("Remove model from LFU cache: {:?}", model_path);
GenericInferenceType::VoiceConversion => {
// TODO
}
// GenericDownloadType::HifiGan => {
// let results = process_hifigan_vocoder(
// job_state,
// job,
// &temp_dir,
// &download_filename,
// &mut redis_logger,
// ).await?;
// entity_token = results.entity_token.clone();
// entity_type = results.entity_type.clone();
// }
// GenericDownloadType::HifiGanRocketVc => {
// let results = process_hifigan_softvc_vocoder(
// job_state,
// job,
// &temp_dir,
// &download_filename,
// &mut redis_logger,
// ).await?;
// entity_token = results.entity_token.clone();
// entity_type = results.entity_type.clone();
// }
// GenericDownloadType::RocketVc => {
// let results = process_softvc_model(
// job_state,
// job,
// &temp_dir,
// &download_filename,
// &mut redis_logger,
// ).await?;
// entity_token = results.entity_token.clone();
// entity_type = results.entity_type.clone();
// }
// GenericDownloadType::Tacotron2 => {
// let results = process_tacotron_model(
// job_state,
// job,
// &temp_dir,
// &download_filename,
// &mut redis_logger,
// ).await?;
// entity_token = results.entity_token.clone();
// entity_type = results.entity_type.clone();
// }
}
let output_audio_fs_path = temp_dir.path().join("output.wav");
let output_metadata_fs_path = temp_dir.path().join("metadata.json");
let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json");
info!("Running TTS inference...");
info!("Expected output audio filename: {:?}", &output_audio_fs_path);
info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
info!("Expected output metadata filename: {:?}", &output_metadata_fs_path);
if let Some(model_path) = maybe_unload_model_path.as_deref() {
warn!("Unload model from sidecar: {:?}", &model_path);
}
let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution;
if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() {
pretrained_vocoder = VocoderType::from_str(default_vocoder)
.map_err(|e| ProcessSingleJobError::Other(e))?;
}
info!("With pretrained vocoder: {:?}", pretrained_vocoder);
let text_pipeline_type_or_guess = model_record.text_pipeline_type
.as_deref()
.and_then(|pipeline_type|
TextPipelineType::from_str(pipeline_type).ok())// NB: If there's an error deserializing, turn it to None.
.unwrap_or_else(||
guess_text_pipeline_heuristic(Some(model_record.created_at)));
info!("With text pipeline type `{:?} ` (or guess: {:?})",
&model_record.text_pipeline_type,
&text_pipeline_type_or_guess);
let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path {
None => {
info!("using pretrained HiFi-GAN vocoder");
pretrained_hifigan_vocoder_model_fs_path
},
Some(custom_vocoder_fs_path) => {
info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path);
custom_vocoder_fs_path
},
};
// NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to
// roughly 12 seconds max. Here we map seconds to decoder steps.
let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds);
job_args.http_clients.tts_inference_sidecar_client.request_inference(
&cleaned_inference_text,
max_decoder_steps,
&tts_synthesizer_fs_path,
pretrained_vocoder,
&text_pipeline_type_or_guess.to_str(),
&hifigan_vocoder_model_fs_path_to_use,
&hifigan_superres_vocoder_model_fs_path,
&waveglow_vocoder_model_fs_path,
&output_audio_fs_path,
&output_spectrogram_fs_path,
&output_metadata_fs_path,
maybe_unload_model_path)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
// ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
info!("Checking that output files exist...");
check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
let file_metadata = read_metadata_file(&output_metadata_fs_path)
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_metadata_fs_path);
// ==================== UPLOAD AUDIO TO BUCKET ==================== //
job_progress_reporter.log_status("uploading result")
.map_err(|e| ProcessSingleJobError::Other(e))?;
let audio_result_object_path = job_args.bucket_path_unifier.tts_inference_wav_audio_output_path(
&job.uuid_idempotency_token); // TODO: Don't use this!
info!("Audio destination bucket path: {:?}", &audio_result_object_path);
info!("Uploading audio...");
job_args.public_bucket_client.upload_filename_with_content_type(
&audio_result_object_path,
&output_audio_fs_path,
"audio/wav")
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_audio_fs_path);
// ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== //
let spectrogram_result_object_path = job_args.bucket_path_unifier.tts_inference_spectrogram_output_path(
&job.uuid_idempotency_token); // TODO: Don't use this!
info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path);
info!("Uploading spectrogram...");
job_args.public_bucket_client.upload_filename_with_content_type(
&spectrogram_result_object_path,
&output_spectrogram_fs_path,
"application/json")
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_spectrogram_fs_path);
// ==================== DELETE DOWNLOADED FILE ==================== //
// NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
safe_delete_temp_directory(&temp_dir);
// ==================== SAVE RECORDS ==================== //
let text_hash = hash_string_sha2(&cleaned_inference_text)
.map_err(|e| ProcessSingleJobError::Other(e))?;
let worker_name = job_args.get_worker_name();
info!("Saving tts inference record...");
let (id, inference_result_token) = insert_tts_result(
&job_args.mysql_pool,
job,
&text_hash,
pretrained_vocoder,
&audio_result_object_path,
&spectrogram_result_object_path,
file_metadata.file_size_bytes,
file_metadata.duration_millis.unwrap_or(0),
job_args.worker_details.is_on_prem,
&worker_name,
job_args.worker_details.is_debug_worker)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
// =====================================================
info!("Marking job complete...");
mark_tts_inference_job_done(
&job_args.mysql_pool,
job.id,
true,
Some(&inference_result_token),
&worker_name)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
mark_generic_inference_job_successfully_done(
&job_dependencies.mysql_pool,
job,
entity_type.as_deref(),
entity_token.as_deref(),
).await?;
info!("TTS Done. Original text was: {}", &job.raw_inference_text);
info!("Saved model record: {} - {}", job.id.0, &job.inference_job_token);
job_args.firehose_publisher.tts_inference_finished(
job.maybe_creator_user_token.as_deref(),
&model_record.model_token,
&inference_result_token)
job_dependencies.firehose_publisher.publish_generic_download_finished(
&job.creator_user_token,
entity_token.as_deref())
.await
.map_err(|e| {
error!("error publishing event: {:?}", e);
ProcessSingleJobError::Other(anyhow!("error publishing event"))
warn!("error publishing event: {:?}", e);
anyhow!("error publishing event")
})?;
job_progress_reporter.log_status("done")
.map_err(|e| ProcessSingleJobError::Other(e))?;
redis_logger.log_status("done")?;
info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}",
job.id, id, &inference_result_token);
info!("Job done: {}", job.id.0);
let duration = start.elapsed();
since_creation_span.set_attribute("status", "success");
since_creation_span.set_duration(duration);
job_iteration_span.set_attribute("status", "success");
job_iteration_span.set_duration(duration);
Ok((since_creation_span, job_iteration_span))
}
fn get_timestamp_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
Ok(())
}
crates/service/job/inference_job/src/job_steps/process_single_job_old.rs
+433
View File
@@ -0,0 +1,433 @@
// NB: Incrementally getting rid of build warnings...
#![forbid(unused_imports)]
#![forbid(unused_mut)]
#![forbid(unused_variables)]
use anyhow::anyhow;
use container_common::filesystem::check_file_exists::check_file_exists;
use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
use container_common::hashing::hash_string_sha2::hash_string_sha2;
use container_common::token::random_uuid::generate_random_uuid;
use crate::job_steps::job_dependencies::JobDependencies;
use crate::job_steps::process_single_job_error::ProcessSingleJobError;
use crate::job_types::tts::seconds_to_decoder_steps::seconds_to_decoder_steps;
use crate::util::download_file_from_bucket::maybe_download_file_from_bucket;
use database_queries::column_types::vocoder_type::VocoderType;
use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done;
use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock;
use database_queries::queries::tts::tts_models::get_tts_model_for_inference::TtsModelForInferenceRecord;
use database_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result;
use errors::AnyhowResult;
use log::{warn, info, error};
use newrelic_telemetry::Span;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH, Instant};
use tempdir::TempDir;
use tts_common::clean_symbols::clean_symbols;
use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic;
use tts_common::text_pipelines::text_pipeline_type::TextPipelineType;
#[derive(Deserialize, Default)]
struct FileMetadata {
pub duration_millis: Option<u64>,
pub mimetype: Option<String>,
pub file_size_bytes: u64,
}
fn read_metadata_file(filename: &PathBuf) -> AnyhowResult<FileMetadata> {
let mut file = File::open(filename)?;
let mut buffer = String::new();
file.read_to_string(&mut buffer)?;
Ok(serde_json::from_str(&buffer)?)
}
pub async fn process_single_job_old(
job_args: &JobDependencies,
job: &AvailableTtsInferenceJob,
model_record: &TtsModelForInferenceRecord,
) -> Result<(Span, Span), ProcessSingleJobError> {
let start = Instant::now();
let span_id = generate_random_uuid();
let trace_id = generate_random_uuid();
let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or("");
let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis())
.name("single job execution")
.attribute("user_token", maybe_user_token)
.service_name("tts-inference-job");
let span_id = generate_random_uuid();
let trace_id = generate_random_uuid();
let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000;
let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp)
.name("job since creation")
.attribute("user_token", maybe_user_token)
.service_name("tts-inference-job");
let mut job_progress_reporter = job_args
.job_progress_reporter
.new_tts_inference(&job.inference_job_token)
.map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
// ==================== ATTEMPT TO GRAB JOB LOCK ==================== //
info!("Attempting to grab lock for job: {}", job.inference_job_token);
let lock_acquired =
mark_tts_inference_job_pending_and_grab_lock(&job_args.mysql_pool, job.id)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
if !lock_acquired {
warn!("Could not acquire job lock for: {:?}", &job.id);
let duration = start.elapsed();
since_creation_span.set_attribute("status", "failure");
since_creation_span.set_duration(duration);
job_iteration_span.set_attribute("status", "failure");
job_iteration_span.set_duration(duration);
return Ok((since_creation_span, job_iteration_span));
}
info!("Lock acquired for job: {}", job.inference_job_token);
// ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== //
let waveglow_vocoder_model_fs_path = {
let waveglow_vocoder_model_filename = job_args.waveglow_vocoder_model_filename.clone();
let waveglow_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename);
let waveglow_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename);
maybe_download_file_from_bucket(
"waveglow vocoder model",
&waveglow_vocoder_model_fs_path,
&waveglow_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (1 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
waveglow_vocoder_model_fs_path
};
// ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== //
let pretrained_hifigan_vocoder_model_fs_path = {
let hifigan_vocoder_model_filename = job_args.hifigan_vocoder_model_filename.clone();
let hifigan_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename);
let hifigan_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename);
maybe_download_file_from_bucket(
"hifigan vocoder model",
&hifigan_vocoder_model_fs_path,
&hifigan_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (2 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
hifigan_vocoder_model_fs_path
};
// ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== //
let hifigan_superres_vocoder_model_fs_path = {
let hifigan_superres_vocoder_model_filename = job_args.hifigan_superres_vocoder_model_filename.clone();
let hifigan_superres_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename);
let hifigan_superres_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename);
maybe_download_file_from_bucket(
"hifigan superres vocoder model",
&hifigan_superres_vocoder_model_fs_path,
&hifigan_superres_vocoder_model_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading vocoder (3 of 3)",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
hifigan_superres_vocoder_model_fs_path
};
// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== //
let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder {
None => None,
Some(vocoder) => {
let custom_vocoder_fs_path = job_args.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token);
let custom_vocoder_object_path = job_args.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash);
maybe_download_file_from_bucket(
"custom vocoder",
&custom_vocoder_fs_path,
&custom_vocoder_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading user vocoder",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
Some(custom_vocoder_fs_path)
}
};
// ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== //
let tts_synthesizer_fs_path = {
let tts_synthesizer_fs_path = job_args.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token);
let tts_synthesizer_object_path = job_args.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash);
maybe_download_file_from_bucket(
"synthesizer",
&tts_synthesizer_fs_path,
&tts_synthesizer_object_path,
&job_args.private_bucket_client,
&mut job_progress_reporter,
"downloading synthesizer",
job.id.0,
&job_args.scoped_temp_dir_creator,
).await?;
tts_synthesizer_fs_path
};
// ==================== Preprocess text ==================== //
let cleaned_inference_text = clean_symbols(&job.raw_inference_text);
// ==================== WRITE TEXT TO FILE ==================== //
info!("Creating tempdir for inference results.");
let temp_dir = format!("temp_tts_inference_{}", job.id.0);
// NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
let temp_dir = TempDir::new(&temp_dir)
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
let text_input_fs_path = temp_dir.path().join("inference_input.txt");
std::fs::write(&text_input_fs_path, &cleaned_inference_text)
.map_err(|e| ProcessSingleJobError::from_io_error(e))?;
// ==================== RUN INFERENCE ==================== //
job_progress_reporter.log_status("running inference")
.map_err(|e| ProcessSingleJobError::Other(e))?;
// TODO: Fix this.
let maybe_unload_model_path = job_args
.virtual_model_lfu
.insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or(""))
.map_err(|e| ProcessSingleJobError::Other(e))?;
if let Some(model_path) = maybe_unload_model_path.as_deref() {
warn!("Remove model from LFU cache: {:?}", model_path);
}
let output_audio_fs_path = temp_dir.path().join("output.wav");
let output_metadata_fs_path = temp_dir.path().join("metadata.json");
let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json");
info!("Running TTS inference...");
info!("Expected output audio filename: {:?}", &output_audio_fs_path);
info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
info!("Expected output metadata filename: {:?}", &output_metadata_fs_path);
if let Some(model_path) = maybe_unload_model_path.as_deref() {
warn!("Unload model from sidecar: {:?}", &model_path);
}
let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution;
if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() {
pretrained_vocoder = VocoderType::from_str(default_vocoder)
.map_err(|e| ProcessSingleJobError::Other(e))?;
}
info!("With pretrained vocoder: {:?}", pretrained_vocoder);
let text_pipeline_type_or_guess = model_record.text_pipeline_type
.as_deref()
.and_then(|pipeline_type|
TextPipelineType::from_str(pipeline_type).ok())// NB: If there's an error deserializing, turn it to None.
.unwrap_or_else(||
guess_text_pipeline_heuristic(Some(model_record.created_at)));
info!("With text pipeline type `{:?} ` (or guess: {:?})",
&model_record.text_pipeline_type,
&text_pipeline_type_or_guess);
let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path {
None => {
info!("using pretrained HiFi-GAN vocoder");
pretrained_hifigan_vocoder_model_fs_path
},
Some(custom_vocoder_fs_path) => {
info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path);
custom_vocoder_fs_path
},
};
// NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to
// roughly 12 seconds max. Here we map seconds to decoder steps.
let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds);
job_args.http_clients.tts_inference_sidecar_client.request_inference(
&cleaned_inference_text,
max_decoder_steps,
&tts_synthesizer_fs_path,
pretrained_vocoder,
&text_pipeline_type_or_guess.to_str(),
&hifigan_vocoder_model_fs_path_to_use,
&hifigan_superres_vocoder_model_fs_path,
&waveglow_vocoder_model_fs_path,
&output_audio_fs_path,
&output_spectrogram_fs_path,
&output_metadata_fs_path,
maybe_unload_model_path)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
// ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
info!("Checking that output files exist...");
check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?;
let file_metadata = read_metadata_file(&output_metadata_fs_path)
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_metadata_fs_path);
// ==================== UPLOAD AUDIO TO BUCKET ==================== //
job_progress_reporter.log_status("uploading result")
.map_err(|e| ProcessSingleJobError::Other(e))?;
let audio_result_object_path = job_args.bucket_path_unifier.tts_inference_wav_audio_output_path(
&job.uuid_idempotency_token); // TODO: Don't use this!
info!("Audio destination bucket path: {:?}", &audio_result_object_path);
info!("Uploading audio...");
job_args.public_bucket_client.upload_filename_with_content_type(
&audio_result_object_path,
&output_audio_fs_path,
"audio/wav")
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_audio_fs_path);
// ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== //
let spectrogram_result_object_path = job_args.bucket_path_unifier.tts_inference_spectrogram_output_path(
&job.uuid_idempotency_token); // TODO: Don't use this!
info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path);
info!("Uploading spectrogram...");
job_args.public_bucket_client.upload_filename_with_content_type(
&spectrogram_result_object_path,
&output_spectrogram_fs_path,
"application/json")
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
safe_delete_temp_file(&output_spectrogram_fs_path);
// ==================== DELETE DOWNLOADED FILE ==================== //
// NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
safe_delete_temp_directory(&temp_dir);
// ==================== SAVE RECORDS ==================== //
let text_hash = hash_string_sha2(&cleaned_inference_text)
.map_err(|e| ProcessSingleJobError::Other(e))?;
let worker_name = job_args.get_worker_name();
info!("Saving tts inference record...");
let (id, inference_result_token) = insert_tts_result(
&job_args.mysql_pool,
job,
&text_hash,
pretrained_vocoder,
&audio_result_object_path,
&spectrogram_result_object_path,
file_metadata.file_size_bytes,
file_metadata.duration_millis.unwrap_or(0),
job_args.worker_details.is_on_prem,
&worker_name,
job_args.worker_details.is_debug_worker)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
info!("Marking job complete...");
mark_tts_inference_job_done(
&job_args.mysql_pool,
job.id,
true,
Some(&inference_result_token),
&worker_name)
.await
.map_err(|e| ProcessSingleJobError::Other(e))?;
info!("TTS Done. Original text was: {}", &job.raw_inference_text);
job_args.firehose_publisher.tts_inference_finished(
job.maybe_creator_user_token.as_deref(),
&model_record.model_token,
&inference_result_token)
.await
.map_err(|e| {
error!("error publishing event: {:?}", e);
ProcessSingleJobError::Other(anyhow!("error publishing event"))
})?;
job_progress_reporter.log_status("done")
.map_err(|e| ProcessSingleJobError::Other(e))?;
info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}",
job.id, id, &inference_result_token);
let duration = start.elapsed();
since_creation_span.set_attribute("status", "success");
since_creation_span.set_duration(duration);
job_iteration_span.set_attribute("status", "success");
job_iteration_span.set_duration(duration);
Ok((since_creation_span, job_iteration_span))
}
fn get_timestamp_millis() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
crates/service/job/inference_job/src/job_steps/select_job_batch.rs
+23
View File
@@ -0,0 +1,23 @@
use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::{list_available_generic_inference_jobs, ListAvailableGenericInferenceJobArgs};
use enums::workers::generic_inference_type::GenericInferenceType;
use sqlx::MySqlPool;
use std::collections::BTreeSet;
pub struct SelectJobBatchArgs <'a> {
pub num_records: u32,
pub is_debug_worker: bool,
pub sort_by_priority: bool,
pub maybe_scope_by_job_type: Option<BTreeSet<GenericInferenceType>>,
pub mysql_pool: &'a MySqlPool,
}
pub async fn select_job_batch(args: SelectJobBatchArgs<'_>) {
let _results = list_available_generic_inference_jobs(ListAvailableGenericInferenceJobArgs {
num_records: args.num_records,
is_debug_worker: args.is_debug_worker,
sort_by_priority: args.sort_by_priority,
maybe_scope_by_job_type: args.maybe_scope_by_job_type,
mysql_pool: args.mysql_pool,
}).await;
}
crates/service/job/inference_job/src/main.rs
+2 -2
View File
@@ -38,7 +38,7 @@ use crate::http_clients::tts_sidecar_health_check_client::TtsSidecarHealthCheckC
use crate::job_steps::job_dependencies::JobWorkerDetails;
use crate::job_steps::job_dependencies::{JobDependencies, JobCaches, JobHttpClients};
use crate::job_steps::job_stats::JobStats;
use crate::job_steps::main_loop::main_loop;
use crate::job_steps::main_loop_old::main_loop_old;
use crate::script_execution::tacotron_inference_command::TacotronInferenceCommand;
use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
use database_queries::mediators::firehose_publisher::FirehosePublisher;
@@ -328,7 +328,7 @@ async fn main() -> AnyhowResult<()> {
maybe_minimum_priority,
};
main_loop(inferencer).await;
main_loop_old(inferencer).await;
Ok(())
}