[ENH]: add manual compaction method to frontend

Tiltfile

@@ -147,7 +147,7 @@ k8s_resource('sysdb', resource_deps=['sysdb-migration-sysdb-migration'], labels=
 k8s_resource('frontend-service', resource_deps=['sysdb', 'logservice'],labels=["chroma"], port_forwards='8000:8000')
 k8s_resource('rust-frontend-service', resource_deps=['sysdb', 'logservice'], labels=["chroma"], port_forwards='3000:8000')
 k8s_resource('query-service', resource_deps=['sysdb'], labels=["chroma"], port_forwards='50053:50051')
-k8s_resource('compaction-service', resource_deps=['sysdb'], labels=["chroma"])
+k8s_resource('compaction-service', resource_deps=['sysdb'], labels=["chroma"], port_forwards='50054:50051')
 
 # I have no idea why these need their own lines but the others don't.
 k8s_resource('jaeger', resource_deps=['k8s_setup'], labels=["observability"])

rust/frontend/sample_configs/distributed.yaml

@@ -21,6 +21,9 @@ log:
     port: 50051
     connect_timeout_ms: 5000
     request_timeout_ms: 5000
+compaction_client:
+  Grpc:
+    url: "compactionservice.chroma:50051" # todo
 executor:
   distributed:
     connections_per_node: 5

rust/frontend/src/compaction_client/compaction_client.rs

@@ -0,0 +1,57 @@
+use super::{
+    grpc_compaction_client::{GrpcCompactionClient, GrpcCompactionClientConfig},
+    local_compaction_client::LocalCompactionClient,
+};
+use chroma_config::{registry::Registry, Configurable};
+use chroma_error::ChromaError;
+use chroma_system::System;
+use chroma_types::{ManualCompactionError, ManualCompactionRequest, ManualCompactionResponse};
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone)]
+pub enum CompactionClient {
+    Local(LocalCompactionClient),
+    Grpc(GrpcCompactionClient),
+}
+
+impl CompactionClient {
+    pub async fn manually_compact(
+        &mut self,
+        request: ManualCompactionRequest,
+    ) -> Result<ManualCompactionResponse, ManualCompactionError> {
+        match self {
+            CompactionClient::Local(client) => client.manually_compact(request).await,
+            CompactionClient::Grpc(client) => client.manually_compact(request).await,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub enum CompactionClientConfig {
+    Local,
+    Grpc(GrpcCompactionClientConfig),
+}
+
+#[async_trait::async_trait]
+impl Configurable<(CompactionClientConfig, System)> for CompactionClient {
+    async fn try_from_config(
+        (config, system): &(CompactionClientConfig, System),
+        registry: &Registry,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        match config {
+            CompactionClientConfig::Local => {
+                let client =
+                    LocalCompactionClient::try_from_config(&((), system.clone()), registry).await?;
+                Ok(CompactionClient::Local(client))
+            }
+            CompactionClientConfig::Grpc(grpc_config) => {
+                let client = GrpcCompactionClient::try_from_config(
+                    &(grpc_config.clone(), system.clone()),
+                    registry,
+                )
+                .await?;
+                Ok(CompactionClient::Grpc(client))
+            }
+        }
+    }
+}

rust/frontend/src/compaction_client/grpc_compaction_client.rs

@@ -0,0 +1,90 @@
+use chroma_config::{registry::Registry, Configurable};
+use chroma_error::ChromaError;
+use chroma_sysdb::SysDb;
+use chroma_system::System;
+use chroma_types::{
+    chroma_proto::{self, compactor_client::CompactorClient},
+    CollectionUuid, ManualCompactionError, ManualCompactionRequest, ManualCompactionResponse,
+};
+use serde::{Deserialize, Serialize};
+use thiserror::Error;
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct GrpcCompactionClientConfig {
+    pub url: String,
+}
+
+#[derive(Debug, Clone)]
+pub(super) struct GrpcCompactionClient {
+    sysdb: SysDb,
+    client: CompactorClient<tonic::transport::Channel>,
+}
+
+impl GrpcCompactionClient {
+    pub async fn manually_compact(
+        &mut self,
+        request: ManualCompactionRequest,
+    ) -> Result<ManualCompactionResponse, ManualCompactionError> {
+        async fn get_version(sysdb: &mut SysDb, collection_id: CollectionUuid) -> i32 {
+            let mut collection = sysdb
+                .get_collections(Some(collection_id), None, None, None, None, 0)
+                .await
+                .unwrap(); // todo
+            let collection = collection.pop().unwrap(); // todo
+            return collection.version;
+        }
+
+        let version_before_compaction = get_version(&mut self.sysdb, request.collection_id).await;
+
+        self.client
+            .compact(chroma_proto::CompactionRequest {
+                ids: Some(chroma_proto::CollectionIds {
+                    ids: vec![request.collection_id.to_string()],
+                }),
+            })
+            .await
+            .unwrap();
+
+        loop {
+            let version_after_compaction =
+                get_version(&mut self.sysdb, request.collection_id).await;
+            if version_after_compaction > version_before_compaction {
+                break;
+            }
+            // todo
+            tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+        }
+
+        Ok(ManualCompactionResponse {})
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum FromConfigError {
+    #[error("Failed to create gRPC client: {0}")]
+    GrpcClient(#[from] tonic::transport::Error),
+}
+
+impl ChromaError for FromConfigError {
+    fn code(&self) -> chroma_error::ErrorCodes {
+        match self {
+            FromConfigError::GrpcClient(_) => chroma_error::ErrorCodes::Internal,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl Configurable<(GrpcCompactionClientConfig, System)> for GrpcCompactionClient {
+    async fn try_from_config(
+        (config, _system): &(GrpcCompactionClientConfig, System),
+        registry: &Registry,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        let client = CompactorClient::connect(config.url.clone())
+            .await
+            .map_err(|err| FromConfigError::GrpcClient(err).boxed())?;
+        Ok(Self {
+            client,
+            sysdb: registry.get().map_err(|err| err.boxed())?,
+        })
+    }
+}

rust/frontend/src/compaction_client/local_compaction_client.rs

@@ -0,0 +1,45 @@
+use chroma_config::{registry::Registry, Configurable};
+use chroma_error::ChromaError;
+use chroma_log::LocalCompactionManager;
+use chroma_system::{ComponentHandle, System};
+use chroma_types::{ManualCompactionError, ManualCompactionRequest, ManualCompactionResponse};
+
+#[derive(Debug, Clone)]
+pub(super) struct LocalCompactionClient {
+    handle: ComponentHandle<LocalCompactionManager>,
+}
+
+impl LocalCompactionClient {
+    pub async fn manually_compact(
+        &mut self,
+        request: ManualCompactionRequest,
+    ) -> Result<ManualCompactionResponse, ManualCompactionError> {
+        // let compact_message = CompactionMessage {
+        //     collection_id: request.collection_id,
+        //     start_offset: start_log_offset,
+        //     total_records,
+        // };
+        // self.handle
+        //     .request(compact_message, None)
+        //     .await
+        //     .unwrap()
+        //     .unwrap();
+        // todo
+
+        Ok(ManualCompactionResponse {})
+    }
+}
+
+#[async_trait::async_trait]
+impl Configurable<((), System)> for LocalCompactionClient {
+    async fn try_from_config(
+        (_config, _system): &((), System),
+        registry: &Registry,
+    ) -> Result<Self, Box<dyn ChromaError>> {
+        // Assume the registry has a compaction manager handle
+        let handle = registry
+            .get::<ComponentHandle<LocalCompactionManager>>()
+            .map_err(|err| err.boxed())?;
+        Ok(Self { handle })
+    }
+}

rust/frontend/src/compaction_client/mod.rs

@@ -0,0 +1,3 @@
+pub mod compaction_client;
+mod grpc_compaction_client;
+mod local_compaction_client;

rust/frontend/src/config.rs

@@ -1,4 +1,5 @@
 use crate::{
+    compaction_client::compaction_client::CompactionClientConfig,
     executor::config::{ExecutorConfig, LocalExecutorConfig},
     CollectionsWithSegmentsProviderConfig,
 };
@@ -61,6 +62,7 @@ pub struct FrontendConfig {
     pub log: LogConfig,
     #[serde(default = "default_executor_config")]
     pub executor: ExecutorConfig,
+    pub compaction_client: CompactionClientConfig,
 }
 
 fn default_otel_service_name() -> String {
@@ -164,8 +166,8 @@ mod tests {
         let sysdb_config = config.frontend.sysdb;
         let sysdb_config = match sysdb_config {
             chroma_sysdb::SysDbConfig::Grpc(grpc_sys_db_config) => grpc_sys_db_config,
-            chroma_sysdb::SysDbConfig::Sqlite(_) => {
-                panic!("Expected grpc sysdb config, got sqlite sysdb config")
+            _ => {
+                panic!("Expected grpc sysdb config")
             }
         };
         assert_eq!(sysdb_config.host, "sysdb.chroma");

rust/frontend/src/impls/service_based_frontend.rs

@@ -1,6 +1,6 @@
 use crate::{
-    config::FrontendConfig, executor::Executor, types::errors::ValidationError,
-    CollectionsWithSegmentsProvider,
+    compaction_client::compaction_client::CompactionClient, config::FrontendConfig,
+    executor::Executor, types::errors::ValidationError, CollectionsWithSegmentsProvider,
 };
 use backon::Retryable;
 use chroma_config::{registry, Configurable};
@@ -25,9 +25,10 @@ use chroma_types::{
     GetCollectionsError, GetDatabaseError, GetDatabaseRequest, GetDatabaseResponse, GetRequest,
     GetResponse, GetTenantError, GetTenantRequest, GetTenantResponse, HealthCheckResponse,
     HeartbeatError, HeartbeatResponse, Include, ListCollectionsRequest, ListCollectionsResponse,
-    ListDatabasesError, ListDatabasesRequest, ListDatabasesResponse, Metadata, Operation,
-    OperationRecord, QueryError, QueryRequest, QueryResponse, ResetError, ResetResponse, Segment,
-    SegmentScope, SegmentType, SegmentUuid, SingleNodeHnswParameters, UpdateCollectionError,
+    ListDatabasesError, ListDatabasesRequest, ListDatabasesResponse, ManualCompactionError,
+    ManualCompactionRequest, ManualCompactionResponse, Metadata, Operation, OperationRecord,
+    QueryError, QueryRequest, QueryResponse, ResetError, ResetResponse, Segment, SegmentScope,
+    SegmentType, SegmentUuid, SingleNodeHnswParameters, UpdateCollectionError,
     UpdateCollectionRecordsError, UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse,
     UpdateCollectionRequest, UpdateCollectionResponse, UpsertCollectionRecordsError,
     UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse, Where,
@@ -54,6 +55,7 @@ pub struct ServiceBasedFrontend {
     executor: Executor,
     log_client: Log,
     sysdb_client: SysDb,
+    compaction_client: CompactionClient,
     collections_with_segments_provider: CollectionsWithSegmentsProvider,
     max_batch_size: u32,
     metrics: Arc<Metrics>,
@@ -63,6 +65,7 @@ impl ServiceBasedFrontend {
     pub fn new(
         allow_reset: bool,
         sysdb_client: SysDb,
+        compaction_client: CompactionClient,
         collections_with_segments_provider: CollectionsWithSegmentsProvider,
         log_client: Log,
         executor: Executor,
@@ -84,6 +87,7 @@ impl ServiceBasedFrontend {
             executor,
             log_client,
             sysdb_client,
+            compaction_client,
             collections_with_segments_provider,
             max_batch_size,
             metrics,
@@ -194,6 +198,13 @@ impl ServiceBasedFrontend {
         Ok(ResetResponse {})
     }
 
+    pub async fn manually_compact(
+        &mut self,
+        request: ManualCompactionRequest,
+    ) -> Result<ManualCompactionResponse, ManualCompactionError> {
+        self.compaction_client.manually_compact(request).await
+    }
+
     pub async fn create_tenant(
         &mut self,
         CreateTenantRequest { name, .. }: CreateTenantRequest,
@@ -1089,9 +1100,16 @@ impl Configurable<(FrontendConfig, System)> for ServiceBasedFrontend {
         let executor =
             Executor::try_from_config(&(config.executor.clone(), system.clone()), registry).await?;
 
+        let compaction_client = CompactionClient::try_from_config(
+            &(config.compaction_client.clone(), system.clone()),
+            registry,
+        )
+        .await?;
+
         Ok(ServiceBasedFrontend::new(
             config.allow_reset,
             sysdb,
+            compaction_client,
             collections_with_segments_provider,
             log,
             executor,

rust/frontend/src/lib.rs

@@ -2,6 +2,7 @@ use std::{path::Path, sync::Arc};
 
 mod ac;
 pub mod auth;
+pub mod compaction_client; // todo
 pub mod config;
 #[allow(dead_code)]
 pub mod executor;

rust/python_bindings/src/bindings.rs

@@ -103,6 +103,8 @@ impl Bindings {
             collections_with_segments_provider: collection_cache_config,
             log: log_config,
             executor: executor_config,
+            compaction_client:
+                chroma_frontend::compaction_client::compaction_client::CompactionClientConfig::Local,
         };
 
         let frontend = runtime.block_on(async {

rust/types/src/api_types.rs

@@ -115,6 +115,25 @@ impl ChromaError for ResetError {
     }
 }
 
+#[non_exhaustive]
+#[derive(Validate)]
+pub struct ManualCompactionRequest {
+    pub collection_id: CollectionUuid,
+}
+
+impl ManualCompactionRequest {
+    pub fn try_new(collection_id: CollectionUuid) -> Result<Self, ChromaValidationError> {
+        let request = Self { collection_id };
+        request.validate().map_err(ChromaValidationError::from)?;
+        Ok(request)
+    }
+}
+
+pub struct ManualCompactionResponse {}
+
+#[derive(Error, Debug)]
+pub enum ManualCompactionError {}
+
 #[derive(Serialize, ToSchema)]
 pub struct ChecklistResponse {
     pub max_batch_size: u32,

rust/worker/src/compactor/compaction_manager.rs

@@ -340,6 +340,10 @@ impl Handler<OneOffCompactionMessage> for CompactionManager {
             "One-off collections queued: {:?}",
             self.scheduler.get_oneoff_collections()
         );
+
+        tracing::info!("CompactionManager: Performing scheduled compaction");
+        let ids = self.compact_batch().await;
+        self.hnsw_index_provider.purge_by_id(&ids).await;
     }
 }