diff --git a/config.yaml b/config.yaml index ceb12a0..e3de209 100755 --- a/config.yaml +++ b/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 diff --git a/pyinfra/flask.py b/pyinfra/flask.py index 81e8b94..0a91bc3 100644 --- a/pyinfra/flask.py +++ b/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() diff --git a/requirements.txt b/requirements.txt index c13ae24..e5d78e9 100755 --- a/requirements.txt +++ b/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 diff --git a/src/serve.py b/src/serve.py index 7373e57..048999f 100644 --- a/src/serve.py +++ b/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