Add Loop Runtime on GKE with isolated Sandbox Service workers

braintrust/README.md

@@ -18,6 +18,16 @@ The `braintrust-secrets` secret must contain the following keys:
 | `GCS_ACCESS_KEY_ID` | Google HMAC Access ID string | Valid S3 API Key Id (only required if `cloud` is `google` and if `enableGcsAuth` is `false`) |
 | `GCS_SECRET_ACCESS_KEY` | Google HMAC Secret string | Valid S3 Secret string (only required if `cloud` is `google` and if `enableGcsAuth` is `false`) |
 
+## Optional: Loop Runtime
+
+Set `loopRuntime.enabled: true` to run Loop Runtime on GKE (`cloud: google`).
+On AWS it runs on ECS through the Terraform module. Sandboxes run on Sandbox
+Service workers in a separate GKE cluster; see
+[examples/gke-loop-runtime](examples/gke-loop-runtime/README.md). The Terraform
+outputs it needs come from
+[terraform-google-braintrust-data-plane#28](https://github.com/braintrustdata/terraform-google-braintrust-data-plane/pull/28),
+which is not merged yet.
+
 ## Azure Key Vault Driver Integration
 
 If you're using Azure, the Azure Key Vault CSI driver is default enabled and will automatically sync secrets from Azure Key Vault into Kubernetes. This eliminates the need to manually create and manage the `braintrust-secrets` Kubernetes secret.

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

@@ -0,0 +1,58 @@
+# Loop Runtime on GKE
+
+Loop Runtime runs in the primary GKE cluster. Its sandboxes run on Sandbox
+Service workers in a separate isolated cluster. Both clusters come from the
+GCP Terraform module with:
+
+```hcl
+gke_isolated_workers = {
+  network_cidr     = "10.64.0.0/16" # must not overlap existing ranges
+  authorized_cidrs = ["192.0.2.10/32"] # replace with operator and CI egress ranges
+  discovery        = { dns_name = "sandbox.<deployment>.internal" }
+}
+```
+
+`<...>` values below come from `terraform output gke_isolated_workers` and
+`terraform output loop_runtime_service_account`.
+
+Use compatible API, Brainstore, and AI Gateway images for the runtime build.
+The API must report `ff_version` 82 or newer, and the Rust gateway must support
+`x-bt-provider-model-map`.
+Deploy the gateway with `aiGateway.enabled=true` first and wait for it to become
+ready before applying this example, which turns on `aiGateway.useGateway`.
+The TypeScript API proxy alone cannot restore the provider model IDs sent by Codex.
+
+1. Create the mTLS Secrets in both clusters:
+
+   ```sh
+   braintrust/examples/gke-loop-runtime/create-sandbox-tls.sh "<discovery.hostname>" <worker> <primary>
+   ```
+
+2. Install the workers. The configured Sandbox Service image prepares local SSDs before starting the worker:
+
+   ```sh
+   helm --kube-context <worker> upgrade --install sandbox-service \
+     https://braintrust-cf-us-east-1.s3.amazonaws.com/sandbox-service/1f4b6826c61064575b3db956b67625133257bd8a6e49e4ea86cdb55012dd53a7/sandbox-service-0.1.0.tgz \
+     --namespace braintrust-sandbox --values braintrust/examples/gke-loop-runtime/sandbox-service.yaml --wait
+   ```
+
+3. Install this chart with your GKE values plus `values.yaml` from this directory.
+   For a Standard primary cluster, set `loopRuntime.nodeSelector` to
+   `{"braintrust/workload": "services"}` to use the services pool:
+
+   ```sh
+   helm --kube-context <primary> upgrade --install braintrust ./braintrust --namespace braintrust \
+     --values <your GKE values> --values braintrust/examples/gke-loop-runtime/values.yaml --wait
+   ```
+
+Workers run at a fixed count. One `c4-standard-8-lssd` worker holds about seven
+1 vCPU / 2 GiB sandboxes; beyond that, new sandboxes fail with
+`CapacityExhausted`. Raise `replicas` in `sandbox-service.yaml` to add workers.
+
+Each runtime pod claims at most two conversations and has a 256Mi in-memory cache budget.
+Sandbox administration is disabled by default. Set `loopRuntime.enableSandboxAdmin: true`
+to allow authorized sysadmins to inspect and terminate sandboxes.
+
+To shut down, set the worker release to `replicas: 0`, wait for the workers to
+drain, then disable `loopRuntime.enabled` and uninstall the worker release. Its
+pre-delete hook refuses to uninstall until workers drain.

braintrust/examples/gke-loop-runtime/create-sandbox-tls.sh

@@ -0,0 +1,72 @@
+#!/bin/bash
+# Creates the mTLS Secrets for Loop Runtime and Sandbox Service, matching the
+# AWS module: separate server and client CAs, P-256 keys, ten-year validity.
+#
+# Usage: create-sandbox-tls.sh <worker hostname> <worker kube context> <runtime kube context>
+#   worker hostname: Terraform output gke_isolated_workers.discovery.hostname
+set -euo pipefail
+
+if [ "$#" -ne 3 ]; then
+  echo "usage: $0 <worker hostname> <worker kube context> <runtime kube context>" >&2
+  exit 2
+fi
+HOSTNAME_SAN="$1"
+WORKER_CONTEXT="$2"
+RUNTIME_CONTEXT="$3"
+WORKER_NAMESPACE="${WORKER_NAMESPACE:-braintrust-sandbox}"
+RUNTIME_NAMESPACE="${RUNTIME_NAMESPACE:-braintrust}"
+DAYS=3650
+
+for context in "$WORKER_CONTEXT" "$RUNTIME_CONTEXT"; do
+  kubectl --context "$context" get namespaces >/dev/null
+done
+# sandbox-tls is written last, so a failed run can be rerun.
+if kubectl --context "$WORKER_CONTEXT" -n "$WORKER_NAMESPACE" get secret sandbox-tls >/dev/null 2>&1; then
+  echo "sandbox-tls already exists in $WORKER_CONTEXT/$WORKER_NAMESPACE; refusing to replace it." >&2
+  exit 1
+fi
+
+dir="$(mktemp -d)"
+trap 'rm -rf "$dir"' EXIT
+cd "$dir"
+
+for purpose in server client; do
+  openssl genpkey -algorithm EC -pkeyopt ec_paramgen_curve:P-256 -out "$purpose-ca.key" 2>/dev/null
+  openssl req -x509 -new -sha256 -key "$purpose-ca.key" -days "$DAYS" \
+    -subj "/CN=sandbox-$purpose-ca" \
+    -addext "basicConstraints=critical,CA:TRUE,pathlen:0" \
+    -addext "keyUsage=critical,keyCertSign,cRLSign" \
+    -out "$purpose-ca.pem"
+  openssl genpkey -algorithm EC -pkeyopt ec_paramgen_curve:P-256 -out "$purpose.key" 2>/dev/null
+  {
+    echo "basicConstraints=critical,CA:FALSE"
+    echo "keyUsage=critical,digitalSignature"
+    echo "authorityKeyIdentifier=keyid"
+    if [ "$purpose" = server ]; then
+      echo "extendedKeyUsage=serverAuth"
+      echo "subjectAltName=DNS:$HOSTNAME_SAN"
+    else
+      echo "extendedKeyUsage=clientAuth"
+    fi
+  } >"$purpose.ext"
+  openssl req -new -key "$purpose.key" -subj "/CN=sandbox-$purpose" |
+    openssl x509 -req -sha256 -days "$DAYS" -CA "$purpose-ca.pem" -CAkey "$purpose-ca.key" \
+      -set_serial "0x$(openssl rand -hex 16)" -extfile "$purpose.ext" -out "$purpose.pem" 2>/dev/null
+done
+
+apply_secret() {
+  kubectl --context "$1" create namespace "$2" --dry-run=client -o yaml | kubectl --context "$1" apply --server-side -f - >/dev/null
+  kubectl --context "$1" -n "$2" create secret generic "$3" "${@:4}" --dry-run=client -o yaml |
+    kubectl --context "$1" apply --server-side -f -
+  kubectl --context "$1" -n "$2" annotate secret "$3" kubectl.kubernetes.io/last-applied-configuration-
+}
+
+# The worker controller and Loop Runtime present the client certificate and trust the server CA.
+for target in "$RUNTIME_CONTEXT $RUNTIME_NAMESPACE" "$WORKER_CONTEXT $WORKER_NAMESPACE"; do
+  # shellcheck disable=SC2086
+  apply_secret $target sandbox-client-tls \
+    --from-file=client.pem=client.pem --from-file=client.key=client.key --from-file=server-ca.pem=server-ca.pem
+done
+# Workers present the server certificate and trust the client CA.
+apply_secret "$WORKER_CONTEXT" "$WORKER_NAMESPACE" sandbox-tls \
+  --from-file=cert.pem=server.pem --from-file=key.pem=server.key --from-file=ca_cert.pem=client-ca.pem

braintrust/examples/gke-loop-runtime/sandbox-service.yaml

@@ -0,0 +1,27 @@
+# Values for the Sandbox Service worker chart, installed in the isolated worker cluster.
+# Every <...> value comes from the GCP Terraform output gke_isolated_workers.
+hostname: "<discovery.hostname>"
+image: public.ecr.aws/braintrust/sandbox-service:38b05599cdc5e4c3a213bbf0718211b397183a7a
+# Fixed worker capacity. The worker node pool autoscales, so raising this is a values change.
+replicas: 1
+gitResources: true
+nodePrep:
+  enabled: true
+prewarmTemplates:
+  - image: public.ecr.aws/braintrust/codex-runtime@sha256:490ab364fd6a787f73957d3d46bc5e173fb418c47d53bcf7a2d3afba8b419049
+    resources:
+      vcpu_count: 1
+      memory_mib: 2048
+    mounts: []
+    durable_file_systems: []
+    policy:
+      networking:
+        type: disabled
+      credentials: []
+    default_workdir: /home/exo/workspace
+externalDns:
+  enabled: true
+  provider: google
+  # One entry per argument in discovery.external_dns_extra_args.
+  extraArgs:
+    - "<argument>"

braintrust/examples/gke-loop-runtime/values.yaml

@@ -0,0 +1,18 @@
+# Layer on top of your existing GKE values (for example, ../google-autopilot/values.yaml).
+# The sandbox and service-account values come from the GCP Terraform module outputs.
+aiGateway:
+  enabled: true
+  useGateway: true
+
+loopRuntime:
+  enabled: true
+  serviceAccount:
+    googleServiceAccount: "<loop_runtime_service_account>"
+  sandbox:
+    dnsUrl: "<gke_isolated_workers.discovery.endpoint>"
+    egressGateway:
+      address: "<gke_isolated_workers.loop_egress.address>"
+      # One entry per range in gke_isolated_workers.loop_egress.source_ranges.
+      sourceRanges:
+        - "<range>"
+      url: "<gke_isolated_workers.loop_egress.url>"

braintrust/templates/_helpers.tpl

@@ -146,6 +146,61 @@ Internal cluster URL for the AI Gateway service.
 http://{{ .Values.aiGateway.service.name | default .Values.aiGateway.name }}.{{ include "braintrust.namespace" . }}:{{ .Values.aiGateway.service.port }}
 {{- end -}}
 
+{{/*
+Whether Loop Runtime is enabled. Values reused from a chart release that
+predates Loop Runtime have no loopRuntime map, which means disabled.
+*/}}
+{{- define "braintrust.loopRuntime.enabled" -}}
+{{- if (.Values.loopRuntime | default dict).enabled -}}true{{- end -}}
+{{- end -}}
+
+{{/* Validate the Loop Runtime configuration. */}}
+{{- define "braintrust.loopRuntime.validate" -}}
+{{- if include "braintrust.loopRuntime.enabled" . -}}
+{{- range $key := list "sandbox" "image" "service" "serviceAccount" "resources" "annotations" -}}
+{{- if not (kindIs "map" (get $.Values.loopRuntime $key)) -}}
+{{- fail "loopRuntime defaults are missing. Values reused from a chart release that predates Loop Runtime omit them; upgrade with --reset-then-reuse-values." -}}
+{{- end -}}
+{{- end -}}
+{{- if ne .Values.cloud "google" -}}
+{{- fail "loopRuntime.enabled requires cloud: google. On AWS, Loop Runtime runs on ECS through the Terraform module." -}}
+{{- end -}}
+{{- $sandbox := .Values.loopRuntime.sandbox -}}
+{{- $egress := $sandbox.egressGateway | default dict -}}
+{{- $_ := required "objectStorage.google.brainstoreBucket is required" .Values.objectStorage.google.brainstoreBucket -}}
+{{- $_ := required "objectStorage.google.apiBucket is required" .Values.objectStorage.google.apiBucket -}}
+{{- $_ := required "loopRuntime.image.repository is required" .Values.loopRuntime.image.repository -}}
+{{- $_ := required "loopRuntime.image.tag is required" (.Values.loopRuntime.image.tag | default "" | toString) -}}
+{{- $_ := required "loopRuntime.serviceAccount.googleServiceAccount is required" .Values.loopRuntime.serviceAccount.googleServiceAccount -}}
+{{- $_ := required "loopRuntime.sandbox.dnsUrl is required" $sandbox.dnsUrl -}}
+{{- $_ := required "loopRuntime.sandbox.image is required" $sandbox.image -}}
+{{- $_ := required "loopRuntime.sandbox.egressGateway.address is required" $egress.address -}}
+{{- $_ := required "loopRuntime.sandbox.egressGateway.url is required" $egress.url -}}
+{{- if not $egress.sourceRanges -}}
+{{- fail "loopRuntime.sandbox.egressGateway.sourceRanges is required" -}}
+{{- end -}}
+{{- 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 }}
+{{- end -}}
+
+{{/* Product-owned Loop Runtime routes for an Istio VirtualService. */}}
+{{- define "braintrust.loopRuntimeVirtualServiceRoutes" -}}
+- match:
+    - uri:
+        exact: /loop/runtime
+    - uri:
+        prefix: /loop/runtime/
+  route:
+    - destination:
+        host: {{ .Values.loopRuntime.service.name | default .Values.loopRuntime.name }}
+        port:
+          number: {{ .Values.loopRuntime.service.port }}
+{{- end -}}
+
 {{/*
 Validate API autoscaling prerequisites (GKE + AutoscalingMetric CRD).
 */}}

braintrust/templates/api-configmap.yaml

@@ -15,6 +15,7 @@
 {{- if and (gt (len $btqlStrictOrgIds) 0) (gt (len $btqlBestEffortOrgIds) 0) -}}
 {{- fail "Set either btqlAuditLogsStrictOrgIds or btqlAuditLogsBestEffortOrgIds, not both." -}}
 {{- end -}}
+{{- include "braintrust.loopRuntime.validate" . -}}
 {{- range $btqlStrictOrgIds -}}
 {{- if not (regexMatch "^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$" .) -}}
 {{- fail (printf "btqlAuditLogsStrictOrgIds: %q is not a valid UUID" .) -}}
@@ -113,6 +114,9 @@ data:
   {{- if .Values.aiGateway.useGateway }}
   GATEWAY_URL: {{ include "braintrust.aiGatewayInternalUrl" . | quote }}
   {{- end }}
+  {{- if include "braintrust.loopRuntime.enabled" . }}
+  LOOP_RUNTIME_URL: {{ include "braintrust.loopRuntimeInternalUrl" . | quote }}
+  {{- end }}
   {{- if $btqlStrictOrgIds }}
   BTQL_AUDIT_LOGS_STRICT_ORG_IDS: {{ $btqlStrictOrgIds | join "," | quote }}
   {{- end }}

braintrust/templates/loop-runtime-configmap.yaml

@@ -0,0 +1,68 @@
+{{- include "braintrust.loopRuntime.validate" . -}}
+{{- if include "braintrust.loopRuntime.enabled" . -}}
+{{- $bucket := printf "gs://%s" .Values.objectStorage.google.brainstoreBucket -}}
+{{- $sandbox := .Values.loopRuntime.sandbox -}}
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: {{ .Values.loopRuntime.name }}
+  namespace: {{ include "braintrust.namespace" . }}
+  {{- with (mergeOverwrite (deepCopy .Values.global.labels) (deepCopy .Values.loopRuntime.labels)) }}
+  labels:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+  {{- with .Values.loopRuntime.annotations.configmap }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+# Matches the ECS Loop Runtime environment in the AWS Terraform module.
+data:
+  BRAINTRUST_LEGACY_IDS: "true"
+  LOOP_RUNTIME_HOST: "0.0.0.0"
+  LOOP_RUNTIME_PORT: {{ .Values.loopRuntime.service.port | quote }}
+  LOOP_RUNTIME_DRAIN_TIMEOUT_SECONDS: "90"
+  LOOP_RUNTIME_REDIS_NAMESPACE: "loop-runtime:{{ .Release.Name }}"
+  LOOP_RUNTIME_STATE_DIR: /mnt/tmp/loop-runtime
+  {{- /* The runtime refuses an empty ORG_NAME; the API treats empty as every org. */}}
+  ORG_NAME: {{ .Values.global.orgName | default "" | toString | trim | default "*" | quote }}
+  BRAINTRUST_API_URL: {{ include "braintrust.apiInternalUrl" . | quote }}
+  LOOP_RUNTIME_AI_PROXY_URL: {{ include "braintrust.apiAiProxyInternalUrl" . | quote }}
+  LOOP_RUNTIME_BRAINSTORE_READER_URL: "http://{{ .Values.brainstore.fastreader.service.name | default .Values.brainstore.fastreader.name }}.{{ include "braintrust.namespace" . }}:{{ .Values.brainstore.fastreader.service.port }}"
+  BRAINSTORE_VERBOSE: "1"
+  BRAINSTORE_INDEX_URI: "{{ $bucket }}/brainstore/index"
+  BRAINSTORE_REALTIME_WAL_URI: "{{ $bucket }}/brainstore/wal"
+  BRAINSTORE_CODE_BUNDLE_URI: "gs://{{ .Values.objectStorage.google.apiBucket }}"
+  BRAINSTORE_CACHE_DIR: /mnt/tmp/brainstore
+  BRAINSTORE_OBJECT_STORE_CACHE_FILE_SIZE: 5Gi
+  BRAINSTORE_OBJECT_STORE_CACHE_MEMORY_LIMIT: 256Mi
+  LOOP_RUNTIME_MAX_CLAIMED_CONVERSATIONS: "2"
+  BRAINSTORE_CONTROL_PLANE_TELEMETRY: {{ .Values.global.controlPlaneTelemetry | quote }}
+  NO_COLOR: "1"
+  EXO_SANDBOX_SERVICE_DNS_URL: {{ $sandbox.dnsUrl | quote }}
+  EXO_SANDBOX_SERVICE_IMAGE: {{ $sandbox.image | quote }}
+  EXO_SANDBOX_VCPU_COUNT: {{ $sandbox.vcpuCount | quote }}
+  EXO_SANDBOX_MEMORY_MIB: {{ $sandbox.memoryMib | quote }}
+  SANDBOX_SERVICE_TLS_CERT: /etc/sandbox-tls/client.pem
+  SANDBOX_SERVICE_TLS_KEY: /etc/sandbox-tls/client.key
+  SANDBOX_SERVICE_TLS_CA: /etc/sandbox-tls/server-ca.pem
+  LOOP_RUNTIME_SANDBOX_FORWARD_PROXY_LISTEN: "0.0.0.0:4002"
+  LOOP_RUNTIME_SANDBOX_EGRESS_GATEWAY_URL: {{ $sandbox.egressGateway.url | quote }}
+  LOOP_RUNTIME_SANDBOX_SERVICE_EGRESS_GATEWAY_URL: {{ $sandbox.egressGateway.url | quote }}
+  LOOP_RUNTIME_SANDBOX_EGRESS_MODE: restricted
+  LOOP_RUNTIME_ENABLE_SANDBOX_ADMIN: {{ .Values.loopRuntime.enableSandboxAdmin | default false | quote }}
+  {{- with (.Values.global.allowedOrgIds | default "" | toString | trim) }}
+  ALLOWED_ORG_IDS: {{ . | quote }}
+  {{- end }}
+  {{- if eq .Values.brainstore.locksBackend "objectStorage" }}
+  BRAINSTORE_LOCKS_URI: "{{ $bucket }}/brainstore/locks"
+  {{- end }}
+  {{- if .Values.brainstoreWalFooterVersion }}
+  BRAINSTORE_WAL_FOOTER_VERSION: {{ .Values.brainstoreWalFooterVersion | quote }}
+  {{- end }}
+  {{- if .Values.skipPgForBrainstoreObjects }}
+  SKIP_PG_FOR_BRAINSTORE_OBJECTS: {{ .Values.skipPgForBrainstoreObjects | quote }}
+  {{- end }}
+  {{- if or .Values.brainstoreWalFooterVersion .Values.skipPgForBrainstoreObjects }}
+  BRAINSTORE_WAL_USE_EFFICIENT_FORMAT: "true"
+  {{- end }}
+{{- end }}