From f53102dabfc99610f41a18cec06d640f1e2a8339 Mon Sep 17 00:00:00 2001 From: Abhishek Kumar Date: Sat, 27 Jun 2026 17:53:31 +0530 Subject: [PATCH] chore: drain active calls before rolling updates --- api/routes/main.py | 19 +++++++ api/services/pipecat/active_calls.py | 35 +++++++++++++ api/services/pipecat/run_pipeline.py | 11 ++++ api/tests/test_active_calls.py | 53 +++++++++++++++++++ scripts/rolling_update.sh | 76 ++++++++++++++++++++++++---- 5 files changed, 183 insertions(+), 11 deletions(-) create mode 100644 api/services/pipecat/active_calls.py create mode 100644 api/tests/test_active_calls.py diff --git a/api/routes/main.py b/api/routes/main.py index c7c589e0..d1e2abcf 100644 --- a/api/routes/main.py +++ b/api/routes/main.py @@ -125,3 +125,22 @@ async def health() -> HealthResponse: STACK_PUBLISHABLE_CLIENT_KEY if is_stack else None ), ) + + +class ActiveCallsResponse(BaseModel): + active_calls: int + + +@router.get("/health/active-calls", response_model=ActiveCallsResponse) +async def active_calls() -> ActiveCallsResponse: + """In-flight call count for THIS worker — the drain signal for deploys. + + A deploy orchestrator polls this per worker and waits for zero before + sending SIGTERM, because uvicorn force-closes live call WebSockets (close + code 1012) on SIGTERM and would cut calls mid-conversation otherwise. The + count is per-process: one uvicorn per VM port (scripts/rolling_update.sh) + or per Kubernetes pod (preStop hook). See api/services/pipecat/active_calls.py. + """ + from api.services.pipecat.active_calls import active_call_count + + return ActiveCallsResponse(active_calls=active_call_count()) diff --git a/api/services/pipecat/active_calls.py b/api/services/pipecat/active_calls.py new file mode 100644 index 00000000..c9cd3e7d --- /dev/null +++ b/api/services/pipecat/active_calls.py @@ -0,0 +1,35 @@ +"""In-process registry of active pipeline runs (live voice calls). + +Each uvicorn worker tracks the calls it is currently running so a deploy +orchestrator can *drain* the worker before stopping it: poll the count, wait for +zero, then send SIGTERM. Sending SIGTERM while calls are live makes uvicorn +force-close their WebSockets (close code 1012), which cuts the calls instead of +letting them finish — so the wait has to happen first. + +The registry is deliberately per-process. That is exactly the unit that gets +drained: one uvicorn process per VM port (see ``scripts/rolling_update.sh``) or +one uvicorn process per Kubernetes pod (drained via a ``preStop`` hook). The +count is exposed read-only at ``GET /api/v1/health/active-calls`` and is also a +natural autoscaling signal (concurrent calls per worker). + +Access is single-threaded (asyncio event loop), so no lock is needed. A set of +run ids — rather than a bare counter — keeps register/unregister idempotent and +makes the in-flight runs inspectable for debugging. +""" + +_active_run_ids: set[int] = set() + + +def register_active_call(workflow_run_id: int) -> None: + """Mark a pipeline run as active in this worker.""" + _active_run_ids.add(workflow_run_id) + + +def unregister_active_call(workflow_run_id: int) -> None: + """Mark a pipeline run as finished in this worker.""" + _active_run_ids.discard(workflow_run_id) + + +def active_call_count() -> int: + """Number of pipeline runs currently active in this worker.""" + return len(_active_run_ids) diff --git a/api/services/pipecat/run_pipeline.py b/api/services/pipecat/run_pipeline.py index 2786d348..32ec2fa0 100644 --- a/api/services/pipecat/run_pipeline.py +++ b/api/services/pipecat/run_pipeline.py @@ -11,6 +11,10 @@ from api.services.integrations import ( IntegrationRuntimeContext, create_runtime_sessions, ) +from api.services.pipecat.active_calls import ( + register_active_call, + unregister_active_call, +) from api.services.pipecat.audio_config import AudioConfig, create_audio_config from api.services.pipecat.event_handlers import ( register_audio_data_handler, @@ -870,6 +874,12 @@ async def _run_pipeline( register_audio_data_handler(audio_buffer, workflow_run_id, in_memory_audio_buffer) + # Mark this run active so a deploy can drain it before stopping this worker: + # the orchestrator polls GET /api/v1/health/active-calls and waits for zero + # before sending SIGTERM (uvicorn force-closes live call WebSockets on + # SIGTERM, so the wait has to come first). See api/services/pipecat/ + # active_calls.py and scripts/rolling_update.sh. + register_active_call(workflow_run_id) try: # Run the pipeline await run_pipeline_worker(task) @@ -877,6 +887,7 @@ async def _run_pipeline( except asyncio.CancelledError: logger.warning("Received CancelledError in _run_pipeline") finally: + unregister_active_call(workflow_run_id) # Close MCP sessions here, not in engine.cleanup(). The anyio cancel # scopes opened by MCPClient.start() in engine.initialize() are # task-affine; this finally runs in the same task as initialize(), diff --git a/api/tests/test_active_calls.py b/api/tests/test_active_calls.py new file mode 100644 index 00000000..c3d6cb3c --- /dev/null +++ b/api/tests/test_active_calls.py @@ -0,0 +1,53 @@ +"""Unit tests for the per-worker active-call registry (deploy draining). + +The registry backs GET /api/v1/health/active-calls, which scripts/rolling_update.sh +(and a k8s preStop hook) polls to wait for live calls to finish before stopping a +worker. The guarantees that matter for draining: register/unregister are +idempotent, and the count only reaches zero when every registered run is gone. +""" + +from api.services.pipecat import active_calls + + +def setup_function(): + # Module-level state — start each test from an empty registry. + active_calls._active_run_ids.clear() + + +def test_starts_empty(): + assert active_calls.active_call_count() == 0 + + +def test_register_counts_distinct_runs(): + active_calls.register_active_call(1) + active_calls.register_active_call(2) + assert active_calls.active_call_count() == 2 + + +def test_register_is_idempotent(): + # Registering the same run twice must not double-count, or the count could + # never drain to zero. + active_calls.register_active_call(1) + active_calls.register_active_call(1) + assert active_calls.active_call_count() == 1 + + +def test_unregister_removes_run(): + active_calls.register_active_call(1) + active_calls.register_active_call(2) + active_calls.unregister_active_call(1) + assert active_calls.active_call_count() == 1 + + +def test_unregister_unknown_run_is_a_noop(): + # discard() semantics: unregistering a run that was never registered (or was + # already removed) is safe and cannot push the count negative. + active_calls.unregister_active_call(999) + assert active_calls.active_call_count() == 0 + + +def test_full_lifecycle_drains_to_zero(): + active_calls.register_active_call(42) + assert active_calls.active_call_count() == 1 + active_calls.unregister_active_call(42) + assert active_calls.active_call_count() == 0 diff --git a/scripts/rolling_update.sh b/scripts/rolling_update.sh index 244933fd..be1dd6de 100755 --- a/scripts/rolling_update.sh +++ b/scripts/rolling_update.sh @@ -28,6 +28,7 @@ NGINX_UPSTREAM_TEMPLATE="$BASE_DIR/nginx/dograh_upstream.conf.template" NGINX_UPSTREAM_CONF="/etc/nginx/conf.d/dograh_upstream.conf" HEALTH_CHECK_ENDPOINT="/api/v1/health" +ACTIVE_CALLS_ENDPOINT="/api/v1/health/active-calls" # Load environment if [[ -f "$ENV_FILE" ]]; then @@ -40,7 +41,9 @@ FASTAPI_WORKERS=${FASTAPI_WORKERS:-$CPU_CORES} ARQ_WORKERS=${ARQ_WORKERS:-1} # Tuning knobs (override via environment) -DRAIN_TIMEOUT=${DRAIN_TIMEOUT:-300} # seconds to wait for old workers to drain +DRAIN_TIMEOUT=${DRAIN_TIMEOUT:-300} # seconds to wait for active calls to finish +DRAIN_INTERVAL=${DRAIN_INTERVAL:-5} # seconds between active-call drain polls +STOP_TIMEOUT=${STOP_TIMEOUT:-30} # seconds to wait for drained workers to exit after SIGTERM HEALTH_MAX_ATTEMPTS=${HEALTH_MAX_ATTEMPTS:-30} # per-worker health-check retries HEALTH_INTERVAL=${HEALTH_INTERVAL:-2} # seconds between health-check retries @@ -96,6 +99,20 @@ kill_process_tree() { fi } +# Active in-progress call count for a single worker, via its health endpoint. +# A worker that is unreachable (already exited) or returns no/garbage body +# reports 0, so it never blocks the drain. +count_active_calls_on_port() { + local port=$1 + local body n + body=$(curl -s --max-time 3 \ + "http://127.0.0.1:${port}${ACTIVE_CALLS_ENDPOINT}" 2>/dev/null || true) + n=$(printf '%s' "$body" \ + | grep -o '"active_calls"[[:space:]]*:[[:space:]]*[0-9]\+' \ + | grep -o '[0-9]\+$' || true) + printf '%s' "${n:-0}" +} + ############################################################################### ### ROLLBACK ############################################################################### @@ -366,9 +383,46 @@ log_info "nginx reloaded — traffic now routed to band $NEW_BAND" ### PHASE 5: DRAIN OLD WORKERS ############################################################################### -log_info "=== Phase 5: Draining old workers (band $OLD_BAND, timeout ${DRAIN_TIMEOUT}s) ===" +# nginx (Phase 4) already routes new calls to the new band, so the old band only +# holds calls still in progress. Wait for those to finish BEFORE signalling the +# workers: SIGTERM makes uvicorn force-close live call WebSockets (close code +# 1012), cutting calls mid-conversation. So we poll each old worker's in-flight +# call count and only stop once it reaches zero (or DRAIN_TIMEOUT elapses). -# Collect old worker PIDs +log_info "=== Phase 5a: Draining active calls from band $OLD_BAND (timeout ${DRAIN_TIMEOUT}s) ===" + +drain_start=$(date +%s) +while true; do + active=0 + for ((w = 0; w < FASTAPI_WORKERS; w++)); do + port=$((OLD_BASE + w)) + # Only poll workers still alive; an exited worker holds no calls. + pidfile="$RUN_DIR/uvicorn_${port}.pid" + if [[ -f "$pidfile" ]] && kill -0 "$(<"$pidfile")" 2>/dev/null; then + active=$((active + $(count_active_calls_on_port "$port"))) + fi + done + + if [[ $active -eq 0 ]]; then + log_info "Band $OLD_BAND fully drained — no active calls" + break + fi + + elapsed=$(( $(date +%s) - drain_start )) + if [[ $elapsed -ge $DRAIN_TIMEOUT ]]; then + log_warn "Drain timeout reached (${DRAIN_TIMEOUT}s) with $active active call(s) still running — stopping anyway." + break + fi + + log_info " Waiting for $active active call(s) to finish... (${elapsed}s / ${DRAIN_TIMEOUT}s)" + sleep "$DRAIN_INTERVAL" +done + +log_info "=== Phase 5b: Stopping old workers (band $OLD_BAND, timeout ${STOP_TIMEOUT}s) ===" + +# Calls are drained — now signal the workers and reap them. A drained worker +# exits within a second or two of SIGTERM; STOP_TIMEOUT bounds stragglers (e.g. +# a call that outlived DRAIN_TIMEOUT) before we force-kill. OLD_PIDS=() for ((w = 0; w < FASTAPI_WORKERS; w++)); do port=$((OLD_BASE + w)) @@ -385,7 +439,7 @@ for ((w = 0; w < FASTAPI_WORKERS; w++)); do done if [[ ${#OLD_PIDS[@]} -gt 0 ]]; then - start_time=$(date +%s) + stop_start=$(date +%s) while true; do all_dead=true @@ -397,13 +451,13 @@ if [[ ${#OLD_PIDS[@]} -gt 0 ]]; then done if $all_dead; then - log_info "All old workers exited gracefully" + log_info "All old workers exited" break fi - elapsed=$(( $(date +%s) - start_time )) - if [[ $elapsed -ge $DRAIN_TIMEOUT ]]; then - log_warn "Drain timeout reached (${DRAIN_TIMEOUT}s). Force-killing remaining old workers." + elapsed=$(( $(date +%s) - stop_start )) + if [[ $elapsed -ge $STOP_TIMEOUT ]]; then + log_warn "Stop timeout reached (${STOP_TIMEOUT}s). Force-killing remaining old workers." for pid in "${OLD_PIDS[@]}"; do if kill -0 "$pid" 2>/dev/null; then kill_process_tree "$pid" "-KILL" @@ -414,11 +468,11 @@ if [[ ${#OLD_PIDS[@]} -gt 0 ]]; then break fi - log_info " Waiting for old workers to drain... (${elapsed}s / ${DRAIN_TIMEOUT}s)" - sleep 5 + log_info " Waiting for old workers to exit... (${elapsed}s / ${STOP_TIMEOUT}s)" + sleep 2 done else - log_warn "No old worker PIDs to drain" + log_warn "No old worker PIDs to stop" fi ###############################################################################