Trim uids

pallets/admin-utils/src/benchmarking.rs

@@ -346,5 +346,16 @@ mod benchmarks {
 		_(RawOrigin::Root, 5u16/*version*/)/*sudo_set_commit_reveal_version()*/;
     }
 
+    #[benchmark]
+    fn sudo_trim_to_max_allowed_uids() {
+        pallet_subtensor::Pallet::<T>::init_new_network(
+            1u16.into(), /*netuid*/
+            1u16,        /*sudo_tempo*/
+        );
+
+        #[extrinsic_call]
+		_(RawOrigin::Root, 1u16.into()/*netuid*/, 4097u16/*max_allowed_uids*/)/*sudo_trim_to_max_allowed_uids()*/;
+    }
+
     //impl_benchmark_test_suite!(AdminUtils, crate::mock::new_test_ext(), crate::mock::Test);
 }

pallets/admin-utils/src/lib.rs

@@ -668,6 +668,7 @@ pub mod pallet {
             min_burn: TaoCurrency,
         ) -> DispatchResult {
             pallet_subtensor::Pallet::<T>::ensure_subnet_owner_or_root(origin.clone(), netuid)?;
+
             ensure!(
                 pallet_subtensor::Pallet::<T>::if_subnet_exist(netuid),
                 Error::<T>::SubnetDoesNotExist
@@ -699,6 +700,7 @@ pub mod pallet {
             max_burn: TaoCurrency,
         ) -> DispatchResult {
             pallet_subtensor::Pallet::<T>::ensure_subnet_owner_or_root(origin.clone(), netuid)?;
+
             ensure!(
                 pallet_subtensor::Pallet::<T>::if_subnet_exist(netuid),
                 Error::<T>::SubnetDoesNotExist
@@ -1623,6 +1625,31 @@ pub mod pallet {
             log::debug!("CKBurnSet( burn: {burn:?} ) ");
             Ok(())
         }
+
+        /// Sets the maximum allowed UIDs for a subnet
+        #[pallet::call_index(74)]
+        #[pallet::weight(Weight::from_parts(15_000_000, 0)
+        .saturating_add(<T as frame_system::Config>::DbWeight::get().reads(1_u64))
+        .saturating_add(<T as frame_system::Config>::DbWeight::get().writes(1_u64)))]
+        pub fn sudo_trim_to_max_allowed_uids(
+            origin: OriginFor<T>,
+            netuid: NetUid,
+            max_n: u16,
+        ) -> DispatchResult {
+            pallet_subtensor::Pallet::<T>::ensure_subnet_owner_or_root(origin.clone(), netuid)?;
+            if let Ok(RawOrigin::Signed(who)) = origin.into() {
+                ensure!(
+                    pallet_subtensor::Pallet::<T>::passes_rate_limit_on_subnet(
+                        &TransactionType::SetMaxAllowedUIDS,
+                        &who,
+                        netuid,
+                    ),
+                    pallet_subtensor::Error::<T>::TxRateLimitExceeded
+                );
+            }
+            pallet_subtensor::Pallet::<T>::trim_to_max_allowed_uids(netuid, max_n)?;
+            Ok(())
+        }
     }
 }
 

pallets/admin-utils/src/tests/mod.rs

@@ -5,7 +5,10 @@ use frame_support::{
     traits::Hooks,
 };
 use frame_system::Config;
-use pallet_subtensor::{Error as SubtensorError, SubnetOwner, Tempo, WeightsVersionKeyRateLimit};
+use pallet_subtensor::{
+    Error as SubtensorError, MaxRegistrationsPerBlock, Rank, SubnetOwner,
+    TargetRegistrationsPerInterval, Tempo, WeightsVersionKeyRateLimit, *,
+};
 // use pallet_subtensor::{migrations, Event};
 use pallet_subtensor::Event;
 use sp_consensus_grandpa::AuthorityId as GrandpaId;
@@ -2069,3 +2072,106 @@ fn test_sudo_set_max_burn() {
         );
     });
 }
+
+#[test]
+fn test_trim_to_max_allowed_uids() {
+    new_test_ext().execute_with(|| {
+        let netuid = NetUid::from(1);
+        add_network(netuid, 10);
+        MaxRegistrationsPerBlock::<Test>::insert(netuid, 256);
+        TargetRegistrationsPerInterval::<Test>::insert(netuid, 256);
+
+        // Add some neurons
+        let max_n = 32;
+        for i in 1..=max_n {
+            let n = i * 1000;
+            register_ok_neuron(netuid, U256::from(n), U256::from(n + i), 0);
+        }
+
+        // Run some block to ensure stake weights are set
+        run_to_block(20);
+
+        // Normal case
+        let new_max_n = 20;
+        assert_ok!(AdminUtils::sudo_trim_to_max_allowed_uids(
+            <<Test as Config>::RuntimeOrigin>::root(),
+            netuid,
+            new_max_n
+        ));
+
+        // Ensure storage has been trimmed
+        assert_eq!(MaxAllowedUids::<Test>::get(netuid), new_max_n);
+        assert_eq!(Rank::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Trust::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Active::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Emission::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Consensus::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Incentive::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(Dividends::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(LastUpdate::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(PruningScores::<Test>::get(netuid).len(), new_max_n as usize);
+        assert_eq!(
+            ValidatorTrust::<Test>::get(netuid).len(),
+            new_max_n as usize
+        );
+        assert_eq!(
+            ValidatorPermit::<Test>::get(netuid).len(),
+            new_max_n as usize
+        );
+        assert_eq!(StakeWeight::<Test>::get(netuid).len(), new_max_n as usize);
+
+        for uid in max_n..new_max_n {
+            assert!(!Keys::<Test>::contains_key(netuid, uid));
+            assert!(!BlockAtRegistration::<Test>::contains_key(netuid, uid));
+            assert!(!Weights::<Test>::contains_key(netuid, uid));
+            assert!(!Bonds::<Test>::contains_key(netuid, uid));
+        }
+
+        for uid in 0..max_n {
+            assert!(
+                Weights::<Test>::get(netuid, uid)
+                    .iter()
+                    .all(|(target_uid, _)| *target_uid < new_max_n),
+                "Found a weight with target_uid >= new_max_n"
+            );
+            assert!(
+                Bonds::<Test>::get(netuid, uid)
+                    .iter()
+                    .all(|(target_uid, _)| *target_uid < new_max_n),
+                "Found a bond with target_uid >= new_max_n"
+            );
+        }
+
+        assert_eq!(SubnetworkN::<Test>::get(netuid), new_max_n);
+
+        // Non existent subnet
+        assert_err!(
+            AdminUtils::sudo_trim_to_max_allowed_uids(
+                <<Test as Config>::RuntimeOrigin>::root(),
+                NetUid::from(42),
+                new_max_n
+            ),
+            pallet_subtensor::Error::<Test>::SubNetworkDoesNotExist
+        );
+
+        // New max n less than lower bound
+        assert_err!(
+            AdminUtils::sudo_trim_to_max_allowed_uids(
+                <<Test as Config>::RuntimeOrigin>::root(),
+                netuid,
+                15
+            ),
+            pallet_subtensor::Error::<Test>::InvalidValue
+        );
+
+        // New max n greater than upper bound
+        assert_err!(
+            AdminUtils::sudo_trim_to_max_allowed_uids(
+                <<Test as Config>::RuntimeOrigin>::root(),
+                netuid,
+                SubtensorModule::get_max_allowed_uids(netuid) + 1
+            ),
+            pallet_subtensor::Error::<Test>::InvalidValue
+        );
+    });
+}

pallets/subtensor/src/lib.rs

@@ -1498,7 +1498,7 @@ pub mod pallet {
     /// ==== Subnetwork Consensus Storage  ====
     /// =======================================
     #[pallet::storage] // --- DMAP ( netuid ) --> stake_weight | weight for stake used in YC.
-    pub(super) type StakeWeight<T: Config> =
+    pub type StakeWeight<T: Config> =
         StorageMap<_, Identity, NetUid, Vec<u16>, ValueQuery, EmptyU16Vec<T>>;
     #[pallet::storage]
     /// --- DMAP ( netuid, hotkey ) --> uid

pallets/subtensor/src/subnets/uids.rs

@@ -16,17 +16,6 @@ impl<T: Config> Pallet<T> {
         }
     }
 
-    /// Resets the trust, emission, consensus, incentive, dividends of the neuron to default
-    pub fn clear_neuron(netuid: NetUid, neuron_uid: u16) {
-        let neuron_index: usize = neuron_uid.into();
-        Emission::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0.into()));
-        Trust::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
-        Consensus::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
-        Incentive::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
-        Dividends::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
-        Bonds::<T>::remove(netuid, neuron_uid); // Remove bonds for Validator.
-    }
-
     /// Replace the neuron under this uid.
     pub fn replace_neuron(
         netuid: NetUid,
@@ -107,6 +96,136 @@ impl<T: Config> Pallet<T> {
         IsNetworkMember::<T>::insert(new_hotkey.clone(), netuid, true); // Fill network is member.
     }
 
+    /// Clears (sets to default) the neuron map values fot a neuron when it is
+    /// removed from the subnet
+    pub fn clear_neuron(netuid: NetUid, neuron_uid: u16) {
+        let neuron_index: usize = neuron_uid.into();
+        Emission::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0.into()));
+        Trust::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
+        Consensus::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
+        Incentive::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
+        Dividends::<T>::mutate(netuid, |v| Self::set_element_at(v, neuron_index, 0));
+        Bonds::<T>::remove(netuid, neuron_uid); // Remove bonds for Validator.
+    }
+
+    pub fn trim_to_max_allowed_uids(netuid: NetUid, max_n: u16) -> DispatchResult {
+        // Reasonable limits
+        ensure!(
+            Self::if_subnet_exist(netuid),
+            Error::<T>::SubNetworkDoesNotExist
+        );
+        ensure!(max_n > 16, Error::<T>::InvalidValue);
+        ensure!(
+            max_n <= Self::get_max_allowed_uids(netuid),
+            Error::<T>::InvalidValue
+        );
+
+        // Set the value.
+        MaxAllowedUids::<T>::insert(netuid, max_n);
+
+        // Check if we need to trim.
+        let current_n: u16 = Self::get_subnetwork_n(netuid);
+
+        // We need to trim, get rid of values between max_n and current_n.
+        if current_n > max_n {
+            let ranks: Vec<u16> = Rank::<T>::get(netuid);
+            let trimmed_ranks: Vec<u16> = ranks.into_iter().take(max_n as usize).collect();
+            Rank::<T>::insert(netuid, trimmed_ranks);
+
+            let trust: Vec<u16> = Trust::<T>::get(netuid);
+            let trimmed_trust: Vec<u16> = trust.into_iter().take(max_n as usize).collect();
+            Trust::<T>::insert(netuid, trimmed_trust);
+
+            let active: Vec<bool> = Active::<T>::get(netuid);
+            let trimmed_active: Vec<bool> = active.into_iter().take(max_n as usize).collect();
+            Active::<T>::insert(netuid, trimmed_active);
+
+            let emission: Vec<AlphaCurrency> = Emission::<T>::get(netuid);
+            let trimmed_emission: Vec<AlphaCurrency> =
+                emission.into_iter().take(max_n as usize).collect();
+            Emission::<T>::insert(netuid, trimmed_emission);
+
+            let consensus: Vec<u16> = Consensus::<T>::get(netuid);
+            let trimmed_consensus: Vec<u16> = consensus.into_iter().take(max_n as usize).collect();
+            Consensus::<T>::insert(netuid, trimmed_consensus);
+
+            let incentive: Vec<u16> = Incentive::<T>::get(netuid);
+            let trimmed_incentive: Vec<u16> = incentive.into_iter().take(max_n as usize).collect();
+            Incentive::<T>::insert(netuid, trimmed_incentive);
+
+            let dividends: Vec<u16> = Dividends::<T>::get(netuid);
+            let trimmed_dividends: Vec<u16> = dividends.into_iter().take(max_n as usize).collect();
+            Dividends::<T>::insert(netuid, trimmed_dividends);
+
+            let lastupdate: Vec<u64> = LastUpdate::<T>::get(netuid);
+            let trimmed_lastupdate: Vec<u64> =
+                lastupdate.into_iter().take(max_n as usize).collect();
+            LastUpdate::<T>::insert(netuid, trimmed_lastupdate);
+
+            let pruning_scores: Vec<u16> = PruningScores::<T>::get(netuid);
+            let trimmed_pruning_scores: Vec<u16> =
+                pruning_scores.into_iter().take(max_n as usize).collect();
+            PruningScores::<T>::insert(netuid, trimmed_pruning_scores);
+
+            let vtrust: Vec<u16> = ValidatorTrust::<T>::get(netuid);
+            let trimmed_vtrust: Vec<u16> = vtrust.into_iter().take(max_n as usize).collect();
+            ValidatorTrust::<T>::insert(netuid, trimmed_vtrust);
+
+            let vpermit: Vec<bool> = ValidatorPermit::<T>::get(netuid);
+            let trimmed_vpermit: Vec<bool> = vpermit.into_iter().take(max_n as usize).collect();
+            ValidatorPermit::<T>::insert(netuid, trimmed_vpermit);
+
+            let stake_weight: Vec<u16> = StakeWeight::<T>::get(netuid);
+            let trimmed_stake_weight: Vec<u16> =
+                stake_weight.into_iter().take(max_n as usize).collect();
+            StakeWeight::<T>::insert(netuid, trimmed_stake_weight);
+
+            for uid in max_n..current_n {
+                // Trim UIDs and Keys by removing entries with UID >= max_n (since UIDs are 0-indexed)
+                // UIDs range from 0 to current_n-1, so we remove UIDs from max_n to current_n-1
+                if let Ok(hotkey) = Keys::<T>::try_get(netuid, uid) {
+                    Uids::<T>::remove(netuid, &hotkey);
+                    // Remove IsNetworkMember association for the hotkey
+                    IsNetworkMember::<T>::remove(&hotkey, netuid);
+                    // Remove last hotkey emission for the hotkey
+                    LastHotkeyEmissionOnNetuid::<T>::remove(&hotkey, netuid);
+                    // Remove alpha dividends for the hotkey
+                    AlphaDividendsPerSubnet::<T>::remove(netuid, &hotkey);
+                    // Remove tao dividends for the hotkey
+                    TaoDividendsPerSubnet::<T>::remove(netuid, &hotkey);
+                    // Trim axons, certificates, and prometheus info for removed hotkeys
+                    Axons::<T>::remove(netuid, &hotkey);
+                    NeuronCertificates::<T>::remove(netuid, &hotkey);
+                    Prometheus::<T>::remove(netuid, &hotkey);
+                }
+                #[allow(unknown_lints)]
+                Keys::<T>::remove(netuid, uid);
+                // Remove block at registration for the uid
+                BlockAtRegistration::<T>::remove(netuid, uid);
+                // Remove entire weights and bonds entries for removed UIDs
+                Weights::<T>::remove(netuid, uid);
+                Bonds::<T>::remove(netuid, uid);
+            }
+
+            // Trim weight and bond connections to removed UIDs for remaining neurons
+            // UIDs 0 to max_n-1 are kept, so we iterate through these valid UIDs
+            for uid in 0..max_n {
+                Weights::<T>::mutate(netuid, uid, |weights| {
+                    weights.retain(|(target_uid, _)| *target_uid < max_n);
+                });
+                Bonds::<T>::mutate(netuid, uid, |bonds| {
+                    bonds.retain(|(target_uid, _)| *target_uid < max_n);
+                });
+            }
+
+            // Update the subnetwork size
+            SubnetworkN::<T>::insert(netuid, max_n);
+        }
+
+        // --- Ok and done.
+        Ok(())
+    }
+
     /// Returns true if the uid is set on the network.
     ///
     pub fn is_uid_exist_on_network(netuid: NetUid, uid: u16) -> bool {

pallets/subtensor/src/utils/rate_limiting.rs

@@ -11,6 +11,7 @@ pub enum TransactionType {
     RegisterNetwork,
     SetWeightsVersionKey,
     SetSNOwnerHotkey,
+    SetMaxAllowedUIDS,
 }
 
 /// Implement conversion from TransactionType to u16
@@ -23,6 +24,7 @@ impl From<TransactionType> for u16 {
             TransactionType::RegisterNetwork => 3,
             TransactionType::SetWeightsVersionKey => 4,
             TransactionType::SetSNOwnerHotkey => 5,
+            TransactionType::SetMaxAllowedUIDS => 6,
         }
     }
 }
@@ -36,6 +38,7 @@ impl From<u16> for TransactionType {
             3 => TransactionType::RegisterNetwork,
             4 => TransactionType::SetWeightsVersionKey,
             5 => TransactionType::SetSNOwnerHotkey,
+            6 => TransactionType::SetMaxAllowedUIDS,
             _ => TransactionType::Unknown,
         }
     }
@@ -50,7 +53,7 @@ impl<T: Config> Pallet<T> {
             TransactionType::SetChildren => 150, // 30 minutes
             TransactionType::SetChildkeyTake => TxChildkeyTakeRateLimit::<T>::get(),
             TransactionType::RegisterNetwork => NetworkRateLimit::<T>::get(),
-
+            TransactionType::SetMaxAllowedUIDS => 7200 * 30,
             TransactionType::Unknown => 0, // Default to no limit for unknown types (no limit)
             _ => 0,
         }
@@ -62,7 +65,6 @@ impl<T: Config> Pallet<T> {
             TransactionType::SetWeightsVersionKey => (Tempo::<T>::get(netuid) as u64)
                 .saturating_mul(WeightsVersionKeyRateLimit::<T>::get()),
             TransactionType::SetSNOwnerHotkey => DefaultSetSNOwnerHotkeyRateLimit::<T>::get(),
-
             _ => Self::get_rate_limit(tx_type),
         }
     }
@@ -89,7 +91,6 @@ impl<T: Config> Pallet<T> {
         let block: u64 = Self::get_current_block_as_u64();
         let limit: u64 = Self::get_rate_limit_on_subnet(tx_type, netuid);
         let last_block: u64 = Self::get_last_transaction_block_on_subnet(hotkey, netuid, tx_type);
-
         Self::check_passes_rate_limit(limit, block, last_block)
     }