Fix immune owner UIDs

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -428,7 +428,7 @@ impl<T: Config> Pallet<T> {
         (prop_alpha_dividends, tao_dividends)
     }
 
-    fn get_immune_owner_hotkeys(netuid: NetUid, coldkey: &T::AccountId) -> Vec<T::AccountId> {
+    fn get_owner_hotkeys(netuid: NetUid, coldkey: &T::AccountId) -> Vec<T::AccountId> {
         // Gather (block, uid, hotkey) only for hotkeys that have a UID and a registration block.
         let mut triples: Vec<(u64, u16, T::AccountId)> = OwnedHotkeys::<T>::get(coldkey)
             .into_iter()
@@ -445,28 +445,19 @@ impl<T: Config> Pallet<T> {
         // Recent registration is priority so that we can let older keys expire (get non-immune)
         triples.sort_by(|(b1, u1, _), (b2, u2, _)| b2.cmp(b1).then(u1.cmp(u2)));
 
-        // Keep first ImmuneOwnerUidsLimit
-        let limit = ImmuneOwnerUidsLimit::<T>::get(netuid).into();
-        if triples.len() > limit {
-            triples.truncate(limit);
-        }
-
         // Project to just hotkeys
-        let mut immune_hotkeys: Vec<T::AccountId> =
+        let mut owner_hotkeys: Vec<T::AccountId> =
             triples.into_iter().map(|(_, _, hk)| hk).collect();
 
         // Insert subnet owner hotkey in the beginning of the list if valid and not
         // already present
         if let Ok(owner_hk) = SubnetOwnerHotkey::<T>::try_get(netuid) {
-            if Uids::<T>::get(netuid, &owner_hk).is_some() && !immune_hotkeys.contains(&owner_hk) {
-                immune_hotkeys.insert(0, owner_hk);
-                if immune_hotkeys.len() > limit {
-                    immune_hotkeys.truncate(limit);
-                }
+            if Uids::<T>::get(netuid, &owner_hk).is_some() && !owner_hotkeys.contains(&owner_hk) {
+                owner_hotkeys.insert(0, owner_hk);
             }
         }
 
-        immune_hotkeys
+        owner_hotkeys
     }
 
     pub fn distribute_dividends_and_incentives(
@@ -498,7 +489,7 @@ impl<T: Config> Pallet<T> {
 
         // Distribute mining incentives.
         let subnet_owner_coldkey = SubnetOwner::<T>::get(netuid);
-        let owner_hotkeys = Self::get_immune_owner_hotkeys(netuid, &subnet_owner_coldkey);
+        let owner_hotkeys = Self::get_owner_hotkeys(netuid, &subnet_owner_coldkey);
         log::debug!("incentives: owner hotkeys: {owner_hotkeys:?}");
         for (hotkey, incentive) in incentives {
             log::debug!("incentives: hotkey: {incentive:?}");

pallets/subtensor/src/subnets/registration.rs

@@ -388,6 +388,47 @@ impl<T: Config> Pallet<T> {
         real_hash
     }
 
+    fn get_immune_owner_hotkeys(netuid: NetUid, coldkey: &T::AccountId) -> Vec<T::AccountId> {
+        // Gather (block, uid, hotkey) only for hotkeys that have a UID and a registration block.
+        let mut triples: Vec<(u64, u16, T::AccountId)> = OwnedHotkeys::<T>::get(coldkey)
+            .into_iter()
+            .filter_map(|hotkey| {
+                // Uids must exist, filter_map ignores hotkeys without UID
+                Uids::<T>::get(netuid, &hotkey).map(|uid| {
+                    let block = BlockAtRegistration::<T>::get(netuid, uid);
+                    (block, uid, hotkey)
+                })
+            })
+            .collect();
+
+        // Sort by BlockAtRegistration (descending), then by uid (ascending)
+        // Recent registration is priority so that we can let older keys expire (get non-immune)
+        triples.sort_by(|(b1, u1, _), (b2, u2, _)| b2.cmp(b1).then(u1.cmp(u2)));
+
+        // Keep first ImmuneOwnerUidsLimit
+        let limit = ImmuneOwnerUidsLimit::<T>::get(netuid).into();
+        if triples.len() > limit {
+            triples.truncate(limit);
+        }
+
+        // Project to just hotkeys
+        let mut immune_hotkeys: Vec<T::AccountId> =
+            triples.into_iter().map(|(_, _, hk)| hk).collect();
+
+        // Insert subnet owner hotkey in the beginning of the list if valid and not
+        // already present
+        if let Ok(owner_hk) = SubnetOwnerHotkey::<T>::try_get(netuid) {
+            if Uids::<T>::get(netuid, &owner_hk).is_some() && !immune_hotkeys.contains(&owner_hk) {
+                immune_hotkeys.insert(0, owner_hk);
+                if immune_hotkeys.len() > limit {
+                    immune_hotkeys.truncate(limit);
+                }
+            }
+        }
+
+        immune_hotkeys
+    }
+
     /// Determine which peer to prune from the network by finding the element with the lowest pruning score out of
     /// immunity period. If there is a tie for lowest pruning score, the neuron registered earliest is pruned.
     /// If all neurons are in immunity period, the neuron with the lowest pruning score is pruned. If there is a tie for
@@ -411,13 +452,14 @@ impl<T: Config> Pallet<T> {
             return 0; // If there are no neurons in this network.
         }
 
+        // Get the list of immortal (top-k by registration time of owner owned) keys
+        let subnet_owner_coldkey = SubnetOwner::<T>::get(netuid);
+        let immortal_hotkeys = Self::get_immune_owner_hotkeys(netuid, &subnet_owner_coldkey);
         for neuron_uid in 0..neurons_n {
-            // Do not deregister the owner's hotkey from the `SubnetOwnerHotkey` map
+            // Do not deregister the owner's owned hotkeys
             if let Ok(hotkey) = Self::get_hotkey_for_net_and_uid(netuid, neuron_uid) {
-                if let Ok(top_sn_owner_hotkey) = SubnetOwnerHotkey::<T>::try_get(netuid) {
-                    if top_sn_owner_hotkey == hotkey {
-                        continue;
-                    }
+                if immortal_hotkeys.contains(&hotkey) {
+                    continue;
                 }
             }
 

pallets/subtensor/src/tests/coinbase.rs

@@ -1857,62 +1857,6 @@ fn test_incentive_to_subnet_owners_hotkey_is_burned() {
     });
 }
 
-// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::coinbase::test_incentive_to_subnet_owners_hotkey_is_burned_with_limit --exact --show-output --nocapture
-#[test]
-fn test_incentive_to_subnet_owners_hotkey_is_burned_with_limit() {
-    new_test_ext(1).execute_with(|| {
-        let subnet_owner_ck = U256::from(0);
-        let subnet_owner_hk = U256::from(1);
-
-        // Other hk owned by owner
-        let other_hk = U256::from(3);
-        Owner::<Test>::insert(other_hk, subnet_owner_ck);
-        OwnedHotkeys::<Test>::insert(subnet_owner_ck, vec![subnet_owner_hk, other_hk]);
-
-        let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
-        Uids::<Test>::insert(netuid, other_hk, 1);
-
-        // Set the burn key limit to 1 - testing the limits
-        ImmuneOwnerUidsLimit::<Test>::insert(netuid, 1);
-
-        let pending_tao: u64 = 1_000_000_000;
-        let pending_alpha = AlphaCurrency::ZERO; // None to valis
-        let owner_cut = AlphaCurrency::ZERO;
-        let mut incentives: BTreeMap<U256, AlphaCurrency> = BTreeMap::new();
-
-        // Give incentive to other_hk
-        incentives.insert(other_hk, 10_000_000.into());
-
-        // Give incentives to subnet_owner_hk
-        incentives.insert(subnet_owner_hk, 10_000_000.into());
-
-        // Verify stake before
-        let subnet_owner_stake_before =
-            SubtensorModule::get_stake_for_hotkey_on_subnet(&subnet_owner_hk, netuid);
-        assert_eq!(subnet_owner_stake_before, 0.into());
-        let other_stake_before = SubtensorModule::get_stake_for_hotkey_on_subnet(&other_hk, netuid);
-        assert_eq!(other_stake_before, 0.into());
-
-        // Distribute dividends and incentives
-        SubtensorModule::distribute_dividends_and_incentives(
-            netuid,
-            owner_cut,
-            incentives,
-            BTreeMap::new(),
-            BTreeMap::new(),
-        );
-
-        // Verify stake after
-        let subnet_owner_stake_after =
-            SubtensorModule::get_stake_for_hotkey_on_subnet(&subnet_owner_hk, netuid);
-        assert_eq!(subnet_owner_stake_after, 0.into());
-        let other_stake_after = SubtensorModule::get_stake_for_hotkey_on_subnet(&other_hk, netuid);
-
-        // Testing the limit - should be not burned
-        assert!(other_stake_after > 0.into());
-    });
-}
-
 // Test that if number of sn owner hotkeys is greater than ImmuneOwnerUidsLimit, then the ones with
 // higher BlockAtRegistration are used to burn
 // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::coinbase::test_burn_key_sorting --exact --show-output --nocapture
@@ -1949,9 +1893,6 @@ fn test_burn_key_sorting() {
         Uids::<Test>::insert(netuid, other_hk_2, 3);
         Uids::<Test>::insert(netuid, other_hk_3, 2);
 
-        // Set the burn key limit to 3 because we also have sn owner
-        ImmuneOwnerUidsLimit::<Test>::insert(netuid, 3);
-
         let pending_tao: u64 = 1_000_000_000;
         let pending_alpha = AlphaCurrency::ZERO; // None to valis
         let owner_cut = AlphaCurrency::ZERO;
@@ -1979,15 +1920,15 @@ fn test_burn_key_sorting() {
             SubtensorModule::get_stake_for_hotkey_on_subnet(&subnet_owner_hk, netuid);
         assert_eq!(subnet_owner_stake_after, 0.into());
 
-        // Testing the limits - HK1 and HK3 should be burned, HK2 should be not burned
+        // No burn limits, all HKs should be burned
         let other_stake_after_1 =
             SubtensorModule::get_stake_for_hotkey_on_subnet(&other_hk_1, netuid);
         let other_stake_after_2 =
             SubtensorModule::get_stake_for_hotkey_on_subnet(&other_hk_2, netuid);
         let other_stake_after_3 =
             SubtensorModule::get_stake_for_hotkey_on_subnet(&other_hk_3, netuid);
         assert_eq!(other_stake_after_1, 0.into());
-        assert!(other_stake_after_2 > 0.into());
+        assert_eq!(other_stake_after_2, 0.into());
         assert_eq!(other_stake_after_3, 0.into());
     });
 }

pallets/subtensor/src/tests/registration.rs

@@ -1,5 +1,6 @@
 #![allow(clippy::unwrap_used)]
 
+use crate::*;
 use approx::assert_abs_diff_eq;
 use frame_support::dispatch::DispatchInfo;
 use frame_support::sp_runtime::{DispatchError, transaction_validity::TransactionSource};
@@ -1335,6 +1336,92 @@ fn test_registration_get_uid_to_prune_none_in_immunity_period() {
     });
 }
 
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::registration::test_registration_get_uid_to_prune_owner_immortality --exact --show-output --nocapture
+#[test]
+fn test_registration_get_uid_to_prune_owner_immortality() {
+    new_test_ext(1).execute_with(|| {
+        [
+            // Burn key limit to 1 - testing the limits
+            // Other owner's hotkey is pruned because there's only 1 immune key and
+            // pruning score of owner key is lower
+            (1, 1),
+            // Burn key limit to 2 - both owner keys are immune
+            (2, 2),
+        ]
+        .iter()
+        .for_each(|(limit, uid_to_prune)| {
+            let subnet_owner_ck = U256::from(0);
+            let subnet_owner_hk = U256::from(1);
+
+            // Other hk owned by owner
+            let other_owner_hk = U256::from(2);
+            Owner::<Test>::insert(other_owner_hk, subnet_owner_ck);
+            OwnedHotkeys::<Test>::insert(subnet_owner_ck, vec![subnet_owner_hk, other_owner_hk]);
+
+            // Another hk not owned by owner
+            let non_owner_hk = U256::from(3);
+
+            let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
+            BlockAtRegistration::<Test>::insert(netuid, 1, 1);
+            BlockAtRegistration::<Test>::insert(netuid, 2, 2);
+            Uids::<Test>::insert(netuid, other_owner_hk, 1);
+            Uids::<Test>::insert(netuid, non_owner_hk, 2);
+            Keys::<Test>::insert(netuid, 1, other_owner_hk);
+            Keys::<Test>::insert(netuid, 2, non_owner_hk);
+            ImmunityPeriod::<Test>::insert(netuid, 1);
+            SubnetworkN::<Test>::insert(netuid, 3);
+
+            step_block(10);
+
+            ImmuneOwnerUidsLimit::<Test>::insert(netuid, *limit);
+
+            // Set lower pruning score to sn owner keys
+            PruningScores::<Test>::insert(netuid, vec![0, 0, 1]);
+
+            assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), *uid_to_prune);
+        });
+    });
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::registration::test_registration_get_uid_to_prune_owner_immortality_all_immune --exact --show-output --nocapture
+#[test]
+fn test_registration_get_uid_to_prune_owner_immortality_all_immune() {
+    new_test_ext(1).execute_with(|| {
+        let limit = 2;
+        let uid_to_prune = 2;
+        let subnet_owner_ck = U256::from(0);
+        let subnet_owner_hk = U256::from(1);
+
+        // Other hk owned by owner
+        let other_owner_hk = U256::from(2);
+        Owner::<Test>::insert(other_owner_hk, subnet_owner_ck);
+        OwnedHotkeys::<Test>::insert(subnet_owner_ck, vec![subnet_owner_hk, other_owner_hk]);
+
+        // Another hk not owned by owner
+        let non_owner_hk = U256::from(3);
+
+        let netuid = add_dynamic_network(&subnet_owner_hk, &subnet_owner_ck);
+        BlockAtRegistration::<Test>::insert(netuid, 0, 12);
+        BlockAtRegistration::<Test>::insert(netuid, 1, 11);
+        BlockAtRegistration::<Test>::insert(netuid, 2, 10);
+        Uids::<Test>::insert(netuid, other_owner_hk, 1);
+        Uids::<Test>::insert(netuid, non_owner_hk, 2);
+        Keys::<Test>::insert(netuid, 1, other_owner_hk);
+        Keys::<Test>::insert(netuid, 2, non_owner_hk);
+        ImmunityPeriod::<Test>::insert(netuid, 100);
+        SubnetworkN::<Test>::insert(netuid, 3);
+
+        step_block(20);
+
+        ImmuneOwnerUidsLimit::<Test>::insert(netuid, limit);
+
+        // Set lower pruning score to sn owner keys
+        PruningScores::<Test>::insert(netuid, vec![0, 0, 1]);
+
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), uid_to_prune);
+    });
+}
+
 #[test]
 fn test_registration_pruning() {
     new_test_ext(1).execute_with(|| {

runtime/src/lib.rs

@@ -220,7 +220,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: 315,
+    spec_version: 316,
     impl_version: 1,
     apis: RUNTIME_API_VERSIONS,
     transaction_version: 1,