improve & automate slash handling (#250)

* improve & automate slash handling

Author: JannikSt

Makefile

@@ -47,7 +47,6 @@ setup:
 	make start-compute-pool
 
 setup-dev-env:
-	make set-min-stake-amount
 	make create-training-domain
 	make create-synth-data-domain
 	make mint-ai-tokens-to-federator

discovery/src/chainsync/sync.rs

@@ -29,6 +29,7 @@ impl ChainSync {
             contracts,
         }
     }
+
     async fn sync_single_node(
         node_store: Arc<NodeStore>,
         contracts: Arc<Contracts>,
@@ -47,16 +48,6 @@ impl ChainSync {
             anyhow::anyhow!("Invalid node address")
         })?;
 
-        // Handle potential errors from async calls
-        let is_blacklisted = contracts
-            .compute_pool
-            .is_node_blacklisted(node.node.compute_pool_id, node_address)
-            .await
-            .map_err(|e| {
-                eprintln!("Error checking if node is blacklisted: {}", e);
-                anyhow::anyhow!("Failed to check blacklist status")
-            })?;
-
         let node_info = contracts
             .compute_registry
             .get_node(provider_address, node_address)
@@ -79,6 +70,16 @@ impl ChainSync {
         n.is_active = is_active;
         n.is_validated = is_validated;
         n.is_provider_whitelisted = provider_info.is_whitelisted;
+
+        // Handle potential errors from async calls
+        let is_blacklisted = contracts
+            .compute_pool
+            .is_node_blacklisted(node.node.compute_pool_id, node_address)
+            .await
+            .map_err(|e| {
+                eprintln!("Error checking if node is blacklisted: {}", e);
+                anyhow::anyhow!("Failed to check blacklist status")
+            })?;
         n.is_blacklisted = is_blacklisted;
         match node_store.update_node(n) {
             Ok(_) => (),

discovery/src/main.rs

@@ -58,6 +58,7 @@ async fn main() -> Result<()> {
             .with_ai_token()
             .with_prime_network()
             .with_compute_pool()
+            .with_stake_manager()
             .build()
             .unwrap(),
     );

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

@@ -6,6 +6,7 @@ use crate::web3::contracts::helpers::utils::get_selector;
 use crate::web3::wallet::Wallet;
 use alloy::dyn_abi::DynSolValue;
 use alloy::primitives::{Address, U256};
+use anyhow::Result;
 
 #[derive(Clone)]
 pub struct ComputeRegistryContract {
@@ -57,12 +58,11 @@ impl ComputeRegistryContract {
             provider_response.first().unwrap().as_uint().unwrap().0,
         ))
     }
-
     pub async fn get_node(
         &self,
         provider_address: Address,
         node_address: Address,
-    ) -> Result<(bool, bool), Box<dyn std::error::Error>> {
+    ) -> anyhow::Result<(bool, bool)> {
         let get_node_selector = get_selector("getNode(address,address)");
 
         let node_response = self
@@ -83,14 +83,14 @@ impl ComputeRegistryContract {
             let node_subkey = node_tuple[1].as_address().unwrap();
 
             if node_provider != provider_address || node_subkey != node_address {
-                return Err("Node does not match provider or subkey".into());
+                return Err(anyhow::anyhow!("Node does not match provider or subkey"));
             }
 
             let active = node_tuple[5].as_bool().unwrap();
             let validated = node_tuple[6].as_bool().unwrap();
             Ok((active, validated))
         } else {
-            Err("Node is not registered".into())
+            Err(anyhow::anyhow!("Node is not registered"))
         }
     }
 }

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

@@ -197,6 +197,7 @@ impl PrimeNetworkContract {
 
         Ok(invalidate_work_tx)
     }
+
     pub async fn get_validator_role(&self) -> Result<Vec<Address>, Error> {
         let hash = keccak256(b"VALIDATOR_ROLE");
         let value = DynSolValue::FixedBytes(hash, 32);

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

@@ -30,6 +30,7 @@ impl StakeManagerContract {
             .0;
         Ok(minimum)
     }
+
     pub async fn get_stake(&self, staker: Address) -> Result<U256, Box<dyn std::error::Error>> {
         let result = self
             .instance

validator/src/main.rs

@@ -2,7 +2,7 @@ pub mod store;
 pub mod validators;
 use actix_web::{web, App, HttpResponse, HttpServer, Responder};
 use alloy::primitives::utils::Unit;
-use alloy::primitives::U256;
+use alloy::primitives::{Address, U256};
 use anyhow::{Context, Result};
 use clap::Parser;
 use log::{debug, LevelFilter};
@@ -245,8 +245,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 };
                 info!(
                     "Synthetic validator has penalty: {} ({})",
-                    penalty,
-                    Unit::ETHER.wei()
+                    penalty, args.validator_penalty
                 );
 
                 Some(SyntheticDataValidator::new(
@@ -358,7 +357,70 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 }
             };
 
-            if let Err(e) = hardware_validator.validate_nodes(nodes).await {
+            // Ensure nodes have enough stake
+            let mut nodes_with_enough_stake = Vec::new();
+            let stake_manager = match contracts.stake_manager.as_ref() {
+                Some(manager) => manager,
+                None => {
+                    error!("Stake manager contract not initialized");
+                    continue;
+                }
+            };
+
+            let mut provider_stake_cache: std::collections::HashMap<String, (U256, U256)> =
+                std::collections::HashMap::new();
+            for node in nodes {
+                let provider_address_str = &node.node.provider_address;
+                let provider_address = match Address::from_str(provider_address_str) {
+                    Ok(address) => address,
+                    Err(e) => {
+                        error!(
+                            "Failed to parse provider address {}: {}",
+                            provider_address_str, e
+                        );
+                        continue;
+                    }
+                };
+
+                let (stake, required_stake) =
+                    if let Some(&cached_info) = provider_stake_cache.get(provider_address_str) {
+                        cached_info
+                    } else {
+                        let stake = stake_manager
+                            .get_stake(provider_address)
+                            .await
+                            .unwrap_or_default();
+                        let total_compute = contracts
+                            .compute_registry
+                            .get_provider_total_compute(provider_address)
+                            .await
+                            .unwrap_or_default();
+                        let required_stake = stake_manager
+                            .calculate_stake(U256::from(0), total_compute)
+                            .await
+                            .unwrap_or_default();
+
+                        provider_stake_cache
+                            .insert(provider_address_str.clone(), (stake, required_stake));
+                        (stake, required_stake)
+                    };
+
+                if stake >= required_stake {
+                    nodes_with_enough_stake.push(node);
+                } else {
+                    info!(
+                        "Node {} has insufficient stake: {} (required: {})",
+                        node.node.id,
+                        stake / Unit::ETHER.wei(),
+                        required_stake / Unit::ETHER.wei()
+                    );
+                }
+            }
+
+            if let Err(e) = hardware_validator
+                .validate_nodes(nodes_with_enough_stake)
+                .await
+            {
                 error!("Error validating nodes: {:#}", e);
             }
         }

validator/src/validators/synthetic_data.rs

@@ -157,6 +157,7 @@ impl SyntheticDataValidator {
         let data = hex::decode(work_key)
             .map_err(|e| Error::msg(format!("Failed to decode hex work key: {}", e)))?;
         info!("Invalidating work: {}", work_key);
+
         match self
             .prime_network
             .invalidate_work(self.pool_id, self.penalty, data)

worker/src/api/routes/invite.rs

@@ -41,7 +41,6 @@ pub async fn invite_node(
     {
         Ok(result) => {
             Console::success(&format!("Successfully joined compute pool: {:?}", result));
-            Console::info("Starting to send heartbeats now.", "");
         }
         Err(err) => {
             Console::error(&format!("Error joining compute pool: {:?}", err));
@@ -61,7 +60,6 @@ pub async fn invite_node(
         )
     };
 
-    println!("Starting heartbeat service with endpoint: {}", endpoint);
     let _ = app_state.heartbeat_service.start(endpoint).await;
 
     HttpResponse::Accepted().json(json!({