From 0cee5e5fccd51ebbaea2856f4f4474bd603c7573 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 7 Sep 2024 03:01:23 -0400 Subject: [PATCH] update cached model usage counts --- Cargo.lock | 16 ++++ crates/lib/mysql_queries/Cargo.toml | 5 +- .../queries/model_weight_usage_counts/mod.rs | 1 + .../sum_model_weight_usage_count_for_model.rs | 59 ++++++++++++ .../src/queries/model_weights/edit/mod.rs | 1 + .../update_model_weight_cached_usage_count.rs | 41 +++++++++ ...weight_tokens_for_cached_usage_backfill.rs | 47 ++++++++++ .../src/queries/model_weights/list/mod.rs | 1 + crates/service/job/analytics_job/src/main.rs | 6 +- .../job/analytics_job/src/tasks/mod.rs | 3 +- .../update_model_usage_counts_table.rs | 2 +- .../mod.rs | 1 + ...update_model_weights_cached_usage_field.rs | 91 +++++++++++++++++++ 13 files changed, 268 insertions(+), 6 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/model_weight_usage_counts/sum_model_weight_usage_count_for_model.rs create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/edit/update_model_weight_cached_usage_count.rs create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/list/list_all_model_weight_tokens_for_cached_usage_backfill.rs create mode 100644 crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/mod.rs create mode 100644 crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/update_model_weights_cached_usage_field.rs diff --git a/Cargo.lock b/Cargo.lock index 1bf5013959..06c9311a14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -913,6 +913,17 @@ dependencies = [ "zeroize", ] +[[package]] +name = "bigdecimal" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6773ddc0eafc0e509fb60e48dff7f450f8e674a0686ae8605e8d9901bd5eefa" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "billing_component" version = "0.0.1" @@ -4398,6 +4409,7 @@ dependencies = [ "idempotency", "log", "my-workspace-hack", + "num-traits", "once_cell", "rand 0.8.5", "regex", @@ -6555,6 +6567,7 @@ checksum = "d84b0a3c3739e220d94b3239fd69fb1f74bc36e16643423bd99de3b43c21bfbd" dependencies = [ "ahash", "atoi", + "bigdecimal", "byteorder", "bytes 1.6.0", "chrono", @@ -6640,6 +6653,7 @@ checksum = "e37195395df71fd068f6e2082247891bc11e3289624bbc776a0cdfa1ca7f1ea4" dependencies = [ "atoi", "base64 0.21.7", + "bigdecimal", "bitflags 2.5.0", "byteorder", "bytes 1.6.0", @@ -6683,6 +6697,7 @@ checksum = "d6ac0ac3b7ccd10cc96c7ab29791a7dd236bd94021f31eec7ba3d46a74aa1c24" dependencies = [ "atoi", "base64 0.21.7", + "bigdecimal", "bitflags 2.5.0", "byteorder", "chrono", @@ -6701,6 +6716,7 @@ dependencies = [ "log", "md-5 0.10.6", "memchr", + "num-bigint", "once_cell", "rand 0.8.5", "serde", diff --git a/crates/lib/mysql_queries/Cargo.toml b/crates/lib/mysql_queries/Cargo.toml index 0ae689ba63..8878fe67d7 100644 --- a/crates/lib/mysql_queries/Cargo.toml +++ b/crates/lib/mysql_queries/Cargo.toml @@ -46,8 +46,11 @@ serde_json = "1.0.108" toml = "0.8.14" serial_test = "3.0.0" +# BigNum +num-traits = "0.2.19" + tokio = { version = "1.38.0", features = ["full"] } rand = "0.8" -sqlx = { version = "0.7.2", features = [ "mysql", "runtime-tokio", "tls-rustls", "chrono" ] } +sqlx = { version = "0.7.2", features = [ "mysql", "runtime-tokio", "tls-rustls", "chrono", "bigdecimal" ] } diff --git a/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/mod.rs b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/mod.rs index 9068cbf4fd..8ab60e7671 100644 --- a/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/mod.rs @@ -1 +1,2 @@ +pub mod sum_model_weight_usage_count_for_model; pub mod upsert_model_weight_usage_count_for_date; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/sum_model_weight_usage_count_for_model.rs b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/sum_model_weight_usage_count_for_model.rs new file mode 100644 index 0000000000..2d0f7a5278 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/sum_model_weight_usage_count_for_model.rs @@ -0,0 +1,59 @@ +use anyhow::anyhow; +use num_traits::cast::ToPrimitive; +use sqlx::{Executor, MySql}; +use sqlx::types::BigDecimal; + +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +pub struct UsageCount { + pub total_usage_count: u64, +} + +pub async fn sum_model_weight_usage_count_for_models<'e, 'c, E>( + model_token: &'e ModelWeightToken, + mysql_executor: E +) + -> AnyhowResult> +where E: 'e + Executor<'c, Database = MySql> +{ + + let result = sqlx::query_as!( + RawUsageCount, + r#" +SELECT + sum(usage_count) as total_usage_count +FROM + model_weight_usage_counts +WHERE + token = ? +GROUP BY token +LIMIT 1 + "#, + model_token.as_str() + ) + .fetch_one(mysql_executor) + .await; + + let result = match result { + Ok(result) => result, + Err(err) => return match err { + sqlx::Error::RowNotFound => Ok(None), + _ => Err(anyhow!("Error querying for IP ban: {:?}", err)), + } + }; + + // TODO(bt,2024-09-07): It sucks that we have to use BigDecimal ... + let count = result.total_usage_count + .map(|count| count.to_u64()) + .flatten() + .map(|count| UsageCount { + total_usage_count: count, + }); + + Ok(count) +} + +pub struct RawUsageCount { + total_usage_count: Option, +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/edit/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/edit/mod.rs index c4c7bf689e..54d4509015 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/edit/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/edit/mod.rs @@ -1,2 +1,3 @@ pub mod set_model_weight_cover_image; +pub mod update_model_weight_cached_usage_count; pub mod update_weight; diff --git a/crates/lib/mysql_queries/src/queries/model_weights/edit/update_model_weight_cached_usage_count.rs b/crates/lib/mysql_queries/src/queries/model_weights/edit/update_model_weight_cached_usage_count.rs new file mode 100644 index 0000000000..df85556778 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/edit/update_model_weight_cached_usage_count.rs @@ -0,0 +1,41 @@ +use std::marker::PhantomData; + +use sqlx::{Executor, MySql}; + +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +pub struct Args<'e, 'c, E> +where E: 'e + Executor<'c, Database = MySql> +{ + /// The model we're updating. + pub model_weight_token: &'e ModelWeightToken, + + /// The new cached usage count. + pub usage_count: u64, + + pub mysql_executor: E, + + // TODO: Not sure if this works to tell the compiler we need the lifetime annotation. + // See: https://doc.rust-lang.org/std/marker/struct.PhantomData.html#unused-lifetime-parameters + pub phantom: PhantomData<&'c E>, +} + +pub async fn update_model_weight_cached_usage_count<'e, 'c : 'e, E>( + args: Args<'e, 'c, E>, +) -> AnyhowResult<()> where E: 'e + Executor<'c, Database = MySql> { + + let _query_result = sqlx::query!( + r#" + UPDATE model_weights + SET + cached_usage_count = ? + WHERE token = ? + LIMIT 1 + "#, + args.usage_count, + args.model_weight_token, + ).execute(args.mysql_executor).await?; + + Ok(()) +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/list_all_model_weight_tokens_for_cached_usage_backfill.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/list_all_model_weight_tokens_for_cached_usage_backfill.rs new file mode 100644 index 0000000000..65d768854d --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/list_all_model_weight_tokens_for_cached_usage_backfill.rs @@ -0,0 +1,47 @@ +use chrono::{DateTime, Utc}; +use sqlx::{Executor, MySql}; + +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +pub struct ModelWeightInfo { + pub token: ModelWeightToken, + pub cached_usage_count: u64, +} + +/// Get records updated since a particular date. +/// NB: This doesn't batch yet at our scale, but in the future we'll likely need/want to batch, +/// and then we'll need another way to cursor in the event that batch_size < # records updated +/// within a given time quantum. +pub async fn list_all_model_weight_tokens_for_cached_usage_backfill<'e, 'c, E>( + mysql_executor: E, +) -> AnyhowResult> + where E: 'e + Executor<'c, Database=MySql> +{ + let query = sqlx::query_as!( + RawRecord, + r#" + SELECT + token as `token: tokens::tokens::model_weights::ModelWeightToken`, + cached_usage_count + FROM model_weights + "#, + ); + + let results = query.fetch_all(mysql_executor).await?; + + let results = results.into_iter() + .map(|record| ModelWeightInfo { + token: record.token, + cached_usage_count: record.cached_usage_count, + }) + .collect(); + + Ok(results) +} + +struct RawRecord { + token: ModelWeightToken, + cached_usage_count: u64, +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs index 68f17efe79..f4b83832be 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs @@ -1,4 +1,5 @@ pub mod list_all_model_weight_tokens_for_backfill; +pub mod list_all_model_weight_tokens_for_cached_usage_backfill; pub mod list_featured_weights; pub mod list_model_weight_tokens_updated_since; pub mod list_model_weights_for_elastic_search_backfill_using_cursor; diff --git a/crates/service/job/analytics_job/src/main.rs b/crates/service/job/analytics_job/src/main.rs index 7e78aed56e..31b17dd4e4 100644 --- a/crates/service/job/analytics_job/src/main.rs +++ b/crates/service/job/analytics_job/src/main.rs @@ -25,8 +25,8 @@ use config::shared_constants::DEFAULT_RUST_LOG; use errors::AnyhowResult; use crate::job_state::{JobState, SleepConfigs}; -use crate::tasks::calculate_old_model_analytics::calculate_old_model_analytics_loop::calculate_old_model_analytics_loop; use crate::tasks::update_model_usage_counts_table::update_model_usage_counts_table::update_model_usage_counts_table; +use crate::tasks::update_model_weights_cached_usage_field::update_model_weights_cached_usage_field::update_model_weights_cached_usage_field; pub mod job_state; pub mod tasks; @@ -62,13 +62,13 @@ async fn main() -> AnyhowResult<()> { let job_state_1 = job_state.clone(); let handle_1 = tokio::task::spawn(async move { - let _r = calculate_old_model_analytics_loop(job_state_1).await; + let _r = update_model_usage_counts_table(job_state_1).await; }); let job_state_2 = job_state.clone(); let handle_2 = tokio::task::spawn(async move { - let _r = update_model_usage_counts_table(job_state_2).await; + let _r = update_model_weights_cached_usage_field(job_state_2).await; }); futures::future::join_all([ diff --git a/crates/service/job/analytics_job/src/tasks/mod.rs b/crates/service/job/analytics_job/src/tasks/mod.rs index fababc4ff1..8b6ff486c9 100644 --- a/crates/service/job/analytics_job/src/tasks/mod.rs +++ b/crates/service/job/analytics_job/src/tasks/mod.rs @@ -1,2 +1,3 @@ pub mod calculate_old_model_analytics; -pub mod update_model_usage_counts_table; \ No newline at end of file +pub mod update_model_usage_counts_table; +pub mod update_model_weights_cached_usage_field; \ No newline at end of file diff --git a/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs index c248c79ba8..01b7407cbd 100644 --- a/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs +++ b/crates/service/job/analytics_job/src/tasks/update_model_usage_counts_table/update_model_usage_counts_table.rs @@ -29,13 +29,13 @@ pub async fn update_model_usage_counts_table(job_state: JobState) -> AnyhowResul } async fn connect_and_calculate_loop(job_state: &JobState) -> AnyhowResult<()> { + // NB: Don't re-establish this connection unless something goes wrong. let mut connection = job_state.mysql_pool.acquire().await?; loop { calculate_with_connection(&mut connection, job_state).await?; } } - async fn calculate_with_connection( connection: &mut PoolConnection, job_state: &JobState diff --git a/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/mod.rs b/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/mod.rs new file mode 100644 index 0000000000..20ae34195a --- /dev/null +++ b/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/mod.rs @@ -0,0 +1 @@ +pub mod update_model_weights_cached_usage_field; \ No newline at end of file diff --git a/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/update_model_weights_cached_usage_field.rs b/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/update_model_weights_cached_usage_field.rs new file mode 100644 index 0000000000..cafc8833ff --- /dev/null +++ b/crates/service/job/analytics_job/src/tasks/update_model_weights_cached_usage_field/update_model_weights_cached_usage_field.rs @@ -0,0 +1,91 @@ +use std::time::Duration; + +use log::{error, info}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; + +use errors::AnyhowResult; +use mysql_queries::queries::model_weight_usage_counts::sum_model_weight_usage_count_for_model::sum_model_weight_usage_count_for_models; +use mysql_queries::queries::model_weights::edit::update_model_weight_cached_usage_count::{Args, update_model_weight_cached_usage_count}; +use mysql_queries::queries::model_weights::list::list_all_model_weight_tokens_for_cached_usage_backfill::list_all_model_weight_tokens_for_cached_usage_backfill; + +use crate::job_state::JobState; + +pub async fn update_model_weights_cached_usage_field(job_state: JobState) -> AnyhowResult<()> { + loop { + info!("update model weights cached usage field"); + + match connect_and_calculate_loop(&job_state).await { + Ok(_) => { + tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)).await; + } + Err(e) => { + error!("error: {:?}", e); + tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_error_wait_millis)).await; + } + } + } +} + +async fn connect_and_calculate_loop(job_state: &JobState) -> AnyhowResult<()> { + // NB: Don't re-establish this connection unless something goes wrong. + let mut connection = job_state.mysql_pool.acquire().await?; + loop { + calculate_with_connection(&mut connection, job_state).await?; + } +} + +async fn calculate_with_connection( + connection: &mut PoolConnection, + job_state: &JobState +) -> AnyhowResult<()> { + + let backfill_tokens = list_all_model_weight_tokens_for_cached_usage_backfill(&mut **connection) + .await?; + + info!("Records found: {}", backfill_tokens.len()); + + let mut skipped_sparse_updates = 0; + let mut skipped_noop_updates = 0; + + for backfill_token in backfill_tokens { + info!("Updating cached usages for model: {:?}", backfill_token.token); + + let usage_count = sum_model_weight_usage_count_for_models( + &backfill_token.token, + &mut **connection, + ).await?; + + let usage_count = usage_count + .map(|usage| usage.total_usage_count); + + let usage_count = match usage_count { + None | Some(0) => { + skipped_sparse_updates += 1; + continue; + } + Some(count) => count, + }; + + if usage_count == backfill_token.cached_usage_count { + skipped_noop_updates += 1; + continue; + } + + info!("Token: {} Uses: {}", backfill_token.token, usage_count); + + update_model_weight_cached_usage_count(Args { + model_weight_token: &backfill_token.token, + usage_count, + mysql_executor: &mut **connection, + phantom: Default::default(), + }).await?; + } + + info!("Skipped sparse updates: {}", skipped_sparse_updates); + info!("Skipped no-op updates: {}", skipped_noop_updates); + + tokio::time::sleep(Duration::from_millis(job_state.sleep_config.between_job_batch_wait_millis)).await; + + Ok(()) +}