From 8563a339a1de54b6e3ef827777fa6a788ce557f1 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Tue, 18 Feb 2025 19:27:56 +0200 Subject: [PATCH 1/6] Add methods for retrieving all neuron certificates on a given netuid. --- bittensor/core/async_subtensor.py | 34 +++++++++++++++++++++++++++++++ bittensor/core/subtensor.py | 25 +++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index 8e357aa141..6fa56c4bd4 100644 --- a/bittensor/core/async_subtensor.py +++ b/bittensor/core/async_subtensor.py @@ -1479,6 +1479,40 @@ async def get_neuron_certificate( return None return None + async def get_all_neuron_certificates( + self, + netuid: int, + block: Optional[int] = None, + block_hash: Optional[str] = None, + reuse_block: bool = False, + ) -> dict[str, Certificate]: + """ + Retrieves the TLS certificates for neurons within a specified subnet (netuid) of the Bittensor network. + + Arguments: + netuid: The unique identifier of the subnet. + block: The blockchain block number for the query. + block_hash: The hash of the block to retrieve the parameter from. Do not specify if using block or + reuse_block. + reuse_block: Whether to use the last-used block. Do not set if using block_hash or block. + + Returns: + {ss58: Certificate} for the key/Certificate pairs on the subnet + + This function is used for certificate discovery for setting up mutual tls communication between neurons. + """ + query_certificates = await self.query_map( + module="SubtensorModule", + name="NeuronCertificates", + block=block, + block_hash=block_hash, + reuse_block=reuse_block, + ) + output = {} + async for key, item in query_certificates: + output[decode_account_id(key)] = Certificate(item.value) + return output + async def get_neuron_for_pubkey_and_subnet( self, hotkey_ss58: str, diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index 03c32a6072..d055e873a2 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -1123,6 +1123,31 @@ def get_neuron_certificate( return None return None + def get_all_neuron_certificates( + self, netuid: int, block: Optional[int] = None + ) -> dict[str, Certificate]: + """ + Retrieves the TLS certificates for neurons within a specified subnet (netuid) of the Bittensor network. + + Arguments: + netuid: The unique identifier of the subnet. + block: The blockchain block number for the query. + + Returns: + {ss58: Certificate} for the key/Certificate pairs on the subnet + + This function is used for certificate discovery for setting up mutual tls communication between neurons. + """ + query_certificates = self.query_map( + module="SubtensorModule", + name="NeuronCertificates", + block=block, + ) + output = {} + for key, item in query_certificates: + output[decode_account_id(key)] = Certificate(item.value) + return output + def get_neuron_for_pubkey_and_subnet( self, hotkey_ss58: str, netuid: int, block: Optional[int] = None ) -> Optional["NeuronInfo"]: From b5a8d5c1e68bfd70a9d70d464b4db3aac0ac6fbb Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Tue, 18 Feb 2025 19:29:08 +0200 Subject: [PATCH 2/6] Add netuid --- bittensor/core/async_subtensor.py | 1 + bittensor/core/subtensor.py | 1 + 2 files changed, 2 insertions(+) diff --git a/bittensor/core/async_subtensor.py b/bittensor/core/async_subtensor.py index 6fa56c4bd4..d83c373140 100644 --- a/bittensor/core/async_subtensor.py +++ b/bittensor/core/async_subtensor.py @@ -1504,6 +1504,7 @@ async def get_all_neuron_certificates( query_certificates = await self.query_map( module="SubtensorModule", name="NeuronCertificates", + params=[netuid], block=block, block_hash=block_hash, reuse_block=reuse_block, diff --git a/bittensor/core/subtensor.py b/bittensor/core/subtensor.py index d055e873a2..e2d9ffda19 100644 --- a/bittensor/core/subtensor.py +++ b/bittensor/core/subtensor.py @@ -1141,6 +1141,7 @@ def get_all_neuron_certificates( query_certificates = self.query_map( module="SubtensorModule", name="NeuronCertificates", + params=[netuid], block=block, ) output = {} From bf5e391640fdf5bfb9f02fbc26fe3c87f179b42f Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Tue, 18 Feb 2025 19:34:10 +0200 Subject: [PATCH 3/6] Added test --- tests/e2e_tests/test_neuron_certificate.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/e2e_tests/test_neuron_certificate.py b/tests/e2e_tests/test_neuron_certificate.py index 245f9bc9fa..674a4ce27c 100644 --- a/tests/e2e_tests/test_neuron_certificate.py +++ b/tests/e2e_tests/test_neuron_certificate.py @@ -47,5 +47,8 @@ async def test_neuron_certificate(subtensor, alice_wallet): ) == encoded_certificate ) + all_certs_query = subtensor.get_all_neuron_certificates(netuid=netuid) + assert alice_wallet.hotkey.ss58_address in all_certs_query.keys() + assert all_certs_query[alice_wallet.hotkey.ss58_address] == encoded_certificate logging.info("✅ Passed test_neuron_certificate") From d93bc7dc71eeeff74d9e914908aa38055002fa21 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Tue, 18 Feb 2025 19:48:54 +0200 Subject: [PATCH 4/6] Added unit test --- tests/unit_tests/test_subtensor.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/unit_tests/test_subtensor.py b/tests/unit_tests/test_subtensor.py index 385bb45501..b6d63d72f7 100644 --- a/tests/unit_tests/test_subtensor.py +++ b/tests/unit_tests/test_subtensor.py @@ -3084,3 +3084,16 @@ def test_set_subnet_identity(mocker, subtensor): wait_for_inclusion=False, ) assert result == mocked_extrinsic.return_value + + +def test_get_all_neuron_certificates(mocker, subtensor): + fake_netuid = 12 + mocked_query_map_subtensor = mocker.MagicMock() + mocker.patch.object(subtensor.substrate, "query_map", mocked_query_map_subtensor) + subtensor.get_all_neuron_certificates(fake_netuid) + mocked_query_map_subtensor.assert_called_once_with( + module="SubtensorModule", + storage_function="NeuronCertificates", + params=[fake_netuid], + block_hash=None, + ) From 1c85c5309afbad429d2834085895f5d9fcc80dc4 Mon Sep 17 00:00:00 2001 From: Benjamin Himes Date: Tue, 18 Feb 2025 19:57:52 +0200 Subject: [PATCH 5/6] ut --- tests/unit_tests/test_async_subtensor.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index 83a29f3298..11b57dec23 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -2676,3 +2676,18 @@ async def test_set_subnet_identity(mocker, subtensor): wait_for_inclusion=False, ) assert result == mocked_extrinsic.return_value + + +@pytest.mark.asyncio +async def test_get_all_neuron_certificates(mocker, subtensor): + fake_netuid = 12 + mocked_query_map_subtensor = mocker.AsyncMock() + mocker.patch.object(subtensor.substrate, "query_map", mocked_query_map_subtensor) + await subtensor.get_all_neuron_certificates(fake_netuid) + mocked_query_map_subtensor.assert_called_once_with( + module="SubtensorModule", + storage_function="NeuronCertificates", + params=[fake_netuid], + block_hash=None, + reuse_block_hash=False, + ) From 7c144a45e45c5b1c5e024eb3cb7d57d9f4c83914 Mon Sep 17 00:00:00 2001 From: BD Himes <37844818+thewhaleking@users.noreply.github.com> Date: Tue, 18 Feb 2025 20:00:59 +0200 Subject: [PATCH 6/6] Update tests/unit_tests/test_async_subtensor.py Co-authored-by: Roman <167799377+roman-opentensor@users.noreply.github.com> --- tests/unit_tests/test_async_subtensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_async_subtensor.py b/tests/unit_tests/test_async_subtensor.py index 11b57dec23..b328044dbc 100644 --- a/tests/unit_tests/test_async_subtensor.py +++ b/tests/unit_tests/test_async_subtensor.py @@ -2684,7 +2684,7 @@ async def test_get_all_neuron_certificates(mocker, subtensor): mocked_query_map_subtensor = mocker.AsyncMock() mocker.patch.object(subtensor.substrate, "query_map", mocked_query_map_subtensor) await subtensor.get_all_neuron_certificates(fake_netuid) - mocked_query_map_subtensor.assert_called_once_with( + mocked_query_map_subtensor.assert_awaited_once_with( module="SubtensorModule", storage_function="NeuronCertificates", params=[fake_netuid],