Skip to content

Commit

Permalink
Merge pull request #1142 from opentensor/fix/fix-hotkey-swap-alpha
Browse files Browse the repository at this point in the history 
Fix hotkey swap alpha
  • Loading branch information
@sam0x17
sam0x17 authored Jan 9, 2025
2 parents d9279fb + f3ed7d9 commit 789b3c7
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 71 deletions.
55 changes: 31 additions & 24 deletions pallets/subtensor/src/swap/swap_hotkey.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,36 @@ impl<T: Config> Pallet<T> {
}
});

// 11. Swap Alpha
// Alpha( hotkey, coldkey, netuid ) -> alpha
let old_alpha_values: Vec<((T::AccountId, u16), U64F64)> =
Alpha::<T>::iter_prefix((old_hotkey,)).collect();
// Clear the entire old prefix here.
let _ = Alpha::<T>::clear_prefix((old_hotkey,), old_alpha_values.len() as u32, None);
weight.saturating_accrue(T::DbWeight::get().reads(old_alpha_values.len() as u64));
weight.saturating_accrue(T::DbWeight::get().writes(old_alpha_values.len() as u64));

// Insert the new alpha values.
for ((coldkey, netuid), alpha) in old_alpha_values {
let new_alpha = Alpha::<T>::get((new_hotkey, &coldkey, netuid));
Alpha::<T>::insert(
(new_hotkey, &coldkey, netuid),
new_alpha.saturating_add(alpha),
);
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));

// Swap StakingHotkeys.
// StakingHotkeys( coldkey ) --> Vec<hotkey> -- the hotkeys that the coldkey stakes.
let mut staking_hotkeys = StakingHotkeys::<T>::get(&coldkey);
weight.saturating_accrue(T::DbWeight::get().reads(1));
if staking_hotkeys.contains(&old_hotkey) {
staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey);
staking_hotkeys.push(new_hotkey.clone());
StakingHotkeys::<T>::insert(&coldkey, staking_hotkeys);
weight.saturating_accrue(T::DbWeight::get().writes(1));
}
}

// 11. Swap Stake.
// Stake( hotkey, coldkey ) -> stake -- the stake that the hotkey controls on behalf of the coldkey.
let stakes: Vec<(T::AccountId, u64)> = Stake::<T>::iter_prefix(old_hotkey).collect();
Expand All @@ -322,30 +352,7 @@ impl<T: Config> Pallet<T> {
&coldkey,
new_stake_value.saturating_add(old_stake_amount),
);
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));

// 3.1 Swap Alpha
for netuid in Self::get_all_subnet_netuids() {
// Get the stake on the old (hot,coldkey) account.
let old_alpha: U64F64 = Alpha::<T>::get((&old_hotkey, coldkey.clone(), netuid));
// Get the stake on the new (hot,coldkey) account.
let new_alpha: U64F64 = Alpha::<T>::get((&new_hotkey, coldkey.clone(), netuid));
// Add the stake to new account.
Alpha::<T>::insert(
(&new_hotkey, coldkey.clone(), netuid),
new_alpha.saturating_add(old_alpha),
);
// Remove the value from the old account.
Alpha::<T>::remove((&old_hotkey, coldkey.clone(), netuid));
}

// Swap StakingHotkeys.
// StakingHotkeys( coldkey ) --> Vec<hotkey> -- the hotkeys that the coldkey stakes.
let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey.clone());
staking_hotkeys.retain(|hk| *hk != *old_hotkey && *hk != *new_hotkey);
staking_hotkeys.push(new_hotkey.clone());
StakingHotkeys::<T>::insert(coldkey.clone(), staking_hotkeys);
weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
weight.saturating_accrue(T::DbWeight::get().writes(1));
}

// 12. Swap ChildKeys.
Expand Down
170 changes: 123 additions & 47 deletions pallets/subtensor/src/tests/swap_hotkey.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -613,56 +613,132 @@ fn test_swap_hotkey_with_no_stake() {
});
}

// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --nocapture
// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::swap_hotkey::test_swap_hotkey_with_multiple_coldkeys_and_subnets --exact --show-output
#[test]
fn test_swap_hotkey_with_multiple_coldkeys_and_subnets() {
new_test_ext(1).execute_with(|| {
assert!(false);

// let old_hotkey = U256::from(1);
// let new_hotkey = U256::from(2);
// let coldkey1 = U256::from(3);
// let coldkey2 = U256::from(4);
// let netuid1 = 0;
// let netuid2 = 1;
// let mut weight = Weight::zero();

// // Set up initial state
// add_network(netuid1, 0, 1);
// add_network(netuid2, 0, 1);
// Owner::<Test>::insert(old_hotkey, coldkey1);
// Stake::<Test>::insert(old_hotkey, coldkey1, 100);
// Stake::<Test>::insert(old_hotkey, coldkey2, 200);
// IsNetworkMember::<Test>::insert(old_hotkey, netuid1, true);
// IsNetworkMember::<Test>::insert(old_hotkey, netuid2, true);
// TotalHotkeyStake::<Test>::insert(old_hotkey, 300);

// assert_ok!(SubtensorModule::perform_hotkey_swap(
// &old_hotkey,
// &new_hotkey,
// &coldkey1,
// &mut weight
// ));

// // Check ownership transfer
// assert!(!Owner::<Test>::contains_key(old_hotkey));
// assert_eq!(Owner::<Test>::get(new_hotkey), coldkey1);

// // Check stake transfer
// assert_eq!(Stake::<Test>::get(new_hotkey, coldkey1), 100);
// assert_eq!(Stake::<Test>::get(new_hotkey, coldkey2), 200);
// assert!(!Stake::<Test>::contains_key(old_hotkey, coldkey1));
// assert!(!Stake::<Test>::contains_key(old_hotkey, coldkey2));

// // Check subnet membership transfer
// assert!(IsNetworkMember::<Test>::get(new_hotkey, netuid1));
// assert!(IsNetworkMember::<Test>::get(new_hotkey, netuid2));
// assert!(!IsNetworkMember::<Test>::get(old_hotkey, netuid1));
// assert!(!IsNetworkMember::<Test>::get(old_hotkey, netuid2));

// // Check total stake transfer
// assert_eq!(TotalHotkeyStake::<Test>::get(new_hotkey), 300);
// assert!(!TotalHotkeyStake::<Test>::contains_key(old_hotkey));
let old_hotkey = U256::from(1);
let new_hotkey = U256::from(2);
let coldkey1 = U256::from(3);
let coldkey2 = U256::from(4);
let netuid1 = 1;
let netuid2 = 2;
let stake = 1_000_000_u64;
let mut weight = Weight::zero();

// Set up initial state
add_network(netuid1, 0, 1);
add_network(netuid2, 0, 1);
register_ok_neuron(netuid1, old_hotkey, coldkey1, 1234);
register_ok_neuron(netuid2, old_hotkey, coldkey1, 1234);

// Add balance to both coldkeys
SubtensorModule::add_balance_to_coldkey_account(&coldkey1, stake + 1_000);
SubtensorModule::add_balance_to_coldkey_account(&coldkey2, stake + 1_000);

// Stake with coldkey1
assert_ok!(SubtensorModule::add_stake(
<<Test as Config>::RuntimeOrigin>::signed(coldkey1),
old_hotkey,
netuid1,
stake
));

// Stake with coldkey2 also
assert_ok!(SubtensorModule::add_stake(
<<Test as Config>::RuntimeOrigin>::signed(coldkey2),
old_hotkey,
netuid2,
stake
));

let ck1_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&old_hotkey,
&coldkey1,
netuid1,
);
let ck2_stake = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&old_hotkey,
&coldkey2,
netuid2,
);
assert!(ck1_stake > 0);
assert!(ck2_stake > 0);
let total_hk_stake = SubtensorModule::get_total_stake_for_hotkey(&old_hotkey);
assert!(total_hk_stake > 0);

assert_ok!(SubtensorModule::perform_hotkey_swap(
&old_hotkey,
&new_hotkey,
&coldkey1,
&mut weight
));

// Check ownership transfer
assert_eq!(
SubtensorModule::get_owning_coldkey_for_hotkey(&new_hotkey),
coldkey1
);
assert!(!SubtensorModule::get_owned_hotkeys(&coldkey2).contains(&new_hotkey));

// Check stake transfer
assert_eq!(
SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&new_hotkey,
&coldkey1,
netuid1
),
ck1_stake
);
assert_eq!(
SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&new_hotkey,
&coldkey2,
netuid2
),
ck2_stake
);
assert_eq!(
SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&old_hotkey,
&coldkey1,
netuid1
),
0
);
assert_eq!(
SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
&old_hotkey,
&coldkey2,
netuid2
),
0
);

// Check subnet membership transfer
assert!(SubtensorModule::is_hotkey_registered_on_network(
netuid1,
&new_hotkey
));
assert!(SubtensorModule::is_hotkey_registered_on_network(
netuid2,
&new_hotkey
));
assert!(!SubtensorModule::is_hotkey_registered_on_network(
netuid1,
&old_hotkey
));
assert!(!SubtensorModule::is_hotkey_registered_on_network(
netuid2,
&old_hotkey
));

// Check total stake transfer
assert_eq!(
SubtensorModule::get_total_stake_for_hotkey(&new_hotkey),
total_hk_stake
);
assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&old_hotkey), 0);
});
}

Expand Down

0 comments on commit 789b3c7

Please sign in to comment.