diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 3af83c73..dab31586 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -104,6 +104,7 @@ subnets, mechanisms as subnet_mechanisms, ) +from bittensor_cli.src.commands.axon import axon from bittensor_cli.version import __version__, __version_as_int__ try: @@ -847,6 +848,7 @@ def __init__(self): self.liquidity_app = typer.Typer(epilog=_epilog) self.crowd_app = typer.Typer(epilog=_epilog) self.utils_app = typer.Typer(epilog=_epilog) + self.axon_app = typer.Typer(epilog=_epilog) self.proxy_app = typer.Typer(epilog=_epilog) # config alias @@ -946,6 +948,14 @@ def __init__(self): no_args_is_help=True, ) + # axon app + self.app.add_typer( + self.axon_app, + name="axon", + short_help="Axon serving commands", + no_args_is_help=True, + ) + # proxy app self.app.add_typer( self.proxy_app, @@ -1037,6 +1047,10 @@ def __init__(self): "verify", rich_help_panel=HELP_PANELS["WALLET"]["OPERATIONS"] )(self.wallet_verify) + # axon commands + self.axon_app.command("reset")(self.axon_reset) + self.axon_app.command("set")(self.axon_set) + # stake commands self.stake_app.command( "add", rich_help_panel=HELP_PANELS["STAKE"]["STAKE_MGMT"] @@ -4078,6 +4092,166 @@ def wallet_swap_coldkey( ) ) + def axon_reset( + self, + netuid: int = Options.netuid, + wallet_name: Optional[str] = Options.wallet_name, + wallet_path: Optional[str] = Options.wallet_path, + wallet_hotkey: Optional[str] = Options.wallet_hotkey, + network: Optional[list[str]] = Options.network, + prompt: bool = Options.prompt, + wait_for_inclusion: bool = Options.wait_for_inclusion, + wait_for_finalization: bool = Options.wait_for_finalization, + quiet: bool = Options.quiet, + verbose: bool = Options.verbose, + json_output: bool = Options.json_output, + ): + """ + Reset the axon information for a neuron on the network. + + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, + indicating the neuron is no longer serving. + + USAGE + + The command requires you to specify the netuid where the neuron is registered. + It will reset the axon information for the hotkey associated with the wallet. + + EXAMPLE + + [green]$[/green] btcli axon reset --netuid 1 --wallet-name my_wallet --wallet-hotkey my_hotkey + + [bold]NOTE[/bold]: This command is used to stop serving on a specific subnet. The neuron will + remain registered but will not be reachable by other neurons until a new axon is set. + """ + self.verbosity_handler(quiet, verbose, json_output) + + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.HOTKEY], + validate=WV.WALLET_AND_HOTKEY, + ) + + subtensor = self.initialize_chain(network) + + logger.debug( + "args:\n" + f"netuid: {netuid}\n" + f"wallet: {wallet}\n" + f"prompt: {prompt}\n" + f"wait_for_inclusion: {wait_for_inclusion}\n" + f"wait_for_finalization: {wait_for_finalization}\n" + ) + + return self._run_command( + axon.reset( + wallet=wallet, + subtensor=subtensor, + netuid=netuid, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + json_output=json_output, + ) + ) + + def axon_set( + self, + netuid: int = Options.netuid, + ip: str = typer.Option( + ..., + "--ip", + help="IP address to set for the axon (e.g., '192.168.1.1')", + prompt="Enter the IP address for the axon (e.g., '192.168.1.1' or '2001:db8::1')", + ), + port: int = typer.Option( + ..., + "--port", + help="Port number to set for the axon (0-65535)", + prompt="Enter the port number for the axon (0-65535)", + ), + ip_type: int = typer.Option( + 4, + "--ip-type", + help="IP type (4 for IPv4, 6 for IPv6)", + ), + protocol: int = typer.Option( + 4, + "--protocol", + help="Protocol version", + ), + wallet_name: Optional[str] = Options.wallet_name, + wallet_path: Optional[str] = Options.wallet_path, + wallet_hotkey: Optional[str] = Options.wallet_hotkey, + network: Optional[list[str]] = Options.network, + prompt: bool = Options.prompt, + wait_for_inclusion: bool = Options.wait_for_inclusion, + wait_for_finalization: bool = Options.wait_for_finalization, + quiet: bool = Options.quiet, + verbose: bool = Options.verbose, + json_output: bool = Options.json_output, + ): + """ + Set the axon information for a neuron on the network. + + This command configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + USAGE + + The command requires you to specify the netuid, IP address, and port number. + It will set the axon information for the hotkey associated with the wallet. + + EXAMPLE + + [green]$[/green] btcli axon set --netuid 1 --ip 192.168.1.100 --port 8091 --wallet-name my_wallet --wallet-hotkey my_hotkey + + [bold]NOTE[/bold]: This command is used to advertise your serving endpoint on the network. + Make sure the IP and port are accessible from the internet if you want other neurons to connect. + """ + self.verbosity_handler(quiet, verbose, json_output) + + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.HOTKEY], + validate=WV.WALLET_AND_HOTKEY, + ) + + subtensor = self.initialize_chain(network) + + logger.debug( + "args:\n" + f"netuid: {netuid}\n" + f"ip: {ip}\n" + f"port: {port}\n" + f"ip_type: {ip_type}\n" + f"protocol: {protocol}\n" + f"wallet: {wallet}\n" + f"prompt: {prompt}\n" + f"wait_for_inclusion: {wait_for_inclusion}\n" + f"wait_for_finalization: {wait_for_finalization}\n" + ) + + return self._run_command( + axon.set_axon( + wallet=wallet, + subtensor=subtensor, + netuid=netuid, + ip=ip, + port=port, + ip_type=ip_type, + protocol=protocol, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + json_output=json_output, + ) + ) + # Stake def get_auto_stake( diff --git a/bittensor_cli/src/bittensor/extrinsics/serving.py b/bittensor_cli/src/bittensor/extrinsics/serving.py new file mode 100644 index 00000000..4e88b694 --- /dev/null +++ b/bittensor_cli/src/bittensor/extrinsics/serving.py @@ -0,0 +1,251 @@ +""" +Extrinsics for serving operations (axon management). +""" + +import typing +from typing import Optional + +from bittensor_wallet import Wallet +from rich.prompt import Confirm + +from bittensor_cli.src.bittensor.utils import ( + console, + err_console, + format_error_message, + unlock_key, + print_extrinsic_id, +) +from bittensor_cli.src.bittensor.networking import int_to_ip + +if typing.TYPE_CHECKING: + from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface + + +def ip_to_int(ip_str: str) -> int: + """ + Converts an IP address string to its integer representation. + + Args: + ip_str: IP address string (e.g., "192.168.1.1") + + Returns: + Integer representation of the IP address + """ + import netaddr + + return int(netaddr.IPAddress(ip_str)) + + +async def reset_axon_extrinsic( + subtensor: "SubtensorInterface", + wallet: Wallet, + netuid: int, + prompt: bool = False, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, +) -> tuple[bool, str, Optional[str]]: + """ + Resets the axon information for a neuron on the network. + + This effectively removes the serving endpoint by setting the IP to 0.0.0.0 + and port to 0, indicating the neuron is no longer serving. + + Args: + subtensor: The subtensor interface to use for the extrinsic + wallet: The wallet containing the hotkey to reset the axon for + netuid: The network UID where the neuron is registered + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + + Returns: + Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) + """ + # Unlock the hotkey + if not ( + unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False) + ).success: + return False, unlock_status.message, None + + # Prompt for confirmation if requested + if prompt: + if not Confirm.ask( + f"Do you want to reset the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " + f"on netuid [bold]{netuid}[/bold]?" + ): + return False, "User cancelled the operation", None + + with console.status( + f":satellite: Resetting axon on [white]netuid {netuid}[/white]..." + ): + try: + # Compose the serve_axon call with reset values (IP: 0.0.0.0, port: 1) + # Note: Port must be >= 1 as chain rejects port 0 as invalid + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="serve_axon", + call_params={ + "netuid": netuid, + "version": 0, + "ip": ip_to_int("0.0.0.0"), + "port": 1, + "ip_type": 4, # IPv4 + "protocol": 4, + "placeholder1": 0, + "placeholder2": 0, + }, + ) + + # Sign with hotkey and submit the extrinsic + extrinsic = await subtensor.substrate.create_signed_extrinsic( + call=call, + keypair=wallet.hotkey, + ) + response = await subtensor.substrate.submit_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: + console.print( + ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" + ) + return True, "Not waiting for finalization or inclusion.", None + + success = await response.is_success + if not success: + error_msg = format_error_message(await response.error_message) + err_console.print(f":cross_mark: [red]Failed[/red]: {error_msg}") + return False, error_msg, None + else: + ext_id = await response.get_extrinsic_identifier() + await print_extrinsic_id(response) + console.print( + ":white_heavy_check_mark: [dark_sea_green3]Axon reset successfully[/dark_sea_green3]" + ) + return True, "Axon reset successfully", ext_id + + except Exception as e: + error_message = format_error_message(e) + err_console.print( + f":cross_mark: [red]Failed to reset axon: {error_message}[/red]" + ) + return False, error_message, None + + +async def set_axon_extrinsic( + subtensor: "SubtensorInterface", + wallet: Wallet, + netuid: int, + ip: str, + port: int, + ip_type: int = 4, + protocol: int = 4, + prompt: bool = False, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, +) -> tuple[bool, str, Optional[str]]: + """ + Sets the axon information for a neuron on the network. + + This configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + Args: + subtensor: The subtensor interface to use for the extrinsic + wallet: The wallet containing the hotkey to set the axon for + netuid: The network UID where the neuron is registered + ip: The IP address to set (e.g., "192.168.1.1") + port: The port number to set + ip_type: IP type (4 for IPv4, 6 for IPv6) + protocol: Protocol version (default: 4) + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + + Returns: + Tuple of (success: bool, message: str, extrinsic_id: Optional[str]) + """ + # Validate port + if not (0 <= port <= 65535): + return False, f"Invalid port number: {port}. Must be between 0 and 65535.", None + + # Validate IP address + try: + ip_int = ip_to_int(ip) + except Exception as e: + return False, f"Invalid IP address: {ip}. Error: {str(e)}", None + + # Unlock the hotkey + if not ( + unlock_status := unlock_key(wallet, unlock_type="hot", print_out=False) + ).success: + return False, unlock_status.message, None + + # Prompt for confirmation if requested + if prompt: + if not Confirm.ask( + f"Do you want to set the axon for hotkey [bold]{wallet.hotkey.ss58_address}[/bold] " + f"on netuid [bold]{netuid}[/bold] to [bold]{ip}:{port}[/bold]?" + ): + return False, "User cancelled the operation", None + + with console.status( + f":satellite: Setting axon on [white]netuid {netuid}[/white] to [white]{ip}:{port}[/white]..." + ): + try: + # Compose the serve_axon call + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="serve_axon", + call_params={ + "netuid": netuid, + "version": 0, + "ip": ip_int, + "port": port, + "ip_type": ip_type, + "protocol": protocol, + "placeholder1": 0, + "placeholder2": 0, + }, + ) + + # Sign with hotkey and submit the extrinsic + extrinsic = await subtensor.substrate.create_signed_extrinsic( + call=call, + keypair=wallet.hotkey, + ) + response = await subtensor.substrate.submit_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: + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" + ) + return True, "Not waiting for finalization or inclusion.", None + + success = await response.is_success + if not success: + error_msg = format_error_message(await response.error_message) + err_console.print(f":cross_mark: [red]Failed[/red]: {error_msg}") + return False, error_msg, None + else: + ext_id = await response.get_extrinsic_identifier() + await print_extrinsic_id(response) + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Axon set successfully to {ip}:{port}[/dark_sea_green3]" + ) + return True, f"Axon set successfully to {ip}:{port}", ext_id + + except Exception as e: + error_message = format_error_message(e) + err_console.print( + f":cross_mark: [red]Failed to set axon: {error_message}[/red]" + ) + return False, error_message, None diff --git a/bittensor_cli/src/commands/axon/__init__.py b/bittensor_cli/src/commands/axon/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bittensor_cli/src/commands/axon/axon.py b/bittensor_cli/src/commands/axon/axon.py new file mode 100644 index 00000000..95db9134 --- /dev/null +++ b/bittensor_cli/src/commands/axon/axon.py @@ -0,0 +1,133 @@ +""" +Axon commands for managing neuron serving endpoints. +""" + +import json +from typing import TYPE_CHECKING + +from bittensor_wallet import Wallet + +from bittensor_cli.src.bittensor.utils import ( + console, + err_console, + json_console, +) +from bittensor_cli.src.bittensor.extrinsics.serving import ( + reset_axon_extrinsic, + set_axon_extrinsic, +) + +if TYPE_CHECKING: + from bittensor_cli.src.bittensor.subtensor_interface import SubtensorInterface + + +async def reset( + wallet: Wallet, + subtensor: "SubtensorInterface", + netuid: int, + json_output: bool, + prompt: bool, + wait_for_inclusion: bool, + wait_for_finalization: bool, +): + """ + Reset the axon information for a neuron on the network. + + This command removes the serving endpoint by setting the IP to 0.0.0.0 and port to 1, + indicating the neuron is no longer serving. + + Args: + wallet: The wallet containing the hotkey to reset the axon for + subtensor: The subtensor interface to use for the extrinsic + netuid: The network UID where the neuron is registered + json_output: Whether to output results in JSON format + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + """ + success, message, ext_id = await reset_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if json_output: + json_console.print( + json.dumps( + { + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + } + ) + ) + elif not success: + err_console.print(f"[red]Failed to reset axon: {message}[/red]") + + +async def set_axon( + wallet: Wallet, + subtensor: "SubtensorInterface", + netuid: int, + ip: str, + port: int, + ip_type: int, + protocol: int, + json_output: bool, + prompt: bool, + wait_for_inclusion: bool, + wait_for_finalization: bool, +): + """ + Set the axon information for a neuron on the network. + + This command configures the serving endpoint for a neuron by specifying its IP address + and port, allowing other neurons to connect to it. + + Args: + wallet: The wallet containing the hotkey to set the axon for + subtensor: The subtensor interface to use for the extrinsic + netuid: The network UID where the neuron is registered + ip: IP address to set for the axon + port: Port number to set for the axon + ip_type: IP type (4 for IPv4, 6 for IPv6) + protocol: Protocol version + json_output: Whether to output results in JSON format + prompt: Whether to prompt for confirmation before submitting + wait_for_inclusion: Whether to wait for the extrinsic to be included in a block + wait_for_finalization: Whether to wait for the extrinsic to be finalized + """ + success, message, ext_id = await set_axon_extrinsic( + subtensor=subtensor, + wallet=wallet, + netuid=netuid, + ip=ip, + port=port, + ip_type=ip_type, + protocol=protocol, + prompt=prompt, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + ) + + if json_output: + json_console.print( + json.dumps( + { + "success": success, + "message": message, + "extrinsic_identifier": ext_id, + "netuid": netuid, + "hotkey": wallet.hotkey.ss58_address, + "ip": ip, + "port": port, + } + ) + ) + elif not success: + err_console.print(f"[red]Failed to set axon: {message}[/red]") diff --git a/tests/e2e_tests/test_axon.py b/tests/e2e_tests/test_axon.py new file mode 100644 index 00000000..20763b91 --- /dev/null +++ b/tests/e2e_tests/test_axon.py @@ -0,0 +1,521 @@ +""" +End-to-end tests for axon commands. + +Verify commands: +* btcli axon reset +* btcli axon set +""" + +import json + +import pytest +import re + +from tests.e2e_tests.utils import execute_turn_off_hyperparam_freeze_window + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_reset_and_set(local_chain, wallet_setup): + """ + Test axon reset and set commands end-to-end. + + This test: + 1. Creates a subnet + 2. Registers a neuron + 3. Sets the axon information + 4. Verifies the axon is set correctly + 5. Resets the axon + 6. Verifies the axon is reset (0.0.0.0:1 - not serving) + """ + wallet_path_alice = "//Alice" + netuid = 2 + + # Create wallet for Alice + keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( + wallet_path_alice + ) + + execute_turn_off_hyperparam_freeze_window(local_chain, wallet_alice) + + # Register a subnet with sudo as Alice + result = exec_command_alice( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--subnet-name", + "Test Axon Subnet", + "--repo", + "https://github.com/test/axon-subnet", + "--contact", + "test@opentensor.dev", + "--url", + "https://testaxon.com", + "--discord", + "test#1234", + "--description", + "Test subnet for axon e2e testing", + "--logo-url", + "https://testaxon.com/logo.png", + "--additional-info", + "Axon test subnet", + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" + + # Register neuron on the subnet + result = exec_command_alice( + command="subnets", + sub_command="register", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + "--json-output", + ], + ) + json_result = json.loads(result.stdout) + assert json_result["success"] is True, (result.stdout, result.stdout) + + # Set serving rate limit to 0 to allow immediate axon updates + result = exec_command_alice( + command="sudo", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--netuid", + str(netuid), + "--param", + "serving_rate_limit", + "--value", + "0", + "--no-prompt", + "--json-output", + ], + ) + result_json = json.loads(result.stdout) + assert result_json["success"] is True, ( + f"Setting serving_rate_limit failed: {result.stdout}" + ) + + # Set axon information + test_ip = "192.168.1.100" + test_port = 8091 + + result = exec_command_alice( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--netuid", + str(netuid), + "--ip", + test_ip, + "--port", + str(test_port), + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon set failed: {result.stdout}" + assert ( + "successfully" in result.stdout.lower() or "success" in result.stdout.lower() + ), f"Success message not found in output: {result.stdout}" + + # Verify axon is set by checking wallet overview + result = exec_command_alice( + command="wallet", + sub_command="overview", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--netuid", + str(netuid), + ], + ) + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" + + # Check that axon column shows an IP (not "none") + # The overview should show the axon info in the AXON column + lines = result.stdout.split("\n") + axon_found = False + for line in lines: + # Look for a line with the neuron info that has an IP address in the AXON column + if wallet_alice.hotkey_str[:8] in line and "none" not in line.lower(): + # Check if there's an IP-like pattern in the line + if re.search(r"\d+\.\d+\.\d+\.\d+:\d+", line): + axon_found = True + break + + assert axon_found, f"Axon not set correctly in overview: {result.stdout}" + + # Reset axon + result = exec_command_alice( + command="axon", + sub_command="reset", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon reset failed: {result.stdout}" + assert "Axon reset successfully" in result.stdout, ( + f"Success message not found in output: {result.stdout}" + ) + + # Verify axon is reset by checking wallet overview + result = exec_command_alice( + command="wallet", + sub_command="overview", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--netuid", + str(netuid), + ], + ) + assert result.exit_code == 0, f"Wallet overview failed: {result.stdout}" + + # Check that axon column shows the reset ip:port value (0.0.0.0:1) after reset + lines = result.stdout.split("\n") + axon_reset = False + for line in lines: + if wallet_alice.hotkey_str[:8] in line and "0.0.0.0:1" in line.lower(): + axon_reset = True + break + + assert axon_reset, f"Axon not reset correctly in overview: {result.stdout}" + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_set_with_ipv6(local_chain, wallet_setup): + """ + Test setting axon with IPv6 address. + """ + wallet_path_bob = "//Bob" + wallet_path_alice = "//Alice" + netuid = 2 + + # Create wallet for Bob + keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup( + wallet_path_bob + ) + _, wallet_alice, *_ = wallet_setup(wallet_path_alice) + + execute_turn_off_hyperparam_freeze_window(local_chain, wallet_alice) + + # Register a subnet with sudo as Bob + result = exec_command_bob( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--subnet-name", + "Test IPv6 Subnet", + "--repo", + "https://github.com/test/ipv6-subnet", + "--contact", + "ipv6@opentensor.dev", + "--url", + "https://testipv6.com", + "--discord", + "ipv6#5678", + "--description", + "Test subnet for IPv6 axon testing", + "--logo-url", + "https://testipv6.com/logo.png", + "--additional-info", + "IPv6 test subnet", + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Subnet creation failed: {result.stdout}" + + # Register neuron on the subnet + result = exec_command_bob( + command="subnets", + sub_command="register", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + "--json-output", + ], + ) + json_result = json.loads(result.stdout) + assert json_result["success"] is True, (result.stdout, result.stdout) + + # Set serving rate limit to 0 to allow immediate axon updates + result = exec_command_bob( + command="sudo", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--netuid", + str(netuid), + "--param", + "serving_rate_limit", + "--value", + "0", + "--no-prompt", + "--json-output", + ], + ) + result_json = json.loads(result.stdout) + assert result_json["success"] is True, ( + f"Setting serving_rate_limit failed: {result.stdout}" + ) + + # Set axon with IPv6 address + test_ipv6 = "2001:db8::1" + test_port = 8092 + + result = exec_command_bob( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_bob, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_bob.name, + "--wallet-hotkey", + wallet_bob.hotkey_str, + "--netuid", + str(netuid), + "--ip", + test_ipv6, + "--port", + str(test_port), + "--ip-type", + "6", # IPv6 + "--no-prompt", + ], + ) + + assert result.exit_code == 0, f"Axon set with IPv6 failed: {result.stdout}" + assert f"Axon set successfully to {test_ipv6}:{test_port}" in result.stdout, ( + f"Success message not found in output: {result.stdout}" + ) + + +@pytest.mark.parametrize("local_chain", [None], indirect=True) +def test_axon_set_invalid_inputs(local_chain, wallet_setup): + """ + Test axon set with invalid inputs to ensure proper error handling. + """ + wallet_path_charlie = "//Charlie" + netuid = 1 + + # Create wallet for Charlie + keypair_charlie, wallet_charlie, wallet_path_charlie, exec_command_charlie = ( + wallet_setup(wallet_path_charlie) + ) + + # Register a subnet + result = exec_command_charlie( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--subnet-name", + "Test Invalid Inputs Subnet", + "--repo", + "https://github.com/test/invalid-subnet", + "--contact", + "invalid@opentensor.dev", + "--url", + "https://testinvalid.com", + "--discord", + "invalid#9999", + "--description", + "Test subnet for invalid inputs testing", + "--logo-url", + "https://testinvalid.com/logo.png", + "--additional-info", + "Invalid inputs test subnet", + "--no-prompt", + ], + ) + assert result.exit_code == 0 + + # Register neuron + result = exec_command_charlie( + command="subnets", + sub_command="register", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--no-prompt", + ], + ) + assert result.exit_code == 0 + + # Set serving rate limit to 0 to allow immediate axon updates + result = exec_command_charlie( + command="sudo", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--param", + "serving_rate_limit", + "--value", + "0", + "--no-prompt", + ], + ) + assert result.exit_code == 0, f"Setting serving_rate_limit failed: {result.stdout}" + + invalid_port = "70000" + # Test with invalid port (too high) + result = exec_command_charlie( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--ip", + "192.168.1.1", + "--port", + invalid_port, # Invalid port + "--no-prompt", + ], + ) + + # Should fail with invalid port + assert ( + f"Failed to set axon: Invalid port number: {invalid_port}. Must be between 0 and 65535." + in result.stderr + ), f"Expected error for invalid port, got: {result.stdout}\n{result.stderr}" + + invalid_ip_address = "invalid.ip.address" + # Test with invalid IP + result = exec_command_charlie( + command="axon", + sub_command="set", + extra_args=[ + "--wallet-path", + wallet_path_charlie, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_charlie.name, + "--wallet-hotkey", + wallet_charlie.hotkey_str, + "--netuid", + str(netuid), + "--ip", + invalid_ip_address, + "--port", + "8091", + "--no-prompt", + ], + ) + + assert ( + "Failed to set axon: Invalid IP address: invalid.ip.address. " + "Error: failed to detect a valid IP address from 'invalid.ip.address'" + ) in result.stderr, f"Expected error for invalid IP, got: {result.stderr}" diff --git a/tests/e2e_tests/utils.py b/tests/e2e_tests/utils.py index effd4cef..d4658954 100644 --- a/tests/e2e_tests/utils.py +++ b/tests/e2e_tests/utils.py @@ -1,3 +1,4 @@ +import asyncio import importlib import inspect import os @@ -5,6 +6,7 @@ import shutil import subprocess import sys +import time from typing import TYPE_CHECKING, Optional, Protocol from bittensor_wallet import Keypair, Wallet @@ -437,3 +439,15 @@ async def turn_off_hyperparam_freeze_window( ) return await response.is_success, await response.error_message + + +def execute_turn_off_hyperparam_freeze_window( + local_chain: "AsyncSubstrateInterface", wallet: Wallet +): + try: + asyncio.run(turn_off_hyperparam_freeze_window(local_chain, wallet)) + time.sleep(3) + except ValueError: + print( + "Skipping turning off hyperparams freeze window. This indicates the call does not exist on the chain you are testing." + ) diff --git a/tests/unit_tests/test_axon_commands.py b/tests/unit_tests/test_axon_commands.py new file mode 100644 index 00000000..4ed73f9e --- /dev/null +++ b/tests/unit_tests/test_axon_commands.py @@ -0,0 +1,551 @@ +""" +Unit tests for axon commands (reset and set). +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, Mock, patch +from bittensor_wallet import Wallet + +from bittensor_cli.src.bittensor.extrinsics.serving import ( + reset_axon_extrinsic, + set_axon_extrinsic, + ip_to_int, +) + + +class TestIpToInt: + """Tests for IP address to integer conversion.""" + + def test_ipv4_conversion(self): + """Test IPv4 address conversion.""" + assert ip_to_int("0.0.0.0") == 0 + assert ip_to_int("127.0.0.1") == 2130706433 + assert ip_to_int("192.168.1.1") == 3232235777 + assert ip_to_int("255.255.255.255") == 4294967295 + + def test_ipv6_conversion(self): + """Test IPv6 address conversion.""" + # IPv6 loopback + result = ip_to_int("::1") + assert result == 1 + + # IPv6 address + result = ip_to_int("2001:db8::1") + assert result > 0 + + def test_invalid_ip_raises_error(self): + """Test that invalid IP addresses raise errors.""" + with pytest.raises(Exception): + ip_to_int("invalid.ip.address") + + with pytest.raises(Exception): + ip_to_int("256.256.256.256") + + +class TestResetAxonExtrinsic: + """Tests for reset_axon_extrinsic function.""" + + @pytest.mark.asyncio + async def test_reset_axon_success(self): + """Test successful axon reset.""" + # Setup mocks + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) + mock_response = MagicMock() + + # is_success is a property that returns a coroutine + async def mock_is_success(): + return True + + mock_response.is_success = mock_is_success() + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with ( + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), + ): + mock_unlock.return_value = MagicMock(success=True) + + # Execute + success, message, ext_id = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + ) + + # Verify + assert success is True + assert "successfully" in message.lower() + assert ext_id == "0x123" + + # Verify compose_call was called with correct parameters + mock_subtensor.substrate.compose_call.assert_called_once() + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_module"] == "SubtensorModule" + assert call_args[1]["call_function"] == "serve_axon" + assert call_args[1]["call_params"]["netuid"] == 1 + assert call_args[1]["call_params"]["ip"] == 0 # 0.0.0.0 as int + assert call_args[1]["call_params"]["port"] == 1 # Port 1, not 0 + assert call_args[1]["call_params"]["ip_type"] == 4 + + @pytest.mark.asyncio + async def test_reset_axon_unlock_failure(self): + """Test axon reset when hotkey unlock fails.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock( + success=False, message="Failed to unlock hotkey" + ) + + success, message, ext_id = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + ) + + assert success is False + assert "unlock" in message.lower() + assert ext_id is None + + @pytest.mark.asyncio + async def test_reset_axon_user_cancellation(self): + """Test axon reset when user cancels prompt.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with ( + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.Confirm" + ) as mock_confirm, + ): + mock_unlock.return_value = MagicMock(success=True) + mock_confirm.ask.return_value = False + + success, message, ext_id = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=True, + ) + + assert success is False + assert "cancelled" in message.lower() + assert ext_id is None + + @pytest.mark.asyncio + async def test_reset_axon_extrinsic_failure(self): + """Test axon reset when extrinsic submission fails.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) + mock_response = MagicMock() + + async def mock_is_success(): + return False + + mock_response.is_success = mock_is_success() + mock_response.error_message = AsyncMock(return_value="Network error") + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock(success=True) + + success, message, ext_id = await reset_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + prompt=False, + ) + + assert success is False + assert len(message) > 0 + assert ext_id is None + + +class TestSetAxonExtrinsic: + """Tests for set_axon_extrinsic function.""" + + @pytest.mark.asyncio + async def test_set_axon_success(self): + """Test successful axon set.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) + mock_response = MagicMock() + + async def mock_is_success(): + return True + + mock_response.is_success = mock_is_success() + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with ( + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), + ): + mock_unlock.return_value = MagicMock(success=True) + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + ip_type=4, + protocol=4, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + ) + + assert success is True + assert "successfully" in message.lower() + assert ext_id == "0x123" + assert "192.168.1.100:8091" in message + + # Verify compose_call was called with correct parameters + mock_subtensor.substrate.compose_call.assert_called_once() + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_module"] == "SubtensorModule" + assert call_args[1]["call_function"] == "serve_axon" + assert call_args[1]["call_params"]["netuid"] == 1 + assert call_args[1]["call_params"]["port"] == 8091 + assert call_args[1]["call_params"]["ip_type"] == 4 + assert call_args[1]["call_params"]["protocol"] == 4 + + @pytest.mark.asyncio + async def test_set_axon_invalid_port(self): + """Test axon set with invalid port number.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + # Test port too high + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=70000, + prompt=False, + ) + + assert success is False + assert "Invalid port" in message + assert ext_id is None + + # Test negative port + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=-1, + prompt=False, + ) + + assert success is False + assert "Invalid port" in message + assert ext_id is None + + @pytest.mark.asyncio + async def test_set_axon_invalid_ip(self): + """Test axon set with invalid IP address.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="invalid.ip.address", + port=8091, + prompt=False, + ) + + assert success is False + assert "Invalid IP" in message + assert ext_id is None + + @pytest.mark.asyncio + async def test_set_axon_unlock_failure(self): + """Test axon set when hotkey unlock fails.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock( + success=False, message="Failed to unlock hotkey" + ) + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=False, + ) + + assert success is False + assert "unlock" in message.lower() + assert "Failed to unlock hotkey" in message + + @pytest.mark.asyncio + async def test_set_axon_user_cancellation(self): + """Test axon set when user cancels prompt.""" + mock_subtensor = MagicMock() + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with ( + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.Confirm" + ) as mock_confirm, + ): + mock_unlock.return_value = MagicMock(success=True) + mock_confirm.ask.return_value = False + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=True, + ) + + assert success is False + assert "cancelled" in message.lower() + + @pytest.mark.asyncio + async def test_set_axon_with_ipv6(self): + """Test axon set with IPv6 address.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock(return_value="mock_call") + mock_subtensor.substrate.create_signed_extrinsic = AsyncMock( + return_value="mock_extrinsic" + ) + mock_response = MagicMock() + + async def mock_is_success(): + return True + + mock_response.is_success = mock_is_success() + mock_response.get_extrinsic_identifier = AsyncMock(return_value="0x123") + mock_subtensor.substrate.submit_extrinsic = AsyncMock( + return_value=mock_response + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with ( + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock, + patch( + "bittensor_cli.src.bittensor.extrinsics.serving.print_extrinsic_id", + new_callable=AsyncMock, + ), + ): + mock_unlock.return_value = MagicMock(success=True) + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="2001:db8::1", + port=8091, + ip_type=6, + protocol=4, + prompt=False, + ) + + assert success is True + assert "successfully" in message.lower() + assert ext_id == "0x123" + # Verify ip_type was set to 6 + call_args = mock_subtensor.substrate.compose_call.call_args + assert call_args[1]["call_params"]["ip_type"] == 6 + + @pytest.mark.asyncio + async def test_set_axon_exception_handling(self): + """Test axon set handles exceptions gracefully.""" + mock_subtensor = MagicMock() + mock_subtensor.substrate.compose_call = AsyncMock( + side_effect=Exception("Unexpected error") + ) + + mock_wallet = MagicMock(spec=Wallet) + mock_wallet.hotkey.ss58_address = ( + "5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY" + ) + + with patch( + "bittensor_cli.src.bittensor.extrinsics.serving.unlock_key" + ) as mock_unlock: + mock_unlock.return_value = MagicMock(success=True) + + success, message, ext_id = await set_axon_extrinsic( + subtensor=mock_subtensor, + wallet=mock_wallet, + netuid=1, + ip="192.168.1.100", + port=8091, + prompt=False, + ) + + assert success is False + assert len(message) > 0 + assert ext_id is None + + +class TestAxonCLICommands: + """Tests for CLI command handlers.""" + + @patch("bittensor_cli.cli.axon") + def test_axon_reset_command_handler(self, mock_axon): + """Test axon reset CLI command handler.""" + from bittensor_cli.cli import CLIManager + + cli_manager = CLIManager() + mock_axon.reset = AsyncMock(return_value=None) + + with ( + patch.object(cli_manager, "verbosity_handler"), + patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, + patch.object(cli_manager, "initialize_chain") as mock_init_chain, + patch.object(cli_manager, "_run_command") as mock_run_command, + ): + mock_wallet = Mock() + mock_wallet_ask.return_value = mock_wallet + mock_subtensor = Mock() + mock_init_chain.return_value = mock_subtensor + + cli_manager.axon_reset( + netuid=1, + wallet_name="test_wallet", + wallet_path="/tmp/test", + wallet_hotkey="test_hotkey", + network=None, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + quiet=False, + verbose=False, + json_output=False, + ) + + # Verify wallet_ask was called correctly + mock_wallet_ask.assert_called_once() + + # Verify _run_command was called + mock_run_command.assert_called_once() + + @patch("bittensor_cli.cli.axon") + def test_axon_set_command_handler(self, mock_axon): + """Test axon set CLI command handler.""" + from bittensor_cli.cli import CLIManager + + cli_manager = CLIManager() + mock_axon.set_axon = AsyncMock(return_value=None) + + with ( + patch.object(cli_manager, "verbosity_handler"), + patch.object(cli_manager, "wallet_ask") as mock_wallet_ask, + patch.object(cli_manager, "initialize_chain") as mock_init_chain, + patch.object(cli_manager, "_run_command") as mock_run_command, + ): + mock_wallet = Mock() + mock_wallet_ask.return_value = mock_wallet + mock_subtensor = Mock() + mock_init_chain.return_value = mock_subtensor + + cli_manager.axon_set( + netuid=1, + ip="192.168.1.100", + port=8091, + ip_type=4, + protocol=4, + wallet_name="test_wallet", + wallet_path="/tmp/test", + wallet_hotkey="test_hotkey", + network=None, + prompt=False, + wait_for_inclusion=True, + wait_for_finalization=False, + quiet=False, + verbose=False, + json_output=False, + ) + + # Verify wallet_ask was called correctly + mock_wallet_ask.assert_called_once() + + # Verify _run_command was called + mock_run_command.assert_called_once()