diff --git a/bittensor/core/axon.py b/bittensor/core/axon.py index ce967acb80..378315f77d 100644 --- a/bittensor/core/axon.py +++ b/bittensor/core/axon.py @@ -59,7 +59,7 @@ from bittensor.core.stream import StreamingSynapse from bittensor.core.synapse import Synapse, TerminalInfo from bittensor.core.threadpool import PriorityThreadPoolExecutor -from bittensor.utils import networking +from bittensor.utils import networking, Certificate from bittensor.utils.axon_utils import allowed_nonce_window_ns, calculate_diff_seconds from bittensor.utils.btlogging import logging @@ -807,7 +807,12 @@ def stop(self) -> "Axon": self.started = False return self - def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon": + def serve( + self, + netuid: int, + subtensor: Optional["Subtensor"] = None, + certificate: Optional[Certificate] = None, + ) -> "Axon": """ Serves the Axon on the specified subtensor connection using the configured wallet. This method registers the Axon with a specific subnet within the Bittensor network, identified by the ``netuid``. @@ -832,7 +837,7 @@ def serve(self, netuid: int, subtensor: Optional["Subtensor"] = None) -> "Axon": 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: "Synapse"): diff --git a/bittensor/core/chain_data/__init__.py b/bittensor/core/chain_data/__init__.py index 68936a6b5f..760eaa3354 100644 --- a/bittensor/core/chain_data/__init__.py +++ b/bittensor/core/chain_data/__init__.py @@ -11,6 +11,7 @@ from .ip_info import IPInfo from .neuron_info import NeuronInfo from .neuron_info_lite import NeuronInfoLite +from .neuron_certificate import NeuronCertificate from .prometheus_info import PrometheusInfo from .proposal_vote_data import ProposalVoteData from .scheduled_coldkey_swap_info import ScheduledColdkeySwapInfo diff --git a/bittensor/core/chain_data/neuron_certificate.py b/bittensor/core/chain_data/neuron_certificate.py new file mode 100644 index 0000000000..a20f377d38 --- /dev/null +++ b/bittensor/core/chain_data/neuron_certificate.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass +from typing import List + +from bittensor.core.chain_data.utils import from_scale_encoding, ChainDataType +from bittensor.utils import Certificate + + +# Dataclasses for chain data. +@dataclass +class NeuronCertificate: + """ + Dataclass for neuron certificate. + """ + + certificate: Certificate + + @classmethod + def from_vec_u8(cls, vec_u8: List[int]) -> "NeuronCertificate": + """Returns a NeuronCertificate object from a ``vec_u8``.""" + return from_scale_encoding(vec_u8, ChainDataType.NeuronCertificate) diff --git a/bittensor/core/chain_data/utils.py b/bittensor/core/chain_data/utils.py index 9c21c9d22e..1218b9ea56 100644 --- a/bittensor/core/chain_data/utils.py +++ b/bittensor/core/chain_data/utils.py @@ -22,6 +22,7 @@ class ChainDataType(Enum): SubnetHyperparameters = 8 ScheduledColdkeySwapInfo = 9 AccountId = 10 + NeuronCertificate = 11 def from_scale_encoding( @@ -178,6 +179,12 @@ def from_scale_encoding_using_type_string( ["pruning_score", "Compact"], ], }, + "NeuronCertificate": { + "type": "struct", + "type_mapping": [ + ["certificate", "Vec"], + ], + }, "axon_info": { "type": "struct", "type_mapping": [ diff --git a/bittensor/core/extrinsics/serving.py b/bittensor/core/extrinsics/serving.py index 0408ca86b8..864726b97c 100644 --- a/bittensor/core/extrinsics/serving.py +++ b/bittensor/core/extrinsics/serving.py @@ -20,7 +20,12 @@ from bittensor.core.errors import MetadataError from bittensor.core.extrinsics.utils import submit_extrinsic from bittensor.core.settings import version_as_int -from bittensor.utils import format_error_message, networking as net, unlock_key +from bittensor.utils import ( + format_error_message, + networking as net, + unlock_key, + Certificate, +) from bittensor.utils.btlogging import logging from bittensor.utils.networking import ensure_connected @@ -57,9 +62,15 @@ def do_serve_axon( This function is crucial for initializing and announcing a neuron's ``Axon`` service on the network, 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" + 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(call=call, keypair=wallet.hotkey) @@ -90,6 +101,7 @@ def serve_extrinsic( placeholder2: int = 0, wait_for_inclusion: bool = False, wait_for_finalization=True, + certificate: Optional[Certificate] = None, ) -> bool: """Subscribes a Bittensor endpoint to the subtensor chain. @@ -124,6 +136,7 @@ def serve_extrinsic( "protocol": protocol, "placeholder1": placeholder1, "placeholder2": placeholder2, + "certificate": certificate, } logging.debug("Checking axon ...") neuron = subtensor.get_neuron_for_pubkey_and_subnet( @@ -182,6 +195,7 @@ def serve_axon_extrinsic( axon: "Axon", wait_for_inclusion: bool = False, wait_for_finalization: bool = True, + certificate: Optional[Certificate] = None, ) -> bool: """Serves the axon to the network. @@ -224,6 +238,7 @@ def serve_axon_extrinsic( protocol=4, wait_for_inclusion=wait_for_inclusion, wait_for_finalization=wait_for_finalization, + certificate=certificate, ) return serve_success diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index a080a9bac7..656a513afe 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -60,6 +60,7 @@ ss58_to_vec_u8, u16_normalized_float, hex_to_bytes, + Certificate, ) from bittensor.utils.balance import Balance from bittensor.utils.btlogging import logging @@ -992,6 +993,7 @@ def serve_axon( axon: "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 is used to set up the Axon, a key component of a neuron that handles incoming queries and data processing tasks. @@ -1008,7 +1010,7 @@ def serve_axon( By registering an Axon, the neuron becomes an active part of the network's distributed 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 ) # metagraph @@ -1149,6 +1151,41 @@ def get_neuron_for_pubkey_and_subnet( block=block, ) + 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], + ) + try: + serialized_certificate = certificate.serialize() + if serialized_certificate: + return ( + chr(serialized_certificate["algorithm"]) + + serialized_certificate["public_key"] + ) + except AttributeError: + return None + return None + @networking.ensure_connected def neuron_for_uid( self, uid: Optional[int], netuid: int, block: Optional[int] = None diff --git a/bittensor/core/types.py b/bittensor/core/types.py index 9fd2b4d052..577df5b6ba 100644 --- a/bittensor/core/types.py +++ b/bittensor/core/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 bittensor.utils import Certificate class AxonServeCallParams(TypedDict): @@ -26,6 +27,7 @@ class AxonServeCallParams(TypedDict): port: int ip_type: int netuid: int + certificate: Optional[Certificate] class PrometheusServeCallParams(TypedDict): diff --git a/bittensor/utils/__init__.py b/bittensor/utils/__init__.py index 02010526c7..90bd593a9f 100644 --- a/bittensor/utils/__init__.py +++ b/bittensor/utils/__init__.py @@ -40,6 +40,8 @@ U16_MAX = 65535 U64_MAX = 18446744073709551615 +Certificate = str + UnlockStatus = namedtuple("UnlockStatus", ["success", "message"]) diff --git a/bittensor/utils/mock/subtensor_mock.py b/bittensor/utils/mock/subtensor_mock.py index 817be08434..39095a7772 100644 --- a/bittensor/utils/mock/subtensor_mock.py +++ b/bittensor/utils/mock/subtensor_mock.py @@ -30,6 +30,7 @@ PrometheusInfo, AxonInfo, ) +from bittensor.core.types import AxonServeCallParams, PrometheusServeCallParams from bittensor.core.errors import ChainQueryError from bittensor.core.subtensor import Subtensor from bittensor.utils import RAOPERTAO, u16_normalized_float @@ -39,26 +40,6 @@ __GLOBAL_MOCK_STATE__ = {} -class AxonServeCallParams(TypedDict): - """Axon serve chain call parameters.""" - - version: int - ip: int - port: int - ip_type: int - netuid: int - - -class PrometheusServeCallParams(TypedDict): - """Prometheus serve chain call parameters.""" - - version: int - ip: int - port: int - ip_type: int - netuid: int - - BlockNumber = int diff --git a/tests/e2e_tests/test_neuron_certificate.py b/tests/e2e_tests/test_neuron_certificate.py new file mode 100644 index 0000000000..807640e74b --- /dev/null +++ b/tests/e2e_tests/test_neuron_certificate.py @@ -0,0 +1,60 @@ +import pytest +from bittensor.core.subtensor import Subtensor +from bittensor.core.axon import Axon +from bittensor.utils.btlogging import logging +from tests.e2e_tests.utils.chain_interactions import ( + wait_interval, + register_subnet, +) +from tests.e2e_tests.utils.e2e_test_utils import ( + setup_wallet, +) + + +@pytest.mark.asyncio +async def test_neuron_certificate(local_chain): + """ + Tests the metagraph + + Steps: + 1. Register a subnet through Alice + 2. Serve Alice axon with neuron certificate + 3. Verify neuron certificate can be retrieved + Raises: + AssertionError: If any of the checks or verifications fail + """ + logging.info("Testing neuron_certificate") + netuid = 1 + + # Register root as Alice - the subnet owner and validator + alice_keypair, alice_wallet = setup_wallet("//Alice") + register_subnet(local_chain, alice_wallet) + + # Verify subnet created successfully + assert local_chain.query( + "SubtensorModule", "NetworksAdded", [netuid] + ).serialize(), "Subnet wasn't created successfully" + + subtensor = Subtensor(network="ws://localhost:9945") + + # Register Alice as a neuron on the subnet + assert subtensor.burned_register( + alice_wallet, netuid + ), "Unable to register Alice as a neuron" + + # Serve Alice's axon with a certificate + axon = Axon(wallet=alice_wallet) + encoded_certificate = "?FAKE_ALICE_CERT" + axon.serve(netuid=netuid, subtensor=subtensor, certificate=encoded_certificate) + + await wait_interval(tempo=1, subtensor=subtensor, netuid=netuid) + + # Verify we are getting the correct certificate + assert ( + subtensor.get_neuron_certificate( + netuid=netuid, hotkey=alice_keypair.ss58_address + ) + == encoded_certificate + ) + + logging.info("✅ Passed test_neuron_certificate") diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py index 602a3027d8..fa7e190dc5 100644 --- a/tests/unit_tests/test_subtensor.py +++ b/tests/unit_tests/test_subtensor.py @@ -27,13 +27,35 @@ from bittensor.core.chain_data import SubnetHyperparameters from bittensor.core.settings import version_as_int from bittensor.core.subtensor import Subtensor, logging -from bittensor.utils import u16_normalized_float, u64_normalized_float +from bittensor.utils import u16_normalized_float, u64_normalized_float, Certificate from bittensor.utils.balance import Balance U16_MAX = 65535 U64_MAX = 18446744073709551615 +@pytest.fixture +def fake_call_params(): + return call_params() + + +def call_params(): + return { + "version": "1.0", + "ip": "0.0.0.0", + "port": 9090, + "ip_type": 4, + "netuid": 1, + "certificate": None, + } + + +def call_params_with_certificate(): + params = call_params() + params["certificate"] = Certificate("fake_cert") + return params + + def test_serve_axon_with_external_ip_set(): internal_ip: str = "192.0.2.146" external_ip: str = "2001:0db8:85a3:0000:0000:8a2e:0370:7334" @@ -1189,6 +1211,7 @@ def test_serve_axon(subtensor, mocker): fake_axon = mocker.MagicMock() fake_wait_for_inclusion = False fake_wait_for_finalization = True + fake_certificate = None mocked_serve_axon_extrinsic = mocker.patch.object( subtensor_module, "serve_axon_extrinsic" @@ -1206,6 +1229,7 @@ def test_serve_axon(subtensor, mocker): fake_axon, fake_wait_for_inclusion, fake_wait_for_finalization, + fake_certificate, ) assert result == mocked_serve_axon_extrinsic.return_value @@ -1400,11 +1424,19 @@ def test_neuron_for_uid_success(subtensor, mocker): assert result == mocked_neuron_from_vec_u8.return_value -def test_do_serve_axon_is_success(subtensor, mocker): +@pytest.mark.parametrize( + ["fake_call_params", "expected_call_function"], + [ + (call_params(), "serve_axon"), + (call_params_with_certificate(), "serve_axon_tls"), + ], +) +def test_do_serve_axon_is_success( + subtensor, mocker, fake_call_params, expected_call_function +): """Successful do_serve_axon call.""" # Prep fake_wallet = mocker.MagicMock() - fake_call_params = mocker.MagicMock() fake_wait_for_inclusion = True fake_wait_for_finalization = True @@ -1421,7 +1453,7 @@ def test_do_serve_axon_is_success(subtensor, mocker): # Asserts subtensor.substrate.compose_call.assert_called_once_with( call_module="SubtensorModule", - call_function="serve_axon", + call_function=expected_call_function, call_params=fake_call_params, ) @@ -1437,14 +1469,14 @@ def test_do_serve_axon_is_success(subtensor, mocker): ) subtensor.substrate.submit_extrinsic.return_value.process_events.assert_called_once() - assert result == (True, None) + assert result[0] is True + assert result[1] is None -def test_do_serve_axon_is_not_success(subtensor, mocker): +def test_do_serve_axon_is_not_success(subtensor, mocker, fake_call_params): """Unsuccessful do_serve_axon call.""" # Prep fake_wallet = mocker.MagicMock() - fake_call_params = mocker.MagicMock() fake_wait_for_inclusion = True fake_wait_for_finalization = True @@ -1483,11 +1515,10 @@ def test_do_serve_axon_is_not_success(subtensor, mocker): ) -def test_do_serve_axon_no_waits(subtensor, mocker): +def test_do_serve_axon_no_waits(subtensor, mocker, fake_call_params): """Unsuccessful do_serve_axon call.""" # Prep fake_wallet = mocker.MagicMock() - fake_call_params = mocker.MagicMock() fake_wait_for_inclusion = False fake_wait_for_finalization = False