Skip to content

Commit 472d178

Browse files
andreea-popescu-reefandreea-popescu-reef
committed
add neuron certificate handling
1 parent 09e0270 commit 472d178

File tree

7 files changed

+230
-2
lines changed

7 files changed

+230
-2
lines changed

‎pallets/subtensor/rpc/src/lib.rs

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,13 @@ pub trait SubtensorCustomApi<BlockHash> {
4141
fn get_neurons(&self, netuid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
4242
#[method(name = "neuronInfo_getNeuron")]
4343
fn get_neuron(&self, netuid: u16, uid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
44-
44+
#[method(name = "neuronInfo_getNeuronCertificate")]
45+
fn get_neuron_certificate(
46+
&self,
47+
netuid: u16,
48+
uid: u16,
49+
at: Option<BlockHash>,
50+
) -> RpcResult<Vec<u8>>;
4551
#[method(name = "subnetInfo_getSubnetInfo")]
4652
fn get_subnet_info(&self, netuid: u16, at: Option<BlockHash>) -> RpcResult<Vec<u8>>;
4753
#[method(name = "subnetInfo_getSubnetsInfo")]
@@ -212,6 +218,25 @@ where
212218
})
213219
}
214220

221+
fn get_neuron_certificate(
222+
&self,
223+
netuid: u16,
224+
uid: u16,
225+
at: Option<<Block as BlockT>::Hash>,
226+
) -> RpcResult<Vec<u8>> {
227+
let api = self.client.runtime_api();
228+
let at = at.unwrap_or_else(|| self.client.info().best_hash);
229+
230+
api.get_neuron_certificate(at, netuid, uid).map_err(|e| {
231+
CallError::Custom(ErrorObject::owned(
232+
Error::RuntimeError.into(),
233+
"Unable to get neuron certificate.",
234+
Some(e.to_string()),
235+
))
236+
.into()
237+
})
238+
}
239+
215240
fn get_subnet_info(
216241
&self,
217242
netuid: u16,

‎pallets/subtensor/runtime-api/src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ sp_api::decl_runtime_apis! {
1414
pub trait NeuronInfoRuntimeApi {
1515
fn get_neurons(netuid: u16) -> Vec<u8>;
1616
fn get_neuron(netuid: u16, uid: u16) -> Vec<u8>;
17+
fn get_neuron_certificate(netuid: u16, uid: u16) -> Vec<u8>;
1718
fn get_neurons_lite(netuid: u16) -> Vec<u8>;
1819
fn get_neuron_lite(netuid: u16, uid: u16) -> Vec<u8>;
1920
}

‎pallets/subtensor/src/lib.rs

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ mod benchmarks;
3535
// =========================
3636
mod block_step;
3737

38+
mod certificate;
3839
mod epoch;
3940
mod math;
4041
mod registration;
@@ -57,6 +58,7 @@ pub mod migration;
5758
#[frame_support::pallet]
5859
pub mod pallet {
5960

61+
use crate::certificate::Certificate;
6062
use frame_support::{
6163
dispatch::GetDispatchInfo,
6264
inherent::Vec,
@@ -522,6 +524,14 @@ pub mod pallet {
522524
pub placeholder2: u8, // --- Axon proto placeholder 1.
523525
}
524526

527+
// --- Struct for NeuronCertificate.
528+
pub type NeuronCertificateOf = NeuronCertificate;
529+
530+
#[derive(Decode, Encode, Default, TypeInfo, PartialEq, Eq, Clone, Debug)]
531+
pub struct NeuronCertificate {
532+
pub certificate: Certificate,
533+
}
534+
525535
// --- Struct for Prometheus.
526536
pub type PrometheusInfoOf = PrometheusInfo;
527537
#[derive(Encode, Decode, Default, TypeInfo, Clone, PartialEq, Eq, Debug)]
@@ -570,6 +580,16 @@ pub mod pallet {
570580
PrometheusInfoOf,
571581
OptionQuery,
572582
>;
583+
#[pallet::storage] // --- MAP ( netuid, hotkey ) --> certificate
584+
pub(super) type NeuronCertificates<T: Config> = StorageDoubleMap<
585+
_,
586+
Identity,
587+
u16,
588+
Blake2_128Concat,
589+
T::AccountId,
590+
NeuronCertificateOf,
591+
OptionQuery,
592+
>;
573593

574594
// =======================================
575595
// ==== Subnetwork Hyperparam storage ====
@@ -1446,6 +1466,91 @@ pub mod pallet {
14461466
protocol,
14471467
placeholder1,
14481468
placeholder2,
1469+
None,
1470+
)
1471+
}
1472+
1473+
// ---- Serves or updates axon /promethteus information for the neuron associated with the caller. If the caller is
1474+
// already registered the metadata is updated. If the caller is not registered this call throws NotRegistered.
1475+
//
1476+
// # Args:
1477+
// * 'origin': (<T as frame_system::Config>Origin):
1478+
// - The signature of the caller.
1479+
//
1480+
// * 'netuid' (u16):
1481+
// - The u16 network identifier.
1482+
//
1483+
// * 'version' (u64):
1484+
// - The bittensor version identifier.
1485+
//
1486+
// * 'ip' (u64):
1487+
// - The endpoint ip information as a u128 encoded integer.
1488+
//
1489+
// * 'port' (u16):
1490+
// - The endpoint port information as a u16 encoded integer.
1491+
//
1492+
// * 'ip_type' (u8):
1493+
// - The endpoint ip version as a u8, 4 or 6.
1494+
//
1495+
// * 'protocol' (u8):
1496+
// - UDP:1 or TCP:0
1497+
//
1498+
// * 'placeholder1' (u8):
1499+
// - Placeholder for further extra params.
1500+
//
1501+
// * 'placeholder2' (u8):
1502+
// - Placeholder for further extra params.
1503+
//
1504+
// * 'certificate' (Vec<u8>):
1505+
// - TLS certificate
1506+
//
1507+
// # Event:
1508+
// * AxonServed;
1509+
// - On successfully serving the axon info.
1510+
//
1511+
// # Raises:
1512+
// * 'NetworkDoesNotExist':
1513+
// - Attempting to set weights on a non-existent network.
1514+
//
1515+
// * 'NotRegistered':
1516+
// - Attempting to set weights from a non registered account.
1517+
//
1518+
// * 'InvalidIpType':
1519+
// - The ip type is not 4 or 6.
1520+
//
1521+
// * 'InvalidIpAddress':
1522+
// - The numerically encoded ip address does not resolve to a proper ip.
1523+
//
1524+
// * 'ServingRateLimitExceeded':
1525+
// - Attempting to set prometheus information withing the rate limit min.
1526+
//
1527+
#[pallet::call_index(40)]
1528+
#[pallet::weight((Weight::from_ref_time(19_000_000)
1529+
.saturating_add(T::DbWeight::get().reads(2))
1530+
.saturating_add(T::DbWeight::get().writes(1)), DispatchClass::Normal, Pays::No))]
1531+
pub fn serve_axon_tls(
1532+
origin: OriginFor<T>,
1533+
netuid: u16,
1534+
version: u32,
1535+
ip: u128,
1536+
port: u16,
1537+
ip_type: u8,
1538+
protocol: u8,
1539+
placeholder1: u8,
1540+
placeholder2: u8,
1541+
certificate: Certificate,
1542+
) -> DispatchResult {
1543+
Self::do_serve_axon(
1544+
origin,
1545+
netuid,
1546+
version,
1547+
ip,
1548+
port,
1549+
ip_type,
1550+
protocol,
1551+
placeholder1,
1552+
placeholder2,
1553+
Some(certificate),
14491554
)
14501555
}
14511556

@@ -1869,6 +1974,10 @@ where
18691974
let transaction_fee = 0;
18701975
Ok((CallType::Serve, transaction_fee, who.clone()))
18711976
}
1977+
Some(Call::serve_axon_tls { .. }) => {
1978+
let transaction_fee = 0;
1979+
Ok((CallType::Serve, transaction_fee, who.clone()))
1980+
}
18721981
Some(Call::register_network { .. }) => {
18731982
let transaction_fee = 0;
18741983
Ok((CallType::RegisterNetwork, transaction_fee, who.clone()))

‎pallets/subtensor/src/neuron_info.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,28 @@ impl<T: Config> Pallet<T> {
170170
neuron
171171
}
172172

173+
pub fn get_neuron_certificate(netuid: u16, uid: u16) -> Option<NeuronCertificate> {
174+
if !Self::if_subnet_exist(netuid) {
175+
return None;
176+
}
177+
178+
let _hotkey = Self::get_hotkey_for_net_and_uid(netuid, uid);
179+
let hotkey;
180+
if _hotkey.is_err() {
181+
return None;
182+
} else {
183+
// No error, hotkey was registered
184+
hotkey = _hotkey.expect("Hotkey should exist");
185+
}
186+
187+
if Self::has_neuron_certificate(netuid, &hotkey) {
188+
let certificate = NeuronCertificates::<T>::get(netuid, hotkey).unwrap();
189+
Some(certificate)
190+
} else {
191+
None
192+
}
193+
}
194+
173195
fn get_neuron_lite_subnet_exists(netuid: u16, uid: u16) -> Option<NeuronInfoLite<T>> {
174196
let _hotkey = Self::get_hotkey_for_net_and_uid(netuid, uid);
175197
let hotkey;

‎pallets/subtensor/src/serving.rs

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
use super::*;
2+
use crate::certificate::Certificate;
23
use frame_support::inherent::Vec;
34
use frame_support::sp_std::vec;
45

@@ -63,6 +64,7 @@ impl<T: Config> Pallet<T> {
6364
protocol: u8,
6465
placeholder1: u8,
6566
placeholder2: u8,
67+
certificate: Option<Certificate>,
6668
) -> dispatch::DispatchResult {
6769
// --- 1. We check the callers (hotkey) signature.
6870
let hotkey_id = ensure_signed(origin)?;
@@ -88,6 +90,11 @@ impl<T: Config> Pallet<T> {
8890
Error::<T>::ServingRateLimitExceeded
8991
);
9092

93+
// --- 5. Check certificate
94+
if let Some(certificate) = certificate {
95+
NeuronCertificates::<T>::insert(netuid, hotkey_id.clone(), NeuronCertificate{certificate})
96+
}
97+
9198
// --- 6. We insert the axon meta.
9299
prev_axon.block = Self::get_current_block_as_u64();
93100
prev_axon.version = version;
@@ -241,6 +248,10 @@ impl<T: Config> Pallet<T> {
241248
return Axons::<T>::contains_key(netuid, hotkey);
242249
}
243250

251+
pub fn has_neuron_certificate(netuid: u16, hotkey: &T::AccountId) -> bool {
252+
return NeuronCertificates::<T>::contains_key(netuid, hotkey);
253+
}
254+
244255
pub fn has_prometheus_info(netuid: u16, hotkey: &T::AccountId) -> bool {
245256
return Prometheus::<T>::contains_key(netuid, hotkey);
246257
}

‎pallets/subtensor/tests/serving.rs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,56 @@ fn test_serving_ok() {
9595
});
9696
}
9797

98+
#[test]
99+
fn test_serving_tls_ok() {
100+
new_test_ext().execute_with(|| {
101+
let hotkey_account_id = U256::from(1);
102+
let uid: u16 = 0;
103+
let netuid: u16 = 1;
104+
let tempo: u16 = 13;
105+
let version: u32 = 2;
106+
let ip: u128 = 1676056785;
107+
let port: u16 = 128;
108+
let ip_type: u8 = 4;
109+
let modality: u16 = 0;
110+
let protocol: u8 = 0;
111+
let placeholder1: u8 = 0;
112+
let placeholder2: u8 = 0;
113+
let certificate: Vec<u8> = "CERT".as_bytes().to_vec();
114+
add_network(netuid, tempo, modality);
115+
register_ok_neuron(netuid, hotkey_account_id, U256::from(66), 0);
116+
assert_ok!(SubtensorModule::serve_axon_tls(
117+
<<Test as Config>::RuntimeOrigin>::signed(hotkey_account_id),
118+
netuid,
119+
version,
120+
ip,
121+
port,
122+
ip_type,
123+
protocol,
124+
placeholder1,
125+
placeholder2,
126+
certificate.clone()
127+
));
128+
let stored_certificate = SubtensorModule::get_neuron_certificate(netuid, uid);
129+
assert_eq!(stored_certificate.unwrap().certificate, certificate);
130+
let new_certificate = "UPDATED_CERT".as_bytes().to_vec();
131+
assert_ok!(SubtensorModule::serve_axon_tls(
132+
<<Test as Config>::RuntimeOrigin>::signed(hotkey_account_id),
133+
netuid,
134+
version,
135+
ip,
136+
port,
137+
ip_type,
138+
protocol,
139+
placeholder1,
140+
placeholder2,
141+
new_certificate.clone()
142+
));
143+
let stored_certificate = SubtensorModule::get_neuron_certificate(netuid, uid);
144+
assert_eq!(stored_certificate.unwrap().certificate, new_certificate)
145+
});
146+
}
147+
98148
#[test]
99149
fn test_serving_set_metadata_update() {
100150
new_test_ext().execute_with(|| {

‎runtime/src/lib.rs

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion {
121121
// `spec_version`, and `authoring_version` are the same between Wasm and native.
122122
// This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use
123123
// the compatible custom types.
124-
spec_version: 144,
124+
spec_version: 145,
125125
impl_version: 1,
126126
apis: RUNTIME_API_VERSIONS,
127127
transaction_version: 1,
@@ -1328,6 +1328,16 @@ impl_runtime_apis! {
13281328
vec![]
13291329
}
13301330
}
1331+
1332+
fn get_neuron_certificate(netuid: u16, uid: u16) -> Vec<u8> {
1333+
let _result = SubtensorModule::get_neuron_certificate(netuid, uid);
1334+
if _result.is_some() {
1335+
let result = _result.expect("Could not get Certificate");
1336+
result.encode()
1337+
} else {
1338+
vec![]
1339+
}
1340+
}
13311341
}
13321342

13331343
impl subtensor_custom_rpc_runtime_api::SubnetInfoRuntimeApi<Block> for Runtime {

0 commit comments

Comments
 (0)