Add KEDA scaling for SmithDB compaction workers

charts/langsmith/README.md

@@ -1011,13 +1011,15 @@ Two things worth planning for before you enable it:
 | smithdb.compaction.service.annotations | object | `{}` |  |
 | smithdb.compaction.service.labels | object | `{}` |  |
 | smithdb.compaction.service.port | int | `8071` |  |
-| smithdb.compactionWorker.autoscaling.enabled | bool | `true` |  |
-| smithdb.compactionWorker.autoscaling.maxReplicas | int | `10` |  |
-| smithdb.compactionWorker.autoscaling.minReplicas | int | `1` |  |
-| smithdb.compactionWorker.autoscaling.scaleDownStabilizationWindowSeconds | int | `300` |  |
-| smithdb.compactionWorker.autoscaling.scalePodCount | int | `1` |  |
-| smithdb.compactionWorker.autoscaling.scaleUpStabilizationWindowSeconds | int | `120` |  |
-| smithdb.compactionWorker.autoscaling.targetCPUUtilizationPercentage | int | `60` |  |
+| smithdb.compactionWorker.autoscaling.hpa.enabled | bool | `true` | Enable HPA scaling. |
+| smithdb.compactionWorker.autoscaling.hpa.maxReplicas | int | `10` |  |
+| smithdb.compactionWorker.autoscaling.hpa.minReplicas | int | `1` |  |
+| smithdb.compactionWorker.autoscaling.hpa.scaleDownStabilizationWindowSeconds | int | `300` |  |
+| smithdb.compactionWorker.autoscaling.hpa.scalePodCount | int | `1` |  |
+| smithdb.compactionWorker.autoscaling.hpa.scaleUpStabilizationWindowSeconds | int | `120` |  |
+| smithdb.compactionWorker.autoscaling.hpa.targetCPUUtilizationPercentage | int | `60` |  |
+| smithdb.compactionWorker.autoscaling.keda | object | `{"annotations":{},"cooldownPeriod":300,"enabled":false,"initialCooldownPeriod":0,"labels":{},"maxReplicaCount":10,"minReplicaCount":1,"pendingJobsThreshold":"60","pollingInterval":30,"scaleDownStabilizationWindowSeconds":300,"scalePodCount":1,"scaleUpStabilizationWindowSeconds":120,"targetCPUUtilizationPercentage":60}` | KEDA scaling based on pending compaction jobs and CPU. Requires KEDA to be installed in the cluster. |
+| smithdb.compactionWorker.autoscaling.keda.pendingJobsThreshold | string | `"60"` | Number of pending compaction jobs per worker. |
 | smithdb.compactionWorker.containerPort | int | `9000` |  |
 | smithdb.compactionWorker.deployment.affinity | object | `{}` |  |
 | smithdb.compactionWorker.deployment.annotations | object | `{}` |  |

charts/langsmith/templates/smithdb/compaction-worker-deployment.yaml

@@ -18,7 +18,7 @@ metadata:
       {{- toYaml . | nindent 4 }}
     {{- end }}
 spec:
-  {{- if not (dig "autoscaling" "enabled" false .Values.smithdb.compactionWorker) }}
+  {{- if not (or .Values.smithdb.compactionWorker.autoscaling.hpa.enabled .Values.smithdb.compactionWorker.autoscaling.keda.enabled) }}
   replicas: {{ .Values.smithdb.compactionWorker.deployment.replicas }}
   {{- end }}
   selector:

charts/langsmith/templates/smithdb/compaction-worker-hpa.yaml

@@ -1,4 +1,4 @@
-{{- if and .Values.smithdb.enabled (dig "autoscaling" "enabled" false .Values.smithdb.compactionWorker) }}
+{{- if and .Values.smithdb.enabled .Values.smithdb.compactionWorker.autoscaling.hpa.enabled }}
 apiVersion: autoscaling/v2
 kind: HorizontalPodAutoscaler
 metadata:
@@ -12,26 +12,26 @@ spec:
     apiVersion: apps/v1
     kind: Deployment
     name: {{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compactionWorker") }}
-  minReplicas: {{ .Values.smithdb.compactionWorker.autoscaling.minReplicas }}
-  maxReplicas: {{ .Values.smithdb.compactionWorker.autoscaling.maxReplicas }}
+  minReplicas: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.minReplicas }}
+  maxReplicas: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.maxReplicas }}
   behavior:
     scaleUp:
-      stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.scaleUpStabilizationWindowSeconds }}
+      stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scaleUpStabilizationWindowSeconds }}
       policies:
         - type: Pods
-          value: {{ .Values.smithdb.compactionWorker.autoscaling.scalePodCount }}
-          periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.scaleUpStabilizationWindowSeconds }}
+          value: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scalePodCount }}
+          periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scaleUpStabilizationWindowSeconds }}
     scaleDown:
-      stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.scaleDownStabilizationWindowSeconds }}
+      stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scaleDownStabilizationWindowSeconds }}
       policies:
         - type: Pods
-          value: {{ .Values.smithdb.compactionWorker.autoscaling.scalePodCount }}
-          periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.scaleDownStabilizationWindowSeconds }}
+          value: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scalePodCount }}
+          periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.scaleDownStabilizationWindowSeconds }}
   metrics:
     - type: Resource
       resource:
         name: cpu
         target:
           type: Utilization
-          averageUtilization: {{ .Values.smithdb.compactionWorker.autoscaling.targetCPUUtilizationPercentage }}
+          averageUtilization: {{ .Values.smithdb.compactionWorker.autoscaling.hpa.targetCPUUtilizationPercentage }}
 {{- end }}

charts/langsmith/templates/smithdb/compaction-worker-scaled-object.yaml

@@ -0,0 +1,55 @@
+{{- if and .Values.smithdb.enabled .Values.smithdb.compactionWorker.autoscaling.keda.enabled }}
+apiVersion: keda.sh/v1alpha1
+kind: ScaledObject
+metadata:
+  name: {{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compactionWorker") }}
+  namespace: {{ .Values.namespace | default .Release.Namespace }}
+  labels:
+    {{- include "langsmith.labels" . | nindent 4 }}
+    app.kubernetes.io/component: {{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compactionWorker") }}
+    {{- with .Values.smithdb.compactionWorker.autoscaling.keda.labels }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+  annotations:
+    scaledobject.keda.sh/transfer-hpa-ownership: "true"
+    {{- with .Values.smithdb.compactionWorker.autoscaling.keda.annotations }}
+    {{- toYaml . | nindent 4 }}
+    {{- end }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: {{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compactionWorker") }}
+  pollingInterval: {{ .Values.smithdb.compactionWorker.autoscaling.keda.pollingInterval }}
+  cooldownPeriod: {{ .Values.smithdb.compactionWorker.autoscaling.keda.cooldownPeriod }}
+  initialCooldownPeriod: {{ .Values.smithdb.compactionWorker.autoscaling.keda.initialCooldownPeriod }}
+  minReplicaCount: {{ .Values.smithdb.compactionWorker.autoscaling.keda.minReplicaCount }}
+  maxReplicaCount: {{ .Values.smithdb.compactionWorker.autoscaling.keda.maxReplicaCount }}
+  advanced:
+    horizontalPodAutoscalerConfig:
+      name: {{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compactionWorker") }}
+      behavior:
+        scaleUp:
+          stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scaleUpStabilizationWindowSeconds }}
+          policies:
+            - type: Pods
+              value: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scalePodCount }}
+              periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scaleUpStabilizationWindowSeconds }}
+        scaleDown:
+          stabilizationWindowSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scaleDownStabilizationWindowSeconds }}
+          policies:
+            - type: Pods
+              value: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scalePodCount }}
+              periodSeconds: {{ .Values.smithdb.compactionWorker.autoscaling.keda.scaleDownStabilizationWindowSeconds }}
+  triggers:
+    - type: metrics-api
+      metadata:
+        targetValue: {{ .Values.smithdb.compactionWorker.autoscaling.keda.pendingJobsThreshold | quote }}
+        url: "http://{{ include "langsmith.smithdb.componentName" (dict "root" . "component" "compaction") }}.{{ .Values.namespace | default .Release.Namespace }}.svc.{{ .Values.clusterDomain }}:{{ .Values.smithdb.compaction.containerPort }}/metrics/jobs"
+        valueLocation: "pending"
+        format: "json"
+    - type: cpu
+      metricType: Utilization
+      metadata:
+        value: {{ .Values.smithdb.compactionWorker.autoscaling.keda.targetCPUUtilizationPercentage | quote }}
+{{- end }}

charts/langsmith/templates/validate.yaml

@@ -543,6 +543,7 @@ AWS Marketplace Helm template verification).
   "playground" .Values.playground.autoscaling
   "queue" .Values.queue.autoscaling
   "ingestQueue" .Values.ingestQueue.autoscaling
+  "smithdb.compactionWorker" .Values.smithdb.compactionWorker.autoscaling
 -}}
 {{- range $serviceName, $autoscaling := $autoscalingServices -}}
 {{- /* Validate that both HPA and KEDA are not enabled at the same time */ -}}

charts/langsmith/tests/langsmith_smithdb_test.yaml

@@ -15,6 +15,7 @@ templates:
   - smithdb/compaction-deployment.yaml
   - smithdb/compaction-worker-deployment.yaml
   - smithdb/compaction-worker-hpa.yaml
+  - smithdb/compaction-worker-scaled-object.yaml
   - smithdb/cluster-manager-deployment.yaml
   - smithdb/migration-job.yaml
   - smithdb/taskdb-postgres-secret.yaml
@@ -774,6 +775,60 @@ tests:
           path: spec.template.spec.containers[0].resources.requests.ephemeral-storage
           value: "100Gi"
 
+  - it: should enable compaction-worker KEDA scaling from pending jobs and CPU
+    values:
+      - ./values/smithdb-enabled.yaml
+    set:
+      smithdb.compactionWorker.autoscaling.hpa.enabled: false
+      smithdb.compactionWorker.autoscaling.keda.enabled: true
+    template: smithdb/compaction-worker-scaled-object.yaml
+    asserts:
+      - isKind:
+          of: ScaledObject
+      - equal:
+          path: spec.minReplicaCount
+          value: 1
+      - equal:
+          path: spec.maxReplicaCount
+          value: 10
+      - equal:
+          path: spec.triggers[0].type
+          value: metrics-api
+      - equal:
+          path: spec.triggers[0].metadata.url
+          value: http://RELEASE-NAME-langsmith-smithdb-compaction.NAMESPACE.svc.cluster.local:8070/metrics/jobs
+      - equal:
+          path: spec.triggers[0].metadata.valueLocation
+          value: pending
+      - equal:
+          path: spec.triggers[0].metadata.targetValue
+          value: "60"
+      - equal:
+          path: spec.triggers[1].type
+          value: cpu
+
+  - it: should omit compaction-worker deployment replicas when KEDA is enabled
+    values:
+      - ./values/smithdb-enabled.yaml
+    set:
+      smithdb.compactionWorker.autoscaling.hpa.enabled: false
+      smithdb.compactionWorker.autoscaling.keda.enabled: true
+    template: smithdb/compaction-worker-deployment.yaml
+    asserts:
+      - notExists:
+          path: spec.replicas
+
+  - it: should render compaction-worker replicas when autoscaling is disabled
+    values:
+      - ./values/smithdb-enabled.yaml
+    set:
+      smithdb.compactionWorker.autoscaling.hpa.enabled: false
+    template: smithdb/compaction-worker-deployment.yaml
+    asserts:
+      - equal:
+          path: spec.replicas
+          value: 2
+
   - it: should expose compaction worker concurrency as a value
     values:
       - ./values/smithdb-enabled.yaml

charts/langsmith/tests/langsmith_smithdb_validation_test.yaml

@@ -2,6 +2,17 @@ suite: test smithdb phase validation
 templates:
   - validate.yaml
 tests:
+  - it: should reject compaction-worker HPA and KEDA being enabled together
+    values:
+      - ./values/smithdb-enabled.yaml
+    set:
+      config.existingSecretName: langsmith-secret
+      smithdb.compactionWorker.autoscaling.hpa.enabled: true
+      smithdb.compactionWorker.autoscaling.keda.enabled: true
+    asserts:
+      - failedTemplate:
+          errorMessage: "Only one of smithdb.compactionWorker.autoscaling.hpa.enabled or smithdb.compactionWorker.autoscaling.keda.enabled can be enabled at the same time."
+
   - it: should reject smithdb without query replicas
     values:
       - ./values/smithdb-enabled.yaml
@@ -203,4 +214,3 @@ tests:
     asserts:
       - failedTemplate:
           errorMessage: "smithdb.commonInitContainers[0] must be a container object."
-

charts/langsmith/values.yaml

@@ -1734,13 +1734,31 @@ smithdb:
       labels: {}
       annotations: {}
     autoscaling:
-      enabled: true
-      minReplicas: 1
-      maxReplicas: 10
-      targetCPUUtilizationPercentage: 60
-      scaleUpStabilizationWindowSeconds: 120
-      scaleDownStabilizationWindowSeconds: 300
-      scalePodCount: 1
+      hpa:
+        # -- Enable HPA scaling.
+        enabled: true
+        minReplicas: 1
+        maxReplicas: 10
+        targetCPUUtilizationPercentage: 60
+        scaleUpStabilizationWindowSeconds: 120
+        scaleDownStabilizationWindowSeconds: 300
+        scalePodCount: 1
+      # -- KEDA scaling based on pending compaction jobs and CPU. Requires KEDA to be installed in the cluster.
+      keda:
+        enabled: false
+        labels: {}
+        annotations: {}
+        pollingInterval: 30
+        cooldownPeriod: 300
+        initialCooldownPeriod: 0
+        minReplicaCount: 1
+        maxReplicaCount: 10
+        targetCPUUtilizationPercentage: 60
+        scaleUpStabilizationWindowSeconds: 120
+        scaleDownStabilizationWindowSeconds: 300
+        scalePodCount: 1
+        # -- Number of pending compaction jobs per worker.
+        pendingJobsThreshold: "60"
 
   clusterManager:
     name: "cluster-manager"