Skip to content

Async unittests for bittensor/core/extrinsics/async_registration.py #2445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
basfroman merged 6 commits into from
Nov 20, 2024
Merged

Async unittests for bittensor/core/extrinsics/async_registration.py #2445

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
350fe7a
make async register_extrinsic consistent with sync version
Nov 20, 2024
db1e40c
add test for `async_registration._do_pow_register`
Nov 20, 2024
6c7f10b
add test for `async_registration.register_extrinsic`
Nov 20, 2024
d299571
remove unused extrinsic
Nov 20, 2024
f4de7de
remove unused import
Nov 20, 2024
023b359
review comments + some comments/fix
Nov 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
284 changes: 94 additions & 190 deletions bittensor/core/extrinsics/async_registration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,20 +7,19 @@
"""

import asyncio
import time
from typing import Optional, Union, TYPE_CHECKING

from bittensor_wallet import Wallet

from bittensor.utils import format_error_message, unlock_key
from bittensor.utils import format_error_message
from bittensor.utils.btlogging import logging
from bittensor.utils.registration import log_no_torch_error, create_pow_async

# For annotation and lazy import purposes
if TYPE_CHECKING:
import torch
from bittensor.core.async_subtensor import AsyncSubtensor
from bittensor.utils.registration import POWSolution
from bittensor.utils.registration.pow import POWSolution
else:
from bittensor.utils.registration.pow import LazyLoadedTorch

Expand All @@ -35,6 +34,65 @@ class MaxAttemptsException(Exception):
"""Raised when the POW Solver has reached the max number of attempts."""


async def _do_pow_register(
subtensor: "AsyncSubtensor",
netuid: int,
wallet: "Wallet",
pow_result: "POWSolution",
wait_for_inclusion: bool = False,
wait_for_finalization: bool = True,
) -> tuple[bool, Optional[str]]:
"""Sends a (POW) register extrinsic to the chain.

Args:
subtensor (bittensor.core.async_subtensor.AsyncSubtensor): The subtensor to send the extrinsic to.
netuid (int): The subnet to register on.
wallet (bittensor.wallet): The wallet to register.
pow_result (POWSolution): The PoW result to register.
wait_for_inclusion (bool): If ``True``, waits for the extrinsic to be included in a block. Default to `False`.
wait_for_finalization (bool): If ``True``, waits for the extrinsic to be finalized. Default to `True`.

Returns:
success (bool): ``True`` if the extrinsic was included in a block.
error (Optional[str]): ``None`` on success or not waiting for inclusion/finalization, otherwise the error message.
"""
# create extrinsic call
call = await subtensor.substrate.compose_call(
call_module="SubtensorModule",
call_function="register",
call_params={
"netuid": netuid,
"block_number": pow_result.block_number,
"nonce": pow_result.nonce,
"work": [int(byte_) for byte_ in pow_result.seal],
"hotkey": wallet.hotkey.ss58_address,
"coldkey": wallet.coldkeypub.ss58_address,
},
)
extrinsic = await subtensor.substrate.create_signed_extrinsic(
call=call, keypair=wallet.hotkey
)
response = await subtensor.substrate.submit_extrinsic(
extrinsic=extrinsic,
wait_for_inclusion=wait_for_inclusion,
wait_for_finalization=wait_for_finalization,
)

# We only wait here if we expect finalization.
if not wait_for_finalization and not wait_for_inclusion:
return True, None

# process if registration successful, try again if pow is still valid
await response.process_events()
if not await response.is_success:
return False, format_error_message(
error_message=await response.error_message, substrate=subtensor.substrate
)
# Successful registration
else:
return True, None


async def register_extrinsic(
subtensor: "AsyncSubtensor",
wallet: "Wallet",
Expand Down Expand Up @@ -85,19 +143,24 @@ async def register_extrinsic(
hotkey_ss58=wallet.hotkey.ss58_address,
netuid=netuid,
)

if not neuron.is_null:
logging.debug(
f"Wallet <green>{wallet}</green> is already registered on subnet <blue>{neuron.netuid}</blue> with uid<blue>{neuron.uid}</blue>."
)
return True

logging.debug(
f"Registration hotkey: <blue>{wallet.hotkey.ss58_address}</blue>, <green>Public</green> coldkey: <blue>{wallet.coldkey.ss58_address}</blue> in the network: <blue>{subtensor.network}</blue>."
)

if not torch:
log_no_torch_error()
return False

# Attempt rolling registration.
attempts = 1
pow_result: Optional["POWSolution"]

while True:
logging.info(
f":satellite: <magenta>Registering...</magenta> <blue>({attempts}/{max_allowed_attempts})</blue>"
Expand All @@ -107,10 +170,10 @@ async def register_extrinsic(
if not torch.cuda.is_available():
return False
pow_result = await create_pow_async(
subtensor,
wallet,
netuid,
output_in_place,
subtensor=subtensor,
wallet=wallet,
netuid=netuid,
output_in_place=output_in_place,
cuda=cuda,
dev_id=dev_id,
tpb=tpb,
Expand All @@ -120,10 +183,10 @@ async def register_extrinsic(
)
else:
pow_result = await create_pow_async(
subtensor,
wallet,
netuid,
output_in_place,
subtensor=subtensor,
wallet=wallet,
netuid=netuid,
output_in_place=output_in_place,
cuda=cuda,
num_processes=num_processes,
update_interval=update_interval,
Expand All @@ -147,49 +210,28 @@ async def register_extrinsic(
logging.info(":satellite: <magenta>Submitting POW...</magenta>")
# check if pow result is still valid
while not await pow_result.is_stale_async(subtensor=subtensor):
call = await subtensor.substrate.compose_call(
call_module="SubtensorModule",
call_function="register",
call_params={
"netuid": netuid,
"block_number": pow_result.block_number,
"nonce": pow_result.nonce,
"work": [int(byte_) for byte_ in pow_result.seal],
"hotkey": wallet.hotkey.ss58_address,
"coldkey": wallet.coldkeypub.ss58_address,
},
)
extrinsic = await subtensor.substrate.create_signed_extrinsic(
call=call, keypair=wallet.hotkey
)
response = await subtensor.substrate.submit_extrinsic(
extrinsic,
result: tuple[bool, Optional[str]] = await _do_pow_register(
subtensor=subtensor,
netuid=netuid,
wallet=wallet,
pow_result=pow_result,
wait_for_inclusion=wait_for_inclusion,
wait_for_finalization=wait_for_finalization,
)
if not wait_for_finalization and not wait_for_inclusion:
success, err_msg = True, ""
else:
await response.process_events()
success = await response.is_success
if not success:
success, err_msg = (
False,
format_error_message(
await response.error_message,
substrate=subtensor.substrate,
),

success, err_msg = result
if not success:
# Look error here
# https://github.com/opentensor/subtensor/blob/development/pallets/subtensor/src/errors.rs

if "HotKeyAlreadyRegisteredInSubNet" in err_msg:
logging.info(
f":white_heavy_check_mark: <green>Already Registered on subnet:</green> <blue>{netuid}</blue>."
)
# Look error here
# https://github.com/opentensor/subtensor/blob/development/pallets/subtensor/src/errors.rs
return True

if "HotKeyAlreadyRegisteredInSubNet" in err_msg:
logging.info(
f":white_heavy_check_mark: <green>Already Registered on subnet:</green> <blue>{netuid}</blue>."
)
return True
logging.error(f":cross_mark: <red>Failed</red>: {err_msg}")
await asyncio.sleep(0.5)
logging.error(f":cross_mark: <red>Failed</red>: {err_msg}")
await asyncio.sleep(0.5)

# Successful registration, final check for neuron and pubkey
if success:
Expand All @@ -212,154 +254,16 @@ async def register_extrinsic(
# Exited loop because pow is no longer valid.
logging.error("<red>POW is stale.</red>")
# Try again.
continue
# continue

if attempts < max_allowed_attempts:
# Failed registration, retry pow
attempts += 1
logging.error(
f":satellite: <magenta>Failed registration, retrying pow ...</magenta> <blue>({attempts}/{max_allowed_attempts})</blue>"
)
continue
else:
# Failed to register after max attempts.
logging.error("<red>No more attempts.</red>")
return False


async def run_faucet_extrinsic(
subtensor: "AsyncSubtensor",
wallet: "Wallet",
wait_for_inclusion: bool = False,
wait_for_finalization: bool = True,
max_allowed_attempts: int = 3,
output_in_place: bool = True,
cuda: bool = False,
dev_id: int = 0,
tpb: int = 256,
num_processes: Optional[int] = None,
update_interval: Optional[int] = None,
log_verbose: bool = False,
max_successes: int = 3,
) -> tuple[bool, str]:
"""Runs a continual POW to get a faucet of TAO on the test net.

Args:
subtensor: The subtensor interface object used to run the extrinsic
wallet: Bittensor wallet object.
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.
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.
max_allowed_attempts: Maximum number of attempts to register the wallet.
output_in_place: Whether to output logging data as the process runs.
cuda: If `True`, the wallet should be registered using CUDA device(s).
dev_id: The CUDA device id to use
tpb: The number of threads per block (CUDA).
num_processes: The number of processes to use to register.
update_interval: The number of nonces to solve between updates.
log_verbose: If `True`, the registration process will log more information.
max_successes: The maximum number of successful faucet runs for the wallet.

Returns:
`True` if extrinsic was finalized or included in the block. If we did not wait for finalization/inclusion, the response is also `True`
"""

if not torch:
log_no_torch_error()
return False, "Requires torch"

# Unlock coldkey
if not (unlock := unlock_key(wallet)).success:
return False, unlock.message

# Get previous balance.
old_balance = await subtensor.get_balance(wallet.coldkeypub.ss58_address)

# Attempt rolling registration.
attempts = 1
successes = 1
while True:
try:
pow_result = None
while pow_result is None or await pow_result.is_stale_async(
subtensor=subtensor
):
# Solve latest POW.
if cuda:
if not torch.cuda.is_available():
return False, "CUDA is not available."
pow_result = await create_pow_async(
subtensor,
wallet,
-1,
output_in_place,
cuda=cuda,
dev_id=dev_id,
tpb=tpb,
num_processes=num_processes,
update_interval=update_interval,
log_verbose=log_verbose,
)
else:
pow_result = await create_pow_async(
subtensor,
wallet,
-1,
output_in_place,
cuda=cuda,
num_processes=num_processes,
update_interval=update_interval,
log_verbose=log_verbose,
)
call = await subtensor.substrate.compose_call(
call_module="SubtensorModule",
call_function="faucet",
call_params={
"block_number": pow_result.block_number,
"nonce": pow_result.nonce,
"work": [int(byte_) for byte_ in pow_result.seal],
},
)
extrinsic = await subtensor.substrate.create_signed_extrinsic(
call=call, keypair=wallet.coldkey
)
response = await subtensor.substrate.submit_extrinsic(
extrinsic,
wait_for_inclusion=wait_for_inclusion,
wait_for_finalization=wait_for_finalization,
)

# process if registration successful, try again if pow is still valid
await response.process_events()
if not await response.is_success:
logging.error(
f":cross_mark: <red>Failed</red>: {format_error_message(error_message=await response.error_message, substrate=subtensor.substrate)}"
)
if attempts == max_allowed_attempts:
raise MaxAttemptsException
attempts += 1
# Wait a bit before trying again
time.sleep(1)

# Successful registration
else:
new_balance = await subtensor.get_balance(
wallet.coldkeypub.ss58_address
)
logging.info(
f"Balance: <blue>{old_balance[wallet.coldkeypub.ss58_address]}</blue> :arrow_right: <green>{new_balance[wallet.coldkeypub.ss58_address]}</green>"
)
old_balance = new_balance

if successes == max_successes:
raise MaxSuccessException

attempts = 1 # Reset attempts on success
successes += 1

except KeyboardInterrupt:
return True, "Done"

except MaxSuccessException:
return True, f"Max successes reached: {3}"

except MaxAttemptsException:
return False, f"Max attempts reached: {max_allowed_attempts}"
Loading