distribute embeddings among agents

Author: mcharytoniuk

src/balancer/agent_controller.rs

@@ -24,8 +24,7 @@ use crate::balancer::agent_controller_update_result::AgentControllerUpdateResult
 use crate::balancer::chat_template_override_sender_collection::ChatTemplateOverrideSenderCollection;
 use crate::balancer::embedding_sender_collection::EmbeddingSenderCollection;
 use crate::balancer::generate_tokens_sender_collection::GenerateTokensSenderCollection;
-use crate::balancer::handles_agent_request::HandlesAgentRequest;
-use crate::balancer::manages_senders::ManagesSenders as _;
+use crate::balancer::handles_agent_streaming_response::HandlesAgentStreamingResponse;
 use crate::balancer::manages_senders::ManagesSenders;
 use crate::balancer::manages_senders_controller::ManagesSendersController;
 use crate::balancer::model_metadata_sender_collection::ModelMetadataSenderCollection;
@@ -235,7 +234,7 @@ impl AgentController {
 }
 
 #[async_trait]
-impl HandlesAgentRequest<ContinueFromConversationHistoryParams> for AgentController {
+impl HandlesAgentStreamingResponse<ContinueFromConversationHistoryParams> for AgentController {
     type SenderCollection = GenerateTokensSenderCollection;
 
     async fn handle(
@@ -256,7 +255,7 @@ impl HandlesAgentRequest<ContinueFromConversationHistoryParams> for AgentControl
 }
 
 #[async_trait]
-impl HandlesAgentRequest<ContinueFromRawPromptParams> for AgentController {
+impl HandlesAgentStreamingResponse<ContinueFromRawPromptParams> for AgentController {
     type SenderCollection = GenerateTokensSenderCollection;
 
     async fn handle(
@@ -277,7 +276,7 @@ impl HandlesAgentRequest<ContinueFromRawPromptParams> for AgentController {
 }
 
 #[async_trait]
-impl HandlesAgentRequest<GenerateEmbeddingBatchParams> for AgentController {
+impl HandlesAgentStreamingResponse<GenerateEmbeddingBatchParams> for AgentController {
     type SenderCollection = EmbeddingSenderCollection;
 
     async fn handle(

src/balancer/handles_agent_request.rs renamed to src/balancer/handles_agent_streaming_response.rs

@@ -6,7 +6,7 @@ use crate::balancer::manages_senders::ManagesSenders;
 use crate::balancer::manages_senders_controller::ManagesSendersController;
 
 #[async_trait]
-pub trait HandlesAgentRequest<TParams>
+pub trait HandlesAgentStreamingResponse<TParams>
 where
     TParams: Into<AgentJsonRpcRequest>,
 {

src/balancer/inference_service/controls_inference_endpoint.rs

@@ -13,7 +13,7 @@ use crate::agent::jsonrpc::Request as AgentJsonRpcRequest;
 use crate::balancer::agent_controller::AgentController;
 use crate::balancer::buffered_request_agent_wait_result::BufferedRequestAgentWaitResult;
 use crate::balancer::buffered_request_manager::BufferedRequestManager;
-use crate::balancer::handles_agent_request::HandlesAgentRequest;
+use crate::balancer::handles_agent_streaming_response::HandlesAgentStreamingResponse;
 use crate::balancer::inference_service::configuration::Configuration as InferenceServiceConfiguration;
 use crate::balancer::inference_service::http_route::api::ws_inference_socket::client::Message as OutgoingMessage;
 use crate::balancer::inference_service::http_route::api::ws_inference_socket::client::Response as OutgoingResponse;
@@ -39,8 +39,8 @@ pub trait ControlsInferenceEndpoint {
     ) -> Result<()>
     where
         TParams: Debug + Into<AgentJsonRpcRequest> + Send,
-        AgentController: HandlesAgentRequest<TParams>,
-        <<AgentController as HandlesAgentRequest<TParams>>::SenderCollection as ManagesSenders>::Value: Into<OutgoingResponse> + StreamableResult,
+        AgentController: HandlesAgentStreamingResponse<TParams>,
+        <<AgentController as HandlesAgentStreamingResponse<TParams>>::SenderCollection as ManagesSenders>::Value: Into<OutgoingResponse> + StreamableResult,
     {
         match Self::wait_for_agent_controller(
             buffered_request_manager.clone(),
@@ -55,6 +55,8 @@ pub trait ControlsInferenceEndpoint {
                     match agent_controller.handle(request_id.clone(), params).await {
                         Ok(receive_response_controller) => receive_response_controller,
                         Err(err) => {
+                            error!("Failed to handle request {request_id:?}: {err}");
+
                             Self::respond_with_error(
                                 JsonRpcError {
                                     code: 500,

src/balancer/inference_service/http_route/api/post_generate_embedding_batch.rs

@@ -10,6 +10,7 @@ use actix_web::Responder;
 use async_trait::async_trait;
 use bytes::Bytes;
 use futures::stream::StreamExt;
+use log::error;
 use nanoid::nanoid;
 use tokio::sync::broadcast;
 use tokio::sync::mpsc;
@@ -26,10 +27,10 @@ pub fn register(cfg: &mut web::ServiceConfig) {
     cfg.service(respond);
 }
 
-struct GenerateEmbeddingBatchController {}
+struct Controller {}
 
 #[async_trait]
-impl ControlsInferenceEndpoint for GenerateEmbeddingBatchController {
+impl ControlsInferenceEndpoint for Controller {
     type SessionController = ChunkForwardingSessionController;
 }
 
@@ -57,12 +58,36 @@ async fn respond(
     let (connection_close_tx, mut connection_close_rx) = broadcast::channel::<()>(1);
     let (chunk_tx, chunk_rx) = mpsc::unbounded_channel();
 
-    let request_batches = params.chunk_by_input_size(
-        agent_desired_state.inference_parameters.batch_n_tokens
-            * CHARACTERS_PER_TOKEN_APPROXIMATELY,
-    );
+    // Distribute the embeddings evenly across the available agents
+    params
+        .chunk_by_input_size(
+            agent_desired_state.inference_parameters.batch_n_tokens
+                * CHARACTERS_PER_TOKEN_APPROXIMATELY,
+        )
+        .for_each(|batch| {
+            let buffered_request_manager_clone = app_data.buffered_request_manager.clone();
+            let chunk_tx_clone = chunk_tx.clone();
+            let connection_close_tx_clone = connection_close_tx.clone();
+            let inference_service_configuration_clone =
+                app_data.inference_service_configuration.clone();
 
-    rt::spawn(async move {});
+            rt::spawn(async move {
+                if let Err(err) = Controller::request_from_agent(
+                    buffered_request_manager_clone,
+                    connection_close_tx_clone,
+                    inference_service_configuration_clone,
+                    batch,
+                    nanoid!(),
+                    ChunkForwardingSessionController {
+                        chunk_tx: chunk_tx_clone,
+                    },
+                )
+                .await
+                {
+                    error!("Failed to handle request: {err}");
+                }
+            });
+        });
 
     let stream = UnboundedReceiverStream::new(chunk_rx)
         .map(|chunk: String| Ok::<_, Error>(Bytes::from(format!("{chunk}\n"))))

src/balancer/inference_service/http_route/api/ws_inference_socket/client/response.rs

@@ -1,15 +1,23 @@
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::embedding_result::EmbeddingResult;
 use crate::generated_token_result::GeneratedTokenResult;
 
 #[derive(Deserialize, Serialize)]
 pub enum Response {
+    Embedding(EmbeddingResult),
     GeneratedToken(GeneratedTokenResult),
     Timeout,
     TooManyBufferedRequests,
 }
 
+impl From<EmbeddingResult> for Response {
+    fn from(result: EmbeddingResult) -> Self {
+        Response::Embedding(result)
+    }
+}
+
 impl From<GeneratedTokenResult> for Response {
     fn from(result: GeneratedTokenResult) -> Self {
         Response::GeneratedToken(result)

src/balancer/mod.rs

@@ -13,7 +13,7 @@ pub mod chat_template_override_sender_collection;
 mod controls_manages_senders_endpoint;
 pub mod embedding_sender_collection;
 pub mod generate_tokens_sender_collection;
-mod handles_agent_request;
+mod handles_agent_streaming_response;
 mod http_route;
 pub mod inference_service;
 pub mod management_service;

src/embedding_result.rs

@@ -2,9 +2,16 @@ use serde::Deserialize;
 use serde::Serialize;
 
 use crate::embedding::Embedding;
+use crate::streamable_result::StreamableResult;
 
 #[derive(Debug, Deserialize, Serialize)]
 pub enum EmbeddingResult {
     Done,
     Embedding(Embedding),
 }
+
+impl StreamableResult for EmbeddingResult {
+    fn is_done(&self) -> bool {
+        matches!(self, EmbeddingResult::Done)
+    }
+}

src/request_params/generate_embedding_batch_params/chunk_by_input_size_iter.rs

@@ -0,0 +1,45 @@
+use super::GenerateEmbeddingBatchParams;
+use crate::embedding_input_document::EmbeddingInputDocument;
+use crate::embedding_normalization_method::EmbeddingNormalizationMethod;
+
+pub struct ChunkByInputSizeIter<'embedding_batch> {
+    pub chunk_size: usize,
+    pub current_index: usize,
+    pub input_batch: &'embedding_batch [EmbeddingInputDocument],
+    pub normalization_method: &'embedding_batch EmbeddingNormalizationMethod,
+}
+
+impl<'embedding_batch> Iterator for ChunkByInputSizeIter<'embedding_batch> {
+    type Item = GenerateEmbeddingBatchParams;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.current_index >= self.input_batch.len() {
+            return None;
+        }
+
+        let mut current_batch = Vec::new();
+        let mut current_size = 0;
+
+        while self.current_index < self.input_batch.len() {
+            let input = &self.input_batch[self.current_index];
+            let input_size = input.content.chars().count();
+
+            if current_size + input_size > self.chunk_size && !current_batch.is_empty() {
+                break;
+            }
+
+            current_batch.push(input.clone());
+            current_size += input_size;
+            self.current_index += 1;
+        }
+
+        if current_batch.is_empty() {
+            None
+        } else {
+            Some(GenerateEmbeddingBatchParams {
+                input_batch: current_batch,
+                normalization_method: self.normalization_method.clone(),
+            })
+        }
+    }
+}

src/request_params/generate_embedding_batch_params.rs renamed to src/request_params/generate_embedding_batch_params/mod.rs

@@ -1,6 +1,9 @@
+mod chunk_by_input_size_iter;
+
 use serde::Deserialize;
 use serde::Serialize;
 
+use self::chunk_by_input_size_iter::ChunkByInputSizeIter;
 use crate::embedding_input_document::EmbeddingInputDocument;
 use crate::embedding_normalization_method::EmbeddingNormalizationMethod;
 
@@ -12,37 +15,16 @@ pub struct GenerateEmbeddingBatchParams {
 
 impl GenerateEmbeddingBatchParams {
     /// Input size is the total number of characters in the resulting batches.
-    pub fn chunk_by_input_size(&self, chunk_size: usize) -> Vec<GenerateEmbeddingBatchParams> {
-        let mut batches: Vec<GenerateEmbeddingBatchParams> = Vec::new();
-        let mut current_batch: Vec<EmbeddingInputDocument> = Vec::new();
-        let mut current_size: usize = 0;
-
-        for input in &self.input_batch {
-            let input_size = input.content.chars().count();
-
-            current_batch.push(input.clone());
-            current_size += input_size;
-
-            if current_size + input_size > chunk_size {
-                if !current_batch.is_empty() {
-                    batches.push(GenerateEmbeddingBatchParams {
-                        input_batch: current_batch.clone(),
-                        normalization_method: self.normalization_method.clone(),
-                    });
-                }
-                current_batch.clear();
-                current_size = 0;
-            }
+    pub fn chunk_by_input_size<'embedding>(
+        &'embedding self,
+        chunk_size: usize,
+    ) -> ChunkByInputSizeIter<'embedding> {
+        ChunkByInputSizeIter {
+            input_batch: &self.input_batch,
+            normalization_method: &self.normalization_method,
+            chunk_size,
+            current_index: 0,
         }
-
-        if !current_batch.is_empty() {
-            batches.push(GenerateEmbeddingBatchParams {
-                input_batch: current_batch,
-                normalization_method: self.normalization_method.clone(),
-            });
-        }
-
-        batches
     }
 }
 
@@ -70,7 +52,8 @@ mod tests {
             normalization_method: EmbeddingNormalizationMethod::None,
         };
 
-        let batches = params.chunk_by_input_size(10);
+        let batches = params.chunk_by_input_size(10).collect::<Vec<_>>();
+
         assert_eq!(batches.len(), 2);
         assert_eq!(batches[0].input_batch.len(), 2);
         assert_eq!(batches[0].input_batch[0].id, "1");