available node logic added via api call, update workflows, trait for compute

Author: erhant

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "dkn-compute"
-version = "0.1.3"
+version = "0.1.4"
 edition = "2021"
 license = "Apache-2.0"
 readme = "README.md"
@@ -12,6 +12,7 @@ parking_lot = "0.12.2"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 async-trait = "0.1.81"
+reqwest = "0.12.5"
 
 # utilities
 base64 = "0.22.0"
@@ -34,8 +35,7 @@ sha3 = "0.10.8"
 fastbloom-rs = "0.5.9"
 
 # workflows
-ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows", rev = "274b26e" }
-ollama-rs = "0.2.0"
+ollama-workflows = { git = "https://github.com/andthattoo/ollama-workflows", rev = "25467d2" }
 
 # peer-to-peer
 libp2p = { version = "0.53", features = [

Cargo.toml

@@ -1,6 +1,6 @@
 use std::time::SystemTime;
 
-use ollama_rs::{
+use ollama_workflows::ollama_rs::{
     generation::completion::{request::GenerationRequest, GenerationResponse},
     Ollama,
 };

examples/common/ollama.rs

@@ -1,4 +1,4 @@
-use ollama_rs::Ollama;
+use ollama_workflows::ollama_rs::Ollama;
 
 const DEFAULT_OLLAMA_HOST: &str = "http://127.0.0.1";
 const DEFAULT_OLLAMA_PORT: u16 = 11434;

src/config/ollama.rs

@@ -1,4 +1,4 @@
-use ollama_rs::error::OllamaError;
+use ollama_workflows::ollama_rs::error::OllamaError;
 
 /// Alias for `Result<T, NodeError>`.
 pub type NodeResult<T> = std::result::Result<T, NodeError>;
@@ -96,6 +96,15 @@ impl From<libp2p::gossipsub::SubscriptionError> for NodeError {
     }
 }
 
+impl From<reqwest::Error> for NodeError {
+    fn from(value: reqwest::Error) -> Self {
+        Self {
+            message: value.to_string(),
+            source: "reqwest".to_string(),
+        }
+    }
+}
+
 impl From<libp2p::gossipsub::PublishError> for NodeError {
     fn from(value: libp2p::gossipsub::PublishError) -> Self {
         Self {

src/errors/mod.rs

@@ -1,5 +1,19 @@
+use async_trait::async_trait;
+use libp2p::gossipsub::MessageAcceptance;
+
 mod pingpong;
-pub use pingpong::HandlesPingpong;
+pub use pingpong::PingpongHandler;
 
 mod workflow;
-pub use workflow::HandlesWorkflow;
+pub use workflow::WorkflowHandler;
+
+use crate::{errors::NodeResult, p2p::P2PMessage, DriaComputeNode};
+
+#[async_trait]
+pub trait ComputeHandler {
+    async fn handle_compute(
+        node: &mut DriaComputeNode,
+        message: P2PMessage,
+        result_topic: &str,
+    ) -> NodeResult<MessageAcceptance>;
+}

src/handlers/mod.rs

@@ -1,10 +1,15 @@
 use crate::{
     errors::NodeResult, node::DriaComputeNode, p2p::P2PMessage, utils::get_current_time_nanos,
 };
+use async_trait::async_trait;
 use libp2p::gossipsub::MessageAcceptance;
 use ollama_workflows::{Model, ModelProvider};
 use serde::{Deserialize, Serialize};
 
+use super::ComputeHandler;
+
+pub struct PingpongHandler;
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 struct PingpongPayload {
     uuid: String,
@@ -18,19 +23,10 @@ struct PingpongResponse {
     pub(crate) timestamp: u128,
 }
 
-/// A ping-pong is a message sent by a node to indicate that it is alive.
-/// Compute nodes listen to `pong` topic, and respond to `ping` topic.
-pub trait HandlesPingpong {
-    fn handle_heartbeat(
-        &mut self,
-        message: P2PMessage,
-        result_topic: &str,
-    ) -> NodeResult<MessageAcceptance>;
-}
-
-impl HandlesPingpong for DriaComputeNode {
-    fn handle_heartbeat(
-        &mut self,
+#[async_trait]
+impl ComputeHandler for PingpongHandler {
+    async fn handle_compute(
+        node: &mut DriaComputeNode,
         message: P2PMessage,
         result_topic: &str,
     ) -> NodeResult<MessageAcceptance> {
@@ -53,15 +49,15 @@ impl HandlesPingpong for DriaComputeNode {
         // respond
         let response_body = PingpongResponse {
             uuid: pingpong.uuid.clone(),
-            models: self.config.model_config.models.clone(),
+            models: node.config.model_config.models.clone(),
             timestamp: get_current_time_nanos(),
         };
         let response = P2PMessage::new_signed(
             serde_json::json!(response_body).to_string(),
             result_topic,
-            &self.config.secret_key,
+            &node.config.secret_key,
         );
-        self.publish(response)?;
+        node.publish(response)?;
 
         // accept message, someone else may be included in the filter
         Ok(MessageAcceptance::Accept)

src/handlers/pingpong.rs

@@ -9,6 +9,10 @@ use crate::p2p::P2PMessage;
 use crate::utils::get_current_time_nanos;
 use crate::utils::payload::{TaskRequest, TaskRequestPayload};
 
+use super::ComputeHandler;
+
+pub struct WorkflowHandler;
+
 #[derive(Debug, Deserialize)]
 struct WorkflowPayload {
     /// Workflow object to be parsed.
@@ -23,18 +27,9 @@ struct WorkflowPayload {
 }
 
 #[async_trait]
-pub trait HandlesWorkflow {
-    async fn handle_workflow(
-        &mut self,
-        message: P2PMessage,
-        result_topic: &str,
-    ) -> NodeResult<MessageAcceptance>;
-}
-
-#[async_trait]
-impl HandlesWorkflow for DriaComputeNode {
-    async fn handle_workflow(
-        &mut self,
+impl ComputeHandler for WorkflowHandler {
+    async fn handle_compute(
+        node: &mut DriaComputeNode,
         message: P2PMessage,
         result_topic: &str,
     ) -> NodeResult<MessageAcceptance> {
@@ -55,7 +50,7 @@ impl HandlesWorkflow for DriaComputeNode {
         }
 
         // check task inclusion via the bloom filter
-        if !task.filter.contains(&self.config.address)? {
+        if !task.filter.contains(&node.config.address)? {
             log::info!(
                 "Task {} does not include this node within the filter.",
                 task.task_id
@@ -75,7 +70,7 @@ impl HandlesWorkflow for DriaComputeNode {
         };
 
         // read model / provider from the task
-        let (model_provider, model) = self
+        let (model_provider, model) = node
             .config
             .model_config
             .get_any_matching_model(task.input.model)?;
@@ -85,8 +80,8 @@ impl HandlesWorkflow for DriaComputeNode {
         let executor = if model_provider == ModelProvider::Ollama {
             Executor::new_at(
                 model,
-                &self.config.ollama_config.host,
-                self.config.ollama_config.port,
+                &node.config.ollama_config.host,
+                node.config.ollama_config.port,
             )
         } else {
             Executor::new(model)
@@ -98,7 +93,7 @@ impl HandlesWorkflow for DriaComputeNode {
             .map(|prompt| Entry::try_value_or_str(&prompt));
         let result: Option<String>;
         tokio::select! {
-            _ = self.cancellation.cancelled() => {
+            _ = node.cancellation.cancelled() => {
                 log::info!("Received cancellation, quitting all tasks.");
                 return Ok(MessageAcceptance::Accept)
             },
@@ -113,7 +108,7 @@ impl HandlesWorkflow for DriaComputeNode {
         let result = result.ok_or::<String>(format!("No result for task {}", task.task_id))?;
 
         // publish the result
-        self.send_result(result_topic, &task.public_key, &task.task_id, result)?;
+        node.send_result(result_topic, &task.public_key, &task.task_id, result)?;
 
         // accept message, someone else may be included in the filter
         Ok(MessageAcceptance::Accept)

src/handlers/workflow.rs

@@ -18,7 +18,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     config.check_services().await?;
 
     // launch the node
-    let mut node = DriaComputeNode::new(config, CancellationToken::new())?;
+    let mut node = DriaComputeNode::new(config, CancellationToken::new()).await?;
     node.launch().await?;
 
     Ok(())