diff --git a/.idea/dictionaries/bt.xml b/.idea/dictionaries/bt.xml index 4e9d82228f..a03480e222 100644 --- a/.idea/dictionaries/bt.xml +++ b/.idea/dictionaries/bt.xml @@ -3,6 +3,7 @@ arpabet backfill + backfilling blendshapes canonicalized canonicalizing diff --git a/Cargo.lock b/Cargo.lock index e6d6b57ab8..cbcbac1ff6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1814,6 +1814,7 @@ dependencies = [ "anyhow", "chrono", "clap", + "datetimes", "dotenv", "easyenv", "enums", diff --git a/crates/cli/production/db-backfill/Cargo.toml b/crates/cli/production/db-backfill/Cargo.toml index 1762017036..18e2ee33ba 100644 --- a/crates/cli/production/db-backfill/Cargo.toml +++ b/crates/cli/production/db-backfill/Cargo.toml @@ -20,6 +20,7 @@ normal = ["my-workspace-hack"] my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } # Internal +datetimes = { path = "../../../lib/data/datetimes" } easyenv = { path = "../../../lib/easyenv" } enums = { path = "../../../_types/enums" } errors = { path = "../../../lib/errors" } diff --git a/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs b/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs index 85d704d56b..aebb335414 100644 --- a/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs +++ b/crates/cli/production/db-backfill/src/operations/calculate_model_weights_usages/run_migration.rs @@ -2,39 +2,66 @@ use anyhow::anyhow; use chrono::NaiveDate; use log::info; use sqlx::{MySql, Pool}; - +use datetimes::CHRONO_DATETIME_UNIX_EPOCH; use errors::AnyhowResult; use mysql_queries::queries::model_weight_usage_counts::upsert_model_weight_usage_count_for_date::{Args, upsert_model_weight_usage_count_for_date}; use mysql_queries::queries::model_weights::count::count_model_use_using_media_files_on_date::count_model_use_using_media_files_on_date; +use mysql_queries::queries::model_weights::list::list_model_weight_tokens_updated_since::list_model_weight_tokens_updated_since; use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::model_weights::ModelWeightToken; -use crate::operations::calculate_model_weights_usages::sub_args::parse_cli_sub_args; +use crate::operations::calculate_model_weights_usages::sub_args::{parse_cli_sub_args, SubArgs}; pub async fn run_migration(mysql: Pool) -> AnyhowResult<()> { let args = parse_cli_sub_args()?; - info!("args: {:?}", args); + info!("Backfill args: {:?}", args); - let model_token = args.model_token - .ok_or_else(|| anyhow!("no token provided"))?; + let model_tokens = match args.model_token.as_ref() { + Some(model_token) => vec![model_token.clone()], + None => get_all_model_weight_tokens(mysql.clone()).await?, + }; + for token in model_tokens.iter() { + backfill_token(&mysql, &args, &token).await?; + } + + info!("Model token count: {:?}", model_tokens.len()); + Ok(()) +} + +async fn get_all_model_weight_tokens(mysql: Pool) -> AnyhowResult> { + let epoch = *CHRONO_DATETIME_UNIX_EPOCH; + let results = list_model_weight_tokens_updated_since(&mysql, &epoch).await?; + Ok(results.into_iter() + .map(|result| result.token) + .collect()) +} + +async fn backfill_token(mysql: &Pool, args: &SubArgs, model_token: &ModelWeightToken) -> AnyhowResult<()> { let date = args.start_date .ok_or_else(|| anyhow!("invalid start date"))?; - let count = count_model_use_using_media_files_on_date(&mysql, &model_token, date).await?; + info!("Backfilling token: {:?} for date: {:?}", model_token, date); - info!("count: {:?}", count.record_count); + let count = count_model_use_using_media_files_on_date( + mysql, &model_token, date).await?; + + if count.record_count == 0 { + info!("Count: {:?} (skipping)", count.record_count); + return Ok(()); + } else { + info!("Count: {:?}", count.record_count); + } upsert_model_weight_usage_count_for_date(Args { model_token: &model_token, date, usage_count: count.record_count, - mysql_executor: &mysql, + insert_on_zero: false, + mysql_executor: mysql, phantom: Default::default(), }).await?; - - println!("calculate_model_weights_usages: run migration"); - Ok(()) } diff --git a/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/upsert_model_weight_usage_count_for_date.rs b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/upsert_model_weight_usage_count_for_date.rs index 71e7c7961d..f6cfc37b36 100644 --- a/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/upsert_model_weight_usage_count_for_date.rs +++ b/crates/lib/mysql_queries/src/queries/model_weight_usage_counts/upsert_model_weight_usage_count_for_date.rs @@ -1,6 +1,7 @@ use std::marker::PhantomData; use chrono::NaiveDate; +use log::info; use sqlx::{Executor, MySql}; use errors::AnyhowResult; @@ -13,6 +14,12 @@ where E: 'e + Executor<'c, Database = MySql> pub date: NaiveDate, pub usage_count: u64, + /// We should only need a sparse matrix for our histogram. + /// There will be a lot of "0" usage count "holes", and we should skip creating + /// records for them. This is just an extra explicit check to make sure we don't + /// insert unless we really want to do so. + pub insert_on_zero: bool, + pub mysql_executor: E, // TODO: Not sure if this works to tell the compiler we need the lifetime annotation. @@ -25,6 +32,10 @@ pub async fn upsert_model_weight_usage_count_for_date<'e, 'c : 'e, E>( ) -> AnyhowResult<()> where E: 'e + Executor<'c, Database = MySql> { + if !args.insert_on_zero && args.usage_count == 0 { + info!("Skipping insert for zero usage count"); + return Ok(()); + } let query = sqlx::query!( r#"