neuron pruning changes, initial tests

Author: VectorChat

pallets/subtensor/src/registration.rs

@@ -423,65 +423,57 @@ impl<T: Config> Pallet<T> {
     }
 
     /// Determine which peer to prune from the network by finding the element with the lowest pruning score out of
-    /// immunity period. If all neurons are in immunity period, return node with lowest prunning score.
-    /// This function will always return an element to prune.
+    /// 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
+    /// the lowest pruning score, the immune neuron registered earliest is pruned.
+    /// Ties for earliest registration are broken by the neuron with the lowest uid.
     pub fn get_neuron_to_prune(netuid: u16) -> u16 {
         let mut min_score: u16 = u16::MAX;
-        let mut min_score_in_immunity_period = u16::MAX;
-        let mut uid_with_min_score = 0;
-        let mut uid_with_min_score_in_immunity_period: u16 = 0;
+        let mut min_score_in_immunity: u16 = u16::MAX;
+        let mut earliest_registration: u64 = u64::MAX;
+        let mut earliest_registration_in_immunity: u64 = u64::MAX;
+        let mut uid_to_prune: u16 = 0;
+        let mut uid_to_prune_in_immunity: u16 = 0;
+        let mut found_non_immune = false;
 
         let neurons_n = Self::get_subnetwork_n(netuid);
         if neurons_n == 0 {
             return 0; // If there are no neurons in this network.
         }
 
-        let current_block: u64 = Self::get_current_block_as_u64();
-        let immunity_period: u64 = Self::get_immunity_period(netuid) as u64;
-        for neuron_uid_i in 0..neurons_n {
-            let pruning_score: u16 = Self::get_pruning_score_for_uid(netuid, neuron_uid_i);
+        for neuron_uid in 0..neurons_n {
+            let pruning_score: u16 = Self::get_pruning_score_for_uid(netuid, neuron_uid);
             let block_at_registration: u64 =
-                Self::get_neuron_block_at_registration(netuid, neuron_uid_i);
-            #[allow(clippy::comparison_chain)]
-            if min_score == pruning_score {
-                if current_block.saturating_sub(block_at_registration) < immunity_period {
-                    //neuron is in immunity period
-                    if min_score_in_immunity_period > pruning_score {
-                        min_score_in_immunity_period = pruning_score;
-                        uid_with_min_score_in_immunity_period = neuron_uid_i;
-                    }
-                } else {
-                    uid_with_min_score = neuron_uid_i;
+                Self::get_neuron_block_at_registration(netuid, neuron_uid);
+            let is_immune = Self::get_neuron_is_immune(netuid, neuron_uid);
+
+            if is_immune {
+                if pruning_score < min_score_in_immunity
+                    || (pruning_score == min_score_in_immunity
+                        && block_at_registration < earliest_registration_in_immunity)
+                {
+                    min_score_in_immunity = pruning_score;
+                    earliest_registration_in_immunity = block_at_registration;
+                    uid_to_prune_in_immunity = neuron_uid;
                 }
-            }
-            // Find min pruning score.
-            else if min_score > pruning_score {
-                if current_block.saturating_sub(block_at_registration) < immunity_period {
-                    //neuron is in immunity period
-                    if min_score_in_immunity_period > pruning_score {
-                        min_score_in_immunity_period = pruning_score;
-                        uid_with_min_score_in_immunity_period = neuron_uid_i;
-                    }
-                } else {
+            } else {
+                found_non_immune = true;
+                if pruning_score < min_score
+                    || (pruning_score == min_score && block_at_registration < earliest_registration)
+                {
                     min_score = pruning_score;
-                    uid_with_min_score = neuron_uid_i;
+                    earliest_registration = block_at_registration;
+                    uid_to_prune = neuron_uid;
                 }
             }
         }
-        if min_score == u16::MAX {
-            //all neuorns are in immunity period
-            Self::set_pruning_score_for_uid(
-                netuid,
-                uid_with_min_score_in_immunity_period,
-                u16::MAX,
-            );
-            uid_with_min_score_in_immunity_period
+
+        if found_non_immune {
+            Self::set_pruning_score_for_uid(netuid, uid_to_prune, u16::MAX);
+            uid_to_prune
         } else {
-            // We replace the pruning score here with u16 max to ensure that all peers always have a
-            // pruning score. In the event that every peer has been pruned this function will prune
-            // the last element in the network continually.
-            Self::set_pruning_score_for_uid(netuid, uid_with_min_score, u16::MAX);
-            uid_with_min_score
+            Self::set_pruning_score_for_uid(netuid, uid_to_prune_in_immunity, u16::MAX);
+            uid_to_prune_in_immunity
         }
     }
 

pallets/subtensor/src/utils.rs

@@ -461,6 +461,13 @@ impl<T: Config> Pallet<T> {
         Self::deposit_event(Event::ImmunityPeriodSet(netuid, immunity_period));
     }
 
+    pub fn get_neuron_is_immune(netuid: u16, uid: u16) -> bool {
+        let registered_at = Self::get_neuron_block_at_registration(netuid, uid);
+        let current_block = Self::get_current_block_as_u64();
+        let immunity_period = Self::get_immunity_period(netuid);
+        current_block.saturating_sub(registered_at) < u64::from(immunity_period)
+    }
+
     pub fn get_min_allowed_weights(netuid: u16) -> u16 {
         MinAllowedWeights::<T>::get(netuid)
     }

pallets/subtensor/tests/registration.rs

@@ -1,6 +1,9 @@
 #![allow(clippy::unwrap_used)]
 
+use std::u16;
+
 use frame_support::traits::Currency;
+use substrate_fixed::types::extra::True;
 
 use crate::mock::*;
 use frame_support::dispatch::{DispatchClass, DispatchInfo, GetDispatchInfo, Pays};
@@ -538,6 +541,118 @@ fn test_burn_adjustment() {
     });
 }
 
+#[test]
+fn test_burn_registration_pruning_scenarios() {
+    new_test_ext(1).execute_with(|| {
+        let netuid: u16 = 1;
+        let tempo: u16 = 13;
+        let burn_cost = 1000;
+        let coldkey_account_id = U256::from(667);
+        let max_allowed_uids = 6;
+        let immunity_period = 5000;
+
+        SubtensorModule::set_burn(netuid, burn_cost);
+        SubtensorModule::set_max_allowed_uids(netuid, max_allowed_uids);
+        SubtensorModule::set_target_registrations_per_interval(netuid, max_allowed_uids);
+        SubtensorModule::set_immunity_period(netuid, immunity_period);
+
+        // SubtensorModule::set_immunity_period(netuid, immunity_period);
+
+        add_network(netuid, tempo, 0);
+
+        let mint_balance = burn_cost * u64::from(max_allowed_uids) + 1_000_000_000;
+        SubtensorModule::add_balance_to_coldkey_account(&coldkey_account_id, mint_balance);
+
+        // Register first half of neurons
+        for i in 0..3 {
+            assert_ok!(SubtensorModule::burned_register(
+                <<Test as Config>::RuntimeOrigin>::signed(coldkey_account_id),
+                netuid,
+                U256::from(i)
+            ));
+            step_block(1);
+        }
+
+        // Note: pruning score is set to u16::MAX after getting neuron to prune
+
+        // 1. Test all immune neurons
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 0), true);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 1), true);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 2), true);
+
+        SubtensorModule::set_pruning_score_for_uid(netuid, 0, 100);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, 75);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, 50);
+
+        // The immune neuron with the lowest score should be pruned
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 2);
+
+        // 2. Test tie-breaking for immune neurons
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, 50);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, 50);
+
+        // Should get the oldest neuron
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 1);
+
+        // 3. Test no immune neurons
+        step_block(immunity_period);
+
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 0), false);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 1), false);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 2), false);
+
+        SubtensorModule::set_pruning_score_for_uid(netuid, 0, 100);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, 50);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, 75);
+
+        // The non-immune neuron with the lowest score should be pruned
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 1);
+
+        // 4. Test tie-breaking for non-immune neurons
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, 50);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, 50);
+
+        // Should get the oldest non-immune neuron
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 1);
+
+        // 5. Test mixed immunity
+        // Register second batch of neurons (these will be non-immune)
+        for i in 3..6 {
+            assert_ok!(SubtensorModule::burned_register(
+                <<Test as Config>::RuntimeOrigin>::signed(coldkey_account_id),
+                netuid,
+                U256::from(i)
+            ));
+            step_block(1);
+        }
+
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 3), true);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 4), true);
+        assert_eq!(SubtensorModule::get_neuron_is_immune(netuid, 5), true);
+
+        // Set pruning scores for all neurons
+        SubtensorModule::set_pruning_score_for_uid(netuid, 0, 75); // non-immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, 50); // non-immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, 60); // non-immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 3, 40); // immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 4, 55); // immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 5, 45); // immune
+
+        // The non-immune neuron with the lowest score should be pruned
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 1);
+
+        // If we remove the lowest non-immune neuron, it should choose the next lowest non-immune
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, u16::MAX);
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 2);
+
+        // If we make all non-immune neurons have high scores, it should choose the oldest non-immune neuron
+        SubtensorModule::set_pruning_score_for_uid(netuid, 0, u16::MAX);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 1, u16::MAX);
+        SubtensorModule::set_pruning_score_for_uid(netuid, 2, u16::MAX);
+        assert_eq!(SubtensorModule::get_neuron_to_prune(netuid), 0);
+    });
+}
+
 #[test]
 fn test_registration_too_many_registrations_per_block() {
     new_test_ext(1).execute_with(|| {