diff --git a/migrations/2021-05-15-212913_inference_jobs/up.sql b/migrations/2021-05-15-212913_inference_jobs/up.sql index 37b9759ed7..ad9f95a2ea 100644 --- a/migrations/2021-05-15-212913_inference_jobs/up.sql +++ b/migrations/2021-05-15-212913_inference_jobs/up.sql @@ -1,5 +1,7 @@ -# noinspection SqlResolveForFile -# noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + CREATE TABLE tts_inference_jobs ( id BIGINT(20) NOT NULL AUTO_INCREMENT, diff --git a/migrations/2021-05-30-010200_firehose/up.sql b/migrations/2021-05-30-010200_firehose/up.sql index 374c8a34ee..bfc5e3c6cc 100644 --- a/migrations/2021-05-30-010200_firehose/up.sql +++ b/migrations/2021-05-30-010200_firehose/up.sql @@ -1,5 +1,6 @@ -# noinspection SqlResolveForFile -# noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile CREATE TABLE firehose_entries ( -- Not used for anything except replication. @@ -9,18 +10,6 @@ CREATE TABLE firehose_entries ( token VARCHAR(32) NOT NULL, -- The type of the event - event_type VARCHAR(32) NOT NULL, - - -- The target user - maybe_target_user_token VARCHAR(32) DEFAULT NULL, - - -- The target "entity", which varies by event type. - maybe_target_entity_token VARCHAR(32) DEFAULT NULL, - - created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, - updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, - - -- Settings -- DO NOT REORDER. -- event_type ENUM( -- 'not-set', @@ -40,12 +29,28 @@ CREATE TABLE firehose_entries ( -- 'twitch_subscribe', -- 'twitch_follow' -- ) NOT NULL DEFAULT 'not-set', + event_type VARCHAR(32) NOT NULL, + + -- The target user + maybe_target_user_token VARCHAR(32) DEFAULT NULL, + + -- The target "entity", which varies by event type. + -- eg. a model token or template token + maybe_target_entity_token VARCHAR(32) DEFAULT NULL, + + -- If something was created, this might be the token (if we know it) + -- eg. an inference job's token + maybe_created_entity_token VARCHAR(32) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, -- INDICES -- PRIMARY KEY (id), UNIQUE KEY (token), KEY fk_maybe_target_user_token (maybe_target_user_token), KEY fk_maybe_target_entity_token (maybe_target_entity_token), + KEY fk_maybe_created_entity_token (maybe_created_entity_token), KEY index_event_type (event_type) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/src/common_queries/firehose_publisher.rs b/src/common_queries/firehose_publisher.rs index 235f2462d6..c4de611150 100644 --- a/src/common_queries/firehose_publisher.rs +++ b/src/common_queries/firehose_publisher.rs @@ -62,7 +62,43 @@ impl FirehosePublisher { let _record_id = self.insert( FirehoseEvent::UserSignUp, Some(user_token), - Some(user_token) + Some(user_token), + Some(user_token) + ).await?; + Ok(()) + } + + pub async fn enqueue_tts_model_upload(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsModelUploadStarted, + Some(user_token), + None, + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn publish_tts_model_upload_finished(&self, user_token: &str, model_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsModelUploadCompleted, + Some(user_token), + Some(model_token), + Some(model_token) + ).await?; + Ok(()) + } + + pub async fn enqueue_tts_inference( + &self, + maybe_user_token: Option<&str>, + job_token: &str, + model_token: &str + ) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsInferenceStarted, + maybe_user_token, + Some(model_token), + Some(job_token) ).await?; Ok(()) } @@ -71,15 +107,17 @@ impl FirehosePublisher { let _record_id = self.insert( FirehoseEvent::W2lTemplateUploadStarted, Some(user_token), + None, Some(job_token) ).await?; Ok(()) } - pub async fn enqueue_w2l_inference(&self, maybe_user_token: Option<&str>, job_token: &str) -> AnyhowResult<()> { + pub async fn enqueue_w2l_inference(&self, maybe_user_token: Option<&str>, job_token: &str, template_token: &str) -> AnyhowResult<()> { let _record_id = self.insert( FirehoseEvent::W2lInferenceStarted, maybe_user_token, + Some(template_token), Some(job_token) ).await?; Ok(()) @@ -89,16 +127,18 @@ impl FirehosePublisher { let _record_id = self.insert( FirehoseEvent::W2lTemplateUploadCompleted, Some(user_token), - Some(template_token) + Some(template_token), + Some(template_token) ).await?; Ok(()) } - pub async fn w2l_inference_finished(&self, maybe_user_token: Option<&str>, job_token: &str) -> AnyhowResult<()> { + pub async fn w2l_inference_finished(&self, maybe_user_token: Option<&str>, job_token: &str, result_token: &str) -> AnyhowResult<()> { let _record_id = self.insert( FirehoseEvent::W2lInferenceCompleted, maybe_user_token, - Some(job_token) + Some(job_token), // TODO: This could be template_token + Some(result_token) ).await?; Ok(()) } @@ -109,7 +149,8 @@ impl FirehosePublisher { &self, event_type: FirehoseEvent, user_token: Option<&str>, - entity_token: Option<&str> + entity_token: Option<&str>, + created_entity_token: Option<&str> ) -> AnyhowResult { let token = random_prefix_crockford_token("EV", 32)?; @@ -120,12 +161,14 @@ SET token = ?, event_type = ?, maybe_target_user_token = ?, - maybe_target_entity_token = ? + maybe_target_entity_token = ?, + maybe_created_entity_token = ? "#, token, event_type.to_db_value(), user_token, entity_token, + created_entity_token ) .execute(&self.mysql_pool) .await; diff --git a/src/http_server/endpoints/tts/enqueue_infer_tts.rs b/src/http_server/endpoints/tts/enqueue_infer_tts.rs index 052ba872b7..3af94f3831 100644 --- a/src/http_server/endpoints/tts/enqueue_infer_tts.rs +++ b/src/http_server/endpoints/tts/enqueue_infer_tts.rs @@ -130,7 +130,7 @@ SET creator_set_visibility = ?, status = "pending" "#, - job_token, + &job_token, request.uuid_idempotency_token.clone(), model_token, inference_text, @@ -172,6 +172,13 @@ SET info!("new w2l inference job id: {}", record_id); + server_state.firehose_publisher.enqueue_tts_inference(maybe_user_token.as_deref(), &job_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + InferTtsError::ServerError + })?; + let response = InferTtsSuccessResponse { success: true, }; diff --git a/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs b/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs index eaa69e8ee0..c9dbe57804 100644 --- a/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs +++ b/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs @@ -200,6 +200,13 @@ SET info!("new model upload job id: {}", record_id); + server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadTtsModelError::ServerError + })?; + let response = UploadTtsModelSuccessResponse { success: true, job_token: job_token.to_string(), diff --git a/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs b/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs index e196ae6035..5e66b444fb 100644 --- a/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs +++ b/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs @@ -341,7 +341,7 @@ SET // .map_err(|err| HttpResponse::InternalServerError() // .body(format!("Redis Err: {:?}", err)))?; - server_state.firehose_publisher.enqueue_w2l_inference(maybe_user_token.as_deref(), &job_token) + server_state.firehose_publisher.enqueue_w2l_inference(maybe_user_token.as_deref(), &job_token, &template_token) .await .map_err(|e| { warn!("error publishing event: {:?}", e); diff --git a/src/job_queries/tts_download_job_queries.rs b/src/job_queries/tts_download_job_queries.rs index 49b2a649f7..b7562d2730 100644 --- a/src/job_queries/tts_download_job_queries.rs +++ b/src/job_queries/tts_download_job_queries.rs @@ -115,7 +115,7 @@ pub async fn insert_tts_model(pool: &MySqlPool, job: &TtsUploadJobRecord, private_bucket_hash: &str, private_bucket_object_name: &str) - -> AnyhowResult + -> AnyhowResult<(u64, String)> { let model_token = random_prefix_crockford_token("TTS_MDL:", 32)?; let updatable_slug = model_token.clone(); @@ -136,7 +136,7 @@ SET private_bucket_hash = ?, private_bucket_object_name = ? "#, - model_token, + &model_token, updatable_slug, job.title.to_string(), job.creator_user_token.clone(), @@ -158,5 +158,5 @@ SET } }; - Ok(record_id) + Ok((record_id, model_token)) } \ No newline at end of file diff --git a/src/tts_download_job.rs b/src/tts_download_job.rs index e1a41b8e74..1436b02f41 100644 --- a/src/tts_download_job.rs +++ b/src/tts_download_job.rs @@ -6,7 +6,10 @@ #![warn(unused_must_use)] //#![allow(warnings)] +#[macro_use] extern crate serde_derive; + pub mod buckets; +pub mod common_queries; pub mod job_queries; pub mod script_execution; pub mod shared_constants; @@ -17,6 +20,7 @@ use chrono::Utc; use crate::buckets::bucket_client::BucketClient; use crate::buckets::bucket_paths::hash_to_bucket_path; use crate::buckets::file_hashing::get_file_hash; +use crate::common_queries::firehose_publisher::FirehosePublisher; use crate::job_queries::tts_download_job_queries::TtsUploadJobRecord; use crate::job_queries::tts_download_job_queries::insert_tts_model; use crate::job_queries::tts_download_job_queries::mark_tts_upload_job_done; @@ -53,6 +57,7 @@ struct Downloader { pub download_temp_directory: PathBuf, pub mysql_pool: MySqlPool, pub bucket_client: BucketClient, + pub firehose_publisher: FirehosePublisher, pub google_drive_downloader: GoogleDriveDownloadCommand, // Command to run pub download_script: String, @@ -116,6 +121,10 @@ async fn main() -> AnyhowResult<()> { .connect(&db_connection_string) .await?; + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + }; + let downloader = Downloader { download_temp_directory: temp_directory, mysql_pool, @@ -123,6 +132,7 @@ async fn main() -> AnyhowResult<()> { download_script, google_drive_downloader, bucket_root_tts_model_uploads: bucket_root.to_string(), + firehose_publisher, }; main_loop(downloader).await; @@ -223,7 +233,7 @@ async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> Anyho downloader.bucket_client.upload_filename(&object_name, &file_path).await?; info!("Saving model record..."); - let id = insert_tts_model( + let (id, model_token) = insert_tts_model( &downloader.mysql_pool, job, &private_bucket_hash, @@ -232,6 +242,13 @@ async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> Anyho info!("Saved model record: {}", id); + downloader.firehose_publisher.publish_tts_model_upload_finished(&job.creator_user_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + info!("Job done: {}", job.id); mark_tts_upload_job_done(&downloader.mysql_pool, job, true).await?; diff --git a/src/w2l_inference_job.rs b/src/w2l_inference_job.rs index 3fc88c61a7..e2475a7d84 100644 --- a/src/w2l_inference_job.rs +++ b/src/w2l_inference_job.rs @@ -535,6 +535,7 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An inferencer.firehose_publisher.w2l_inference_finished( job.maybe_creator_user_token.as_deref(), + &job.inference_job_token, &inference_result_token) .await .map_err(|e| {