Merge pull request #187 from firstbatchxyz/erhant/better-pull-and-env

feat: add custom env option & better pulls

Author: erhant

Cargo.lock

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

Cargo.toml

@@ -5,39 +5,19 @@ ifneq (,$(wildcard ./.env))
 endif
 
 ###############################################################################
-.PHONY: launch #       | Run with INFO logs in release mode
-launch:
-		cargo run --release --bin dkn-compute
-
-.PHONY: run #          | Run with INFO logs
-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 --bin dkn-compute
 
-.PHONY: trace #        | Run with TRACE logs
-trace:
-		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 --bin dkn-compute
-
-.PHONY: profile-mem #  | Profile memory usage with instruments
-profile-mem:
-	  DKN_EXIT_TIMEOUT=120 cargo instruments --profile=profiling -t Allocations --bin dkn-compute
+.PHONY: trace #        | Run with TRACE logs
+trace:
+		RUST_LOG=warn,dkn_compute=trace,libp2p=debug \
+		cargo run --bin dkn-compute
 
 .PHONY: ollama-versions
 ollama-versions:

Makefile

@@ -65,11 +65,13 @@ If you would like to run the node from source (which is really handy during deve
 make help
 ```
 
-You will need OpenSSL installed as well, see shorthand commands [here](https://github.com/sfackler/rust-openssl/issues/855#issuecomment-450057552). While running Ollama elsewhere (if you are using it) or with an OpenAI API key provided, you can run the compute node with:
+You will need OpenSSL installed, see shorthand commands [here](https://github.com/sfackler/rust-openssl/issues/855#issuecomment-450057552).
 
 ```sh
-make run      # info-level logs
-make debug    # debug-level logs
+cargo run
+
+# specify custom .env file
+DKN_COMPUTE_ENV=./path/to/.env cargo run
 ```
 
 If you have a valid `.env` file, you can run the latest Docker image via compose as well:
@@ -112,6 +114,32 @@ make lint   # clippy
 make format # rustfmt
 ```
 
+### Profiling
+
+We have scripts to profile both CPU and Memory usage. A special build is created for profiling, via a custom `profiling` feature, such that the output inherits `release` mode but also has debug symbols.
+
+Furthermore, the profiling build will exit automatically after a certain time, as if CTRL+C has been pressed. This is needed by the memory profiling tool in particular.
+
+**CPU Profiling**: To create a [flamegraph](https://crates.io/crates/flamegraph) of the application, the command below will create a profiling build that inherits `release` mode, except with debug information:
+
+```sh
+DKN_EXIT_TIMEOUT=120 cargo flamegraph --root --profile=profiling --bin dkn-compute
+```
+
+> [!NOTE]
+>
+> CPU profiling may require super-user access.
+
+**Memory Profiling**: To profile memory usage, we make use of [cargo-instruments](https://crates.io/crates/cargo-instruments):
+
+```sh
+DKN_EXIT_TIMEOUT=120 cargo instruments --profile=profiling -t Allocations --bin dkn-compute
+```
+
+> [!TIP]
+>
+> You can adjust the profiling duration via the `DKN_EXIT_TIMEOUT` variable, which takes a number of seconds until termination.
+
 ## License
 
 This project is licensed under the [Apache License 2.0](https://opensource.org/license/Apache-2.0).

README.md

@@ -65,8 +65,9 @@ impl PingpongHandler {
 
         log::info!("Received a {} for: {}", "ping".blue(), pingpong.uuid);
 
-        // record ping moment
+        // record ping
         node.last_pinged_at = Instant::now();
+        node.num_pings += 1;
 
         // respond
         let response_body = PingpongResponse {

compute/src/gossipsub/pingpong.rs

@@ -7,7 +7,9 @@ use workers::task::TaskWorker;
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    let dotenv_result = dotenvy::dotenv();
+    // load a particular environment file specified by DKN_COMPUTE_ENV, or `.env` by default
+    let env_path = env::var("DKN_COMPUTE_ENV").unwrap_or_else(|_| ".env".to_string());
+    let dotenv_result = dotenvy::from_path(&env_path);
 
     env_logger::builder()
         .format_timestamp(Some(env_logger::TimestampPrecision::Millis))
@@ -33,8 +35,8 @@ async fn main() -> Result<()> {
 
     // log about env usage
     match dotenv_result {
-        Ok(path) => log::info!("Loaded .env file at: {}", path.display()),
-        Err(e) => log::warn!("Could not load .env file: {}", e),
+        Ok(_) => log::info!("Loaded environment file from {}", env_path),
+        Err(e) => log::warn!("Could not load environment file from {}: {}", env_path, e),
     }
 
     // task tracker for multiple threads

compute/src/main.rs

@@ -22,10 +22,6 @@ impl DriaComputeNode {
     pub(crate) async fn handle_diagnostic_refresh(&self) {
         let mut diagnostics = vec![format!("Diagnostics (v{}):", DRIA_COMPUTE_NODE_VERSION)];
 
-        // if we have not received pings for a while, we are considered offline
-        let is_offline =
-            self.last_pinged_at < Instant::now() - Duration::from_secs(PING_LIVENESS_SECS);
-
         // print peer counts
         match self.p2p.peer_counts().await {
             Ok((mesh, all)) => diagnostics.push(format!(
@@ -74,14 +70,27 @@ impl DriaComputeNode {
         ));
 
         // add network status as well
-        diagnostics.push(format!(
-            "Node Status: {}",
-            if is_offline {
-                "OFFLINE".bold().red()
-            } else {
-                "ONLINE".bold().green()
-            }
-        ));
+        // if we have not received pings for a while, we are considered offline
+        let is_offline = Instant::now().duration_since(self.last_pinged_at)
+            > Duration::from_secs(PING_LIVENESS_SECS);
+        if self.num_pings == 0 {
+            // if we didnt have any pings, we might still be connecting
+            diagnostics.push(format!("Node Status: {}", "CONNECTING".yellow()));
+        } else {
+            diagnostics.push(format!(
+                "Node Status: {}",
+                if is_offline {
+                    "OFFLINE".red()
+                } else {
+                    "ONLINE".green()
+                }
+            ));
+        }
+
+        // add pings per second
+        let elapsed = Instant::now().duration_since(self.started_at).as_secs_f64();
+        let pings_per_second = self.num_pings as f64 / elapsed; // elapsed is always > 0
+        diagnostics.push(format!("Pings/sec: {:.3}", pings_per_second));
 
         log::info!("{}", diagnostics.join("\n  "));
 

compute/src/node/diagnostic.rs

@@ -33,7 +33,11 @@ pub struct DriaComputeNode {
     pub p2p: DriaP2PCommander,
     /// The last time the node was pinged by the network.
     /// If this is too much, we can say that the node is not reachable by RPC.
-    pub last_pinged_at: Instant,
+    pub(crate) last_pinged_at: Instant,
+    /// Number of pings received.
+    pub(crate) num_pings: u64,
+    /// The time the node was started.
+    pub(crate) started_at: Instant,
     /// Gossipsub message receiver, used by peer-to-peer client in a separate thread.
     ///
     /// It will publish messages sent to this channel to the network.
@@ -143,6 +147,8 @@ impl DriaComputeNode {
                 initial_steps,
                 spec_collector: SpecCollector::new(model_names),
                 last_pinged_at: Instant::now(),
+                num_pings: 0,
+                started_at: Instant::now(),
             },
             p2p_client,
             task_batch_worker,