[ENH]: Add Rust task client and execution operators (#5686)

## Description of changes

This change is second in a series of changes to implement the TaskRunner. This PR does the actual running of tasks and makes the required modifications to the Compaction operators for this to work.

- Improvements & Bug fixes

- New functionality
    - Added `TaskUuid`, `NonceUuid`, and `JobUuid` types to work alongside `CollectionUuid`​ in the compactor code.
    - Created new operators for task execution:
        - `PrepareTaskOperator`: Prepares a task for execution. Resolves a taskuuid to a Task object. Updates `next_nonce`​ if needed.
        - `ExecuteTaskOperator`: Executes task logic on input records
        - `FinishTaskOperator`: Updates task state after execution
- Added `TaskExecutor` trait and `CountTask` implementation
- Enhanced `GetCollectionAndSegmentsOperator` to fetch both input and output collections
- Updated `RegisterOperator` to handle task-based compactions

## Test plan

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

## Migration plan

No migrations needed. The changes maintain backward compatibility with existing code paths while adding new functionality for task execution.

## Observability plan

## Documentation Changes

No user-facing API changes requiring documentation updates.

Author: tanujnay112

Cargo.toml

@@ -51,7 +51,7 @@ tracing = { version = "0.1" }
 tracing-bunyan-formatter = "0.3"
 tracing-opentelemetry = "0.28.0"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
-uuid = { version = "1.11.0", features = ["v4", "fast-rng", "macro-diagnostics", "serde"] }
+uuid = { version = "1.11.0", features = ["v4", "v7", "fast-rng", "macro-diagnostics", "serde"] }
 utoipa = { version = "5", features = ["macros", "axum_extras", "debug", "uuid"] }
 sqlx = { version = "0.8.3", features = ["runtime-tokio", "sqlite", "postgres", "chrono"] }
 sha2 = "0.10.8"

Tiltfile

@@ -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', resource_deps=['minio-deployment'])
+k8s_resource('rust-log-service', labels=["chroma"], port_forwards=['50054:50051', '50052:50052'], 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')

go/pkg/sysdb/coordinator/task.go

@@ -2,6 +2,7 @@ package coordinator
 
 import (
 	"context"
+	"math"
 	"strings"
 	"time"
 
@@ -195,6 +196,15 @@ func (s *Coordinator) GetTaskByName(ctx context.Context, req *coordinatorpb.GetT
 		}
 	}
 
+	// Validate completion_offset is non-negative before converting to uint64
+	if task.CompletionOffset < 0 {
+		log.Error("GetTaskByName: invalid completion_offset",
+			zap.String("task_id", task.ID.String()),
+			zap.Int64("completion_offset", task.CompletionOffset))
+		return nil, status.Errorf(codes.Internal,
+			"task has invalid completion_offset: %d", task.CompletionOffset)
+	}
+
 	// Convert task to response with nested Task message
 	taskProto := &coordinatorpb.Task{
 		TaskId:               task.ID.String(),
@@ -203,13 +213,15 @@ func (s *Coordinator) GetTaskByName(ctx context.Context, req *coordinatorpb.GetT
 		InputCollectionId:    task.InputCollectionID,
 		OutputCollectionName: task.OutputCollectionName,
 		Params:               paramsStruct,
-		CompletionOffset:     task.CompletionOffset,
+		CompletionOffset:     uint64(task.CompletionOffset),
 		MinRecordsForTask:    uint64(task.MinRecordsForTask),
 		TenantId:             task.TenantID,
 		DatabaseId:           task.DatabaseID,
 		NextRunAt:            uint64(task.NextRun.UnixMicro()),
 		LowestLiveNonce:      "",
 		NextNonce:            task.NextNonce.String(),
+		CreatedAt:            uint64(task.CreatedAt.UnixMicro()),
+		UpdatedAt:            uint64(task.UpdatedAt.UnixMicro()),
 	}
 	// Add lowest_live_nonce if it's set
 	if task.LowestLiveNonce != nil {
@@ -269,6 +281,15 @@ func (s *Coordinator) GetTaskByUuid(ctx context.Context, req *coordinatorpb.GetT
 		}
 	}
 
+	// Validate completion_offset is non-negative before converting to uint64
+	if task.CompletionOffset < 0 {
+		log.Error("GetTaskByUuid: invalid completion_offset",
+			zap.String("task_id", task.ID.String()),
+			zap.Int64("completion_offset", task.CompletionOffset))
+		return nil, status.Errorf(codes.Internal,
+			"task has invalid completion_offset: %d", task.CompletionOffset)
+	}
+
 	// Convert task to response with nested Task message
 	taskProto := &coordinatorpb.Task{
 		TaskId:               task.ID.String(),
@@ -277,13 +298,15 @@ func (s *Coordinator) GetTaskByUuid(ctx context.Context, req *coordinatorpb.GetT
 		InputCollectionId:    task.InputCollectionID,
 		OutputCollectionName: task.OutputCollectionName,
 		Params:               paramsStruct,
-		CompletionOffset:     task.CompletionOffset,
+		CompletionOffset:     uint64(task.CompletionOffset),
 		MinRecordsForTask:    uint64(task.MinRecordsForTask),
 		TenantId:             task.TenantID,
 		DatabaseId:           task.DatabaseID,
 		NextRunAt:            uint64(task.NextRun.UnixMicro()),
 		LowestLiveNonce:      "",
 		NextNonce:            task.NextNonce.String(),
+		CreatedAt:            uint64(task.CreatedAt.UnixMicro()),
+		UpdatedAt:            uint64(task.UpdatedAt.UnixMicro()),
 	}
 	// Add lowest_live_nonce if it's set
 	if task.LowestLiveNonce != nil {
@@ -480,16 +503,34 @@ func (s *Coordinator) AdvanceTask(ctx context.Context, req *coordinatorpb.Advanc
 		return nil, status.Errorf(codes.InvalidArgument, "invalid task_run_nonce: %v", err)
 	}
 
-	advanceTask, err := s.catalog.metaDomain.TaskDb(ctx).AdvanceTask(taskID, taskRunNonce, *req.CompletionOffset, *req.NextRunDelaySecs)
+	// Validate completion_offset fits in int64 before storing in database
+	if *req.CompletionOffset > uint64(math.MaxInt64) { // math.MaxInt64
+		log.Error("AdvanceTask: completion_offset too large",
+			zap.Uint64("completion_offset", *req.CompletionOffset))
+		return nil, status.Errorf(codes.InvalidArgument,
+			"completion_offset too large: %d", *req.CompletionOffset)
+	}
+	completionOffsetInt64 := int64(*req.CompletionOffset)
+
+	advanceTask, err := s.catalog.metaDomain.TaskDb(ctx).AdvanceTask(taskID, taskRunNonce, completionOffsetInt64, *req.NextRunDelaySecs)
 	if err != nil {
 		log.Error("AdvanceTask failed", zap.Error(err), zap.String("task_id", taskID.String()))
 		return nil, err
 	}
 
+	// Validate completion_offset from database is non-negative before converting to uint64
+	if advanceTask.CompletionOffset < 0 {
+		log.Error("AdvanceTask: invalid completion_offset from database",
+			zap.String("task_id", taskID.String()),
+			zap.Int64("completion_offset", advanceTask.CompletionOffset))
+		return nil, status.Errorf(codes.Internal,
+			"task has invalid completion_offset: %d", advanceTask.CompletionOffset)
+	}
+
 	return &coordinatorpb.AdvanceTaskResponse{
 		NextRunNonce:     advanceTask.NextNonce.String(),
 		NextRunAt:        uint64(advanceTask.NextRun.UnixMilli()),
-		CompletionOffset: advanceTask.CompletionOffset,
+		CompletionOffset: uint64(advanceTask.CompletionOffset),
 	}, nil
 }
 

go/pkg/sysdb/grpc/collection_service.go

@@ -3,6 +3,7 @@ package grpc
 import (
 	"context"
 	"encoding/json"
+	"math"
 
 	"github.com/chroma-core/chroma/go/pkg/grpcutils"
 
@@ -607,6 +608,14 @@ func (s *Server) FlushCollectionCompactionAndTask(ctx context.Context, req *coor
 		return nil, grpcutils.BuildInternalGrpcError(err.Error())
 	}
 
+	// Validate completion_offset fits in int64 before storing in database
+	if taskUpdate.CompletionOffset > uint64(math.MaxInt64) {
+		log.Error("FlushCollectionCompactionAndTask: completion_offset too large",
+			zap.Uint64("completion_offset", taskUpdate.CompletionOffset))
+		return nil, grpcutils.BuildInternalGrpcError("completion_offset too large")
+	}
+	completionOffsetSigned := int64(taskUpdate.CompletionOffset)
+
 	segmentCompactionInfo := make([]*model.FlushSegmentCompaction, 0, len(flushReq.SegmentCompactionInfo))
 	for _, flushSegmentCompaction := range flushReq.SegmentCompactionInfo {
 		segmentID, err := types.ToUniqueID(&flushSegmentCompaction.SegmentId)
@@ -640,7 +649,7 @@ func (s *Server) FlushCollectionCompactionAndTask(ctx context.Context, req *coor
 		flushCollectionCompaction,
 		taskID,
 		taskRunNonce,
-		taskUpdate.CompletionOffset,
+		completionOffsetSigned,
 	)
 	if err != nil {
 		log.Error("FlushCollectionCompactionAndTask failed", zap.Error(err), zap.String("collection_id", flushReq.CollectionId), zap.String("task_id", taskUpdate.TaskId))
@@ -667,7 +676,13 @@ func (s *Server) FlushCollectionCompactionAndTask(ctx context.Context, req *coor
 		res.NextRun = timestamppb.New(*flushCollectionInfo.TaskNextRun)
 	}
 	if flushCollectionInfo.TaskCompletionOffset != nil {
-		res.CompletionOffset = *flushCollectionInfo.TaskCompletionOffset
+		// Validate completion_offset is non-negative before converting to uint64
+		if *flushCollectionInfo.TaskCompletionOffset < 0 {
+			log.Error("FlushCollectionCompactionAndTask: invalid completion_offset",
+				zap.Int64("completion_offset", *flushCollectionInfo.TaskCompletionOffset))
+			return nil, grpcutils.BuildInternalGrpcError("task has invalid completion_offset")
+		}
+		res.CompletionOffset = uint64(*flushCollectionInfo.TaskCompletionOffset)
 	}
 
 	return res, nil

go/pkg/sysdb/metastore/db/dao/task.go

@@ -138,7 +138,7 @@ func (s *taskDb) AdvanceTask(taskID uuid.UUID, taskRunNonce uuid.UUID, completio
 	// Bump next_nonce to mark a new run, but don't touch lowest_live_nonce yet
 	// lowest_live_nonce will be updated later by finish_task when verification completes
 	next_run := now.Add(time.Duration(nextRunDelaySecs) * time.Second)
-	result := s.db.Model(&dbmodel.Task{}).Where("task_id = ?", taskID).Where("is_deleted = false").Where("next_nonce = ?", taskRunNonce).Where("completion_offset < ?", completionOffset).UpdateColumns(map[string]interface{}{
+	result := s.db.Model(&dbmodel.Task{}).Where("task_id = ?", taskID).Where("is_deleted = false").Where("next_nonce = ?", taskRunNonce).Where("completion_offset <= ?", completionOffset).UpdateColumns(map[string]interface{}{
 		"completion_offset": completionOffset,
 		"next_run":          next_run,
 		"last_run":          now,

idl/chromadb/proto/coordinator.proto

@@ -325,7 +325,7 @@ message FlushCollectionCompactionResponse {
 message TaskUpdateInfo {
   string task_id = 1;
   string task_run_nonce = 2;
-  int64 completion_offset = 3;
+  uint64 completion_offset = 3;
 }
 
 // Combined request to flush collection compaction and update task atomically in a single transaction
@@ -341,7 +341,7 @@ message FlushCollectionCompactionAndTaskResponse {
   // Updated task fields from database (authoritative)
   string next_nonce = 4;
   google.protobuf.Timestamp next_run = 5;
-  int64 completion_offset = 6;
+  uint64 completion_offset = 6;
 }
 
 // Used for serializing contents in collection version history file.
@@ -588,13 +588,15 @@ message Task {
   string output_collection_name = 5;
   optional string output_collection_id = 6;
   optional google.protobuf.Struct params = 7;
-  int64 completion_offset = 8;
+  uint64 completion_offset = 8;
   uint64 min_records_for_task = 9;
   string tenant_id = 10;
   string database_id = 11;
   uint64 next_run_at = 12;
   string lowest_live_nonce = 13;
   string next_nonce = 14;
+  uint64 created_at = 15;
+  uint64 updated_at = 16;
 }
 
 message GetTaskByNameResponse {
@@ -623,14 +625,14 @@ message AdvanceTaskRequest {
   optional string collection_id = 1;
   optional string task_id = 2;
   optional string task_run_nonce = 3;
-  optional int64 completion_offset = 4;
+  optional uint64 completion_offset = 4;
   optional uint64 next_run_delay_secs = 5;
 }
 
 message AdvanceTaskResponse {
   string next_run_nonce = 1;
   uint64 next_run_at = 2;
-  int64 completion_offset = 3;
+  uint64 completion_offset = 3;
 }
 
 message FinishTaskRequest {

rust/s3heap-service/src/lib.rs

@@ -269,7 +269,7 @@ impl HeapTender {
                         let schedule = Schedule {
                             triggerable,
                             next_scheduled,
-                            nonce: s.task_run_nonce,
+                            nonce: s.task_run_nonce.0,
                         };
                         Ok(Some(schedule))
                     } else {

rust/s3heap-service/src/scheduler.rs

@@ -75,7 +75,7 @@ impl HeapScheduler for SysDbScheduler {
                         partitioning: schedule.collection_id.0.into(),
                         scheduling: schedule.task_id.into(),
                     },
-                    nonce: schedule.task_run_nonce,
+                    nonce: schedule.task_run_nonce.0,
                     next_scheduled: when_to_run,
                 });
             }

rust/sysdb/src/bin/chroma-task-manager.rs

@@ -169,6 +169,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             println!("Operator: {:?}", task.operator_name);
             println!("Input Collection: {:?}", task.input_collection_id);
             println!("Output Collection Name: {:?}", task.output_collection_name);
+            println!("Output Collection ID: {:?}", task.output_collection_id);
             println!("Params: {:?}", task.params);
             println!("Completion Offset: {:?}", task.completion_offset);
             println!("Min Records: {:?}", task.min_records_for_task);
@@ -198,7 +199,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 collection_id: Some(collection_id),
                 task_id: Some(task_id),
                 task_run_nonce: Some(task_run_nonce),
-                completion_offset: Some(completion_offset.try_into().unwrap()),
+                completion_offset: Some(completion_offset),
                 next_run_delay_secs: Some(next_run_delay_secs),
             };