update cached model usage counts

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",

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" ] }
 
 

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;

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<Option<UsageCount>>
+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<BigDecimal>,
+}

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;

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(())
+}

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<Vec<ModelWeightInfo>>
+    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,
+}

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;

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([

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;
+pub mod update_model_usage_counts_table;
+pub mod update_model_weights_cached_usage_field;

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<MySql>,
   job_state: &JobState

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;

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<MySql>,
+  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(())
+}