✨ Is994/resource tracking service auto populate missing data (#4420)

Author: matusdrobuliak66

services/docker-compose.yml

@@ -198,6 +198,8 @@ services:
       - PROMETHEUS_USERNAME=${RESOURCE_USAGE_TRACKER_PROMETHEUS_USERNAME}
       - PROMETHEUS_PASSWORD=${RESOURCE_USAGE_TRACKER_PROMETHEUS_PASSWORD}
       - RESOURCE_USAGE_TRACKER_LOGLEVEL=${LOG_LEVEL:-INFO}
+      - REDIS_HOST=${REDIS_HOST}
+      - REDIS_PORT=${REDIS_PORT}
 
   static-webserver:
     image: ${DOCKER_REGISTRY:-itisfoundation}/static-webserver:${DOCKER_IMAGE_TAG:-latest}

services/resource-usage-tracker/requirements/_test.in

@@ -16,6 +16,7 @@ coverage
 docker
 faker
 httpx
+fakeredis[lua]
 pytest
 pytest-asyncio
 pytest-cov

services/resource-usage-tracker/requirements/_test.txt

@@ -14,6 +14,14 @@ anyio==3.6.2
     #   httpcore
 asgi-lifespan==2.1.0
     # via -r _test.in
+async-timeout==4.0.2
+    # via
+    #   -c _base.txt
+    #   redis
+asyncpg==0.27.0
+    # via
+    #   -c _base.txt
+    #   sqlalchemy
 certifi==2023.5.7
     # via
     #   -c _base.txt
@@ -34,6 +42,8 @@ exceptiongroup==1.1.1
     # via pytest
 faker==18.10.1
     # via -r _test.in
+fakeredis[lua]==2.15.0
+    # via -r _test.in
 greenlet==2.0.2
     # via
     #   -c _base.txt
@@ -60,6 +70,8 @@ idna==3.4
     #   requests
 iniconfig==2.0.0
     # via pytest
+lupa==1.14.1
+    # via fakeredis
 mako==1.2.4
     # via
     #   -c ../../../requirements/constraints.txt
@@ -110,6 +122,11 @@ python-dotenv==1.0.0
     # via
     #   -c _base.txt
     #   -r _test.in
+redis==4.5.5
+    # via
+    #   -c ../../../requirements/constraints.txt
+    #   -c _base.txt
+    #   fakeredis
 requests==2.30.0
     # via
     #   -c _base.txt
@@ -131,7 +148,9 @@ sniffio==1.3.0
     #   asgi-lifespan
     #   httpcore
     #   httpx
-sqlalchemy[asyncio,mypy,postgresql_psycopg2binary]==1.4.48
+sortedcontainers==2.4.0
+    # via fakeredis
+sqlalchemy[mypy,postgresql_asyncpg,postgresql_psycopg2binary]==1.4.48
     # via
     #   -c ../../../requirements/constraints.txt
     #   -c _base.txt

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/core/application.py

@@ -14,6 +14,7 @@
 from ..api.routes import setup_api_routes
 from ..modules.db import setup as setup_db
 from ..modules.prometheus import setup as setup_prometheus_api_client
+from ..modules.redis import setup as setup_redis
 from ..resource_tracker import setup as setup_background_task
 from .settings import ApplicationSettings
 
@@ -46,6 +47,7 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     setup_prometheus_api_client(app)
     if settings.RESOURCE_USAGE_TRACKER_POSTGRES:
         setup_db(app)
+    setup_redis(app)
 
     setup_background_task(app)
 

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/core/settings.py

@@ -8,6 +8,7 @@
 from settings_library.basic_types import BuildTargetEnum, LogLevel, VersionTag
 from settings_library.postgres import PostgresSettings
 from settings_library.prometheus import PrometheusSettings
+from settings_library.redis import RedisSettings
 from settings_library.utils_logging import MixinLoggingSettings
 
 from .._meta import API_VERSION, API_VTAG, PROJECT_NAME
@@ -88,6 +89,8 @@ class MinimalApplicationSettings(_BaseApplicationSettings):
         auto_default_from_env=True
     )
 
+    RESOURCE_USAGE_TRACKER_REDIS: RedisSettings = Field(auto_default_from_env=True)
+
 
 class ApplicationSettings(MinimalApplicationSettings):
     """Web app's environment variables

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/modules/db/repositories/resource_tracker.py

@@ -1,7 +1,8 @@
 import logging
+from datetime import datetime
 
+import sqlalchemy as sa
 from simcore_postgres_database.models.resource_tracker import resource_tracker_container
-from sqlalchemy import func
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 
 from ....models.resource_tracker_container import ContainerResourceUsage
@@ -11,7 +12,16 @@
 
 
 class ResourceTrackerRepository(BaseRepository):
-    async def upsert_resource_tracker_container_data_(
+    async def get_prometheus_last_scraped_timestamp(self) -> datetime | None:
+        async with self.db_engine.begin() as conn:
+            max_last_scraped_timestamp: datetime | None = await conn.scalar(
+                sa.select(
+                    sa.func.max(resource_tracker_container.c.prometheus_last_scraped)
+                )
+            )
+            return max_last_scraped_timestamp
+
+    async def upsert_resource_tracker_container_data(
         self, data: ContainerResourceUsage
     ) -> None:
         async with self.db_engine.begin() as conn:
@@ -27,27 +37,27 @@ async def upsert_resource_tracker_container_data_(
                 container_cpu_usage_seconds_total=data.container_cpu_usage_seconds_total,
                 prometheus_created=data.prometheus_created.datetime,
                 prometheus_last_scraped=data.prometheus_last_scraped.datetime,
-                modified=func.now(),
+                modified=sa.func.now(),
             )
 
             on_update_stmt = insert_stmt.on_conflict_do_update(
                 index_elements=[
                     resource_tracker_container.c.container_id,
                 ],
                 set_={
-                    "container_cpu_usage_seconds_total": func.greatest(
+                    "container_cpu_usage_seconds_total": sa.func.greatest(
                         resource_tracker_container.c.container_cpu_usage_seconds_total,
                         insert_stmt.excluded.container_cpu_usage_seconds_total,
                     ),
-                    "prometheus_created": func.least(
+                    "prometheus_created": sa.func.least(
                         resource_tracker_container.c.prometheus_created,
                         insert_stmt.excluded.prometheus_created,
                     ),
-                    "prometheus_last_scraped": func.greatest(
+                    "prometheus_last_scraped": sa.func.greatest(
                         resource_tracker_container.c.prometheus_last_scraped,
                         insert_stmt.excluded.prometheus_last_scraped,
                     ),
-                    "modified": func.now(),
+                    "modified": sa.func.now(),
                 },
             )
 

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/modules/redis.py

@@ -0,0 +1,29 @@
+import logging
+from typing import cast
+
+from fastapi import FastAPI
+from servicelib.redis import RedisClientSDK
+from settings_library.redis import RedisDatabase, RedisSettings
+
+logger = logging.getLogger(__name__)
+
+
+def setup(app: FastAPI) -> None:
+    async def on_startup() -> None:
+        app.state.redis_client_sdk = None
+        settings: RedisSettings = app.state.settings.RESOURCE_USAGE_TRACKER_REDIS
+        redis_locks_dsn = settings.build_redis_dsn(RedisDatabase.LOCKS)
+        app.state.redis_client_sdk = client = RedisClientSDK(redis_locks_dsn)
+        await client.setup()
+
+    async def on_shutdown() -> None:
+        redis_client_sdk: None | RedisClientSDK = app.state.redis_client_sdk
+        if redis_client_sdk:
+            await redis_client_sdk.shutdown()
+
+    app.add_event_handler("startup", on_startup)
+    app.add_event_handler("shutdown", on_shutdown)
+
+
+def get_redis_client(app: FastAPI) -> RedisClientSDK:
+    return cast(RedisClientSDK, app.state.redis_client_sdk)

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/resource_tracker.py

@@ -3,9 +3,11 @@
 
 from fastapi import FastAPI
 from servicelib.background_task import start_periodic_task, stop_periodic_task
+from servicelib.redis_utils import exclusive
 from settings_library.prometheus import PrometheusSettings
 
 from .core.settings import ApplicationSettings
+from .modules.redis import get_redis_client
 from .resource_tracker_core import collect_container_resource_usage_task
 
 _TASK_NAME = "periodic_prometheus_polling"
@@ -23,8 +25,12 @@ async def _startup() -> None:
         if not settings:
             _logger.warning("Prometheus API client is de-activated in the settings")
             return
+        lock_key = f"{app.title}:collect_container_resource_usage_task_lock"
+        lock_value = "locked"
         app.state.resource_tracker_task = start_periodic_task(
-            collect_container_resource_usage_task,
+            exclusive(get_redis_client(app), lock_key=lock_key, lock_value=lock_value)(
+                collect_container_resource_usage_task
+            ),
             interval=app_settings.RESOURCE_USAGE_TRACKER_EVALUATION_INTERVAL_SEC,
             task_name=_TASK_NAME,
             app=app,

services/resource-usage-tracker/src/simcore_service_resource_usage_tracker/resource_tracker_core.py

@@ -1,11 +1,13 @@
 import asyncio
 import json
 import logging
+from datetime import datetime, timedelta, timezone
 from typing import Any
 
 import arrow
 from fastapi import FastAPI
 from prometheus_api_client import PrometheusConnect
+from pydantic import BaseModel
 from simcore_service_resource_usage_tracker.modules.prometheus import (
     get_prometheus_api_client,
 )
@@ -16,30 +18,58 @@
 _logger = logging.getLogger(__name__)
 
 
+_PAST_X_MINUTES = 30  # in promql query the first part: [30m:1m]
+_RESOLUTION_MINUTES = 1  # in promql query the second part: [30m:1m]
+_PROME_QUERY_PARAMS_TIMEDELTA_MINUTES = (
+    25  # should be less than _PAST_X_MINUTES (so there is small overlap)
+)
+
+
+class _PromQueryParameters(BaseModel):
+    image_regex: str
+    scrape_timestamp: datetime
+
+    class Config:
+        validation = False
+
+
 def _prometheus_sync_client_custom_query(
-    prometheus_client: PrometheusConnect, promql_cpu_query: str
+    prometheus_client: PrometheusConnect,
+    promql_cpu_query: str,
+    scrape_timestamp: datetime,
 ) -> list[dict]:
-    _logger.info("Querying prometheus with: %s", promql_cpu_query)
-    data: list[dict] = prometheus_client.custom_query(promql_cpu_query)
+    rfc3339: str = scrape_timestamp.isoformat("T")
+    _logger.info("Querying prometheus at <%s> with: <%s>", rfc3339, promql_cpu_query)
+    data: list[dict] = prometheus_client.custom_query(
+        promql_cpu_query, params={"time": rfc3339}
+    )
     return data
 
 
-async def _scrape_and_upload_container_resource_usage(
+async def _scrape_container_resource_usage(
     prometheus_client: PrometheusConnect,
-    resource_tracker_repo: ResourceTrackerRepository,
     image_regex: str,
-) -> None:
+    scrape_timestamp: datetime = datetime.now(tz=timezone.utc),
+) -> list[ContainerResourceUsage]:
     # Query CPU seconds
-    promql_cpu_query = f"sum without (cpu) (container_cpu_usage_seconds_total{{image=~'{image_regex}'}})[30m:1m]"
+    promql_cpu_query = f"sum without (cpu) (container_cpu_usage_seconds_total{{image=~'{image_regex}'}})[{_PAST_X_MINUTES}m:{_RESOLUTION_MINUTES}m]"
     containers_cpu_seconds_usage: list[
         dict
     ] = await asyncio.get_event_loop().run_in_executor(
-        None, _prometheus_sync_client_custom_query, prometheus_client, promql_cpu_query
+        None,
+        _prometheus_sync_client_custom_query,
+        prometheus_client,
+        promql_cpu_query,
+        scrape_timestamp,
     )
     _logger.info(
-        "Received %s containers from Prometheus", len(containers_cpu_seconds_usage)
+        "Received <%s> containers from Prometheus for image <%s> at <%s>",
+        len(containers_cpu_seconds_usage),
+        image_regex,
+        scrape_timestamp,
     )
 
+    data: list[ContainerResourceUsage] = []
     for item in containers_cpu_seconds_usage:
         # Prepare metric
         metric: dict[str, Any] = item["metric"]
@@ -87,21 +117,79 @@ async def _scrape_and_upload_container_resource_usage(
             prometheus_last_scraped=arrow.get(last_value[0]),
         )
 
-        await resource_tracker_repo.upsert_resource_tracker_container_data_(
-            container_resource_usage
+        data.append(container_resource_usage)
+
+    return data
+
+
+def _prepare_prom_query_parameters(
+    machine_fqdn: str,
+    prometheus_last_scraped_timestamp: datetime | None,
+    current_timestamp: datetime = datetime.now(tz=timezone.utc),
+) -> list[_PromQueryParameters]:
+    image_regex = f"registry.{machine_fqdn}/simcore/services/dynamic/jupyter-smash:.*"
+
+    # When we start this service for a first time (meaning there is empty database) we start from current timestamp
+    scrape_timestamp = (
+        prometheus_last_scraped_timestamp
+        if prometheus_last_scraped_timestamp
+        else current_timestamp
+    )
+
+    data = []
+    while scrape_timestamp < current_timestamp - timedelta(
+        minutes=_PROME_QUERY_PARAMS_TIMEDELTA_MINUTES
+    ):
+        data.append(
+            _PromQueryParameters(
+                image_regex=image_regex, scrape_timestamp=scrape_timestamp
+            )
+        )
+        scrape_timestamp = scrape_timestamp + timedelta(
+            minutes=_PROME_QUERY_PARAMS_TIMEDELTA_MINUTES
         )
+    data.append(
+        _PromQueryParameters(
+            image_regex=image_regex, scrape_timestamp=current_timestamp
+        )
+    )
+    return data
 
 
 async def collect_container_resource_usage(
     prometheus_client: PrometheusConnect,
     resource_tracker_repo: ResourceTrackerRepository,
     machine_fqdn: str,
 ) -> None:
-    await _scrape_and_upload_container_resource_usage(
-        prometheus_client=prometheus_client,
-        resource_tracker_repo=resource_tracker_repo,
-        image_regex=f"registry.{machine_fqdn}/simcore/services/dynamic/jupyter-smash:.*",
+    prometheus_last_scraped_timestamp: datetime | None = (
+        await resource_tracker_repo.get_prometheus_last_scraped_timestamp()
     )
+    prom_query_params: list[_PromQueryParameters] = _prepare_prom_query_parameters(
+        machine_fqdn, prometheus_last_scraped_timestamp
+    )
+
+    for i, parameter in enumerate(prom_query_params):
+        _logger.info(
+            "Collecting %s/%s with parameter <%s>",
+            i + 1,
+            len(prom_query_params),
+            parameter,
+        )
+        data: list[ContainerResourceUsage] = await _scrape_container_resource_usage(
+            prometheus_client=prometheus_client,
+            image_regex=parameter.image_regex,
+            scrape_timestamp=parameter.scrape_timestamp,
+        )
+
+        # Upload to the database
+        await asyncio.gather(
+            *[
+                resource_tracker_repo.upsert_resource_tracker_container_data(
+                    container_resource_usage
+                )
+                for container_resource_usage in data
+            ]
+        )
 
 
 async def collect_container_resource_usage_task(app: FastAPI) -> None:

services/resource-usage-tracker/tests/unit/api/test__oas_spec.py

@@ -11,6 +11,7 @@
 def test_openapi_json_is_in_sync_with_app_oas(
     disabled_database: None,
     disabled_prometheus: None,
+    mocked_redis_server: None,
     client: TestClient,
     project_slug_dir: Path,
 ):