Augments LabsAugments ADK

Horizontal Scaling

A single replica of a served agent works correctly out of the box. Adding replicas requires a shared backend for any state that must survive across replicas or process restarts.

Single-replica default

The generated Kubernetes Deployment sets replicas: 1. A single replica:

  • Runs the SQLite A2A task store entirely in-process (pass --task-db to augments serve to make it durable across restarts).
  • Uses any SessionStore the developer wires into the app (in-memory or SQLite-backed via the session extra).
  • Satisfies the container contract: binds 0.0.0.0:$PORT, reads config from env, handles SIGTERM with a 45-second drain window, exposes /healthz and /readyz.

For many workloads β€” especially those where agent turns are short β€” a single replica with a durable SQLite store is sufficient.

Scaling to multiple replicas

When traffic requires more than one replica, two shared backends must be configured before increasing the replica count. Both are provided by the ADK and wired via augments serve flags.

A2A task store β€” PostgresTaskStore

PostgresTaskStore (augments.adk.a2a.postgres_task_store) is a durable, shared A2A executor task store backed by a PostgreSQL a2a_tasks table reached through a psycopg async connection pool. Because every replica talks to the same database, background task state and A2AContinuationToken resumption survive across replicas β€” which a per-pod SQLite file cannot provide.

Wire it via --task-dsn:

augments serve \
  --agent my_pkg.agents:assistant \
  --card card.json \
  --task-dsn "postgresql://user:pass@db-host/agents"

Install the required extras:

pip install 'augments-adk[a2a,a2a-postgres]'

PostgresTaskStore also runs recover_on_startup: any task a prior process left non-terminal is marked FAILED before the server accepts requests, so clients receive a clean error and can resubmit.

[!NOTE] --task-dsn wires the ADK's own executor task store (background task snapshots for restart recovery and shared state across replicas). The a2a-sdk request-handler's own task store is separate; for full A2A wire-protocol task lookups across replicas an operator may additionally supply a shared a2a-sdk store via build_starlette_app. REST sessions via --session-dsn are fully shared with no further configuration.

--task-db (SQLite file) and --task-dsn (Postgres DSN) are mutually exclusive; choose one.

Session store β€” PostgresMultiSessions

PostgresMultiSessions (augments.adk.session.postgres_multi_sessions) is the Postgres counterpart of SQLiteMultiSessions. It owns a psycopg async connection pool and makes REST session history fully shared across replicas β€” any replica can continue any conversation from the last persisted message.

Wire it via --session-dsn:

augments serve \
  --agent my_pkg.agents:assistant \
  --session-dsn "postgresql://user:pass@db-host/agents"

Install the required extra:

pip install 'augments-adk[session-postgres]'

--session-db (SQLite file) and --session-dsn (Postgres DSN) are mutually exclusive; choose one.

Combined multi-replica invocation

For a horizontally-scaled deployment, point both DSNs at the same Postgres instance and set replicas > 1:

augments serve \
  --agent my_pkg.agents:assistant \
  --card card.json \
  --task-dsn "$PG_DSN" \
  --session-dsn "$PG_DSN"
pip install 'augments-adk[a2a,a2a-postgres,session-postgres]'

Enabling the HPA

The generated Kubernetes Deployment includes a HorizontalPodAutoscaler manifest set to minReplicas: 1 and maxReplicas: 4. It is disabled (no metrics target is configured) until the operator enables it.

To activate autoscaling once a shared backend is in place:

  1. Configure --task-dsn and --session-dsn in the Deployment's container args as described above.
  2. Set a CPU or custom metric target in deploy/k8s/HPA.yaml.
  3. Apply the updated manifests:
kubectl apply -k deploy/k8s

A typical CPU-based target:

# deploy/k8s/HPA.yaml
spec:
  minReplicas: 1
  maxReplicas: 4
  metrics:
    - type: Resource
      resource:
        name: cpu
        target:
          type: Utilization
          averageUtilization: 60

Cloud Run and App Runner

Both Cloud Run and App Runner scale horizontally by default, managing replicas automatically. The same shared-backend requirement applies: any stateful resource (A2A task store, session store) must use a shared backend when more than one instance can be active simultaneously. Use --task-dsn and --session-dsn as described above.

For Cloud Run, set --min-instances to control cold-start behaviour:

augments deploy cloud-run \
  --agent my_pkg.agents:assistant \
  --project my-project \
  --region us-central1 \
  --min-instances 1

Drain window sizing

In-flight agent turns can outlast the default SIGTERM β†’ kill window of 30 seconds (especially for multi-turn conversations, tool calls, and LLM providers with high latency). The generated Kubernetes manifests set terminationGracePeriodSeconds: 45. Increase this value for agents with longer expected turn durations:

spec:
  template:
    spec:
      terminationGracePeriodSeconds: 120  # adjust for your agent's turn budget

Cloud Run allows configuring the request timeout and drain timeout via the service configuration.

See also