add prometheus memory peak monitoring, combine report with analysis report

config.yaml

@@ -1,5 +1,6 @@
 service:
   logging_level: $LOGGING_LEVEL_ROOT|DEBUG  # Logging level for service logger
+  monitoring_enabled: $MONITORING_ENABLED|True # if app is doing prometheus monitoring or not
 
 probing_webserver:
   host: $PROBING_WEBSERVER_HOST|"0.0.0.0"  # Probe webserver address

pyinfra/flask.py

@@ -1,8 +1,10 @@
+import logging
+
 import requests
 from flask import Flask, jsonify
 from retry import retry
 from waitress import serve
-
+from prometheus_client import generate_latest, CollectorRegistry
 from pyinfra.config import CONFIG
 
 
@@ -24,6 +26,7 @@ def run_probing_webserver(app, host=None, port=None, mode=None):
 
 
 def set_up_probing_webserver():
+    # TODO: implement meaningful checks
     app = Flask(__name__)
 
     @app.route("/ready", methods=["GET"])
@@ -39,14 +42,21 @@ def set_up_probing_webserver():
         return resp
 
     @app.route("/prometheus", methods=["GET"])
-    def get_analysis_prometheus_endpoint():
-        @retry(requests.exceptions.ConnectionError, tries=3, delay=5, jitter=(1, 3))
-        def inner():
-            prom_endpoint = f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus"
-            metric = requests.get(prom_endpoint)
+    def prometheus():
+        analysis_metrics = get_metrics_from_analysis_endpoint()
+        combined_metrics = f"{analysis_metrics}{generate_latest(PROM_REGISTER).decode()}"
+        return combined_metrics
+
+    def get_metrics_from_analysis_endpoint():
+        try:
+            metric = requests.get(f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus")
             metric.raise_for_status()
             return metric.text
-
-        return inner()
+        except requests.exceptions.ConnectionError as err:
+            logging.warning(f"Got no metrics from analysis prometheus endpoint: {err}")
+            return "# WARNING No Analysis metrics\n"
 
     return app
+
+
+PROM_REGISTER = CollectorRegistry()

requirements.txt

@@ -7,6 +7,7 @@ waitress==2.0.0
 azure-core==1.22.1
 azure-storage-blob==12.9.0
 requests==2.27.1
+prometheus-client~=0.13.1
 # dev
 docker-compose==1.29.2
 tqdm==4.62.3

src/serve.py

@@ -1,12 +1,14 @@
 import logging
+import tracemalloc
 from multiprocessing import Process
 
 import requests
+from prometheus_client import Gauge
 from retry import retry
 
 from pyinfra.config import CONFIG, make_art
 from pyinfra.exceptions import AnalysisFailure, ConsumerError
-from pyinfra.flask import run_probing_webserver, set_up_probing_webserver
+from pyinfra.flask import run_probing_webserver, set_up_probing_webserver, PROM_REGISTER
 from pyinfra.queue.consumer import Consumer
 from pyinfra.queue.queue_manager.pika_queue_manager import PikaQueueManager
 from pyinfra.storage.storages import get_storage
@@ -37,9 +39,13 @@ def make_callback(analysis_endpoint):
 
 
 def main():
-    # TODO: implement meaningful checks
-    logging.info(make_art())
+    if CONFIG.service.monitoring_enabled:
+        tracemalloc.start()
+        memory_peak_metric = Gauge("pyinfra_memory_peak_per_call", "Pyinfra memory peak when processing a message in Mb")
+        PROM_REGISTER.register(memory_peak_metric)
+
     webserver = Process(target=run_probing_webserver, args=(set_up_probing_webserver(),))
+    logging.info(make_art())
     logging.info("Starting webserver...")
     webserver.start()
 
@@ -55,6 +61,9 @@ def main():
         try:
             consumer = Consumer(visitor, queue_manager)
             consumer.consume_and_publish()
+            if CONFIG.service.monitoring_enabled:
+                _, peak = tracemalloc.get_traced_memory()
+                memory_peak_metric.set(peak / 10 ** 6)
         except Exception as err:
             raise ConsumerError from err