Merge pull request #481 from opentensor/feat/thewhaleking/multiple-netuid-staking

Allows for staking to multiple netuids in one btcli command

Author: thewhaleking

bittensor_cli/cli.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import asyncio
 import curses
+import copy
 import importlib
 import json
 import os.path
@@ -10,7 +11,7 @@
 import traceback
 import warnings
 from pathlib import Path
-from typing import Coroutine, Optional
+from typing import Coroutine, Optional, Union
 from dataclasses import fields
 
 import rich
@@ -89,6 +90,23 @@ class Options:
     Re-usable typer args
     """
 
+    @classmethod
+    def edit_help(cls, option_name: str, help_text: str):
+        """
+        Edits the `help` attribute of a copied given Typer option in this class, returning
+        the modified Typer option.
+
+        Args:
+            option_name: the name of the option (e.g. "wallet_name")
+            help_text: New help text to be used (e.g. "Wallet's name")
+
+        Returns:
+            Modified Typer Option with new help text.
+        """
+        copied_attr = copy.copy(getattr(cls, option_name))
+        setattr(copied_attr, "help", help_text)
+        return copied_attr
+
     wallet_name = typer.Option(
         None,
         "--wallet-name",
@@ -3202,7 +3220,11 @@ def stake_add(
             help="When set, this command stakes to all hotkeys associated with the wallet. Do not use if specifying "
             "hotkeys in `--include-hotkeys`.",
         ),
-        netuid: Optional[int] = Options.netuid_not_req,
+        netuids: Optional[str] = Options.edit_help(
+            "netuids",
+            "Netuid(s) to for which to add stake. Specify multiple netuids by separating with a comma, e.g."
+            "`btcli st add -n 1,2,3",
+        ),
         all_netuids: bool = Options.all_netuids,
         wallet_name: str = Options.wallet_name,
         wallet_path: str = Options.wallet_path,
@@ -3242,29 +3264,52 @@ def stake_add(
         6. Stake all balance to a subnet:
             [green]$[/green] btcli stake add --all --netuid 3
 
+        7. Stake the same amount to multiple subnets:
+            [green]$[/green] btcli stake add --amount 100 --netuids 4,5,6
+
         [bold]Safe Staking Parameters:[/bold]
         • [blue]--safe[/blue]: Enables rate tolerance checks
         • [blue]--tolerance[/blue]: Maximum % rate change allowed (0.05 = 5%)
         • [blue]--partial[/blue]: Complete partial stake if rates exceed tolerance
 
         """
+        netuids = netuids or []
         self.verbosity_handler(quiet, verbose, json_output)
         safe_staking = self.ask_safe_staking(safe_staking)
         if safe_staking:
             rate_tolerance = self.ask_rate_tolerance(rate_tolerance)
             allow_partial_stake = self.ask_partial_stake(allow_partial_stake)
             console.print("\n")
-        netuid = get_optional_netuid(netuid, all_netuids)
+
+        if netuids:
+            netuids = parse_to_list(
+                netuids, int, "Netuids must be ints separated by commas", False
+            )
+        else:
+            netuid_ = get_optional_netuid(None, all_netuids)
+            netuids = [netuid_] if netuid_ else None
+        if netuids:
+            for netuid_ in netuids:
+                # ensure no negative netuids make it into our list
+                validate_netuid(netuid_)
 
         if stake_all and amount:
             print_error(
                 "Cannot specify an amount and 'stake-all'. Choose one or the other."
             )
-            raise typer.Exit()
+            return
 
         if stake_all and not amount:
             if not Confirm.ask("Stake all the available TAO tokens?", default=False):
-                raise typer.Exit()
+                return
+
+        if (
+            stake_all
+            and (isinstance(netuids, list) and len(netuids) > 1)
+            or (netuids is None)
+        ):
+            print_error("Cannot stake all to multiple subnets.")
+            return
 
         if all_hotkeys and include_hotkeys:
             print_error(
@@ -3285,9 +3330,10 @@ def stake_add(
                     "Enter the [blue]wallet name[/blue]",
                     default=self.config.get("wallet_name") or defaults.wallet.name,
                 )
-            if netuid is not None:
+            if netuids is not None:
                 hotkey_or_ss58 = Prompt.ask(
-                    "Enter the [blue]wallet hotkey[/blue] name or [blue]ss58 address[/blue] to stake to [dim](or Press Enter to view delegates)[/dim]",
+                    "Enter the [blue]wallet hotkey[/blue] name or [blue]ss58 address[/blue] to stake to [dim]"
+                    "(or Press Enter to view delegates)[/dim]",
                 )
             else:
                 hotkey_or_ss58 = Prompt.ask(
@@ -3299,10 +3345,18 @@ def stake_add(
                 wallet = self.wallet_ask(
                     wallet_name, wallet_path, wallet_hotkey, ask_for=[WO.NAME, WO.PATH]
                 )
+                if len(netuids) > 1:
+                    netuid_ = IntPrompt.ask(
+                        "Enter the netuid for which to show delegates",
+                        choices=[str(x) for x in netuids],
+                    )
+                else:
+                    netuid_ = netuids[0]
+
                 selected_hotkey = self._run_command(
                     subnets.show(
                         subtensor=self.initialize_chain(network),
-                        netuid=netuid,
+                        netuid=netuid_,
                         sort=False,
                         max_rows=12,
                         prompt=False,
@@ -3312,7 +3366,7 @@ def stake_add(
                 )
                 if not selected_hotkey:
                     print_error("No delegate selected. Exiting.")
-                    raise typer.Exit()
+                    return
                 include_hotkeys = selected_hotkey
             elif is_valid_ss58_address(hotkey_or_ss58):
                 wallet = self.wallet_ask(
@@ -3373,8 +3427,8 @@ def stake_add(
             )
             if free_balance == Balance.from_tao(0):
                 print_error("You dont have any balance to stake.")
-                raise typer.Exit()
-            if netuid is not None:
+                return
+            if netuids:
                 amount = FloatPrompt.ask(
                     f"Amount to [{COLORS.G.SUBHEAD_MAIN}]stake (TAO τ)"
                 )
@@ -3396,7 +3450,7 @@ def stake_add(
             add_stake.stake_add(
                 wallet,
                 self.initialize_chain(network),
-                netuid,
+                netuids,
                 stake_all,
                 amount,
                 prompt,
@@ -4796,12 +4850,9 @@ def subnets_list(
     def subnets_price(
         self,
         network: Optional[list[str]] = Options.network,
-        netuids: str = typer.Option(
-            None,
-            "--netuids",
-            "--netuid",
-            "-n",
-            help="Netuid(s) to show the price for.",
+        netuids: str = Options.edit_help(
+            "netuids",
+            "Netuids to show the price for. Separate multiple netuids with a comma, for example: `-n 0,1,2`.",
         ),
         interval_hours: int = typer.Option(
             24,

bittensor_cli/src/commands/stake/add.py

@@ -31,7 +31,7 @@
 async def stake_add(
     wallet: Wallet,
     subtensor: "SubtensorInterface",
-    netuid: Optional[int],
+    netuids: Optional[list[int]],
     stake_all: bool,
     amount: float,
     prompt: bool,
@@ -48,7 +48,7 @@ async def stake_add(
     Args:
         wallet: wallet object
         subtensor: SubtensorInterface object
-        netuid: the netuid to stake to (None indicates all subnets)
+        netuids: the netuids to stake to (None indicates all subnets)
         stake_all: whether to stake all available balance
         amount: specified amount of balance to stake
         prompt: whether to prompt the user
@@ -233,9 +233,7 @@ async def stake_extrinsic(
                 return True
 
     netuids = (
-        [int(netuid)]
-        if netuid is not None
-        else await subtensor.get_all_subnet_netuids()
+        netuids if netuids is not None else await subtensor.get_all_subnet_netuids()
     )
 
     hotkeys_to_stake_to = _get_hotkeys_to_stake_to(
@@ -445,10 +443,10 @@ def _prompt_stake_amount(
     while True:
         amount_input = Prompt.ask(
             f"\nEnter the amount to {action_name}"
-            f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{Balance.get_unit(netuid)}[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] "
-            f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}](max: {current_balance})[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] "
+            f"[{COLOR_PALETTE.S.STAKE_AMOUNT}]{Balance.get_unit(netuid)}[/{COLOR_PALETTE.S.STAKE_AMOUNT}] "
+            f"[{COLOR_PALETTE.S.STAKE_AMOUNT}](max: {current_balance})[/{COLOR_PALETTE.S.STAKE_AMOUNT}] "
             f"or "
-            f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]'all'[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}] "
+            f"[{COLOR_PALETTE.S.STAKE_AMOUNT}]'all'[/{COLOR_PALETTE.S.STAKE_AMOUNT}] "
             f"for entire balance"
         )
 
@@ -463,7 +461,7 @@ def _prompt_stake_amount(
             if amount > current_balance.tao:
                 console.print(
                     f"[red]Amount exceeds available balance of "
-                    f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{current_balance}[/{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]"
+                    f"[{COLOR_PALETTE.S.STAKE_AMOUNT}]{current_balance}[/{COLOR_PALETTE.S.STAKE_AMOUNT}]"
                     f"[/red]"
                 )
                 continue
@@ -542,10 +540,10 @@ def _define_stake_table(
         Table: An initialized rich Table object with appropriate columns
     """
     table = Table(
-        title=f"\n[{COLOR_PALETTE['GENERAL']['HEADER']}]Staking to:\n"
-        f"Wallet: [{COLOR_PALETTE['GENERAL']['COLDKEY']}]{wallet.name}[/{COLOR_PALETTE['GENERAL']['COLDKEY']}], "
-        f"Coldkey ss58: [{COLOR_PALETTE['GENERAL']['COLDKEY']}]{wallet.coldkeypub.ss58_address}[/{COLOR_PALETTE['GENERAL']['COLDKEY']}]\n"
-        f"Network: {subtensor.network}[/{COLOR_PALETTE['GENERAL']['HEADER']}]\n",
+        title=f"\n[{COLOR_PALETTE.G.HEADER}]Staking to:\n"
+        f"Wallet: [{COLOR_PALETTE.G.CK}]{wallet.name}[/{COLOR_PALETTE.G.CK}], "
+        f"Coldkey ss58: [{COLOR_PALETTE.G.CK}]{wallet.coldkeypub.ss58_address}[/{COLOR_PALETTE.G.CK}]\n"
+        f"Network: {subtensor.network}[/{COLOR_PALETTE.G.HEADER}]\n",
         show_footer=True,
         show_edge=False,
         header_style="bold white",
@@ -609,9 +607,13 @@ def _print_table_and_slippage(table: Table, max_slippage: float, safe_staking: b
 
     # Greater than 5%
     if max_slippage > 5:
-        message = f"[{COLOR_PALETTE['STAKE']['SLIPPAGE_TEXT']}]-------------------------------------------------------------------------------------------------------------------\n"
-        message += f"[bold]WARNING:[/bold]  The slippage on one of your operations is high: [{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}]{max_slippage} %[/{COLOR_PALETTE['STAKE']['SLIPPAGE_PERCENT']}], this may result in a loss of funds.\n"
-        message += "-------------------------------------------------------------------------------------------------------------------\n"
+        message = (
+            f"[{COLOR_PALETTE.S.SLIPPAGE_TEXT}]" + ("-" * 115) + "\n"
+            f"[bold]WARNING:[/bold]  The slippage on one of your operations is high: "
+            f"[{COLOR_PALETTE.S.SLIPPAGE_PERCENT}]{max_slippage} %[/{COLOR_PALETTE.S.SLIPPAGE_PERCENT}], "
+            f"this may result in a loss of funds.\n" + ("-" * 115) + "\n"
+        )
+
         console.print(message)
 
     # Table description

tests/e2e_tests/test_staking_sudo.py

@@ -36,6 +36,7 @@ def test_staking(local_chain, wallet_setup):
     """
     print("Testing staking and sudo commands🧪")
     netuid = 2
+    multiple_netuids = [2, 3]
     wallet_path_alice = "//Alice"
 
     # Create wallet for Alice
@@ -91,7 +92,42 @@ def test_staking(local_chain, wallet_setup):
     )
     result_output = json.loads(result.stdout)
     assert result_output["success"] is True
-    assert result_output["netuid"] == 2
+    assert result_output["netuid"] == netuid
+
+    # Register another subnet with sudo as Alice
+    result_for_second_repo = 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 Subnet",
+            "--repo",
+            "https://github.com/username/repo",
+            "--contact",
+            "[email protected]",
+            "--url",
+            "https://testsubnet.com",
+            "--discord",
+            "alice#1234",
+            "--description",
+            "A test subnet for e2e testing",
+            "--additional-info",
+            "Created by Alice",
+            "--no-prompt",
+            "--json-output",
+        ],
+    )
+    result_output_second = json.loads(result_for_second_repo.stdout)
+    assert result_output_second["success"] is True
+    assert result_output_second["netuid"] == multiple_netuids[1]
 
     # Register Alice in netuid = 1 using her hotkey
     register_subnet = exec_command_alice(
@@ -192,7 +228,7 @@ def test_staking(local_chain, wallet_setup):
     assert get_identity_output["additional"] == sn_add_info
 
     # Add stake to Alice's hotkey
-    add_stake = exec_command_alice(
+    add_stake_single = exec_command_alice(
         command="stake",
         sub_command="add",
         extra_args=[
@@ -216,10 +252,10 @@ def test_staking(local_chain, wallet_setup):
             "144",
         ],
     )
-    assert "✅ Finalized" in add_stake.stdout, add_stake.stderr
+    assert "✅ Finalized" in add_stake_single.stdout, add_stake_single.stderr
 
     # Execute stake show for Alice's wallet
-    show_stake = exec_command_alice(
+    show_stake_adding_single = exec_command_alice(
         command="stake",
         sub_command="list",
         extra_args=[
@@ -235,7 +271,8 @@ def test_staking(local_chain, wallet_setup):
 
     # Assert correct stake is added
     cleaned_stake = [
-        re.sub(r"\s+", " ", line) for line in show_stake.stdout.splitlines()
+        re.sub(r"\s+", " ", line)
+        for line in show_stake_adding_single.stdout.splitlines()
     ]
     stake_added = cleaned_stake[8].split("│")[3].strip().split()[0]
     assert Balance.from_tao(float(stake_added)) >= Balance.from_tao(90)
@@ -284,6 +321,36 @@ def test_staking(local_chain, wallet_setup):
     )
     assert "✅ Finalized" in remove_stake.stdout
 
+    add_stake_multiple = exec_command_alice(
+        command="stake",
+        sub_command="add",
+        extra_args=[
+            "--netuids",
+            ",".join(str(x) for x in multiple_netuids),
+            "--wallet-path",
+            wallet_path_alice,
+            "--wallet-name",
+            wallet_alice.name,
+            "--hotkey",
+            wallet_alice.hotkey_str,
+            "--chain",
+            "ws://127.0.0.1:9945",
+            "--amount",
+            "100",
+            "--tolerance",
+            "0.1",
+            "--partial",
+            "--no-prompt",
+            "--era",
+            "144",
+        ],
+    )
+    assert "✅ Finalized" in add_stake_multiple.stdout, add_stake_multiple.stderr
+    for netuid_ in multiple_netuids:
+        assert f"Stake added to netuid: {netuid_}" in add_stake_multiple.stdout, (
+            add_stake_multiple.stderr
+        )
+
     # Fetch the hyperparameters of the subnet
     hyperparams = exec_command_alice(
         command="sudo",