feat(scaling): enable horizontal scaling for API and worker pods

- Add unauthenticated /api/diagnostic/healthz/live and /healthz/ready
  probe endpoints for Kubernetes liveness/readiness checks
- Separate Celery Beat into dedicated beat service in docker-compose.yaml
- Remove container_name from api and worker services to allow scaling
- Create Helm beat-deployment.yaml for standalone Beat scheduler pod
- Remove -B flag from worker-deployment.yaml so workers can scale safely
- Add beat section and fix probe paths in Helm values.yaml
- Add tests for the new probe endpoints

Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
copilot-swe-agent[bot]
2026-03-16 23:52:56 +00:00
parent b3dce16838
commit f75b125992
6 changed files with 282 additions and 9 deletions
+61
View File
@@ -21,6 +21,67 @@ _DEFAULT_REDIS_URL = "redis://localhost:6379/0"
router = APIRouter()
# ---------------------------------------------------------------------------
# Unauthenticated probe endpoints for Kubernetes liveness / readiness checks.
# These intentionally skip authentication so that kubelet can reach them
# without credentials. They live under /diagnostic/healthz/* so that the
# existing authenticated /diagnostic/health endpoint is unaffected.
# ---------------------------------------------------------------------------
@router.get("/diagnostic/healthz/live")
async def liveness_probe() -> JSONResponse:
"""Lightweight liveness probe for Kubernetes.
Returns **200 OK** as long as the process is running. Kubernetes uses
this to decide whether to *restart* the container — it should therefore
be as cheap as possible and **never** check external dependencies.
**Authentication:** None (designed for kubelet probes).
"""
return JSONResponse(content={"status": "ok"}, status_code=200)
@router.get("/diagnostic/healthz/ready")
async def readiness_probe() -> JSONResponse:
"""Readiness probe for Kubernetes.
Verifies that the application can serve traffic by checking the database
and Redis. Kubernetes uses this to decide whether to *route traffic* to
the pod.
Returns **200 OK** when all critical subsystems are reachable, or
**503 Service Unavailable** when the database is down.
**Authentication:** None (designed for kubelet probes).
"""
checks: dict[str, dict[str, str]] = {}
db_ok = False
# ── Database check ─────────────────────────────────────────────────
try:
with engine.connect() as conn:
conn.execute(text("SELECT 1"))
checks["database"] = {"status": "ok"}
db_ok = True
except Exception as exc:
logger.warning("Readiness probe: database check failed: %s", exc)
checks["database"] = {"status": "error", "detail": str(exc)}
# ── Redis check ────────────────────────────────────────────────────
try:
redis_url = settings.redis_url or _DEFAULT_REDIS_URL
r = redis_lib.from_url(redis_url, socket_connect_timeout=2, socket_timeout=2)
r.ping()
checks["redis"] = {"status": "ok"}
except Exception as exc:
logger.warning("Readiness probe: Redis check failed: %s", exc)
checks["redis"] = {"status": "error", "detail": str(exc)}
http_status = 503 if not db_ok else 200
overall = "ready" if db_ok else "not_ready"
return JSONResponse(content={"status": overall, "checks": checks}, status_code=http_status)
@router.get("/diagnostic/health")
@require_login
+25 -4
View File
@@ -3,7 +3,7 @@ services:
build:
context: .
dockerfile: Dockerfile
container_name: document_api
# No container_name — allows `docker compose up --scale api=N`
restart: always
# We'll keep the code in /app, but set working_dir to the shared data directory
@@ -24,7 +24,7 @@ services:
depends_on:
- redis
- worker
- beat
# Mount the shared working directory for data
volumes:
@@ -34,13 +34,14 @@ services:
build:
context: .
dockerfile: Dockerfile
container_name: document_worker
# No container_name — allows `docker compose up --scale worker=N`
restart: always
# same shared working directory
working_dir: /workdir
command: ["celery", "-A", "app.celery_worker", "worker", "-B", "--loglevel=info", "-Q", "document_processor,default,celery"]
# Workers process tasks only — no -B flag (Beat runs in the dedicated beat service)
command: ["celery", "-A", "app.celery_worker", "worker", "--loglevel=info", "-Q", "document_processor,default,celery"]
env_file:
- .env
environment:
@@ -54,6 +55,26 @@ services:
volumes:
- /var/docparse/workdir:/workdir
# Dedicated Celery Beat scheduler — exactly one instance must run at all times.
# Beat publishes periodic tasks to the Redis broker; workers pick them up.
# Do NOT scale this service (replicas must stay at 1).
beat:
build:
context: .
dockerfile: Dockerfile
container_name: document_beat
restart: always
working_dir: /workdir
command: ["celery", "-A", "app.celery_worker", "beat", "--loglevel=info"]
env_file:
- .env
environment:
- PYTHONPATH=/app
depends_on:
- redis
volumes:
- /var/docparse/workdir:/workdir
gotenberg:
image: gotenberg/gotenberg:latest
container_name: gotenberg
@@ -0,0 +1,83 @@
{{- /*
Celery Beat scheduler — publishes periodic tasks to the broker.
Exactly ONE replica must run; never scale this deployment.
*/ -}}
apiVersion: apps/v1
kind: Deployment
metadata:
name: {{ include "docuelevate.fullname" . }}-beat
namespace: {{ .Release.Namespace }}
labels:
{{- include "docuelevate.labels" . | nindent 4 }}
app.kubernetes.io/component: beat
spec:
replicas: 1
strategy:
type: Recreate # Prevent two Beat instances from running simultaneously
selector:
matchLabels:
{{- include "docuelevate.selectorLabels" . | nindent 6 }}
app.kubernetes.io/component: beat
template:
metadata:
labels:
{{- include "docuelevate.selectorLabels" . | nindent 8 }}
app.kubernetes.io/component: beat
{{- with .Values.beat.podAnnotations }}
annotations:
{{- toYaml . | nindent 8 }}
{{- end }}
spec:
serviceAccountName: {{ include "docuelevate.serviceAccountName" . }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.beat.podSecurityContext }}
securityContext:
{{- toYaml . | nindent 8 }}
{{- end }}
containers:
- name: beat
image: {{ include "docuelevate.image" . }}
imagePullPolicy: {{ .Values.image.pullPolicy }}
command:
- celery
- -A
- app.celery_worker
- beat
- --loglevel=info
envFrom:
- configMapRef:
name: {{ include "docuelevate.fullname" . }}-config
- secretRef:
name: {{ include "docuelevate.fullname" . }}-secret
{{- with .Values.beat.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
resources:
{{- toYaml .Values.beat.resources | nindent 12 }}
volumeMounts:
- name: workdir
mountPath: /workdir
volumes:
- name: workdir
{{- if .Values.workdir.persistence.enabled }}
persistentVolumeClaim:
claimName: {{ .Values.workdir.persistence.existingClaim | default (printf "%s-workdir" (include "docuelevate.fullname" .)) }}
{{- else }}
emptyDir: {}
{{- end }}
{{- with .Values.beat.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.beat.affinity }}
affinity:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.beat.tolerations }}
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
@@ -42,7 +42,6 @@ spec:
- -A
- app.celery_worker
- worker
- -B
- --loglevel=info
- -Q
- document_processor,default,celery
+33 -4
View File
@@ -121,10 +121,10 @@ api:
type: ClusterIP
port: 8000
# Liveness / readiness probes
# Liveness / readiness probes (unauthenticated endpoints for kubelet)
livenessProbe:
httpGet:
path: /api/health
path: /api/diagnostic/healthz/live
port: 8000
initialDelaySeconds: 30
periodSeconds: 20
@@ -132,7 +132,7 @@ api:
readinessProbe:
httpGet:
path: /api/health
path: /api/diagnostic/healthz/ready
port: 8000
initialDelaySeconds: 15
periodSeconds: 10
@@ -191,7 +191,36 @@ worker:
drop: ["ALL"]
# ---------------------------------------------------------------------------
# Shared workdir volume (api + worker mount the same PVC)
# Celery Beat scheduler (singleton — always exactly 1 replica)
# Beat publishes periodic tasks; workers consume them from the broker.
# ---------------------------------------------------------------------------
beat:
resources:
requests:
cpu: 100m
memory: 256Mi
limits:
cpu: 500m
memory: 512Mi
podAnnotations: {}
nodeSelector: {}
tolerations: []
affinity: {}
podSecurityContext:
runAsNonRoot: true
runAsUser: 1000
fsGroup: 1000
securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: false
capabilities:
drop: ["ALL"]
# ---------------------------------------------------------------------------
# Shared workdir volume (api + worker + beat mount the same PVC)
# ---------------------------------------------------------------------------
workdir:
persistence:
+80
View File
@@ -5,6 +5,86 @@ from unittest.mock import MagicMock, patch
import pytest
@pytest.mark.unit
class TestLivenessProbe:
"""Tests for GET /api/diagnostic/healthz/live (unauthenticated)."""
def test_liveness_returns_200(self, client):
"""Liveness probe always returns 200 OK."""
response = client.get("/api/diagnostic/healthz/live")
assert response.status_code == 200
data = response.json()
assert data["status"] == "ok"
@pytest.mark.unit
class TestReadinessProbe:
"""Tests for GET /api/diagnostic/healthz/ready (unauthenticated)."""
def test_readiness_returns_200_when_all_ok(self, client):
"""Readiness probe returns 200 when database and Redis are reachable."""
with (
patch("app.api.diagnostic.engine") as mock_engine,
patch("app.api.diagnostic.redis_lib") as mock_redis,
):
mock_conn = MagicMock()
mock_engine.connect.return_value.__enter__ = MagicMock(return_value=mock_conn)
mock_engine.connect.return_value.__exit__ = MagicMock(return_value=False)
mock_redis_inst = MagicMock()
mock_redis.from_url.return_value = mock_redis_inst
response = client.get("/api/diagnostic/healthz/ready")
assert response.status_code == 200
data = response.json()
assert data["status"] == "ready"
assert data["checks"]["database"]["status"] == "ok"
def test_readiness_returns_503_when_database_fails(self, client):
"""Readiness probe returns 503 when database is unreachable."""
with (
patch("app.api.diagnostic.engine") as mock_engine,
patch("app.api.diagnostic.redis_lib") as mock_redis,
):
mock_engine.connect.side_effect = Exception("DB unavailable")
mock_redis_inst = MagicMock()
mock_redis.from_url.return_value = mock_redis_inst
response = client.get("/api/diagnostic/healthz/ready")
assert response.status_code == 503
data = response.json()
assert data["status"] == "not_ready"
assert data["checks"]["database"]["status"] == "error"
def test_readiness_returns_200_when_redis_fails(self, client):
"""Readiness remains 200 when only Redis is down (non-critical)."""
with (
patch("app.api.diagnostic.engine") as mock_engine,
patch("app.api.diagnostic.redis_lib") as mock_redis,
):
mock_conn = MagicMock()
mock_engine.connect.return_value.__enter__ = MagicMock(return_value=mock_conn)
mock_engine.connect.return_value.__exit__ = MagicMock(return_value=False)
mock_redis.from_url.return_value = MagicMock()
mock_redis.from_url.return_value.ping.side_effect = Exception("Connection refused")
response = client.get("/api/diagnostic/healthz/ready")
assert response.status_code == 200
data = response.json()
assert data["status"] == "ready"
assert data["checks"]["redis"]["status"] == "error"
def test_readiness_contains_checks_keys(self, client):
"""Readiness response always contains database and redis checks."""
response = client.get("/api/diagnostic/healthz/ready")
data = response.json()
assert "checks" in data
assert "database" in data["checks"]
assert "redis" in data["checks"]
@pytest.mark.unit
class TestHealthEndpoint:
"""Tests for GET /api/diagnostic/health endpoint."""