[WIP] Deploy 2 chromas in tiltfile

Author: jasonvigil

Tiltfile

@@ -1,3 +1,7 @@
+# Enable to test multi-region deployment
+config.define_bool('multi_region')
+cfg = config.parse()
+
 update_settings(max_parallel_updates=6)
 
 # *:ci images are defined in .github/actions/tilt/docker-bake.hcl and used for .github/actions/tilt/action.yaml.
@@ -157,10 +161,21 @@ k8s_yaml(
     'helm template --set-file rustFrontendService.configuration=rust/frontend/sample_configs/distributed.yaml,rustLogService.configuration=rust/worker/chroma_config.yaml,heapTenderService.configuration=rust/worker/chroma_config.yaml,compactionService.configuration=rust/worker/chroma_config.yaml,queryService.configuration=rust/worker/chroma_config.yaml,garbageCollector.configuration=rust/worker/chroma_config.yaml --values k8s/distributed-chroma/values.yaml,k8s/distributed-chroma/values.dev.yaml k8s/distributed-chroma'
   ),
 )
+
+k8s_yaml(
+  local(
+    'helm template --set-file rustFrontendService.configuration=rust/frontend/sample_configs/distributed2.yaml,rustLogService.configuration=rust/worker/chroma_config2.yaml,heapTenderService.configuration=rust/worker/chroma_config2.yaml,compactionService.configuration=rust/worker/chroma_config2.yaml,queryService.configuration=rust/worker/chroma_config2.yaml,garbageCollector.configuration=rust/worker/chroma_config2.yaml --values k8s/distributed-chroma/values2.yaml,k8s/distributed-chroma/values2.dev.yaml k8s/distributed-chroma'
+  ),
+)
+
 watch_file('rust/frontend/sample_configs/distributed.yaml')
+watch_file('rust/frontend/sample_configs/distributed2.yaml')
 watch_file('rust/worker/chroma_config.yaml')
+watch_file('rust/worker/chroma_config2.yaml')
 watch_file('k8s/distributed-chroma/values.yaml')
 watch_file('k8s/distributed-chroma/values.dev.yaml')
+watch_file('k8s/distributed-chroma/values2.yaml')
+watch_file('k8s/distributed-chroma/values2.dev.yaml')
 watch_file('k8s/distributed-chroma/*.yaml')
 
 
@@ -177,60 +192,151 @@ k8s_yaml([
   'k8s/test/prometheus.yaml',
   'k8s/test/test-memberlist-cr.yaml',
   'k8s/test/postgres.yaml',
+  'k8s/test/postgres2.yaml',
 ])
 
 # Lots of things assume the cluster is in a basic state. Get it into a basic
 # state before deploying anything else.
 k8s_resource(
   objects=[
-    'pod-watcher:Role',
     'memberlists.chroma.cluster:CustomResourceDefinition',
-    'query-service-memberlist:MemberList',
-    'compaction-service-memberlist:MemberList',
-    'garbage-collection-service-memberlist:MemberList',
-    'rust-log-service-memberlist:MemberList',
-
-    'sysdb-serviceaccount:ServiceAccount',
-    'sysdb-serviceaccount-rolebinding:RoleBinding',
-    'sysdb-query-service-memberlist-binding:RoleBinding',
-    'sysdb-compaction-service-memberlist-binding:RoleBinding',
-
-    'query-service-serviceaccount:ServiceAccount',
-    'query-service-serviceaccount-rolebinding:RoleBinding',
-    'query-service-memberlist-readerwriter:Role',
-    'query-service-query-service-memberlist-binding:RoleBinding',
-    'query-service-memberlist-readerwriter-binding:RoleBinding',
-
-    'compaction-service-memberlist-readerwriter:Role',
-    'compaction-service-compaction-service-memberlist-binding:RoleBinding',
-    'compaction-service-memberlist-readerwriter-binding:RoleBinding',
-    'compaction-service-serviceaccount:ServiceAccount',
-    'compaction-service-serviceaccount-rolebinding:RoleBinding',
-
-    'test-memberlist:MemberList',
-    'test-memberlist-reader:Role',
-    'test-memberlist-reader-binding:RoleBinding',
-    'lease-watcher:Role',
-    'rust-frontend-service-config:ConfigMap',
+
+    'pod-watcher:Role:chroma',
+    'query-service-memberlist:MemberList:chroma',
+    'compaction-service-memberlist:MemberList:chroma',
+    'garbage-collection-service-memberlist:MemberList:chroma',
+    'rust-log-service-memberlist:MemberList:chroma',
+
+    'sysdb-serviceaccount:ServiceAccount:chroma',
+    'sysdb-serviceaccount-rolebinding:RoleBinding:chroma',
+    'sysdb-query-service-memberlist-binding:RoleBinding:chroma',
+    'sysdb-compaction-service-memberlist-binding:RoleBinding:chroma',
+
+    'query-service-serviceaccount:ServiceAccount:chroma',
+    'query-service-serviceaccount-rolebinding:RoleBinding:chroma',
+    'query-service-memberlist-readerwriter:Role:chroma',
+    'query-service-query-service-memberlist-binding:RoleBinding:chroma',
+    'query-service-memberlist-readerwriter-binding:RoleBinding:chroma',
+
+    'compaction-service-memberlist-readerwriter:Role:chroma',
+    'compaction-service-compaction-service-memberlist-binding:RoleBinding:chroma',
+    'compaction-service-memberlist-readerwriter-binding:RoleBinding:chroma',
+    'compaction-service-serviceaccount:ServiceAccount:chroma',
+    'compaction-service-serviceaccount-rolebinding:RoleBinding:chroma',
+
+    'test-memberlist:MemberList:chroma',
+    'test-memberlist-reader:Role:chroma',
+    'test-memberlist-reader-binding:RoleBinding:chroma',
+    'lease-watcher:Role:chroma',
+    'rust-frontend-service-config:ConfigMap:chroma',
   ],
   new_name='k8s_setup',
   labels=["infrastructure"],
 )
 
+# Lots of things assume the cluster is in a basic state. Get it into a basic
+# state before deploying anything else.
+k8s_resource(
+  objects=[
+    'pod-watcher:Role:chroma2',
+    'query-service-memberlist:MemberList:chroma2',
+    'compaction-service-memberlist:MemberList:chroma2',
+    'garbage-collection-service-memberlist:MemberList:chroma2',
+    'rust-log-service-memberlist:MemberList:chroma2',
+
+    'sysdb-serviceaccount:ServiceAccount:chroma2',
+    'sysdb-serviceaccount-rolebinding:RoleBinding:chroma2',
+    'sysdb-query-service-memberlist-binding:RoleBinding:chroma2',
+    'sysdb-compaction-service-memberlist-binding:RoleBinding:chroma2',
+
+    'query-service-serviceaccount:ServiceAccount:chroma2',
+    'query-service-serviceaccount-rolebinding:RoleBinding:chroma2',
+    'query-service-memberlist-readerwriter:Role:chroma2',
+    'query-service-query-service-memberlist-binding:RoleBinding:chroma2',
+    'query-service-memberlist-readerwriter-binding:RoleBinding:chroma2',
+
+    'compaction-service-memberlist-readerwriter:Role:chroma2',
+    'compaction-service-compaction-service-memberlist-binding:RoleBinding:chroma2',
+    'compaction-service-memberlist-readerwriter-binding:RoleBinding:chroma2',
+    'compaction-service-serviceaccount:ServiceAccount:chroma2',
+    'compaction-service-serviceaccount-rolebinding:RoleBinding:chroma2',
+
+    'lease-watcher:Role:chroma2',
+    'rust-frontend-service-config:ConfigMap:chroma2',
+  ],
+  new_name='k8s_setup2',
+  labels=["infrastructure2"],
+)
+
 # Production Chroma
-k8s_resource('postgres', resource_deps=['k8s_setup'], labels=["infrastructure"], port_forwards='5432:5432')
+k8s_resource('postgres:deployment:chroma', 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('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')
-k8s_resource('compaction-service', resource_deps=['sysdb'], labels=["chroma"])
+k8s_resource('sysdb-migration-latest:job:chroma', resource_deps=['postgres:deployment:chroma'], labels=["infrastructure"])
+k8s_resource('rust-log-service:statefulset:chroma', labels=["chroma"], port_forwards=['50054:50051', '50052:50052'], resource_deps=['minio-deployment'])
+k8s_resource('sysdb:deployment:chroma', resource_deps=['sysdb-migration-latest:job:chroma'], labels=["chroma"], port_forwards='50051:50051')
+k8s_resource('rust-frontend-service:deployment:chroma', resource_deps=['sysdb:deployment:chroma', 'rust-log-service:statefulset:chroma'], labels=["chroma"], port_forwards='8000:8000')
+k8s_resource('query-service:statefulset:chroma', resource_deps=['sysdb:deployment:chroma'], labels=["chroma"], port_forwards='50053:50051')
+k8s_resource('compaction-service:statefulset:chroma', resource_deps=['sysdb:deployment:chroma'], labels=["chroma"])
+k8s_resource('garbage-collector:statefulset:chroma', resource_deps=['k8s_setup', 'minio-deployment'], labels=["chroma"], port_forwards='50055:50055')
 k8s_resource('load-service', resource_deps=['k8s_setup'], labels=["infrastructure"], port_forwards='3001:3001')
+
+# Production Chroma 2
+k8s_resource('postgres:deployment:chroma2', resource_deps=['k8s_setup2'], labels=["infrastructure2"], port_forwards='6432: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:job:chroma2', resource_deps=['postgres:deployment:chroma2'], labels=["infrastructure2"])
+k8s_resource('rust-log-service:statefulset:chroma2', labels=["chroma2"], port_forwards=['60054:50051', '60052:50052'], resource_deps=['minio-deployment'])
+k8s_resource('sysdb:deployment:chroma2', resource_deps=['sysdb-migration-latest:job:chroma2'], labels=["chroma2"], port_forwards='60051:50051')
+k8s_resource('rust-frontend-service:deployment:chroma2', resource_deps=['sysdb:deployment:chroma2', 'rust-log-service:statefulset:chroma2'], labels=["chroma2"], port_forwards='8001:8000')
+k8s_resource('query-service:statefulset:chroma2', resource_deps=['sysdb:deployment:chroma2'], labels=["chroma2"], port_forwards='60053:50051')
+k8s_resource('compaction-service:statefulset:chroma2', resource_deps=['sysdb:deployment:chroma2'], labels=["chroma2"])
+k8s_resource('garbage-collector:statefulset:chroma2', resource_deps=['k8s_setup2', 'minio-deployment'], labels=["chroma2"], port_forwards='60055:50055')
+
+# Observability
 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', 'minio-deployment'], labels=["chroma"], port_forwards='50055:50055')
+
 # Local S3
 k8s_resource('minio-deployment', resource_deps=['k8s_setup'], labels=["debug"], port_forwards=['9000:9000', '9005:9005'])
+
+
+# Set the enabled resources
+# - Basic resources are always enabled.
+# - Multi-region resources are only enabled if the multi-region flag is set.
+config.clear_enabled_resources()
+
+groups = {
+  'basic': [
+    'k8s_setup',
+    'postgres:deployment:chroma',
+    'sysdb-migration-latest:job:chroma',
+    'rust-log-service:statefulset:chroma',
+    'sysdb:deployment:chroma',
+    'rust-frontend-service:deployment:chroma',
+    'query-service:statefulset:chroma',
+    'compaction-service:statefulset:chroma',
+    'load-service',
+    'garbage-collector:statefulset:chroma',
+    'jaeger',
+    'grafana',
+    'prometheus',
+    'otel-collector',
+    'minio-deployment',
+  ],
+  'multi_region': [
+    'postgres:deployment:chroma2',
+    'sysdb-migration-latest:job:chroma2',
+    'rust-log-service:statefulset:chroma2',
+    'sysdb:deployment:chroma2',
+    'rust-frontend-service:deployment:chroma2',
+    'query-service:statefulset:chroma2',
+    'compaction-service:statefulset:chroma2',
+    'garbage-collector:statefulset:chroma2',
+  ],
+}
+
+if cfg.get('multi_region'):
+  config.set_enabled_resources(groups['basic'] + groups['multi_region'])
+else:
+  config.set_enabled_resources(groups['basic'])

go/pkg/memberlist_manager/memberlist_store.go

@@ -144,7 +144,7 @@ func (s *CRMemberlistStore) UpdateMemberlist(ctx context.Context, memberlist Mem
 	log.Debug("Updating memberlist store", zap.Any("memberlist", memberlist))
 	unstructured := memberlist.toCr(s.coordinatorNamespace, s.memberlistCustomResource, resourceVersion)
 	log.Debug("Setting memberlist to unstructured object", zap.Any("unstructured", unstructured))
-	_, err := s.dynamicClient.Resource(gvr).Namespace("chroma").Update(context.Background(), unstructured, metav1.UpdateOptions{})
+	_, err := s.dynamicClient.Resource(gvr).Namespace(s.coordinatorNamespace).Update(context.Background(), unstructured, metav1.UpdateOptions{})
 	if err != nil {
 		return err
 	}

k8s/distributed-chroma/values.yaml

@@ -20,7 +20,7 @@ sysdb:
   replicaCount: 1
   env:
     - name: OPTL_TRACING_ENDPOINT
-      value: 'value: "jaeger:4317"'
+      value: 'value: "jaeger.chroma.svc.cluster.local:4317"'
   resources:
     limits:
       cpu: '2000m'

k8s/distributed-chroma/values2.dev.yaml

@@ -0,0 +1,36 @@
+sysdb:
+  flags:
+    version-file-enabled: true
+    s3-endpoint: "http://minio.chroma.svc.cluster.local:9000"
+    s3-access-key-id: "minio"
+    s3-secret-access-key: "minio123"
+    s3-force-path-style: true
+    create-bucket-if-not-exists: true
+    kubernetes-namespace: chroma2
+rustFrontendService:
+  # We have to specify the command, because the Dockerfile uses the CLI since its shared with
+  # single node, so in values.dev we pass the CONFIG_PATH into the chroma run command
+  command: '["chroma", "run", "$(CONFIG_PATH)"]'
+  otherEnvConfig: |
+    - name: CHROMA_ALLOW_RESET
+      value: "true"
+    - name: RUST_BACKTRACE
+      value: 'value: "1"'
+
+queryService:
+  env:
+    - name: RUST_BACKTRACE
+      value: 'value: "1"'
+  jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19"
+
+compactionService:
+  env:
+    - name: RUST_BACKTRACE
+      value: 'value: "1"'
+  jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19"
+
+rustLogService:
+  replicaCount: 1
+
+garbageCollector:
+  jemallocConfig: "prof:true,prof_active:true,lg_prof_sample:19"

k8s/distributed-chroma/values2.yaml

@@ -0,0 +1,87 @@
+# Default values for distributed-chroma.
+# Strongly prefer single quotes.
+namespace: 'chroma2'
+rustFrontendService:
+  image:
+    repository: 'rust-frontend-service'
+    tag: 'latest'
+  replicaCount: 1
+  resources:
+    limits:
+      cpu: '2000m'
+      memory: '1Gi'
+    requests:
+      cpu: '1000m'
+      memory: '512Mi'
+sysdb:
+  image:
+    repository: 'sysdb'
+    tag: 'latest'
+  replicaCount: 1
+  env:
+    - name: OPTL_TRACING_ENDPOINT
+      value: 'value: "jaeger.chroma:4317"'
+  resources:
+    limits:
+      cpu: '2000m'
+      memory: '1Gi'
+    requests:
+      cpu: '1000m'
+      memory: '512Mi'
+rustLogService:
+  image:
+    repository: 'rust-log-service'
+    tag: 'latest'
+  cache:
+    hostPath: '/local/cache/chroma-log-service'
+    type: DirectoryOrCreate
+    mountPath: '/cache/'
+    additionalVolumes: []
+queryService:
+  image:
+    repository: 'query-service'
+    tag: 'latest'
+  env:
+  cache:
+    hostPath: '/local/cache/chroma-query-service'
+    type: DirectoryOrCreate
+    mountPath: '/cache/'
+    additionalVolumes: []
+  replicaCount: 2
+compactionService:
+  image:
+    repository: 'compaction-service'
+    tag: 'latest'
+  env:
+  cache:
+    hostPath: '/local/cache/chroma-compaction-service'
+    type: DirectoryOrCreate
+    mountPath: '/cache/'
+    additionalVolumes: []
+  replicaCount: 1
+sysdbMigration:
+  image:
+    repository: 'sysdb-migration'
+    tag: 'latest'
+  username: chroma
+  password: chroma
+  netloc: postgres
+  port: 5432
+  dbName: sysdb
+  sslmode: disable
+# Add the garbage collector configuration
+garbageCollector:
+  image:
+    repository: 'garbage-collector'
+    tag: 'latest'
+  replicaCount: 1
+  resources:
+    limits:
+      cpu: '200m'
+      memory: '256Mi'
+    requests:
+      cpu: '100m'
+      memory: '128Mi'
+  cache:
+    hostPath: '/local/cache/chroma-garbage-collector'
+    mountPath: '/cache/'

k8s/test/postgres-service2.yaml

@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: postgres-lb
+  namespace: chroma2
+spec:
+  ports:
+    - name: postgres-port
+      port: 5432
+      targetPort: 5432
+  selector:
+    app: postgres
+  type: LoadBalancer