From 06393fc9c6c8c7a2962173bafc25faa2a535a059 Mon Sep 17 00:00:00 2001
From: Nils Mechtel <nils.mech@gmail.com>
Date: Tue, 10 Dec 2024 16:40:42 +0100
Subject: [PATCH] add probes

---
 bioimageio_colab/register_sam_service.py | 38 +++++++++++++++++++++++-
 pyproject.toml                           |  2 +-
 test/test_model_service.py               |  4 +--
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/bioimageio_colab/register_sam_service.py b/bioimageio_colab/register_sam_service.py
index ee8ab9a..57756ce 100644
--- a/bioimageio_colab/register_sam_service.py
+++ b/bioimageio_colab/register_sam_service.py
@@ -208,6 +208,22 @@ def clear_cache(embedding_cache: TTLCache, context: dict = None) -> bool:
 def hello(context: dict = None) -> str:
     return "Welcome to the Interactive Segmentation service!"
 
+def ping(context: dict = None) -> str:
+    return "pong"
+
+def check_readiness():
+    # Check if the service is ready
+    return {"status": "ok"}
+
+async def check_liveness(colab_client, sid):
+    # Check if the service is alive
+    service = await colab_client.get_service(sid)
+    alive = await service.ping() == "pong"
+    if alive:
+        return {"status": "ok"}
+    else:
+        return {"success": "false", "detail": "Service is not alive"}
+
 
 async def register_service(args: dict) -> None:
     """
@@ -243,6 +259,7 @@ async def register_service(args: dict) -> None:
             },
             # Exposed functions:
             "hello": hello,
+            "ping": ping,
             # **Run segmentation**
             # Params:
             # - model name
@@ -269,7 +286,26 @@ async def register_service(args: dict) -> None:
     sid = service_info["id"]
     logger.info(f"Registered service with ID: {sid}")
     logger.info(
-        f"Test the service here: {args.server_url}/{args.workspace_name}/services/{args.service_id}/hello"
+        f"Test the service here: {args.server_url}/{args.workspace_name}/services/{sid.split('/')[1]}/hello"
+    )
+
+    # Register a probe for the service
+    probe_info = await colab_client.register_service({
+        "name": "Probes service",
+        "id": "probes",
+        "config": {
+            "visibility": "public"
+        },
+        "readiness": check_readiness,
+        "liveness": partial(check_liveness, colab_client, sid)
+    })
+    pid = probe_info["id"]
+    logger.info(f"Registered probe with ID: {pid}")
+    logger.info(
+        f"Test the readiness probe here: {args.server_url}/{args.workspace_name}/services/{pid.split('/')[1]}/readiness"
+    )
+    logger.info(
+        f"Test the liveness probe here: {args.server_url}/{args.workspace_name}/services/{pid.split('/')[1]}/liveness"
     )
 
 
diff --git a/pyproject.toml b/pyproject.toml
index 5bc71ca..8add3c0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "bioimageio-colab"
-version = "0.1.4"
+version = "0.1.5"
 readme = "README.md"
 description = "Collaborative image annotation and model training with human in the loop."
 dependencies = [
diff --git a/test/test_model_service.py b/test/test_model_service.py
index be49ecf..88843c5 100644
--- a/test/test_model_service.py
+++ b/test/test_model_service.py
@@ -10,10 +10,10 @@
 
 
 def test_service_available():
-    service_url = f"{SERVER_URL}/{WORKSPACE_NAME}/services/{SERVICE_ID}/hello"
+    service_url = f"{SERVER_URL}/{WORKSPACE_NAME}/services/{SERVICE_ID}/ping"
     response = requests.get(service_url)
     assert response.status_code == 200
-    assert response.json() == "Welcome to the Interactive Segmentation service!"
+    assert response.json() == "pong"
 
 def test_get_service():
     client = connect_to_server({"server_url": SERVER_URL, "method_timeout": 5})