bittensor-cli 8.4.3__py3-none-any.whl → 9.0.0rc2__py3-none-any.whl

bittensor_cli/src/bittensor/extrinsics/registration.py

@@ -18,15 +18,19 @@ import typing
 from typing import Optional
 import subprocess
 
+import backoff
 from bittensor_wallet import Wallet
+from bittensor_wallet.errors import KeyFileError
 from Crypto.Hash import keccak
 import numpy as np
 from rich.prompt import Confirm
 from rich.console import Console
 from rich.status import Status
-from substrateinterface.exceptions import SubstrateRequestException
+from async_substrate_interface.errors import SubstrateRequestException
 
+from bittensor_cli.src import COLOR_PALETTE
 from bittensor_cli.src.bittensor.chain_data import NeuronInfo
+from bittensor_cli.src.bittensor.balances import Balance
 from bittensor_cli.src.bittensor.utils import (
     console,
     err_console,
@@ -35,8 +39,6 @@ from bittensor_cli.src.bittensor.utils import (
     get_human_readable,
     print_verbose,
     print_error,
-    unlock_key,
-    hex_to_bytes,
 )
 
 if typing.TYPE_CHECKING:
@@ -436,7 +438,7 @@ async def is_hotkey_registered(
     subtensor: "SubtensorInterface", netuid: int, hotkey_ss58: str
 ) -> bool:
     """Checks to see if the hotkey is registered on a given netuid"""
-    _result = await subtensor.substrate.query(
+    _result = await subtensor.query(
         module="SubtensorModule",
         storage_function="Uids",
         params=[netuid, hotkey_ss58],
@@ -487,22 +489,18 @@ async def register_extrinsic(
     """
 
     async def get_neuron_for_pubkey_and_subnet():
-        uid = await subtensor.substrate.query(
+        uid = await subtensor.query(
             "SubtensorModule", "Uids", [netuid, wallet.hotkey.ss58_address]
         )
         if uid is None:
             return NeuronInfo.get_null_neuron()
 
-        params = [netuid, uid]
-        json_body = await subtensor.substrate.rpc_request(
-            method="neuronInfo_getNeuron",
-            params=params,
+        result = await subtensor.neuron_for_uid(
+            uid=uid,
+            netuid=netuid,
+            block_hash=subtensor.substrate.last_block_hash,
         )
-
-        if not (result := json_body.get("result", None)):
-            return NeuronInfo.get_null_neuron()
-
-        return NeuronInfo.from_vec_u8(bytes(result))
+        return result
 
     print_verbose("Checking subnet status")
     if not await subtensor.subnet_exists(netuid):
@@ -526,9 +524,9 @@ async def register_extrinsic(
     if prompt:
         if not Confirm.ask(
             f"Continue Registration?\n"
-            f"  hotkey ({wallet.hotkey_str}):\t[bold white]{wallet.hotkey.ss58_address}[/bold white]\n"
-            f"  coldkey ({wallet.name}):\t[bold white]{wallet.coldkeypub.ss58_address}[/bold white]\n"
-            f"  network:\t\t[bold white]{subtensor.network}[/bold white]"
+            f"  hotkey [{COLOR_PALETTE['GENERAL']['HOTKEY']}]({wallet.hotkey_str})[/{COLOR_PALETTE['GENERAL']['HOTKEY']}]:\t[{COLOR_PALETTE['GENERAL']['HOTKEY']}]{wallet.hotkey.ss58_address}[/{COLOR_PALETTE['GENERAL']['HOTKEY']}]\n"
+            f"  coldkey [{COLOR_PALETTE['GENERAL']['COLDKEY']}]({wallet.name})[/{COLOR_PALETTE['GENERAL']['COLDKEY']}]:\t[{COLOR_PALETTE['GENERAL']['COLDKEY']}]{wallet.coldkeypub.ss58_address}[/{COLOR_PALETTE['GENERAL']['COLDKEY']}]\n"
+            f"  network:\t\t[{COLOR_PALETTE['GENERAL']['LINKS']}]{subtensor.network}[/{COLOR_PALETTE['GENERAL']['LINKS']}]\n"
         ):
             return False
 
@@ -581,7 +579,7 @@ async def register_extrinsic(
             )
             if is_registered:
                 err_console.print(
-                    f":white_heavy_check_mark: [green]Already registered on netuid:{netuid}[/green]"
+                    f":white_heavy_check_mark: [dark_sea_green3]Already registered on netuid:{netuid}[/dark_sea_green3]"
                 )
                 return True
 
@@ -618,15 +616,18 @@ async def register_extrinsic(
                         if not success:
                             success, err_msg = (
                                 False,
-                                format_error_message(await response.error_message),
+                                format_error_message(
+                                    await response.error_message,
+                                    substrate=subtensor.substrate,
+                                ),
                             )
                             # Look error here
                             # https://github.com/opentensor/subtensor/blob/development/pallets/subtensor/src/errors.rs
 
                             if "HotKeyAlreadyRegisteredInSubNet" in err_msg:
                                 console.print(
-                                    f":white_heavy_check_mark: [green]Already Registered on "
-                                    f"[bold]subnet:{netuid}[/bold][/green]"
+                                    f":white_heavy_check_mark: [dark_sea_green3]Already Registered on "
+                                    f"[bold]subnet:{netuid}[/bold][/dark_sea_green3]"
                                 )
                                 return True
                             err_console.print(
@@ -644,7 +645,7 @@ async def register_extrinsic(
                         )
                         if is_registered:
                             console.print(
-                                ":white_heavy_check_mark: [green]Registered[/green]"
+                                ":white_heavy_check_mark: [dark_sea_green3]Registered[/dark_sea_green3]"
                             )
                             return True
                         else:
@@ -671,6 +672,117 @@ async def register_extrinsic(
             return False
 
 
+async def burned_register_extrinsic(
+    subtensor: "SubtensorInterface",
+    wallet: Wallet,
+    netuid: int,
+    old_balance: Balance,
+    wait_for_inclusion: bool = True,
+    wait_for_finalization: bool = True,
+    prompt: bool = False,
+) -> bool:
+    """Registers the wallet to chain by recycling TAO.
+
+    :param subtensor: The SubtensorInterface object to use for the call, initialized
+    :param wallet: Bittensor wallet object.
+    :param netuid: The `netuid` of the subnet to register on.
+    :param old_balance: The wallet balance prior to the registration burn.
+    :param wait_for_inclusion: If set, waits for the extrinsic to enter a block before returning `True`, or returns
+                               `False` if the extrinsic fails to enter the block within the timeout.
+    :param wait_for_finalization: If set, waits for the extrinsic to be finalized on the chain before returning `True`,
+                                  or returns `False` if the extrinsic fails to be finalized within the timeout.
+    :param prompt: If `True`, the call waits for confirmation from the user before proceeding.
+
+    :return: Flag is `True` if extrinsic was finalized or included in the block. If we did not wait for
+             finalization/inclusion, the response is `True`.
+    """
+
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        err_console.print("Error decrypting coldkey (possibly incorrect password)")
+        return False
+
+    with console.status(
+        f":satellite: Checking Account on [bold]subnet:{netuid}[/bold]...",
+        spinner="aesthetic",
+    ) as status:
+        my_uid = await subtensor.query(
+            "SubtensorModule", "Uids", [netuid, wallet.hotkey.ss58_address]
+        )
+
+        print_verbose("Checking if already registered", status)
+        neuron = await subtensor.neuron_for_uid(
+            uid=my_uid,
+            netuid=netuid,
+            block_hash=subtensor.substrate.last_block_hash,
+        )
+
+    if not neuron.is_null:
+        console.print(
+            ":white_heavy_check_mark: [dark_sea_green3]Already Registered[/dark_sea_green3]:\n"
+            f"uid: [{COLOR_PALETTE['GENERAL']['NETUID_EXTRA']}]{neuron.uid}[/{COLOR_PALETTE['GENERAL']['NETUID_EXTRA']}]\n"
+            f"netuid: [{COLOR_PALETTE['GENERAL']['NETUID']}]{neuron.netuid}[/{COLOR_PALETTE['GENERAL']['NETUID']}]\n"
+            f"hotkey: [{COLOR_PALETTE['GENERAL']['HOTKEY']}]{neuron.hotkey}[/{COLOR_PALETTE['GENERAL']['HOTKEY']}]\n"
+            f"coldkey: [{COLOR_PALETTE['GENERAL']['COLDKEY']}]{neuron.coldkey}[/{COLOR_PALETTE['GENERAL']['COLDKEY']}]"
+        )
+        return True
+
+    with console.status(
+        ":satellite: Recycling TAO for Registration...", spinner="aesthetic"
+    ):
+        call = await subtensor.substrate.compose_call(
+            call_module="SubtensorModule",
+            call_function="burned_register",
+            call_params={
+                "netuid": netuid,
+                "hotkey": wallet.hotkey.ss58_address,
+            },
+        )
+        success, err_msg = await subtensor.sign_and_send_extrinsic(
+            call, wallet, wait_for_inclusion, wait_for_finalization
+        )
+
+    if not success:
+        err_console.print(f":cross_mark: [red]Failed[/red]: {err_msg}")
+        await asyncio.sleep(0.5)
+        return False
+    # Successful registration, final check for neuron and pubkey
+    else:
+        with console.status(":satellite: Checking Balance...", spinner="aesthetic"):
+            block_hash = await subtensor.substrate.get_chain_head()
+            new_balance, netuids_for_hotkey, my_uid = await asyncio.gather(
+                subtensor.get_balance(
+                    wallet.coldkeypub.ss58_address,
+                    block_hash=block_hash,
+                    reuse_block=False,
+                ),
+                subtensor.get_netuids_for_hotkey(
+                    wallet.hotkey.ss58_address, block_hash=block_hash
+                ),
+                subtensor.query(
+                    "SubtensorModule", "Uids", [netuid, wallet.hotkey.ss58_address]
+                ),
+            )
+
+        console.print(
+            "Balance:\n"
+            f"  [blue]{old_balance}[/blue] :arrow_right: [{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_balance}[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]"
+        )
+
+        if len(netuids_for_hotkey) > 0:
+            console.print(
+                f":white_heavy_check_mark: [green]Registered on netuid {netuid} with UID {my_uid}[/green]"
+            )
+            return True
+        else:
+            # neuron not found, try again
+            err_console.print(
+                ":cross_mark: [red]Unknown error. Neuron not found.[/red]"
+            )
+            return False
+
+
 async def run_faucet_extrinsic(
     subtensor: "SubtensorInterface",
     wallet: Wallet,
@@ -684,7 +796,7 @@ async def run_faucet_extrinsic(
     tpb: int = 256,
     num_processes: Optional[int] = None,
     update_interval: Optional[int] = None,
-    log_verbose: bool = True,
+    log_verbose: bool = False,
     max_successes: int = 3,
 ) -> tuple[bool, str]:
     r"""Runs a continual POW to get a faucet of TAO on the test net.
@@ -723,8 +835,10 @@ async def run_faucet_extrinsic(
         return False, "Requires torch"
 
     # Unlock coldkey
-    if not (unlock_status := unlock_key(wallet, print_out=False)).success:
-        return False, unlock_status.message
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        return False, "There was an error unlocking your coldkey"
 
     # Get previous balance.
     old_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)
@@ -732,12 +846,8 @@ async def run_faucet_extrinsic(
     # Attempt rolling registration.
     attempts = 1
     successes = 1
-    pow_result: Optional[POWSolution]
     while True:
         try:
-            account_nonce = await subtensor.substrate.get_account_nonce(
-                wallet.coldkey.ss58_address
-            )
             pow_result = None
             while pow_result is None or await pow_result.is_stale(subtensor=subtensor):
                 # Solve latest POW.
@@ -746,7 +856,7 @@ async def run_faucet_extrinsic(
                         if prompt:
                             err_console.print("CUDA is not available.")
                         return False, "CUDA is not available."
-                    pow_result = await create_pow(
+                    pow_result: Optional[POWSolution] = await create_pow(
                         subtensor,
                         wallet,
                         -1,
@@ -759,7 +869,7 @@ async def run_faucet_extrinsic(
                         log_verbose=log_verbose,
                     )
                 else:
-                    pow_result = await create_pow(
+                    pow_result: Optional[POWSolution] = await create_pow(
                         subtensor,
                         wallet,
                         -1,
@@ -779,7 +889,7 @@ async def run_faucet_extrinsic(
                 },
             )
             extrinsic = await subtensor.substrate.create_signed_extrinsic(
-                call=call, keypair=wallet.coldkey, nonce=account_nonce
+                call=call, keypair=wallet.coldkey
             )
             response = await subtensor.substrate.submit_extrinsic(
                 extrinsic,
@@ -792,7 +902,7 @@ async def run_faucet_extrinsic(
             if not await response.is_success:
                 err_console.print(
                     f":cross_mark: [red]Failed[/red]: "
-                    f"{format_error_message(await response.error_message)}"
+                    f"{format_error_message(await response.error_message, subtensor.substrate)}"
                 )
                 if attempts == max_allowed_attempts:
                     raise MaxAttemptsException
@@ -806,8 +916,8 @@ async def run_faucet_extrinsic(
                     wallet.coldkeypub.ss58_address
                 )
                 console.print(
-                    f"Balance: [blue]{old_balance[wallet.coldkeypub.ss58_address]}[/blue] :arrow_right:"
-                    f" [green]{new_balance[wallet.coldkeypub.ss58_address]}[/green]"
+                    f"Balance: [blue]{old_balance}[/blue] :arrow_right:"
+                    f" [green]{new_balance}[/green]"
                 )
                 old_balance = new_balance
 
@@ -864,7 +974,7 @@ async def _check_for_newest_block_and_update(
         block_number, difficulty, block_hash = await _get_block_with_retry(
             subtensor=subtensor, netuid=netuid
         )
-        block_bytes = hex_to_bytes(block_hash)
+        block_bytes = bytes.fromhex(block_hash[2:])
 
         update_curr_block(
             curr_diff,
@@ -912,14 +1022,11 @@ async def _block_solver(
     limit = int(math.pow(2, 256)) - 1
 
     # Establish communication queues
-    ## See the _Solver class for more information on the queues.
+    # See the _Solver class for more information on the queues.
     stop_event = Event()
     stop_event.clear()
 
     solution_queue = Queue()
-    if cuda:
-        num_processes = len(dev_id)
-
     finished_queues = [Queue() for _ in range(num_processes)]
     check_block = Lock()
 
@@ -928,7 +1035,8 @@ async def _block_solver(
     )
 
     if cuda:
-        ## Create a worker per CUDA device
+        # Create a worker per CUDA device
+        num_processes = len(dev_id)
         solvers = [
             _CUDASolver(
                 i,
@@ -971,7 +1079,7 @@ async def _block_solver(
         subtensor=subtensor, netuid=netuid
     )
 
-    block_bytes = hex_to_bytes(block_hash)
+    block_bytes = bytes.fromhex(block_hash[2:])
     old_block_number = block_number
     # Set to current block
     _update_curr_block(
@@ -1246,9 +1354,11 @@ def _terminate_workers_and_wait_for_exit(
             worker.terminate()
 
 
+# TODO verify this works with async
+@backoff.on_exception(backoff.constant, Exception, interval=1, max_tries=3)
 async def _get_block_with_retry(
     subtensor: "SubtensorInterface", netuid: int
-) -> tuple[int, int, str]:
+) -> tuple[int, int, bytes]:
     """
     Gets the current block number, difficulty, and block hash from the substrate node.
 
@@ -1260,9 +1370,10 @@ async def _get_block_with_retry(
     :raises Exception: If the block hash is None.
     :raises ValueError: If the difficulty is None.
     """
-    block = await subtensor.substrate.get_block()
-    block_hash = block["header"]["hash"]
-    block_number = block["header"]["number"]
+    block_number = await subtensor.substrate.get_block_number(None)
+    block_hash = await subtensor.substrate.get_block_hash(
+        block_number
+    )  # TODO check if I need to do all this
     try:
         difficulty = (
             1_000_000
@@ -1635,7 +1746,10 @@ async def swap_hotkey_extrinsic(
         )
         return False
 
-    if not unlock_key(wallet).success:
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        err_console.print("Error decrypting coldkey (possibly incorrect password)")
         return False
 
     if prompt:

bittensor_cli/src/bittensor/extrinsics/root.py

@@ -21,12 +21,13 @@ import time
 from typing import Union, List, TYPE_CHECKING
 
 from bittensor_wallet import Wallet, Keypair
+from bittensor_wallet.errors import KeyFileError
 import numpy as np
 from numpy.typing import NDArray
 from rich.prompt import Confirm
 from rich.table import Table, Column
 from scalecodec import ScaleBytes, U16, Vec
-from substrateinterface.exceptions import SubstrateRequestException
+from async_substrate_interface.errors import SubstrateRequestException
 
 from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface
 from bittensor_cli.src.bittensor.extrinsics.registration import is_hotkey_registered
@@ -36,7 +37,6 @@ from bittensor_cli.src.bittensor.utils import (
     u16_normalized_float,
     print_verbose,
     format_error_message,
-    unlock_key,
 )
 
 if TYPE_CHECKING:
@@ -306,7 +306,10 @@ async def root_register_extrinsic(
              the response is `True`.
     """
 
-    if not unlock_key(wallet).success:
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        err_console.print("Error decrypting coldkey (possibly incorrect password)")
         return False
 
     print_verbose(f"Checking if hotkey ({wallet.hotkey_str}) is registered on root")
@@ -339,7 +342,7 @@ async def root_register_extrinsic(
 
         # Successful registration, final check for neuron and pubkey
         else:
-            uid = await subtensor.substrate.query(
+            uid = await subtensor.query(
                 module="SubtensorModule",
                 storage_function="Uids",
                 params=[0, wallet.hotkey.ss58_address],
@@ -416,7 +419,7 @@ async def set_root_weights_extrinsic(
         else:
             return False, await response.error_message
 
-    my_uid = await subtensor.substrate.query(
+    my_uid = await subtensor.query(
         "SubtensorModule", "Uids", [0, wallet.hotkey.ss58_address]
     )
 
@@ -424,7 +427,10 @@ async def set_root_weights_extrinsic(
         err_console.print("Your hotkey is not registered to the root network")
         return False
 
-    if not unlock_key(wallet).success:
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        err_console.print("Error decrypting coldkey (possibly incorrect password)")
         return False
 
     # First convert types.
@@ -486,11 +492,11 @@ async def set_root_weights_extrinsic(
                 console.print(":white_heavy_check_mark: [green]Finalized[/green]")
                 return True
             else:
-                fmt_err = format_error_message(error_message)
+                fmt_err = format_error_message(error_message, subtensor.substrate)
                 err_console.print(f":cross_mark: [red]Failed[/red]: {fmt_err}")
                 return False
 
     except SubstrateRequestException as e:
-        fmt_err = format_error_message(e)
+        fmt_err = format_error_message(e, subtensor.substrate)
         err_console.print(":cross_mark: [red]Failed[/red]: error:{}".format(fmt_err))
         return False

bittensor_cli/src/bittensor/extrinsics/transfer.py

@@ -1,8 +1,9 @@
 import asyncio
 
 from bittensor_wallet import Wallet
+from bittensor_wallet.errors import KeyFileError
 from rich.prompt import Confirm
-from substrateinterface.exceptions import SubstrateRequestException
+from async_substrate_interface.errors import SubstrateRequestException
 
 from bittensor_cli.src import NETWORK_EXPLORER_MAP
 from bittensor_cli.src.bittensor.balances import Balance
@@ -14,8 +15,6 @@ from bittensor_cli.src.bittensor.utils import (
     format_error_message,
     get_explorer_url_for_network,
     is_valid_bittensor_address_or_public_key,
-    print_error,
-    unlock_key,
 )
 
 
@@ -24,7 +23,6 @@ async def transfer_extrinsic(
     wallet: Wallet,
     destination: str,
     amount: Balance,
-    transfer_all: bool = False,
     wait_for_inclusion: bool = True,
     wait_for_finalization: bool = False,
     keep_alive: bool = True,
@@ -36,7 +34,6 @@ async def transfer_extrinsic(
     :param wallet: Bittensor wallet object to make transfer from.
     :param destination: Destination public key address (ss58_address or ed25519) of recipient.
     :param amount: Amount to stake as Bittensor balance.
-    :param transfer_all: Whether to transfer all funds from this wallet to the destination address.
     :param wait_for_inclusion: If set, waits for the extrinsic to enter a block before returning `True`,
                                or returns `False` if the extrinsic fails to enter the block within the timeout.
     :param wait_for_finalization:  If set, waits for the extrinsic to be finalized on the chain before returning
@@ -64,14 +61,14 @@ async def transfer_extrinsic(
                 call=call, keypair=wallet.coldkeypub
             )
         except SubstrateRequestException as e:
-            payment_info = {"partialFee": int(2e7)}  # assume  0.02 Tao
+            payment_info = {"partial_fee": int(2e7)}  # assume  0.02 Tao
             err_console.print(
-                f":cross_mark: [red]Failed to get payment info[/red]:[bold white]\n"
-                f"  {format_error_message(e)}[/bold white]\n"
-                f"  Defaulting to default transfer fee: {payment_info['partialFee']}"
+                f":cross_mark: [red]Failed to get payment info[/red]:\n"
+                f"  [bold white]{format_error_message(e)}[/bold white]\n"
+                f"  Defaulting to default transfer fee: {payment_info['partial_fee']}"
             )
 
-        return Balance.from_rao(payment_info["partialFee"])
+        return Balance.from_rao(payment_info["partial_fee"])
 
     async def do_transfer() -> tuple[bool, str, str]:
         """
@@ -104,7 +101,7 @@ async def transfer_extrinsic(
             return (
                 False,
                 "",
-                format_error_message(await response.error_message),
+                format_error_message(await response.error_message, subtensor.substrate),
             )
 
     # Validate destination address.
@@ -115,7 +112,10 @@ async def transfer_extrinsic(
         return False
     console.print(f"[dark_orange]Initiating transfer on network: {subtensor.network}")
     # Unlock wallet coldkey.
-    if not unlock_key(wallet).success:
+    try:
+        wallet.unlock_coldkey()
+    except KeyFileError:
+        err_console.print("Error decrypting coldkey (possibly incorrect password)")
         return False
 
     # Check balance.
@@ -126,13 +126,12 @@ async def transfer_extrinsic(
         # check existential deposit and fee
         print_verbose("Fetching existential and fee", status)
         block_hash = await subtensor.substrate.get_chain_head()
-        account_balance_, existential_deposit = await asyncio.gather(
+        account_balance, existential_deposit = await asyncio.gather(
             subtensor.get_balance(
                 wallet.coldkeypub.ss58_address, block_hash=block_hash
             ),
             subtensor.get_existential_deposit(block_hash=block_hash),
         )
-        account_balance = account_balance_[wallet.coldkeypub.ss58_address]
         fee = await get_transfer_fee()
 
     if not keep_alive:
@@ -140,12 +139,6 @@ async def transfer_extrinsic(
         existential_deposit = Balance(0)
 
     # Check if we have enough balance.
-    if transfer_all is True:
-        amount = account_balance - fee - existential_deposit
-        if amount < Balance(0):
-            print_error("Not enough balance to transfer")
-            return False
-
     if account_balance < (amount + fee + existential_deposit):
         err_console.print(
             ":cross_mark: [bold red]Not enough balance[/bold red]:\n\n"
@@ -194,7 +187,7 @@ async def transfer_extrinsic(
             )
             console.print(
                 f"Balance:\n"
-                f"  [blue]{account_balance}[/blue] :arrow_right: [green]{new_balance[wallet.coldkeypub.ss58_address]}[/green]"
+                f"  [blue]{account_balance}[/blue] :arrow_right: [green]{new_balance}[/green]"
             )
             return True
 

bittensor_cli/src/bittensor/minigraph.py

@@ -3,7 +3,7 @@ import asyncio
 import numpy as np
 from numpy.typing import NDArray
 
-from bittensor_cli.src.bittensor.chain_data import NeuronInfo
+from bittensor_cli.src.bittensor.chain_data import NeuronInfo, SubnetState
 from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface
 from bittensor_cli.src.bittensor.utils import (
     convert_root_weight_uids_and_vals_to_tensor,
@@ -18,6 +18,7 @@ class MiniGraph:
         netuid: int,
         neurons: list[NeuronInfo],
         subtensor: "SubtensorInterface",
+        subnet_state: "SubnetState",
         block: int,
     ):
         self.neurons = neurons
@@ -62,12 +63,14 @@ class MiniGraph:
         self.validator_trust = self._create_tensor(
             [neuron.validator_trust for neuron in self.neurons], dtype=np.float32
         )
-        self.total_stake = self._create_tensor(
-            [neuron.total_stake.tao for neuron in self.neurons], dtype=np.float32
-        )
-        self.stake = self._create_tensor(
-            [neuron.stake for neuron in self.neurons], dtype=np.float32
+
+        # Fetch stakes from subnet_state until we get updated data in NeuronInfo
+        global_stake_list, local_stake_list, stake_weights_list = self._process_stakes(
+            neurons, subnet_state
         )
+        self.global_stake = self._create_tensor(global_stake_list, dtype=np.float32)
+        self.local_stake = self._create_tensor(local_stake_list, dtype=np.float32)
+        self.stake_weights = self._create_tensor(stake_weights_list, dtype=np.float32)
 
     async def __aenter__(self):
         if not self.weights:
@@ -120,6 +123,41 @@ class MiniGraph:
                 [neuron.bonds for neuron in self.neurons], "bonds"
             )
 
+    def _process_stakes(
+        self,
+        neurons: list[NeuronInfo],
+        subnet_state: SubnetState,
+    ) -> tuple[list[float], list[float], list[float]]:
+        """
+        Processes the global_stake, local_stake, and stake_weights based on the neuron's hotkey.
+
+        Args:
+            neurons (List[NeuronInfo]): List of neurons.
+            subnet_state (SubnetState): The subnet state containing stake information.
+
+        Returns:
+            tuple[list[float], list[float], list[float]]: Lists of global_stake, local_stake, and stake_weights.
+        """
+        global_stake_list = []
+        local_stake_list = []
+        stake_weights_list = []
+        hotkey_to_index = {
+            hotkey: idx for idx, hotkey in enumerate(subnet_state.hotkeys)
+        }
+
+        for neuron in neurons:
+            idx = hotkey_to_index.get(neuron.hotkey)
+            if idx is not None:
+                global_stake_list.append(subnet_state.global_stake[idx].tao)
+                local_stake_list.append(subnet_state.local_stake[idx].tao)
+                stake_weights_list.append(subnet_state.stake_weight[idx])
+            else:
+                global_stake_list.append(0.0)
+                local_stake_list.append(0.0)
+                stake_weights_list.append(0.0)
+
+        return global_stake_list, local_stake_list, stake_weights_list
+
     def _process_weights_or_bonds(self, data, attribute: str) -> NDArray:
         """
         Processes the raw weights or bonds data and converts it into a structured tensor format. This method handles
@@ -177,7 +215,7 @@ class MiniGraph:
         """
 
         async def get_total_subnets():
-            _result = await self.subtensor.substrate.query(
+            _result = await self.subtensor.query(
                 module="SubtensorModule",
                 storage_function="TotalNetworks",
                 params=[],
@@ -186,7 +224,7 @@ class MiniGraph:
             return _result
 
         async def get_subnets():
-            _result = await self.subtensor.substrate.query(
+            _result = await self.subtensor.query(
                 module="SubtensorModule",
                 storage_function="TotalNetworks",
             )