feat(server): add Prometheus metrics endpoint

This commit adds a Prometheus metrics endpoint (`/metrics`) that allows Prometheus
to scrape stream metrics, such as FPS andaverage FPS per stream.

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

@@ -25,12 +25,13 @@
 from aiortc.rtcrtpsender import RTCRtpSender
 from aiortc.codecs import h264
 from pipeline import Pipeline
-from utils import patch_loop_datagram, StreamStats, add_prefix_to_app_routes
+from utils import patch_loop_datagram, add_prefix_to_app_routes
+from metrics import MetricsManager, StreamStatsManager
 import time
 
 logger = logging.getLogger(__name__)
-logging.getLogger('aiortc.rtcrtpsender').setLevel(logging.WARNING)
-logging.getLogger('aiortc.rtcrtpreceiver').setLevel(logging.WARNING)
+logging.getLogger("aiortc.rtcrtpsender").setLevel(logging.WARNING)
+logging.getLogger("aiortc.rtcrtpreceiver").setLevel(logging.WARNING)
 
 
 MAX_BITRATE = 2000000
@@ -45,7 +46,9 @@ class VideoStreamTrack(MediaStreamTrack):
         track (MediaStreamTrack): The underlying media stream track.
         pipeline (Pipeline): The processing pipeline to apply to each video frame.
     """
+
     kind = "video"
+
     def __init__(self, track: MediaStreamTrack, pipeline: Pipeline):
         """Initialize the VideoStreamTrack.
 
@@ -63,6 +66,7 @@ def __init__(self, track: MediaStreamTrack, pipeline: Pipeline):
         self._fps_loop_start_time = time.monotonic()
         self._fps = 0.0
         self._fps_measurements = deque(maxlen=60)
+        self._average_fps = 0.0
         self._running_event = asyncio.Event()
 
         asyncio.create_task(self.collect_frames())
@@ -88,18 +92,36 @@ async def _calculate_fps_loop(self):
                 current_time = time.monotonic()
                 if self._last_fps_calculation_time is not None:
                     time_diff = current_time - self._last_fps_calculation_time
-                    self._fps = self._fps_interval_frame_count / time_diff
+                    self._fps = (
+                        self._fps_interval_frame_count / time_diff
+                        if time_diff > 0
+                        else 0.0
+                    )
                     self._fps_measurements.append(
                         {
                             "timestamp": current_time - self._fps_loop_start_time,
                             "fps": self._fps,
                         }
                     )  # Store the FPS measurement with timestamp
 
-                # Reset start_time and frame_count for the next interval.
+                    # Store the average FPS over the last minute.
+                    self._average_fps = (
+                        sum(m["fps"] for m in self._fps_measurements)
+                        / len(self._fps_measurements)
+                        if self._fps_measurements
+                        else self._fps
+                    )
+
+                # Reset tracking variables for the next interval.
                 self._last_fps_calculation_time = current_time
                 self._fps_interval_frame_count = 0
-            await asyncio.sleep(1)  # Calculate FPS every second.
+
+            # Update Prometheus metrics if enabled.
+            app["metrics_manager"].update_metrics(
+                self.track.id, self._fps, self._average_fps
+            )
+
+            await asyncio.sleep(1)  # Calculate FPS every second
 
     @property
     async def fps(self) -> float:
@@ -129,11 +151,7 @@ async def average_fps(self) -> float:
             The average FPS over the last minute.
         """
         async with self._lock:
-            if not self._fps_measurements:
-                return 0.0
-            return sum(
-                measurement["fps"] for measurement in self._fps_measurements
-            ) / len(self._fps_measurements)
+            return self._average_fps
 
     @property
     async def last_fps_calculation_time(self) -> float:
@@ -159,6 +177,7 @@ async def recv(self):
 
 class AudioStreamTrack(MediaStreamTrack):
     kind = "audio"
+
     def __init__(self, track: MediaStreamTrack, pipeline):
         super().__init__()
         self.track = track
@@ -257,30 +276,29 @@ async def offer(request):
     @pc.on("datachannel")
     def on_datachannel(channel):
         if channel.label == "control":
+
             @channel.on("message")
             async def on_message(message):
                 try:
                     params = json.loads(message)
 
                     if params.get("type") == "get_nodes":
                         nodes_info = await pipeline.get_nodes_info()
-                        response = {
-                            "type": "nodes_info",
-                            "nodes": nodes_info
-                        }
+                        response = {"type": "nodes_info", "nodes": nodes_info}
                         channel.send(json.dumps(response))
                     elif params.get("type") == "update_prompts":
                         if "prompts" not in params:
-                            logger.warning("[Control] Missing prompt in update_prompt message")
+                            logger.warning(
+                                "[Control] Missing prompt in update_prompt message"
+                            )
                             return
                         await pipeline.update_prompts(params["prompts"])
-                        response = {
-                            "type": "prompts_updated",
-                            "success": True
-                        }
+                        response = {"type": "prompts_updated", "success": True}
                         channel.send(json.dumps(response))
                     else:
-                        logger.warning("[Server] Invalid message format - missing required fields")
+                        logger.warning(
+                            "[Server] Invalid message format - missing required fields"
+                        )
                 except json.JSONDecodeError:
                     logger.error("[Server] Invalid JSON received")
                 except Exception as e:
@@ -385,12 +403,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.",
+    )
     args = parser.parse_args()
 
     logging.basicConfig(
         level=args.log_level.upper(),
-        format='%(asctime)s [%(levelname)s] %(message)s',
-        datefmt='%H:%M:%S'
+        format="%(asctime)s [%(levelname)s] %(message)s",
+        datefmt="%H:%M:%S",
     )
 
     app = web.Application()
@@ -408,11 +432,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()
+    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,42 @@
+"""Prometheus metrics utilities."""
+
+from prometheus_client import Gauge, generate_latest
+from aiohttp import web
+
+
+class MetricsManager:
+    """Manages Prometheus metrics collection.
+
+    Attributes:
+        fps_gauge: A Prometheus gauge for the current frames per second.
+        avg_fps_gauge: A Prometheus gauge for the average frames per second per minute.
+    """
+
+    def __init__(self):
+        self._enabled = False
+        self._fps_gauge = Gauge(
+            "stream_fps", "Frames per second of the stream", ["stream_id"]
+        )
+        self._avg_fps_gauge = Gauge(
+            "stream_avg_fps", "Average frames per second per minute", ["stream_id"]
+        )
+
+    def enable(self):
+        """Enable Prometheus metrics collection."""
+        self._enabled = True
+
+    def update_metrics(self, stream_id: str, fps: float, avg_fps: float):
+        """Update Prometheus metrics for a given stream.
+
+        Args:
+            stream_id: The ID of the stream.
+            fps: The current frames per second.
+            avg_fps: The average frames per second per minute.
+        """
+        if self._enabled:
+            self._fps_gauge.labels(stream_id=stream_id).set(fps)
+            self._avg_fps_gauge.labels(stream_id=stream_id).set(avg_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:
+    """Manages real-time video stream statistics retrieval."""
+
+    def __init__(self, app: web.Application):
+        """Initializes the StreamMetrics class.
+
+        Args:
+            app: The web application instance storing video streams under the
+                "video_tracks" key.
+        """
+        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.last_fps_calculation_time,
+            "fps": await video_track.fps,
+            "minute_avg_fps": await video_track.average_fps,
+            "minute_fps_array": await video_track.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_track = self._app["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),
+        )