Merge pull request #66 from eth-cscs/health-check-refactoring

Refactored health checks

Author: theely

f7t-api-config.local-env-tests.yaml

@@ -67,4 +67,7 @@ storage:
     max_part_size: 1073741824 # 1G
     parallel_runs: 3
     tmp_folder: "tmp"
-  max_ops_file_size: 1048576 # 1M
+  max_ops_file_size: 1048576 # 1M
+  probing:
+    interval: 60
+    timeout: 10

f7t-api-config.local-env.yaml

@@ -91,4 +91,5 @@ storage:
     tmp_folder: "tmp"
   max_ops_file_size: 1048576 # 1M
   probing:
+    interval: 60
     timeout: 10

src/firecrest/config.py

@@ -12,6 +12,7 @@
     BaseModel,
     ConfigDict,
     Field,
+    SecretStr,
     field_validator,
 )
 from pydantic_settings import (
@@ -75,57 +76,6 @@ def to_json(self):
         }
 
 
-class StorageProbing(CamelModel):
-    """Probing configuration to check availability of the storage system."""
-
-    timeout: int = Field(
-        ..., description="Timeout for storage health probing in seconds."
-    )
-
-
-class Storage(BaseModel):
-    """Object storage configuration, including credentials, endpoints, and upload behavior."""
-
-    name: str = Field(..., description="Name identifier for the storage.")
-    private_url: str = Field(
-        ..., description="Private/internal endpoint URL for the storage."
-    )
-    public_url: str = Field(..., description="Public/external URL for the storage.")
-    access_key_id: str = Field(
-        ..., description="Access key ID for S3-compatible storage."
-    )
-    secret_access_key: LoadFileSecretStr = Field(
-        ...,
-        description=(
-            "Secret access key for storage. You can give directly the "
-            "content or the file path using `'secret_file:/path/to/file'`."
-        ),
-    )
-    region: str = Field(..., description="Region of the storage bucket.")
-    ttl: int = Field(..., description="Time-to-live (in seconds) for generated URLs.")
-    tenant: Optional[str] = Field(
-        None, description="Optional tenant identifier for multi-tenant setups."
-    )
-    multipart: MultipartUpload = Field(
-        default_factory=MultipartUpload,
-        description="Settings for multipart upload, including chunk size and concurrency.",
-    )
-    bucket_lifecycle_configuration: BucketLifecycleConfiguration = Field(
-        default_factory=BucketLifecycleConfiguration,
-        description="Lifecycle policy settings for auto-deleting files after a given number of days.",
-    )
-    max_ops_file_size: int = Field(
-        5 * 1024 * 1024,
-        description=(
-            "Maximum file size (in bytes) allowed for direct upload and "
-            "download. Larger files will go through the staging area."
-        ),
-    )
-    probing: Optional[StorageProbing] = Field(
-        None, description="Configuration for probing storage availability."
-    )
-
-
 class SchedulerType(str, Enum):
     """Supported job scheduler types."""
 
@@ -230,7 +180,7 @@ class HealthCheckException(BaseServiceHealth):
     pass
 
 
-class ClusterProbing(CamelModel):
+class Probing(CamelModel):
     """Cluster monitoring attributes."""
 
     interval: int = Field(
@@ -239,6 +189,53 @@ class ClusterProbing(CamelModel):
     timeout: int = Field(..., description="Maximum time in seconds allowed per check.")
 
 
+class Storage(BaseModel):
+    """Object storage configuration, including credentials, endpoints, and upload behavior."""
+
+    name: str = Field(..., description="Name identifier for the storage.")
+    private_url: SecretStr = Field(
+        ..., description="Private/internal endpoint URL for the storage."
+    )
+    public_url: str = Field(..., description="Public/external URL for the storage.")
+    access_key_id: SecretStr = Field(
+        ..., description="Access key ID for S3-compatible storage."
+    )
+    secret_access_key: LoadFileSecretStr = Field(
+        ...,
+        description=(
+            "Secret access key for storage. You can give directly the "
+            "content or the file path using `'secret_file:/path/to/file'`."
+        ),
+    )
+    region: str = Field(..., description="Region of the storage bucket.")
+    ttl: int = Field(..., description="Time-to-live (in seconds) for generated URLs.")
+    tenant: Optional[str] = Field(
+        None, description="Optional tenant identifier for multi-tenant setups."
+    )
+    multipart: MultipartUpload = Field(
+        default_factory=MultipartUpload,
+        description="Settings for multipart upload, including chunk size and concurrency.",
+    )
+    bucket_lifecycle_configuration: BucketLifecycleConfiguration = Field(
+        default_factory=BucketLifecycleConfiguration,
+        description="Lifecycle policy settings for auto-deleting files after a given number of days.",
+    )
+    max_ops_file_size: int = Field(
+        5 * 1024 * 1024,
+        description=(
+            "Maximum file size (in bytes) allowed for direct upload and "
+            "download. Larger files will go through the staging area."
+        ),
+    )
+    probing: Optional[Probing] = Field(
+        None, description="Configuration for probing storage availability."
+    )
+    servicesHealth: Optional[List[S3ServiceHealth]] = Field(
+        None,
+        description="Optional health information for different services in the cluster.",
+    )
+
+
 class FileSystem(CamelModel):
     """Defines a cluster file system and its type."""
 
@@ -306,14 +303,13 @@ class HPCCluster(CamelModel):
             SchedulerServiceHealth
             | FilesystemServiceHealth
             | SSHServiceHealth
-            | S3ServiceHealth
             | HealthCheckException
         ]
     ] = Field(
         None,
         description="Optional health information for different services in the cluster.",
     )
-    probing: ClusterProbing = Field(
+    probing: Probing = Field(
         ..., description="Probing configuration for monitoring the cluster."
     )
     file_systems: List[FileSystem] = Field(

src/firecrest/dependencies.py

@@ -321,14 +321,14 @@ def __init__(
         if settings.storage:
             self.url = settings.storage.public_url
             if connection == S3ClientConnectionType.private:
-                self.url = settings.storage.private_url
+                self.url = settings.storage.private_url.get_secret_value()
 
     async def __call__(self):
         async with get_session().create_client(
             "s3",
             region_name=settings.storage.region,
             aws_secret_access_key=settings.storage.secret_access_key.get_secret_value(),
-            aws_access_key_id=settings.storage.access_key_id,
+            aws_access_key_id=settings.storage.access_key_id.get_secret_value(),
             endpoint_url=self.url,
             config=AioConfig(signature_version="s3v4"),
         ) as client:

src/firecrest/main.py

@@ -17,6 +17,8 @@
 from fastapi.exceptions import RequestValidationError
 
 from starlette.exceptions import HTTPException as StarletteHTTPException
+from firecrest.status.health_check.health_checker_cluster import ClusterHealthChecker
+from firecrest.status.health_check.health_checker_storage import StorageHealthChecker
 from starlette_context import plugins
 from starlette_context.middleware import RawContextMiddleware
 
@@ -34,9 +36,6 @@
 )
 from lib.ssh_clients.ssh_keygen_client import SSHKeygenClient
 from firecrest.dependencies import SSHClientDependency
-from firecrest.status.health_check.health_checker import (
-    SchedulerHealthChecker,
-)
 
 # routers
 from firecrest.status.router import (
@@ -111,10 +110,16 @@ async def schedule_tasks(scheduler: AsyncScheduler):
     for cluster in plugin_settings.clusters:
         if cluster.probing:
             await scheduler.add_schedule(
-                SchedulerHealthChecker(cluster).check,
+                ClusterHealthChecker(cluster).check,
                 IntervalTrigger(seconds=cluster.probing.interval),
                 id=f"check-cluster-{cluster.name}",
             )
+    if settings.storage and settings.storage.probing:
+        await scheduler.add_schedule(
+            StorageHealthChecker(settings.storage).check,
+            IntervalTrigger(seconds=settings.storage.probing.interval),
+            id="check-storage",
+        )
     await scheduler.add_schedule(
         SSHClientDependency.prune_client_pools,
         IntervalTrigger(seconds=5),

src/firecrest/status/health_check/checks/health_check_base.py

@@ -11,7 +11,7 @@
 
 class HealthCheckBase(ABC):
 
-    def __init__(self, system: HPCCluster):
+    def __init__(self, system: HPCCluster = None):
         self.system = system
 
     async def check(self) -> BaseServiceHealth:

src/firecrest/status/health_check/checks/health_check_s3.py

@@ -4,7 +4,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from firecrest.config import (
-    HPCCluster,
     S3ServiceHealth,
 )
 from firecrest.dependencies import (
@@ -16,8 +15,8 @@
 
 class S3HealthCheck(HealthCheckBase):
 
-    def __init__(self, system: HPCCluster, timeout: int):
-        super().__init__(system)
+    def __init__(self, timeout: int):
+        super().__init__()
         self.timeout = timeout
 
     async def execute_check(self) -> S3ServiceHealth:

src/firecrest/status/health_check/health_checker.py renamed to src/firecrest/status/health_check/health_checker_cluster.py

@@ -12,7 +12,6 @@
 from firecrest.status.health_check.checks.health_check_filesystem import (
     FilesystemHealthCheck,
 )
-from firecrest.status.health_check.checks.health_check_s3 import S3HealthCheck
 from firecrest.status.health_check.checks.health_check_scheduler import (
     SchedulerHealthCheck,
 )
@@ -25,7 +24,7 @@
 from firecrest.plugins import settings
 
 
-class SchedulerHealthChecker:
+class ClusterHealthChecker:
 
     scheduler_client: SchedulerBaseClient = None
     cluster: HPCCluster = None
@@ -78,20 +77,12 @@ async def check(self) -> None:
                 )
                 checks += [filesystemCheck.check()]
 
-            if settings.storage and settings.storage.probing:
-                s3Check = S3HealthCheck(
-                    system=self.cluster, timeout=settings.storage.probing.timeout
-                )
-                checks += [s3Check.check()]
-
             results = await asyncio.gather(*checks, return_exceptions=True)
             self.cluster.servicesHealth = results
         except Exception as ex:
-            error_message = (
-                f"HealthChecker execution failed with error: {ex.__class__.__name__}"
-            )
+            error_message = f"Cluster HealthChecker execution failed with error: {ex.__class__.__name__}"
             if len(str(ex)) > 0:
-                error_message = f"HealthChecker execution failed with error: {ex.__class__.__name__} - {str(ex)}"
+                error_message = f"Cluster HealthChecker execution failed with error: {ex.__class__.__name__} - {str(ex)}"
             exception = HealthCheckException(service_type="exception")
             exception.healthy = False
             exception.last_checked = time.time()

src/firecrest/status/health_check/health_checker_storage.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025, ETH Zurich. All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import asyncio
+import time
+
+
+from firecrest.config import HealthCheckException, Storage
+from firecrest.status.health_check.checks.health_check_s3 import S3HealthCheck
+from firecrest.plugins import settings
+
+
+class StorageHealthChecker:
+
+    storage: Storage = None
+
+    def __init__(self, storage: Storage):
+        self.storage = storage
+
+    async def check(self) -> None:
+        try:
+            checks = []
+            s3Check = S3HealthCheck(timeout=settings.storage.probing.timeout)
+            checks += [s3Check.check()]
+
+            results = await asyncio.gather(*checks, return_exceptions=True)
+            self.storage.servicesHealth = results
+        except Exception as ex:
+            error_message = f"Storage HealthChecker execution failed with error: {ex.__class__.__name__}"
+            if len(str(ex)) > 0:
+                error_message = f"Storage HealthChecker execution failed with error: {ex.__class__.__name__} - {str(ex)}"
+            exception = HealthCheckException(service_type="exception")
+            exception.healthy = False
+            exception.last_checked = time.time()
+            exception.message = error_message
+            self.storage.servicesHealth = [exception]
+            # Note: raising the exception might not be handled well by apscheduler.
+            # Instead consider printing the exceotion with: traceback.print_exception(ex)
+            raise ex

src/firecrest/status/models.py

@@ -4,10 +4,10 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from datetime import datetime
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 # configs
-from firecrest.config import HPCCluster
+from firecrest.config import HPCCluster, Storage
 
 # models
 from lib.models import CamelModel
@@ -25,6 +25,7 @@ class GetLiveness(CamelModel):
 
 class GetSystemsResponse(CamelModel):
     systems: List[HPCCluster]
+    storage: Optional[Storage] = None
 
 
 class GetNodesResponse(CamelModel):

src/firecrest/status/router.py

@@ -62,7 +62,7 @@
     response_description="System list returned successfully",
 )
 async def get_systems() -> Any:
-    return {"systems": settings.clusters}
+    return {"systems": settings.clusters, "storage": settings.storage}
 
 
 @router_on_systen.get(

tests/conftest.py

@@ -66,8 +66,8 @@ async def s3_client():
         "s3",
         region_name=settings.storage.region,
         aws_secret_access_key=settings.storage.secret_access_key,
-        aws_access_key_id=settings.storage.access_key_id,
-        endpoint_url=settings.storage.private_url,
+        aws_access_key_id=settings.storage.access_key_id.get_secret_value(),
+        endpoint_url=settings.storage.private_url.get_secret_value(),
         config=AioConfig(signature_version="s3v4"),
     ) as client:
         global_s3_client = client

tests/health_check_test.py

@@ -9,7 +9,7 @@
 import pytest
 from aioresponses import aioresponses
 from pytest_httpx import HTTPXMock
-from firecrest.status.health_check.health_checker import SchedulerHealthChecker
+from firecrest.status.health_check.health_checker_cluster import ClusterHealthChecker
 from lib.auth.authN.OIDC_token_auth import OIDCTokenAuth
 from lib.models.apis.api_auth_model import ApiAuthModel
 from tests import mocked_api_responses
@@ -67,7 +67,7 @@ async def test_health_check(
             body=json.dumps(mocked_nodes_get_response),
         )
 
-        health_checker = SchedulerHealthChecker(
+        health_checker = ClusterHealthChecker(
             slurm_cluster_with_api_config, token_decoder=TokenDecoderMock()
         )
         await health_checker.check()