add task report to heartbeat, doc fixes

Author: erhant

Makefile

@@ -1,6 +1,6 @@
 # load .env
 ifneq (,$(wildcard ./.env))
-		include .env
+		include ./.env
 		export
 endif
 

compute/src/handlers/pingpong.rs

@@ -22,6 +22,8 @@ struct PingpongResponse {
     pub(crate) uuid: String,
     pub(crate) models: Vec<(ModelProvider, Model)>,
     pub(crate) timestamp: u128,
+    /// Number of tasks in the channel currently, `single` and `batch`.
+    pub(crate) tasks: [usize; 2],
 }
 
 impl PingpongHandler {
@@ -64,6 +66,7 @@ impl PingpongHandler {
             uuid: pingpong.uuid.clone(),
             models: node.config.workflows.models.clone(),
             timestamp: get_current_time_nanos(),
+            tasks: node.task_count(),
         };
 
         // publish message

compute/src/node.rs

@@ -122,6 +122,14 @@ impl DriaComputeNode {
         Ok(())
     }
 
+    /// Returns the task count within the channels, `single` and `batch`.
+    pub fn task_count(&self) -> [usize; 2] {
+        [
+            self.workflow_single_tx.max_capacity() - self.workflow_single_tx.capacity(),
+            self.workflow_batch_tx.max_capacity() - self.workflow_batch_tx.capacity(),
+        ]
+    }
+
     /// Publishes a given message to the network w.r.t the topic of it.
     ///
     /// Internally, identity is attached to the the message which is then JSON serialized to bytes

p2p/README.md

@@ -12,7 +12,92 @@ dkn-p2p = { git = "https://github.com/firstbatchxyz/dkn-compute-node" }
 
 ## Usage
 
-The P2P client is expected to be run within a separate thread, and it has two types of interactions:
+The peer-to-peer client, when created, returns 3 things:
+
+- `client`: the actual peer-to-peer client that should be run in a **separate thread**.
+- `commander`: a small client that exposes peer-to-peer functions with oneshot channels, so that we can communicate with the client in another thread.
+- `channel`: a message channel receiver, it is expected to handle GossipSub messages that are handled & sent by the client.
+
+### Client
+
+Here is an example where we create the said entities:
+
+```rs
+use dkn_p2p::{DriaP2PClient, DriaP2PProtocol};
+
+// your wallet, or something random maybe
+let keypair = Keypair::generate_secp256k1();
+
+// your listen address
+let addr = Multiaddr::from_str("/ip4/0.0.0.0/tcp/4001")?;
+
+// static bootstrap & relay & rpc addresses
+let bootstraps = vec![Multiaddr::from_str(
+    "some-multiaddrs-here"
+)?];
+let relays = vec![Multiaddr::from_str(
+    "some-multiaddrs-here"
+)?];
+let rpcs = vec![Multiaddr::from_str(
+    "some-multiaddrs-here"
+)?];
+
+let protocol = "0.2";
+
+// `new` returns 3 things:
+// - p2p client itself, to be given to a thread
+// - p2p commander, a small client to be able to speak with the p2p in another thread
+// - `msg_rx`, the channel to listen for gossipsub messages
+let (client, mut commander, mut msg_rx) = DriaP2PClient::new(
+  keypair,
+  addr,
+  bootstraps,
+  relays,
+  rpc,
+  protocol
+)?;
+```
+
+Now, you can give the peer-to-peer client to a thread and store its handle:
+
+```rs
+let task_handle = tokio::spawn(async move { client.run().await });
+```
+
+This task handle should be `await`'ed at the end of the program to ensure thread has exited correctly.
+
+### Commander
+
+You can communicate with this thread using the `commander` entity. For example, here is how one would subscribe to a topic:
+
+```rs
+commander
+  .subscribe("your-topic")
+  .await
+  .expect("could not subscribe");
+```
+
+### Channel
+
+The message channel should be handled with `recv` (or `recv_many` to process in batches) to process the GossipSub messages.
+
+```rs
+loop {
+    match msg_rx.recv().await {
+        Some(msg) => {
+            todo!("handle stuff")
+        }
+        None => {
+            todo!("channel closed");
+            break
+        }
+    }
+}
+```
+
+### Interactions
+
+Here is how the whole thing works in a bit more detail:
 
 - **Events**: When a message is received within the Swarm event handler, it is returned via a `mpsc` channel. Here, the p2p is `Sender` and your application must be the `Receiver`. The client handles many events, and only sends GossipSub message receipts via this channel so that the application can handle them however they would like.
 
@@ -54,35 +139,3 @@ sequenceDiagram
   P ->> C: o_tx.send(output)
   deactivate P
 ```
-
-<!--
-
-FIXME: REMOVE COMMENTS
-
-You can create the client as follows:
-
-```rs
-use dkn_p2p::DriaP2PClient;
-
-// your wallet, or something random maybe
-let keypair = Keypair::generate_secp256k1();
-
-// your listen address
-let addr = Multiaddr::from_str("/ip4/0.0.0.0/tcp/4001")?;
-
-// static bootstrap & relay addresses
-let bootstraps = vec![Multiaddr::from_str(
-    "some-multiaddrs-here"
-)?];
-let relays = vec![Multiaddr::from_str(
-    "some-multiaddrs-here"
-)?];
-
-// protocol version number, usually derived as `{major}.{minor}`
-let version = "0.2";
-
-// create the client!
-let mut client = DriaP2PClient::new(keypair, addr, &bootstraps, &relays, "0.2")?;
-```
-
-Then, you can use its underlying functions, such as `subscribe`, `process_events` and `unsubscribe`. In particular, `process_events` handles all p2p events and returns a GossipSub message when it is received. -->

p2p/tests/listen_test.rs

@@ -29,7 +29,7 @@ async fn test_listen_topic_once() -> Result<()> {
     .expect("could not create p2p client");
 
     // spawn task
-    let p2p_task = tokio::spawn(async move { client.run().await });
+    let task_handle = tokio::spawn(async move { client.run().await });
 
     // subscribe to the given topic
     commander
@@ -61,7 +61,7 @@ async fn test_listen_topic_once() -> Result<()> {
     msg_rx.close();
 
     log::info!("Waiting for p2p task to finish...");
-    p2p_task.await?;
+    task_handle.await?;
 
     log::info!("Done!");
     Ok(())