diff --git a/idl/chromadb/proto/fn_consumer.proto b/idl/chromadb/proto/fn_consumer.proto new file mode 100644 index 00000000000..4c11125ef48 --- /dev/null +++ b/idl/chromadb/proto/fn_consumer.proto @@ -0,0 +1,18 @@ +syntax = "proto3"; + +package chroma; + +message ListFnConsumerInProgressJobsRequest {} + +message FnConsumerInProgressJobInfo { + string fn_id = 1; + int64 expires_at_epoch_secs = 2; +} + +message ListFnConsumerInProgressJobsResponse { + repeated FnConsumerInProgressJobInfo jobs = 1; +} + +service FnConsumer { + rpc ListInProgressJobs(ListFnConsumerInProgressJobsRequest) returns (ListFnConsumerInProgressJobsResponse) {} +} diff --git a/rust/memberlist/src/client_manager.rs b/rust/memberlist/src/client_manager.rs index 8d656d60262..b0338fe31fe 100644 --- a/rust/memberlist/src/client_manager.rs +++ b/rust/memberlist/src/client_manager.rs @@ -7,8 +7,9 @@ use chroma_error::ChromaError; use chroma_system::{Component, ComponentContext, Handler}; use chroma_tracing::GrpcClientTraceService; use chroma_types::chroma_proto::{ - compactor_client::CompactorClient, heap_tender_service_client::HeapTenderServiceClient, - log_service_client::LogServiceClient, query_executor_client::QueryExecutorClient, + compactor_client::CompactorClient, fn_consumer_client::FnConsumerClient, + heap_tender_service_client::HeapTenderServiceClient, log_service_client::LogServiceClient, + query_executor_client::QueryExecutorClient, }; use parking_lot::RwLock; use std::{ @@ -561,6 +562,20 @@ impl ClientFactory } } +impl ClientFactory + for FnConsumerClient> +{ + fn new_from_channel(channel: GrpcClientTraceService) -> Self { + FnConsumerClient::new(channel) + } + fn max_encoding_message_size(self, max_size: usize) -> Self { + self.max_encoding_message_size(max_size) + } + fn max_decoding_message_size(self, max_size: usize) -> Self { + self.max_decoding_message_size(max_size) + } +} + #[cfg(test)] mod test { use super::super::memberlist_provider::Member; diff --git a/rust/types/build.rs b/rust/types/build.rs index 09da39f2934..6647f548db4 100644 --- a/rust/types/build.rs +++ b/rust/types/build.rs @@ -9,6 +9,7 @@ fn main() -> Result<(), Box> { "idl/chromadb/proto/query_executor.proto", "idl/chromadb/proto/garbage_collector.proto", "idl/chromadb/proto/fault_injection.proto", + "idl/chromadb/proto/fn_consumer.proto", "idl/chromadb/proto/workqueue.proto", ]; diff --git a/rust/worker/src/fn_consumer/fn_consumer_manager.rs b/rust/worker/src/fn_consumer/fn_consumer_manager.rs index 626bc25c1c3..8198fdcb9e6 100644 --- a/rust/worker/src/fn_consumer/fn_consumer_manager.rs +++ b/rust/worker/src/fn_consumer/fn_consumer_manager.rs @@ -16,7 +16,7 @@ use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::time::{Duration, SystemTime}; use thiserror::Error; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, oneshot}; use tracing::{instrument, span}; use crate::compactor::config::CompactorConfig; @@ -46,6 +46,35 @@ impl InProgressFn { } } +#[derive(Debug, PartialEq, Eq)] +pub struct InProgressFnEntry { + pub fn_id: AttachedFunctionUuid, + pub expires_at_epoch_secs: i64, +} + +#[derive(Debug)] +pub struct ListInProgressJobsMessage { + pub response_tx: oneshot::Sender>, +} + +fn snapshot_in_progress_jobs( + in_progress: &HashMap, +) -> Vec { + let mut entries: Vec<_> = in_progress + .iter() + .map(|(fn_id, job)| InProgressFnEntry { + fn_id: *fn_id, + expires_at_epoch_secs: job + .expires_at + .duration_since(std::time::UNIX_EPOCH) + .map(|duration| duration.as_secs() as i64) + .unwrap_or(0), + }) + .collect(); + entries.sort_unstable_by_key(|entry| entry.fn_id.to_string()); + entries +} + #[derive(Error, Debug)] pub enum DispatchError { #[error("Dispatcher not initialized")] @@ -92,6 +121,18 @@ struct FnDispatchCompletion { result: FnDispatchOutput, } +fn drain_dispatch_completions( + receiver: &mut mpsc::UnboundedReceiver, + in_progress: &mut HashMap, +) -> Vec { + let mut completions = Vec::new(); + while let Ok(completion) = receiver.try_recv() { + in_progress.remove(&completion.fn_id); + completions.push(completion); + } + completions +} + #[derive(Clone)] pub struct FnConsumerContext { pub system: System, @@ -307,9 +348,10 @@ impl FnConsumerManager { } fn process_completions(&mut self) { - while let Ok(completion) = self.dispatch_awaiter_completion_channel.try_recv() { - self.in_progress.remove(&completion.fn_id); - + for completion in drain_dispatch_completions( + &mut self.dispatch_awaiter_completion_channel, + &mut self.in_progress, + ) { match completion.result { Ok(FnDispatchOutcome::Completed) => { tracing::debug!( @@ -592,12 +634,91 @@ impl Handler for FnConsumerManager { } } +#[async_trait] +impl Handler for FnConsumerManager { + type Result = (); + + async fn handle(&mut self, message: ListInProgressJobsMessage, _ctx: &ComponentContext) { + self.process_completions(); + let entries = snapshot_in_progress_jobs(&self.in_progress); + if let Err(entries) = message.response_tx.send(entries) { + tracing::warn!( + job_count = entries.len(), + "Failed to send fn-consumer in-progress jobs response" + ); + } + } +} + #[cfg(test)] mod tests { use super::*; use tokio::sync::oneshot; use tokio::time::{timeout, Duration}; + #[test] + fn snapshots_in_progress_jobs() { + let first_fn_id = AttachedFunctionUuid::new(); + let second_fn_id = AttachedFunctionUuid::new(); + let mut in_progress = HashMap::new(); + in_progress.insert( + first_fn_id, + InProgressFn { + expires_at: std::time::UNIX_EPOCH + Duration::from_secs(20), + expiry_logged: false, + }, + ); + in_progress.insert( + second_fn_id, + InProgressFn { + expires_at: std::time::UNIX_EPOCH + Duration::from_secs(10), + expiry_logged: false, + }, + ); + + let entries = snapshot_in_progress_jobs(&in_progress); + assert_eq!(entries.len(), 2); + assert!(entries + .windows(2) + .all(|pair| pair[0].fn_id.to_string() < pair[1].fn_id.to_string())); + assert!(entries + .iter() + .any(|entry| { entry.fn_id == first_fn_id && entry.expires_at_epoch_secs == 20 })); + assert!(entries + .iter() + .any(|entry| { entry.fn_id == second_fn_id && entry.expires_at_epoch_secs == 10 })); + } + + #[test] + fn snapshots_empty_in_progress_jobs() { + assert!(snapshot_in_progress_jobs(&HashMap::new()).is_empty()); + } + + #[test] + fn draining_completions_removes_finished_jobs() { + let fn_id = AttachedFunctionUuid::new(); + let mut in_progress = HashMap::from([( + fn_id, + InProgressFn { + expires_at: std::time::UNIX_EPOCH + Duration::from_secs(20), + expiry_logged: false, + }, + )]); + let (completion_tx, mut completion_rx) = mpsc::unbounded_channel(); + completion_tx + .send(FnDispatchCompletion { + fn_id, + batch_size: 1, + result: Ok(FnDispatchOutcome::Completed), + }) + .unwrap(); + + let completions = drain_dispatch_completions(&mut completion_rx, &mut in_progress); + + assert_eq!(completions.len(), 1); + assert!(in_progress.is_empty()); + } + #[tokio::test] async fn dispatch_awaiter_completes_later_tasks_while_one_is_running() { let (task_tx, task_rx) = mpsc::channel(2); diff --git a/rust/worker/src/fn_consumer/grpc.rs b/rust/worker/src/fn_consumer/grpc.rs new file mode 100644 index 00000000000..4ba094353c7 --- /dev/null +++ b/rust/worker/src/fn_consumer/grpc.rs @@ -0,0 +1,51 @@ +use async_trait::async_trait; +use chroma_system::ComponentHandle; +use chroma_types::chroma_proto::{ + fn_consumer_server::{FnConsumer, FnConsumerServer}, + FnConsumerInProgressJobInfo, ListFnConsumerInProgressJobsRequest, + ListFnConsumerInProgressJobsResponse, +}; +use tonic::{Request, Response, Status}; + +use super::fn_consumer_manager::{FnConsumerManager, ListInProgressJobsMessage}; + +pub struct FnConsumerGrpcServer { + manager: ComponentHandle, +} + +impl FnConsumerGrpcServer { + pub fn new(manager: ComponentHandle) -> Self { + Self { manager } + } + + pub fn into_service(self) -> FnConsumerServer { + FnConsumerServer::new(self) + } +} + +#[async_trait] +impl FnConsumer for FnConsumerGrpcServer { + async fn list_in_progress_jobs( + &self, + _request: Request, + ) -> Result, Status> { + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + self.manager + .receiver() + .send(ListInProgressJobsMessage { response_tx }, None) + .await + .map_err(|error| Status::internal(error.to_string()))?; + + let jobs = response_rx + .await + .map_err(|error| Status::internal(format!("Failed to receive response: {error}")))? + .into_iter() + .map(|entry| FnConsumerInProgressJobInfo { + fn_id: entry.fn_id.to_string(), + expires_at_epoch_secs: entry.expires_at_epoch_secs, + }) + .collect(); + + Ok(Response::new(ListFnConsumerInProgressJobsResponse { jobs })) + } +} diff --git a/rust/worker/src/fn_consumer/mod.rs b/rust/worker/src/fn_consumer/mod.rs index a106b7a8da4..25281c88595 100644 --- a/rust/worker/src/fn_consumer/mod.rs +++ b/rust/worker/src/fn_consumer/mod.rs @@ -1,5 +1,6 @@ pub mod config; pub mod fn_consumer_manager; +mod grpc; pub mod server; pub use server::fn_consumer_service_entrypoint; diff --git a/rust/worker/src/fn_consumer/server.rs b/rust/worker/src/fn_consumer/server.rs index 0aa2d397851..f2661c20041 100644 --- a/rust/worker/src/fn_consumer/server.rs +++ b/rust/worker/src/fn_consumer/server.rs @@ -1,5 +1,6 @@ use crate::config::RootConfig; use crate::fn_consumer::fn_consumer_manager::FnConsumerManager; +use crate::fn_consumer::grpc::FnConsumerGrpcServer; use crate::work_queue::work_queue_client::WorkQueueClient; use chroma_blockstore::provider::BlockfileProvider; use chroma_config::registry::Registry; @@ -172,10 +173,16 @@ pub async fn fn_consumer_service_entrypoint() { spann_provider, ); manager.set_dispatcher(dispatcher_handle); - let _manager_handle = system.start_component(manager); + let manager_handle = system.start_component(manager); // Create health service for readiness probe - let (_health_reporter, health_service) = tonic_health::server::health_reporter(); + let (health_reporter, health_service) = tonic_health::server::health_reporter(); + health_reporter + .set_serving::>() + .await; + let fn_consumer_service = FnConsumerGrpcServer::new(manager_handle).into_service(); let addr = format!("0.0.0.0:{}", service_config.my_port) .parse() @@ -186,6 +193,7 @@ pub async fn fn_consumer_service_entrypoint() { // Start server (this blocks forever) Server::builder() .add_service(health_service) + .add_service(fn_consumer_service) .serve(addr) .await .expect("Failed to start fn-consumer service");