diff --git a/bioimageio_colab/models/sam_deployment.py b/bioimageio_colab/models/sam_deployment.py index 3267ce1..e914e1c 100644 --- a/bioimageio_colab/models/sam_deployment.py +++ b/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 diff --git a/bioimageio_colab/register_sam_service.py b/bioimageio_colab/register_sam_service.py index 2e58ac5..a7f400e 100644 --- a/bioimageio_colab/register_sam_service.py +++ b/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" ) diff --git a/pyproject.toml b/pyproject.toml index 668febc..c1be771 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tests/test_model_service.py b/tests/test_model_service.py index 58ee538..1301239 100644 --- a/tests/test_model_service.py +++ b/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"