[ENH]: Refactor compactor

Author: tanujnay112

Cargo.lock

@@ -462,7 +462,7 @@ impl<'log_data, 'segment_data: 'log_data> HydratedMaterializedLogRecord<'log_dat
             | MaterializedLogOperation::OverwriteExisting
             | MaterializedLogOperation::UpdateExisting => {
                 (self.get_user_id().len()
-                    + size_of_val(self.merged_embeddings_ref())
+                    + self.merged_embeddings_ref().len() * size_of::<f32>()
                     + logical_size_of_metadata(&merged_metadata)
                     + self
                         .merged_document_ref()

rust/segment/src/types.rs

@@ -2,7 +2,6 @@ use super::scheduler::Scheduler;
 use super::scheduler_policy::LasCompactionTimeSchedulerPolicy;
 use super::OneOffCompactMessage;
 use super::RebuildMessage;
-use crate::compactor::tasks::SchedulableFunction;
 use crate::compactor::types::{ListDeadJobsMessage, ScheduledCompactMessage};
 use crate::config::CompactionServiceConfig;
 use crate::execution::operators::purge_dirty_log::PurgeDirtyLog;
@@ -13,8 +12,7 @@ use crate::execution::operators::repair_log_offsets::RepairLogOffsets;
 use crate::execution::operators::repair_log_offsets::RepairLogOffsetsError;
 use crate::execution::operators::repair_log_offsets::RepairLogOffsetsInput;
 use crate::execution::operators::repair_log_offsets::RepairLogOffsetsOutput;
-use crate::execution::orchestration::CompactOrchestrator;
-use crate::execution::orchestration::CompactionResponse;
+use crate::execution::orchestration::compact::{compact, CompactionResponse};
 use async_trait::async_trait;
 use chroma_blockstore::provider::BlockfileProvider;
 use chroma_config::assignment::assignment_policy::AssignmentPolicy;
@@ -27,11 +25,9 @@ use chroma_memberlist::memberlist_provider::Memberlist;
 use chroma_segment::spann_provider::SpannProvider;
 use chroma_storage::Storage;
 use chroma_sysdb::SysDb;
-use chroma_system::wrap;
-use chroma_system::Dispatcher;
-use chroma_system::Orchestrator;
-use chroma_system::TaskResult;
-use chroma_system::{Component, ComponentContext, ComponentHandle, Handler, System};
+use chroma_system::{
+    wrap, Component, ComponentContext, ComponentHandle, Dispatcher, Handler, System, TaskResult,
+};
 use chroma_types::{CollectionUuid, JobId};
 use futures::stream::FuturesUnordered;
 use futures::FutureExt;
@@ -118,8 +114,6 @@ pub(crate) struct CompactionManager {
 pub(crate) enum CompactionError {
     #[error("Failed to compact")]
     FailedToCompact,
-    #[error("Failed to execute task")]
-    FailedToExecuteTask,
     #[error("Heap service is not initialized for task based compaction")]
     HeapServiceNotInitialized,
 }
@@ -128,7 +122,6 @@ impl ChromaError for CompactionError {
     fn code(&self) -> ErrorCodes {
         match self {
             CompactionError::FailedToCompact => ErrorCodes::Internal,
-            CompactionError::FailedToExecuteTask => ErrorCodes::Internal,
             CompactionError::HeapServiceNotInitialized => ErrorCodes::InvalidArgument,
         }
     }
@@ -211,70 +204,33 @@ impl CompactionManager {
         let compact_awaiter_channel = &self.compact_awaiter_channel;
         self.scheduler.schedule().await;
 
-        match self.mode {
-            ExecutionMode::Compaction => {
-                let jobs: Vec<_> = self.scheduler.get_jobs().cloned().collect();
-                for job in jobs {
-                    let instrumented_span = span!(
-                        parent: None,
-                        tracing::Level::INFO,
-                        "Compacting job",
-                        collection_id = ?job.collection_id
-                    );
-                    Span::current()
-                        .add_link(instrumented_span.context().span().span_context().clone());
-
-                    let future = self
-                        .context
-                        .clone()
-                        .compact(job.collection_id, false)
-                        .instrument(instrumented_span);
-                    if let Err(e) = compact_awaiter_channel
-                        .send(CompactionTask {
-                            job_id: job.collection_id.into(),
-                            future: Box::pin(future),
-                        })
-                        .await
-                    {
-                        tracing::error!(
-                            collection_id = ?job.collection_id,
-                            error = ?e,
-                            "Failed to send start scheduled compaction task"
-                        );
-                    }
-                }
-            }
-            ExecutionMode::AttachedFunction => {
-                let tasks_iter = self.scheduler.get_tasks_scheduled_for_execution().clone();
-                for task in tasks_iter {
-                    let instrumented_span = span!(
-                        parent: None,
-                        tracing::Level::INFO,
-                        "Compacting task",
-                        collection_id = ?task.collection_id
-                    );
-                    Span::current()
-                        .add_link(instrumented_span.context().span().span_context().clone());
-
-                    let future = self
-                        .context
-                        .clone()
-                        .execute_task(task.clone())
-                        .instrument(instrumented_span);
-                    if let Err(e) = compact_awaiter_channel
-                        .send(CompactionTask {
-                            job_id: task.task_id.into(),
-                            future: Box::pin(future),
-                        })
-                        .await
-                    {
-                        tracing::error!(
-                            task_id = ?task.task_id,
-                            error = ?e,
-                            "Failed to start scheduled task run"
-                        );
-                    }
-                }
+        let jobs: Vec<_> = self.scheduler.get_jobs().cloned().collect();
+        for job in jobs {
+            let instrumented_span = span!(
+                parent: None,
+                tracing::Level::INFO,
+                "Compacting job",
+                collection_id = ?job.collection_id
+            );
+            Span::current().add_link(instrumented_span.context().span().span_context().clone());
+
+            let future = self
+                .context
+                .clone()
+                .compact(job.collection_id, false)
+                .instrument(instrumented_span);
+            if let Err(e) = compact_awaiter_channel
+                .send(CompactionTask {
+                    job_id: job.collection_id.into(),
+                    future: Box::pin(future),
+                })
+                .await
+            {
+                tracing::error!(
+                    collection_id = ?job.collection_id,
+                    error = ?e,
+                    "Failed to send start scheduled compaction task"
+                );
             }
         }
     }
@@ -380,17 +336,17 @@ impl CompactionManager {
         while let Ok(resp) = compact_awaiter_completion_channel.try_recv() {
             match resp.result {
                 Ok(ref compaction_response) => match compaction_response {
-                    CompactionResponse::Success { job_id } => {
-                        if job_id != &resp.job_id.0 {
+                    CompactionResponse::Success { job_id, .. } => {
+                        if job_id != &resp.job_id {
                             tracing::event!(Level::ERROR, name = "mismatched job ids in result", lhs =? *job_id, rhs =? resp.job_id);
                         }
                         self.scheduler.succeed_job(resp.job_id);
                     }
                     CompactionResponse::RequireCompactionOffsetRepair {
-                        collection_id,
+                        job_id: collection_id,
                         witnessed_offset_in_sysdb,
                     } => {
-                        if collection_id.0 != resp.job_id.0 {
+                        if *collection_id != resp.job_id {
                             tracing::event!(Level::ERROR, name = "mismatched job ids in result", lhs =? *collection_id, rhs =? resp.job_id);
                             self.scheduler.fail_job(resp.job_id);
                         } else {
@@ -434,8 +390,11 @@ impl CompactionManagerContext {
             }
         };
 
-        let orchestrator = CompactOrchestrator::new(
-            collection_id, // input_collection_id
+        // fetch data to compact -> execute_task/compact -> register
+        // Use the compact function to handle the entire orchestration process
+        let compact_result = compact(
+            self.system.clone(),
+            collection_id,
             rebuild,
             self.fetch_log_batch_size,
             self.max_compaction_size,
@@ -445,67 +404,17 @@ impl CompactionManagerContext {
             self.blockfile_provider.clone(),
             self.hnsw_index_provider.clone(),
             self.spann_provider.clone(),
-            dispatcher,
-            None,
-        );
-
-        match orchestrator.run(self.system.clone()).await {
-            Ok(result) => {
-                tracing::info!("Compaction Job completed: {:?}", result);
-                return Ok(result);
-            }
-            Err(e) => {
-                if e.should_trace_error() {
-                    tracing::error!("Compaction Job failed: {:?}", e);
-                }
-                return Err(Box::new(e));
-            }
-        }
-    }
-
-    async fn execute_task(self, task: SchedulableFunction) -> CompactionOutput {
-        tracing::info!("Executing task {}", task.task_id);
-        let dispatcher = match self.dispatcher {
-            Some(ref dispatcher) => dispatcher.clone(),
-            None => {
-                tracing::error!("No dispatcher found");
-                return Err(Box::new(CompactionError::FailedToExecuteTask));
-            }
-        };
+            dispatcher.clone(),
+        )
+        .await;
 
-        let orchestrator = CompactOrchestrator::new_for_attached_function(
-            task.collection_id,
-            false,
-            self.fetch_log_batch_size,
-            self.max_compaction_size,
-            self.max_partition_size,
-            self.log.clone(),
-            self.sysdb.clone(),
-            self.heap_service.ok_or_else(|| {
-                Box::new(CompactionError::HeapServiceNotInitialized) as Box<dyn ChromaError>
-            })?,
-            self.blockfile_provider.clone(),
-            self.hnsw_index_provider.clone(),
-            self.spann_provider.clone(),
-            dispatcher,
-            None,
-            task.task_id,
-            task.nonce,
-        );
-        match orchestrator.run(self.system.clone()).await {
-            Ok(result) => {
-                tracing::info!(
-                    " Attached Function {} completed: {:?}",
-                    task.task_id,
-                    result
-                );
-                Ok(result)
-            }
+        match compact_result {
+            Ok(response) => Ok(response),
             Err(e) => {
                 if e.should_trace_error() {
-                    tracing::error!(" Attached Function {} failed: {:?}", task.task_id, e);
+                    tracing::error!("Compaction failed: {:?}", e);
                 }
-                Err(Box::new(e))
+                Err(e)
             }
         }
     }
@@ -566,11 +475,9 @@ impl Configurable<(CompactionServiceConfig, System)> for CompactionManager {
         let job_expiry_seconds = config.compactor.job_expiry_seconds;
         let max_failure_count = config.compactor.max_failure_count;
         let scheduler = Scheduler::new(
-            ExecutionMode::Compaction, // Default to Compaction mode
             my_ip,
             log.clone(),
             sysdb.clone(),
-            storage.clone(),
             policy,
             max_concurrent_jobs,
             min_compaction_size,
@@ -678,11 +585,9 @@ pub(crate) async fn attach_functionrunner_manager(
     };
 
     let scheduler = Scheduler::new(
-        ExecutionMode::AttachedFunction, // Taskrunner mode
         my_ip,
         log.clone(),
         sysdb.clone(),
-        storage.clone(),
         policy,
         task_config.max_concurrent_jobs,
         0, // min_compaction_size not used for tasks
@@ -1265,6 +1170,7 @@ mod tests {
         while let Some(entry) = entries.next_entry().await.expect("Failed to read next dir") {
             let path = entry.path();
             let metadata = entry.metadata().await.expect("Failed to read metadata");
+            println!("{}", path.display());
 
             if metadata.is_dir() {
                 assert!(path.ends_with("tenant"));

rust/worker/src/compactor/compaction_manager.rs

@@ -2,7 +2,6 @@ mod compaction_manager;
 pub(crate) mod config;
 mod scheduler;
 mod scheduler_policy;
-mod tasks;
 mod types;
 
 pub(crate) use compaction_manager::*;