[ENH] Clean up client manager into manager/assigner - make log client use it (#4640)

## Description of changes

_Summarize the changes made by this PR._
- Remove `set_system` pattern in lieu of passing system into
try_from_config - aligning with the common. pattern since we touched
this
- Refactor client manager into client manager + assigner thats generic
over its clients
- Use this in the old call site and in the log, the log uses assignment
for write path, and reads all for read path


## Test plan

_How are these changes tested?_
- [x] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Documentation Changes
None

Author: HammadB

Cargo.lock

@@ -69,7 +69,7 @@ tokio = { version = "1.41", features = ["fs", "macros", "rt-multi-thread"] }
 tokio-util = "0.7.12"
 tonic = "0.12"
 tonic-health = "0.12.3"
-tower = "0.5"
+tower = { version = "0.4.13", features = ["discover"] }
 backon = "1.3.0"
 tracing = { version = "0.1" }
 tracing-bunyan-formatter = "0.3"

Cargo.toml

@@ -28,7 +28,7 @@ regex = { workspace = true }
 futures = { workspace = true }
 backon = { workspace = true }
 tracing-opentelemetry = { workspace = true }
-tower = { version = "0.4.13", features = ["discover"] }
+tower = { workspace = true }
 mdac = { workspace = true }
 opentelemetry.workspace = true
 validator = { workspace = true }

rust/frontend/Cargo.toml

@@ -1,28 +1,29 @@
-use super::client_manager::NodeNameToClient;
-use super::{client_manager::ClientManager, config};
+use super::config;
 use async_trait::async_trait;
 use backon::ExponentialBuilder;
 use backon::Retryable;
 use chroma_config::registry;
 use chroma_config::{assignment::assignment_policy::AssignmentPolicy, Configurable};
 use chroma_error::ChromaError;
+use chroma_memberlist::client_manager::ClientAssigner;
+use chroma_memberlist::client_manager::{ClientManager, ClientOptions};
 use chroma_memberlist::{
     config::MemberlistProviderConfig,
     memberlist_provider::{CustomResourceMemberlistProvider, MemberlistProvider},
 };
 use chroma_system::System;
+use chroma_types::chroma_proto::query_executor_client::QueryExecutorClient;
 use chroma_types::SegmentType;
 use chroma_types::{
-    chroma_proto::query_executor_client::QueryExecutorClient,
     operator::{CountResult, GetResult, KnnBatchResult},
     plan::{Count, Get, Knn},
-    CollectionUuid, ExecutorError,
+    ExecutorError,
 };
 use rand::seq::SliceRandom;
-use std::cmp::min;
 use tonic::Request;
 
-type Client = QueryExecutorClient<chroma_tracing::GrpcTraceService<tonic::transport::Channel>>;
+// Convenience type alias for the gRPC query client used by the DistributedExecutor
+type QueryClient = QueryExecutorClient<chroma_tracing::GrpcTraceService<tonic::transport::Channel>>;
 
 /// A distributed executor that routes requests to the appropriate node based on the assignment policy
 /// # Fields
@@ -36,8 +37,7 @@ type Client = QueryExecutorClient<chroma_tracing::GrpcTraceService<tonic::transp
 /// outside.
 #[derive(Clone, Debug)]
 pub struct DistributedExecutor {
-    node_name_to_client: NodeNameToClient,
-    assignment_policy: Box<dyn AssignmentPolicy>,
+    client_assigner: ClientAssigner<QueryClient>,
     replication_factor: usize,
     backoff: ExponentialBuilder,
 }
@@ -50,13 +50,13 @@ impl Configurable<(config::DistributedExecutorConfig, System)> for DistributedEx
     ) -> Result<Self, Box<dyn ChromaError>> {
         let assignment_policy =
             Box::<dyn AssignmentPolicy>::try_from_config(&config.assignment, registry).await?;
-        let node_name_to_client = NodeNameToClient::default();
+        let client_assigner = ClientAssigner::new(assignment_policy, config.replication_factor);
         let client_manager = ClientManager::new(
-            node_name_to_client.clone(),
+            client_assigner.clone(),
             config.connections_per_node,
             config.connect_timeout_ms,
             config.request_timeout_ms,
-            config.max_query_service_response_size_bytes,
+            ClientOptions::new(Some(config.max_query_service_response_size_bytes)),
         );
         let client_manager_handle = system.start_component(client_manager);
 
@@ -75,8 +75,7 @@ impl Configurable<(config::DistributedExecutorConfig, System)> for DistributedEx
         let retry_config = &config.retry;
         let backoff = retry_config.into();
         Ok(Self {
-            node_name_to_client,
-            assignment_policy,
+            client_assigner,
             replication_factor: config.replication_factor,
             backoff,
         })
@@ -97,7 +96,17 @@ impl DistributedExecutor {
 impl DistributedExecutor {
     ///////////////////////// Plan Operations /////////////////////////
     pub async fn count(&mut self, plan: Count) -> Result<CountResult, ExecutorError> {
-        let clients = self.clients(plan.scan.collection_and_segments.collection.collection_id)?;
+        let clients = self
+            .client_assigner
+            .clients(
+                &plan
+                    .scan
+                    .collection_and_segments
+                    .collection
+                    .collection_id
+                    .to_string(),
+            )
+            .map_err(|e| ExecutorError::Internal(e.boxed()))?;
         let plan: chroma_types::chroma_proto::CountPlan = plan.clone().try_into()?;
         let res = (|| async {
             choose_client(clients.as_slice())?
@@ -111,7 +120,17 @@ impl DistributedExecutor {
     }
 
     pub async fn get(&mut self, plan: Get) -> Result<GetResult, ExecutorError> {
-        let clients = self.clients(plan.scan.collection_and_segments.collection.collection_id)?;
+        let clients = self
+            .client_assigner
+            .clients(
+                &plan
+                    .scan
+                    .collection_and_segments
+                    .collection
+                    .collection_id
+                    .to_string(),
+            )
+            .map_err(|e| ExecutorError::Internal(e.boxed()))?;
         let res = (|| async {
             choose_client(clients.as_slice())?
                 .get(Request::new(plan.clone().try_into()?))
@@ -124,7 +143,17 @@ impl DistributedExecutor {
     }
 
     pub async fn knn(&mut self, plan: Knn) -> Result<KnnBatchResult, ExecutorError> {
-        let clients = self.clients(plan.scan.collection_and_segments.collection.collection_id)?;
+        let clients = self
+            .client_assigner
+            .clients(
+                &plan
+                    .scan
+                    .collection_and_segments
+                    .collection
+                    .collection_id
+                    .to_string(),
+            )
+            .map_err(|e| ExecutorError::Internal(e.boxed()))?;
         let res = (|| async {
             choose_client(clients.as_slice())?
                 .knn(Request::new(plan.clone().try_into()?))
@@ -137,38 +166,15 @@ impl DistributedExecutor {
     }
 
     pub async fn is_ready(&self) -> bool {
-        !self.node_name_to_client.read().is_empty()
+        !self.client_assigner.is_empty()
     }
+}
 
-    ///////////////////////// Helpers /////////////////////////
-
-    /// Get the gRPC clients for the given collection id by performing the assignment policy
-    /// # Arguments
-    /// - `collection_id` - The collection id for which the client is to be fetched
-    /// # Returns
-    /// - The gRPC clients for the given collection id in the order of the assignment policy
-    /// # Errors
-    /// - If no client is found for the given collection id
-    /// - If the assignment policy fails to assign the collection id
-    fn clients(&mut self, collection_id: CollectionUuid) -> Result<Vec<Client>, ExecutorError> {
-        let node_name_to_client_guard = self.node_name_to_client.read();
-        let members: Vec<String> = node_name_to_client_guard.keys().cloned().collect();
-        let target_replication_factor = min(self.replication_factor, members.len());
-        self.assignment_policy.set_members(members);
-        let assigned = self
-            .assignment_policy
-            .assign(&collection_id.to_string(), target_replication_factor)?;
-        let clients = assigned
-            .iter()
-            .map(|node_name| {
-                node_name_to_client_guard
-                    .get(node_name)
-                    .ok_or_else(|| ExecutorError::NoClientFound(node_name.clone()))
-                    .cloned()
-            })
-            .collect::<Result<Vec<_>, _>>()?;
-        Ok(clients)
-    }
+fn choose_client(clients: &[QueryClient]) -> Result<QueryClient, tonic::Status> {
+    Ok(clients
+        .choose(&mut rand::thread_rng())
+        .ok_or(no_clients_found_status())?
+        .clone())
 }
 
 fn is_retryable_error(e: &tonic::Status) -> bool {
@@ -181,10 +187,3 @@ fn is_retryable_error(e: &tonic::Status) -> bool {
 fn no_clients_found_status() -> tonic::Status {
     tonic::Status::internal("No clients found")
 }
-
-fn choose_client(clients: &[Client]) -> Result<Client, tonic::Status> {
-    Ok(clients
-        .choose(&mut rand::thread_rng())
-        .ok_or(no_clients_found_status())?
-        .clone())
-}

rust/frontend/src/executor/distributed.rs

@@ -7,7 +7,6 @@ use distributed::DistributedExecutor;
 use local::LocalExecutor;
 
 //////////////////////// Exposed Modules ////////////////////////
-pub(super) mod client_manager;
 pub mod config;
 // TODO: This should be private once we fix dep injection
 mod distributed;

rust/frontend/src/executor/mod.rs

@@ -1458,7 +1458,7 @@ impl Configurable<(FrontendConfig, System)> for ServiceBasedFrontend {
         };
 
         let sysdb = SysDb::try_from_config(&config.sysdb, registry).await?;
-        let mut log = Log::try_from_config(&config.log, registry).await?;
+        let mut log = Log::try_from_config(&(config.log.clone(), system.clone()), registry).await?;
         let max_batch_size = log.get_max_batch_size().await?;
 
         // Create compation manager and pass handle to log service if configured

rust/frontend/src/impls/service_based_frontend.rs

@@ -35,6 +35,7 @@ chroma-system = { workspace = true }
 chroma-sysdb = { workspace = true }
 chroma-types = { workspace = true }
 chroma-sqlite = { workspace = true }
+chroma-memberlist = { workspace = true }
 wal3 = { workspace = true }
 figment.workspace = true
 

rust/log/Cargo.toml

@@ -1,3 +1,4 @@
+use chroma_memberlist::config::CustomResourceMemberlistProviderConfig;
 use serde::{Deserialize, Serialize};
 
 #[derive(Deserialize, Clone, Serialize, Debug)]
@@ -28,6 +29,10 @@ pub struct GrpcLogConfig {
     pub use_alt_for_collections: Vec<String>,
     #[serde(default = "Option::default")]
     pub alt_host_threshold: Option<String>,
+    #[serde(default = "GrpcLogConfig::default_memberlist_provider")]
+    pub memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig,
+    #[serde(default = "GrpcLogConfig::default_assignment")]
+    pub assignment: chroma_config::assignment::config::AssignmentPolicyConfig,
 }
 
 impl GrpcLogConfig {
@@ -54,6 +59,24 @@ impl GrpcLogConfig {
     fn default_max_decoding_message_size() -> usize {
         32_000_000
     }
+
+    fn default_memberlist_provider() -> chroma_memberlist::config::MemberlistProviderConfig {
+        chroma_memberlist::config::MemberlistProviderConfig::CustomResource(
+            CustomResourceMemberlistProviderConfig {
+                kube_namespace: "chroma".to_string(),
+                memberlist_name: "rust-log-service-memberlist".to_string(),
+                queue_size: 100,
+            },
+        )
+    }
+
+    fn default_assignment() -> chroma_config::assignment::config::AssignmentPolicyConfig {
+        chroma_config::assignment::config::AssignmentPolicyConfig::RendezvousHashing(
+            chroma_config::assignment::config::RendezvousHashingAssignmentPolicyConfig {
+                hasher: chroma_config::assignment::config::HasherType::Murmur3,
+            },
+        )
+    }
 }
 
 impl Default for GrpcLogConfig {
@@ -69,6 +92,8 @@ impl Default for GrpcLogConfig {
             use_alt_for_tenants: vec![],
             use_alt_for_collections: vec![],
             alt_host_threshold: None,
+            memberlist_provider: GrpcLogConfig::default_memberlist_provider(),
+            assignment: GrpcLogConfig::default_assignment(),
         }
     }
 }