feat: Integrate with Prometheus and Grafana

Introduce a basic observability stack using Prometheus for metrics
collection and Grafana for visualisation. Prometheus is configured to
scrape metrics from the CogStack Model Gateway, Ripper, and Scheduler,
while Grafana is set up with some basic dashboards for each service. At
this stage, our instrumentation (and visualisations) focus on sheding
some light on the total number of requests received and processed, along
with information about their outcome (success or failure), the model and
task requested. In the future, we should expand our set up for more
detailed insights such as the time taken to process requests. Relevant
metadata are already provided as labels in the collected metrics (e.g.
'created_at', 'started_at', and 'finished_at' timestamps for each task),
so we can easily extend our dashboards to integrate these.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

README.md

@@ -171,4 +171,3 @@ The project is still under active development. In the future we will be focusing
 * **Smart scheduling**: Implement a more sophisticated scheduling algorithm that takes into account
   the state of the cluster.
 * **CI/CD**: Set up a continuous integration and deployment pipeline for the project.
-* **Monitoring**: Integrate with Prometheus and Grafana.

cogstack_model_gateway/common/config.py

@@ -30,7 +30,9 @@
         "CMG_QUEUE_NAME": "cmg_tasks",
         "CMG_QUEUE_HOST": "rabbitmq",
         "CMG_QUEUE_PORT": "5672",
+        "CMG_SCHEDULER_METRICS_PORT": "8001",
         "CMG_SCHEDULER_MAX_CONCURRENT_TASKS": "1",
+        "CMG_RIPPER_METRICS_PORT": "8002",
     },
 }
 

cogstack_model_gateway/gateway/Dockerfile

@@ -19,6 +19,7 @@ RUN curl -sSL https://install.python-poetry.org | python3 -
 
 WORKDIR /app
 
+COPY cogstack_model_gateway/gateway/gunicorn.conf.py ./
 COPY pyproject.toml poetry.lock ./
 RUN poetry install --only main --no-root --no-directory
 
@@ -30,4 +31,4 @@ RUN poetry install --only main
 
 EXPOSE 8000
 
-CMD ["poetry", "run", "fastapi", "run", "--workers", "4", "cogstack_model_gateway/gateway/main.py"]
+CMD ["poetry", "run", "gunicorn", "-c", "gunicorn.conf.py", "cogstack_model_gateway.gateway.main:app"]

cogstack_model_gateway/gateway/gunicorn.conf.py

@@ -0,0 +1,16 @@
+from prometheus_client import multiprocess
+
+bind = "0.0.0.0:8000"
+
+workers = 4
+
+worker_class = "uvicorn.workers.UvicornWorker"
+
+
+def child_exit(server, worker):
+    """Mark the Prometheus metrics for this worker as dead before a worker exits.
+
+    This function is called by Gunicorn when a worker process exits. It marks the process as dead in
+    the Prometheus multiprocess registry, allowing MultiProcessCollector to ignore its old files.
+    """
+    multiprocess.mark_process_dead(worker.pid)

cogstack_model_gateway/gateway/main.py

@@ -1,20 +1,36 @@
 import logging
+import os
 from contextlib import asynccontextmanager
 
 import urllib3
 from fastapi import FastAPI
+from prometheus_client import CollectorRegistry, make_asgi_app, multiprocess
 
 from cogstack_model_gateway.common.config import load_config
 from cogstack_model_gateway.common.db import DatabaseManager
 from cogstack_model_gateway.common.logging import configure_logging
 from cogstack_model_gateway.common.object_store import ObjectStoreManager
 from cogstack_model_gateway.common.queue import QueueManager
 from cogstack_model_gateway.common.tasks import TaskManager
+from cogstack_model_gateway.gateway.prometheus.metrics import gateway_requests_total
 from cogstack_model_gateway.gateway.routers import models, tasks
 
 log = logging.getLogger("cmg.gateway")
 
 
+def make_metrics_app():
+    """Create a registry for each process and aggregate metrics with MultiProcessCollector."""
+    if "PROMETHEUS_MULTIPROC_DIR" not in os.environ:
+        raise RuntimeError(
+            "Environment variable PROMETHEUS_MULTIPROC_DIR is not set. Please set it to a directory"
+            " where Prometheus can store metrics data for multiprocess mode, e.g. /tmp/prometheus/"
+        )
+    os.makedirs(os.environ["PROMETHEUS_MULTIPROC_DIR"], exist_ok=True)
+    registry = CollectorRegistry()
+    multiprocess.MultiProcessCollector(registry)
+    return make_asgi_app(registry=registry)
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Setup gateway and initialize database, object store, queue, and task manager connections."""
@@ -73,6 +89,15 @@ async def lifespan(app: FastAPI):
 app.include_router(models.router)
 app.include_router(tasks.router)
 
+app.mount("/metrics", make_metrics_app())
+
+
+@app.middleware("http")
+async def prometheus_request_counter(request, call_next):
+    response = await call_next(request)
+    gateway_requests_total.labels(method=request.method, endpoint=request.url.path).inc()
+    return response
+
 
 @app.get("/")
 async def root():

cogstack_model_gateway/gateway/prometheus/metrics.py

@@ -0,0 +1,19 @@
+from prometheus_client import Counter
+
+gateway_requests_total = Counter(
+    "gateway_requests_total",
+    "Total number of HTTP requests to the CogStack Model Gateway",
+    ["method", "endpoint"],
+)
+
+gateway_models_deployed_total = Counter(
+    "gateway_models_deployed_total",
+    "Total number of models deployed via the CogStack Model Gateway",
+    ["model", "model_uri"],
+)
+
+gateway_tasks_processed_total = Counter(
+    "gateway_tasks_processed_total",
+    "Total number of tasks processed through the CogStack Model Gateway",
+    ["model", "task"],
+)

cogstack_model_gateway/gateway/routers/models.py

@@ -14,6 +14,10 @@
 from cogstack_model_gateway.common.tracking import TrackingClient
 from cogstack_model_gateway.gateway.core.models import get_running_models, run_model_container
 from cogstack_model_gateway.gateway.core.priority import calculate_task_priority
+from cogstack_model_gateway.gateway.prometheus.metrics import (
+    gateway_models_deployed_total,
+    gateway_tasks_processed_total,
+)
 from cogstack_model_gateway.gateway.routers.utils import (
     get_cms_url,
     get_content_type,
@@ -133,6 +137,7 @@ async def get_models(
 )
 async def get_model_info(model_name: str):
     """Get information about a running model server through its `/info` API."""
+    gateway_tasks_processed_total.labels(model=model_name, task="info").inc()
     # FIXME: Enable SSL verification when certificates are properly set up
     response = requests.get(get_cms_url(model_name, "info"), verify=False)
     if response.status_code == 404:
@@ -220,6 +225,8 @@ async def deploy_model(
             status_code=500, detail=f"Failed to deploy model '{model_name}': {str(e)}"
         )
 
+    gateway_models_deployed_total.labels(model=model_name, model_uri=model_uri).inc()
+
     log.info(f"Model '{model_name}' deployed successfully with container ID {container.id}")
     return {
         "message": f"Model '{model_name}' deployed successfully",
@@ -335,4 +342,6 @@ async def execute_task(
     qm: QueueManager = config.queue_manager
     qm.publish(task_dict, priority)
 
+    gateway_tasks_processed_total.labels(model=model_name, task=task).inc()
+
     return {"uuid": task_uuid, "status": "Task submitted successfully"}

cogstack_model_gateway/ripper/main.py

@@ -8,6 +8,7 @@
 import docker
 from dateutil import parser
 from docker.models.containers import Container
+from prometheus_client import start_http_server
 
 from cogstack_model_gateway.common.containers import (
     IS_MODEL_LABEL,
@@ -16,8 +17,10 @@
     TTL_LABEL,
 )
 from cogstack_model_gateway.common.logging import configure_logging
+from cogstack_model_gateway.ripper.prometheus.metrics import containers_purged_total
 
 PURGE_INTERVAL = int(os.getenv("CMG_RIPPER_INTERVAL") or 60)
+METRICS_PORT = int(os.getenv("CMG_RIPPER_METRICS_PORT") or 8002)
 
 log = logging.getLogger("cmg.ripper")
 
@@ -27,6 +30,7 @@ def stop_and_remove_container(container: Container):
     log.info(f"Stopping and removing expired container: {container.name}")
     container.stop()
     container.remove()
+    containers_purged_total.inc()
 
 
 def purge_expired_containers():
@@ -71,6 +75,9 @@ def purge_expired_containers():
 def main():
     """Run the ripper service."""
     configure_logging()
+
+    start_http_server(METRICS_PORT)
+
     purge_expired_containers()
 
 

cogstack_model_gateway/ripper/prometheus/metrics.py

@@ -0,0 +1,6 @@
+from prometheus_client import Counter
+
+containers_purged_total = Counter(
+    "ripper_containers_purged_total",
+    "Total number of Docker containers purged by the Ripper",
+)

cogstack_model_gateway/scheduler/main.py

@@ -1,7 +1,9 @@
 import logging
 import sys
 
-from cogstack_model_gateway.common.config import load_config
+from prometheus_client import start_http_server
+
+from cogstack_model_gateway.common.config import Config, load_config
 from cogstack_model_gateway.common.db import DatabaseManager
 from cogstack_model_gateway.common.logging import configure_logging
 from cogstack_model_gateway.common.object_store import ObjectStoreManager
@@ -12,12 +14,11 @@
 log = logging.getLogger("cmg.scheduler")
 
 
-def initialize_connections() -> tuple[
-    DatabaseManager, ObjectStoreManager, QueueManager, TaskManager
-]:
+def initialize_connections(
+    config: Config,
+) -> tuple[DatabaseManager, ObjectStoreManager, QueueManager, TaskManager]:
     """Initialize database, object store, queue, and task manager connections for the scheduler."""
     log.info("Initializing database and queue connections")
-    config = load_config()
     dbm = DatabaseManager(
         user=config.cmg.db_user,
         password=config.cmg.db_password,
@@ -65,7 +66,10 @@ def main():
     urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
     configure_logging()
-    connections = initialize_connections()
+    config = load_config()
+    connections = initialize_connections(config)
+
+    start_http_server(int(config.cmg.scheduler_metrics_port))
 
     scheduler = Scheduler(
         task_object_store_manager=connections[1],