Repository: avast/ep-stats (language: python)

Issue:
prometheus metric fix



Discussion / hints:

Fix patch:
diff --git a/requirements.txt b/requirements.txt
index ff08dc2..1537488 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,6 +4,7 @@
 #
 #    pip-compile
 #
+
 anyio==3.7.1
     # via starlette
 asgiref==3.7.2
@@ -30,6 +31,8 @@ pandas==2.0.3
     #   statsmodels
 patsy==0.5.3
     # via statsmodels
+prometheus-client==0.17.1
+    # via ep-stats (setup.py)
 pydantic==1.10.12
     # via
     #   ep-stats (setup.py)
@@ -52,8 +55,6 @@ sniffio==1.3.0
     # via anyio
 starlette==0.27.0
     # via fastapi
-statsd==3.3.0
-    # via ep-stats (setup.py)
 statsmodels==0.13.5
     # via ep-stats (setup.py)
 typing-extensions==4.7.1
diff --git a/setup.cfg b/setup.cfg
index c6bf6a3..cf4ac00 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -27,13 +27,12 @@ setup_requires =
 install_requires =
     pandas>=1.3.5
     scipy>=1.7.3
-    statsd==3.3.0
     pyparsing==2.4.6
-    statsd==3.3.0
     fastapi==0.95.2
     uvicorn==0.17.6
     pydantic>=1.8.2
     statsmodels==0.13.5
+    prometheus-client>=0.17.1
 
 tests_require =
     pytest==7.0.1
diff --git a/src/epstats/main.py b/src/epstats/main.py
index 82cca6b..56c99b8 100644
--- a/src/epstats/main.py
+++ b/src/epstats/main.py
@@ -1,18 +1,15 @@
 import logging
 from concurrent.futures import ThreadPoolExecutor
-from statsd import StatsClient
 from pydantic import BaseSettings
 
 from .toolkit.testing import TestDaoFactory, TestData
 from .server import get_api, serve, ApiSettings
+from prometheus_client import make_asgi_app
 
 
 class Settings(BaseSettings):
     api: ApiSettings = ApiSettings()
 
-    statsd_host: str = "localhost"
-    statsd_port: int = 8888
-
     evaluation_pool_size: int = 5
 
 
@@ -40,16 +37,9 @@ def get_executor_pool():
         pass
 
 
-def get_statsd():
-    try:
-        prefix = f"{settings.api.app_name}.{settings.api.app_env}"
-        statsd = StatsClient(settings.statsd_host, settings.statsd_port, prefix=prefix)
-        yield statsd
-    finally:
-        pass
-
-
-api = get_api(settings.api, get_dao, get_executor_pool, get_statsd)
+api = get_api(settings.api, get_dao, get_executor_pool)
+metrics_app = make_asgi_app()
+api.mount("/metrics", metrics_app)
 
 
 def main():
diff --git a/src/epstats/prometheus.py b/src/epstats/prometheus.py
new file mode 100644
index 0000000..7b96c83
--- /dev/null
+++ b/src/epstats/prometheus.py
@@ -0,0 +1,9 @@
+from prometheus_client import Counter, Summary, REGISTRY
+from typing import Union
+
+
+def get_prometheus_metric(metric_name: str, metric_type: type) -> Union[Counter, Summary]:
+    try:
+        return metric_type(metric_name, "")
+    except ValueError:
+        return REGISTRY._names_to_collectors[metric_name]
diff --git a/src/epstats/server/api.py b/src/epstats/server/api.py
index 812d2c9..ddd33d8 100644
--- a/src/epstats/server/api.py
+++ b/src/epstats/server/api.py
@@ -1,8 +1,7 @@
 import logging
 from typing import Dict
 import uvicorn
-from statsd import StatsClient
-from fastapi import FastAPI, Depends
+from fastapi import FastAPI
 from fastapi.exception_handlers import (
     http_exception_handler,
     request_validation_exception_handler,
@@ -16,7 +15,7 @@
 from .api_settings import ApiSettings
 
 
-def get_api(settings: ApiSettings, get_dao, get_executor_pool, get_statsd) -> FastAPI:
+def get_api(settings: ApiSettings, get_dao, get_executor_pool) -> FastAPI:
     api = FastAPI(
         title=settings.app_title,
         description=settings.app_description,
@@ -43,12 +42,11 @@ async def custom_request_validation_exception_handler(request, exc):
         return await request_validation_exception_handler(request, exc)
 
     @api.get("/health", tags=["Health"])
-    async def readiness_liveness_probe(statsd: StatsClient = Depends(get_statsd)):
-        statsd.incr("requests.health")
+    async def readiness_liveness_probe():
         return {"message": "ep-stats-api is ready"}
 
-    api.include_router(get_evaluate_router(get_dao, get_executor_pool, get_statsd))
-    api.include_router(get_sample_size_calculation_router(get_executor_pool, get_statsd))
+    api.include_router(get_evaluate_router(get_dao, get_executor_pool))
+    api.include_router(get_sample_size_calculation_router(get_executor_pool))
 
     return api
 
diff --git a/src/epstats/server/api_evaluate.py b/src/epstats/server/api_evaluate.py
index 53a0f41..3fe99cb 100644
--- a/src/epstats/server/api_evaluate.py
+++ b/src/epstats/server/api_evaluate.py
@@ -1,10 +1,10 @@
 import logging
 import numpy as np
 from fastapi import APIRouter, Depends, HTTPException
-from statsd import StatsClient
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
+from ..prometheus import get_prometheus_metric, Summary, Counter
 from ..toolkit import Experiment as EvExperiment
 from ..toolkit import Dao
 
@@ -13,19 +13,25 @@
 
 
 _logger = logging.getLogger("epstats")
+evaluation_duration_metric = get_prometheus_metric("evaluation_duration_seconds", Summary)
+query_duration_metric = get_prometheus_metric("query_duration_seconds", Summary)
+stats_computation_duration_metric = get_prometheus_metric("stats_computation_duration_seconds", Summary)
+evaluation_errors_metric = get_prometheus_metric("evaluation_errors_total", Counter)
+evaluation_successes_metric = get_prometheus_metric("evaluation_successes_total", Counter)
+evaluation_requests_metric = get_prometheus_metric("evaluation_requests_total", Counter)
 
 
-def get_evaluate_router(get_dao, get_executor_pool, get_statsd) -> APIRouter:
-    def _evaluate(experiment: EvExperiment, dao: Dao, statsd: StatsClient):
+def get_evaluate_router(get_dao, get_executor_pool) -> APIRouter:
+    def _evaluate(experiment: EvExperiment, dao: Dao):
         try:
-            with statsd.timer("timing.evaluation"):
+            with evaluation_duration_metric.time():
                 _logger.debug(f"Loading goals for experiment [{experiment.id}]")
-                with statsd.timer("timing.query"):
+                with query_duration_metric.time():
                     goals = dao.get_agg_goals(experiment).sort_values(["exp_variant_id", "goal"])
                     _logger.info(f"Retrieved {len(goals)} goals in experiment [{experiment.id}]")
-                with statsd.timer("timing.stats"):
+                with stats_computation_duration_metric.time():
                     evaluation = experiment.evaluate_agg(goals)
-                    statsd.incr("evaluations")
+                    evaluation_successes_metric.inc()
                 _logger.info(
                     f"Evaluation response: [{experiment.id}]",
                     {
@@ -40,7 +46,7 @@ def _evaluate(experiment: EvExperiment, dao: Dao, statsd: StatsClient):
         except Exception as e:
             _logger.error(f"Cannot evaluate experiment [{experiment.id}] because of {e}")
             _logger.exception(e)
-            statsd.incr("errors.experiment")
+            evaluation_errors_metric.inc()
             raise HTTPException(
                 status_code=500,
                 detail=f"Cannot evaluate experiment [{experiment.id}] because of {e}",
@@ -53,14 +59,13 @@ async def evaluate_experiment(
         experiment: Experiment,
         evaluation_pool: ThreadPoolExecutor = Depends(get_executor_pool),
         dao: Dao = Depends(get_dao),
-        statsd: StatsClient = Depends(get_statsd),
     ):
         """
         Evaluates single `Experiment`.
         """
         _logger.info(f"Evaluation request: [{experiment.id}]", experiment.dict())
-        statsd.incr("requests.evaluate")
+        evaluation_requests_metric.inc()
         loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(evaluation_pool, _evaluate, experiment.to_experiment(statsd), dao, statsd)
+        return await loop.run_in_executor(evaluation_pool, _evaluate, experiment.to_experiment(), dao)
 
     return router
diff --git a/src/epstats/server/api_sample_size_calculation.py b/src/epstats/server/api_sample_size_calculation.py
index a801b3d..976ad80 100644
--- a/src/epstats/server/api_sample_size_calculation.py
+++ b/src/epstats/server/api_sample_size_calculation.py
@@ -1,6 +1,5 @@
 import logging
 from fastapi import APIRouter, Depends, HTTPException
-from statsd import StatsClient
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 
@@ -13,8 +12,8 @@
 _logger = logging.getLogger("epstats")
 
 
-def get_sample_size_calculation_router(get_executor_pool, get_statsd) -> APIRouter:
-    def _sample_size_calculation(data: SampleSizeCalculationData, statsd: StatsClient):
+def get_sample_size_calculation_router(get_executor_pool) -> APIRouter:
+    def _sample_size_calculation(data: SampleSizeCalculationData):
         try:
 
             if data.std is None:
@@ -29,7 +28,6 @@ def _sample_size_calculation(data: SampleSizeCalculationData, statsd: StatsClien
         except Exception as e:
             _logger.error(f"Cannot calculate the sample size because of: '{e}'")
             _logger.exception(e)
-            statsd.incr("errors.sample_size_calculation")
             raise HTTPException(
                 status_code=500,
                 detail=f"Cannot calculate the sample size because of: '{e}'",
@@ -41,14 +39,12 @@ def _sample_size_calculation(data: SampleSizeCalculationData, statsd: StatsClien
     async def sample_size_calculation(
         data: SampleSizeCalculationData,
         evaluation_pool: ThreadPoolExecutor = Depends(get_executor_pool),
-        statsd: StatsClient = Depends(get_statsd),
     ):
         """
         Calculates sample size based on `data`.
         """
         _logger.info(f"Calling the sample size calculation with {data.json()}")
-        statsd.incr("requests.sample_size_calculation")
         loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(evaluation_pool, _sample_size_calculation, data, statsd)
+        return await loop.run_in_executor(evaluation_pool, _sample_size_calculation, data)
 
     return router
diff --git a/src/epstats/server/req.py b/src/epstats/server/req.py
index cb4806c..9128501 100644
--- a/src/epstats/server/req.py
+++ b/src/epstats/server/req.py
@@ -2,7 +2,6 @@
 from pydantic import BaseModel, validator, root_validator, Field
 from pyparsing import ParseException
 from datetime import datetime
-from statsd import StatsClient
 from inspect import signature
 
 from ..toolkit import Experiment as EvExperiment, Filter as EvFilter, FilterScope
@@ -340,7 +339,7 @@ def check_date_from_to(cls, values):
 
         return values
 
-    def to_experiment(self, statsd: StatsClient):
+    def to_experiment(self):
         metrics = [m.to_metric() for m in self.metrics]
         checks = [c.to_check() for c in self.checks]
         return EvExperiment(
@@ -353,7 +352,6 @@ def to_experiment(self, statsd: StatsClient):
             date_for=self.date_for,
             unit_type=self.unit_type,
             variants=self.variants,
-            statsd=statsd,
             filters=[f.to_filter() for f in self.filters] if self.filters else [],
             query_parameters=self.query_parameters,
         )
diff --git a/src/epstats/toolkit/experiment.py b/src/epstats/toolkit/experiment.py
index 1ef831e..b52d694 100644
--- a/src/epstats/toolkit/experiment.py
+++ b/src/epstats/toolkit/experiment.py
@@ -6,7 +6,6 @@
 from collections import Counter
 from typing import Optional
 from datetime import datetime
-from statsd import StatsClient
 from dataclasses import dataclass
 
 from .metric import Metric, SimpleMetric
@@ -15,6 +14,10 @@
 from .parser import EpGoal, UnitType, AggType, Goal
 
 from .statistics import Statistics, DEFAULT_CONFIDENCE_LEVEL, DEFAULT_POWER
+from ..prometheus import get_prometheus_metric, Counter as PrometheusCounter
+
+
+check_evaluation_errors_metric = get_prometheus_metric("check_evaluation_errors_total", PrometheusCounter)
 
 
 class Evaluation:
@@ -138,7 +141,6 @@ def __init__(
         date_for: Optional[str] = None,
         confidence_level: float = DEFAULT_CONFIDENCE_LEVEL,
         variants: Optional[List[str]] = None,
-        statsd: StatsClient = StatsClient(),
         filters: Optional[List[Filter]] = None,
         query_parameters: dict = {},
     ):
@@ -171,7 +173,6 @@ def __init__(
             )
         ]
         self._update_dimension_to_value()
-        self.statsd = statsd
         self.filters = filters if filters is not None else []
         self.query_parameters = query_parameters
 
@@ -543,7 +544,7 @@ def _evaluate_checks(self, goals: pd.DataFrame, check_fce) -> pd.DataFrame:
                 res.append(r)
             except Exception as e:
                 self._logger.warning(f"Cannot evaluate check [{c.id} in experiment [{self.id}] because of {e}")
-                self.statsd.incr("errors.check")
+                check_evaluation_errors_metric.inc()
 
         c = pd.concat(res, axis=0) if res != [] else pd.DataFrame([], columns=Evaluation.check_columns())
         c["timestamp"] = round(get_utc_timestamp(datetime.now()).timestamp())


Test patch:
diff --git a/tests/epstats/server/depend.py b/tests/epstats/server/depend.py
index 2db2311..94d25b3 100644
--- a/tests/epstats/server/depend.py
+++ b/tests/epstats/server/depend.py
@@ -1,4 +1,3 @@
-from statsd import StatsClient
 from concurrent.futures import ThreadPoolExecutor
 
 from src.epstats.toolkit.testing import TestDaoFactory, TestData
@@ -14,14 +13,6 @@ def get_test_dao():
         dao.close()
 
 
-def get_test_statsd():
-    statsd = StatsClient("localhost", port="8888", prefix="epstats")
-    try:
-        yield statsd
-    finally:
-        pass
-
-
 def get_test_executor_pool():
     evaluation_pool = ThreadPoolExecutor(max_workers=1)
     try:
diff --git a/tests/epstats/server/test_api_evaluate.py b/tests/epstats/server/test_api_evaluate.py
index 6b6a9a0..d2b8acb 100644
--- a/tests/epstats/server/test_api_evaluate.py
+++ b/tests/epstats/server/test_api_evaluate.py
@@ -9,15 +9,14 @@
 )
 
 from src.epstats.main import api
-from src.epstats.main import get_dao, get_statsd, get_executor_pool
+from src.epstats.main import get_dao, get_executor_pool
 from src.epstats.server.res import Result
 
-from .depend import get_test_dao, get_test_executor_pool, get_test_statsd, dao_factory
+from .depend import get_test_dao, get_test_executor_pool, dao_factory
 
 
 client = TestClient(api)
 api.dependency_overrides[get_dao] = get_test_dao
-api.dependency_overrides[get_statsd] = get_test_statsd
 api.dependency_overrides[get_executor_pool] = get_test_executor_pool
 
 
@@ -294,6 +293,12 @@ def test_metric_with_minimum_effect():
     assert_experiment(resp.json(), dao_factory.get_dao(), 1, 0)
 
 
+def test_prometheus_metrics():
+    prometheus_resp = client.get("/metrics")
+    assert prometheus_resp.status_code == 200
+    assert "evaluation_duration_seconds" in prometheus_resp.text
+
+
 def assert_experiment(target, test_dao: TestDao, expected_metrics: int, expected_checks: int = 1) -> None:
     result = Result(**target)
     assert len(result.metrics) == expected_metrics
diff --git a/tests/epstats/server/test_api_sample_size_calculation.py b/tests/epstats/server/test_api_sample_size_calculation.py
index 3393fc9..10a5aee 100644
--- a/tests/epstats/server/test_api_sample_size_calculation.py
+++ b/tests/epstats/server/test_api_sample_size_calculation.py
@@ -3,13 +3,12 @@
 from fastapi.testclient import TestClient
 
 from src.epstats.main import api
-from src.epstats.main import get_statsd, get_executor_pool
+from src.epstats.main import get_executor_pool
 
-from .depend import get_test_executor_pool, get_test_statsd
+from .depend import get_test_executor_pool
 
 
 client = TestClient(api)
-api.dependency_overrides[get_statsd] = get_test_statsd
 api.dependency_overrides[get_executor_pool] = get_test_executor_pool
 
 
diff --git a/tests/epstats/server/test_req.py b/tests/epstats/server/test_req.py
index 7ee1103..8ba5111 100644
--- a/tests/epstats/server/test_req.py
+++ b/tests/epstats/server/test_req.py
@@ -1,14 +1,13 @@
 from fastapi.testclient import TestClient
 
 from src.epstats.main import api
-from src.epstats.main import get_dao, get_statsd, get_executor_pool
+from src.epstats.main import get_dao, get_executor_pool
 
-from .depend import get_test_dao, get_test_executor_pool, get_test_statsd
+from .depend import get_test_dao, get_test_executor_pool
 
 
 client = TestClient(api)
 api.dependency_overrides[get_dao] = get_test_dao
-api.dependency_overrides[get_statsd] = get_test_statsd
 api.dependency_overrides[get_executor_pool] = get_test_executor_pool