feat(server): add Prometheus metrics endpoint (#134)

This commit introduces a `/metrics` endpoint for Prometheus to scrape stream
metrics, including FPS and average FPS per stream. Additionally, it adds the
`--stream-id-label` argument, allowing users to optionally include the `stream-id`
label in Prometheus metrics.

Co-authored-by: jpotter92 <git@hjpotter92.email>

Author: rickstaa

requirements.txt

@@ -3,4 +3,5 @@ comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@ce3583ad42c024b8f060d0
 aiortc
 aiohttp
 toml
-twilio
+twilio
+prometheus_client

server/app.py

@@ -24,7 +24,9 @@
 from aiortc.rtcrtpsender import RTCRtpSender
 from pipeline import Pipeline
 from twilio.rest import Client
-from utils import FPSMeter, StreamStats, add_prefix_to_app_routes, patch_loop_datagram
+from utils import patch_loop_datagram, add_prefix_to_app_routes, FPSMeter
+from metrics import MetricsManager, StreamStatsManager
+import time
 
 logger = logging.getLogger(__name__)
 logging.getLogger("aiortc.rtcrtpsender").setLevel(logging.WARNING)
@@ -56,7 +58,9 @@ def __init__(self, track: MediaStreamTrack, pipeline: Pipeline):
         super().__init__()
         self.track = track
         self.pipeline = pipeline
-        self.fps_meter = FPSMeter()
+        self.fps_meter = FPSMeter(
+            metrics_manager=app["metrics_manager"], track_id=track.id
+        )
 
         asyncio.create_task(self.collect_frames())
 
@@ -312,6 +316,18 @@ async def on_shutdown(app: web.Application):
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set the logging level",
     )
+    parser.add_argument(
+        "--monitor",
+        default=False,
+        action="store_true",
+        help="Start a Prometheus metrics endpoint for monitoring.",
+    )
+    parser.add_argument(
+        "--stream-id-label",
+        default=False,
+        action="store_true",
+        help="Include stream ID as a label in Prometheus metrics.",
+    )
     args = parser.parse_args()
 
     logging.basicConfig(
@@ -335,11 +351,23 @@ async def on_shutdown(app: web.Application):
     app.router.add_post("/prompt", set_prompt)
 
     # Add routes for getting stream statistics.
-    stream_stats = StreamStats(app)
-    app.router.add_get("/streams/stats", stream_stats.collect_all_stream_metrics)
+    stream_stats_manager = StreamStatsManager(app)
     app.router.add_get(
-        "/stream/{stream_id}/stats", stream_stats.collect_stream_metrics_by_id
+        "/streams/stats", stream_stats_manager.collect_all_stream_metrics
     )
+    app.router.add_get(
+        "/stream/{stream_id}/stats", stream_stats_manager.collect_stream_metrics_by_id
+    )
+
+    # Add Prometheus metrics endpoint.
+    app["metrics_manager"] = MetricsManager(include_stream_id=args.stream_id_label)
+    if args.monitor:
+        app["metrics_manager"].enable()
+        logger.info(
+            f"Monitoring enabled - Prometheus metrics available at: "
+            f"http://{args.host}:{args.port}/metrics"
+        )
+        app.router.add_get("/metrics", app["metrics_manager"].metrics_handler)
 
     # Add hosted platform route prefix.
     # NOTE: This ensures that the local and hosted experiences have consistent routes.

server/metrics/__init__.py

@@ -0,0 +1,2 @@
+from .prometheus_metrics import MetricsManager
+from .stream_stats import StreamStatsManager

server/metrics/prometheus_metrics.py

@@ -0,0 +1,44 @@
+"""Prometheus metrics utilities."""
+
+from prometheus_client import Gauge, generate_latest
+from aiohttp import web
+from typing import Optional
+
+
+class MetricsManager:
+    """Manages Prometheus metrics collection."""
+
+    def __init__(self, include_stream_id: bool = False):
+        """Initializes the MetricsManager class.
+
+        Args:
+            include_stream_id: Whether to include the stream ID as a label in the metrics.
+        """
+        self._enabled = False
+        self._include_stream_id = include_stream_id
+
+        base_labels = ["stream_id"] if include_stream_id else []
+        self._fps_gauge = Gauge(
+            "stream_fps", "Frames per second of the stream", base_labels
+        )
+
+    def enable(self):
+        """Enable Prometheus metrics collection."""
+        self._enabled = True
+
+    def update_fps_metrics(self, fps: float, stream_id: Optional[str] = None):
+        """Update Prometheus metrics for a given stream.
+
+        Args:
+            fps: The current frames per second.
+            stream_id: The ID of the stream.
+        """
+        if self._enabled:
+            if self._include_stream_id:
+                self._fps_gauge.labels(stream_id=stream_id or "").set(fps)
+            else:
+                self._fps_gauge.set(fps)
+
+    async def metrics_handler(self, _):
+        """Handle Prometheus metrics endpoint."""
+        return web.Response(body=generate_latest(), content_type="text/plain")

server/metrics/stream_stats.py

@@ -0,0 +1,76 @@
+"""Handles real-time video stream statistics (non-Prometheus, JSON API)."""
+
+from typing import Any, Dict
+import json
+from aiohttp import web
+from aiortc import MediaStreamTrack
+
+
+class StreamStatsManager:
+    """Handles real-time video stream statistics collection."""
+
+    def __init__(self, app: web.Application):
+        """Initializes the StreamMetrics class.
+
+        Args:
+            app: The web application instance storing stream tracks.
+        """
+        self._app = app
+
+    async def collect_video_metrics(
+        self, video_track: MediaStreamTrack
+    ) -> Dict[str, Any]:
+        """Collects real-time statistics for a video track.
+
+        Args:
+            video_track: The video stream track instance.
+
+        Returns:
+            A dictionary containing FPS-related statistics.
+        """
+        return {
+            "timestamp": await video_track.fps_meter.last_fps_calculation_time,
+            "fps": await video_track.fps_meter.fps,
+            "minute_avg_fps": await video_track.fps_meter.average_fps,
+            "minute_fps_array": await video_track.fps_meter.fps_measurements,
+        }
+
+    async def collect_all_stream_metrics(self, _) -> web.Response:
+        """Retrieves real-time metrics for all active video streams.
+
+        Returns:
+            A JSON response containing FPS statistics for all streams.
+        """
+        video_tracks = self._app.get("video_tracks", {})
+        all_stats = {
+            stream_id: await self.collect_video_metrics(track)
+            for stream_id, track in video_tracks.items()
+        }
+
+        return web.Response(
+            content_type="application/json",
+            text=json.dumps(all_stats),
+        )
+
+    async def collect_stream_metrics_by_id(self, request: web.Request) -> web.Response:
+        """Retrieves real-time metrics for a specific video stream by ID.
+
+        Args:
+            request: The HTTP request containing the stream ID.
+
+        Returns:
+            A JSON response with stream metrics or an error message.
+        """
+        stream_id = request.match_info.get("stream_id")
+        video_tracks = self._app.get("video_tracks", {})
+        video_track = video_tracks.get(stream_id)
+
+        if video_track:
+            stats = await self.collect_video_metrics(video_track)
+        else:
+            stats = {"error": "Stream not found"}
+
+        return web.Response(
+            content_type="application/json",
+            text=json.dumps(stats),
+        )