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:
@@ -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
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user