Feat/use tx pool for set weights (#2534)

* use the tx pool for setting/committing weights

* add supports rpc method with cache

* handle when state call is not possible (unlikely)

* increment next index by one for correct nonce

* use for set sync set weights also

* oops, define in subtensor for sync methods

* make tests less specific

* try e2e test for commit weights

* modify comments on test

* add root set test helper

* try upping submit timeout

* add set weights e2e test

* use high timeout on tests

* add awaits and log netuid

* chore: ruff

* use sudo_set_admin_utils instead

* don't increment next index for nonce

* dont await sudo_set_admin_utils

* lower test ext timeout to 12

* enable reg during test

* use set hyp values instead

* sleep after setting reg allowed

* fix test_set_weights_uses_next_nonce

* Use asyncstdlib for lru_cache of `AsyncSubstrateInterface.supports_rpc_method`

* Increases sleep for txs in test_set_weights_uses_next_nonce

---------

Co-authored-by: Benjamin Himes <[email protected]>
Co-authored-by: ibraheem-opentensor <[email protected]>
Co-authored-by: Benjamin Himes <[email protected]>

bittensor/core/extrinsics/async_weights.py

@@ -58,11 +58,17 @@ async def _do_set_weights(
             "version_key": version_key,
         },
     )
+
+    next_nonce = await subtensor.substrate.get_account_next_index(
+        wallet.hotkey.ss58_address
+    )
+
     # Period dictates how long the extrinsic will stay as part of waiting pool
     extrinsic = await subtensor.substrate.create_signed_extrinsic(
         call=call,
         keypair=wallet.hotkey,
         era={"period": 5},
+        nonce=next_nonce,
     )
     response = await subtensor.substrate.submit_extrinsic(
         extrinsic,
@@ -180,9 +186,15 @@ async def _do_commit_weights(
             "commit_hash": commit_hash,
         },
     )
+
+    next_nonce = await subtensor.substrate.get_account_next_index(
+        wallet.hotkey.ss58_address
+    )
+
     extrinsic = await subtensor.substrate.create_signed_extrinsic(
         call=call,
         keypair=wallet.hotkey,
+        nonce=next_nonce,
     )
     response = await subtensor.substrate.submit_extrinsic(
         substrate=subtensor.substrate,

bittensor/core/extrinsics/commit_weights.py

@@ -66,9 +66,11 @@ def do_commit_weights(
             "commit_hash": commit_hash,
         },
     )
+    next_nonce = self.get_account_next_index(wallet.hotkey.ss58_address)
     extrinsic = self.substrate.create_signed_extrinsic(
         call=call,
         keypair=wallet.hotkey,
+        nonce=next_nonce,
     )
     response = submit_extrinsic(
         subtensor=self,

bittensor/core/extrinsics/set_weights.py

@@ -76,11 +76,13 @@ def do_set_weights(
             "version_key": version_key,
         },
     )
+    next_nonce = self.get_account_next_index(wallet.hotkey.ss58_address)
     # Period dictates how long the extrinsic will stay as part of waiting pool
     extrinsic = self.substrate.create_signed_extrinsic(
         call=call,
         keypair=wallet.hotkey,
         era={"period": period},
+        nonce=next_nonce,
     )
     response = submit_extrinsic(
         self,

bittensor/core/subtensor.py

@@ -661,6 +661,16 @@ def query_module(
             ),
         )
 
+    @networking.ensure_connected
+    def get_account_next_index(self, address: str) -> int:
+        """
+        Returns the next nonce for an account, taking into account the transaction pool.
+        """
+        if not self.substrate.supports_rpc_method("account_nextIndex"):
+            raise Exception("account_nextIndex not supported")
+
+        return self.substrate.rpc_request("account_nextIndex", [address])["result"]
+
     # Common subtensor methods =========================================================================================
     def metagraph(
         self, netuid: int, lite: bool = True, block: Optional[int] = None

bittensor/utils/async_substrate_interface.py

@@ -13,6 +13,7 @@
 from hashlib import blake2b
 from typing import Optional, Any, Union, Callable, Awaitable, cast, TYPE_CHECKING
 
+import asyncstdlib as a
 from async_property import async_property
 from bittensor_wallet import Keypair
 from bt_decode import PortableRegistry, decode as decode_by_type_string, MetadataV15
@@ -1704,6 +1705,24 @@ def make_payload(id_: str, method: str, params: list) -> dict:
             "payload": {"jsonrpc": "2.0", "method": method, "params": params},
         }
 
+    @a.lru_cache(maxsize=512)  # RPC methods are unlikely to change often
+    async def supports_rpc_method(self, name: str) -> bool:
+        """
+        Check if substrate RPC supports given method
+        Parameters
+        ----------
+        name: name of method to check
+
+        Returns
+        -------
+        bool
+        """
+        result = await self.rpc_request("rpc_methods", []).get("result")
+        if result:
+            self.config["rpc_methods"] = result.get("methods", [])
+
+        return name in self.config["rpc_methods"]
+
     async def rpc_request(
         self,
         method: str,
@@ -2296,10 +2315,33 @@ async def get_account_nonce(self, account_address: str) -> int:
         Returns:
             Nonce for given account address
         """
-        nonce_obj = await self.runtime_call(
-            "AccountNonceApi", "account_nonce", [account_address]
-        )
-        return nonce_obj.value
+        if await self.supports_rpc_method("state_call"):
+            nonce_obj = await self.runtime_call(
+                "AccountNonceApi", "account_nonce", [account_address]
+            )
+            return nonce_obj
+        else:
+            response = await self.query(
+                module="System", storage_function="Account", params=[account_address]
+            )
+            return response["nonce"]
+
+    async def get_account_next_index(self, account_address: str) -> int:
+        """
+        Returns next index for the given account address, taking into account the transaction pool.
+
+        Args:
+            account_address: SS58 formatted address
+
+        Returns:
+            Next index for the given account address
+        """
+        if not await self.supports_rpc_method("account_nextIndex"):
+            # Unlikely to happen, this is a common RPC method
+            raise Exception("account_nextIndex not supported")
+
+        nonce_obj = await self.rpc_request("account_nextIndex", [account_address])
+        return nonce_obj["result"]
 
     async def get_metadata_constant(self, module_name, constant_name, block_hash=None):
         """

requirements/prod.txt

@@ -1,6 +1,7 @@
 wheel
 setuptools~=70.0.0
 aiohttp~=3.9
+asyncstdlib~=3.13.0
 async-property==0.2.2
 bittensor-cli
 bt-decode==0.4.0

tests/e2e_tests/test_commit_weights.py

@@ -2,6 +2,9 @@
 
 import numpy as np
 import pytest
+
+import asyncio
+
 from bittensor.core.subtensor import Subtensor
 from bittensor.utils.balance import Balance
 from bittensor.utils.weight_utils import convert_weights_and_uids_for_emit
@@ -171,3 +174,158 @@ async def test_commit_and_reveal_weights_legacy(local_chain):
         weight_vals[0] == revealed_weights.value[0][1]
     ), f"Incorrect revealed weights. Expected: {weights[0]}, Actual: {revealed_weights.value[0][1]}"
     print("✅ Passed test_commit_and_reveal_weights")
+
+
+@pytest.mark.asyncio
+async def test_commit_weights_uses_next_nonce(local_chain):
+    """
+    Tests that commiting weights doesn't re-use a nonce in the transaction pool.
+
+    Steps:
+        1. Register a subnet through Alice
+        2. Register Alice's neuron and add stake
+        3. Enable commit-reveal mechanism on the subnet
+        4. Lower the commit_reveal interval and rate limit
+        5. Commit weights three times
+        6. Assert that all commits succeeded
+    Raises:
+        AssertionError: If any of the checks or verifications fail
+    """
+    netuid = 1
+    utils.EXTRINSIC_SUBMISSION_TIMEOUT = 12  # handle fast blocks
+    print("Testing test_commit_and_reveal_weights")
+    # Register root as Alice
+    keypair, alice_wallet = setup_wallet("//Alice")
+    assert register_subnet(local_chain, alice_wallet), "Unable to register the subnet"
+
+    # Verify subnet 1 created successfully
+    assert local_chain.query(
+        "SubtensorModule", "NetworksAdded", [1]
+    ).serialize(), "Subnet wasn't created successfully"
+
+    subtensor = Subtensor(network="ws://localhost:9945")
+
+    # Register Alice to the subnet
+    assert subtensor.burned_register(
+        alice_wallet, netuid
+    ), "Unable to register Alice as a neuron"
+
+    # Stake to become to top neuron after the first epoch
+    add_stake(local_chain, alice_wallet, Balance.from_tao(100_000))
+
+    # Enable commit_reveal on the subnet
+    assert sudo_set_hyperparameter_bool(
+        local_chain,
+        alice_wallet,
+        "sudo_set_commit_reveal_weights_enabled",
+        True,
+        netuid,
+    ), "Unable to enable commit reveal on the subnet"
+
+    assert subtensor.get_subnet_hyperparameters(
+        netuid=netuid,
+    ).commit_reveal_weights_enabled, "Failed to enable commit/reveal"
+
+    # Lower the commit_reveal interval
+    assert sudo_set_hyperparameter_values(
+        local_chain,
+        alice_wallet,
+        call_function="sudo_set_commit_reveal_weights_interval",
+        call_params={"netuid": netuid, "interval": "1"},
+        return_error_message=True,
+    )
+
+    assert (
+        subtensor.get_subnet_hyperparameters(
+            netuid=netuid
+        ).commit_reveal_weights_interval
+        == 1
+    ), "Failed to set commit/reveal periods"
+
+    assert (
+        subtensor.weights_rate_limit(netuid=netuid) > 0
+    ), "Weights rate limit is below 0"
+    # Lower the rate limit
+    assert sudo_set_hyperparameter_values(
+        local_chain,
+        alice_wallet,
+        call_function="sudo_set_weights_set_rate_limit",
+        call_params={"netuid": netuid, "weights_set_rate_limit": "0"},
+        return_error_message=True,
+    )
+
+    assert (
+        subtensor.get_subnet_hyperparameters(netuid=netuid).weights_rate_limit == 0
+    ), "Failed to set weights_rate_limit"
+    assert subtensor.weights_rate_limit(netuid=netuid) == 0
+
+    # Commit-reveal values
+    uids = np.array([0], dtype=np.int64)
+    weights = np.array([0.1], dtype=np.float32)
+    salt = [18, 179, 107, 0, 165, 211, 141, 197]
+    weight_uids, weight_vals = convert_weights_and_uids_for_emit(
+        uids=uids, weights=weights
+    )
+
+    # Make a second salt
+    salt2 = salt.copy()
+    salt2[0] += 1  # Increment the first byte to produce a different commit hash
+
+    # Make a third salt
+    salt3 = salt.copy()
+    salt3[0] += 2  # Increment the first byte to produce a different commit hash
+
+    # Commit all three salts
+    success, message = subtensor.commit_weights(
+        alice_wallet,
+        netuid,
+        salt=salt,
+        uids=weight_uids,
+        weights=weight_vals,
+        wait_for_inclusion=False,  # Don't wait for inclusion, we are testing the nonce when there is a tx in the pool
+        wait_for_finalization=False,
+    )
+
+    assert success is True
+
+    success, message = subtensor.commit_weights(
+        alice_wallet,
+        netuid,
+        salt=salt2,
+        uids=weight_uids,
+        weights=weight_vals,
+        wait_for_inclusion=False,
+        wait_for_finalization=False,
+    )
+
+    assert success is True
+
+    # Commit the third salt
+    success, message = subtensor.commit_weights(
+        alice_wallet,
+        netuid,
+        salt=salt3,
+        uids=weight_uids,
+        weights=weight_vals,
+        wait_for_inclusion=False,
+        wait_for_finalization=False,
+    )
+
+    assert success is True
+
+    # Wait a few blocks
+    await asyncio.sleep(2)  # Wait for the txs to be included in the chain
+
+    # Query the WeightCommits storage map for all three salts
+    weight_commits = subtensor.query_module(
+        module="SubtensorModule",
+        name="WeightCommits",
+        params=[netuid, alice_wallet.hotkey.ss58_address],
+    )
+    # Assert that the committed weights are set correctly
+    assert weight_commits.value is not None, "Weight commit not found in storage"
+    commit_hash, commit_block, reveal_block, expire_block = weight_commits.value[0]
+    assert commit_block > 0, f"Invalid block number: {commit_block}"
+
+    # Check for three commits in the WeightCommits storage map
+    assert len(weight_commits.value) == 3, "Expected 3 weight commits"

tests/e2e_tests/test_incentive.py

@@ -15,6 +15,7 @@
     templates_repo,
 )
 from bittensor.utils.balance import Balance
+from bittensor.core.extrinsics import utils
 from bittensor.core.extrinsics.set_weights import do_set_weights
 from bittensor.core.metagraph import Metagraph
 
@@ -40,6 +41,8 @@ async def test_incentive(local_chain):
     print("Testing test_incentive")
     netuid = 1
 
+    utils.EXTRINSIC_SUBMISSION_TIMEOUT = 12  # handle fast blocks
+
     # Register root as Alice - the subnet owner and validator
     alice_keypair, alice_wallet = setup_wallet("//Alice")
     register_subnet(local_chain, alice_wallet)