[ENH]: Integrate task operators into compaction (#5687)

## Description of changes

This change implements the TaskRunner that pulls scheduled Tasks from the heap and runs them. This reuses a large part of the CompactionManager's code with the following additions:

- The CompactionManager takes in a strongly typed binary enum to specify whether it is intended to service Tasks or usual collection Compactions. All the changes are meant to take effect in Tasks mode.
- The scheduler regularly pulls tasks from the S3 heap and adds them to its internal queue that the overlying CompactionManager polls.
- CompactionOrchestrator has been refactored to not just take in an input_collection to compact but an input_collection and output_collection. Its preexisting operators have been refactored as such as well.
- For each scheduled task the Task-related CompactionManager finds, it kicks off a CompactionOrchestrator with PrepareTask as the initial operator. PrepareTask does the following:
    - It takes in a task uuid and task nonce to resolve a SysDB Task object. It creates an output collection if needed.
    - It marks the beginning of this nonce's execution by generating a new value for `next_nonce`​ in the `Tasks`​ table in the SysDB.
        - It detects if this step had already been done before by seeing whether `Tasks.lowest_live_nonce != Tasks.next_nonce`​.
            - If the two are not equal, the previously described increment to `next_nonce`​ does not occur. This case implies that there was an incomplete execution of `lowest_live_nonce`​. PrepareTask compares the `scout_logs`​ result of the given input collection with `Tasks.completion_offset`​ to determine whether this Task needs to be re-executed or whether it can skip to the FinishTask operator.
- The Register operator has been edited to update the current Task's completion offset in the same transaction where the output collection's records are flushed to SysDB.
- A FinishTask operator has been added to run after the Register operator in Task-related CompactionOrchestrators. This operator rechecks scout_logs for the given input collection. If the resulting offset is too far ahead of the `completion_offset`​ that we sent to SysDB in Register, we will schedule a new item in the heap to do this task again at a later time. We then unconditionally set `lowest_live_nonce`​ to `next_nonce`​ in the SysDB, marking `lowest_live_nonce`​ as completed and allowing all heap entries below the new value of `lowest_live_nonce`​ to be pruned.
- Improvements & Bug fixes
    - ^^^^
- New functionality
    - The TaskRunner's operation as a separate system component in the Compactor node is disabled by default it can be enabled by the following configuration in the compaction_service config.

```
  task_runner:
    enabled: true
```

Relevant TODOs after this:

- Adjust scheduler to be aware of `lowest_live_nonce`
- Send a new heap item in `FinishTask`​ if the `scout_logs`​ check deems it necessary.
- Adjust `create_task`​ to do a 2-phase commit and adjust the TaskRunner logic to handle partially created tasks. `create_task`​ must also send over an initial backfill item to the heap if its input collection existed before this Task.
- Validate operator parameters and pull in operator data during TaskRunner execution.
- TaskRunner only pulls data from compacted segments of the input collection. It is possible to also pull log data in this process.
- Support for incremental tasks.

## Test plan

_How are these changes tested?_

An integration test has been added in [compact.rs](http://compact.rs).

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js, `cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the_ [_docs section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: tanujnay112

Cargo.lock

@@ -237,7 +237,7 @@ k8s_resource(
 k8s_resource('postgres', resource_deps=['k8s_setup'], labels=["infrastructure"], port_forwards='5432:5432')
 # Jobs are suffixed with the image tag to ensure they are unique. In this context, the image tag is defined in k8s/distributed-chroma/values.yaml.
 k8s_resource('sysdb-migration-latest', resource_deps=['postgres'], labels=["infrastructure"])
-k8s_resource('rust-log-service', labels=["chroma"], port_forwards=['50054:50051', '50052:50052'], resource_deps=['minio-deployment'])
+k8s_resource('rust-log-service', labels=["chroma"], port_forwards=['50054:50051'], resource_deps=['minio-deployment'])
 k8s_resource('sysdb', resource_deps=['sysdb-migration-latest'], labels=["chroma"], port_forwards='50051:50051')
 k8s_resource('rust-frontend-service', resource_deps=['sysdb', 'rust-log-service'], labels=["chroma"], port_forwards='8000:8000')
 k8s_resource('query-service', resource_deps=['sysdb'], labels=["chroma"], port_forwards='50053:50051')

Tiltfile

@@ -7,6 +7,7 @@
 """
 
 import chromadb
+import time
 
 # Connect to Chroma server
 client = chromadb.HttpClient(host="localhost", port=8000)
@@ -60,6 +61,8 @@
 print("Task is now registered and will run on new data!")
 print("=" * 60)
 
+time.sleep(10)
+
 # Add more documents to trigger task execution
 print("\nAdding more documents...")
 collection.add(

examples/task_api_example.py

@@ -1388,7 +1388,10 @@ impl LogServer {
         if records.len() != pull_logs.batch_size as usize
             || (!records.is_empty() && records[0].log_offset != pull_logs.start_from_offset)
         {
-            return Err(Status::not_found("Some entries have been purged"));
+            return Err(Status::not_found(format!(
+                "Some entries have been purged {} versus {}",
+                records[0].log_offset, pull_logs.start_from_offset
+            )));
         }
         Ok(Response::new(PullLogsResponse { records }))
     }

rust/log-service/src/lib.rs

@@ -14,6 +14,7 @@ serde_json = { workspace = true }
 tokio = { workspace = true }
 tonic = { workspace = true }
 tonic-health = { workspace = true }
+tower = { workspace = true }
 tracing = { workspace = true }
 
 chroma-config = { workspace = true }

rust/s3heap-service/Cargo.toml

@@ -6,6 +6,26 @@ use uuid::Uuid;
 
 use crate::CollectionUuid;
 
+define_uuid_newtype!(
+    /// JobId is a wrapper around Uuid to provide a unified type for job identifiers.
+    /// Jobs can be either collection compaction jobs or task execution jobs.
+    JobId,
+    new_v4
+);
+
+// Custom From implementations for JobId
+impl From<CollectionUuid> for JobId {
+    fn from(collection_uuid: CollectionUuid) -> Self {
+        JobId(collection_uuid.0)
+    }
+}
+
+impl From<TaskUuid> for JobId {
+    fn from(task_uuid: TaskUuid) -> Self {
+        JobId(task_uuid.0)
+    }
+}
+
 define_uuid_newtype!(
     /// TaskUuid is a wrapper around Uuid to provide a type for task identifiers.
     #[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]

rust/types/src/task.rs

@@ -79,6 +79,7 @@ shuttle = { workspace = true }
 rand = { workspace = true }
 rand_xorshift = { workspace = true }
 tempfile = { workspace = true }
+reqwest = { workspace = true, features = ["json"] }
 
 chroma-benchmark = { workspace = true }