monitor as subcrate

Author: erhant

Cargo.lock

@@ -1,9 +1,6 @@
 [workspace]
 resolver = "2"
-members = ["compute", "p2p", "workflows", "utils"]
-# compute node is the default member, until Oracle comes in
-# then, a Launcher will be the default member
-default-members = ["compute"]
+members = ["compute", "p2p", "workflows", "utils", "monitor"]
 
 [workspace.package]
 edition = "2021"

Cargo.toml

@@ -7,31 +7,38 @@ endif
 ###############################################################################
 .PHONY: launch #       | Run with INFO logs in release mode
 launch:
-		RUST_LOG=none,dkn_compute=info,dkn_workflows=info,dkn_p2p=info cargo run --release
+		RUST_LOG=none,dkn_compute=info,dkn_workflows=info,dkn_p2p=info \
+		cargo run --release --bin dkn-compute
 
 .PHONY: run #          | Run with INFO logs
 run:
-		cargo run
+		cargo run --bin dkn-compute
+
+.PHONY: monitor #      | Run monitor node with INFO logs
+monitor:
+		cargo run --bin dkn-monitor
 
 .PHONY: debug #        | Run with DEBUG logs with INFO log-level workflows
 debug:
-		RUST_LOG=warn,dkn_compute=debug,dkn_workflows=debug,dkn_p2p=debug,ollama_workflows=info cargo run
+		RUST_LOG=warn,dkn_compute=debug,dkn_workflows=debug,dkn_p2p=debug,ollama_workflows=info \
+		cargo run --bin dkn-compute
 
 .PHONY: trace #        | Run with TRACE logs
 trace:
-		RUST_LOG=warn,dkn_compute=trace,libp2p=debug cargo run
+		RUST_LOG=warn,dkn_compute=trace,libp2p=debug \
+		cargo run --bin dkn-compute
 
 .PHONY: build #        | Build
 build:
 		cargo build --workspace
 
 .PHONY: profile-cpu #  | Profile CPU usage with flamegraph
 profile-cpu:
-	  DKN_EXIT_TIMEOUT=120 cargo flamegraph --root --profile=profiling
+	  DKN_EXIT_TIMEOUT=120 cargo flamegraph --root --profile=profiling --bin dkn-compute
 
 .PHONY: profile-mem #  | Profile memory usage with instruments
 profile-mem:
-	  DKN_EXIT_TIMEOUT=120 cargo instruments --profile=profiling -t Allocations
+	  DKN_EXIT_TIMEOUT=120 cargo instruments --profile=profiling -t Allocations --bin dkn-compute
 
 .PHONY: ollama-versions
 ollama-versions:

Makefile

@@ -1,5 +1,5 @@
 mod pingpong;
-pub use pingpong::PingpongHandler;
+pub use pingpong::*;
 
 mod workflow;
-pub use workflow::{WorkflowHandler, WorkflowPayload};
+pub use workflow::*;

compute/src/handlers/mod.rs

@@ -1,10 +1,9 @@
-pub(crate) mod config;
-pub(crate) mod handlers;
-pub(crate) mod monitor;
-pub(crate) mod node;
-pub(crate) mod payloads;
-pub(crate) mod utils;
-pub(crate) mod workers;
+pub mod config;
+pub mod handlers;
+pub mod node;
+pub mod payloads;
+pub mod utils;
+pub mod workers;
 
 /// Crate version of the compute node.
 /// This value is attached within the published messages.
@@ -14,5 +13,3 @@ pub use utils::refresh_dria_nodes;
 
 pub use config::DriaComputeNodeConfig;
 pub use node::DriaComputeNode;
-
-pub use monitor::DriaMonitorNode;

compute/src/lib.rs

@@ -8,11 +8,11 @@ pub struct TaskRequestPayload<T> {
     /// The unique identifier of the task.
     pub task_id: String,
     /// The deadline of the task in nanoseconds.
-    pub(crate) deadline: u128,
+    pub deadline: u128,
     /// The input to the compute function.
-    pub(crate) input: T,
+    pub input: T,
     /// The Bloom filter of the task.
-    pub(crate) filter: TaskFilter,
+    pub filter: TaskFilter,
     /// The public key of the requester, in hexadecimals.
-    pub(crate) public_key: String,
+    pub public_key: String,
 }

compute/src/payloads/request.rs

@@ -12,22 +12,22 @@ use serde::{Deserialize, Serialize};
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct DriaMessage {
     /// Base64 encoded payload, stores the main result.
-    pub(crate) payload: String,
+    pub payload: String,
     /// The topic of the message, derived from `TopicHash`
     ///
     /// NOTE: This can be obtained via `TopicHash` in GossipSub
-    pub(crate) topic: String,
+    pub topic: String,
     /// The version of the Dria Compute Node
     ///
     /// NOTE: This can be obtained via Identify protocol version
-    pub(crate) version: String,
+    pub version: String,
     /// Identity protocol string of the Dria Compute Node
     #[serde(default)]
-    pub(crate) identity: String,
+    pub identity: String,
     /// The timestamp of the message, in nanoseconds
     ///
     /// NOTE: This can be obtained via `DataTransform` in GossipSub
-    pub(crate) timestamp: u128,
+    pub timestamp: u128,
 }
 
 /// 65-byte signature as hex characters take up 130 characters.
@@ -79,7 +79,7 @@ impl DriaMessage {
     }
 
     /// Decodes and parses the base64 payload into JSON for the provided type `T`.
-    pub(crate) fn parse_payload<T: for<'a> Deserialize<'a>>(&self, signed: bool) -> Result<T> {
+    pub fn parse_payload<T: for<'a> Deserialize<'a>>(&self, signed: bool) -> Result<T> {
         let payload = self.decode_payload()?;
 
         let body = if signed {

compute/src/utils/message.rs

@@ -0,0 +1,35 @@
+[package]
+name = "dkn-monitor"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+readme = "README.md"
+authors = ["Erhan Tezcan <[email protected]>"]
+
+[dependencies]
+# async stuff
+tokio-util.workspace = true
+tokio.workspace = true
+async-trait.workspace = true
+
+# serialize & deserialize
+serde.workspace = true
+serde_json.workspace = true
+
+# http & networking
+reqwest.workspace = true
+
+# utilities
+dotenvy.workspace = true
+
+# logging & errors
+env_logger.workspace = true
+log.workspace = true
+eyre.workspace = true
+
+libsecp256k1 = "0.7.1"
+
+# dria subcrates
+dkn-p2p = { path = "../p2p" }
+dkn-utils = { path = "../utils" }
+dkn-compute = { path = "../compute" }

monitor/Cargo.toml

@@ -1,33 +1,38 @@
-use dkn_compute::{refresh_dria_nodes, DriaMonitorNode};
+use dkn_compute::refresh_dria_nodes;
 use dkn_p2p::{
     libp2p_identity::Keypair, DriaNetworkType, DriaNodes, DriaP2PClient, DriaP2PProtocol,
 };
 use tokio_util::sync::CancellationToken;
 
+mod node;
+use node::DriaMonitorNode;
+
 #[tokio::main]
 async fn main() -> eyre::Result<()> {
     dotenvy::dotenv().expect("could not load .env");
 
     env_logger::builder()
         .filter(None, log::LevelFilter::Off)
-        .filter_module("dkn_p2p", log::LevelFilter::Info)
+        .filter_module("dkn_p2p", log::LevelFilter::Warn)
         .filter_module("dkn_compute", log::LevelFilter::Info)
-        .filter_module("monitor", log::LevelFilter::Info)
+        .filter_module("dkn_monitor", log::LevelFilter::Info)
         .parse_default_env() // reads RUST_LOG variable
         .init();
 
-    log::info!("Starting Dria Task Monitor");
-
-    let network = DriaNetworkType::Pro;
+    let network = std::env::var("DKN_NETWORK")
+        .map(|s| DriaNetworkType::from(s.as_str()))
+        .unwrap_or(DriaNetworkType::Pro);
     let mut nodes = DriaNodes::new(network);
     refresh_dria_nodes(&mut nodes).await?;
 
     // setup p2p client
+    let listen_addr = "/ip4/0.0.0.0/tcp/4069".parse()?;
+    log::info!("Listen Address: {}", listen_addr);
     let keypair = Keypair::generate_secp256k1();
     log::info!("PeerID: {}", keypair.public().to_peer_id());
     let (client, commander, msg_rx) = DriaP2PClient::new(
-        Keypair::generate_secp256k1(),
-        "/ip4/0.0.0.0/tcp/4069".parse()?,
+        keypair,
+        listen_addr,
         nodes.bootstrap_nodes.into_iter(),
         nodes.relay_nodes.into_iter(),
         nodes.rpc_nodes.into_iter(),
@@ -54,6 +59,11 @@ async fn main() -> eyre::Result<()> {
     });
 
     // create monitor node
+    log::info!(
+        "Monitoring {} network (protocol: {}).",
+        network,
+        network.protocol_name()
+    );
     let mut monitor = DriaMonitorNode::new(commander, msg_rx);
 
     // setup monitor

compute/src/bin/monitor.rs renamed to monitor/src/main.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use crate::{
+use dkn_compute::{
     handlers::{WorkflowHandler, WorkflowPayload},
     payloads::{TaskRequestPayload, TaskResponsePayload},
     utils::DriaMessage,
@@ -56,6 +56,9 @@ impl DriaMonitorNode {
         self.p2p.shutdown().await?;
         self.msg_rx.close();
 
+        // print tasks one final time
+        self.handle_task_print();
+
         Ok(())
     }
 
@@ -65,16 +68,16 @@ impl DriaMonitorNode {
 
         loop {
             tokio::select! {
+                // handle gossipsub message
                 message = self.msg_rx.recv() => match message {
                     Some(message) => match self.handle_message(message).await {
                         Ok(_) => {}
                         Err(e) => log::error!("Error handling message: {:?}", e),
                     }
-                    None => break, // channel closed
+                    None => break, // channel closed, we can return now
                 },
-                _ = task_print_interval.tick() => {
-                    log::info!("Current seen tasks: {:#?}", self.tasks.keys().collect::<Vec<_>>());
-                }
+                // print task counts
+                _ = task_print_interval.tick() => self.handle_task_print(),
                 _ = token.cancelled() => break,
             }
         }
@@ -116,4 +119,26 @@ impl DriaMonitorNode {
         }
         Ok(())
     }
+
+    fn handle_task_print(&self) {
+        let seen_task_ids = self.tasks.keys().collect::<Vec<_>>();
+        let seen_result_ids = self.results.keys().collect::<Vec<_>>();
+
+        // print the tasks that have not been responded to
+        let pending_tasks = seen_task_ids
+            .iter()
+            .filter(|id| !seen_result_ids.contains(*id))
+            .map(|id| self.tasks.get(*id).unwrap())
+            .collect::<Vec<_>>();
+
+        log::info!(
+            "Pending tasks ({} / {}): {:#?}",
+            pending_tasks.len(),
+            self.tasks.len(),
+            pending_tasks
+                .iter()
+                .map(|t| t.task_id.clone())
+                .collect::<Vec<_>>()
+        );
+    }
 }