chore: drain active calls before rolling updates

This commit is contained in:
Abhishek Kumar 2026-06-27 17:53:31 +05:30
parent 557de72b9c
commit f53102dabf
5 changed files with 183 additions and 11 deletions

View file

@ -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)

View file

@ -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(),