add neuron certificate discovery

bittensor/axon.py

@@ -34,6 +34,7 @@
 import warnings
 from inspect import signature, Signature, Parameter
 from typing import List, Optional, Tuple, Callable, Any, Dict, Awaitable
+from .utils import Certificate
 
 import uvicorn
 from fastapi import APIRouter, Depends, FastAPI
@@ -832,7 +833,10 @@ def stop(self) -> "bittensor.axon":
         return self
 
     def serve(
-        self, netuid: int, subtensor: Optional[bittensor.subtensor] = None
+        self,
+        netuid: int,
+        subtensor: Optional[bittensor.subtensor] = None,
+        certificate: Optional[Certificate] = None,
     ) -> "bittensor.axon":
         """
         Serves the Axon on the specified subtensor connection using the configured wallet. This method
@@ -858,7 +862,7 @@ def serve(
             to start receiving and processing requests from other neurons.
         """
         if subtensor is not None and hasattr(subtensor, "serve_axon"):
-            subtensor.serve_axon(netuid=netuid, axon=self)
+            subtensor.serve_axon(netuid=netuid, axon=self, certificate=certificate)
         return self
 
     async def default_verify(self, synapse: bittensor.Synapse):

bittensor/chain_data.py

@@ -34,6 +34,7 @@
 from .utils import networking as net, RAOPERTAO, U16_NORMALIZED_FLOAT
 from .utils.balance import Balance
 from .utils.registration import torch, use_torch
+from .utils import Certificate
 
 custom_rpc_type_registry = {
     "types": {
@@ -73,6 +74,12 @@
                 ["total_daily_return", "Compact<u64>"],
             ],
         },
+        "NeuronCertificate": {
+            "type": "struct",
+            "type_mapping": [
+                ["certificate", "Vec<u8>"],
+            ],
+        },
         "NeuronInfo": {
             "type": "struct",
             "type_mapping": [
@@ -333,6 +340,7 @@ class ChainDataType(Enum):
     IPInfo = 7
     SubnetHyperparameters = 8
     ScheduledColdkeySwapInfo = 9
+    NeuronCertificate = 10
 
 
 def from_scale_encoding(
@@ -540,6 +548,21 @@ def from_weights_bonds_and_neuron_lite(
         return cls(**n_dict)
 
 
+# Dataclasses for chain data.
+@dataclass
+class NeuronCertificate:
+    r"""
+    Dataclass for neuron certificate.
+    """
+
+    certificate: Certificate
+
+    @classmethod
+    def from_vec_u8(cls, vec_u8: List[int]) -> "NeuronCertificate":
+        r"""Returns a NeuronCertificate object from a ``vec_u8``."""
+        return from_scale_encoding(vec_u8, ChainDataType.NeuronCertificate)
+
+
 @dataclass
 class NeuronInfoLite:
     """Dataclass for neuron metadata, but without the weights and bonds."""

bittensor/commands/stake.py

@@ -44,19 +44,25 @@
 
 
 def get_netuid(
-    cli: "bittensor.cli", subtensor: "bittensor.subtensor"
+    cli: "bittensor.cli", subtensor: "bittensor.subtensor", prompt: bool = True
 ) -> Tuple[bool, int]:
     """Retrieve and validate the netuid from the user or configuration."""
     console = Console()
-    if not cli.config.is_set("netuid"):
-        try:
-            cli.config.netuid = int(Prompt.ask("Enter netuid"))
-        except ValueError:
-            console.print(
-                "[red]Invalid input. Please enter a valid integer for netuid.[/red]"
-            )
-            return False, -1
+    if not cli.config.is_set("netuid") and prompt:
+        cli.config.netuid = Prompt.ask("Enter netuid")
+    try:
+        cli.config.netuid = int(cli.config.netuid)
+    except ValueError:
+        console.print(
+            "[red]Invalid input. Please enter a valid integer for netuid.[/red]"
+        )
+        return False, -1
     netuid = cli.config.netuid
+    if netuid < 0 or netuid > 65535:
+        console.print(
+            "[red]Invalid input. Please enter a valid integer for netuid in subnet range.[/red]"
+        )
+        return False, -1
     if not subtensor.subnet_exists(netuid=netuid):
         console.print(
             "[red]Network with netuid {} does not exist. Please try again.[/red]".format(
@@ -1136,10 +1142,27 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
         wallet = bittensor.wallet(config=cli.config)
 
         # check all
-        if not cli.config.is_set("all"):
-            exists, netuid = get_netuid(cli, subtensor)
-            if not exists:
-                return
+        if cli.config.is_set("all"):
+            cli.config.netuid = None
+            cli.config.all = True
+        elif cli.config.is_set("netuid"):
+            if cli.config.netuid == "all":
+                cli.config.all = True
+            else:
+                cli.config.netuid = int(cli.config.netuid)
+                exists, netuid = get_netuid(cli, subtensor)
+                if not exists:
+                    return
+        else:
+            netuid_input = Prompt.ask("Enter netuid or 'all'", default="all")
+            if netuid_input == "all":
+                cli.config.netuid = None
+                cli.config.all = True
+            else:
+                cli.config.netuid = int(netuid_input)
+                exists, netuid = get_netuid(cli, subtensor, False)
+                if not exists:
+                    return
 
         # get parent hotkey
         hotkey = get_hotkey(wallet, cli.config)
@@ -1148,11 +1171,7 @@ def _run(cli: "bittensor.cli", subtensor: "bittensor.subtensor"):
             return
 
         try:
-            netuids = (
-                subtensor.get_all_subnet_netuids()
-                if cli.config.is_set("all")
-                else [netuid]
-            )
+            netuids = subtensor.get_all_subnet_netuids() if cli.config.all else [netuid]
             hotkey_stake = GetChildrenCommand.get_parent_stake_info(
                 console, subtensor, hotkey
             )
@@ -1236,7 +1255,7 @@ def add_args(parser: argparse.ArgumentParser):
         parser = parser.add_parser(
             "get_children", help="""Get child hotkeys on subnet."""
         )
-        parser.add_argument("--netuid", dest="netuid", type=int, required=False)
+        parser.add_argument("--netuid", dest="netuid", type=str, required=False)
         parser.add_argument("--hotkey", dest="hotkey", type=str, required=False)
         parser.add_argument(
             "--all",
@@ -1294,7 +1313,7 @@ def render_table(
 
         # Add columns to the table with specific styles
         table.add_column("Index", style="bold yellow", no_wrap=True, justify="center")
-        table.add_column("ChildHotkey", style="bold green")
+        table.add_column("Child Hotkey", style="bold green")
         table.add_column("Proportion", style="bold cyan", no_wrap=True, justify="right")
         table.add_column(
             "Childkey Take", style="bold blue", no_wrap=True, justify="right"

bittensor/extrinsics/serving.py

@@ -26,6 +26,8 @@
 import bittensor.utils.networking as net
 from bittensor.utils import format_error_message
 from ..errors import MetadataError
+from typing import Optional
+from ..utils import Certificate
 
 
 def serve_extrinsic(
@@ -40,6 +42,7 @@ def serve_extrinsic(
     wait_for_inclusion: bool = False,
     wait_for_finalization=True,
     prompt: bool = False,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     r"""Subscribes a Bittensor endpoint to the subtensor chain.
 
@@ -50,6 +53,8 @@ def serve_extrinsic(
             Endpoint host port i.e., ``192.122.31.4``.
         port (int):
             Endpoint port number i.e., ``9221``.
+        certificate (str):
+            TLS Certificate
         protocol (int):
             An ``int`` representation of the protocol.
         netuid (int):
@@ -81,6 +86,7 @@ def serve_extrinsic(
         "protocol": protocol,
         "placeholder1": placeholder1,
         "placeholder2": placeholder2,
+        "certificate": certificate,
     }
     bittensor.logging.debug("Checking axon ...")
     neuron = subtensor.get_neuron_for_pubkey_and_subnet(
@@ -148,6 +154,7 @@ def serve_axon_extrinsic(
     wait_for_inclusion: bool = False,
     wait_for_finalization: bool = True,
     prompt: bool = False,
+    certificate: Optional[Certificate] = None,
 ) -> bool:
     r"""Serves the axon to the network.
 
@@ -200,6 +207,7 @@ def serve_axon_extrinsic(
         protocol=4,
         wait_for_inclusion=wait_for_inclusion,
         wait_for_finalization=wait_for_finalization,
+        certificate=certificate,
     )
     return serve_success
 

bittensor/subtensor.py

@@ -45,6 +45,7 @@
 from bittensor.utils import torch, weight_utils, format_error_message
 from .chain_data import (
     DelegateInfoLite,
+    NeuronCertificate,
     NeuronInfo,
     DelegateInfo,
     PrometheusInfo,
@@ -115,6 +116,7 @@
     ss58_to_vec_u8,
     U64_NORMALIZED_FLOAT,
     networking,
+    Certificate,
 )
 from .utils.balance import Balance
 from .utils.registration import POWSolution
@@ -1814,6 +1816,7 @@ def serve(
         placeholder2: int = 0,
         wait_for_inclusion: bool = False,
         wait_for_finalization=True,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers a neuron's serving endpoint on the Bittensor network. This function announces the
@@ -1850,6 +1853,7 @@ def serve(
             placeholder2,
             wait_for_inclusion,
             wait_for_finalization,
+            certificate=certificate,
         )
 
     def serve_axon(
@@ -1858,6 +1862,7 @@ def serve_axon(
         axon: "bittensor.axon",
         wait_for_inclusion: bool = False,
         wait_for_finalization: bool = True,
+        certificate: Optional[Certificate] = None,
     ) -> bool:
         """
         Registers an Axon serving endpoint on the Bittensor network for a specific neuron. This function
@@ -1877,7 +1882,12 @@ def serve_axon(
         computing infrastructure, contributing to the collective intelligence of Bittensor.
         """
         return serve_axon_extrinsic(
-            self, netuid, axon, wait_for_inclusion, wait_for_finalization
+            self,
+            netuid,
+            axon,
+            wait_for_inclusion,
+            wait_for_finalization,
+            certificate=certificate,
         )
 
     def _do_serve_axon(
@@ -1904,11 +1914,17 @@ def _do_serve_axon(
         enhancing the decentralized computation capabilities of Bittensor.
         """
 
+        if call_params["certificate"] is None:
+            del call_params["certificate"]
+            call_function = "serve_axon"
+        else:
+            call_function = "serve_axon_tls"
+
         @retry(delay=1, tries=3, backoff=2, max_delay=4, logger=_logger)
         def make_substrate_call_with_retry():
             call = self.substrate.compose_call(
                 call_module="SubtensorModule",
-                call_function="serve_axon",
+                call_function=call_function,
                 call_params=call_params,
             )
             extrinsic = self.substrate.create_signed_extrinsic(
@@ -5144,6 +5160,36 @@ def make_substrate_call_with_retry():
 
         return NeuronInfo.from_vec_u8(result)
 
+    def get_neuron_certificate(
+        self, hotkey: str, netuid: int, block: Optional[int] = None
+    ) -> Optional[Certificate]:
+        """
+        Retrieves the TLS certificate for a specific neuron identified by its unique identifier (UID)
+        within a specified subnet (netuid) of the Bittensor network.
+        Args:
+            hotkey (str): The hotkey to query.
+            netuid (int): The unique identifier of the subnet.
+            block (Optional[int], optional): The blockchain block number for the query.
+
+        Returns:
+            Optional[Certificate]: the certificate of the neuron if found, ``None`` otherwise.
+
+        This function is used for certificate discovery for setting up mutual tls communication between neurons
+        """
+
+        certificate = self.query_module(
+            module="SubtensorModule",
+            name="NeuronCertificates",
+            block=block,
+            params=[netuid, hotkey],
+        )
+        if not hasattr(certificate, "serialize"):
+            return None
+        certificate = certificate.serialize()
+        if not certificate:
+            return None
+        return certificate.get("certificate", None)
+
     def neurons(self, netuid: int, block: Optional[int] = None) -> List[NeuronInfo]:
         """
         Retrieves a list of all neurons within a specified subnet of the Bittensor network. This function

bittensor/types.py

@@ -15,7 +15,8 @@
 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 # DEALINGS IN THE SOFTWARE.
 
-from typing import TypedDict
+from typing import TypedDict, Optional
+from .utils import Certificate
 
 
 class AxonServeCallParams(TypedDict):
@@ -28,6 +29,7 @@ class AxonServeCallParams(TypedDict):
     port: int
     ip_type: int
     netuid: int
+    certificate: Optional[Certificate]
 
 
 class PrometheusServeCallParams(TypedDict):

bittensor/utils/__init__.py

@@ -31,6 +31,8 @@
 U16_MAX = 65535
 U64_MAX = 18446744073709551615
 
+Certificate = str
+
 
 def ss58_to_vec_u8(ss58_address: str) -> List[int]:
     ss58_bytes: bytes = bittensor.utils.ss58_address_to_bytes(ss58_address)