Add conversation-based Loop Runtime autoscaling on GKE

braintrust/examples/gke-loop-runtime/README.md

@@ -59,6 +59,18 @@ for more capacity or pod redundancy.
 With `helm upgrade --reuse-values`, add `--set loopRuntime.replicas=1`
 to replace a retained count of two.
 Each runtime pod claims at most two conversations and has a 256Mi in-memory cache budget.
+
+For conversation-based runtime autoscaling, set `loopRuntime.autoscaling.enabled: true`.
+This uses the same GKE custom-metrics prerequisites as `api.autoscaling` and defaults
+to 1–50 pods at 50% claimed conversation capacity, with a five-minute downscale window.
+The application must expose the gauge `loop_runtime_claimed_conversation_utilization`
+(0–100 percent) at `/metrics` on the runtime HTTP port. `metricsPath`, `minReplicas`,
+`maxReplicas`, and `targetConversationUtilization` are configurable under `loopRuntime.autoscaling`.
+An image without that endpoint or gauge cannot drive metric-based scaling; there is no image-version gate.
+The HPA owns replicas while enabled; disabling it restores `loopRuntime.replicas`.
+Enabling it on an existing Deployment can initially reset the count to one.
+This scales runtime pods only; provision worker replicas separately for the expected sandbox capacity.
+
 Sandbox administration is disabled by default. Set `loopRuntime.enableSandboxAdmin: true`
 to allow authorized sysadmins to inspect and terminate sandboxes.
 

braintrust/templates/_helpers.tpl

@@ -201,6 +201,32 @@ predates Loop Runtime have no loopRuntime map, which means disabled.
 {{- end -}}
 {{- end -}}
 
+{{/* Autoscaling defaults also cover partial maps supplied with --reuse-values. */}}
+{{- define "braintrust.loopRuntime.autoscaling" -}}
+{{- $runtime := .Values.loopRuntime | default dict -}}
+{{- $defaults := dict "enabled" false "minReplicas" 1 "maxReplicas" 50 "metricsPath" "/metrics" "targetConversationUtilization" "50" -}}
+{{- $config := mergeOverwrite (deepCopy $defaults) (deepCopy ($runtime.autoscaling | default dict)) -}}
+{{- range $key, $value := $defaults -}}
+{{- if eq (index $config $key) nil -}}{{- $_ := set $config $key $value -}}{{- end -}}
+{{- end -}}
+{{- if and (include "braintrust.loopRuntime.enabled" .) $config.enabled -}}
+{{- include "braintrust.loopRuntime.validate" . -}}
+{{- if not (.Capabilities.APIVersions.Has "autoscaling.gke.io/v1beta1") -}}
+{{- fail "loopRuntime.autoscaling requires the AutoscalingMetric API (autoscaling.gke.io/v1beta1). For helm template, pass --api-versions=autoscaling.gke.io/v1beta1." -}}
+{{- end -}}
+{{- if or (not (regexMatch "^[1-9][0-9]*$" (toString $config.minReplicas))) (not (regexMatch "^[1-9][0-9]*$" (toString $config.maxReplicas))) (lt (int $config.maxReplicas) (int $config.minReplicas)) -}}
+{{- fail "loopRuntime.autoscaling requires positive minReplicas and maxReplicas >= minReplicas" -}}
+{{- end -}}
+{{- if not (and (gt (float64 $config.targetConversationUtilization) 0.0) (le (float64 $config.targetConversationUtilization) 100.0)) -}}
+{{- fail "loopRuntime.autoscaling.targetConversationUtilization must be greater than 0 and at most 100" -}}
+{{- end -}}
+{{- if or (not (kindIs "string" $config.metricsPath)) (not (hasPrefix "/" (toString $config.metricsPath))) -}}
+{{- fail "loopRuntime.autoscaling.metricsPath must be an absolute HTTP path" -}}
+{{- end -}}
+{{- toYaml $config -}}
+{{- end -}}
+{{- end -}}
+
 {{/* Internal cluster URL for the Loop Runtime service. */}}
 {{- define "braintrust.loopRuntimeInternalUrl" -}}
 http://{{ .Values.loopRuntime.service.name | default .Values.loopRuntime.name }}.{{ include "braintrust.namespace" . }}:{{ .Values.loopRuntime.service.port }}

braintrust/templates/loop-runtime-autoscaling-metric.yaml

@@ -0,0 +1,24 @@
+{{- $autoscaling := include "braintrust.loopRuntime.autoscaling" . | fromYaml -}}
+{{- if $autoscaling -}}
+apiVersion: autoscaling.gke.io/v1beta1
+kind: AutoscalingMetric
+metadata:
+  name: {{ .Values.loopRuntime.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  labels:
+    {{- mergeOverwrite (deepCopy .Values.global.labels) (deepCopy .Values.loopRuntime.labels) | toYaml | nindent 4 }}
+spec:
+  metrics:
+    - pod:
+        selector:
+          matchLabels:
+            app: {{ .Values.loopRuntime.name }}
+        containers:
+          - endpoint:
+              port: {{ .Values.loopRuntime.service.port }}
+              path: {{ $autoscaling.metricsPath }}
+            metrics:
+              - gauge:
+                  name: claimed-conversation-utilization
+                  prometheusMetricName: loop_runtime_claimed_conversation_utilization
+{{- end }}

braintrust/templates/loop-runtime-deployment.yaml

@@ -1,5 +1,6 @@
 {{- include "braintrust.loopRuntime.validate" . -}}
 {{- if include "braintrust.loopRuntime.enabled" . -}}
+{{- $autoscaling := include "braintrust.loopRuntime.autoscaling" . | fromYaml -}}
 {{- $podLabels := mergeOverwrite (deepCopy .Values.global.labels) (deepCopy .Values.loopRuntime.labels) (deepCopy .Values.loopRuntime.podLabels) (dict "app" .Values.loopRuntime.name "gke-workload-identity/use" "true") -}}
 apiVersion: apps/v1
 kind: Deployment
@@ -15,7 +16,9 @@ metadata:
     {{- toYaml . | nindent 4 }}
   {{- end }}
 spec:
+  {{- if not $autoscaling }}
   replicas: {{ ternary 1 .Values.loopRuntime.replicas (has .Values.loopRuntime.replicas (list nil "")) }}
+  {{- end }}
   # Start the replacement before draining the current runtime.
   strategy:
     type: RollingUpdate

braintrust/templates/loop-runtime-hpa.yaml

@@ -0,0 +1,28 @@
+{{- $autoscaling := include "braintrust.loopRuntime.autoscaling" . | fromYaml -}}
+{{- if $autoscaling -}}
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: {{ .Values.loopRuntime.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  labels:
+    {{- mergeOverwrite (deepCopy .Values.global.labels) (deepCopy .Values.loopRuntime.labels) | toYaml | nindent 4 }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: {{ .Values.loopRuntime.name }}
+  minReplicas: {{ $autoscaling.minReplicas }}
+  maxReplicas: {{ $autoscaling.maxReplicas }}
+  metrics:
+    - type: Pods
+      pods:
+        metric:
+          name: autoscaling.gke.io|{{ .Values.loopRuntime.name }}|claimed-conversation-utilization
+        target:
+          type: AverageValue
+          averageValue: {{ $autoscaling.targetConversationUtilization | quote }}
+  behavior:
+    scaleDown:
+      stabilizationWindowSeconds: 300
+{{- end }}

braintrust/tests/__fixtures__/reuse-loop-runtime-autoscaling-values.yaml

@@ -0,0 +1,9 @@
+# Nulls remove chart defaults during coalescing, leaving the enabled override
+# to simulate --reuse-values from a chart without autoscaling settings.
+loopRuntime:
+  autoscaling:
+    enabled: null
+    minReplicas: null
+    maxReplicas: null
+    metricsPath: null
+    targetConversationUtilization: null

braintrust/tests/api-autoscaling-crd_test.yaml

@@ -1,6 +1,7 @@
-suite: test API autoscaling CRD requirement
+suite: test GKE autoscaling CRD requirements
 templates:
   - api-hpa.yaml
+  - loop-runtime-hpa.yaml
 # Do not advertise autoscaling.gke.io here — this suite asserts the fail path.
 tests:
   - it: should fail when AutoscalingMetric API is unavailable
@@ -15,3 +16,14 @@ tests:
     asserts:
       - failedTemplate:
           errorPattern: "api\\.autoscaling requires the AutoscalingMetric API \\(autoscaling\\.gke\\.io/v1beta1\\)"
+
+  - it: should require the AutoscalingMetric API for Loop Runtime autoscaling
+    template: loop-runtime-hpa.yaml
+    values:
+      - __fixtures__/base-values.yaml
+      - __fixtures__/loop-runtime-values.yaml
+    set:
+      loopRuntime.autoscaling.enabled: true
+    asserts:
+      - failedTemplate:
+          errorPattern: "loopRuntime\\.autoscaling requires the AutoscalingMetric API"