[ENH] Propagate version through system (#2839)

## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
- Pass the request_version_context into the protobufs and to the rust
code. In the rust orchestrators, we error on mismatches and propagate
this error up to python.
 - New functionality
	 - None

## 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

chromadb/proto/convert.py

@@ -9,6 +9,7 @@
     LogRecord,
     Metadata,
     Operation,
+    RequestVersionContext,
     ScalarEncoding,
     Segment,
     SegmentScope,
@@ -295,3 +296,21 @@ def from_proto_vector_query_result(
         distance=vector_query_result.distance,
         embedding=from_proto_vector(vector_query_result.vector)[0],
     )
+
+
+def from_proto_request_version_context(
+    request_version_context: proto.RequestVersionContext,
+) -> RequestVersionContext:
+    return RequestVersionContext(
+        collection_version=request_version_context.collection_version,
+        log_position=request_version_context.log_position,
+    )
+
+
+def to_proto_request_version_context(
+    request_version_context: RequestVersionContext,
+) -> proto.RequestVersionContext:
+    return proto.RequestVersionContext(
+        collection_version=request_version_context["collection_version"],
+        log_position=request_version_context["log_position"],
+    )

chromadb/segment/impl/metadata/grpc_segment.py

@@ -1,4 +1,5 @@
 from typing import Dict, List, Optional, Sequence
+from chromadb.proto.convert import to_proto_request_version_context
 from chromadb.proto.utils import RetryOnRpcErrorClientInterceptor
 from chromadb.segment import MetadataReader
 from chromadb.config import System
@@ -52,9 +53,11 @@ def count(self, request_version_context: RequestVersionContext) -> int:
         request: pb.CountRecordsRequest = pb.CountRecordsRequest(
             segment_id=self._segment["id"].hex,
             collection_id=self._segment["collection"].hex,
+            version_context=to_proto_request_version_context(request_version_context),
         )
         response: pb.CountRecordsResponse = self._metadata_reader_stub.CountRecords(
-            request, timeout=self._request_timeout_seconds
+            request,
+            timeout=self._request_timeout_seconds,
         )
         return response.count
 
@@ -104,10 +107,12 @@ def get_metadata(
             limit=limit,
             offset=offset,
             include_metadata=include_metadata,
+            version_context=to_proto_request_version_context(request_version_context),
         )
 
         response: pb.QueryMetadataResponse = self._metadata_reader_stub.QueryMetadata(
-            request, timeout=self._request_timeout_seconds
+            request,
+            timeout=self._request_timeout_seconds,
         )
         results: List[MetadataEmbeddingRecord] = []
         for record in response.records:

chromadb/segment/impl/vector/grpc_segment.py

@@ -4,6 +4,7 @@
 from chromadb.proto.convert import (
     from_proto_vector_embedding_record,
     from_proto_vector_query_result,
+    to_proto_request_version_context,
     to_proto_vector,
 )
 from chromadb.proto.utils import RetryOnRpcErrorClientInterceptor
@@ -64,9 +65,11 @@ def get_vectors(
             ids=ids,
             segment_id=self._segment["id"].hex,
             collection_id=self._segment["collection"].hex,
+            version_context=to_proto_request_version_context(request_version_context),
         )
         response: GetVectorsResponse = self._vector_reader_stub.GetVectors(
-            request, timeout=self._request_timeout_seconds
+            request,
+            timeout=self._request_timeout_seconds,
         )
         results: List[VectorEmbeddingRecord] = []
         for vector in response.records:
@@ -89,9 +92,13 @@ def query_vectors(
             include_embeddings=query["include_embeddings"],
             segment_id=self._segment["id"].hex,
             collection_id=self._segment["collection"].hex,
+            version_context=to_proto_request_version_context(
+                query["request_version_context"]
+            ),
         )
         response: QueryVectorsResponse = self._vector_reader_stub.QueryVectors(
-            request, timeout=self._request_timeout_seconds
+            request,
+            timeout=self._request_timeout_seconds,
         )
         results: List[List[VectorQueryResult]] = []
         for result in response.results:

rust/error/src/lib.rs

@@ -10,7 +10,7 @@ pub enum ErrorCodes {
     // CANCELLED indicates the operation was cancelled (typically by the caller).
     Cancelled = 1,
     // UNKNOWN indicates an unknown error.
-    UNKNOWN = 2,
+    Unknown = 2,
     // INVALID_ARGUMENT indicates client specified an invalid argument.
     InvalidArgument = 3,
     // DEADLINE_EXCEEDED means operation expired before completion.
@@ -21,8 +21,6 @@ pub enum ErrorCodes {
     AlreadyExists = 6,
     // PERMISSION_DENIED indicates the caller does not have permission to execute the specified operation.
     PermissionDenied = 7,
-    // UNAUTHENTICATED indicates the request does not have valid authentication credentials for the operation.
-    UNAUTHENTICATED = 16,
     // RESOURCE_EXHAUSTED indicates some resource has been exhausted, perhaps a per-user quota, or perhaps the entire file system is out of space.
     ResourceExhausted = 8,
     // FAILED_PRECONDITION indicates operation was rejected because the system is not in a state required for the operation's execution.
@@ -39,6 +37,10 @@ pub enum ErrorCodes {
     Unavailable = 14,
     // DATA_LOSS indicates unrecoverable data loss or corruption.
     DataLoss = 15,
+    // UNAUTHENTICATED indicates the request does not have valid authentication credentials for the operation.
+    Unauthenticated = 16,
+    // VERSION_MISMATCH indicates a version mismatch. This is not from the gRPC spec and is specific to Chroma.
+    VersionMismatch = 17,
 }
 
 pub trait ChromaError: Error + Send {

rust/worker/src/execution/operators/merge_knn_results.rs

@@ -63,7 +63,7 @@ pub enum MergeKnnResultsOperatorError {}
 
 impl ChromaError for MergeKnnResultsOperatorError {
     fn code(&self) -> ErrorCodes {
-        return ErrorCodes::UNKNOWN;
+        return ErrorCodes::Unknown;
     }
 }
 

rust/worker/src/execution/orchestration/count.rs

@@ -36,6 +36,9 @@ pub(crate) struct CountQueryOrchestrator {
     blockfile_provider: BlockfileProvider,
     // Result channel
     result_channel: Option<tokio::sync::oneshot::Sender<Result<usize, Box<dyn ChromaError>>>>,
+    // Request version context
+    collection_version: u32,
+    log_position: u64,
 }
 
 #[derive(Error, Debug)]
@@ -52,6 +55,8 @@ enum CountQueryOrchestratorError {
     CollectionNotFound(Uuid),
     #[error("Get collection error: {0}")]
     GetCollectionError(#[from] GetCollectionsError),
+    #[error("Collection version mismatch")]
+    CollectionVersionMismatch,
 }
 
 impl ChromaError for CountQueryOrchestratorError {
@@ -65,6 +70,7 @@ impl ChromaError for CountQueryOrchestratorError {
             CountQueryOrchestratorError::SystemTimeError(_) => ErrorCodes::Internal,
             CountQueryOrchestratorError::CollectionNotFound(_) => ErrorCodes::NotFound,
             CountQueryOrchestratorError::GetCollectionError(e) => e.code(),
+            CountQueryOrchestratorError::CollectionVersionMismatch => ErrorCodes::VersionMismatch,
         }
     }
 }
@@ -78,6 +84,8 @@ impl CountQueryOrchestrator {
         sysdb: Box<SysDb>,
         dispatcher: ComponentHandle<Dispatcher>,
         blockfile_provider: BlockfileProvider,
+        collection_version: u32,
+        log_position: u64,
     ) -> Self {
         Self {
             system,
@@ -90,6 +98,8 @@ impl CountQueryOrchestrator {
             dispatcher,
             blockfile_provider,
             result_channel: None,
+            collection_version,
+            log_position,
         }
     }
 
@@ -140,8 +150,19 @@ impl CountQueryOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(CountQueryOrchestratorError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         self.record_segment = Some(record_segment);
         self.collection = Some(collection);
+        self.pull_logs(ctx).await;
     }
 
     // shared
@@ -170,7 +191,9 @@ impl CountQueryOrchestrator {
         let input = PullLogsInput::new(
             collection.id,
             // The collection log position is inclusive, and we want to start from the next log.
-            collection.log_position + 1,
+            // Note that we query using the incoming log position this is critical for correctness
+            // TODO: We should make all the log service code use u64 instead of i64
+            (self.log_position as i64) + 1,
             100,
             None,
             Some(end_timestamp),
@@ -307,7 +330,6 @@ impl Component for CountQueryOrchestrator {
 
     async fn on_start(&mut self, ctx: &crate::system::ComponentContext<Self>) -> () {
         self.start(ctx).await;
-        self.pull_logs(ctx).await;
     }
 }
 

rust/worker/src/execution/orchestration/get_vectors.rs

@@ -43,13 +43,16 @@ enum GetVectorsError {
     TaskSendError(#[from] ChannelError),
     #[error("System time error")]
     SystemTimeError(#[from] std::time::SystemTimeError),
+    #[error("Collection version mismatch")]
+    CollectionVersionMismatch,
 }
 
 impl ChromaError for GetVectorsError {
     fn code(&self) -> ErrorCodes {
         match self {
             GetVectorsError::TaskSendError(e) => e.code(),
             GetVectorsError::SystemTimeError(_) => ErrorCodes::Internal,
+            GetVectorsError::CollectionVersionMismatch => ErrorCodes::VersionMismatch,
         }
     }
 }
@@ -74,6 +77,8 @@ pub struct GetVectorsOrchestrator {
     // Result channel
     result_channel:
         Option<tokio::sync::oneshot::Sender<Result<GetVectorsResult, Box<dyn ChromaError>>>>,
+    collection_version: u32,
+    log_position: u64,
 }
 
 impl GetVectorsOrchestrator {
@@ -86,6 +91,8 @@ impl GetVectorsOrchestrator {
         sysdb: Box<SysDb>,
         dispatcher: ComponentHandle<Dispatcher>,
         blockfile_provider: BlockfileProvider,
+        collection_version: u32,
+        log_position: u64,
     ) -> Self {
         Self {
             state: ExecutionState::Pending,
@@ -100,6 +107,8 @@ impl GetVectorsOrchestrator {
             record_segment: None,
             collection: None,
             result_channel: None,
+            collection_version,
+            log_position,
         }
     }
 
@@ -132,7 +141,9 @@ impl GetVectorsOrchestrator {
         let input = PullLogsInput::new(
             collection.id,
             // The collection log position is inclusive, and we want to start from the next log
-            collection.log_position + 1,
+            // Note that we query using the incoming log position this is critical for correctness
+            // TODO: We should make all the log service code use u64 instead of i64
+            (self.log_position as i64) + 1,
             100,
             None,
             Some(end_timestamp),
@@ -241,6 +252,16 @@ impl Component for GetVectorsOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(GetVectorsError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         let record_segment =
             match get_record_segment_by_collection_id(self.sysdb.clone(), collection_id).await {
                 Ok(segment) => segment,

rust/worker/src/execution/orchestration/hnsw.rs

@@ -88,6 +88,8 @@ enum HnswSegmentQueryError {
     RecordSegmentNotFound(Uuid),
     #[error("Collection has no dimension set")]
     CollectionHasNoDimension,
+    #[error("Collection version mismatch")]
+    CollectionVersionMismatch,
 }
 
 impl ChromaError for HnswSegmentQueryError {
@@ -99,6 +101,7 @@ impl ChromaError for HnswSegmentQueryError {
             HnswSegmentQueryError::GetCollectionError(_) => ErrorCodes::Internal,
             HnswSegmentQueryError::RecordSegmentNotFound(_) => ErrorCodes::NotFound,
             HnswSegmentQueryError::CollectionHasNoDimension => ErrorCodes::InvalidArgument,
+            HnswSegmentQueryError::CollectionVersionMismatch => ErrorCodes::VersionMismatch,
         }
     }
 }
@@ -143,6 +146,9 @@ pub(crate) struct HnswQueryOrchestrator {
     result_channel: Option<
         tokio::sync::oneshot::Sender<Result<Vec<Vec<VectorQueryResult>>, Box<dyn ChromaError>>>,
     >,
+    // Request version context
+    collection_version: u32,
+    log_position: u64,
 }
 
 impl HnswQueryOrchestrator {
@@ -159,6 +165,8 @@ impl HnswQueryOrchestrator {
         hnsw_index_provider: HnswIndexProvider,
         blockfile_provider: BlockfileProvider,
         dispatcher: ComponentHandle<Dispatcher>,
+        collection_version: u32,
+        log_position: u64,
     ) -> Self {
         // Set the merge dependency count to the number of query vectors * 2
         // N for the HNSW query and N for the Brute force query
@@ -203,6 +211,8 @@ impl HnswQueryOrchestrator {
             hnsw_index_provider,
             blockfile_provider,
             result_channel: None,
+            collection_version,
+            log_position,
         }
     }
 
@@ -230,7 +240,9 @@ impl HnswQueryOrchestrator {
         let input = PullLogsInput::new(
             collection.id,
             // The collection log position is inclusive, and we want to start from the next log
-            collection.log_position + 1,
+            // Note that we query using the incoming log position this is critical for correctness
+            // TODO: We should make all the log service code use u64 instead of i64
+            (self.log_position as i64) + 1,
             100,
             None,
             Some(end_timestamp),
@@ -571,6 +583,16 @@ impl Component for HnswQueryOrchestrator {
             }
         };
 
+        // If the collection version does not match the request version then we terminate with an error
+        if collection.version as u32 != self.collection_version {
+            terminate_with_error(
+                self.result_channel.take(),
+                Box::new(HnswSegmentQueryError::CollectionVersionMismatch),
+                ctx,
+            );
+            return;
+        }
+
         // If segment is uninitialized and dimension is not set then we assume
         // that this is a query before any add so return empty response.
         if hnsw_segment.file_path.len() <= 0 && collection.dimension.is_none() {