write cached model usage counts to elasticsearch

.sqlx/query-4bce2b10f0510feca7f5705049a5c0eaae6e1a840146628495827468a38c05dd.json

@@ -1,6 +1,6 @@
 {
   "db_name": "MySQL",
-  "query": "\nSELECT\n    w.id,\n    w.token as `token: tokens::tokens::model_weights::ModelWeightToken`,\n\n    w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n    w.weights_type as `weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n    w.weights_category as `weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n    w.title,\n\n    cover_image.token as `maybe_cover_image_media_file_token: tokens::tokens::media_files::MediaFileToken`,\n    cover_image.public_bucket_directory_hash as maybe_cover_image_public_bucket_hash,\n    cover_image.maybe_public_bucket_prefix as maybe_cover_image_public_bucket_prefix,\n    cover_image.maybe_public_bucket_extension as maybe_cover_image_public_bucket_extension,\n\n    w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,\n    users.username as creator_username,\n    users.display_name as creator_display_name,\n    users.email_gravatar_hash as creator_gravatar_hash,\n\n    featured_items.entity_token IS NOT NULL AS is_featured,\n\n    entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n    entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n    entity_stats.bookmark_count as maybe_bookmark_count,\n\n    extension_tts.ietf_language_tag as maybe_tts_ietf_language_tag,\n    extension_tts.ietf_primary_language_subtag as maybe_tts_ietf_primary_language_subtag,\n\n    extension_vc.ietf_language_tag as maybe_vc_ietf_language_tag,\n    extension_vc.ietf_primary_language_subtag as maybe_vc_ietf_primary_language_subtag,\n\n    w.created_at,\n    w.updated_at,\n    w.user_deleted_at,\n    w.mod_deleted_at,\n\n    NOW() as database_read_time\n\nFROM model_weights as w\n\nJOIN users\n    ON users.token = w.creator_user_token\n\nLEFT OUTER JOIN model_weights_extension_tts_details as extension_tts\n    ON extension_tts.model_weights_token = w.token\nLEFT OUTER JOIN model_weights_extension_voice_conversion_details as extension_vc\n    ON extension_vc.model_weights_token = w.token\nLEFT OUTER JOIN entity_stats\n    ON entity_stats.entity_type = \"model_weight\"\n    AND entity_stats.entity_token = w.token\nLEFT OUTER JOIN media_files as cover_image\n    ON cover_image.token = w.maybe_cover_image_media_file_token\nLEFT OUTER JOIN featured_items\n    ON featured_items.entity_type = \"model_weight\"\n    AND featured_items.entity_token = w.token\n    AND featured_items.deleted_at IS NULL\n\nWHERE\n  w.id > ?\nORDER BY id ASC\nLIMIT ?\n        ",
+  "query": "\nSELECT\n    w.id,\n    w.token as `token: tokens::tokens::model_weights::ModelWeightToken`,\n\n    w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n    w.weights_type as `weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n    w.weights_category as `weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n    w.title,\n\n    cover_image.token as `maybe_cover_image_media_file_token: tokens::tokens::media_files::MediaFileToken`,\n    cover_image.public_bucket_directory_hash as maybe_cover_image_public_bucket_hash,\n    cover_image.maybe_public_bucket_prefix as maybe_cover_image_public_bucket_prefix,\n    cover_image.maybe_public_bucket_extension as maybe_cover_image_public_bucket_extension,\n\n    w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,\n    users.username as creator_username,\n    users.display_name as creator_display_name,\n    users.email_gravatar_hash as creator_gravatar_hash,\n\n    featured_items.entity_token IS NOT NULL AS is_featured,\n\n    entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n    entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n    entity_stats.bookmark_count as maybe_bookmark_count,\n\n    w.cached_usage_count,\n\n    extension_tts.ietf_language_tag as maybe_tts_ietf_language_tag,\n    extension_tts.ietf_primary_language_subtag as maybe_tts_ietf_primary_language_subtag,\n\n    extension_vc.ietf_language_tag as maybe_vc_ietf_language_tag,\n    extension_vc.ietf_primary_language_subtag as maybe_vc_ietf_primary_language_subtag,\n\n    w.created_at,\n    w.updated_at,\n    w.user_deleted_at,\n    w.mod_deleted_at,\n\n    NOW() as database_read_time\n\nFROM model_weights as w\n\nJOIN users\n    ON users.token = w.creator_user_token\n\nLEFT OUTER JOIN model_weights_extension_tts_details as extension_tts\n    ON extension_tts.model_weights_token = w.token\nLEFT OUTER JOIN model_weights_extension_voice_conversion_details as extension_vc\n    ON extension_vc.model_weights_token = w.token\nLEFT OUTER JOIN entity_stats\n    ON entity_stats.entity_type = \"model_weight\"\n    AND entity_stats.entity_token = w.token\nLEFT OUTER JOIN media_files as cover_image\n    ON cover_image.token = w.maybe_cover_image_media_file_token\nLEFT OUTER JOIN featured_items\n    ON featured_items.entity_type = \"model_weight\"\n    AND featured_items.entity_token = w.token\n    AND featured_items.deleted_at IS NULL\n\nWHERE\n  w.id > ?\nORDER BY id ASC\nLIMIT ?\n        ",
   "describe": {
     "columns": [
       {
@@ -185,6 +185,16 @@
       },
       {
         "ordinal": 18,
+        "name": "cached_usage_count",
+        "type_info": {
+          "type": "LongLong",
+          "flags": "NOT_NULL | UNSIGNED",
+          "char_set": 63,
+          "max_size": 20
+        }
+      },
+      {
+        "ordinal": 19,
         "name": "maybe_tts_ietf_language_tag",
         "type_info": {
           "type": "VarString",
@@ -194,7 +204,7 @@
         }
       },
       {
-        "ordinal": 19,
+        "ordinal": 20,
         "name": "maybe_tts_ietf_primary_language_subtag",
         "type_info": {
           "type": "VarString",
@@ -204,7 +214,7 @@
         }
       },
       {
-        "ordinal": 20,
+        "ordinal": 21,
         "name": "maybe_vc_ietf_language_tag",
         "type_info": {
           "type": "VarString",
@@ -214,7 +224,7 @@
         }
       },
       {
-        "ordinal": 21,
+        "ordinal": 22,
         "name": "maybe_vc_ietf_primary_language_subtag",
         "type_info": {
           "type": "VarString",
@@ -224,7 +234,7 @@
         }
       },
       {
-        "ordinal": 22,
+        "ordinal": 23,
         "name": "created_at",
         "type_info": {
           "type": "Timestamp",
@@ -234,7 +244,7 @@
         }
       },
       {
-        "ordinal": 23,
+        "ordinal": 24,
         "name": "updated_at",
         "type_info": {
           "type": "Timestamp",
@@ -244,7 +254,7 @@
         }
       },
       {
-        "ordinal": 24,
+        "ordinal": 25,
         "name": "user_deleted_at",
         "type_info": {
           "type": "Timestamp",
@@ -254,7 +264,7 @@
         }
       },
       {
-        "ordinal": 25,
+        "ordinal": 26,
         "name": "mod_deleted_at",
         "type_info": {
           "type": "Timestamp",
@@ -264,7 +274,7 @@
         }
       },
       {
-        "ordinal": 26,
+        "ordinal": 27,
         "name": "database_read_time",
         "type_info": {
           "type": "Datetime",
@@ -296,6 +306,7 @@
       true,
       true,
       true,
+      false,
       true,
       true,
       true,
@@ -307,5 +318,5 @@
       false
     ]
   },
-  "hash": "9dba7a47b8d1bc79da56354a1d7f75a36f41732bba110b37c72a97aa2408ea4b"
+  "hash": "4bce2b10f0510feca7f5705049a5c0eaae6e1a840146628495827468a38c05dd"
 }

Cargo.lock

@@ -2307,6 +2307,7 @@ dependencies = [
  "log",
  "my-workspace-hack",
  "mysql_queries",
+ "primitives",
  "serde",
  "serde_derive",
  "serde_json",
@@ -2502,6 +2503,7 @@ dependencies = [
  "my-workspace-hack",
  "mysql_queries",
  "once_cell",
+ "primitives",
  "rand 0.8.5",
  "reusable_types",
  "serde",
@@ -4411,6 +4413,7 @@ dependencies = [
  "my-workspace-hack",
  "num-traits",
  "once_cell",
+ "primitives",
  "rand 0.8.5",
  "regex",
  "reusable_types",

_elasticsearch/index_definitions/model_weights_v1.json

@@ -29,6 +29,11 @@
       "ratings_negative_count" : { "type" : "unsigned_long" },
       "bookmark_count" : { "type" : "unsigned_long" },
 
+      "cached_usage_count" : {
+        "type": "integer",
+        "index": true
+      },
+
       "maybe_ietf_language_tag" : { "type" : "keyword" },
       "maybe_ietf_primary_language_subtag" : { "type" : "keyword" },
 

crates/cli/production/elasticsearch-cli/Cargo.toml

@@ -24,6 +24,7 @@ elasticsearch_schema = { path = "../../../lib/elasticsearch_schema" }
 enums = { path = "../../../_types/enums" }
 errors = { path = "../../../lib/errors" }
 mysql_queries = { path = "../../../lib/mysql_queries" }
+primitives = { path = "../../../lib/data/primitives" }
 storyteller_root = { path = "../../../lib/storyteller_root" }
 tokens = { path = "../../../_types/tokens" }
 

crates/cli/production/elasticsearch-cli/src/plans/create_all_model_weight_documents.rs

@@ -11,6 +11,7 @@ use elasticsearch_schema::utils::create_index_if_not_exists::{create_index_if_no
 use enums::by_table::model_weights::weights_category::WeightsCategory;
 use errors::AnyhowResult;
 use mysql_queries::queries::model_weights::list::list_model_weights_for_elastic_search_backfill_using_cursor::{list_model_weights_for_elastic_search_backfill_using_cursor, ModelWeightForElasticsearchRecord};
+use primitives::numerics::u64_to_i32_saturating::u64_to_i32_saturating;
 use storyteller_root::get_storyteller_rust_root;
 
 pub async fn create_all_model_weight_documents(
@@ -86,6 +87,8 @@ async fn create_document_from_record(elasticsearch: &Elasticsearch, record: Mode
     ratings_negative_count: record.maybe_ratings_negative_count.unwrap_or(0),
     bookmark_count: record.maybe_bookmark_count.unwrap_or(0),
 
+    cached_usage_count: Some(u64_to_i32_saturating(record.cached_usage_count)),
+
     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,

crates/lib/data/primitives/Cargo.toml

@@ -19,6 +19,7 @@ normal = ["my-workspace-hack"]
 [dependencies]
 # NB: This is designed to be a central crate with few dependencies.
 # Consider the dependencies here wisely.
+# This is used in MySQL, etc., so dependencies are ill-advised.
 
 # Internal
 #errors = { path = "../../errors" }

crates/lib/data/primitives/src/lib.rs

@@ -20,9 +20,10 @@
 
 pub mod bool_to_str;
 pub mod lazy_any_option_true;
+pub mod numerics;
+pub mod optional_false_to_none;
 pub mod str_to_bool;
 pub mod traits;
 pub mod trim_or_empty;
 pub mod truncate_str;
 pub mod try_str_to_num;
-pub mod optional_false_to_none;

crates/lib/data/primitives/src/numerics/mod.rs

@@ -0,0 +1 @@
+pub mod u64_to_i32_saturating;

crates/lib/data/primitives/src/numerics/u64_to_i32_saturating.rs

@@ -0,0 +1,41 @@
+
+const MAX_I32 : u64 = i32::MAX as u64;
+
+/// Convert a u64 to a positive i32, saturating to i32::MAX if the u64 overflows the container.
+pub fn u64_to_i32_saturating(num: u64) -> i32 {
+  if num > MAX_I32 {
+    return i32::MAX;
+  } else {
+    num as i32
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn zero() {
+    assert_eq!(u64_to_i32_saturating(0), 0);
+  }
+
+  #[test]
+  fn u64_max() {
+    assert_eq!(u64_to_i32_saturating(u64::MAX), i32::MAX);
+  }
+
+  #[test]
+  fn i32_max() {
+    assert_eq!(u64_to_i32_saturating(i32::MAX as u64), i32::MAX);
+  }
+
+  #[test]
+  fn i32_max_plus_one() {
+    assert_eq!(u64_to_i32_saturating(i32::MAX as u64 + 1), i32::MAX);
+  }
+
+  #[test]
+  fn i32_max_minus_one() {
+    assert_eq!(u64_to_i32_saturating(i32::MAX as u64 - 1), i32::MAX - 1);
+  }
+}

crates/lib/elasticsearch_schema/src/documents/model_weight_document.rs

@@ -49,7 +49,7 @@ pub struct ModelWeightDocument {
   pub creator_display_name: String,
   pub creator_gravatar_hash: String,
 
-  // TODO(bt,2024-08-08): This should be migrated to non-nullable soon. Turn this into a strict `bool` and drop the Option<T>
+  // TODO(bt, 2024-08-08): This should be migrated to non-nullable soon. Turn this into a strict `bool` and drop the Option<T>
   pub is_featured: Option<bool>,
 
   // Statistics
@@ -57,6 +57,10 @@ pub struct ModelWeightDocument {
   pub ratings_negative_count: u32,
   pub bookmark_count: u32,
 
+  // Stored as "integer" type
+  // TODO(bt, 2024-09-07): This should be migrated to non-nullable soon. Turn this into a strict `i32` and drop the Option<T>
+  pub cached_usage_count: Option<i32>,
+
   // Fields only used for TTS models and voice conversion models.
   pub maybe_ietf_language_tag: Option<String>,
   pub maybe_ietf_primary_language_subtag: Option<String>,

crates/lib/mysql_queries/Cargo.toml

@@ -23,6 +23,7 @@ composite_identifiers = { path = "../../_types/composite_identifiers" }
 enums = { path = "../../_types/enums" }
 errors = { path = "../errors" }
 idempotency = { path = "../idempotency" }
+primitives = { path = "../data/primitives" }
 reusable_types = { path = "../../_types/reusable_types" }
 tokens = { path = "../../_types/tokens" }
 

crates/lib/mysql_queries/src/queries/model_weights/batch_get/batch_get_model_weights_for_elastic_search_backfill.rs

@@ -46,6 +46,7 @@ pub struct ModelWeightForElasticsearchRecord {
   pub maybe_ratings_positive_count: Option<u32>,
   pub maybe_ratings_negative_count: Option<u32>,
   pub maybe_bookmark_count: Option<u32>,
+  pub cached_usage_count: u64,
 
   // TTS extensions
   pub maybe_tts_ietf_language_tag: Option<String>,
@@ -113,6 +114,7 @@ pub async fn batch_get_model_weights_for_elastic_search_backfill<'e, 'c, E>(
           maybe_ratings_positive_count: model.maybe_ratings_positive_count,
           maybe_ratings_negative_count: model.maybe_ratings_negative_count,
           maybe_bookmark_count: model.maybe_bookmark_count,
+          cached_usage_count: model.cached_usage_count,
           creator_set_visibility: model.creator_set_visibility,
           created_at: model.created_at,
           updated_at: model.updated_at,
@@ -163,6 +165,8 @@ SELECT
     entity_stats.ratings_negative_count as maybe_ratings_negative_count,
     entity_stats.bookmark_count as maybe_bookmark_count,
 
+    w.cached_usage_count,
+
     extension_tts.ietf_language_tag as maybe_tts_ietf_language_tag,
     extension_tts.ietf_primary_language_subtag as maybe_tts_ietf_primary_language_subtag,
 
@@ -247,6 +251,7 @@ struct RawRecord {
   pub maybe_ratings_positive_count: Option<u32>,
   pub maybe_ratings_negative_count: Option<u32>,
   pub maybe_bookmark_count: Option<u32>,
+  pub cached_usage_count: u64,
 
   // TTS extensions
   pub maybe_tts_ietf_language_tag: Option<String>,
@@ -309,6 +314,8 @@ impl FromRow<'_, MySqlRow> for RawRecord {
       maybe_ratings_negative_count: row.try_get("maybe_ratings_negative_count")?,
       maybe_bookmark_count: row.try_get("maybe_bookmark_count")?,
 
+      cached_usage_count: row.try_get("cached_usage_count")?,
+
       maybe_tts_ietf_language_tag: row.try_get("maybe_tts_ietf_language_tag")?,
       maybe_tts_ietf_primary_language_subtag: row.try_get("maybe_tts_ietf_primary_language_subtag")?,
 

crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_elastic_search_backfill_using_cursor.rs

@@ -45,6 +45,7 @@ pub struct ModelWeightForElasticsearchRecord {
   pub maybe_ratings_positive_count: Option<u32>,
   pub maybe_ratings_negative_count: Option<u32>,
   pub maybe_bookmark_count: Option<u32>,
+  pub cached_usage_count: u64,
 
   // TTS extensions
   pub maybe_tts_ietf_language_tag: Option<String>,
@@ -107,6 +108,7 @@ pub async fn list_model_weights_for_elastic_search_backfill_using_cursor(
           maybe_ratings_positive_count: model.maybe_ratings_positive_count,
           maybe_ratings_negative_count: model.maybe_ratings_negative_count,
           maybe_bookmark_count: model.maybe_bookmark_count,
+          cached_usage_count: model.cached_usage_count,
           creator_set_visibility: model.creator_set_visibility,
           created_at: model.created_at,
           updated_at: model.updated_at,
@@ -153,6 +155,8 @@ SELECT
     entity_stats.ratings_negative_count as maybe_ratings_negative_count,
     entity_stats.bookmark_count as maybe_bookmark_count,
 
+    w.cached_usage_count,
+
     extension_tts.ietf_language_tag as maybe_tts_ietf_language_tag,
     extension_tts.ietf_primary_language_subtag as maybe_tts_ietf_primary_language_subtag,
 
@@ -227,6 +231,7 @@ struct RawRecord {
   pub maybe_ratings_positive_count: Option<u32>,
   pub maybe_ratings_negative_count: Option<u32>,
   pub maybe_bookmark_count: Option<u32>,
+  pub cached_usage_count: u64,
 
   // TTS extensions
   pub maybe_tts_ietf_language_tag: Option<String>,

crates/service/job/es-update-job/Cargo.toml

@@ -25,6 +25,7 @@ elasticsearch_schema = { path = "../../../lib/elasticsearch_schema" }
 enums = { path = "../../../_types/enums" }
 errors = { path = "../../../lib/errors" }
 mysql_queries = { path = "../../../lib/mysql_queries" }
+primitives = { path = "../../../lib/data/primitives" }
 reusable_types = { path = "../../../_types/reusable_types" }
 server_environment = { path = "../../../lib/server_environment" }
 storyteller_root = { path = "../../../lib/storyteller_root" }

crates/service/job/es-update-job/src/tasks/model_weights/util/create_model_weight_document_from_record.rs

@@ -8,6 +8,7 @@ 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;
+use primitives::numerics::u64_to_i32_saturating::u64_to_i32_saturating;
 
 pub async fn create_model_weight_document_from_record(
   elasticsearch: &Elasticsearch,
@@ -45,6 +46,8 @@ pub async fn create_model_weight_document_from_record(
     ratings_negative_count: record.maybe_ratings_negative_count.unwrap_or(0),
     bookmark_count: record.maybe_bookmark_count.unwrap_or(0),
 
+    cached_usage_count: Some(u64_to_i32_saturating(record.cached_usage_count)),
+
     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,