Merge pull request #133 from PrimeIntellect-ai/fix/optimize-validator

fix/optimize-validator

Author: mattdf

Cargo.lock

@@ -8,6 +8,7 @@ actix-web = "4.9.0"
 alloy = { version = "0.9.2", features = ["full"] }
 anyhow = "1.0.95"
 clap = { version = "4.5.26", features = ["derive"] }
+directories = "6.0.0"
 env_logger = "0.11.6"
 log = "0.4.25"
 nalgebra = "0.33.2"
@@ -17,4 +18,8 @@ serde = "1.0.217"
 serde_json = "1.0.135"
 shared = { path = "../shared" }
 tokio = "1.43.0"
+toml = "0.8.20"
 url = "2.5.4"
+
+[dev-dependencies]
+tempfile = "=3.14.0"

validator/Cargo.toml

@@ -93,7 +93,7 @@ fn main() {
 
     let pool_id = args.pool_id.clone();
     let mut synthetic_validator = match contracts.synthetic_data_validator.clone() {
-        Some(validator) => SyntheticDataValidator::new(pool_id.unwrap(), validator),
+        Some(validator) => SyntheticDataValidator::new(None, pool_id.unwrap(), validator),
         None => {
             error!("Synthetic data validator not found");
             std::process::exit(1);
@@ -178,7 +178,7 @@ fn main() {
 
 #[cfg(test)]
 mod tests {
-    
+
     use actix_web::{test, App};
     use actix_web::{
         web::{self, post},

validator/src/main.rs

@@ -1,14 +1,37 @@
 use alloy::primitives::U256;
 use anyhow::{Context, Result};
+use directories::ProjectDirs;
+use log::debug;
 use log::{error, info};
+use serde::{Deserialize, Serialize};
+use std::fs;
+use std::path::Path;
+use std::path::PathBuf;
+use toml;
 
 use crate::validators::Validator;
 use shared::web3::contracts::implementations::work_validators::synthetic_data_validator::SyntheticDataWorkValidator;
 
+fn get_default_state_dir() -> Option<String> {
+    ProjectDirs::from("com", "prime", "validator")
+        .map(|proj_dirs| proj_dirs.data_local_dir().to_string_lossy().into_owned())
+}
+
+fn state_filename(pool_id: &str) -> String {
+    format!("work_state_{}.toml", pool_id)
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct PersistedWorkState {
+    pool_id: U256,
+    last_validation_timestamp: U256,
+}
+
 pub struct SyntheticDataValidator {
     pool_id: U256,
     validator: SyntheticDataWorkValidator,
     last_validation_timestamp: U256,
+    state_dir: Option<PathBuf>,
 }
 
 impl Validator for SyntheticDataValidator {
@@ -20,16 +43,84 @@ impl Validator for SyntheticDataValidator {
 }
 
 impl SyntheticDataValidator {
-    pub fn new(pool_id_str: String, validator: SyntheticDataWorkValidator) -> Self {
+    pub fn new(
+        state_dir: Option<String>,
+        pool_id_str: String,
+        validator: SyntheticDataWorkValidator,
+    ) -> Self {
         let pool_id = pool_id_str.parse::<U256>().expect("Invalid pool ID");
+        let default_state_dir = get_default_state_dir();
+        debug!("Default state dir: {:?}", default_state_dir);
+        let state_path = state_dir
+            .map(PathBuf::from)
+            .or_else(|| default_state_dir.map(PathBuf::from));
+        debug!("State path: {:?}", state_path);
+        let mut last_validation_timestamp: Option<_> = None;
+
+        // Try to load state, log info if creating new file
+        if let Some(path) = &state_path {
+            let state_file = path.join(state_filename(&pool_id.to_string()));
+            if !state_file.exists() {
+                debug!(
+                    "No state file found at {:?}, will create on first state change",
+                    state_file
+                );
+            } else if let Ok(Some(loaded_state)) =
+                SyntheticDataValidator::load_state(path, &pool_id.to_string())
+            {
+                debug!("Loaded previous state from {:?}", state_file);
+                last_validation_timestamp = Some(loaded_state.last_validation_timestamp);
+            } else {
+                debug!("Failed to load state from {:?}", state_file);
+            }
+        }
+
+        // if no last time, set it to 24 hours ago, as nothing before that can be invalidated
+        if last_validation_timestamp.is_none() {
+            last_validation_timestamp = Some(U256::from(
+                std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .expect("Failed to get current timestamp")
+                    .as_secs()
+                    .saturating_sub(24 * 60 * 60),
+            ));
+        }
 
         Self {
             pool_id,
             validator,
-            last_validation_timestamp: U256::from(0),
+            last_validation_timestamp: last_validation_timestamp.unwrap(),
+            state_dir: state_path.clone(),
         }
     }
 
+    fn save_state(&self) -> Result<()> {
+        // Get values without block_on
+        let state = PersistedWorkState {
+            pool_id: self.pool_id,
+            last_validation_timestamp: self.last_validation_timestamp,
+        };
+
+        if let Some(ref state_dir) = self.state_dir {
+            fs::create_dir_all(state_dir)?;
+            let state_path = state_dir.join(state_filename(&self.pool_id.to_string()));
+            let toml = toml::to_string_pretty(&state)?;
+            fs::write(&state_path, toml)?;
+            debug!("Saved state to {:?}", state_path);
+        }
+        Ok(())
+    }
+
+    fn load_state(state_dir: &Path, pool_id: &str) -> Result<Option<PersistedWorkState>> {
+        let state_path = state_dir.join(state_filename(pool_id));
+        if state_path.exists() {
+            let contents = fs::read_to_string(state_path)?;
+            let state: PersistedWorkState = toml::from_str(&contents)?;
+            return Ok(Some(state));
+        }
+        Ok(None)
+    }
+
     pub async fn validate_work(&mut self) -> Result<()> {
         info!("Validating work for pool ID: {:?}", self.pool_id);
 
@@ -70,6 +161,8 @@ impl SyntheticDataValidator {
                 .as_secs(),
         );
 
+        self.save_state()?;
+
         Ok(())
     }
 }

validator/src/validators/synthetic_data.rs

@@ -94,10 +94,10 @@ impl HeartbeatService {
                         match Self::send_heartbeat(&client, state.get_endpoint().await, wallet_clone.clone(), docker_service.clone(), metrics_store.clone()).await {
                             Ok(_) => {
                                 state.update_last_heartbeat().await;
-                                Console::success("Synced with orchestrator"); // Updated message to reflect sync
+                                log::info!("Synced with orchestrator"); // Updated message to reflect sync
                             }
                             Err(e) => {
-                                Console::error(&format!("Failed to sync with orchestrator: {:?}", e)); // Updated error message
+                                log::error!("{}", &format!("Failed to sync with orchestrator: {:?}", e)); // Updated error message
                             }
                         }
                     }