service check quickfixes

erhant · erhant · commit de734edc6db5 · 2025-06-01T22:52:27.000+03:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -7,7 +7,7 @@ default-members = ["compute"]
 
 [workspace.package]
 edition = "2021"
-version = "0.6.1"
+version = "0.6.2"
 license = "Apache-2.0"
 readme = "README.md"
 
diff --git a/compose.yml b/compose.yml
@@ -9,14 +9,10 @@ services:
       DKN_WALLET_SECRET_KEY: ${DKN_WALLET_SECRET_KEY}
       DKN_MODELS: ${DKN_MODELS}
       DKN_P2P_LISTEN_ADDR: ${DKN_P2P_LISTEN_ADDR}
-      DKN_RELAY_NODES: ${DKN_RELAY_NODES}
-      DKN_BOOTSTRAP_NODES: ${DKN_BOOTSTRAP_NODES}
       # API Keys
       OPENAI_API_KEY: ${OPENAI_API_KEY}
       GEMINI_API_KEY: ${GEMINI_API_KEY}
       OPENROUTER_API_KEY: ${OPENROUTER_API_KEY}
-      SERPER_API_KEY: ${SERPER_API_KEY}
-      JINA_API_KEY: ${JINA_API_KEY}
       # Ollama
       OLLAMA_HOST: ${OLLAMA_HOST}
       OLLAMA_PORT: ${OLLAMA_PORT}
diff --git a/compute/src/main.rs b/compute/src/main.rs
@@ -92,17 +92,23 @@ async fn main() -> Result<()> {
             log::info!("Service check cancelled, exiting.");
             return Ok(());
         }
-    }?;
-    log::info!(
-        "Using models: {}\n{}",
-        config.executors.get_model_names().join(", "),
-        model_perf
-            .iter()
-            .map(|(model, perf)| format!("{}: {}", model, perf))
-            .collect::<Vec<_>>()
-            .join("\n")
-    );
+    };
 
+    if config.executors.models.is_empty() {
+        return Err(eyre::eyre!(
+            "No valid models left after service checks, exiting."
+        ));
+    } else {
+        log::info!(
+            "Using models: {}\n{}",
+            config.executors.get_model_names().join(", "),
+            model_perf
+                .iter()
+                .map(|(model, perf)| format!("{}: {}", model, perf))
+                .collect::<Vec<_>>()
+                .join("\n")
+        );
+    }
     // create the node
     let batch_size = config.batch_size;
     let (mut node, p2p, worker_batch, worker_single) =
diff --git a/compute/src/utils/specs.rs b/compute/src/utils/specs.rs
@@ -1,10 +1,9 @@
-use std::collections::HashMap;
-
 use dkn_executor::Model;
 use dkn_utils::{
     payloads::{SpecModelPerformance, Specs},
     SemanticVersion,
 };
+use std::collections::HashMap;
 use sysinfo::{CpuRefreshKind, MemoryRefreshKind, RefreshKind};
 
 pub struct SpecCollector {
@@ -30,7 +29,7 @@ impl SpecCollector {
         version: SemanticVersion,
         exec_platform: String,
     ) -> Self {
-        log::debug!("Creating spec collector with version {version} and platform {exec_platform}");
+        log::info!("Creating spec collector with version {version} and platform {exec_platform} and models {models:?}");
         SpecCollector {
             system: sysinfo::System::new_with_specifics(Self::get_refresh_specifics()),
             models,
diff --git a/executor/examples/gemini.rs b/executor/examples/gemini.rs
@@ -7,20 +7,12 @@ async fn main() -> eyre::Result<()> {
     let model = Model::Gemini2_0Flash;
     let models = vec![model];
     let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
-    config.check_services().await?;
-
+    config.check_services().await;
     assert!(config.models.contains(&model));
 
-    // make a request
     let task = dkn_executor::TaskBody::new_prompt("Write a haiku about category theory.", model);
-    let executor = config
-        .get_executor(&task.model)
-        .await
-        .expect("could not get executor");
-    let result = executor
-        .execute(task)
-        .await
-        .expect("failed to execute task");
+    let executor = config.get_executor(&task.model).await?;
+    let result = executor.execute(task).await?;
 
     println!("{}", result);
     Ok(())
diff --git a/executor/examples/ollama.rs b/executor/examples/ollama.rs
@@ -7,20 +7,12 @@ async fn main() -> eyre::Result<()> {
     let model = Model::Llama3_2_1bInstructQ4Km;
     let models = vec![model];
     let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
-    config.check_services().await?;
-
+    config.check_services().await;
     assert!(config.models.contains(&model));
 
-    // make a request
     let task = dkn_executor::TaskBody::new_prompt("Write a haiku about category theory.", model);
-    let executor = config
-        .get_executor(&task.model)
-        .await
-        .expect("could not get executor");
-    let result = executor
-        .execute(task)
-        .await
-        .expect("failed to execute task");
+    let executor = config.get_executor(&task.model).await?;
+    let result = executor.execute(task).await?;
 
     println!("{}", result);
     Ok(())
diff --git a/executor/examples/openai.rs b/executor/examples/openai.rs
@@ -7,20 +7,12 @@ async fn main() -> eyre::Result<()> {
     let model = Model::GPT4o;
     let models = vec![model];
     let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
-    config.check_services().await?;
-
+    config.check_services().await;
     assert!(config.models.contains(&model));
 
-    // make a request
     let task = dkn_executor::TaskBody::new_prompt("Write a haiku about category theory.", model);
-    let executor = config
-        .get_executor(&task.model)
-        .await
-        .expect("could not get executor");
-    let result = executor
-        .execute(task)
-        .await
-        .expect("failed to execute task");
+    let executor = config.get_executor(&task.model).await?;
+    let result = executor.execute(task).await?;
 
     println!("{}", result);
     Ok(())
diff --git a/executor/examples/openrouter.rs b/executor/examples/openrouter.rs
@@ -7,20 +7,12 @@ async fn main() -> eyre::Result<()> {
     let model = Model::OR3_5Sonnet;
     let models = vec![model];
     let mut config = DriaExecutorsManager::new_from_env_for_models(models.into_iter())?;
-    config.check_services().await?;
-
+    config.check_services().await;
     assert!(config.models.contains(&model));
 
-    // make a request
     let task = dkn_executor::TaskBody::new_prompt("Write a haiku about category theory.", model);
-    let executor = config
-        .get_executor(&task.model)
-        .await
-        .expect("could not get executor");
-    let result = executor
-        .execute(task)
-        .await
-        .expect("failed to execute task");
+    let executor = config.get_executor(&task.model).await?;
+    let result = executor.execute(task).await?;
 
     println!("{}", result);
     Ok(())
diff --git a/executor/src/executors/mod.rs b/executor/src/executors/mod.rs
@@ -59,4 +59,13 @@ impl DriaExecutor {
             DriaExecutor::OpenRouter(provider) => provider.check(models).await,
         }
     }
+
+    pub fn name(&self) -> String {
+        match self {
+            DriaExecutor::Ollama(_) => ModelProvider::Ollama.to_string(),
+            DriaExecutor::OpenAI(_) => ModelProvider::OpenAI.to_string(),
+            DriaExecutor::Gemini(_) => ModelProvider::Gemini.to_string(),
+            DriaExecutor::OpenRouter(_) => ModelProvider::OpenRouter.to_string(),
+        }
+    }
 }
diff --git a/executor/src/manager.rs b/executor/src/manager.rs
@@ -35,8 +35,17 @@ impl DriaExecutorsManager {
                 }
                 None => {
                     // create a new executor for the provider, may return an error!
-                    let executor = DriaExecutor::new_from_env(provider)?;
-                    provider_set.insert(provider, (executor, HashSet::from_iter([model])));
+                    match DriaExecutor::new_from_env(provider) {
+                        Ok(executor) => {
+                            provider_set.insert(provider, (executor, HashSet::from_iter([model])));
+                        }
+                        Err(err) => {
+                            log::error!(
+                            "Failed to create executor for {provider}: {err}, {model} will not be supported.",
+                        );
+                            continue; // skip this model if the executor creation failed
+                        }
+                    }
                 }
             }
 
@@ -90,14 +99,27 @@ impl DriaExecutorsManager {
     ///
     /// In the end, bad models are filtered out and we simply check if we are left if any valid models at all.
     /// If there are no models left in the end, an error is thrown.
-    pub async fn check_services(&mut self) -> eyre::Result<HashMap<Model, SpecModelPerformance>> {
+    pub async fn check_services(&mut self) -> HashMap<Model, SpecModelPerformance> {
         log::info!("Checking configured services.");
 
         // check all configured providers & record model performances
         let mut model_perf = HashMap::new();
         for (client, models) in self.providers.values_mut() {
-            let provider_model_perf = client.check(models).await?;
-            model_perf.extend(provider_model_perf);
+            if let Ok(provider_model_perf) = client.check(models).await {
+                model_perf.extend(provider_model_perf);
+            } else {
+                log::warn!(
+                    "Provider {} failed to check services, ignoring its models.",
+                    client.name()
+                );
+                model_perf.extend(
+                    models
+                        .iter()
+                        .map(|m| (m.clone(), SpecModelPerformance::ExecutionFailed)),
+                );
+                // clear models
+                models.clear();
+            }
         }
 
         // obtain the final list of providers & models, removing the providers with no models left
@@ -112,11 +134,13 @@ impl DriaExecutorsManager {
             ok
         });
 
-        // check if we have any models left at all
-        if self.providers.is_empty() {
-            eyre::bail!("No good models found, please check logs for errors.")
-        }
+        // update the models set
+        self.models = self
+            .providers
+            .values()
+            .flat_map(|(_, models)| models.iter().cloned())
+            .collect();
 
-        Ok(model_perf)
+        model_perf
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -59,4 +59,13 @@ impl DriaExecutor {`
`59`	`59`	`DriaExecutor::OpenRouter(provider) => provider.check(models).await,`
`60`	`60`	`}`
`61`	`61`	`}`
	`62`	`+`
	`63`	`+ pub fn name(&self) -> String {`
	`64`	`+ match self {`
	`65`	`+ DriaExecutor::Ollama(_) => ModelProvider::Ollama.to_string(),`
	`66`	`+ DriaExecutor::OpenAI(_) => ModelProvider::OpenAI.to_string(),`
	`67`	`+ DriaExecutor::Gemini(_) => ModelProvider::Gemini.to_string(),`
	`68`	`+ DriaExecutor::OpenRouter(_) => ModelProvider::OpenRouter.to_string(),`
	`69`	`+ }`
	`70`	`+ }`
`62`	`71`	`}`