update recent model weights records

This commit is contained in:
Brandon Thomas committed 2024-08-01 20:03:00 -04:00
1 parent d2e9140b50
commit 62f7279870
13 files changed
+409 -195

No files matched your search

@@ -0,0 +1,30 @@
use anyhow::anyhow;
use chrono::{DateTime, Utc};
use sqlx::MySqlPool;
use errors::AnyhowResult;
pub struct GetDatabaseTimeResult {
pub database_time: DateTime<Utc>,
}
/// Query the DB for its current time
pub async fn get_database_time(
pool: &MySqlPool,
) -> AnyhowResult<GetDatabaseTimeResult> {
let maybe_record = sqlx::query_as!(
GetDatabaseTimeResult,
r#"
SELECT NOW() as `database_time: DateTime<Utc>`
"#,
)
.fetch_one(pool)
.await;
match maybe_record {
Ok(record) => Ok(record),
Err(sqlx::Error::RowNotFound) => Err(anyhow!("query didn't return anything")),
Err(ref err) => Err(anyhow!("database query error: {:?}", err)),
}
}
@@ -0,0 +1 @@
pub mod get_database_time;
@@ -5,6 +5,7 @@ pub mod beta_keys;
pub mod billing;
pub mod browser_session_logs;
pub mod comments;
pub mod database_time;
pub mod email_sender_jobs;
pub mod entity_stats;
pub mod featured_items;
+21 -6
View File
@@ -17,6 +17,7 @@
//#![forbid(warnings)]
use std::sync::Arc;
use elasticsearch::Elasticsearch;
use elasticsearch::http::transport::Transport;
use log::info;
@@ -29,8 +30,10 @@ use config::shared_constants::DEFAULT_RUST_LOG;
use errors::AnyhowResult;
use crate::job_state::{JobState, SleepConfigs};
use crate::tasks::update_all_model_weights::update_all_model_weights::update_all_model_weights;
use crate::tasks::model_weights::update_all_model_weights::update_all_model_weights;
use crate::tasks::model_weights::update_recently_written_model_weights::update_recently_written_model_weights;
use crate::tasks::update_engine_media_files::update_engine_media_files::update_engine_media_files;
pub mod job_state;
pub mod tasks;
@@ -60,23 +63,35 @@ async fn main() -> AnyhowResult<()> {
},
});
info!("Starting thread to update all model weights...");
info!("Starting thread to update recently updated model weights...");
let job_state_1 = job_state.clone();
let handle_1 = tokio::task::spawn(async move {
let _r = update_all_model_weights(job_state_1).await;
let _r = update_recently_written_model_weights(job_state_1).await;
});
info!("Starting thread to update engine media files...");
info!("Starting thread to update all model weights...");
let job_state_2 = job_state.clone();
let handle_2 = tokio::task::spawn(async move {
let _r = update_engine_media_files(job_state_2).await;
let _r = update_all_model_weights(job_state_2).await;
});
futures::future::join_all([handle_1, handle_2]).await;
info!("Starting thread to update engine media files...");
let job_state_3 = job_state.clone();
let handle_3 = tokio::task::spawn(async move {
let _r = update_engine_media_files(job_state_3).await;
});
futures::future::join_all([
handle_1,
handle_2,
handle_3,
]).await;
Ok(())
}
@@ -1,2 +1,2 @@
pub mod update_all_model_weights;
pub mod model_weights;
pub mod update_engine_media_files;
@@ -0,0 +1,3 @@
pub mod update_all_model_weights;
pub mod update_recently_written_model_weights;
pub mod util;
@@ -0,0 +1,100 @@
use std::sync::Arc;
use std::time::Duration;
use chrono::{DateTime, Utc};
use log::{error, info};
use errors::AnyhowResult;
use mysql_queries::queries::model_weights::batch_get::batch_get_model_weights_for_elastic_search_backfill::batch_get_model_weights_for_elastic_search_backfill;
use mysql_queries::queries::model_weights::list::list_model_weight_tokens_updated_since::list_model_weight_tokens_updated_since;
use crate::job_state::JobState;
use crate::tasks::model_weights::util::create_model_weight_document_from_record::create_model_weight_document_from_record;
pub async fn update_all_model_weights(job_state: Arc<JobState>) {
// TODO(bt,2024-02-05): Write this cursor to Redis so job can resume without reindexing everything.
let mut cursor = DateTime::UNIX_EPOCH;
loop {
info!("Main loop; cursor @ {:?}", &cursor);
let result = with_database_main_loop(&mut cursor, &job_state).await;
if let Err(err) = result {
error!("Error in main loop: {:?}", err);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)).await;
}
// NB: Function should never return in success case, but we'll sleep just in case.
}
}
pub async fn with_database_main_loop(updated_at_cursor: &mut DateTime<Utc>, job_state: &JobState) -> AnyhowResult<()> {
info!("Acquiring MySQL connection...");
let mut mysql_connection = job_state.mysql_pool.acquire().await?;
let mut last_successful_update_at = *updated_at_cursor;
loop {
info!("Querying tokens updated since: {:?}", &updated_at_cursor);
let mut maybe_tokens =
list_model_weight_tokens_updated_since(&mut *mysql_connection, &updated_at_cursor).await?;
if maybe_tokens.is_empty() {
info!("No records updated since {:?}", &updated_at_cursor);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_no_updates_wait_millis)).await;
continue;
}
info!("Found {} updated records", maybe_tokens.len());
let mut last_observed_updated_at = *updated_at_cursor;
while !maybe_tokens.is_empty() {
// NB: This list might be very large if we query from (1) the epoch, or (2) there was a large series of updates
let last = 50.min(maybe_tokens.len());
let drained_tokens = maybe_tokens.drain(0..last)
.into_iter()
.map(|record| record.token)
.collect::<Vec<_>>();
let records
= batch_get_model_weights_for_elastic_search_backfill(&mut *mysql_connection, &drained_tokens).await?;
for record in records {
let updated_at = record.updated_at.clone();
create_model_weight_document_from_record(&job_state.elasticsearch, record).await?;
// NB: We don't want to advance the cursor to the current second just yet, because we might be processing
// several records with the exact same "updated_at" timestamp. We do know that due to ordering, we can update
// the timestamp cursor to the second before this timestamp, however.
//
// If we change the cursoring or batch sizes, we'll need to reconsider all of this logic.
if let Some(cursor_at_least) = updated_at.checked_sub_signed(chrono::Duration::seconds(1)) {
// NB: We don't want the cursor to slide backwards (it shouldn't, but the max will save us).
last_successful_update_at = cursor_at_least.max(last_successful_update_at);
}
*updated_at_cursor = last_successful_update_at.max(*updated_at_cursor);
last_observed_updated_at = updated_at;
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_es_writes_wait_millis)).await;
}
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)).await;
}
// NB: The last cursor math put us a second behind the current clock. Here we'll advance to the last record.
// NB: This technically could miss records if we're updating within the same second batch we read from, but that
// seems both unlikely and not worth solving at our present scale.
*updated_at_cursor = last_observed_updated_at.max(*updated_at_cursor);
info!("Up to date at cursor = {:?}", &updated_at_cursor);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_query_wait_millis)).await;
}
}
@@ -0,0 +1,91 @@
use std::ops::Add;
use std::sync::Arc;
use std::time::Duration;
use chrono::{DateTime, Utc};
use log::{error, info};
use errors::AnyhowResult;
use mysql_queries::queries::database_time::get_database_time::get_database_time;
use mysql_queries::queries::model_weights::list::list_model_weight_tokens_updated_since::list_model_weight_tokens_updated_since;
use crate::job_state::JobState;
use crate::tasks::model_weights::util::copy_model_weight_records_to_documents::copy_model_weight_records_to_documents;
pub async fn update_recently_written_model_weights(job_state: Arc<JobState>) {
let mut last_updated_time= DateTime::UNIX_EPOCH;
// TODO(bt,2024-02-05): Write this cursor to Redis so job can resume without reindexing everything.
loop {
match get_database_time(&job_state.mysql_pool).await {
Ok(db_time) => {
last_updated_time = db_time.database_time;
break;
},
Err(err) => {
error!("Error getting database time: {:?}", err);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)).await;
continue;
}
}
}
loop {
info!("Main loop; cursor @ {:?}", &last_updated_time);
let result = with_database_main_loop(&mut last_updated_time, &job_state).await;
if let Err(err) = result {
error!("Error in main loop: {:?}", err);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)).await;
}
// NB: Function should never return in success case, but we'll sleep just in case.
}
}
pub async fn with_database_main_loop(updated_at_cursor: &mut DateTime<Utc>, job_state: &JobState) -> AnyhowResult<()> {
info!("Acquiring MySQL connection...");
let mut mysql_connection = job_state.mysql_pool.acquire().await?;
loop {
info!("Querying tokens updated since: {:?}", &updated_at_cursor);
let maybe_records =
list_model_weight_tokens_updated_since(&mut *mysql_connection, &updated_at_cursor).await?;
if maybe_records.is_empty() {
info!("No records updated since {:?}", &updated_at_cursor);
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_no_updates_wait_millis)).await;
continue;
}
let maybe_tokens = maybe_records.into_iter()
.map(|record| record.token)
.collect::<Vec<_>>();
info!("Found {} updated records", maybe_tokens.len());
let maybe_last_cursor = copy_model_weight_records_to_documents(
maybe_tokens,
&mut *mysql_connection,
&job_state.elasticsearch,
&job_state.sleep_config,
).await?;
if let Some(last_cursor) = maybe_last_cursor {
// NB: If we trigger thousands of records to get updated at the same time, the cursor could get stuck.
// This is a simple way to ensure we don't get stuck, but we could miss out on some updates.
let delta = last_cursor.updated_at.signed_duration_since(*updated_at_cursor);
if delta.num_seconds() == 0 {
// Cursor didn't move, so we need to advance it.
*updated_at_cursor = last_cursor.updated_at.add(Duration::from_secs(1));
} else {
*updated_at_cursor = last_cursor.updated_at;
}
}
tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_query_wait_millis)).await;
}
}
@@ -0,0 +1,63 @@
use std::time::Duration;
use chrono::{DateTime, Utc};
use elasticsearch::Elasticsearch;
use sqlx::MySqlConnection;
use errors::AnyhowResult;
use mysql_queries::queries::model_weights::batch_get::batch_get_model_weights_for_elastic_search_backfill::batch_get_model_weights_for_elastic_search_backfill;
use tokens::tokens::model_weights::ModelWeightToken;
use crate::job_state::SleepConfigs;
use crate::tasks::model_weights::util::create_model_weight_document_from_record::create_model_weight_document_from_record;
pub struct Cursor {
pub token: ModelWeightToken,
pub updated_at: DateTime<Utc>,
}
pub async fn copy_model_weight_records_to_documents(
mut tokens: Vec<ModelWeightToken>,
mysql_connection: &mut MySqlConnection,
elasticsearch: &Elasticsearch,
sleep_config: &SleepConfigs,
) -> AnyhowResult<Option<Cursor>> {
let mut last_updated_record = None;
while !tokens.is_empty() {
// NB: This list might be very large if we query from (1) the epoch, or (2) there was a large series of updates
let last = 50.min(tokens.len());
let drained_tokens = tokens.drain(0..last)
.into_iter()
.collect::<Vec<_>>();
let records =
batch_get_model_weights_for_elastic_search_backfill(
&mut *mysql_connection,
&drained_tokens
).await?;
for record in records {
let new_cursor = Cursor {
token: record.token.clone(),
updated_at: record.updated_at.clone(),
};
create_model_weight_document_from_record(elasticsearch, record).await?;
let last_timestamp = last_updated_record.as_ref()
.map(|cursor: &Cursor| cursor.updated_at)
.unwrap_or_else(|| DateTime::UNIX_EPOCH);
if new_cursor.updated_at > last_timestamp {
last_updated_record = Some(new_cursor);
}
tokio::time::sleep(Duration::from_millis(sleep_config.between_es_writes_wait_millis)).await;
}
tokio::time::sleep(Duration::from_millis(sleep_config.between_job_batch_wait_millis)).await;
}
Ok(last_updated_record)
}
@@ -0,0 +1,96 @@
use anyhow::anyhow;
use elasticsearch::{BulkOperation, BulkParts, Elasticsearch};
use log::{error, info};
use serde_json::Value;
use elasticsearch_schema::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument};
use elasticsearch_schema::traits::document::Document;
use enums::by_table::model_weights::weights_category::WeightsCategory;
use errors::AnyhowResult;
use mysql_queries::queries::model_weights::batch_get::batch_get_model_weights_for_elastic_search_backfill::ModelWeightForElasticsearchRecord;
pub async fn create_model_weight_document_from_record(
elasticsearch: &Elasticsearch,
record: ModelWeightForElasticsearchRecord
) -> AnyhowResult<()> {
info!("Create record for {:?} - {:?}", record.token, record.title);
let is_deleted = record.user_deleted_at.is_some() || record.mod_deleted_at.is_some();
let document = ModelWeightDocument {
token: record.token,
creator_set_visibility: record.creator_set_visibility,
weights_type: record.weights_type,
weights_category: record.weights_category,
title: record.title.to_string(),
title_as_keyword: record.title,
maybe_cover_image_media_file_token: record.maybe_cover_image_media_file_token,
maybe_cover_image_public_bucket_hash: record.maybe_cover_image_public_bucket_hash,
maybe_cover_image_public_bucket_prefix: record.maybe_cover_image_public_bucket_prefix,
maybe_cover_image_public_bucket_extension: record.maybe_cover_image_public_bucket_extension,
creator_user_token: record.creator_user_token,
creator_username: record.creator_username,
creator_display_name: record.creator_display_name,
creator_gravatar_hash: record.creator_gravatar_hash,
ratings_positive_count: record.maybe_ratings_positive_count.unwrap_or(0),
ratings_negative_count: record.maybe_ratings_negative_count.unwrap_or(0),
bookmark_count: record.maybe_bookmark_count.unwrap_or(0),
maybe_ietf_language_tag: match record.weights_category {
WeightsCategory::TextToSpeech => record.maybe_tts_ietf_language_tag,
WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_language_tag,
_ => None,
},
maybe_ietf_primary_language_subtag: match record.weights_category {
WeightsCategory::TextToSpeech => record.maybe_tts_ietf_primary_language_subtag,
WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_primary_language_subtag,
_ => None,
},
created_at: record.created_at,
updated_at: record.updated_at,
user_deleted_at: record.user_deleted_at,
mod_deleted_at: record.mod_deleted_at,
database_read_time: record.database_read_time,
is_deleted,
};
let op : BulkOperation<_> = BulkOperation::index(&document)
.id(document.get_document_id())
.into();
let response = elasticsearch
.bulk(BulkParts::Index(MODEL_WEIGHT_INDEX))
.body(vec![op])
.send()
.await?;
let json: Value = response.json().await?;
let had_errors = json["errors"].as_bool().unwrap_or(false);
if had_errors {
let failed: Vec<&Value> = json["items"]
.as_array()
.unwrap()
.iter()
.filter(|v| !v["error"].is_null())
.collect();
error!("Errors during indexing. Failures: {}", failed.len());
return Err(anyhow!("Errors during indexing. Failures: {}", failed.len()));
}
Ok(())
}
@@ -0,0 +1,2 @@
pub mod copy_model_weight_records_to_documents;
pub mod create_model_weight_document_from_record;
@@ -1 +0,0 @@
pub mod update_all_model_weights;
@@ -1,187 +0,0 @@
use std::sync::Arc;
use std::time::Duration;
use anyhow::anyhow;
use chrono::{DateTime, Utc};
use elasticsearch::{BulkOperation, BulkParts, Elasticsearch};
use log::{error, info};
use serde_json::Value;
use elasticsearch_schema::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument};
use elasticsearch_schema::traits::document::Document;
use enums::by_table::model_weights::weights_category::WeightsCategory;
use errors::AnyhowResult;
use mysql_queries::queries::model_weights::batch_get::batch_get_model_weights_for_elastic_search_backfill::{batch_get_model_weights_for_elastic_search_backfill, ModelWeightForElasticsearchRecord};
use mysql_queries::queries::model_weights::list::list_model_weight_tokens_updated_since::list_model_weight_tokens_updated_since;
use crate::job_state::JobState;
pub async fn update_all_model_weights(job_state: Arc<JobState>) {
// TODO(bt,2024-02-05): Write this cursor to Redis so job can resume without reindexing everything.
let mut cursor = DateTime::UNIX_EPOCH;
loop {
info!("Main loop; cursor @ {:?}", &cursor);
let result = with_database_main_loop(&mut cursor, &job_state).await;
if let Err(err) = result {
error!("Error in main loop: {:?}", err);
std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis));
}
// NB: Function should never return in success case, but we'll sleep just in case.
}
}
pub async fn with_database_main_loop(updated_at_cursor: &mut DateTime<Utc>, job_state: &JobState) -> AnyhowResult<()> {
info!("Acquiring MySQL connection...");
let mut mysql_connection = job_state.mysql_pool.acquire().await?;
let mut last_successful_update_at = *updated_at_cursor;
loop {
info!("Querying tokens updated since: {:?}", &updated_at_cursor);
let mut maybe_tokens =
list_model_weight_tokens_updated_since(&mut *mysql_connection, &updated_at_cursor).await?;
if maybe_tokens.is_empty() {
info!("No records updated since {:?}", &updated_at_cursor);
std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_no_updates_wait_millis));
continue;
}
info!("Found {} updated records", maybe_tokens.len());
let mut last_observed_updated_at = *updated_at_cursor;
while !maybe_tokens.is_empty() {
// NB: This list might be very large if we query from (1) the epoch, or (2) there was a large series of updates
let last = 50.min(maybe_tokens.len());
let drained_tokens = maybe_tokens.drain(0..last)
.into_iter()
.map(|record| record.token)
.collect::<Vec<_>>();
let records
= batch_get_model_weights_for_elastic_search_backfill(&mut *mysql_connection, &drained_tokens).await?;
for record in records {
let updated_at = record.updated_at;
create_document_from_record(&job_state.elasticsearch, record).await?;
// NB: We don't want to advance the cursor to the current second just yet, because we might be processing
// several records with the exact same "updated_at" timestamp. We do know that due to ordering, we can update
// the timestamp cursor to the second before this timestamp, however.
//
// If we change the cursoring or batch sizes, we'll need to reconsider all of this logic.
if let Some(cursor_at_least) = updated_at.checked_sub_signed(chrono::Duration::seconds(1)) {
// NB: We don't want the cursor to slide backwards (it shouldn't, but the max will save us).
last_successful_update_at = cursor_at_least.max(last_successful_update_at);
}
*updated_at_cursor = last_successful_update_at.max(*updated_at_cursor);
last_observed_updated_at = updated_at;
std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_es_writes_wait_millis));
}
std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis));
}
// NB: The last cursor math put us a second behind the current clock. Here we'll advance to the last record.
// NB: This technically could miss records if we're updating within the same second batch we read from, but that
// seems both unlikely and not worth solving at our present scale.
*updated_at_cursor = last_observed_updated_at.max(*updated_at_cursor);
info!("Up to date at cursor = {:?}", &updated_at_cursor);
std::thread::sleep(Duration::from_millis(job_state.sleep_config.between_query_wait_millis));
}
}
async fn create_document_from_record(elasticsearch: &Elasticsearch, record: ModelWeightForElasticsearchRecord) -> AnyhowResult<()> {
info!("Create record for {:?} - {:?}", record.token, record.title);
let is_deleted = record.user_deleted_at.is_some() || record.mod_deleted_at.is_some();
let document = ModelWeightDocument {
token: record.token,
creator_set_visibility: record.creator_set_visibility,
weights_type: record.weights_type,
weights_category: record.weights_category,
title: record.title.to_string(),
title_as_keyword: record.title,
maybe_cover_image_media_file_token: record.maybe_cover_image_media_file_token,
maybe_cover_image_public_bucket_hash: record.maybe_cover_image_public_bucket_hash,
maybe_cover_image_public_bucket_prefix: record.maybe_cover_image_public_bucket_prefix,
maybe_cover_image_public_bucket_extension: record.maybe_cover_image_public_bucket_extension,
creator_user_token: record.creator_user_token,
creator_username: record.creator_username,
creator_display_name: record.creator_display_name,
creator_gravatar_hash: record.creator_gravatar_hash,
ratings_positive_count: record.maybe_ratings_positive_count.unwrap_or(0),
ratings_negative_count: record.maybe_ratings_negative_count.unwrap_or(0),
bookmark_count: record.maybe_bookmark_count.unwrap_or(0),
maybe_ietf_language_tag: match record.weights_category {
WeightsCategory::TextToSpeech => record.maybe_tts_ietf_language_tag,
WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_language_tag,
_ => None,
},
maybe_ietf_primary_language_subtag: match record.weights_category {
WeightsCategory::TextToSpeech => record.maybe_tts_ietf_primary_language_subtag,
WeightsCategory::VoiceConversion => record.maybe_voice_conversion_ietf_primary_language_subtag,
_ => None,
},
created_at: record.created_at,
updated_at: record.updated_at,
user_deleted_at: record.user_deleted_at,
mod_deleted_at: record.mod_deleted_at,
database_read_time: record.database_read_time,
is_deleted,
};
let op : BulkOperation<_> = BulkOperation::index(&document)
.id(document.get_document_id())
.into();
let response = elasticsearch
.bulk(BulkParts::Index(MODEL_WEIGHT_INDEX))
.body(vec![op])
.send()
.await?;
let json: Value = response.json().await?;
let had_errors = json["errors"].as_bool().unwrap_or(false);
if had_errors {
let failed: Vec<&Value> = json["items"]
.as_array()
.unwrap()
.iter()
.filter(|v| !v["error"].is_null())
.collect();
error!("Errors during indexing. Failures: {}", failed.len());
return Err(anyhow!("Errors during indexing. Failures: {}", failed.len()));
}
Ok(())
}