Merge pull request #26 from saichandrapandraju/prometheus-publishing

feat(RHOAIENG-21044): Metrics: Prometheus publishing

Author: ruivieira

pyproject.toml

@@ -20,6 +20,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "pytest>=7.4.2,<9",
+    "pytest-asyncio>=0.26.0,<2",
     "isort>=5.12.0,<7",
     "flake8>=6.1.0,<8",
     "mypy>=1.5.1,<2",

src/endpoints/metrics/metrics_directory.py

@@ -0,0 +1,35 @@
+import logging
+from typing import Callable, Dict
+
+from pandas import DataFrame
+
+from src.service.payloads.metrics.base_metric_request import BaseMetricRequest
+from src.service.prometheus.metric_value_carrier import MetricValueCarrier
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+class MetricsDirectory:
+    def __init__(self) -> None:
+        self.calculator_directory: Dict[
+            str, Callable[[DataFrame, BaseMetricRequest], MetricValueCarrier]
+        ] = {}
+
+    def register(
+        self,
+        name: str,
+        calculator: Callable[[DataFrame, BaseMetricRequest], MetricValueCarrier],
+    ) -> None:
+        if name not in self.calculator_directory:
+            self.calculator_directory[name] = calculator
+            logger.debug(f"Registered calculator for metric: {name}")
+        else:
+            logger.warning(
+                f"Attempted to register duplicate calculator for metric: {name}. "
+                "Ignoring duplicate registration."
+            )
+
+    def get_calculator(
+        self, name: str
+    ) -> Callable[[DataFrame, BaseMetricRequest], MetricValueCarrier]:
+        return self.calculator_directory.get(name)

src/main.py

@@ -1,33 +1,37 @@
-import os
+import asyncio
+import logging
+from contextlib import asynccontextmanager
 
 import uvicorn
 from fastapi import FastAPI, Request, Response
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
+from fastapi_utils.tasks import repeat_every
 from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
-from fastapi.middleware.cors import CORSMiddleware
-import logging
 
 # Endpoint routers
 from src.endpoints.consumer.consumer_endpoint import router as consumer_router
-from src.endpoints.metrics.fairness.group.dir import router as dir_router
+from src.endpoints.data.data_download import router as data_download_router
 from src.endpoints.data.data_upload import router as data_upload_router
 
+# from src.endpoints.explainers import router as explainers_router
+from src.endpoints.explainers.global_explainer import router as explainers_global_router
+from src.endpoints.explainers.local_explainer import router as explainers_local_router
+from src.endpoints.metadata import router as metadata_router
+
 # from src.endpoints.drift_metrics import router as drift_metrics_router
 from src.endpoints.metrics.drift.approx_ks_test import (
     router as drift_approx_ks_test_router,
 )
 from src.endpoints.metrics.drift.fourier_mmd import router as drift_fourier_mmd_router
 from src.endpoints.metrics.drift.ks_test import router as drift_ks_test_router
 from src.endpoints.metrics.drift.meanshift import router as drift_meanshift_router
-
-# from src.endpoints.explainers import router as explainers_router
-from src.endpoints.explainers.global_explainer import router as explainers_global_router
-from src.endpoints.explainers.local_explainer import router as explainers_local_router
+from src.endpoints.metrics.fairness.group.dir import router as dir_router
 from src.endpoints.metrics.fairness.group.spd import router as spd_router
 from src.endpoints.metrics.identity.identity_endpoint import router as identity_router
-from src.endpoints.metadata import router as metadata_router
 from src.endpoints.metrics.metrics_info import router as metrics_info_router
-from src.endpoints.data.data_download import router as data_download_router
+
+from src.service.prometheus.prometheus_scheduler import PrometheusScheduler
 
 try:
     from src.endpoints.evaluation.lm_evaluation_harness import (
@@ -44,10 +48,36 @@
 )
 logger = logging.getLogger(__name__)
 
+prometheus_scheduler = PrometheusScheduler()
+
+
+@repeat_every(
+    seconds=prometheus_scheduler.service_config.get("metrics_schedule", 30),
+    logger=logger,
+    raise_exceptions=False,
+)
+async def schedule_metrics_calculation():
+    await prometheus_scheduler.calculate()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    task = asyncio.create_task(schedule_metrics_calculation())
+
+    yield
+
+    task.cancel()
+    try:
+        await task
+    except asyncio.CancelledError:
+        logger.info("Prometheus metrics calculation task cancelled during shutdown")
+
+
 app = FastAPI(
     title="TrustyAI Service API",
     version="1.0.0rc0",
     description="TrustyAI Service API",
+    lifespan=lifespan,
 )
 
 # CORS

src/service/__init__.py

@@ -1,4 +1,4 @@
 # Place all external packages and service algorithms here
 # The main REST server should only call
 # - modules directly related to the service HTTP, data, prometheus, etc.
-# - This service module
+# - This service module

src/service/constants.py

@@ -4,9 +4,13 @@
 METADATA_SUFFIX = "_metadata"
 PROTECTED_DATASET_SUFFIX = "trustyai_internal_"
 PARTIAL_PAYLOAD_DATASET_NAME = "partial_payloads"
-
+GROUND_TRUTH_SUFFIX = "-ground-truths"
+METADATA_FILENAME = "metadata.json"
+INTERNAL_DATA_FILENAME = "internal_data.csv"
 # Payload parsing
 TRUSTYAI_TAG_PREFIX = "_trustyai"
 SYNTHETIC_TAG = f"{TRUSTYAI_TAG_PREFIX}_synthetic"
 UNLABELED_TAG = f"{TRUSTYAI_TAG_PREFIX}_unlabeled"
-BIAS_IGNORE_PARAM =  "bias-ignore"
+BIAS_IGNORE_PARAM = "bias-ignore"
+# Prometheus constants
+PROMETHEUS_METRIC_PREFIX = "trustyai_"