added service check workflow tests & identify versioning, todo kad versioning

Author: erhant

Cargo.lock

@@ -68,6 +68,7 @@ libp2p = { git = "https://github.com/anilaltuner/rust-libp2p.git", rev = "be2ed5
 libp2p-identity = { version = "0.2.9", features = ["secp256k1"] }
 tracing = { version = "0.1.40" }
 tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+semver = "1.0.23"
 
 
 [dev-dependencies]

Cargo.toml

@@ -28,9 +28,7 @@
 
 ## About
 
-A **Dria Compute Node** is a unit of computation within the Dria Knowledge Network. It's purpose is to process tasks given by the **Dria Admin Node**, and receive rewards for providing correct results.
-
-To get started, [setup](#setup) your envrionment and then see [usage](#usage) to run the node.
+A **Dria Compute Node** is a unit of computation within the Dria Knowledge Network. It's purpose is to process tasks given by the **Dria Admin Node**. To get started, [setup](#setup) your envrionment and then see [usage](#usage) to run the node.
 
 ### Tasks
 

README.md

@@ -9,7 +9,7 @@ use ollama::OllamaConfig;
 use ollama_workflows::ModelProvider;
 use openai::OpenAIConfig;
 
-use std::env;
+use std::{env, time::Duration};
 
 #[derive(Debug, Clone)]
 pub struct DriaComputeNodeConfig {
@@ -108,31 +108,49 @@ impl DriaComputeNodeConfig {
     /// Check if the required compute services are running, e.g. if Ollama
     /// is detected as a provider for the chosen models, it will check that
     /// Ollama is running.
-    pub async fn check_services(&self) -> Result<(), String> {
+    pub async fn check_services(&mut self) -> Result<(), String> {
         log::info!("Checking configured services.");
         let unique_providers = self.model_config.get_providers();
 
+        let mut good_models = Vec::new();
+
         // if Ollama is a provider, check that it is running & Ollama models are pulled (or pull them)
         if unique_providers.contains(&ModelProvider::Ollama) {
             let ollama_models = self
                 .model_config
                 .get_models_for_provider(ModelProvider::Ollama);
-            self.ollama_config
-                .check(ollama_models.into_iter().map(|m| m.to_string()).collect())
-                .await?;
+
+            // ensure that the models are pulled / pull them if not
+            let timeout = Duration::from_secs(30);
+            let good_ollama_models = self.ollama_config.check(ollama_models, timeout).await?;
+            good_models.extend(
+                good_ollama_models
+                    .into_iter()
+                    .map(|m| (ModelProvider::Ollama, m)),
+            );
         }
 
         // if OpenAI is a provider, check that the API key is set
         if unique_providers.contains(&ModelProvider::OpenAI) {
             let openai_models = self
                 .model_config
                 .get_models_for_provider(ModelProvider::OpenAI);
-            self.openai_config
-                .check(openai_models.into_iter().map(|m| m.to_string()).collect())
-                .await?;
+
+            let good_openai_models = self.openai_config.check(openai_models).await?;
+            good_models.extend(
+                good_openai_models
+                    .into_iter()
+                    .map(|m| (ModelProvider::OpenAI, m)),
+            );
         }
 
-        Ok(())
+        // update good models
+        if good_models.is_empty() {
+            return Err("No good models found, please check logs for errors.".into());
+        } else {
+            self.model_config.models = good_models;
+            Ok(())
+        }
     }
 }
 

src/config/mod.rs

@@ -38,6 +38,7 @@ impl ModelConfig {
         Self { models }
     }
 
+    /// Returns the models that belong to a given providers from the config.
     pub fn get_models_for_provider(&self, provider: ModelProvider) -> Vec<Model> {
         self.models
             .iter()

src/config/models.rs

@@ -1,4 +1,6 @@
-use ollama_workflows::ollama_rs::Ollama;
+use std::time::Duration;
+
+use ollama_workflows::{ollama_rs::Ollama, Executor, Model, ProgramMemory, Workflow};
 
 const DEFAULT_OLLAMA_HOST: &str = "http://127.0.0.1";
 const DEFAULT_OLLAMA_PORT: u16 = 11434;
@@ -46,12 +48,11 @@ impl OllamaConfig {
             .unwrap_or(DEFAULT_OLLAMA_PORT);
 
         // Ollama workflows may require specific models to be loaded regardless of the choices
-        let hardcoded_models = HARDCODED_MODELS
-            .into_iter()
-            .map(|s| s.to_string())
-            .collect();
+        let hardcoded_models = HARDCODED_MODELS.iter().map(|s| s.to_string()).collect();
 
-        let auto_pull = std::env::var("OLLAMA_AUTO_PULL").unwrap_or_default() == "true";
+        let auto_pull = std::env::var("OLLAMA_AUTO_PULL")
+            .map(|s| s == "true")
+            .unwrap_or_default();
 
         Self {
             host,
@@ -61,19 +62,20 @@ impl OllamaConfig {
         }
     }
 
-    /// Check if requested models exist.
-    pub async fn check(&self, external_models: Vec<String>) -> Result<(), String> {
+    /// Check if requested models exist in Ollama, and then tests them using a workflow.
+    pub async fn check(
+        &self,
+        external_models: Vec<Model>,
+        test_workflow_timeout: Duration,
+    ) -> Result<Vec<Model>, String> {
         log::info!(
-            "Checking Ollama requirements (auto-pull {})",
-            if self.auto_pull { "on" } else { "off" }
+            "Checking Ollama requirements (auto-pull {}, workflow timeout: {}s)",
+            if self.auto_pull { "on" } else { "off" },
+            test_workflow_timeout.as_secs()
         );
 
         let ollama = Ollama::new(&self.host, self.port);
 
-        // the list of required models is those given in DKN_MODELS and the hardcoded ones
-        let mut required_models = self.hardcoded_models.clone();
-        required_models.extend(external_models);
-
         // fetch local models
         let local_models = match ollama.list_local_models().await {
             Ok(models) => models.into_iter().map(|m| m.name).collect::<Vec<_>>(),
@@ -84,34 +86,139 @@ impl OllamaConfig {
                 }
             }
         };
+        log::info!("Found local Ollama models: {:#?}", local_models);
+
+        // check hardcoded models & pull them if available
+        // these are not used directly by the user, but are needed for the workflows
+        log::debug!("Checking hardcoded models: {:#?}", self.hardcoded_models);
+        for model in &self.hardcoded_models {
+            if !local_models.contains(model) {
+                self.try_pull(&ollama, model.to_owned()).await?;
+            }
+
+            // we dont check workflows for hardcoded models
+        }
+
+        // check external models & pull them if available
+        // and also run a test workflow for them
+        let mut good_models = Vec::new();
+        for model in external_models {
+            if !local_models.contains(&model.to_string()) {
+                self.try_pull(&ollama, model.to_string()).await?;
+            }
+
+            let ok = self
+                .test_workflow(model.clone(), test_workflow_timeout)
+                .await;
+            if ok {
+                good_models.push(model);
+            }
+        }
+
+        log::info!(
+            "Ollama checks are finished, using models: {:#?}",
+            good_models
+        );
+        Ok(good_models)
+    }
 
-        // check that each required model exists here
-        log::debug!("Checking required models: {:#?}", required_models);
-        log::debug!("Found local models: {:#?}", local_models);
-        for model in required_models {
-            if !local_models.iter().any(|m| *m == model) {
-                log::warn!("Model {} not found in Ollama", model);
-                if self.auto_pull {
-                    // if auto-pull is enabled, pull the model
-                    log::info!(
-                        "Downloading missing model {} (this may take a while)",
-                        model
-                    );
-                    let status = ollama
-                        .pull_model(model, false)
-                        .await
-                        .map_err(|e| format!("Error pulling model with Ollama: {}", e))?;
-                    log::debug!("Pulled model with Ollama, final status: {:#?}", status);
+    /// Pulls a model if `auto_pull` exists, otherwise returns an error.
+    async fn try_pull(&self, ollama: &Ollama, model: String) -> Result<(), String> {
+        log::warn!("Model {} not found in Ollama", model);
+        if self.auto_pull {
+            // if auto-pull is enabled, pull the model
+            log::info!(
+                "Downloading missing model {} (this may take a while)",
+                model
+            );
+            let status = ollama
+                .pull_model(model, false)
+                .await
+                .map_err(|e| format!("Error pulling model with Ollama: {}", e))?;
+            log::debug!("Pulled model with Ollama, final status: {:#?}", status);
+            Ok(())
+        } else {
+            // otherwise, give error
+            log::error!("Please download missing model with: ollama pull {}", model);
+            log::error!("Or, set OLLAMA_AUTO_PULL=true to pull automatically.");
+            return Err("Required model not pulled in Ollama.".into());
+        }
+    }
+
+    /// Runs a small workflow to test Ollama Workflows.
+    ///
+    /// This is to see if a given system can execute Ollama workflows for their chosen models,
+    /// e.g. if they have enough RAM/CPU and such.
+    pub async fn test_workflow(&self, model: Model, timeout: Duration) -> bool {
+        // this is the test workflow that we will run
+        // TODO: when Workflow's have `Clone`, we can remove the repetitive parsing here
+        let workflow = serde_json::from_value::<Workflow>(serde_json::json!({
+            "name": "Simple",
+            "description": "This is a simple workflow",
+            "config":{
+                "max_steps": 5,
+                "max_time": 100,
+                "max_tokens": 100,
+                "tools": []
+            },
+            "tasks":[
+                {
+                    "id": "A",
+                    "name": "Random Poem",
+                    "description": "Writes a poem about Kapadokya.",
+                    "prompt": "Please write a poem about Kapadokya.",
+                    "inputs":[],
+                    "operator": "generation",
+                    "outputs":[
+                        {
+                            "type": "write",
+                            "key": "poem",
+                            "value": "__result"
+                        }
+                    ]
+                },
+                {
+                    "id": "__end",
+                    "name": "end",
+                    "description": "End of the task",
+                    "prompt": "End of the task",
+                    "inputs": [],
+                    "operator": "end",
+                    "outputs": []
+                }
+            ],
+            "steps":[
+                {
+                    "source":"A",
+                    "target":"end"
+                }
+            ],
+            "return_value":{
+                "input":{
+                    "type": "read",
+                    "key": "poem"
+                }
+            }
+        }))
+        .expect("Preset workflow should be parsed");
+
+        log::info!("Testing model {}", model);
+        let executor = Executor::new_at(model.clone(), &self.host, self.port);
+        let mut memory = ProgramMemory::new();
+        tokio::select! {
+            _ = tokio::time::sleep(timeout) => {
+                log::warn!("Ignoring model {}: Timeout", model);
+            },
+            result = executor.execute(None, workflow, &mut memory) => {
+                if result.is_empty() {
+                    log::warn!("Ignoring model {}: Empty Result", model);
                 } else {
-                    // otherwise, give error
-                    log::error!("Please download it with: ollama pull {}", model);
-                    log::error!("Or, set OLLAMA_AUTO_PULL=true to pull automatically.");
-                    return Err("Required model not pulled in Ollama.".into());
+                    log::info!("Accepting model {}", model);
+                    return true;
                 }
             }
-        }
+        };
 
-        log::info!("Ollama setup is all good.",);
-        Ok(())
+        return false;
     }
 }

src/config/ollama.rs

@@ -1,5 +1,6 @@
 #![allow(unused)]
 
+use ollama_workflows::Model;
 use serde::Deserialize;
 
 const OPENAI_API_KEY: &str = "OPENAI_API_KEY";
@@ -38,14 +39,18 @@ impl OpenAIConfig {
         Self { api_key }
     }
 
-    /// Check if requested models exist.
-    pub async fn check(&self, models: Vec<String>) -> Result<(), String> {
+    /// Check if requested models exist &
+    ///
+    ///
+    pub async fn check(&self, models: Vec<Model>) -> Result<Vec<Model>, String> {
         log::info!("Checking OpenAI requirements");
 
+        // check API key
         let Some(api_key) = &self.api_key else {
             return Err("OpenAI API key not found".into());
         };
 
+        // fetch models
         let client = reqwest::Client::new();
         let request = client
             .get(OPENAI_MODELS_API)
@@ -58,25 +63,40 @@ impl OpenAIConfig {
             .await
             .map_err(|e| format!("Failed to send request: {}", e))?;
 
+        // parse response
         if response.status().is_client_error() {
             return Err(format!(
                 "Failed to fetch OpenAI models:\n{}",
                 response.text().await.unwrap_or_default()
             ));
         }
+        let openai_models = response
+            .json::<OpenAIModelsResponse>()
+            .await
+            .map_err(|e| e.to_string())?;
 
-        let openai_models = response.json::<OpenAIModelsResponse>().await.unwrap();
+        // check if models exist and select those that are available
+        let mut available_models = Vec::new();
         for requested_model in models {
-            if !openai_models.data.iter().any(|m| m.id == requested_model) {
-                return Err(format!(
-                    "Model {} not found in your OpenAI account.",
+            if !openai_models
+                .data
+                .iter()
+                .any(|m| m.id == requested_model.to_string())
+            {
+                log::warn!(
+                    "Model {} not found in your OpenAI account, ignoring it.",
                     requested_model
-                ));
+                );
+            } else {
+                available_models.push(requested_model);
             }
         }
 
-        log::info!("OpenAI setup is all good.",);
-        Ok(())
+        log::info!(
+            "OpenAI checks are finished, using models: {:#?}",
+            available_models
+        );
+        Ok(available_models)
     }
 }