[ENH] Refactor k8s manifests for configurability (#1808)

## Description of changes *Summarize the changes made by this PR.* - Improvements & Bug fixes - Refactor k8s charts to make everything more configurable, especially in ways we need for Hosted. - Put correct deps into Tiltfile so the cluster comes up happily. ## Test plan *How are these changes tested?* - [ ] Tests pass locally with `pytest` for python, `yarn test` for js ## Documentation Changes *Are all docstrings for user-facing APIs updated if required? Do we need to make documentation changes in the [docs repository](https://github.com/chroma-core/docs)?*
chroma-core · Mar 4, 2024 · bbc676c · bbc676c
1 parent 3fb5c1e
commit bbc676c
Show file tree

Hide file tree

Showing 8 changed files with 92 additions and 29 deletions.
diff --git a/Tiltfile b/Tiltfile
@@ -64,13 +64,16 @@ k8s_resource(
 )
 
 # Production Chroma
+k8s_yaml([
+  'k8s/test/postgres.yaml',
+])
 k8s_resource('postgres', resource_deps=['k8s_setup'], labels=["infrastructure"])
 k8s_resource('pulsar', resource_deps=['k8s_setup'], labels=["infrastructure"], port_forwards=['6650:6650', '8080:8080'])
 k8s_resource('migration', resource_deps=['postgres'], labels=["infrastructure"])
 k8s_resource('logservice', resource_deps=['migration'], labels=["chroma"], port_forwards='50052:50051')
-k8s_resource('frontend-server', resource_deps=['logservice'],labels=["chroma"], port_forwards=8000 )
-k8s_resource('coordinator', resource_deps=['pulsar', 'frontend-server', 'migration'], labels=["chroma"], port_forwards=50051)
-k8s_resource('worker', resource_deps=['coordinator'],labels=["chroma"])
+k8s_resource('coordinator', resource_deps=['pulsar', 'migration'], labels=["chroma"], port_forwards='50051:50051')
+k8s_resource('frontend-server', resource_deps=['pulsar', 'coordinator', 'logservice'],labels=["chroma"], port_forwards='8000:8000')
+k8s_resource('worker', resource_deps=['coordinator', 'pulsar'], labels=["chroma"])
 
 # Extra stuff to make debugging and testing easier
 k8s_yaml([
@@ -93,4 +96,4 @@ k8s_resource(
 )
 
 # Local S3
-k8s_resource('minio-deployment', resource_deps=['k8s_setup'], labels=["debug"], port_forwards=9000)
+k8s_resource('minio-deployment', resource_deps=['k8s_setup'], labels=["debug"], port_forwards='9000:9000')
diff --git a/k8s/distributed-chroma/templates/coordinator.yaml b/k8s/distributed-chroma/templates/coordinator.yaml
@@ -18,10 +18,16 @@ spec:
         - command:
             - "coordinator"
             - "coordinator"
-            - "--pulsar-admin-url=http://pulsar.chroma:8080"
-            - "--pulsar-url=pulsar://pulsar.chroma:6650"
-            - "--notifier-provider=pulsar"
-          image: coordinator
+            {{ range $k, $v := .Values.coordinator.flags }}
+            - "--{{ $k }}={{ $v }}"
+            {{ end }}
+          env:
+            {{ range .Values.coordinator.env }}
+            - name: {{ .name }}
+              # TODO properly use flow control here to check which type of value we need.
+              {{ .value }}
+            {{ end }}
+          image: {{ .Values.coordinator.image }}
           imagePullPolicy: IfNotPresent
           name: coordinator
           ports:

diff --git a/k8s/distributed-chroma/templates/frontend-server.yaml b/k8s/distributed-chroma/templates/frontend-server.yaml
@@ -15,7 +15,7 @@ spec:
     spec:
       containers:
         - name: frontend-server
-          image: server
+          image: {{ .Values.frontend.image }}
           imagePullPolicy: IfNotPresent
           ports:
             - containerPort: 8000
@@ -24,38 +24,48 @@ spec:
               mountPath: /test
           env:
             - name: IS_PERSISTENT
-              value: "TRUE"
+              {{ .Values.frontend.isPersistent }}
             - name: CHROMA_PRODUCER_IMPL
-              value: "chromadb.logservice.logservice.LogService"
+              {{ .Values.frontend.producerImpl }}
             - name: CHROMA_CONSUMER_IMPL
-              value: "chromadb.logservice.logservice.LogService"
+              {{ .Values.frontend.consumerImpl }}
             - name: CHROMA_SEGMENT_MANAGER_IMPL
-              value: "chromadb.segment.impl.manager.distributed.DistributedSegmentManager"
+              {{ .Values.frontend.segmentManagerImpl }}
             - name: PULSAR_BROKER_URL
-              value: "pulsar.chroma"
+              {{ .Values.frontend.pulsarBrokerUrl }}
             - name: PULSAR_BROKER_PORT
-              value: "6650"
+              {{ .Values.frontend.pulsarBrokerPort }}
             - name: PULSAR_ADMIN_PORT
-              value: "8080"
+              {{ .Values.frontend.pulsarAdminPort }}
             - name: ALLOW_RESET
-              value: "TRUE"
+              {{ .Values.frontend.allowReset }}
             - name: CHROMA_SYSDB_IMPL
-              value: "chromadb.db.impl.grpc.client.GrpcSysDB"
+              {{ .Values.frontend.sysdbImpl }}
             - name: CHROMA_SERVER_GRPC_PORT
-              value: "50051"
+              {{ .Values.frontend.serverGrpcPort }}
             - name: CHROMA_COORDINATOR_HOST
-              value: "coordinator.chroma"
+              {{ .Values.frontend.coordinatorHost }}
+            - name: CHROMA_SERVER_AUTH_PROVIDER
+              {{ .Values.frontend.serverAuthProvider }}
+            - name: CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER
+              {{ .Values.frontend.serverAuthCredentialsProvider }}
+            - name: CHROMA_SERVER_AUTHZ_PROVIDER
+              {{ .Values.frontend.serverAuthzProvider }}
+            - name: CHROMA_SERVER_AUTHZ_CONFIG_PROVIDER
+              {{ .Values.frontend.serverAuthzConfigProvider }}
+              {{ .Values.frontend.otherEnvConfig }}
             - name: CHROMA_MEMBERLIST_PROVIDER_IMPL
-              value: "chromadb.segment.impl.distributed.segment_directory.MockMemberlistProvider"
+              {{ .Values.frontend.memberlistProviderImpl }}
             - name: CHROMA_LOGSERVICE_HOST
-              value: "logservice.chroma"
+              {{ .Values.frontend.logServiceHost }}
             - name: CHROMA_LOGSERVICE_PORT
-              value: "50051"
+              {{ .Values.frontend.logServicePort }}
       volumes:
         - name: chroma
           emptyDir: {}
 
 ---
+
 apiVersion: v1
 kind: Service
 metadata:

diff --git a/k8s/distributed-chroma/templates/logservice.yaml b/k8s/distributed-chroma/templates/logservice.yaml
@@ -18,7 +18,7 @@ spec:
         - command:
             - "logservice"
             - "logservice"
-          image: coordinator
+          image: {{ .Values.logService.image }}
           imagePullPolicy: IfNotPresent
           name: logservice
           ports:

diff --git a/k8s/distributed-chroma/templates/migration.yaml b/k8s/distributed-chroma/templates/migration.yaml
@@ -16,7 +16,7 @@ spec:
             - 'apply'
             - '--url'
             - 'postgres://chroma:chroma@postgres:5432/chroma?sslmode=disable'
-          image: migration
+          image: {{ .Values.sysdbMigration.image }}
           imagePullPolicy: IfNotPresent
           name: migration
 ---
diff --git a/k8s/distributed-chroma/templates/pulsar.yaml b/k8s/distributed-chroma/templates/pulsar.yaml
@@ -17,7 +17,6 @@ spec:
         - name: pulsar
           image: apachepulsar/pulsar
           command: [ "/pulsar/bin/pulsar", "standalone" ]
-
           ports:
             - containerPort: 6650
             - containerPort: 8080

diff --git a/k8s/distributed-chroma/values.yaml b/k8s/distributed-chroma/values.yaml
@@ -1,6 +1,49 @@
 # Default values for distributed-chroma.
+# Strongly prefer single quotes.
 
-namespace: "chroma"
+namespace: 'chroma'
+
+frontend:
+  image: 'server'
+
+  # Sometimes users (and us) want to pass values directly as flags. Sometimes, these are
+  # populated from secrets or configMaps. So we let consumers fill these directly.
+  # TODO we could maybe have mutually exclusive pieces of config, e.g. isPersistentValue
+  # and isPersistentFromConfigMap or something.
+  isPersistent: 'value: "TRUE"'
+  producerImpl: 'value: "chromadb.logservice.logservice.LogService"'
+  consumerImpl: 'value: "chromadb.logservice.logservice.LogService"'
+  segmentManagerImpl: 'value: "chromadb.segment.impl.manager.distributed.DistributedSegmentManager"'
+  pulsarBrokerUrl: 'value: "pulsar.chroma"'
+  pulsarBrokerPort: 'value: "6650"'
+  pulsarAdminPort: 'value: "8080"'
+  allowReset: 'value: "TRUE"'
+  sysdbImpl: 'value: "chromadb.db.impl.grpc.client.GrpcSysDB"'
+  serverGrpcPort: 'value: "50051"'
+  coordinatorHost: 'value: "coordinator.chroma"'
+  authProvider: 'value: ""'
+  authCredentialsProvider: 'value: ""'
+  authzProvider: 'value: ""'
+  authzConfigProvider: 'value: ""'
+  memberlistProviderImpl: 'value: "chromadb.segment.impl.distributed.segment_directory.MockMemberlistProvider"'
+  logServiceHost: 'value: "logservice.chroma"'
+  logServicePort: 'value: "50051"'
+  otherEnvConfig: ''
 
 coordinator:
+  image: 'coordinator'
   replicaCount: 1
+  env:
+  flags:
+    pulsar-admin-url: "http://pulsar.chroma:8080"
+    pulsar-url: "pulsar://pulsar.chroma:6650"
+    notifier-provider: "pulsar"
+
+logService:
+  image: 'coordinator'
+
+worker:
+  image: 'worker'
+
+sysdbMigration:
+  image: 'migration'
diff --git a/...istributed-chroma/templates/postgres.yaml → k8s/test/postgres.yaml b/...istributed-chroma/templates/postgres.yaml → k8s/test/postgres.yaml
@@ -2,7 +2,7 @@ apiVersion: apps/v1
 kind: Deployment
 metadata:
   name: postgres
-  namespace: {{ .Values.namespace }}
+  namespace: chroma
 spec:
   replicas: 1
   selector:
@@ -25,12 +25,14 @@ spec:
               value: chroma
           ports:
             - containerPort: 5432
+
 ---
+
 apiVersion: v1
 kind: Service
 metadata:
   name: postgres
-  namespace: {{ .Values.namespace }}
+  namespace: chroma
 spec:
   ports:
     - name: postgres-port