use dummy gen instead of embedding for warm-ups for Ollama

Author: erhant

Cargo.lock

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

Cargo.toml

@@ -1,8 +1,5 @@
 use eyre::{eyre, Context, Result};
-use ollama_rs::generation::{
-    completion::request::GenerationRequest,
-    embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest},
-};
+use ollama_rs::generation::completion::request::GenerationRequest;
 use rig::completion::{Chat, PromptError};
 use rig::providers::ollama;
 use std::time::Duration;
@@ -165,18 +162,19 @@ impl OllamaClient {
     pub async fn test_performance(&self, model: &Model) -> bool {
         log::info!("Testing model {}", model);
 
-        // first generate a dummy embedding to load the model into memory (warm-up)
-        let request = GenerateEmbeddingsRequest::new(
-            model.to_string(),
-            EmbeddingsInput::Single("embedme".into()),
-        );
-        if let Err(err) = self.ollama_rs_client.generate_embeddings(request).await {
-            log::error!("Failed to generate embedding for model {}: {}", model, err);
-            return false;
-        };
-
         let generation_request = GenerationRequest::new(model.to_string(), TEST_PROMPT.to_string());
 
+        // run a dummy generation for warm-up
+        log::debug!("Warming up Ollama for model {}", model);
+        if let Err(e) = self
+            .ollama_rs_client
+            .generate(generation_request.clone())
+            .await
+        {
+            log::warn!("Ignoring model {}: Workflow failed with error {}", model, e);
+            return false;
+        }
+
         // then, run a sample generation with timeout and measure tps
         tokio::select! {
             _ = tokio::time::sleep(PERFORMANCE_TIMEOUT) => {

workflows-v2/src/providers/ollama.rs

@@ -1,12 +1,6 @@
 use eyre::{eyre, Context, Result};
 use ollama_workflows::{
-    ollama_rs::{
-        generation::{
-            completion::request::GenerationRequest,
-            embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest},
-        },
-        Ollama,
-    },
+    ollama_rs::{generation::completion::request::GenerationRequest, Ollama},
     Model,
 };
 use std::env;
@@ -187,18 +181,15 @@ impl OllamaConfig {
     pub async fn test_performance(&self, ollama: &Ollama, model: &Model) -> bool {
         log::info!("Testing model {}", model);
 
-        // first generate a dummy embedding to load the model into memory (warm-up)
-        let request = GenerateEmbeddingsRequest::new(
-            model.to_string(),
-            EmbeddingsInput::Single("embedme".into()),
-        );
-        if let Err(err) = ollama.generate_embeddings(request).await {
-            log::error!("Failed to generate embedding for model {}: {}", model, err);
-            return false;
-        };
-
         let generation_request = GenerationRequest::new(model.to_string(), TEST_PROMPT.to_string());
 
+        // run a dummy generation for warm-up
+        log::debug!("Warming up Ollama for model {}", model);
+        if let Err(e) = ollama.generate(generation_request.clone()).await {
+            log::warn!("Ignoring model {}: Workflow failed with error {}", model, e);
+            return false;
+        }
+
         // then, run a sample generation with timeout and measure tps
         tokio::select! {
             _ = tokio::time::sleep(self.timeout) => {