Clear unknown neuron axon/prom/cert

pallets/subtensor/src/macros/hooks.rs

@@ -162,7 +162,9 @@ mod hooks {
                 // Migrate pending emissions
                 .saturating_add(migrations::migrate_pending_emissions::migrate_pending_emissions::<T>())
                 // Reset unactive subnets
-                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>());
+                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>())
+                // Remove unknown neuron axon, certificate prom
+                .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_remove_unknown_neuron_axon_cert_prom.rs

@@ -0,0 +1,83 @@
+use super::*;
+use crate::HasMigrationRun;
+use frame_support::{traits::Get, weights::Weight};
+use scale_info::prelude::string::String;
+use sp_std::collections::btree_set::BTreeSet;
+
+pub fn migrate_remove_unknown_neuron_axon_cert_prom<T: Config>() -> Weight {
+    let migration_name = b"migrate_remove_neuron_axon_cert_prom".to_vec();
+    let mut weight: Weight = T::DbWeight::get().reads(1);
+
+    // Skip if already executed
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            target: "runtime",
+            "Migration '{}' already run - skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    for network in NetworksAdded::<T>::iter_keys() {
+        weight.saturating_accrue(T::DbWeight::get().reads(1));
+
+        let hotkeys = BTreeSet::from_iter(Uids::<T>::iter_key_prefix(network));
+        weight.saturating_accrue(T::DbWeight::get().reads(hotkeys.len() as u64));
+
+        // Axons
+        let axons = Axons::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        weight.saturating_accrue(T::DbWeight::get().reads(axons.len() as u64));
+        let mut cleaned_axons: u32 = 0;
+        for axon_hotkey in axons {
+            if !hotkeys.contains(&axon_hotkey) {
+                Axons::<T>::remove(network, &axon_hotkey);
+                cleaned_axons = cleaned_axons.saturating_add(1);
+            }
+        }
+        weight.saturating_accrue(T::DbWeight::get().writes(cleaned_axons as u64));
+
+        // Neuron Certificates
+        let certificates = NeuronCertificates::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        weight.saturating_accrue(T::DbWeight::get().reads(certificates.len() as u64));
+        let mut cleaned_certificates: u32 = 0;
+        for certificate_hotkey in certificates {
+            if !hotkeys.contains(&certificate_hotkey) {
+                NeuronCertificates::<T>::remove(network, &certificate_hotkey);
+                cleaned_certificates = cleaned_certificates.saturating_add(1);
+            }
+        }
+        weight.saturating_accrue(T::DbWeight::get().writes(cleaned_certificates as u64));
+
+        // Prometheus
+        let prometheus = Prometheus::<T>::iter_key_prefix(network).collect::<Vec<_>>();
+        weight.saturating_accrue(T::DbWeight::get().reads(prometheus.len() as u64));
+        let mut cleaned_prometheus: u32 = 0;
+        for prometheus_hotkey in prometheus {
+            if !hotkeys.contains(&prometheus_hotkey) {
+                Prometheus::<T>::remove(network, &prometheus_hotkey);
+                cleaned_prometheus = cleaned_prometheus.saturating_add(1);
+            }
+        }
+        weight.saturating_accrue(T::DbWeight::get().writes(cleaned_prometheus as u64));
+
+        log::info!(
+            "Cleaned {cleaned_axons} axons, {cleaned_certificates} neuron certificates, {cleaned_prometheus} prometheus for network {network}"
+        );
+    }
+
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed successfully.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    log::info!("{weight:#?} weight");
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -37,6 +37,7 @@ pub mod migrate_remove_network_modality;
 pub mod migrate_remove_stake_map;
 pub mod migrate_remove_tao_dividends;
 pub mod migrate_remove_total_hotkey_coldkey_stakes_this_interval;
+pub mod migrate_remove_unknown_neuron_axon_cert_prom;
 pub mod migrate_remove_unused_maps_and_values;
 pub mod migrate_remove_zero_total_hotkey_alpha;
 pub mod migrate_reset_bonds_moving_average;

pallets/subtensor/src/tests/migration.rs

@@ -2724,3 +2724,69 @@ fn test_migrate_reset_unactive_sn_idempotence() {
         assert_eq!(TotalIssuance::<Test>::get(), total_issuance_before);
     });
 }
+
+#[test]
+fn test_migrate_remove_unknown_neuron_axon_cert_prom() {
+    use crate::migrations::migrate_remove_unknown_neuron_axon_cert_prom::*;
+    const MIGRATION_NAME: &[u8] = b"migrate_remove_neuron_axon_cert_prom";
+
+    new_test_ext(1).execute_with(|| {
+        setup_for(NetUid::from(2), 64, 1231);
+        setup_for(NetUid::from(42), 256, 15151);
+        setup_for(NetUid::from(99), 1024, 32323);
+        assert!(!HasMigrationRun::<Test>::get(MIGRATION_NAME));
+
+        let w = migrate_remove_unknown_neuron_axon_cert_prom::<Test>();
+        assert!(!w.is_zero(), "Weight must be non-zero");
+
+        assert!(HasMigrationRun::<Test>::get(MIGRATION_NAME));
+        assert_for(NetUid::from(2), 64, 1231);
+        assert_for(NetUid::from(42), 256, 15151);
+        assert_for(NetUid::from(99), 1024, 32323);
+    });
+
+    fn setup_for(netuid: NetUid, uids: u32, items: u32) {
+        NetworksAdded::<Test>::insert(netuid, true);
+
+        for i in 1u32..=uids {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            Uids::<Test>::insert(netuid, hk, i as u16);
+        }
+
+        for i in 1u32..=items {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            Axons::<Test>::insert(netuid, hk, AxonInfo::default());
+            NeuronCertificates::<Test>::insert(netuid, hk, NeuronCertificate::default());
+            Prometheus::<Test>::insert(netuid, hk, PrometheusInfo::default());
+        }
+    }
+
+    fn assert_for(netuid: NetUid, uids: u32, items: u32) {
+        assert_eq!(
+            Axons::<Test>::iter_key_prefix(netuid).count(),
+            uids as usize
+        );
+        assert_eq!(
+            NeuronCertificates::<Test>::iter_key_prefix(netuid).count(),
+            uids as usize
+        );
+        assert_eq!(
+            Prometheus::<Test>::iter_key_prefix(netuid).count(),
+            uids as usize
+        );
+
+        for i in 1u32..=uids {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            assert!(Axons::<Test>::contains_key(netuid, hk));
+            assert!(NeuronCertificates::<Test>::contains_key(netuid, hk));
+            assert!(Prometheus::<Test>::contains_key(netuid, hk));
+        }
+
+        for i in uids + 1u32..=items {
+            let hk = U256::from(netuid.inner() as u32 * 1000 + i);
+            assert!(!Axons::<Test>::contains_key(netuid, hk));
+            assert!(!NeuronCertificates::<Test>::contains_key(netuid, hk));
+            assert!(!Prometheus::<Test>::contains_key(netuid, hk));
+        }
+    }
+}