Merge pull request #1305 from opentensor/fix/allow-replace-sn-owner-hotkey

Use SubnetOwnerHotkey for neuron replace/prune checks

Author: sam0x17

pallets/subtensor/src/subnets/registration.rs

@@ -2,7 +2,6 @@ use super::*;
 use sp_core::{H256, U256};
 use sp_io::hashing::{keccak_256, sha2_256};
 use sp_runtime::Saturating;
-use substrate_fixed::types::I64F64;
 use system::pallet_prelude::BlockNumberFor;
 
 const LOG_TARGET: &str = "runtime::subtensor::registration";
@@ -425,29 +424,11 @@ impl<T: Config> Pallet<T> {
             return 0; // If there are no neurons in this network.
         }
 
-        // Get SN owner top stake hotkey
-        let mut top_stake_sn_owner_hotkey: Option<T::AccountId> = None;
-        let mut max_stake_weight: I64F64 = I64F64::from_num(-1);
         for neuron_uid in 0..neurons_n {
+            // Do not deregister the owner's hotkey from the `SubnetOwnerHotkey` map
             if let Ok(hotkey) = Self::get_hotkey_for_net_and_uid(netuid, neuron_uid) {
-                let coldkey = Self::get_owning_coldkey_for_hotkey(&hotkey);
-                if Self::get_subnet_owner(netuid) != coldkey {
-                    continue;
-                }
-
-                let stake_weights = Self::get_stake_weights_for_hotkey_on_subnet(&hotkey, netuid);
-                if stake_weights.0 > max_stake_weight {
-                    max_stake_weight = stake_weights.0;
-                    top_stake_sn_owner_hotkey = Some(hotkey);
-                }
-            }
-        }
-
-        for neuron_uid in 0..neurons_n {
-            // Do not deregister the owner's top-stake hotkey
-            if let Ok(hotkey) = Self::get_hotkey_for_net_and_uid(netuid, neuron_uid) {
-                if let Some(ref top_sn_owner_hotkey) = top_stake_sn_owner_hotkey {
-                    if top_sn_owner_hotkey == &hotkey {
+                if let Ok(top_sn_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {
+                    if top_sn_owner_hotkey == hotkey {
                         continue;
                     }
                 }

pallets/subtensor/src/subnets/uids.rs

@@ -42,18 +42,19 @@ impl<T: Config> Pallet<T> {
         // 1. Get the old hotkey under this position.
         let old_hotkey: T::AccountId = Keys::<T>::get(netuid, uid_to_replace);
 
-        // Do not deregister the owner
-        let coldkey = Self::get_owning_coldkey_for_hotkey(&old_hotkey);
-        if Self::get_subnet_owner(netuid) == coldkey {
-            log::warn!(
-                "replace_neuron: Skipped replacement because neuron belongs to the subnet owner. \
-                 netuid: {:?}, uid_to_replace: {:?}, new_hotkey: {:?}, owner_coldkey: {:?}",
-                netuid,
-                uid_to_replace,
-                new_hotkey,
-                coldkey
-            );
-            return;
+        // Do not replace owner hotkey from `SubnetOwnerHotkey`
+        if let Ok(sn_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {
+            if sn_owner_hotkey == old_hotkey.clone() {
+                log::warn!(
+                    "replace_neuron: Skipped replacement because neuron is the subnet owner hotkey. \
+                    netuid: {:?}, uid_to_replace: {:?}, new_hotkey: {:?}, owner_hotkey: {:?}",
+                    netuid,
+                    uid_to_replace,
+                    new_hotkey,
+                    sn_owner_hotkey
+                );
+                return;
+            }
         }
 
         // 2. Remove previous set memberships.

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -437,6 +437,18 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // 15. Swap SubnetOwnerHotkey
+        // SubnetOwnerHotkey( netuid ) --> hotkey -- the hotkey that is the owner of the subnet.
+        for netuid in Self::get_all_subnet_netuids() {
+            if let Ok(old_subnet_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {
+                weight.saturating_accrue(T::DbWeight::get().reads(1));
+                if old_subnet_owner_hotkey == *old_hotkey {
+                    SubnetOwnerHotkey::<T>::insert(netuid, new_hotkey);
+                    weight.saturating_accrue(T::DbWeight::get().writes(1));
+                }
+            }
+        }
+
         // Return successful after swapping all the relevant terms.
         Ok(())
     }

pallets/subtensor/src/tests/swap_hotkey.rs

@@ -1359,3 +1359,25 @@ fn test_swap_child_hotkey_childkey_maps() {
         );
     })
 }
+
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_hotkey -- test_swap_hotkey_is_sn_owner_hotkey --exact --nocapture
+#[test]
+fn test_swap_hotkey_is_sn_owner_hotkey() {
+    new_test_ext(1).execute_with(|| {
+        let old_hotkey = U256::from(1);
+        let new_hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let mut weight = Weight::zero();
+
+        // Create dynamic network
+        let netuid = add_dynamic_network(&old_hotkey, &coldkey);
+        // Check for SubnetOwnerHotkey
+        assert_eq!(SubnetOwnerHotkey::<Test>::get(netuid), old_hotkey);
+
+        // Perform the swap
+        SubtensorModule::perform_hotkey_swap(&old_hotkey, &new_hotkey, &coldkey, &mut weight);
+
+        // Check for SubnetOwnerHotkey
+        assert_eq!(SubnetOwnerHotkey::<Test>::get(netuid), new_hotkey);
+    });
+}

pallets/subtensor/src/tests/uids.rs

@@ -5,7 +5,6 @@ use crate::*;
 use frame_support::{assert_err, assert_ok};
 use frame_system::Config;
 use sp_core::U256;
-use substrate_fixed::types::I64F64;
 
 /********************************************
     tests for uids.rs file
@@ -281,10 +280,11 @@ fn test_replace_neuron_subnet_owner_not_replaced() {
 }
 
 #[test]
-fn test_get_neuron_to_prune_owner_not_pruned() {
+fn test_replace_neuron_subnet_owner_not_replaced_if_in_sn_owner_hotkey_map() {
     new_test_ext(1).execute_with(|| {
         let owner_hotkey = U256::from(123);
         let owner_coldkey = U256::from(999);
+        let other_owner_hotkey = U256::from(456);
 
         let netuid = add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
@@ -295,54 +295,64 @@ fn test_get_neuron_to_prune_owner_not_pruned() {
         let owner_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &owner_hotkey)
             .expect("Owner neuron should already be registered by add_dynamic_network");
 
-        let additional_hotkey_1 = U256::from(1000);
-        let additional_coldkey_1 = U256::from(2000);
+        // Register another hotkey for the owner
+        register_ok_neuron(netuid, other_owner_hotkey, owner_coldkey, 0);
+        let other_owner_uid =
+            SubtensorModule::get_uid_for_net_and_hotkey(netuid, &other_owner_hotkey)
+                .expect("Should be registered");
 
+        let additional_hotkey_1 = U256::from(1000);
         let additional_hotkey_2 = U256::from(1001);
-        let additional_coldkey_2 = U256::from(2001);
-
-        register_ok_neuron(netuid, additional_hotkey_1, additional_coldkey_1, 0);
-        let uid_1 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_1)
-            .expect("Should be registered");
 
-        register_ok_neuron(netuid, additional_hotkey_2, additional_coldkey_2, 1);
-        let uid_2 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_2)
-            .expect("Should be registered");
-
-        SubtensorModule::set_pruning_score_for_uid(netuid, owner_uid, 0);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_1, 1);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_2, 2);
-
-        let pruned_uid = SubtensorModule::get_neuron_to_prune(netuid);
+        let current_block = SubtensorModule::get_current_block_as_u64();
+        SubtensorModule::replace_neuron(netuid, owner_uid, &additional_hotkey_1, current_block);
 
-        // - The pruned UID must be `uid_1` (score=1).
-        // - The owner's UID remains unpruned.
+        let still_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &owner_hotkey);
+        assert_ok!(still_uid);
         assert_eq!(
-            pruned_uid, uid_1,
-            "Should prune the neuron with pruning score=1, not the owner (score=0)."
+            still_uid.unwrap(),
+            owner_uid,
+            "Owner's first hotkey should remain registered"
         );
 
-        let pruned_score = SubtensorModule::get_pruning_score_for_uid(netuid, uid_1);
-        assert_eq!(
-            pruned_score,
-            u16::MAX,
-            "Pruned neuron's score should be set to u16::MAX"
-        );
+        let new_key_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_1);
+        assert_err!(new_key_uid, Error::<Test>::HotKeyNotRegisteredInSubNet,);
 
-        let owner_score = SubtensorModule::get_pruning_score_for_uid(netuid, owner_uid);
-        assert_eq!(
-            owner_score, 0,
-            "Owner's pruning score remains 0, indicating it was skipped"
+        // Try to replace the other owner hotkey
+        SubtensorModule::replace_neuron(
+            netuid,
+            other_owner_uid,
+            &additional_hotkey_1,
+            current_block,
         );
+        let still_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &other_owner_hotkey);
+        assert_err!(still_uid, Error::<Test>::HotKeyNotRegisteredInSubNet,); // Was replaced
+
+        // Re-register this hotkey
+        register_ok_neuron(netuid, other_owner_hotkey, owner_coldkey, 0);
+        let _other_owner_uid =
+            SubtensorModule::get_uid_for_net_and_hotkey(netuid, &other_owner_hotkey)
+                .expect("Should be registered");
+
+        // Set this hotkey as the SubnetOwnerHotkey
+        SubnetOwnerHotkey::<Test>::insert(netuid, other_owner_hotkey);
+
+        SubtensorModule::replace_neuron(netuid, owner_uid, &additional_hotkey_2, current_block);
+
+        // The owner's first hotkey should be replaceable; it's not the top-stake owner hotkey
+        let still_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &owner_hotkey);
+        assert_err!(still_uid, Error::<Test>::HotKeyNotRegisteredInSubNet,);
+
+        let new_key_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_2);
+        assert_ok!(new_key_uid);
     });
 }
 
 #[test]
-fn test_get_neuron_to_prune_owner_pruned_if_not_top_stake_owner_hotkey() {
+fn test_get_neuron_to_prune_owner_not_pruned() {
     new_test_ext(1).execute_with(|| {
         let owner_hotkey = U256::from(123);
         let owner_coldkey = U256::from(999);
-        let other_owner_hotkey = U256::from(456);
 
         let netuid = add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
@@ -353,68 +363,54 @@ fn test_get_neuron_to_prune_owner_pruned_if_not_top_stake_owner_hotkey() {
         let owner_uid = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &owner_hotkey)
             .expect("Owner neuron should already be registered by add_dynamic_network");
 
-        // Register another hotkey for the owner
-        register_ok_neuron(netuid, other_owner_hotkey, owner_coldkey, 0);
-        let other_owner_uid =
-            SubtensorModule::get_uid_for_net_and_hotkey(netuid, &other_owner_hotkey)
-                .expect("Should be registered");
-
         let additional_hotkey_1 = U256::from(1000);
         let additional_coldkey_1 = U256::from(2000);
 
         let additional_hotkey_2 = U256::from(1001);
         let additional_coldkey_2 = U256::from(2001);
 
-        register_ok_neuron(netuid, additional_hotkey_1, additional_coldkey_1, 1);
-        let uid_2 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_1)
+        register_ok_neuron(netuid, additional_hotkey_1, additional_coldkey_1, 0);
+        let uid_1 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_1)
             .expect("Should be registered");
 
-        register_ok_neuron(netuid, additional_hotkey_2, additional_coldkey_2, 2);
-        let uid_3 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_2)
+        register_ok_neuron(netuid, additional_hotkey_2, additional_coldkey_2, 1);
+        let uid_2 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_2)
             .expect("Should be registered");
 
         SubtensorModule::set_pruning_score_for_uid(netuid, owner_uid, 0);
-        // Other owner key has pruning score not worse than the owner's first hotkey, but worse than the additional hotkeys
-        SubtensorModule::set_pruning_score_for_uid(netuid, other_owner_uid, 1);
+        SubtensorModule::set_pruning_score_for_uid(netuid, uid_1, 1);
         SubtensorModule::set_pruning_score_for_uid(netuid, uid_2, 2);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_3, 3);
 
         let pruned_uid = SubtensorModule::get_neuron_to_prune(netuid);
-        assert_eq!(pruned_uid, other_owner_uid, "Should prune the owner");
 
-        // Give the owner's other hotkey some stake
-        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
-            &other_owner_hotkey,
-            &owner_coldkey,
-            netuid,
-            1000,
+        // - The pruned UID must be `uid_1` (score=1).
+        // - The owner's UID remains unpruned.
+        assert_eq!(
+            pruned_uid, uid_1,
+            "Should prune the neuron with pruning score=1, not the owner (score=0)."
         );
 
-        // Reset pruning scores
-        SubtensorModule::set_pruning_score_for_uid(netuid, owner_uid, 0);
-        SubtensorModule::set_pruning_score_for_uid(netuid, other_owner_uid, 1);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_2, 2);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_3, 3);
-
-        let pruned_uid = SubtensorModule::get_neuron_to_prune(netuid);
+        let pruned_score = SubtensorModule::get_pruning_score_for_uid(netuid, uid_1);
+        assert_eq!(
+            pruned_score,
+            u16::MAX,
+            "Pruned neuron's score should be set to u16::MAX"
+        );
 
-        // - The pruned UID must be `uid_1` (score=1).
-        // - The owner's UID remains unpruned.
+        let owner_score = SubtensorModule::get_pruning_score_for_uid(netuid, owner_uid);
         assert_eq!(
-            pruned_uid, owner_uid,
-            "Should prune the owner, not the top-stake owner hotkey and not the additional hotkeys"
+            owner_score, 0,
+            "Owner's pruning score remains 0, indicating it was skipped"
         );
     });
 }
 
 #[test]
-fn test_get_neuron_to_prune_owner_pruned_if_not_top_stake_owner_hotkey_chk() {
+fn test_get_neuron_to_prune_owner_pruned_if_not_in_sn_owner_hotkey_map() {
     new_test_ext(1).execute_with(|| {
         let owner_hotkey = U256::from(123);
         let owner_coldkey = U256::from(999);
         let other_owner_hotkey = U256::from(456);
-        let parent_hotkey = U256::from(4567);
-        let parent_coldkey = U256::from(4568);
 
         let netuid = add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
@@ -445,60 +441,26 @@ fn test_get_neuron_to_prune_owner_pruned_if_not_top_stake_owner_hotkey_chk() {
         let uid_3 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &additional_hotkey_2)
             .expect("Should be registered");
 
-        register_ok_neuron(netuid, parent_hotkey, parent_coldkey, 3);
-        let uid_4: u16 = SubtensorModule::get_uid_for_net_and_hotkey(netuid, &parent_hotkey)
-            .expect("Should be registered");
-
-        // Give parent key some stake
-        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
-            &parent_hotkey,
-            &parent_coldkey,
-            netuid,
-            10_000_000,
-        );
-
         SubtensorModule::set_pruning_score_for_uid(netuid, owner_uid, 0);
         // Other owner key has pruning score not worse than the owner's first hotkey, but worse than the additional hotkeys
         SubtensorModule::set_pruning_score_for_uid(netuid, other_owner_uid, 1);
         SubtensorModule::set_pruning_score_for_uid(netuid, uid_2, 2);
         SubtensorModule::set_pruning_score_for_uid(netuid, uid_3, 3);
 
-        // Ensure parent key is not pruned
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_4, 10_000);
-
         let pruned_uid = SubtensorModule::get_neuron_to_prune(netuid);
-        assert_eq!(
-            pruned_uid, other_owner_uid,
-            "Should prune the owner's other hotkey"
-        );
+        assert_eq!(pruned_uid, other_owner_uid, "Should prune the owner");
 
-        // Give the owner's other hotkey some CHK stake; Doesn't need to be much
-        mock_set_children_no_epochs(
-            netuid,
-            &parent_hotkey,
-            &[(
-                I64F64::saturating_from_num(0.1)
-                    .saturating_mul(I64F64::saturating_from_num(u64::MAX))
-                    .saturating_to_num::<u64>(),
-                other_owner_hotkey,
-            )],
-        );
-        // Check stake weight of other_owner_hotkey
-        let stake_weight =
-            SubtensorModule::get_stake_weights_for_hotkey_on_subnet(&other_owner_hotkey, netuid);
-        assert!(stake_weight.0 > 0);
+        // Set the owner's other hotkey as the SubnetOwnerHotkey
+        SubnetOwnerHotkey::<Test>::insert(netuid, other_owner_hotkey);
 
         // Reset pruning scores
         SubtensorModule::set_pruning_score_for_uid(netuid, owner_uid, 0);
         SubtensorModule::set_pruning_score_for_uid(netuid, other_owner_uid, 1);
         SubtensorModule::set_pruning_score_for_uid(netuid, uid_2, 2);
         SubtensorModule::set_pruning_score_for_uid(netuid, uid_3, 3);
-        SubtensorModule::set_pruning_score_for_uid(netuid, uid_4, 10_000);
 
         let pruned_uid = SubtensorModule::get_neuron_to_prune(netuid);
 
-        // - The pruned UID must be `uid_1` (score=1).
-        // - The owner's UID remains unpruned.
         assert_eq!(
             pruned_uid, owner_uid,
             "Should prune the owner, not the top-stake owner hotkey and not the additional hotkeys"

runtime/src/lib.rs

@@ -229,7 +229,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
     //   `spec_version`, and `authoring_version` are the same between Wasm and native.
     // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
     //   the compatible custom types.
-    spec_version: 237,
+    spec_version: 238,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,