different points retrieval logic

Author: erhant

Cargo.lock

@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.5.2"
+version = "0.5.3"
 license = "Apache-2.0"
 readme = "README.md"
 

Cargo.toml

@@ -14,8 +14,13 @@ impl DriaComputeNode {
     /// Runs the main loop of the compute node.
     /// This method is not expected to return until cancellation occurs for the given token.
     pub async fn run(&mut self, cancellation: CancellationToken) {
+        // initialize the points client
+        self.points_client.initialize().await;
+
         /// Duration between refreshing for diagnostic prints.
         const DIAGNOSTIC_REFRESH_INTERVAL_SECS: Duration = Duration::from_secs(45);
+        /// Duration between refreshing for points update.
+        const POINTS_REFRESH_INTERVAL_SECS: Duration = Duration::from_secs(180);
         /// Duration between refreshing the available nodes.
         const RPC_LIVENESS_REFRESH_INTERVAL_SECS: Duration = Duration::from_secs(2 * 60);
         /// Duration between each specs update sent to the RPC.
@@ -28,6 +33,11 @@ impl DriaComputeNode {
             tokio::time::interval(RPC_LIVENESS_REFRESH_INTERVAL_SECS);
         rpc_liveness_refresh_interval.tick().await; // move each one tick
 
+        // tick the first time a bit earlier
+        let mut points_refresh_interval = tokio::time::interval(POINTS_REFRESH_INTERVAL_SECS);
+        points_refresh_interval.tick().await;
+        points_refresh_interval.reset_after(POINTS_REFRESH_INTERVAL_SECS / 12);
+
         // move one tick, and wait at least a third of the diagnostics
         let mut heartbeat_interval = tokio::time::interval(HeartbeatRequester::HEARTBEAT_DEADLINE);
         heartbeat_interval.tick().await;
@@ -68,6 +78,9 @@ impl DriaComputeNode {
                 // check RPC, and get a new one if we are disconnected
                 _ = rpc_liveness_refresh_interval.tick() => self.handle_rpc_liveness_check().await,
 
+                // log points every now and then
+                _ = points_refresh_interval.tick() => self.handle_points_refresh().await,
+
                 // send a heartbeat request to publish liveness info
                 _ = heartbeat_interval.tick() => {
                   if let Err(e) = self.send_heartbeat().await {

compute/src/node/core.rs

@@ -1,7 +1,7 @@
 use colored::Colorize;
 use std::time::Duration;
 
-use crate::{node::rpc::DriaRPC, utils::get_points, DriaComputeNode, DRIA_COMPUTE_NODE_VERSION};
+use crate::{node::rpc::DriaRPC, DriaComputeNode, DRIA_COMPUTE_NODE_VERSION};
 
 /// Number of seconds such that if the last heartbeat ACK is older than this, the node is considered unreachable.
 /// This must be at least greated than the heartbeat interval duration, and the liveness check duration.
@@ -22,17 +22,6 @@ impl DriaComputeNode {
     pub(crate) async fn handle_diagnostic_refresh(&mut self) {
         let mut diagnostics = vec![format!("Diagnostics (v{}):", DRIA_COMPUTE_NODE_VERSION)];
 
-        // print steps
-        if let Ok(steps) = get_points(&self.config.address).await {
-            let earned = steps.score - self.initial_steps;
-            diagnostics.push(format!(
-                "$DRIA Points: {} total, {} earned in this run, within top {}%",
-                steps.score,
-                earned,
-                steps.percentile.unwrap_or("100".to_string())
-            ));
-        }
-
         // completed tasks count is printed as well in debug
         if log::log_enabled!(log::Level::Debug) {
             diagnostics.push(format!(
@@ -137,4 +126,24 @@ impl DriaComputeNode {
             log::debug!("Connection with {} is intact.", self.dria_rpc.peer_id);
         }
     }
+
+    /// Updates the points for the given address.
+    #[inline]
+    pub(crate) async fn handle_points_refresh(&mut self) {
+        // get points from the API
+        match self.points_client.get_points().await {
+            Ok(steps) => {
+                log::info!(
+                    "{}: {} total, {} earned in this run, within top {}%",
+                    "$DRIA Points".purple(),
+                    steps.score,
+                    steps.score - self.points_client.initial,
+                    steps.percentile.unwrap_or("100".to_string())
+                );
+            }
+            Err(err) => {
+                log::error!("Could not get $DRIA points info: {err:?}");
+            }
+        }
+    }
 }

compute/src/node/diagnostic.rs

@@ -9,7 +9,7 @@ use uuid::Uuid;
 
 use crate::{
     config::*,
-    utils::{get_points, SpecCollector},
+    utils::{DriaPointsClient, SpecCollector},
     workers::task::{TaskWorker, TaskWorkerInput, TaskWorkerMetadata, TaskWorkerOutput},
 };
 
@@ -58,8 +58,8 @@ pub struct DriaComputeNode {
     completed_tasks_batch: usize,
     /// Specifications collector.
     spec_collector: SpecCollector,
-    /// Initial steps count.
-    initial_steps: f64,
+    /// Points client.
+    points_client: DriaPointsClient,
 }
 
 impl DriaComputeNode {
@@ -78,7 +78,7 @@ impl DriaComputeNode {
         let keypair = secret_to_keypair(&config.secret_key);
 
         // dial the RPC node
-        let dria_nodes = if let Some(addr) = config.initial_rpc_addr.take() {
+        let dria_rpc = if let Some(addr) = config.initial_rpc_addr.take() {
             log::info!("Using initial RPC address: {}", addr);
             DriaRPC::new(addr, config.network_type).expect("could not get RPC to connect to")
         } else {
@@ -96,7 +96,7 @@ impl DriaComputeNode {
         let (p2p_client, p2p_commander, request_rx) = DriaP2PClient::new(
             keypair,
             config.p2p_listen_addr.clone(),
-            &dria_nodes.addr,
+            &dria_rpc.addr,
             protocol,
         )?;
 
@@ -120,19 +120,15 @@ impl DriaComputeNode {
         };
 
         let model_names = config.workflows.get_model_names();
-
-        let initial_steps = get_points(&config.address)
-            .await
-            .map(|s| s.score)
-            .unwrap_or_default();
+        let points_client = DriaPointsClient::new(&config.address)?;
 
         let spec_collector = SpecCollector::new(model_names.clone(), config.version);
         Ok((
             DriaComputeNode {
                 config,
                 p2p: p2p_commander,
-                dria_rpc: dria_nodes,
-                initial_steps,
+                dria_rpc,
+                points_client,
                 // receivers
                 task_output_rx: publish_rx,
                 reqres_rx: request_rx,

compute/src/node/mod.rs

@@ -3,6 +3,57 @@ use eyre::Context;
 /// Points URL, use with an `address` query parameter.
 const POINTS_API_BASE_URL: &str =
     "https://mainnet.dkn.dria.co/dashboard/supply/v0/leaderboard/steps";
+// TODO: support testnet here?
+
+pub struct DriaPointsClient {
+    pub url: String,
+    client: reqwest::Client,
+    /// The total number of points you have accumulated at the start of the run.
+    pub initial: f64,
+}
+
+impl DriaPointsClient {
+    /// Creates a new `DriaPointsClient` for the given address.
+    pub fn new(address: &str) -> eyre::Result<Self> {
+        const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
+        let url = format!(
+            "{}?address=0x{}",
+            POINTS_API_BASE_URL,
+            address.trim_start_matches("0x")
+        );
+
+        let client = reqwest::Client::builder()
+            .user_agent(USER_AGENT)
+            .build()
+            .wrap_err("could not create Points client")?;
+
+        Ok(Self {
+            url,
+            client,
+            initial: 0.0,
+        })
+    }
+
+    /// Sets the initial points to the current points.
+    ///
+    /// If there is an error, it sets to 0.0.
+    pub async fn initialize(&mut self) {
+        self.initial = self.get_points().await.map(|p| p.score).unwrap_or_default();
+    }
+
+    pub async fn get_points(&self) -> eyre::Result<DriaPoints> {
+        let res = self
+            .client
+            .get(&self.url)
+            .send()
+            .await
+            .wrap_err("could not make request")?;
+        res.json::<DriaPoints>()
+            .await
+            .wrap_err("could not parse response")
+    }
+}
 
 #[derive(Debug, serde::Deserialize)]
 pub struct DriaPoints {
@@ -15,33 +66,15 @@ pub struct DriaPoints {
     pub score: f64,
 }
 
-/// Returns the points for the given address.
-pub async fn get_points(address: &str) -> eyre::Result<DriaPoints> {
-    // the address can have 0x or not, we add it ourselves here
-    let url = format!(
-        "{}?address=0x{}",
-        POINTS_API_BASE_URL,
-        address.trim_start_matches("0x")
-    );
-
-    let res = reqwest::get(&url)
-        .await
-        .wrap_err("could not make request")?;
-    res.json::<DriaPoints>()
-        .await
-        .wrap_err("could not parse response")
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[tokio::test]
     #[ignore = "waiting for API"]
     async fn test_get_points() {
-        let steps = get_points("0xa43536a6032a3907ccf60e8109429ee1047b207c")
-            .await
-            .unwrap();
+        let client = DriaPointsClient::new("0xa43536a6032a3907ccf60e8109429ee1047b207c").unwrap();
+        let steps = client.get_points().await.unwrap();
         assert!(steps.score != 0.0);
     }
 }