Merge branch 'staging' into release/8.4.0

Author: ibraheem-abe

bittensor/core/async_subtensor.py

@@ -831,29 +831,18 @@ async def neurons(
 
         Understanding the distribution and status of neurons within a subnet is key to comprehending the network's decentralized structure and the dynamics of its consensus and governance processes.
         """
-        if not block_hash:
-            if reuse_block:
-                block_hash = self.substrate.last_block_hash
-            else:
-                block_hash = await self.substrate.get_chain_head()
-
-        neurons_lite, weights, bonds = await asyncio.gather(
-            self.neurons_lite(netuid=netuid, block_hash=block_hash),
-            self.weights(netuid=netuid, block_hash=block_hash),
-            self.bonds(netuid=netuid, block_hash=block_hash),
+        hex_bytes_result = await self.query_runtime_api(
+            runtime_api="NeuronInfoRuntimeApi",
+            method="get_neurons",
+            params=[netuid],
+            block_hash=block_hash,
+            reuse_block=reuse_block,
         )
 
-        weights_as_dict = {uid: w for uid, w in weights}
-        bonds_as_dict = {uid: b for uid, b in bonds}
-
-        neurons = [
-            NeuronInfo.from_weights_bonds_and_neuron_lite(
-                neuron_lite, weights_as_dict, bonds_as_dict
-            )
-            for neuron_lite in neurons_lite
-        ]
+        if hex_bytes_result is None:
+            return []
 
-        return neurons
+        return NeuronInfo.list_from_vec_u8(hex_to_bytes(hex_bytes_result))
 
     async def neurons_lite(
         self, netuid: int, block_hash: Optional[str] = None, reuse_block: bool = False

bittensor/core/chain_data/neuron_info.py

@@ -174,3 +174,57 @@ def from_vec_u8(cls, vec_u8: bytes) -> "NeuronInfo":
             ),
             is_null=False,
         )
+
+    @classmethod
+    def list_from_vec_u8(cls, vec_u8: bytes) -> list["NeuronInfo"]:
+        nn = bt_decode.NeuronInfo.decode_vec(bytes(vec_u8))
+
+        def fix(n):
+            stake_dict = process_stake_data(n.stake)
+            total_stake = sum(stake_dict.values()) if stake_dict else Balance(0)
+            axon_info = n.axon_info
+            coldkey = decode_account_id(n.coldkey)
+            hotkey = decode_account_id(n.hotkey)
+            return NeuronInfo(
+                hotkey=hotkey,
+                coldkey=coldkey,
+                uid=n.uid,
+                netuid=n.netuid,
+                active=n.active,
+                stake=total_stake,
+                stake_dict=stake_dict,
+                total_stake=total_stake,
+                rank=u16_normalized_float(n.rank),
+                emission=n.emission / 1e9,
+                incentive=u16_normalized_float(n.incentive),
+                consensus=u16_normalized_float(n.consensus),
+                trust=u16_normalized_float(n.trust),
+                validator_trust=u16_normalized_float(n.validator_trust),
+                dividends=u16_normalized_float(n.dividends),
+                last_update=n.last_update,
+                validator_permit=n.validator_permit,
+                weights=[[e[0], e[1]] for e in n.weights],
+                bonds=[[e[0], e[1]] for e in n.bonds],
+                pruning_score=n.pruning_score,
+                prometheus_info=PrometheusInfo(
+                    block=n.prometheus_info.block,
+                    version=n.prometheus_info.version,
+                    ip=str(netaddr.IPAddress(n.prometheus_info.ip)),
+                    port=n.prometheus_info.port,
+                    ip_type=n.prometheus_info.ip_type,
+                ),
+                axon_info=AxonInfo(
+                    version=axon_info.version,
+                    ip=str(netaddr.IPAddress(axon_info.ip)),
+                    port=axon_info.port,
+                    ip_type=axon_info.ip_type,
+                    placeholder1=axon_info.placeholder1,
+                    placeholder2=axon_info.placeholder2,
+                    protocol=axon_info.protocol,
+                    hotkey=hotkey,
+                    coldkey=coldkey,
+                ),
+                is_null=False,
+            )
+
+        return [fix(n) for n in nn]

bittensor/core/subtensor.py

@@ -216,7 +216,7 @@ def _get_substrate(self):
         """Establishes a connection to the Substrate node using configured parameters."""
         try:
             # Set up params.
-            if not self.websocket:
+            if self.websocket is None or self.websocket.close_code is not None:
                 self.websocket = ws_client.connect(
                     self.chain_endpoint,
                     open_timeout=self._connection_timeout,

bittensor/utils/async_substrate_interface.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import inspect
 import json
 import random
 from collections import defaultdict
@@ -1171,14 +1172,14 @@ async def _get_block_handler(
         include_author: bool = False,
         header_only: bool = False,
         finalized_only: bool = False,
-        subscription_handler: Optional[Callable] = None,
+        subscription_handler: Optional[Callable[[dict], Awaitable[Any]]] = None,
     ):
         try:
             await self.init_runtime(block_hash=block_hash)
         except BlockNotFound:
             return None
 
-        async def decode_block(block_data, block_data_hash=None):
+        async def decode_block(block_data, block_data_hash=None) -> dict[str, Any]:
             if block_data:
                 if block_data_hash:
                     block_data["header"]["hash"] = block_data_hash
@@ -1193,12 +1194,12 @@ async def decode_block(block_data, block_data_hash=None):
 
                 if "extrinsics" in block_data:
                     for idx, extrinsic_data in enumerate(block_data["extrinsics"]):
-                        extrinsic_decoder = extrinsic_cls(
-                            data=ScaleBytes(extrinsic_data),
-                            metadata=self.__metadata,
-                            runtime_config=self.runtime_config,
-                        )
                         try:
+                            extrinsic_decoder = extrinsic_cls(
+                                data=ScaleBytes(extrinsic_data),
+                                metadata=self.__metadata,
+                                runtime_config=self.runtime_config,
+                            )
                             extrinsic_decoder.decode(check_remaining=True)
                             block_data["extrinsics"][idx] = extrinsic_decoder
 
@@ -1314,23 +1315,29 @@ async def decode_block(block_data, block_data_hash=None):
         if callable(subscription_handler):
             rpc_method_prefix = "Finalized" if finalized_only else "New"
 
-            async def result_handler(message, update_nr, subscription_id):
-                new_block = await decode_block({"header": message["params"]["result"]})
+            async def result_handler(
+                message: dict, subscription_id: str
+            ) -> tuple[Any, bool]:
+                reached = False
+                subscription_result = None
+                if "params" in message:
+                    new_block = await decode_block(
+                        {"header": message["params"]["result"]}
+                    )
 
-                subscription_result = subscription_handler(
-                    new_block, update_nr, subscription_id
-                )
+                    subscription_result = await subscription_handler(new_block)
 
-                if subscription_result is not None:
-                    # Handler returned end result: unsubscribe from further updates
-                    self._forgettable_task = asyncio.create_task(
-                        self.rpc_request(
-                            f"chain_unsubscribe{rpc_method_prefix}Heads",
-                            [subscription_id],
+                    if subscription_result is not None:
+                        reached = True
+                        # Handler returned end result: unsubscribe from further updates
+                        self._forgettable_task = asyncio.create_task(
+                            self.rpc_request(
+                                f"chain_unsubscribe{rpc_method_prefix}Heads",
+                                [subscription_id],
+                            )
                         )
-                    )
 
-                return subscription_result
+                return subscription_result, reached
 
             result = await self._make_rpc_request(
                 [
@@ -1343,7 +1350,7 @@ async def result_handler(message, update_nr, subscription_id):
                 result_handler=result_handler,
             )
 
-            return result
+            return result["_get_block_handler"][-1]
 
         else:
             if header_only:
@@ -2770,3 +2777,49 @@ async def close(self):
             await self.ws.shutdown()
         except AttributeError:
             pass
+
+    async def wait_for_block(
+        self,
+        block: int,
+        result_handler: Callable[[dict], Awaitable[Any]],
+        task_return: bool = True,
+    ) -> Union[asyncio.Task, Union[bool, Any]]:
+        """
+        Executes the result_handler when the chain has reached the block specified.
+
+        Args:
+            block: block number
+            result_handler: coroutine executed upon reaching the block number. This can be basically anything, but
+                must accept one single arg, a dict with the block data; whether you use this data or not is entirely
+                up to you.
+            task_return: True to immediately return the result of wait_for_block as an asyncio Task, False to wait
+                for the block to be reached, and return the result of the result handler.
+
+        Returns:
+            Either an asyncio.Task (which contains the running subscription, and whose `result()` will contain the
+                return of the result_handler), or the result itself, depending on `task_return` flag.
+                Note that if your result_handler returns `None`, this method will return `True`, otherwise
+                the return will be the result of your result_handler.
+        """
+
+        async def _handler(block_data: dict[str, Any]):
+            required_number = block
+            number = block_data["header"]["number"]
+            if number >= required_number:
+                return (
+                    r if (r := await result_handler(block_data)) is not None else True
+                )
+
+        args = inspect.getfullargspec(result_handler).args
+        if len(args) != 1:
+            raise ValueError(
+                "result_handler must take exactly one arg: the dict block data."
+            )
+
+        co = self._get_block_handler(
+            self.last_block_hash, subscription_handler=_handler
+        )
+        if task_return is True:
+            return asyncio.create_task(co)
+        else:
+            return await co

bittensor/utils/mock/subtensor_mock.py

@@ -20,9 +20,11 @@
 from hashlib import sha256
 from types import SimpleNamespace
 from typing import Any, Optional, Union, TypedDict
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 from bittensor_wallet import Wallet
+from substrateinterface.base import SubstrateInterface
+from websockets.sync.client import ClientConnection
 
 from bittensor.core.chain_data import (
     NeuronInfo,
@@ -33,6 +35,7 @@
 from bittensor.core.types import AxonServeCallParams, PrometheusServeCallParams
 from bittensor.core.errors import ChainQueryError
 from bittensor.core.subtensor import Subtensor
+import bittensor.core.subtensor as subtensor_module
 from bittensor.utils import RAOPERTAO, u16_normalized_float
 from bittensor.utils.balance import Balance
 
@@ -248,14 +251,22 @@ def setup(self) -> None:
 
             self.network = "mock"
             self.chain_endpoint = "ws://mock_endpoint.bt"
-            self.substrate = MagicMock()
+            self.substrate = MagicMock(autospec=SubstrateInterface)
 
     def __init__(self, *args, **kwargs) -> None:
-        super().__init__()
-        self.__dict__ = __GLOBAL_MOCK_STATE__
-
-        if not hasattr(self, "chain_state") or getattr(self, "chain_state") is None:
-            self.setup()
+        mock_substrate_interface = MagicMock(autospec=SubstrateInterface)
+        mock_websocket = MagicMock(autospec=ClientConnection)
+        mock_websocket.close_code = None
+        with patch.object(
+            subtensor_module,
+            "SubstrateInterface",
+            return_value=mock_substrate_interface,
+        ):
+            super().__init__(websocket=mock_websocket)
+            self.__dict__ = __GLOBAL_MOCK_STATE__
+
+            if not hasattr(self, "chain_state") or getattr(self, "chain_state") is None:
+                self.setup()
 
     def get_block_hash(self, block_id: int) -> str:
         return "0x" + sha256(str(block_id).encode()).hexdigest()[:64]

bittensor/utils/networking.py

@@ -10,7 +10,7 @@
 import netaddr
 import requests
 from retry import retry
-from websocket import WebSocketConnectionClosedException
+from websockets.exceptions import ConnectionClosed
 
 from bittensor.utils.btlogging import logging
 
@@ -193,7 +193,7 @@ def wrapper(self, *args, **kwargs):
 
         try:
             return func(self, *args, **kwargs)
-        except WebSocketConnectionClosedException:
+        except ConnectionClosed:
             logging.console.warning(
                 "WebSocket connection closed. Attempting to reconnect 5 times..."
             )

tests/helpers/helpers.py

@@ -125,6 +125,8 @@ def get_mock_neuron_by_uid(uid: int, **kwargs) -> NeuronInfo:
 
 
 class FakeWebsocket(ClientConnection):
+    close_code = None
+
     def __init__(self, *args, seed, **kwargs):
         protocol = ClientProtocol(parse_uri("ws://127.0.0.1:9945"))
         super().__init__(socket=None, protocol=protocol, **kwargs)

tests/unit_tests/test_async_subtensor.py

@@ -992,41 +992,32 @@ async def test_neurons(subtensor, mocker):
     # Preps
     fake_netuid = 1
     fake_block_hash = "block_hash"
-    fake_neurons = [mocker.Mock(), mocker.Mock()]
-    fake_weights = [(1, [(10, 20), (30, 40)]), (2, [(50, 60), (70, 80)])]
-    fake_bonds = [(1, [(10, 20), (30, 40)]), (2, [(50, 60), (70, 80)])]
-
-    mocked_neurons_lite = mocker.AsyncMock(return_value=fake_neurons)
-    subtensor.neurons_lite = mocked_neurons_lite
-
-    mocked_weights = mocker.AsyncMock(return_value=fake_weights)
-    subtensor.weights = mocked_weights
-
-    mocked_bonds = mocker.AsyncMock(return_value=fake_bonds)
-    subtensor.bonds = mocked_bonds
+    fake_reuse_block_hash = True
 
-    mocked_neuron_info_method = mocker.Mock()
-    async_subtensor.NeuronInfo.from_weights_bonds_and_neuron_lite = (
-        mocked_neuron_info_method
+    mocked_query_runtime_api = mocker.patch.object(
+        subtensor, "query_runtime_api", return_value="NOT NONE"
+    )
+    mocked_hex_to_bytes = mocker.patch.object(async_subtensor, "hex_to_bytes")
+    mocked_neuron_info_list_from_vec_u8 = mocker.patch.object(
+        async_subtensor.NeuronInfo, "list_from_vec_u8"
     )
-
     # Call
-    result = await subtensor.neurons(netuid=fake_netuid, block_hash=fake_block_hash)
+    result = await subtensor.neurons(
+        netuid=fake_netuid,
+        block_hash=fake_block_hash,
+        reuse_block=fake_reuse_block_hash,
+    )
 
     # Asserts
-    mocked_neurons_lite.assert_awaited_once()
-    mocked_neurons_lite.assert_called_once_with(
-        netuid=fake_netuid, block_hash=fake_block_hash
-    )
-    mocked_weights.assert_awaited_once()
-    mocked_weights.assert_called_once_with(
-        netuid=fake_netuid, block_hash=fake_block_hash
+    mocked_query_runtime_api.assert_called_once_with(
+        runtime_api="NeuronInfoRuntimeApi",
+        method="get_neurons",
+        params=[fake_netuid],
+        block_hash=fake_block_hash,
+        reuse_block=fake_reuse_block_hash,
     )
-    mocked_bonds.assert_awaited_once()
-    mocked_bonds.assert_called_once_with(netuid=fake_netuid, block_hash=fake_block_hash)
-    assert result == [
-        mocked_neuron_info_method.return_value for _ in range(len(fake_neurons))
-    ]
+    mocked_hex_to_bytes.assert_called_once_with(mocked_query_runtime_api.return_value)
+    assert result == mocked_neuron_info_list_from_vec_u8.return_value
 
 
 @pytest.mark.parametrize(