feat(api): add prometheus and grafana (#3)

* docs: add contribution instructions
* chore: add pre-commit linting and checking poetry
* feat(api): add prometheus and grafana

There is a decorator to translate API that emits metrics of /translate
request count and response time that tags source and destination language

Author: hspedro

.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+    -   id: trailing-whitespace
+    -   id: end-of-file-fixer
+    -   id: check-yaml
+    -   id: check-added-large-files
+    -   id: check-merge-conflict
+
+-   repo: local
+    hooks:
+    -   id: poetry-check
+        name: Check if Poetry lock file is up to date
+        entry: bash -c 'poetry check && poetry lock'
+        language: system
+        pass_filenames: false
+        files: ^pyproject\.toml$
+
+    -   id: lint
+        name: Run linting checks
+        entry: bash -c 'make lint'
+        language: system
+        pass_filenames: false
+        types: [python]

babeltron/app/main.py

@@ -3,13 +3,14 @@
 from importlib.metadata import version
 from typing import AsyncIterator
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Response
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi_cache import FastAPICache
 from fastapi_cache.backends.inmemory import InMemoryBackend
 from fastapi_cache.backends.redis import RedisBackend
 from redis import asyncio as aioredis
 
+from babeltron.app.monitoring import PrometheusMiddleware, metrics_endpoint
 from babeltron.app.utils import include_routers
 
 try:
@@ -20,7 +21,7 @@
 
 @asynccontextmanager
 async def lifespan(_: FastAPI) -> AsyncIterator[None]:
-    cache_url = os.environ.get("CACHE_URL")
+    cache_url = os.environ.get("CACHE_URL", "")
 
     if cache_url.startswith("in-memory"):
         FastAPICache.init(InMemoryBackend(), prefix="babeltron")
@@ -40,9 +41,9 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
     description="API for machine translation using NLLB models",
     version="0.1.0",
     contact={
-        "name": "Your Name",
-        "url": "https://your-website.com",
-        "email": "your-email@example.com",
+        "name": "Pedro Soares",
+        "url": "https://github.com/hspedro",
+        "email": "pedrofigueiredoc@gmail.com",
     },
     license_info={
         "name": "MIT",
@@ -66,6 +67,16 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
 # Include all routers
 include_routers(app)
 
+# Add Prometheus middleware
+app.add_middleware(PrometheusMiddleware)
+
+
+# Add metrics endpoint
+@app.get("/metrics", include_in_schema=False)
+async def metrics():
+    return Response(content=metrics_endpoint(), media_type="text/plain")
+
+
 if __name__ == "__main__":
     import uvicorn
 

babeltron/app/monitoring.py

@@ -0,0 +1,173 @@
+import time
+
+from fastapi import Request
+from prometheus_client import CollectorRegistry, Counter, Histogram, generate_latest
+from starlette.middleware.base import BaseHTTPMiddleware
+
+# Create a registry
+registry = CollectorRegistry()
+
+# Define metrics
+REQUEST_COUNT = Counter(
+    "http_requests_total",
+    "Total count of HTTP requests",
+    ["method", "endpoint", "status_code"],
+    registry=registry,
+)
+
+REQUEST_LATENCY = Histogram(
+    "http_request_duration_seconds",
+    "HTTP request latency in seconds",
+    ["method", "endpoint"],
+    buckets=(
+        0.01,
+        0.025,
+        0.05,
+        0.075,
+        0.1,
+        0.25,
+        0.5,
+        0.75,
+        1.0,
+        2.5,
+        5.0,
+        7.5,
+        10.0,
+        25.0,
+        50.0,
+    ),
+    registry=registry,
+)
+
+ERROR_COUNT = Counter(
+    "http_request_errors_total",
+    "Total count of HTTP request errors",
+    ["method", "endpoint", "exception_type"],
+    registry=registry,
+)
+
+TRANSLATION_COUNT = Counter(
+    "translation_requests_total",
+    "Total count of translation requests",
+    ["src_lang", "tgt_lang"],
+    registry=registry,
+)
+
+TRANSLATION_LATENCY = Histogram(
+    "translation_duration_seconds",
+    "Translation processing time in seconds",
+    ["src_lang", "tgt_lang"],
+    buckets=(
+        0.05,
+        0.1,
+        0.25,
+        0.5,
+        0.75,
+        1.0,
+        1.5,
+        2.0,
+        2.5,
+        3.0,
+        3.5,
+        4,
+        4.5,
+        5.0,
+        7.5,
+        10.0,
+        15.0,
+        20.0,
+        30.0,
+        60.0,
+    ),
+    registry=registry,
+)
+
+CACHE_HIT_COUNT = Counter(
+    "cache_hits_total", "Total count of cache hits", ["endpoint"], registry=registry
+)
+
+CACHE_MISS_COUNT = Counter(
+    "cache_misses_total", "Total count of cache misses", ["endpoint"], registry=registry
+)
+
+MODEL_LOAD_TIME = Histogram(
+    "model_load_time_seconds",
+    "Time taken to load the model",
+    ["model_size"],
+    buckets=(0.1, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0),
+    registry=registry,
+)
+
+
+class PrometheusMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app=None):
+        super().__init__(app)
+
+    async def dispatch(self, request: Request, call_next):
+        start_time = time.time()
+
+        # Get the route path for the request
+        route = request.url.path
+        method = request.method
+
+        try:
+            response = await call_next(request)
+
+            # Record request count and latency
+            REQUEST_COUNT.labels(
+                method=method, endpoint=route, status_code=response.status_code
+            ).inc()
+            REQUEST_LATENCY.labels(method=method, endpoint=route).observe(
+                time.time() - start_time
+            )
+
+            # Record error if status code is 4xx or 5xx
+            if 400 <= response.status_code < 600:
+                ERROR_COUNT.labels(
+                    method=method,
+                    endpoint=route,
+                    exception_type=f"HTTP{response.status_code}",
+                ).inc()
+
+            return response
+
+        except Exception as e:
+            # Record exception
+            ERROR_COUNT.labels(
+                method=method, endpoint=route, exception_type=type(e).__name__
+            ).inc()
+            raise
+
+
+def track_dynamic_translation_metrics():
+    def decorator(func):
+        from functools import wraps
+
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            request = kwargs.get("request")
+            if not request and args:
+                request = args[0]
+
+            start_time = time.time()
+            src_lang = request.src_lang
+            tgt_lang = request.tgt_lang
+
+            TRANSLATION_COUNT.labels(src_lang=src_lang, tgt_lang=tgt_lang).inc()
+
+            result = await func(*args, **kwargs)
+
+            TRANSLATION_LATENCY.labels(src_lang=src_lang, tgt_lang=tgt_lang).observe(
+                time.time() - start_time
+            )
+
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def metrics_endpoint():
+    """Generate latest metrics in Prometheus format"""
+    return generate_latest(registry)

babeltron/app/routers/translate.py

@@ -6,6 +6,7 @@
 from pydantic import BaseModel, Field
 from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 
+from babeltron.app.monitoring import track_dynamic_translation_metrics
 from babeltron.app.utils import ORJsonCoder, cache_key_builder, get_model_path
 
 router = APIRouter(tags=["Translation"])
@@ -77,6 +78,7 @@ class TranslationResponse(BaseModel):
     status_code=status.HTTP_200_OK,
 )
 @cache(expire=CACHE_TTL_SECONDS, key_builder=cache_key_builder, coder=ORJsonCoder)
+@track_dynamic_translation_metrics()
 async def translate(request: TranslationRequest):
     if model is None or tokenizer is None:
         raise HTTPException(

docker-compose.yml

@@ -12,10 +12,42 @@ services:
     environment:
       - MODEL_PATH=/models
       - CACHE_URL=in-memory
+      - METRICS_ENABLED=true
     restart: unless-stopped
     healthcheck:
       test: ["CMD", "curl", "-f", "http://localhost:8000/healthz"]
       interval: 30s
       timeout: 10s
       retries: 3
       start_period: 40s
+
+  prometheus:
+    image: prom/prometheus:latest
+    container_name: prometheus
+    ports:
+      - "9090:9090"
+    volumes:
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--web.console.libraries=/usr/share/prometheus/console_libraries'
+      - '--web.console.templates=/usr/share/prometheus/consoles'
+    depends_on:
+      - api
+
+  grafana:
+    image: grafana/grafana:latest
+    container_name: grafana
+    ports:
+      - "3000:3000"
+    environment:
+      - GF_SECURITY_ADMIN_USER=admin
+      - GF_SECURITY_ADMIN_PASSWORD=admin
+    volumes:
+      - grafana-storage:/var/lib/grafana
+    depends_on:
+      - prometheus
+
+volumes:
+  grafana-storage:

docs/CONTRIBUTING.md

@@ -0,0 +1,82 @@
+# Contributing
+
+We welcome contributions to Babeltron! This document provides guidelines for setting up your development environment and ensuring code quality.
+
+## Development Setup
+
+1. Clone the repository:
+```bash
+git clone https://github.com/hspedro/babeltron.git
+cd babeltron
+```
+
+2. Install dependencies:
+```bash
+make install
+```
+
+3. Install pre-commit hooks:
+```bash
+make pre-commit-install
+```
+
+## Pre-commit Hooks
+
+Babeltron uses pre-commit hooks to ensure code quality and consistency. These hooks run automatically when you commit changes and check for:
+
+- Outdated Poetry lock files
+- Linting issues
+- Trailing whitespace and file formatting issues
+- YAML syntax errors
+- Merge conflicts
+
+To run the pre-commit hooks manually on all files:
+
+```bash
+pre-commit run --all-files
+```
+
+## Code Style
+
+Babeltron follows these code style guidelines:
+
+- [Black](https://github.com/psf/black) for code formatting
+- [isort](https://pycqa.github.io/isort/) for import sorting
+
+You can run these checks manually with:
+
+```bash
+make lint
+```
+
+## Testing
+
+Before submitting a pull request, make sure all tests pass:
+
+```bash
+make test
+```
+
+To run tests with coverage reporting:
+
+```bash
+make coverage
+```
+
+## Pull Request Process
+
+1. Fork the repository
+2. Create a feature branch (`git checkout -b feature/amazing-feature`)
+3. Make your changes
+4. Run tests and linting checks
+5. Commit your changes (`git commit -m 'Add amazing feature'`)
+6. Push to the branch (`git push origin feature/amazing-feature`)
+7. Open a Pull Request
+
+## Running locally
+
+```bash
+make serve
+```
+
+Will mostly do the trick since it has auto-reload.