From 5d77e7ad4a3c47e4399d393004dfa9a69ed06593 Mon Sep 17 00:00:00 2001 From: Greg Zaitsev Date: Tue, 15 Oct 2024 17:37:42 -0400 Subject: [PATCH] Merge old and new lists of child keys when swapping hotkey --- pallets/subtensor/src/lib.rs | 6 + pallets/subtensor/src/staking/set_children.rs | 6 +- pallets/subtensor/src/swap/swap_hotkey.rs | 78 ++++++- pallets/subtensor/tests/swap_hotkey.rs | 204 ++++++++++++++++++ 4 files changed, 289 insertions(+), 5 deletions(-) diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 6c8da7a69..5c7c2557c 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -208,6 +208,12 @@ pub mod pallet { T::InitialMinDelegateTake::get() } + #[pallet::type_value] + /// Default limit for number of childkeys on one hotkey + pub fn DefaultChildKeyLimit() -> u16 { + 5 + } + #[pallet::type_value] /// Default minimum childkey take. pub fn DefaultMinChildKeyTake() -> u16 { diff --git a/pallets/subtensor/src/staking/set_children.rs b/pallets/subtensor/src/staking/set_children.rs index 9029d58ca..897bc004c 100644 --- a/pallets/subtensor/src/staking/set_children.rs +++ b/pallets/subtensor/src/staking/set_children.rs @@ -1,4 +1,5 @@ use super::*; +use frame_support::traits::Get; impl Pallet { /// ---- The implementation for the extrinsic do_set_child_singular: Sets a single child. @@ -96,7 +97,10 @@ impl Pallet { ); // --- 4.1. Ensure that the number of children does not exceed 5. - ensure!(children.len() <= 5, Error::::TooManyChildren); + ensure!( + children.len() <= DefaultChildKeyLimit::::get() as usize, + Error::::TooManyChildren + ); // --- 5. Ensure that each child is not the hotkey. for (_, child_i) in &children { diff --git a/pallets/subtensor/src/swap/swap_hotkey.rs b/pallets/subtensor/src/swap/swap_hotkey.rs index 54095d5fb..1addde549 100644 --- a/pallets/subtensor/src/swap/swap_hotkey.rs +++ b/pallets/subtensor/src/swap/swap_hotkey.rs @@ -1,6 +1,7 @@ use super::*; use frame_support::weights::Weight; use sp_core::Get; +use substrate_fixed::types::I96F32; impl Pallet { /// Swaps the hotkey of a coldkey account. @@ -82,7 +83,7 @@ impl Pallet { Self::burn_tokens(actual_burn_amount); // 14. Perform the hotkey swap - let _ = Self::perform_hotkey_swap(old_hotkey, new_hotkey, &coldkey, &mut weight); + Self::perform_hotkey_swap(old_hotkey, new_hotkey, &coldkey, &mut weight)?; // 15. Update the last transaction block for the coldkey Self::set_last_tx_block(&coldkey, block); @@ -321,11 +322,80 @@ impl Pallet { // ChildKeys( parent, netuid ) --> Vec<(proportion,child)> -- the child keys of the parent. for netuid in Self::get_all_subnet_netuids() { // Get the children of the old hotkey for this subnet - let my_children: Vec<(u64, T::AccountId)> = ChildKeys::::get(old_hotkey, netuid); + let old_hotkey_children: Vec<(u64, T::AccountId)> = + ChildKeys::::get(old_hotkey, netuid); + // Read old children of the new hotkey + let new_hotkey_children: Vec<(u64, T::AccountId)> = + ChildKeys::::get(new_hotkey, netuid); + // Merged childkey list (initialize with old hotkey children) + let mut merged_childkey_list: Vec<(u64, T::AccountId)> = + ChildKeys::::get(old_hotkey, netuid); + + // Calculate total of proportions between two lists being merged (old and new) + let proportion_total: I96F32 = I96F32::from_num( + old_hotkey_children + .iter() + .map(|(proportion, _)| *proportion as u128) + .sum::(), + ) + .saturating_add(I96F32::from_num( + new_hotkey_children + .iter() + .map(|(proportion, _)| *proportion as u128) + .sum::(), + )); + + // Merge old and new child lists. Do not use HashSet because we only have 5 keys per list, + // and iteration is more efficient. + new_hotkey_children.iter().for_each(|new_hotkey_child| { + if !merged_childkey_list + .iter() + .any(|merged_hotkey_child| merged_hotkey_child.1 == new_hotkey_child.1) + { + merged_childkey_list.push(new_hotkey_child.clone()); + } + }); + + if merged_childkey_list.len() > DefaultChildKeyLimit::::get() as usize { + return Err(Error::::TooManyChildren.into()); + } + + // Update proportions in the merged list + // Re-normalize child proportions if total of proportions exceeds 1.0 + // (because two child lists were merged and we can't go over 1.0 total proportion) + let proportion_normalization_coeff: I96F32 = if proportion_total > u64::MAX { + I96F32::from_num(u64::MAX).saturating_div(proportion_total) + } else { + I96F32::from_num(1) + }; + merged_childkey_list.iter_mut().for_each(|child| { + let proportion_from_old_hotkey = old_hotkey_children + .iter() + .find(|&old_hotkey_child| old_hotkey_child.1 == child.1) + .map(|old_hotkey_child| old_hotkey_child.0) + .unwrap_or(0_u64); + let proportion_from_new_hotkey = new_hotkey_children + .iter() + .find(|&new_hotkey_child| new_hotkey_child.1 == child.1) + .map(|new_hotkey_child| new_hotkey_child.0) + .unwrap_or(0_u64); + + let mut normalized_proportion: I96F32 = + I96F32::from_num(proportion_from_old_hotkey) + .saturating_mul(proportion_normalization_coeff); + normalized_proportion = normalized_proportion.saturating_add( + I96F32::from_num(proportion_from_new_hotkey) + .saturating_mul(proportion_normalization_coeff), + ); + + *child = (normalized_proportion.to_num::(), child.1.clone()); + }); + // Remove the old hotkey's child entries ChildKeys::::remove(old_hotkey, netuid); - // Insert the same child entries for the new hotkey - ChildKeys::::insert(new_hotkey, netuid, my_children); + + // Insert the new hotkey's child entries + ChildKeys::::insert(new_hotkey, netuid, merged_childkey_list); } // 13. Swap ParentKeys. diff --git a/pallets/subtensor/tests/swap_hotkey.rs b/pallets/subtensor/tests/swap_hotkey.rs index 57f206452..f7559fc3c 100644 --- a/pallets/subtensor/tests/swap_hotkey.rs +++ b/pallets/subtensor/tests/swap_hotkey.rs @@ -985,6 +985,210 @@ fn test_swap_child_keys() { }); } +#[test] +fn test_swap_existing_child_keys() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let children1 = vec![(100u64, U256::from(4)), (200u64, U256::from(5))]; + let children2 = vec![(100u64, U256::from(6)), (200u64, U256::from(7))]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children1.clone()); + ChildKeys::::insert(new_hotkey, netuid, children2.clone()); + + // Perform the swap + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight); + + // Verify the swap + assert_eq!( + ChildKeys::::get(new_hotkey, netuid), + [ + (100u64, U256::from(4)), + (200u64, U256::from(5)), + (100u64, U256::from(6)), + (200u64, U256::from(7)), + ] + ); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + }); +} + +#[test] +fn test_swap_existing_child_keys_too_many_children() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let children1 = vec![ + (100u64, U256::from(4)), + (200u64, U256::from(5)), + (200u64, U256::from(6)), + ]; + let children2 = vec![ + (100u64, U256::from(7)), + (200u64, U256::from(8)), + (200u64, U256::from(9)), + ]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children1.clone()); + ChildKeys::::insert(new_hotkey, netuid, children2.clone()); + + // Perform the swap + assert_err!( + SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight), + Error::::TooManyChildren, + ); + + // Verify no change + assert_eq!(ChildKeys::::get(old_hotkey, netuid), children1); + assert_eq!(ChildKeys::::get(new_hotkey, netuid), children2); + }); +} + +#[test] +fn test_swap_existing_child_keys_too_many_children_extrinsic() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 1u16; + let swap_cost = 1_000_000_000u64 * 2; + let children1 = vec![ + (100u64, U256::from(4)), + (200u64, U256::from(5)), + (200u64, U256::from(6)), + ]; + let children2 = vec![ + (100u64, U256::from(7)), + (200u64, U256::from(8)), + (200u64, U256::from(9)), + ]; + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + register_ok_neuron(netuid, old_hotkey, coldkey, 0); + SubtensorModule::add_balance_to_coldkey_account(&coldkey, swap_cost); + ChildKeys::::insert(old_hotkey, netuid, children1.clone()); + ChildKeys::::insert(new_hotkey, netuid, children2.clone()); + + // Perform the swap using high level function + assert_err!( + SubtensorModule::do_swap_hotkey( + <::RuntimeOrigin>::signed(coldkey), + &old_hotkey, + &new_hotkey, + ), + Error::::TooManyChildren + ); + + // Verify no change + assert_eq!(ChildKeys::::get(old_hotkey, netuid), children1); + assert_eq!(ChildKeys::::get(new_hotkey, netuid), children2); + }); +} + +#[test] +fn test_swap_existing_child_keys_proper_merge() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let children1 = vec![ + (100u64, U256::from(4)), + (200u64, U256::from(5)), + (300u64, U256::from(6)), + ]; + let children2 = vec![ + (100u64, U256::from(6)), + (200u64, U256::from(7)), + (300u64, U256::from(8)), + ]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children1.clone()); + ChildKeys::::insert(new_hotkey, netuid, children2.clone()); + + // Perform the swap + assert_ok!(SubtensorModule::perform_hotkey_swap( + &old_hotkey, + &new_hotkey, + &coldkey, + &mut weight + ),); + + // Verify the swap + assert_eq!( + ChildKeys::::get(new_hotkey, netuid), + [ + (100u64, U256::from(4)), + (200u64, U256::from(5)), + (400u64, U256::from(6)), + (200u64, U256::from(7)), + (300u64, U256::from(8)), + ] + ); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + }); +} + +#[test] +fn test_swap_existing_child_keys_proportion_overflow() { + new_test_ext(1).execute_with(|| { + let old_hotkey = U256::from(1); + let new_hotkey = U256::from(2); + let coldkey = U256::from(3); + let netuid = 0u16; + let children1 = vec![ + (u64::MAX / 5, U256::from(4)), + (u64::MAX / 5 * 4, U256::from(5)), + ]; + let children2 = vec![ + (u64::MAX / 5 * 4, U256::from(5)), + (u64::MAX / 5, U256::from(6)), + ]; + let mut weight = Weight::zero(); + + // Initialize ChildKeys for old_hotkey + add_network(netuid, 1, 0); + ChildKeys::::insert(old_hotkey, netuid, children1.clone()); + ChildKeys::::insert(new_hotkey, netuid, children2.clone()); + + // Perform the swap + assert_ok!(SubtensorModule::perform_hotkey_swap( + &old_hotkey, + &new_hotkey, + &coldkey, + &mut weight + ),); + + // Verify the swap + // key 4: max / 5 renormalized = max / 10 + // key 5: 8 max / 5 renormalized = 4 max / 5 + // key 6: max / 5 renormalized = max / 10 + assert_eq!( + ChildKeys::::get(new_hotkey, netuid), + [ + (u64::MAX / 10, U256::from(4)), + (u64::MAX / 5 * 4, U256::from(5)), + (u64::MAX / 10, U256::from(6)), + ] + ); + assert!(ChildKeys::::get(old_hotkey, netuid).is_empty()); + }); +} + // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_parent_keys --exact --nocapture #[test] fn test_swap_parent_keys() {