From d447ad7d8676f1ea2f3e5304b9c3fed77278acd8 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 14 Dec 2023 03:49:58 -0500 Subject: [PATCH] deterministic test tokens it'd be good to use the same token values in tests, dev seed tools, etc. this should be a threadsafe, resettable, deterministic RNG --- crates/_types/tokens/src/deterministic_rng.rs | 4 +- crates/_types/tokens/src/macros.rs | 4 +- .../tokens/tests/media_file_token_tests.rs | 17 +++---- .../src/seeding/model_weights.rs | 44 ++++++++++--------- .../src/seeding/seed_model_weights.rs | 44 +++++++++---------- 5 files changed, 54 insertions(+), 59 deletions(-) diff --git a/crates/_types/tokens/src/deterministic_rng.rs b/crates/_types/tokens/src/deterministic_rng.rs index 90ad47f5ab..affd9fda00 100644 --- a/crates/_types/tokens/src/deterministic_rng.rs +++ b/crates/_types/tokens/src/deterministic_rng.rs @@ -27,11 +27,11 @@ impl DeterministicRng { } /// Reset the deterministic RNG's seed. - pub fn reset_rng(&self) { + pub fn reset_rng(&self, state: u64) { match self.rng.write() { Err(err) => panic!(format!("test panic due to rng failure {err}")), Ok(mut lock) => { - *lock = StdRng::seed_from_u64(0); + *lock = StdRng::seed_from_u64(state); } } } diff --git a/crates/_types/tokens/src/macros.rs b/crates/_types/tokens/src/macros.rs index 4df28db41a..44fc14d6df 100644 --- a/crates/_types/tokens/src/macros.rs +++ b/crates/_types/tokens/src/macros.rs @@ -63,9 +63,9 @@ macro_rules! impl_crockford_generator { /// Constructor for a new token to be used **DETERMINISTICALLY** in tests and seeding. /// This will not use a real RNG, so do not use in production code! /// Once called, it will output the same random numbers in sequence. - pub fn reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously() { + pub fn reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(state: u64) { let deterministic = crate::deterministic_rng::DeterministicRng::get_instance().expect("tests should not fail due to mutex"); - deterministic.reset_rng(); + deterministic.reset_rng(state); } /// Constructor for a new token to be used **DETERMINISTICALLY** in tests and seeding. diff --git a/crates/_types/tokens/tests/media_file_token_tests.rs b/crates/_types/tokens/tests/media_file_token_tests.rs index 9a8704e1b2..722f89c0a5 100644 --- a/crates/_types/tokens/tests/media_file_token_tests.rs +++ b/crates/_types/tokens/tests/media_file_token_tests.rs @@ -23,23 +23,18 @@ mod interface { #[test] fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_1() { - MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + // NB: Using the same reset seed will produce the same results each time + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(0); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); - } - - #[test] - fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_2() { - MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + // NB: Same seed -> same tokens generated + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(0); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); - } - - #[test] - fn generate_for_testing_and_dev_seeding_never_use_in_production_seriously_3() { - MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(); + // Once more... + MediaFileToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(0); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_q8sz47gmfw2zx02snrbz88ns9m16ab"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_ma1xetxrwbh39vg639a9zrq8b9wk6h"); assert_eq!(MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously().as_str(), "m_4tswec8z27wnm01njypx4vmfhgj41e"); diff --git a/crates/cli/dev_database_seed/src/seeding/model_weights.rs b/crates/cli/dev_database_seed/src/seeding/model_weights.rs index 8860d323c7..13e23d6014 100644 --- a/crates/cli/dev_database_seed/src/seeding/model_weights.rs +++ b/crates/cli/dev_database_seed/src/seeding/model_weights.rs @@ -132,7 +132,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT - Gawr Gura MOS 4 "#; - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("Gawr Gura: {}", i); description = format!("Description Number {} The special edition offers significantly better quality with fewer issues. It's the model as initially envisioned.", i); description_rendered_html = rvc_markdown_description; @@ -155,7 +155,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT }, 21..=40 =>{ // TTS - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("dragonfruitGTv1: {}", i); description = format!("Description Number {}:", i); @@ -187,7 +187,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT }, 41..=60 => { // SD 1.5 - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("dragonfruitGTv1: {}", i); description = format!("Description Number {}:", i); @@ -232,7 +232,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT "#; let sdxl_image_token = "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/be706282-2978-42a0-aaa2-73881aad94e9/width=1024/00049-2287632957-1girl,face,curly%20hair,red%20hair,white%20background,.jpeg"; - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("SDXL_Niji_Special Edition: {}", i); description = format!("Description Number {}:", i); @@ -263,7 +263,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT "#; let lora_image_token = "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/123645df-dee2-4239-863a-76a150b09c32/width=1024/00000-2171948503.jpeg"; - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("Gawr Gura LoRA: {}", i); description = format!("Description Number {}:", i); @@ -295,7 +295,7 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT - Gawr Gura MOS 4 "#; - model_weight_token = ModelWeightToken(format!("{}", i)); + model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); title = format!("Gawr Gura: {}", i); description = format!("Description Number {} The special edition offers significantly better quality with fewer issues. It's the model as initially envisioned.", i); description_rendered_html = rvc_markdown_description; @@ -356,11 +356,13 @@ pub async fn seed_weights_for_user_token( ) -> AnyhowResult<()> { info!("Seeding weights..."); - let model_weight_token1 = ModelWeightToken("11".to_string()); - let model_weight_token2 = ModelWeightToken("22".to_string()); - let model_weight_token3 = ModelWeightToken("33".to_string()); - let model_weight_token4 = ModelWeightToken("44".to_string()); - let model_weight_token5 = ModelWeightToken("55".to_string()); + ModelWeightToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(54321); + + let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token4 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token5 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); let model_weights_args = vec![ CreateModelWeightsArgs { @@ -503,20 +505,20 @@ pub async fn seed_weights_for_user_token( pub async fn original_seed_weights(mysql_pool: &Pool,user_token: UserToken) -> AnyhowResult<()> { - let model_weight_token1 = ModelWeightToken("1".to_string()); - let creator_token1 = UserToken("creatorToken1".to_string()); + let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let creator_token1 = UserToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); - let model_weight_token2 = ModelWeightToken("2".to_string()); - let creator_token2 = UserToken("creatorToken2".to_string()); + let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let creator_token2 = UserToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); - let model_weight_token3 = ModelWeightToken("3".to_string()); - let creator_token3 = UserToken("creatorToken3".to_string()); + let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let creator_token3 = UserToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); - let model_weight_token4 = ModelWeightToken("4".to_string()); - let creator_token4 = UserToken("creatorToken4".to_string()); + let model_weight_token4 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let creator_token4 = UserToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); - let model_weight_token5 = ModelWeightToken("5".to_string()); - let creator_token5 = UserToken("creatorToken5".to_string()); + let model_weight_token5 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let creator_token5 = UserToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); let model_weights_args = vec![ CreateModelWeightsArgs { diff --git a/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs b/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs index 72ab9badd3..ac6388b877 100644 --- a/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs +++ b/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs @@ -1,35 +1,31 @@ use std::path::Path; -use log::{ info, warn }; -use sqlx::{ MySql, Pool }; - -use buckets::{ public::{ media_files::bucket_file_path::MediaFileBucketPath, self }, private }; +use log::{info, warn}; +use rand::Rng; +use sqlx::{MySql, Pool}; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use enums::{ + by_table::model_weights::{weights_category::WeightsCategory, weights_types::WeightsType}, + common::visibility::Visibility, +}; +use errors::AnyhowResult; use filesys::file_read_bytes::file_read_bytes; use filesys::file_size::file_size; use filesys::path_to_string::path_to_string; use hashing::sha256::sha256_hash_file::sha256_hash_file; use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; -use mimetypes::mimetype_for_file::get_mimetype_for_file; - -use storyteller_root::get_seed_tool_data_root; -use tokens::tokens::users::UserToken; -use crate::bucket_clients::BucketClients; -use crate::seeding::users::HANASHI_USER_TOKEN; - -use enums::{ - by_table::model_weights::{ weights_category::WeightsCategory, weights_types::WeightsType }, - common::visibility::Visibility, -}; - -use rand::Rng; -use errors::{ anyhow, AnyhowResult }; use mysql_queries::queries::model_weights::create_weight::{ create_weight, - CreateModelWeightsArgs, - self, + CreateModelWeightsArgs + , }; +use storyteller_root::get_seed_tool_data_root; use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +use crate::bucket_clients::BucketClients; +use crate::seeding::users::HANASHI_USER_TOKEN; pub async fn seed_weights_files( mysql_pool: &Pool, @@ -37,6 +33,9 @@ pub async fn seed_weights_files( ) -> AnyhowResult<()> { info!("Seeding weights files..."); + // NB: Stable token generation. + ModelWeightToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(123); + let seeded_weights = [ ( HANASHI_USER_TOKEN, @@ -87,9 +86,8 @@ pub async fn seed_weights_files( let mut rng = rand::thread_rng(); let n: u32 = rng.gen(); - let model_weight_token = ModelWeightToken(format!("{}", n)); // random token - let model_weight_token = model_weight_token; - + let model_weight_token = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let title = model_name; let description = "This is a description"; let description_rendered_html = "This is a description rendered html";