diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 50804e27..d4b2a79a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -37,6 +37,9 @@ jobs: - name: KVM guest resolver handoff run: tests/kvm_guest_dns.sh + - name: Worker drains before sandbox shutdown + run: tests/worker_sandbox_shutdown.sh + - name: Sandbox-runner liveness checks run: tests/sandbox_runner_healthcheck.sh diff --git a/docker-compose.local-dev.yml b/docker-compose.local-dev.yml index 4d8f5a33..9e072008 100644 --- a/docker-compose.local-dev.yml +++ b/docker-compose.local-dev.yml @@ -103,6 +103,7 @@ services: # Local Service API (combined API + Workers, no auth required) # For production, split into api + worker-sandbox pods service: + stop_grace_period: 400s build: context: . dockerfile: service/Dockerfile.local @@ -111,6 +112,7 @@ services: ports: - 3112:3112 environment: + - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-} - LOCAL_MODE=${LOCAL_MODE:-true} - SANDBOX_ENDPOINT=http://sandbox:2000/api/v2 - FILE_SERVER_URL=http://file_server:3000 diff --git a/docker-compose.scalable.yml b/docker-compose.scalable.yml index 49c89148..cd3a6be2 100644 --- a/docker-compose.scalable.yml +++ b/docker-compose.scalable.yml @@ -166,6 +166,7 @@ services: # Each pod runs its own microVM — PYTHON_CONCURRENCY is per-sandbox # ============================================================================= worker-sandbox: + stop_grace_period: 400s build: context: . dockerfile: docker/Dockerfile.worker-sandbox @@ -174,6 +175,7 @@ services: devices: - ${KVM_DEVICE_PATH:-/dev/kvm}:/dev/kvm environment: + - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-} - KVM_ENABLED=${KVM_ENABLED:-true} - LOCAL_MODE=${LOCAL_MODE:-true} # Worker config diff --git a/docker-compose.yaml b/docker-compose.yaml index db517b4c..4b36c1d5 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -58,12 +58,14 @@ services: - tool_call_server service-worker: + stop_grace_period: 400s build: context: . dockerfile: service/Dockerfile target: worker container_name: service-worker environment: + - CODEAPI_SHUTDOWN_TIMEOUT_MS=${CODEAPI_SHUTDOWN_TIMEOUT_MS:-} - LOCAL_MODE=${LOCAL_MODE:-true} - CODEAPI_HARDENED_SANDBOX_MODE=${CODEAPI_HARDENED_SANDBOX_MODE:-true} - CODEAPI_BRIDGE_TOKEN=${CODEAPI_BRIDGE_TOKEN:-} diff --git a/docker/supervisor.sh b/docker/supervisor.sh index f087655d..0f640a0d 100644 --- a/docker/supervisor.sh +++ b/docker/supervisor.sh @@ -38,11 +38,13 @@ cleanup() { echo "Received shutdown signal, stopping processes..." - if [ -n "$SANDBOX_PID" ]; then - kill -TERM "$SANDBOX_PID" 2>/dev/null || true - fi if [ -n "$WORKER_PID" ]; then kill -TERM "$WORKER_PID" 2>/dev/null || true + # Active jobs still need the sandbox while the worker drains. + wait "$WORKER_PID" 2>/dev/null || true + fi + if [ -n "$SANDBOX_PID" ]; then + kill -TERM "$SANDBOX_PID" 2>/dev/null || true fi wait diff --git a/helm/codeapi/Chart.yaml b/helm/codeapi/Chart.yaml index dfcdd45b..6de60a24 100644 --- a/helm/codeapi/Chart.yaml +++ b/helm/codeapi/Chart.yaml @@ -3,7 +3,7 @@ apiVersion: v2 name: codeapi description: A Helm chart for Code Interpreter API - scalable code execution service type: application -version: 0.3.2 # Chart version (bump this when you change the chart) +version: 0.3.3 # Chart version (bump this when you change the chart) appVersion: "2.0.0" # App version (bump this when you change the app) # Keywords for searching diff --git a/helm/codeapi/templates/worker-sandbox-deployment.yaml b/helm/codeapi/templates/worker-sandbox-deployment.yaml index 2819d998..a20733b1 100644 --- a/helm/codeapi/templates/worker-sandbox-deployment.yaml +++ b/helm/codeapi/templates/worker-sandbox-deployment.yaml @@ -121,6 +121,7 @@ spec: labels: {{- include "codeapi.serviceWorker.selectorLabels" . | nindent 8 }} spec: + terminationGracePeriodSeconds: {{ .Values.workerSandbox.serviceWorker.terminationGracePeriodSeconds | default (add 100 (div (add (int .Values.workerSandbox.config.jobTimeout) 999) 1000)) }} {{- if or .Values.workerSandbox.serviceAccount.create .Values.workerSandbox.serviceAccount.name }} serviceAccountName: {{ $serviceWorkerServiceAccount }} {{- end }} diff --git a/helm/codeapi/values.yaml b/helm/codeapi/values.yaml index 4858fd71..3eb25285 100644 --- a/helm/codeapi/values.yaml +++ b/helm/codeapi/values.yaml @@ -207,6 +207,8 @@ workerSandbox: # nodeSelector/tolerations/affinity below. Use sandboxRunner to place only the # untrusted-code tier on KVM-capable nodes. serviceWorker: + # Defaults to jobTimeout (rounded up to seconds) + 100 seconds. + terminationGracePeriodSeconds: null # Defaults to workerSandbox.replicaCount when unset. # Keep this as fixed HA capacity; scale sandboxRunner for execution throughput. replicaCount: null diff --git a/service/src/config.ts b/service/src/config.ts index b3e397b8..135e5585 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -250,6 +250,19 @@ export function hostedAppOperationTimeoutMs(): number { + env.HOSTED_APP_START_TIMEOUT_MS + 30_000; } +export function gracefulShutdownTimeoutMs(hasWorkers: boolean): number { + const defaultTimeout = hasWorkers + ? Math.max( + jobCompletionWaitTimeoutMs(env.JOB_TIMEOUT, env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS), + env.HOSTED_APPS_ENABLED + ? hostedAppOperationTimeoutMs() + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + 30_000 + : 0, + ) + : 30_000; + // Node/Bun timers overflow to an immediate timeout above this limit. + return Math.min(resolvePositiveIntEnv(process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS, defaultTimeout), 2_147_483_647); +} + function configuredNumber(raw: string | undefined, fallback: number): number { return raw == null || raw.trim() === '' ? fallback : Number(raw); } diff --git a/service/src/lifecycle-timeout.test.ts b/service/src/lifecycle-timeout.test.ts new file mode 100644 index 00000000..dffe46fe --- /dev/null +++ b/service/src/lifecycle-timeout.test.ts @@ -0,0 +1,60 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import { env, gracefulShutdownTimeoutMs } from './config'; + +const original = { + JOB_TIMEOUT: env.JOB_TIMEOUT, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, + EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, + HOSTED_APPS_ENABLED: env.HOSTED_APPS_ENABLED, + CHECKPOINT_TIMEOUT_MS: env.CHECKPOINT_TIMEOUT_MS, +}; +const originalOverride = process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS; + +beforeEach(() => { + Object.assign(env, { + JOB_TIMEOUT: 300_000, + LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS: 60_000, + EGRESS_GATEWAY_REVOKE_TIMEOUT_MS: 5_000, + HOSTED_APPS_ENABLED: false, + CHECKPOINT_TIMEOUT_MS: 60_000, + }); + delete process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS; +}); + +afterEach(() => { + Object.assign(env, original); + if (originalOverride === undefined) delete process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS; + else process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = originalOverride; +}); + +test('workers can finish a five-minute job and cleanup while API-only shutdown stays bounded', () => { + expect(gracefulShutdownTimeoutMs(true)).toBe(370_000); + expect(gracefulShutdownTimeoutMs(false)).toBe(30_000); +}); + +test('worker shutdown follows configured execution and cleanup deadlines', () => { + env.JOB_TIMEOUT = 1_200_000; + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS = 90_000; + env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS = 20_000; + expect(gracefulShutdownTimeoutMs(true)).toBe(1_315_000); +}); + +test('hosted-app workers retain their longer operation budget', () => { + env.HOSTED_APPS_ENABLED = true; + expect(gracefulShutdownTimeoutMs(true)).toBe(1_110_000); + env.JOB_TIMEOUT = 1_200_000; + expect(gracefulShutdownTimeoutMs(true)).toBe(1_270_000); + expect(gracefulShutdownTimeoutMs(false)).toBe(30_000); +}); + +test('explicit shutdown overrides remain finite and invalid values use the default', () => { + process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = '420000'; + expect(gracefulShutdownTimeoutMs(true)).toBe(420_000); + expect(gracefulShutdownTimeoutMs(false)).toBe(420_000); + for (const invalid of ['', ' ', '0', '-1', 'NaN', 'Infinity']) { + process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = invalid; + expect(gracefulShutdownTimeoutMs(true)).toBe(370_000); + } + process.env.CODEAPI_SHUTDOWN_TIMEOUT_MS = '2147483648'; + expect(gracefulShutdownTimeoutMs(true)).toBe(2_147_483_647); +}); diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index f3747ceb..f183213c 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -7,7 +7,7 @@ import { closeQueueConnections, } from './queue'; import { validateStartupAuthConfig } from './auth/startup'; -import { env, hostedAppOperationTimeoutMs } from './config'; +import { env, gracefulShutdownTimeoutMs } from './config'; import { validateApiBridgePolicy, validateApiHardenedConfig, @@ -241,9 +241,7 @@ export async function gracefulShutdown(): Promise { const shutdownTimeout = setTimeout(() => { logger.error('Shutdown timeout reached, forcing exit'); process.exit(1); - }, hasWorkers && env.HOSTED_APPS_ENABLED - ? hostedAppOperationTimeoutMs() + env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS + 30_000 - : 30_000); + }, gracefulShutdownTimeoutMs(hasWorkers)); try { if (hasWorkers) { diff --git a/tests/worker_sandbox_shutdown.sh b/tests/worker_sandbox_shutdown.sh new file mode 100755 index 00000000..d91ea1e4 --- /dev/null +++ b/tests/worker_sandbox_shutdown.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +fixture_dir="$(mktemp -d)" +trap 'rm -rf "$fixture_dir"' EXIT + +# Exercise the production cleanup function without booting a MicroVM. +sed -n '/^cleanup() {/,/^trap cleanup SIGTERM SIGINT/p' "$repo_root/docker/supervisor.sh" \ + | sed '$d' > "$fixture_dir/cleanup.sh" +cat > "$fixture_dir/fixture.sh" <<'FIXTURE' +#!/usr/bin/env bash +set -u +CLEANUP_IN_PROGRESS=false +source "$1/cleanup.sh" +fixture_dir="$1" + +sleep 30 & +SANDBOX_PID=$! +( + trap ' + sleep 0.1 + if kill -0 "$SANDBOX_PID" 2>/dev/null \ + && [ "$(awk '\''{print $3}'\'' "/proc/$SANDBOX_PID/stat")" != Z ]; then + touch "$fixture_dir/drained" + fi + exit 0 + ' TERM + touch "$fixture_dir/ready" + while true; do sleep 0.01; done +) & +WORKER_PID=$! +while [ ! -f "$fixture_dir/ready" ]; do sleep 0.01; done +cleanup +FIXTURE + +timeout 5 bash "$fixture_dir/fixture.sh" "$fixture_dir" +test -f "$fixture_dir/drained" +echo 'Worker drained before sandbox shutdown.' + +for entry in docker-compose.yaml:service-worker docker-compose.scalable.yml:worker-sandbox docker-compose.local-dev.yml:service; do + docker compose -f "$repo_root/${entry%:*}" config --format json \ + | python3 -c ' +import json, sys +worker = json.load(sys.stdin)["services"][sys.argv[1]] +assert worker["stop_grace_period"] == "6m40s", worker["stop_grace_period"] +assert "CODEAPI_SHUTDOWN_TIMEOUT_MS" in worker["environment"] +' "${entry#*:}" +done +echo 'Compose worker stop grace permits the default execution and cleanup budget.'