Skip to content

Commit 40fcaf8

Browse files
JannikStJannikSt
authored
support dynamic staking system and ability to automatically increase stake (#136)
* support dynamic staking system and ability to automatically increase stake * remove provider stage arg on worker
1 parent 82a9b4d commit 40fcaf8

File tree

12 files changed

+319
-32
lines changed

12 files changed

+319
-32
lines changed

‎Makefile

+8-6
Original file line numberDiff line numberDiff line change
@@ -73,22 +73,22 @@ watch-discovery:
7373

7474
watch-worker:
7575
set -a; source ${ENV_FILE}; set +a; \
76-
cargo watch -w worker/src -x "run --bin worker -- run --private-key-provider $$PROVIDER_PRIVATE_KEY --private-key-node $$NODE_PRIVATE_KEY --port 8091 --external-ip 0.0.0.0 --compute-pool-id 0 --validator-address $$VALIDATOR_ADDRESS"
76+
cargo watch -w worker/src -x "run --bin worker -- run --private-key-provider $$PROVIDER_PRIVATE_KEY --private-key-node $$NODE_PRIVATE_KEY --port 8091 --external-ip 0.0.0.0 --compute-pool-id $$WORKER_COMPUTE_POOL_ID --validator-address $$VALIDATOR_ADDRESS"
7777

7878
watch-validator:
7979
set -a; source ${ENV_FILE}; set +a; \
80-
cargo watch -w validator/src -x "run --bin validator -- --validator-key $${PRIVATE_KEY_VALIDATOR} --rpc-url $${RPC_URL} --pool-id 0 --work-validation-contract $${WORK_VALIDATION_CONTRACT} --leviticus-url $${LEVITICUS_URL}"
80+
cargo watch -w validator/src -x "run --bin validator -- --validator-key $${PRIVATE_KEY_VALIDATOR} --rpc-url $${RPC_URL} --pool-id 1 --work-validation-contract $${WORK_VALIDATION_CONTRACT} --leviticus-url $${LEVITICUS_URL}"
8181

8282
watch-orchestrator:
8383
set -a; source ${ENV_FILE}; set +a; \
84-
cargo watch -w orchestrator/src -x "run --bin orchestrator -- -r $$RPC_URL -k $$POOL_OWNER_PRIVATE_KEY -d 0 -p 8090 -i 10 -u http://localhost:8090 --s3-credentials $$S3_CREDENTIALS"
84+
cargo watch -w orchestrator/src -x "run --bin orchestrator -- -r $$RPC_URL -k $$POOL_OWNER_PRIVATE_KEY -d 0 -p 8090 -i 10 -u http://localhost:8090 --s3-credentials $$S3_CREDENTIALS --compute-pool-id $$WORKER_COMPUTE_POOL_ID"
8585

8686
build-worker:
8787
cargo build --release --bin worker
8888

8989
run-worker-bin:
9090
set -a; source .env; set +a; \
91-
./target/release/worker run --private-key-provider $$PROVIDER_PRIVATE_KEY --private-key-node $$NODE_PRIVATE_KEY --port 8091 --external-ip 0.0.0.0 --compute-pool-id 0 --validator-address $$VALIDATOR_ADDRESS
91+
./target/release/worker run --private-key-provider $$PROVIDER_PRIVATE_KEY --private-key-node $$NODE_PRIVATE_KEY --port 8091 --external-ip 0.0.0.0 --compute-pool-id $$WORKER_COMPUTE_POOL_ID --validator-address $$VALIDATOR_ADDRESS
9292

9393
SSH_CONNECTION ?= your-ssh-conn string
9494
EXTERNAL_IP ?= 0.0.0.0
@@ -143,8 +143,10 @@ watch-worker-remote: setup-remote setup-tunnel sync-remote
143143
--private-key-node \$$NODE_PRIVATE_KEY \
144144
--port $(PORT) \
145145
--external-ip \$$EXTERNAL_IP \
146-
--compute-pool-id 0 \
147-
--validator-address \$$VALIDATOR_ADDRESS 2>&1 | tee worker.log\""
146+
--compute-pool-id \$$WORKER_COMPUTE_POOL_ID \
147+
--validator-address \$$VALIDATOR_ADDRESS \
148+
2>&1 | tee worker.log\""
149+
148150
# Kill SSH tunnel
149151
.PHONY: kill-tunnel
150152
kill-tunnel:

‎shared/artifacts/abi/compute_registry.json

+38
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,25 @@
510510
],
511511
"stateMutability": "view"
512512
},
513+
{
514+
"type": "function",
515+
"name": "getProviderTotalCompute",
516+
"inputs": [
517+
{
518+
"name": "provider",
519+
"type": "address",
520+
"internalType": "address"
521+
}
522+
],
523+
"outputs": [
524+
{
525+
"name": "",
526+
"type": "uint256",
527+
"internalType": "uint256"
528+
}
529+
],
530+
"stateMutability": "view"
531+
},
513532
{
514533
"type": "function",
515534
"name": "getProviderTotalNodes",
@@ -714,6 +733,25 @@
714733
],
715734
"stateMutability": "view"
716735
},
736+
{
737+
"type": "function",
738+
"name": "providerTotalCompute",
739+
"inputs": [
740+
{
741+
"name": "",
742+
"type": "address",
743+
"internalType": "address"
744+
}
745+
],
746+
"outputs": [
747+
{
748+
"name": "",
749+
"type": "uint256",
750+
"internalType": "uint256"
751+
}
752+
],
753+
"stateMutability": "view"
754+
},
717755
{
718756
"type": "function",
719757
"name": "providers",

‎shared/artifacts/abi/prime_network.json

+37
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,30 @@
113113
"outputs": [],
114114
"stateMutability": "nonpayable"
115115
},
116+
{
117+
"type": "function",
118+
"name": "calculateMinimumStake",
119+
"inputs": [
120+
{
121+
"name": "provider",
122+
"type": "address",
123+
"internalType": "address"
124+
},
125+
{
126+
"name": "computeUnits",
127+
"type": "uint256",
128+
"internalType": "uint256"
129+
}
130+
],
131+
"outputs": [
132+
{
133+
"name": "",
134+
"type": "uint256",
135+
"internalType": "uint256"
136+
}
137+
],
138+
"stateMutability": "view"
139+
},
116140
{
117141
"type": "function",
118142
"name": "computePool",
@@ -317,6 +341,19 @@
317341
],
318342
"stateMutability": "view"
319343
},
344+
{
345+
"type": "function",
346+
"name": "increaseStake",
347+
"inputs": [
348+
{
349+
"name": "additionalStake",
350+
"type": "uint256",
351+
"internalType": "uint256"
352+
}
353+
],
354+
"outputs": [],
355+
"stateMutability": "nonpayable"
356+
},
320357
{
321358
"type": "function",
322359
"name": "invalidateNode",

‎shared/src/web3/contracts/core/builder.rs

+10-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use crate::web3::{
66
implementations::{
77
ai_token_contract::AIToken, compute_pool_contract::ComputePool,
88
compute_registry_contract::ComputeRegistryContract,
9-
prime_network_contract::PrimeNetworkContract,
9+
prime_network_contract::PrimeNetworkContract, stake_manager::StakeManagerContract,
1010
work_validators::synthetic_data_validator::SyntheticDataWorkValidator,
1111
},
1212
},
@@ -20,6 +20,7 @@ pub struct Contracts {
2020
pub ai_token: AIToken,
2121
pub prime_network: PrimeNetworkContract,
2222
pub compute_pool: ComputePool,
23+
pub stake_manager: Option<StakeManagerContract>,
2324
pub synthetic_data_validator: Option<SyntheticDataWorkValidator>,
2425
}
2526

@@ -29,6 +30,7 @@ pub struct ContractBuilder<'a> {
2930
ai_token: Option<AIToken>,
3031
prime_network: Option<PrimeNetworkContract>,
3132
compute_pool: Option<ComputePool>,
33+
stake_manager: Option<StakeManagerContract>,
3234
synthetic_data_validator: Option<SyntheticDataWorkValidator>,
3335
}
3436

@@ -40,6 +42,7 @@ impl<'a> ContractBuilder<'a> {
4042
ai_token: None,
4143
prime_network: None,
4244
compute_pool: None,
45+
stake_manager: None,
4346
synthetic_data_validator: None,
4447
}
4548
}
@@ -76,6 +79,11 @@ impl<'a> ContractBuilder<'a> {
7679
self
7780
}
7881

82+
pub fn with_stake_manager(mut self) -> Self {
83+
self.stake_manager = Some(StakeManagerContract::new(self.wallet, "stake_manager.json"));
84+
self
85+
}
86+
7987
// TODO: This is not ideal yet - now you have to init all contracts all the time
8088
pub fn build(self) -> Result<Contracts, ContractError> {
8189
// Using custom error ContractError
@@ -101,6 +109,7 @@ impl<'a> ContractBuilder<'a> {
101109
None => return Err(ContractError::Other("PrimeNetwork not initialized".into())), // Custom error handling
102110
},
103111
synthetic_data_validator: self.synthetic_data_validator,
112+
stake_manager: self.stake_manager,
104113
})
105114
}
106115
}

‎shared/src/web3/contracts/core/contract.rs

+1
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ impl Contract {
4545
"synthetic_data_work_validator.json" => {
4646
include_abi!("../../../../artifacts/abi/synthetic_data_work_validator.json")
4747
}
48+
"stake_manager.json" => include_abi!("../../../../artifacts/abi/stake_manager.json"),
4849
_ => panic!("Unknown ABI file: {}", path),
4950
};
5051

‎shared/src/web3/contracts/implementations/compute_registry_contract.rs

+17-1
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use super::{
55
use crate::web3::contracts::helpers::utils::get_selector;
66
use crate::web3::wallet::Wallet;
77
use alloy::dyn_abi::DynSolValue;
8-
use alloy::primitives::Address;
8+
use alloy::primitives::{Address, U256};
99

1010
pub struct ComputeRegistryContract {
1111
instance: Contract,
@@ -41,6 +41,22 @@ impl ComputeRegistryContract {
4141
Ok(provider)
4242
}
4343

44+
pub async fn get_provider_total_compute(
45+
&self,
46+
address: Address,
47+
) -> Result<U256, Box<dyn std::error::Error>> {
48+
let provider_response = self
49+
.instance
50+
.instance()
51+
.function("getProviderTotalCompute", &[address.into()])?
52+
.call()
53+
.await?;
54+
55+
Ok(U256::from(
56+
provider_response.first().unwrap().as_uint().unwrap().0,
57+
))
58+
}
59+
4460
pub async fn get_node(
4561
&self,
4662
#[allow(unused_variables)] provider_address: Address,

‎shared/src/web3/contracts/implementations/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@ pub mod ai_token_contract;
22
pub mod compute_pool_contract;
33
pub mod compute_registry_contract;
44
pub mod prime_network_contract;
5+
pub mod stake_manager;
56
pub mod work_validators;

‎shared/src/web3/contracts/implementations/prime_network_contract.rs

+16
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,22 @@ impl PrimeNetworkContract {
3232
Ok(register_tx)
3333
}
3434

35+
pub async fn stake(
36+
&self,
37+
additional_stake: U256,
38+
) -> Result<FixedBytes<32>, Box<dyn std::error::Error>> {
39+
let stake_tx = self
40+
.instance
41+
.instance()
42+
.function("increaseStake", &[additional_stake.into()])?
43+
.send()
44+
.await?
45+
.watch()
46+
.await?;
47+
48+
Ok(stake_tx)
49+
}
50+
3551
pub async fn add_compute_node(
3652
&self,
3753
node_address: Address,

‎shared/src/web3/contracts/implementations/stake_manager.rs

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
use crate::web3::contracts::constants::addresses::STAKE_MANAGER_ADDRESS;
2+
use crate::web3::contracts::core::contract::Contract;
3+
use crate::web3::wallet::Wallet;
4+
use alloy::primitives::{Address, U256};
5+
6+
pub struct StakeManagerContract {
7+
instance: Contract,
8+
}
9+
10+
impl StakeManagerContract {
11+
pub fn new(wallet: &Wallet, abi_file_path: &str) -> Self {
12+
let instance = Contract::new(STAKE_MANAGER_ADDRESS, wallet, abi_file_path);
13+
Self { instance }
14+
}
15+
16+
pub async fn get_stake_minimum(&self) -> Result<U256, Box<dyn std::error::Error>> {
17+
let result = self
18+
.instance
19+
.instance()
20+
.function("getStakeMinimum", &[])?
21+
.call()
22+
.await?;
23+
24+
let minimum: U256 = result
25+
.into_iter()
26+
.next()
27+
.map(|value| value.as_uint().unwrap_or_default())
28+
.unwrap_or_default()
29+
.0;
30+
Ok(minimum)
31+
}
32+
pub async fn get_stake(&self, staker: Address) -> Result<U256, Box<dyn std::error::Error>> {
33+
let result = self
34+
.instance
35+
.instance()
36+
.function("getStake", &[staker.into()])?
37+
.call()
38+
.await?;
39+
println!("Result: {:?}", result);
40+
41+
Ok(result[0].as_uint().unwrap_or_default().0)
42+
}
43+
44+
pub async fn calculate_stake(
45+
&self,
46+
compute_units: U256,
47+
provider_total_compute: U256,
48+
) -> Result<U256, Box<dyn std::error::Error>> {
49+
let min_stake_per_unit = self.get_stake_minimum().await?;
50+
let total_compute = provider_total_compute + compute_units + U256::from(1);
51+
let required_stake = total_compute * min_stake_per_unit;
52+
Ok(required_stake)
53+
}
54+
}

0 commit comments

Comments
 (0)