api/monitoring.py

"""Implements Prometheus metrics for FastAPI application monitoring.

This is adapted from TheClouds/api/monitoring.py for the Hugging Face Space API.
"""

import os

from prometheus_client import Counter
from prometheus_fastapi_instrumentator import Instrumentator, metrics

# Namespace and Subsystem
NAMESPACE = os.environ.get("METRICS_NAMESPACE", "fastapi")
SUBSYSTEM = os.environ.get("METRICS_SUBSYSTEM", "model")

prediction_metric = Counter(
    "prediction_count",
    "Count of predictions by model, language and label",
    labelnames=["language", "model_type"],
    namespace=NAMESPACE,
    subsystem=SUBSYSTEM,
)

instrumentator = Instrumentator(
    should_group_status_codes=True,
    should_ignore_untemplated=True,
    should_instrument_requests_inprogress=True,
    excluded_handlers=["/metrics", "/docs", "/openapi.json", "/health", "/status"],
    inprogress_name="fastapi_inprogress",
    inprogress_labels=True,
)

# ----- Standard Metrics -----
# latency, request/response size, request count, errors
instrumentator.add(
    metrics.request_size(
        should_include_handler=True,
        should_include_method=True,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM,
    )
)
instrumentator.add(
    metrics.response_size(
        should_include_handler=True,
        should_include_method=True,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM,
    )
)
instrumentator.add(
    metrics.latency(
        should_include_handler=True,
        should_include_method=True,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM,
    )
)
instrumentator.add(
    metrics.requests(
        should_include_handler=True,
        should_include_method=True,
        should_include_status=True,
        metric_namespace=NAMESPACE,
        metric_subsystem=SUBSYSTEM,
    )
)