From f75b12599291050f97aa452d980f17547ddd7bd6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 16 Mar 2026 23:52:56 +0000 Subject: [PATCH] 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> --- app/api/diagnostic.py | 61 ++++++++++++++ docker-compose.yaml | 29 ++++++- .../templates/beat-deployment.yaml | 83 +++++++++++++++++++ .../templates/worker-deployment.yaml | 1 - helm/docuelevate/values.yaml | 37 ++++++++- tests/test_diagnostic.py | 80 ++++++++++++++++++ 6 files changed, 282 insertions(+), 9 deletions(-) create mode 100644 helm/docuelevate/templates/beat-deployment.yaml diff --git a/app/api/diagnostic.py b/app/api/diagnostic.py index a329da46..d21a3d39 100644 --- a/app/api/diagnostic.py +++ b/app/api/diagnostic.py @@ -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 diff --git a/docker-compose.yaml b/docker-compose.yaml index 05e1584c..10c370f6 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -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 diff --git a/helm/docuelevate/templates/beat-deployment.yaml b/helm/docuelevate/templates/beat-deployment.yaml new file mode 100644 index 00000000..3ed52251 --- /dev/null +++ b/helm/docuelevate/templates/beat-deployment.yaml @@ -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 }} diff --git a/helm/docuelevate/templates/worker-deployment.yaml b/helm/docuelevate/templates/worker-deployment.yaml index c698b0f0..efb09f1c 100644 --- a/helm/docuelevate/templates/worker-deployment.yaml +++ b/helm/docuelevate/templates/worker-deployment.yaml @@ -42,7 +42,6 @@ spec: - -A - app.celery_worker - worker - - -B - --loglevel=info - -Q - document_processor,default,celery diff --git a/helm/docuelevate/values.yaml b/helm/docuelevate/values.yaml index 7ed056b0..de112df5 100644 --- a/helm/docuelevate/values.yaml +++ b/helm/docuelevate/values.yaml @@ -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: diff --git a/tests/test_diagnostic.py b/tests/test_diagnostic.py index fa47acc1..8d3f917a 100644 --- a/tests/test_diagnostic.py +++ b/tests/test_diagnostic.py @@ -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."""