[ENH]: set up rendezvous hashing for collection -> garbage collector node mapping

Tiltfile

@@ -198,6 +198,7 @@ k8s_resource(
     'memberlists.chroma.cluster:CustomResourceDefinition',
     'query-service-memberlist:MemberList',
     'compaction-service-memberlist:MemberList',
+    'garbage-collection-service-memberlist:MemberList',
 
     'sysdb-serviceaccount:serviceaccount',
     'sysdb-serviceaccount-rolebinding:RoleBinding',
@@ -245,6 +246,6 @@ k8s_resource('jaeger', resource_deps=['k8s_setup'], labels=["observability"])
 k8s_resource('grafana', resource_deps=['k8s_setup'], labels=["observability"])
 k8s_resource('prometheus', resource_deps=['k8s_setup'], labels=["observability"])
 k8s_resource('otel-collector', resource_deps=['k8s_setup'], labels=["observability"])
-k8s_resource('garbage-collector', resource_deps=['k8s_setup'], labels=["chroma"])
+k8s_resource('garbage-collector', resource_deps=['k8s_setup', 'minio-deployment'], labels=["chroma"])
 # Local S3
 k8s_resource('minio-deployment', resource_deps=['k8s_setup'], labels=["debug"], port_forwards=['9000:9000', '9005:9005'])

go/cmd/coordinator/cmd.go

@@ -65,6 +65,10 @@ func init() {
 	Cmd.Flags().StringVar(&conf.CompactionServiceMemberlistName, "compaction-memberlist-name", "compaction-service-memberlist", "Compaction memberlist name")
 	Cmd.Flags().StringVar(&conf.CompactionServicePodLabel, "compaction-pod-label", "compaction-service", "Compaction pod label")
 
+	// Garbage collection service Memberlist
+	Cmd.Flags().StringVar(&conf.GarbageCollectionServiceMemberlistName, "garbage-collection-memberlist-name", "garbage-collection-service-memberlist", "Garbage collection memberlist name")
+	Cmd.Flags().StringVar(&conf.GarbageCollectionServicePodLabel, "garbage-collection-pod-label", "garbage-collection-service", "Garbage collection pod label")
+
 	// S3 config
 	Cmd.Flags().BoolVar(&conf.MetaStoreConfig.CreateBucketIfNotExists, "create-bucket-if-not-exists", false, "Create bucket if not exists")
 	Cmd.Flags().StringVar(&conf.MetaStoreConfig.BucketName, "bucket-name", "chroma-storage", "Bucket name")

go/pkg/sysdb/grpc/server.go

@@ -49,6 +49,10 @@ type Config struct {
 	CompactionServiceMemberlistName string
 	CompactionServicePodLabel       string
 
+	// Garbage collection service memberlist config
+	GarbageCollectionServiceMemberlistName string
+	GarbageCollectionServicePodLabel       string
+
 	// Config for soft deletes.
 	SoftDeleteEnabled          bool
 	SoftDeleteCleanupInterval  time.Duration
@@ -130,6 +134,9 @@ func NewWithGrpcProvider(config Config, provider grpcutils.GrpcProvider) (*Serve
 			return nil, err
 		}
 
+		// Create memberlist manager for garbage collection service
+		garbageCollectionMemberlistManager, err := createMemberlistManager(namespace, config.GarbageCollectionServiceMemberlistName, config.GarbageCollectionServicePodLabel, config.WatchInterval, config.ReconcileInterval, config.ReconcileCount)
+
 		// Start the memberlist manager for query service
 		err = queryMemberlistManager.Start()
 		if err != nil {
@@ -141,6 +148,12 @@ func NewWithGrpcProvider(config Config, provider grpcutils.GrpcProvider) (*Serve
 			return nil, err
 		}
 
+		// Start the memberlist manager for garbage collection service
+		err = garbageCollectionMemberlistManager.Start()
+		if err != nil {
+			return nil, err
+		}
+
 		log.Info("Starting soft delete cleaner", zap.Duration("cleanup_interval", s.softDeleteCleaner.cleanupInterval), zap.Duration("max_age", s.softDeleteCleaner.maxAge), zap.Uint("limit_per_check", s.softDeleteCleaner.limitPerCheck))
 		s.softDeleteCleaner.Start()
 

k8s/distributed-chroma/Chart.yaml

@@ -16,7 +16,7 @@ apiVersion: v2
 name: distributed-chroma
 description: A helm chart for distributed Chroma
 type: application
-version: 0.1.35
+version: 0.1.36
 appVersion: "0.4.24"
 keywords:
   - chroma

k8s/distributed-chroma/templates/garbage-collection-service-memberlist-cr.yaml

@@ -0,0 +1,71 @@
+apiVersion: chroma.cluster/v1
+kind: MemberList
+metadata:
+  name: garbage-collection-service-memberlist
+  namespace: {{ .Values.namespace }}
+spec:
+  members:
+
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: garbage-collection-service-memberlist-readerwriter
+rules:
+- apiGroups:
+    - chroma.cluster
+  resources:
+    - memberlists
+  verbs:
+    - get
+    - list
+    - watch
+    # TODO: FIX THIS LEAKY PERMISSION
+    - create
+    - update
+    - patch
+    - delete
+
+---
+
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: sysdb-garbage-collection-service-memberlist-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: ClusterRole
+  name: garbage-collection-service-memberlist-readerwriter
+subjects:
+- kind: ServiceAccount
+  name: sysdb-serviceaccount
+  namespace: {{ .Values.namespace }}
+
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: garbage-collection-service-garbage-collection-service-memberlist-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: ClusterRole
+  name: garbage-collection-service-memberlist-readerwriter
+subjects:
+- kind: ServiceAccount
+  name: garbage-collector-serviceaccount
+  namespace: {{ .Values.namespace }}
+
+---
+
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRoleBinding
+metadata:
+  name: garbage-collection-service-memberlist-readerwriter-binding
+roleRef:
+  apiGroup: rbac.authorization.k8s.io
+  kind: ClusterRole
+  name: garbage-collection-service-memberlist-readerwriter
+subjects:
+- kind: ServiceAccount
+  name: default
+  namespace: {{ .Values.namespace }}

k8s/distributed-chroma/templates/garbage-collector.yaml

@@ -12,7 +12,7 @@ data:
 {{ end }}
 
 apiVersion: apps/v1
-kind: Deployment
+kind: StatefulSet
 metadata:
   name: garbage-collector
   namespace: {{ .Values.namespace }}
@@ -27,6 +27,7 @@ spec:
     metadata:
       labels:
         app: garbage-collector
+        member-type: garbage-collection-service
     spec:
       serviceAccountName: garbage-collector-serviceaccount
       volumes:
@@ -58,6 +59,10 @@ spec:
           - name: CONFIG_PATH
             value: /config/config.yaml
           {{ end }}
+          - name: CHROMA_GC_MY_MEMBER_ID
+            valueFrom:
+              fieldRef:
+                fieldPath: metadata.name
           {{ range .Values.garbageCollector.env }}
           - name: {{ .name }}
 {{ .value | nindent 12 }}

k8s/distributed-chroma/values.yaml

@@ -140,10 +140,11 @@ garbageCollector:
   configuration: |
     service_name: "garbage-collector"
     otel_endpoint: "http://otel-collector:4317"
-    relative_cutoff_time_seconds: 43200 # GC all versions created at time < now() - relative_cutoff_time_seconds (12 hours)
+    relative_cutoff_time_seconds: 60 # GC all versions created at time < now() - relative_cutoff_time_seconds (1 minute)
     max_collections_to_gc: 1000
-    gc_interval_mins: 120
+    gc_interval_mins: 1
     disallow_collections: []
+    default_mode: "delete"
     sysdb_config:
       host: "sysdb.chroma"
       port: 50051
@@ -156,12 +157,11 @@ garbageCollector:
     storage_config:
       s3:
         bucket: "chroma-storage"
-      # object_store:
-      #   bucket:
-      #     name: "chroma-storage"
-      #     type: "minio"
-      #   credentials: "Minio"
-      #   connect_timeout_ms: 5000
-      #   request_timeout_ms: 30000
-      #   upload_part_size_bytes: 536870912
-      #   download_part_size_bytes: 8388608
+    assignment_policy:
+      rendezvous_hashing:
+        hasher: Murmur3
+    memberlist_provider:
+      custom_resource:
+        kube_namespace: "chroma"
+        memberlist_name: "garbage-collection-service-memberlist"
+        queue_size: 100

rust/garbage_collector/Cargo.toml

@@ -43,6 +43,7 @@ chroma-blockstore = { workspace = true }
 chroma-cache = { workspace = true }
 chroma-index = { workspace = true }
 chroma-segment = { workspace = true }
+chroma-memberlist = { workspace = true }
 chroma-tracing = { workspace = true }
 
 [dev-dependencies]

rust/garbage_collector/garbage_collector_config.yaml

@@ -23,3 +23,12 @@ storage_config:
     request_timeout_ms: 30000 # 1 minute
     upload_part_size_bytes: 536870912 # 512MiB
     download_part_size_bytes: 8388608 # 8MiB
+assignment_policy:
+  rendezvous_hashing:
+    hasher: Murmur3
+my_member_id: "garbage-collector-0"
+memberlist_provider:
+  custom_resource:
+    kube_namespace: "chroma"
+    memberlist_name: "garbage-collection-service-memberlist"
+    queue_size: 100

rust/garbage_collector/src/config.rs

@@ -35,6 +35,9 @@ pub(super) struct GarbageCollectorConfig {
     pub(super) default_mode: CleanupMode,
     #[serde(default)]
     pub(super) tenant_mode_overrides: Option<HashMap<String, CleanupMode>>,
+    pub(super) assignment_policy: chroma_config::assignment::config::AssignmentPolicyConfig,
+    pub(super) memberlist_provider: chroma_memberlist::config::MemberlistProviderConfig,
+    pub my_member_id: String,
 }
 
 impl GarbageCollectorConfig {
@@ -45,9 +48,10 @@ impl GarbageCollectorConfig {
     pub(super) fn load_from_path(path: &str) -> Self {
         // Unfortunately, figment doesn't support environment variables with underscores. So we have to map and replace them.
         // Excluding our own environment variables, which are prefixed with CHROMA_.
-        let mut f = figment::Figment::from(
-            Env::prefixed("CHROMA_GC_").map(|k| k.as_str().replace("__", ".").into()),
-        );
+        let mut f = figment::Figment::from(Env::prefixed("CHROMA_GC_").map(|k| match k {
+            k if k == "my_member_id" => k.into(),
+            k => k.as_str().replace("__", ".").into(),
+        }));
         if std::path::Path::new(path).exists() {
             f = figment::Figment::from(Yaml::file(path)).merge(f);
         }