Added readiness and liveness probes (#8488)

azhavoro · web-flow · commit 5104fd9e8614 · 2024-10-09T18:55:46.000+03:00
diff --git a/changelog.d/20241009_164505_andrey_k8s_probes.md b/changelog.d/20241009_164505_andrey_k8s_probes.md
@@ -0,0 +1,4 @@
+### Added
+
+- \[Helm\] Readiness and liveness probes
+  (<https://github.com/cvat-ai/cvat/pull/8488>)
diff --git a/cvat/apps/health/management/__init__.py b/cvat/apps/health/management/__init__.py
diff --git a/cvat/apps/health/management/commands/__init__.py b/cvat/apps/health/management/commands/__init__.py
diff --git a/cvat/apps/health/management/commands/workerprobe.py b/cvat/apps/health/management/commands/workerprobe.py
@@ -0,0 +1,32 @@
+import os
+import platform
+from datetime import datetime, timedelta
+from django.core.management.base import BaseCommand, CommandError
+from django.conf import settings
+from rq.worker import Worker
+import django_rq
+
+
+class Command(BaseCommand):
+    help = "Check worker liveness in specified queues"
+
+    def add_arguments(self, parser):
+        parser.add_argument("queue_names", nargs="+", type=str)
+
+    def handle(self, *args, **options):
+        hostname = platform.node()
+        for queue_name in options["queue_names"]:
+            if queue_name not in settings.RQ_QUEUES:
+                raise CommandError(f"Queue {queue_name} is not defined")
+
+            connection = django_rq.get_connection(queue_name)
+            workers = [w for w in Worker.all(connection) if queue_name in w.queue_names() and w.hostname == hostname]
+
+            expected_workers = int(os.getenv("NUMPROCS", 1))
+
+            if len(workers) != expected_workers:
+                raise CommandError("Number of registered workers does not match the expected number, " \
+                                    f"actual: {len(workers)}, expected: {expected_workers}")
+            for worker in workers:
+                if datetime.now() - worker.last_heartbeat > timedelta(seconds=worker.worker_ttl):
+                    raise CommandError(f"It seems that worker {worker.name}, pid: {worker.pid} is dead")
diff --git a/helm-chart/Chart.yaml b/helm-chart/Chart.yaml
@@ -16,7 +16,7 @@ type: application
 # This is the chart version. This version number should be incremented each time you make changes
 # to the chart and its templates, including the app version.
 # Versions are expected to follow Semantic Versioning (https://semver.org/)
-version: 0.13.2
+version: 0.14.0
 
 # This is the version number of the application being deployed. This version number should be
 # incremented each time you make changes to the application. Versions are not expected to
diff --git a/helm-chart/templates/_helpers.tpl b/helm-chart/templates/_helpers.tpl
@@ -169,3 +169,18 @@ The name of the service account to use for backend pods
       key: CLICKHOUSE_PASSWORD
 {{- end }}
 {{- end }}
+
+{{- define "cvat.backend.worker.livenessProbe" -}}
+{{- if .livenessProbe.enabled }}
+livenessProbe:
+  exec:
+    command:
+    - python
+    - manage.py
+    - workerprobe
+    {{- range .args }}
+    - {{ . }}
+    {{- end }}
+{{ toYaml (omit .livenessProbe "enabled") | indent 2}}
+{{- end }}
+{{- end }}
diff --git a/helm-chart/templates/cvat_backend/server/deployment.yml b/helm-chart/templates/cvat_backend/server/deployment.yml
@@ -65,6 +65,20 @@ spec:
           {{- end }}
           ports:
           - containerPort: 8080
+          {{- if $localValues.readinessProbe.enabled }}
+          readinessProbe:
+            httpGet:
+              path: /api/server/about
+              port: 8080
+            {{- toYaml (omit $localValues.readinessProbe "enabled") | nindent 12 }}
+          {{- end }}
+          {{- if $localValues.livenessProbe.enabled }}
+          livenessProbe:
+            httpGet:
+              path: /api/server/about
+              port: 8080
+            {{- toYaml (omit $localValues.livenessProbe "enabled") | nindent 12 }}
+          {{- end }}
           volumeMounts:
           {{- if not .Values.cvat.backend.disableDistinctCachePerService }}
           - mountPath: /home/django/data/cache
diff --git a/helm-chart/templates/cvat_backend/utils/deployment.yml b/helm-chart/templates/cvat_backend/utils/deployment.yml
@@ -60,8 +60,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
-          ports:
-          - containerPort: 8080
+          {{- $probeArgs := list "notifications" "cleaning" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           volumeMounts:
           {{- if not .Values.cvat.backend.disableDistinctCachePerService }}
           - mountPath: /home/django/data/cache
diff --git a/helm-chart/templates/cvat_backend/worker_analyticsreports/deployment.yml b/helm-chart/templates/cvat_backend/worker_analyticsreports/deployment.yml
@@ -61,6 +61,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "analytics_reports" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           {{- with concat .Values.cvat.backend.additionalVolumeMounts $localValues.additionalVolumeMounts }}
           volumeMounts:
           {{- toYaml . | nindent 10 }}
diff --git a/helm-chart/templates/cvat_backend/worker_annotation/deployment.yml b/helm-chart/templates/cvat_backend/worker_annotation/deployment.yml
@@ -60,6 +60,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "annotation" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           volumeMounts:
           {{- if not .Values.cvat.backend.disableDistinctCachePerService }}
           - mountPath: /home/django/data/cache
diff --git a/helm-chart/templates/cvat_backend/worker_export/deployment.yml b/helm-chart/templates/cvat_backend/worker_export/deployment.yml
@@ -61,6 +61,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "export" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           volumeMounts:
           {{- if not .Values.cvat.backend.disableDistinctCachePerService }}
           - mountPath: /home/django/data/cache
diff --git a/helm-chart/templates/cvat_backend/worker_import/deployment.yml b/helm-chart/templates/cvat_backend/worker_import/deployment.yml
@@ -60,6 +60,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "import" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           volumeMounts:
           {{- if not .Values.cvat.backend.disableDistinctCachePerService }}
           - mountPath: /home/django/data/cache
diff --git a/helm-chart/templates/cvat_backend/worker_qualityreports/deployment.yml b/helm-chart/templates/cvat_backend/worker_qualityreports/deployment.yml
@@ -60,6 +60,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "quality_reports" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           {{- with concat .Values.cvat.backend.additionalVolumeMounts $localValues.additionalVolumeMounts }}
           volumeMounts:
           {{- toYaml . | nindent 10 }}
diff --git a/helm-chart/templates/cvat_backend/worker_webhooks/deployment.yml b/helm-chart/templates/cvat_backend/worker_webhooks/deployment.yml
@@ -60,6 +60,9 @@ spec:
           {{- with concat .Values.cvat.backend.additionalEnv $localValues.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- $probeArgs := list "webhooks" -}}
+          {{- $probeConfig := dict "args" $probeArgs "livenessProbe" $.Values.cvat.backend.worker.livenessProbe -}}
+          {{ include "cvat.backend.worker.livenessProbe" $probeConfig | indent 10 }}
           {{- with concat .Values.cvat.backend.additionalVolumeMounts $localValues.additionalVolumeMounts }}
           volumeMounts:
           {{- toYaml . | nindent 10 }}
diff --git a/helm-chart/templates/cvat_frontend/deployment.yml b/helm-chart/templates/cvat_frontend/deployment.yml
@@ -45,6 +45,18 @@ spec:
           {{- toYaml . | nindent 10 }}
           {{- end }}
           {{- with .Values.cvat.frontend.additionalVolumeMounts }}
+          {{- if .Values.cvat.frontend.readinessProbe.enabled }}
+          readinessProbe:
+            tcpSocket:
+              port: 80
+            {{- toYaml (omit .Values.cvat.frontend.readinessProbe "enabled") | nindent 12 }}
+          {{- end }}
+          {{- if .Values.cvat.frontend.livenessProbe.enabled }}
+          livenessProbe:
+            tcpSocket:
+              port: 80
+            {{- toYaml (omit .Values.cvat.frontend.livenessProbe "enabled") | nindent 12 }}
+          {{- end }}
           volumeMounts:
           {{- toYaml . | nindent 10 }}
           {{- end }}
diff --git a/helm-chart/templates/cvat_kvrocks/statefulset.yml b/helm-chart/templates/cvat_kvrocks/statefulset.yml
@@ -59,6 +59,25 @@ spec:
           {{- with .Values.cvat.kvrocks.additionalEnv }}
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          #https://github.com/apache/kvrocks/blob/unstable/Dockerfile
+          {{- if .Values.cvat.kvrocks.readinessProbe.enabled }}
+          readinessProbe:
+            exec:
+              command:
+              - /bin/sh
+              - -c
+              - ./bin/redis-cli -p 6666 PING | grep -E '(PONG|NOAUTH)'
+            {{- toYaml (omit .Values.cvat.kvrocks.readinessProbe "enabled") | nindent 12 }}
+          {{- end }}
+          {{- if .Values.cvat.kvrocks.livenessProbe.enabled }}
+          livenessProbe:
+            exec:
+              command:
+              - /bin/sh
+              - -c
+              - ./bin/redis-cli -p 6666 PING | grep -E '(PONG|NOAUTH)'
+            {{- toYaml (omit .Values.cvat.kvrocks.livenessProbe "enabled") | nindent 12 }}
+          {{- end }}
           volumeMounts:
           - name: {{ .Release.Name }}-kvrocks-data
             mountPath: /var/lib/kvrocks/data
diff --git a/helm-chart/templates/cvat_opa/deployment.yml b/helm-chart/templates/cvat_opa/deployment.yml
@@ -53,6 +53,20 @@ spec:
           env:
           {{- toYaml . | nindent 10 }}
           {{- end }}
+          {{- if .Values.cvat.opa.readinessProbe.enabled }}
+          readinessProbe:
+            httpGet:
+              port: 8181
+              path: "/health?bundles"
+            {{- toYaml (omit .Values.cvat.opa.readinessProbe "enabled") | nindent 12 }}
+          {{- end }}
+          {{- if .Values.cvat.opa.livenessProbe.enabled }}
+          livenessProbe:
+            httpGet:
+              port: 8181
+              path: "/health?bundles"
+            {{- toYaml (omit .Values.cvat.opa.livenessProbe "enabled") | nindent 12 }}
+          {{- end }}
           {{- with .Values.cvat.opa.additionalVolumeMounts }}
           volumeMounts:
           {{- toYaml . | nindent 10 }}
diff --git a/helm-chart/values.yaml b/helm-chart/values.yaml
@@ -42,7 +42,21 @@ cvat:
       additionalEnv: []
       additionalVolumes: []
       additionalVolumeMounts: []
+      readinessProbe:
+        enabled: true
+        periodSeconds: 15
+        initialDelaySeconds: 15
+      livenessProbe:
+        enabled: true
+        periodSeconds: 15
+        failureThreshold: 5
+        initialDelaySeconds: 60
     worker:
+      livenessProbe:
+        enabled: true
+        periodSeconds: 120
+        initialDelaySeconds: 30
+        timeoutSeconds: 10
       export:
         replicas: 2
         labels: {}
@@ -172,6 +186,14 @@ cvat:
     # -   mountPath: /tmp
     #     name: tmp
     #     subPath: test
+    readinessProbe:
+      enabled: true
+      periodSeconds: 10
+      failureThreshold: 5
+    livenessProbe:
+      enabled: true
+      periodSeconds: 10
+      failureThreshold: 5
     service:
       type: ClusterIP
       ports:
@@ -216,6 +238,14 @@ cvat:
     #     name: tmp
     #     subPath: test
     composeCompatibleServiceName: true # Sets service name to opa in order to be compatible with Docker Compose. Necessary because changing IAM_OPA_DATA_URL via environment variables in current images. Hinders multiple deployment due to duplicate name
+    readinessProbe:
+      enabled: true
+      periodSeconds: 15
+      initialDelaySeconds: 15
+    livenessProbe:
+      enabled: true
+      periodSeconds: 15
+      initialDelaySeconds: 15
     service:
       type: ClusterIP
       ports:
@@ -266,6 +296,14 @@ cvat:
     # -   mountPath: /tmp
     #     name: tmp
     #     subPath: test
+    readinessProbe:
+      enabled: true
+      periodSeconds: 10
+      initialDelaySeconds: 30
+    livenessProbe:
+      enabled: true
+      periodSeconds: 10
+      initialDelaySeconds: 30
     defaultStorage:
       enabled: true
 #     storageClassName: default
diff --git a/supervisord/utils.conf b/supervisord/utils.conf
@@ -24,13 +24,14 @@ command=%(ENV_HOME)s/wait_for_deps.sh
         -i 30 --path %(ENV_HOME)s
 environment=VECTOR_EVENT_HANDLER="SynchronousLogstashHandler"
 numprocs=1
+autorestart=true
 
 [program:rqworker-notifications]
 command=%(ENV_HOME)s/wait_for_deps.sh
     python3 %(ENV_HOME)s/manage.py rqworker -v 3 notifications
         --worker-class cvat.rqworker.DefaultWorker
 environment=VECTOR_EVENT_HANDLER="SynchronousLogstashHandler",CVAT_POSTGRES_APPLICATION_NAME="cvat:worker:notifications"
-numprocs=1
+numprocs=%(ENV_NUMPROCS)s
 autorestart=true
 
 [program:rqworker-cleaning]
diff --git a/tests/python/shared/fixtures/init.py b/tests/python/shared/fixtures/init.py
@@ -239,12 +239,32 @@ def kube_restore_clickhouse_db():
     )
 
 
+def _get_redis_inmem_keys_to_keep():
+    return ("rq:worker:", "rq:workers", "rq:scheduler_instance:", "rq:queues:")
+
+
 def docker_restore_redis_inmem():
-    docker_exec_redis_inmem(["redis-cli", "-e", "flushall"])
+    docker_exec_redis_inmem(
+        [
+            "sh",
+            "-c",
+            'redis-cli -e --scan --pattern "*" |'
+            'grep -v "' + r"\|".join(_get_redis_inmem_keys_to_keep()) + '" |'
+            "xargs -r redis-cli -e del",
+        ]
+    )
 
 
 def kube_restore_redis_inmem():
-    kube_exec_redis_inmem(["sh", "-c", 'redis-cli -e -a "${REDIS_PASSWORD}" flushall'])
+    kube_exec_redis_inmem(
+        [
+            "sh",
+            "-c",
+            'redis-cli -e -a "${REDIS_PASSWORD}" --scan --pattern "*" |'
+            'grep -v "' + r"\|".join(_get_redis_inmem_keys_to_keep()) + '" |'
+            'xargs -r redis-cli -e -a "${REDIS_PASSWORD}" del',
+        ]
+    )
 
 
 def docker_restore_redis_ondisk():
