Merge pull request #230 from storytold/brandon-run-rvc-svc-migration

migration cleanup

Cargo.lock

@@ -3519,6 +3519,7 @@ dependencies = [
  "easyenv",
  "enums",
  "errors",
+ "filesys",
  "hashing",
  "log",
  "my-workspace-hack",

crates/cli/migration-tool/Cargo.toml

@@ -20,6 +20,7 @@ config = { path = "../../lib/deprecated/config" }
 easyenv = { path = "../../lib/easyenv" }
 enums = { path = "../../_types/enums" }
 errors = { path = "../../lib/errors" }
+filesys = { path = "../../lib/files/filesys" }
 hashing = { path = "../../lib/files/hashing" }
 mysql_queries = { path = "../../lib/mysql_queries" }
 storyteller_root = { path = "../../lib/storyteller_root" }

crates/cli/migration-tool/src/migrations/voice_conversion_to_weights/copy_cloud_files.rs

@@ -4,6 +4,8 @@ use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath;
 use cloud_storage::bucket_path_unifier::BucketPathUnifier;
 use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType;
 use errors::{anyhow, AnyhowResult};
+use filesys::safe_delete_temp_directory::safe_delete_temp_directory;
+use filesys::safe_delete_temp_file::safe_delete_temp_file;
 use hashing::sha256::sha256_hash_file::sha256_hash_file;
 use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_voice_conversion_model::CopiedFileData;
 use mysql_queries::queries::voice_conversion::migration::list_whole_voice_conversion_models_using_cursor::WholeVoiceConversionModelRecord;
@@ -43,11 +45,14 @@ async fn copy_model(model: &WholeVoiceConversionModelRecord, deps: &Deps) -> Any
     VoiceConversionModelType::SoftVc => return Err(anyhow!("we never built softvc models")),
   };
 
-  deps.bucket_development_public.upload_filename_with_content_type(
+  deps.bucket_production_public.upload_filename_with_content_type(
     &new_model_bucket_path.get_full_object_path_str(),
     &model_temp_fs_path,
     "application/octet-stream").await?;
 
+  safe_delete_temp_file(&model_temp_fs_path);
+  safe_delete_temp_directory(&temp_dir);
+
   Ok(CopiedFileData {
     bucket_path: new_model_bucket_path,
     file_sha_hash: file_checksum,
@@ -68,10 +73,13 @@ async fn copy_index_file(model: &WholeVoiceConversionModelRecord, deps: &Deps, b
   let new_model_bucket_path =
       WeightFileBucketPath::rvc_index_file_from_object_hash(bucket_path.get_object_hash());
 
-  deps.bucket_development_public.upload_filename_with_content_type(
+  deps.bucket_production_public.upload_filename_with_content_type(
     &new_model_bucket_path.get_full_object_path_str(),
     &model_temp_fs_path,
     "application/octet-stream").await?;
 
+  safe_delete_temp_file(&model_temp_fs_path);
+  safe_delete_temp_directory(&temp_dir);
+
   Ok(()) // NB: We don't care about the path of the index file.
 }

crates/cli/migration-tool/src/migrations/voice_conversion_to_weights/migrate.rs

@@ -1,9 +1,15 @@
 use std::thread;
 use std::time::Duration;
 
+use log::{info, warn};
+
+use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath;
+use enums::by_table::model_weights::weights_types::WeightsType;
 use errors::AnyhowResult;
-use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_voice_conversion_model::upsert_model_weight_from_voice_conversion_model;
-use mysql_queries::queries::voice_conversion::migration::list_whole_voice_conversion_models_using_cursor::list_whole_voice_conversion_models_using_cursor;
+use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
+use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_voice_conversion_model::{CopiedFileData, upsert_model_weight_from_voice_conversion_model};
+use mysql_queries::queries::voice_conversion::migration::list_whole_voice_conversion_models_using_cursor::{list_whole_voice_conversion_models_using_cursor, WholeVoiceConversionModelRecord};
+use tokens::tokens::model_weights::ModelWeightToken;
 
 use crate::deps::Deps;
 use crate::migrations::voice_conversion_to_weights::copy_cloud_files::copy_cloud_files;
@@ -15,17 +21,19 @@ pub async fn migrate_voice_conversion_to_weights(deps: &Deps) -> AnyhowResult<()
   let mut cursor = 0;
 
   loop {
+    info!("Querying {PAGE_SIZE} models at cursor = {cursor}");
+
     let results
         = list_whole_voice_conversion_models_using_cursor(&deps.mysql_production, PAGE_SIZE, cursor).await?;
 
     if results.is_empty() {
+      warn!("No more results found; exiting.");
       break;
     }
 
     for result in results.iter() {
-      println!("result: {:?}\n\n", result);
-      let copied_data = copy_cloud_files(result, &deps).await?;
-      upsert_model_weight_from_voice_conversion_model(result, &deps.mysql_development, &copied_data).await?;
+      println!("\n\nmigrating over result {:?} : {:?} ...", result.token, result);
+      copy_model_record_and_bucket_files(result, &deps).await?;
     }
 
     if let Some(last_id) = results.last().map(|result| result.id) {
@@ -37,3 +45,48 @@ pub async fn migrate_voice_conversion_to_weights(deps: &Deps) -> AnyhowResult<()
 
   Ok(())
 }
+
+async fn copy_model_record_and_bucket_files(record: &WholeVoiceConversionModelRecord, deps: &Deps) -> AnyhowResult<()> {
+  let mut maybe_copied_data = None;
+
+  match &record.maybe_migration_new_model_weights_token {
+    None => {
+      info!("Copying bucket files for record {:?} ...", record.token);
+      let result = copy_cloud_files(record, &deps).await?;
+      maybe_copied_data = Some(result);
+    }
+    Some(token) => {
+      info!("Existing migrated record; skipping cloud bucket copy...");
+      maybe_copied_data = lookup_bucket_details_from_record(token, &deps).await?;
+    }
+  }
+
+  if let Some(copied_data) = &maybe_copied_data {
+    info!("Upserting records...");
+    upsert_model_weight_from_voice_conversion_model(record, &deps.mysql_production, copied_data).await?;
+  }
+
+  Ok(())
+}
+
+/// Look up the bucket storage details from a previously-migrated record
+async fn lookup_bucket_details_from_record(model_weight_token: &ModelWeightToken, deps: &Deps) -> AnyhowResult<Option<CopiedFileData>> {
+  let maybe_model_weight_record =
+      get_weight_by_token(model_weight_token, true, &deps.mysql_production).await?;
+
+  let model_weight_record = match maybe_model_weight_record {
+    None => return Ok(None),
+    Some(record) => record,
+  };
+
+  let bucket_path = match model_weight_record.weights_type {
+    WeightsType::RvcV2 => WeightFileBucketPath::rvc_model_file_from_object_hash(&model_weight_record.private_bucket_hash),
+    WeightsType::SoVitsSvc => WeightFileBucketPath::svc_model_file_from_object_hash(&model_weight_record.private_bucket_hash),
+    _ => return Ok(None), // NB: Stable Diffusion, etc. aren't valid for migration
+  };
+
+  Ok(Some(CopiedFileData {
+    bucket_path,
+    file_sha_hash: model_weight_record.file_checksum_sha2,
+  }))
+}

crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_voice_conversion_model.rs

@@ -32,8 +32,7 @@ pub async fn upsert_model_weight_from_voice_conversion_model(
   upsert_model_weights_record(record, &model_weight_token, copied_data, &mut transaction).await?;
   upsert_model_weights_extension_record(record, &model_weight_token, &mut transaction).await?;
 
-  // TODO: Don't update source record with new token *unless* we're backfilling the same database.
-  // update_original_record(record, &model_weight_token, &mut transaction).await?;
+  update_original_record(record, &model_weight_token, &mut transaction).await?;
 
   transaction.commit().await?;
 
@@ -227,6 +226,18 @@ pub async fn update_original_record(
   model_weight_token: &ModelWeightToken,
   transaction: &mut Transaction<'_, MySql>
 ) -> AnyhowResult<()> {
+  let query = sqlx::query!(
+        r#"
+UPDATE voice_conversion_models
+SET
+  maybe_migration_new_model_weights_token = ?
+WHERE token = ?
+        "#,
+      model_weight_token,
+      record.token,
+    );
+
+  let _r = query.execute(&mut **transaction).await?;
 
   Ok(())
 }