From f9771427e2f7c363ae775006d61e7cb3f506faf4 Mon Sep 17 00:00:00 2001
From: Mathieu Virbel <mat@meltingrocks.com>
Date: Thu, 30 Nov 2023 19:43:19 +0100
Subject: [PATCH] server: add healthcheck for worker

---
 README.md                              | 14 +++++++++++---
 server/reflector/settings.py           |  3 +++
 server/reflector/worker/app.py         | 15 +++++++++++++++
 server/reflector/worker/healthcheck.py | 18 ++++++++++++++++++
 server/runserver.sh                    |  2 ++
 5 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 server/reflector/worker/healthcheck.py

diff --git a/README.md b/README.md
index cb75c76b..8830c79e 100644
--- a/README.md
+++ b/README.md
@@ -133,17 +133,25 @@ TRANSLATE_URL=https://monadical-sas--reflector-translator-web.modal.run
 ZEPHYR_LLM_URL=https://monadical-sas--reflector-llm-zephyr-web.modal.run
 ```
 
-### Start the project
+### Start the API/Backend
 
-Use:
+Start the API server:
 
 ```bash
 poetry run python3 -m reflector.app
 ```
 
-And start the background worker
+Start the background worker:
 
+```bash
 celery -A reflector.worker.app worker --loglevel=info
+```
+
+For crontab (only healthcheck for now), start the celery beat:
+
+```bash
+celery -A reflector.worker.app beat
+```
 
 #### Using docker
 
diff --git a/server/reflector/settings.py b/server/reflector/settings.py
index 2c68c4e5..d0ddc91a 100644
--- a/server/reflector/settings.py
+++ b/server/reflector/settings.py
@@ -128,5 +128,8 @@ class Settings(BaseSettings):
     # Profiling
     PROFILING: bool = False
 
+    # Healthcheck
+    HEALTHCHECK_URL: str | None = None
+
 
 settings = Settings()
diff --git a/server/reflector/worker/app.py b/server/reflector/worker/app.py
index e1000364..689623ce 100644
--- a/server/reflector/worker/app.py
+++ b/server/reflector/worker/app.py
@@ -1,6 +1,8 @@
+import structlog
 from celery import Celery
 from reflector.settings import settings
 
+logger = structlog.get_logger(__name__)
 app = Celery(__name__)
 app.conf.broker_url = settings.CELERY_BROKER_URL
 app.conf.result_backend = settings.CELERY_RESULT_BACKEND
@@ -8,5 +10,18 @@ app.conf.broker_connection_retry_on_startup = True
 app.autodiscover_tasks(
     [
         "reflector.pipelines.main_live_pipeline",
+        "reflector.worker.healthcheck",
     ]
 )
+
+# crontab
+app.conf.beat_schedule = {}
+
+if settings.HEALTHCHECK_URL:
+    app.conf.beat_schedule["healthcheck_ping"] = {
+        "task": "reflector.worker.healthcheck.healthcheck_ping",
+        "schedule": 60.0 * 10,
+    }
+    logger.info("Healthcheck enabled", url=settings.HEALTHCHECK_URL)
+else:
+    logger.warning("Healthcheck disabled, no url configured")
diff --git a/server/reflector/worker/healthcheck.py b/server/reflector/worker/healthcheck.py
new file mode 100644
index 00000000..e4ce6bc3
--- /dev/null
+++ b/server/reflector/worker/healthcheck.py
@@ -0,0 +1,18 @@
+import httpx
+import structlog
+from celery import shared_task
+from reflector.settings import settings
+
+logger = structlog.get_logger(__name__)
+
+
+@shared_task
+def healthcheck_ping():
+    url = settings.HEALTHCHECK_URL
+    if not url:
+        return
+    try:
+        print("pinging healthcheck url", url)
+        httpx.get(url, timeout=10)
+    except Exception as e:
+        logger.error("healthcheck_ping", error=str(e))
diff --git a/server/runserver.sh b/server/runserver.sh
index b0c3f138..31cce123 100755
--- a/server/runserver.sh
+++ b/server/runserver.sh
@@ -9,6 +9,8 @@ if [ "${ENTRYPOINT}" = "server" ]; then
     python -m reflector.app
 elif [ "${ENTRYPOINT}" = "worker" ]; then
     celery -A reflector.worker.app worker --loglevel=info
+elif [ "${ENTRYPOINT}" = "beat" ]; then
+    celery -A reflector.worker.app beat --loglevel=info
 else
     echo "Unknown command"
 fi