Allow setting CHK immediately as long as start_call has not been executed yet

pallets/subtensor/src/staking/set_children.rs

@@ -547,14 +547,30 @@ impl<T: Config> Pallet<T> {
         let current_block = Self::get_current_block_as_u64();
         TransactionType::SetChildren.set_last_block_on_subnet::<T>(&hotkey, netuid, current_block);
 
+        // Schedule or immediately apply CK
+        Self::schedule_or_apply_ck(netuid, hotkey, children)
+    }
+
+    /// If the start call occured, schedule children, otherwise,
+    /// apply immediately
+    fn schedule_or_apply_ck(
+        netuid: NetUid,
+        hotkey: T::AccountId,
+        children: Vec<(u64, T::AccountId)>,
+    ) -> DispatchResult {
+        if !SubtokenEnabled::<T>::get(netuid) {
+            Self::persist_pending_chidren_ok(netuid, &hotkey, &children);
+            return Ok(());
+        }
+
         // Calculate cool-down block
         let cooldown_block =
             Self::get_current_block_as_u64().saturating_add(PendingChildKeyCooldown::<T>::get());
 
         // Insert or update PendingChildKeys
         PendingChildKeys::<T>::insert(netuid, hotkey.clone(), (children.clone(), cooldown_block));
 
-        // --- 8. Log and return.
+        // Log and return.
         log::trace!(
             "SetChildrenScheduled( netuid:{:?}, cooldown_block:{:?}, hotkey:{:?}, children:{:?} )",
             cooldown_block,
@@ -563,10 +579,10 @@ impl<T: Config> Pallet<T> {
             children.clone()
         );
         Self::deposit_event(Event::SetChildrenScheduled(
-            hotkey.clone(),
+            hotkey,
             netuid,
             cooldown_block,
-            children.clone(),
+            children,
         ));
 
         // Ok and return.
@@ -600,65 +616,47 @@ impl<T: Config> Pallet<T> {
         let current_block = Self::get_current_block_as_u64();
 
         // If the childkey cools down before the subnet start call + PendingChildKeyCooldown:
-        //   - If Start call happened: Postpone to start_block + PendingChildKeyCooldown
-        //   - If Start call didn't happen: Postpone to now + PendingChildKeyCooldown
-        let cooldown_period = PendingChildKeyCooldown::<T>::get();
-        let cooldown_allowed_block =
-            if let Some(first_emission_block) = FirstEmissionBlockNumber::<T>::get(netuid) {
-                let start_call_block = first_emission_block.saturating_sub(1);
-                start_call_block.saturating_add(cooldown_period)
-            } else {
-                current_block.saturating_add(cooldown_period)
-            };
+        //   - If Start call happened: Normal track
+        //   - If Start call didn't happen: Apply immediately
+        // TODO: This check may be removed after all ck are applied after the runtime upgrade
+        let start_call_occured = SubtokenEnabled::<T>::get(netuid);
 
         // Iterate over all pending children of this subnet and set as needed
         let mut to_remove: Vec<T::AccountId> = Vec::new();
-        let mut to_reschedule: Vec<(T::AccountId, Vec<(u64, T::AccountId)>)> = Vec::new();
 
         PendingChildKeys::<T>::iter_prefix(netuid).for_each(
             |(hotkey, (children, cool_down_block))| {
-                if cool_down_block < current_block {
-                    if current_block >= cooldown_allowed_block {
-                        // If child-parent consistency is broken, we will fail setting new children silently
-                        let maybe_relations =
-                            Self::load_relations_from_pending(hotkey.clone(), &children, netuid);
-                        if let Ok(relations) = maybe_relations {
-                            let mut _weight: Weight = T::DbWeight::get().reads(0);
-                            if let Ok(()) = Self::persist_child_parent_relations(
-                                relations,
-                                netuid,
-                                &mut _weight,
-                            ) {
-                                // Log and emit event.
-                                log::trace!(
-                                    "SetChildren( netuid:{:?}, hotkey:{:?}, children:{:?} )",
-                                    hotkey,
-                                    netuid,
-                                    children.clone()
-                                );
-                                Self::deposit_event(Event::SetChildren(
-                                    hotkey.clone(),
-                                    netuid,
-                                    children.clone(),
-                                ));
-                            }
-                        }
-
-                        to_remove.push(hotkey);
-                    } else {
-                        to_remove.push(hotkey.clone());
-                        to_reschedule.push((hotkey, children));
-                    }
+                if (cool_down_block < current_block) || !start_call_occured {
+                    Self::persist_pending_chidren_ok(netuid, &hotkey, &children);
+                    to_remove.push(hotkey);
                 }
             },
         );
 
         for hotkey in to_remove {
             PendingChildKeys::<T>::remove(netuid, hotkey);
         }
+    }
 
-        for (hotkey, children) in to_reschedule {
-            PendingChildKeys::<T>::insert(netuid, hotkey, (children, cooldown_allowed_block));
+    // If child-parent consistency is broken, fail setting new children silently
+    fn persist_pending_chidren_ok(
+        netuid: NetUid,
+        hotkey: &T::AccountId,
+        children: &Vec<(u64, T::AccountId)>,
+    ) {
+        let maybe_relations = Self::load_relations_from_pending(hotkey.clone(), children, netuid);
+        if let Ok(relations) = maybe_relations {
+            let mut _weight: Weight = T::DbWeight::get().reads(0);
+            if let Ok(()) = Self::persist_child_parent_relations(relations, netuid, &mut _weight) {
+                // Log and emit event.
+                log::trace!(
+                    "SetChildren( netuid:{:?}, hotkey:{:?}, children:{:?} )",
+                    hotkey,
+                    netuid,
+                    children.clone()
+                );
+                Self::deposit_event(Event::SetChildren(hotkey.clone(), netuid, children.clone()));
+            }
         }
     }
 

pallets/subtensor/src/tests/children.rs

@@ -2234,6 +2234,7 @@ fn test_do_remove_stake_clears_pending_childkeys() {
         add_network(netuid, 13, 0);
         register_ok_neuron(netuid, hotkey, coldkey, 0);
         SubtensorModule::add_balance_to_coldkey_account(&coldkey, 10_000_000_000_000_u64.into());
+        SubtokenEnabled::<Test>::insert(netuid, true);
 
         let reserve = 1_000_000_000_000_000_u64;
         mock::setup_reserves(netuid, reserve.into(), reserve.into());
@@ -4223,6 +4224,47 @@ fn test_set_child_keys_empty_vector_clears_storage() {
     });
 }
 
+// SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::children::test_set_child_keys_no_start_call_sets_immediately --exact --show-output
+#[test]
+fn test_set_child_keys_no_start_call_sets_immediately() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(1);
+        let hotkey = U256::from(2);
+        let child1 = U256::from(3);
+        let child2 = U256::from(4);
+        let netuid = NetUid::from(1);
+        let proportion1: u64 = 1000;
+        let proportion2: u64 = 2000;
+
+        // Add network and register hotkey
+        add_network(netuid, 13, 0);
+        register_ok_neuron(netuid, hotkey, coldkey, 0);
+
+        // Clear SubtokenEnabled
+        SubtokenEnabled::<Test>::remove(netuid);
+
+        // Set multiple children
+        mock_schedule_children(
+            &coldkey,
+            &hotkey,
+            netuid,
+            &[(proportion1, child1), (proportion2, child2)],
+        );
+
+        // Normally happens on epoch
+        SubtensorModule::do_set_pending_children(netuid);
+
+        // Verify pending map is empty
+        assert!(!PendingChildKeys::<Test>::contains_key(netuid, hotkey));
+
+        // Verify that childkey is set
+        assert_eq!(
+            ChildKeys::<Test>::get(hotkey, netuid),
+            vec![(proportion1, child1), (proportion2, child2)]
+        );
+    });
+}
+
 // Test that the subnet owner can always set weights (owner bypass in check_weights_min_stake)
 // and that do_set_root_validators_for_subnet correctly creates parent-child relationships.
 // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::children::test_root_children_enable_subnet_owner_set_weights --exact --show-output --nocapture
@@ -4349,8 +4391,9 @@ fn test_root_children_enable_subnet_owner_set_weights() {
     });
 }
 
-// Test that register_network automatically schedules root validators as parents of the
-// subnet owner, enabling the owner to set weights after cooldown.
+// Test that register_network automatically sets root validators as parents of the
+// subnet owner, enabling the owner to set weights. Since SubtokenEnabled is false
+// for a new subnet (start_call hasn't executed yet), child keys are applied immediately.
 // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --package pallet-subtensor --lib -- tests::children::test_register_network_schedules_root_validators --exact --show-output --nocapture
 #[test]
 fn test_register_network_schedules_root_validators() {
@@ -4439,21 +4482,7 @@ fn test_register_network_schedules_root_validators() {
             root_stake,
         );
 
-        // --- Verify pending children were scheduled during registration ---
-        assert!(
-            PendingChildKeys::<Test>::contains_key(netuid, root_val_hotkey_1),
-            "Root validator 1 should have pending children on the new subnet"
-        );
-        assert!(
-            PendingChildKeys::<Test>::contains_key(netuid, root_val_hotkey_2),
-            "Root validator 2 should have pending children on the new subnet"
-        );
-
-        // --- Activate pending children ---
-        step_block(1);
-        SubtensorModule::do_set_pending_children(netuid);
-
-        // --- Verify child-parent relationships ---
+        // --- Verify child keys were applied immediately (SubtokenEnabled is false for new subnets) ---
         let children_1 = SubtensorModule::get_children(&root_val_hotkey_1, netuid);
         assert_eq!(
             children_1,