replace probes service

bioimageio_colab/models/sam_deployment.py

@@ -40,7 +40,7 @@ async def _download_model(self, model_path: str, model_url: str) -> None:
                 with open(model_path, "wb") as f:
                     f.write(content)
 
-    @serve.multiplexed(max_num_models_per_replica=2)
+    @serve.multiplexed(max_num_models_per_replica=1)
     async def get_model(self, model_id: str):
         import os
 

bioimageio_colab/register_sam_service.py

@@ -235,27 +235,6 @@ async def compute_image_embedding(
 #         raise e
 
 
-async def check_readiness(client: RemoteService, service_id: str) -> dict:
-    """
-    Readiness probe for the SAM service.
-    """
-    try:
-        services = await client.list_services()
-        service_found = False
-        for service in services:
-            if service["id"] == service_id:
-                service_found = True
-                break
-        assert service_found, f"Service with ID '{service_id}' not found."
-
-        logger.info(f"Service with ID '{service_id}' is ready.")
-
-        return {"status": "ready"}
-    except Exception as e:
-        logger.error(f"Error during readiness probe: {e}")
-        raise e
-
-
 def format_time(last_deployed_time_s, tz: timezone = timezone.utc) -> str:
     # Get the current time
     current_time = datetime.now(tz)
@@ -290,7 +269,9 @@ def format_time(last_deployed_time_s, tz: timezone = timezone.utc) -> str:
     }
 
 
-async def check_liveness(app_name: str, client: RemoteService, service_id: str) -> dict:
+async def deployment_status(
+    app_name: str, service_id: str, registration_time_s: float, context: dict = None
+) -> dict:
     """
     Liveness probe for the SAM service.
     """
@@ -316,13 +297,12 @@ async def check_liveness(app_name: str, client: RemoteService, service_id: str)
                 "replica_states": deployment.replica_states,
             }
 
-        # Check if the service can be accessed
-        service = await client.get_service(service_id)
-        assert await service.ping() == "pong"
-
-        output["service"] = {
+        formatted_time = format_time(registration_time_s)
+        output["hypha_service"] = {
             "status": "RUNNING",
             "service_id": service_id,
+            "last_registered_at": formatted_time["last_deployed_at"],
+            "duration_since": formatted_time["duration_since"],
         }
 
         return output
@@ -393,6 +373,12 @@ async def register_service(args: dict) -> None:
             # Exposed functions:
             "hello": hello,
             "ping": ping,
+            "deployment_status": partial(
+                deployment_status,
+                app_name=app_name,
+                service_id=f"{workspace}/{client_id}:{args.service_id}",
+                registration_time_s=datetime.now(timezone.utc).timestamp(),
+            ),
             "compute_embedding": partial(
                 compute_image_embedding,
                 app_handle=app_handle,
@@ -409,26 +395,6 @@ async def register_service(args: dict) -> None:
     sid = service_info["id"]
     logger.info(f"Service registered with ID: {sid}")
     logger.info(f"Test the service here: {client_base_url}:{args.service_id}/hello")
-
-    # Register probes for the service
-    await client.register_probes(
-        {
-            "readiness": partial(
-                check_readiness,
-                client=client,
-                service_id=sid,
-            ),
-            "liveness": partial(
-                check_liveness,
-                app_name=app_name,
-                client=client,
-                service_id=sid,
-            ),
-        }
-    )
-
-    # This will register probes service where you can accessed via hypha or the HTTP proxy
-    logger.info(f"Probes registered at workspace: {workspace}")
     logger.info(
-        f"Test the liveness probe here: {args.server_url}/{workspace}/services/probes/liveness"
+        f"Check deployment status: {client_base_url}:{args.service_id}/deployment_status"
     )

pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"]
 
 [project]
 name = "bioimageio-colab"
-version = "0.2.3"
+version = "0.2.4"
 readme = "README.md"
 description = "Collaborative image annotation and model training with human in the loop."
 dependencies = [

tests/test_model_service.py

@@ -11,7 +11,7 @@
 IMG_PATH = "./data/example_image.tif"
 
 
-def test_service_is_alive_http_api():
+def test_service_http_api():
     client_str = f"{CLIENT_ID}:" if CLIENT_ID else ""
     service_url = f"{SERVER_URL}/{WORKSPACE_NAME}/services/{client_str}{SERVICE_ID}"
 
@@ -23,16 +23,7 @@ def test_service_is_alive_http_api():
     assert response.status_code == 200
     assert response.json() == "pong"
 
-
-def test_probes_http_api():
-    client_str = f"{CLIENT_ID}:" if CLIENT_ID else ""
-    probes_url = f"{SERVER_URL}/{WORKSPACE_NAME}/services/{client_str}probes"
-
-    response = requests.get(f"{probes_url}/readiness")
-    assert response.status_code == 200
-    assert response.json() == {"status": "ready"}
-
-    response = requests.get(f"{probes_url}/liveness")
+    response = requests.get(f"{service_url}/deployment_status")
     assert response.status_code == 200
     for value in response.json().values():
         assert value["status"] == "RUNNING"