diff --git a/setup.py b/setup.py index 6d109949..ad66a2e9 100644 --- a/setup.py +++ b/setup.py @@ -43,6 +43,7 @@ setuptools.setup( "anthropic", "google-cloud-aiplatform", "pyyaml", + "prometheus-client", ], scripts=[ "scripts/chunker-recursive", diff --git a/trustgraph/base/processor.py b/trustgraph/base/processor.py index 4f035966..f7012921 100644 --- a/trustgraph/base/processor.py +++ b/trustgraph/base/processor.py @@ -2,8 +2,10 @@ import os import argparse import pulsar +import _pulsar import time from pulsar.schema import JsonSchema +from prometheus_client import start_http_server, Histogram, Info, Counter from .. log_level import LogLevel @@ -11,16 +13,23 @@ class BaseProcessor: default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650') - def __init__( - self, - pulsar_host=default_pulsar_host, - log_level=LogLevel.INFO, - ): + def __init__(self, **params): self.client = None - if pulsar_host == None: - pulsar_host = default_pulsar_host + if not hasattr(__class__, "params_metric"): + __class__.params_metric = Info( + 'params', 'Parameters configuration' + ) + + # FIXME: Maybe outputs information it should not + __class__.params_metric.info({ + k: str(params[k]) + for k in params + }) + + pulsar_host = params.get("pulsar_host", self.default_pulsar_host) + log_level = params.get("log_level", LogLevel.INFO) self.pulsar_host = pulsar_host @@ -51,6 +60,20 @@ class BaseProcessor: help=f'Output queue (default: info)' ) + parser.add_argument( + '-M', '--metrics-enabled', + type=bool, + default=True, + help=f'Pulsar host (default: true)', + ) + + parser.add_argument( + '-P', '--metrics-port', + type=int, + default=8000, + help=f'Pulsar host (default: 8000)', + ) + def run(self): raise RuntimeError("Something should have implemented the run method") @@ -69,13 +92,26 @@ class BaseProcessor: args = parser.parse_args() args = vars(args) + if args["metrics_enabled"]: + start_http_server(args["metrics_port"]) + try: p = cls(**args) p.run() + except KeyboardInterrupt: + print("Keyboard interrupt.") + return + + except _pulsar.Interrupted: + print("Pulsar Interrupted.") + return + except Exception as e: + print(type(e)) + print("Exception:", e, flush=True) print("Will retry...", flush=True) @@ -83,19 +119,13 @@ class BaseProcessor: class Consumer(BaseProcessor): - def __init__( - self, - pulsar_host=None, - log_level=LogLevel.INFO, - input_queue="input", - subscriber="subscriber", - input_schema=None, - ): + def __init__(self, **params): - super(Consumer, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - ) + super(Consumer, self).__init__(**params) + + input_queue = params.get("input_queue") + subscriber = params.get("subscriber") + input_schema = params.get("input_schema") if input_schema == None: raise RuntimeError("input_schema must be specified") @@ -144,21 +174,38 @@ class Consumer(BaseProcessor): class ConsumerProducer(BaseProcessor): - def __init__( - self, - pulsar_host=None, - log_level=LogLevel.INFO, - input_queue="input", - output_queue="output", - subscriber="subscriber", - input_schema=None, - output_schema=None, - ): + def __init__(self, **params): - super(ConsumerProducer, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - ) + input_queue = params.get("input_queue") + output_queue = params.get("output_queue") + subscriber = params.get("subscriber") + input_schema = params.get("input_schema") + output_schema = params.get("output_schema") + + if not hasattr(__class__, "request_metric"): + __class__.request_metric = Histogram( + 'request_latency', 'Request latency (seconds)' + ) + + if not hasattr(__class__, "pubsub_metric"): + __class__.pubsub_metric = Info( + 'pubsub', 'Pub/sub configuration' + ) + + if not hasattr(__class__, "processing_metric"): + __class__.processing_metric = Counter( + 'processing_count', 'Processing count', ["status"] + ) + + __class__.pubsub_metric.info({ + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": input_schema.__name__, + "output_schema": output_schema.__name__, + }) + + super(ConsumerProducer, self).__init__(**params) if input_schema == None: raise RuntimeError("input_schema must be specified") @@ -184,11 +231,14 @@ class ConsumerProducer(BaseProcessor): try: - resp = self.handle(msg) + with __class__.request_metric.time(): + resp = self.handle(msg) # Acknowledge successful processing of the message self.consumer.acknowledge(msg) + __class__.processing_metric.labels(status="success").inc() + except Exception as e: print("Exception:", e, flush=True) @@ -196,6 +246,8 @@ class ConsumerProducer(BaseProcessor): # Message failed to be processed self.consumer.negative_acknowledge(msg) + __class__.processing_metric.labels(status="error").inc() + def send(self, msg, properties={}): self.producer.send(msg, properties) @@ -228,18 +280,12 @@ class ConsumerProducer(BaseProcessor): class Producer(BaseProcessor): - def __init__( - self, - pulsar_host=None, - log_level=LogLevel.INFO, - output_queue="output", - output_schema=None, - ): + def __init__(self, **params): - super(Producer, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - ) + output_queue = params.get("output_queue") + output_schema = params.get("output_schema") + + super(Producer, self).__init__(**params) if output_schema == None: raise RuntimeError("output_schema must be specified") diff --git a/trustgraph/llm/ollama_text/llm.py b/trustgraph/llm/ollama_text/llm.py index b522e225..3d256137 100755 --- a/trustgraph/llm/ollama_text/llm.py +++ b/trustgraph/llm/ollama_text/llm.py @@ -5,6 +5,7 @@ Input is prompt, output is response. """ from langchain_community.llms import Ollama +from prometheus_client import Histogram, Info, Counter from ... schema import TextCompletionRequest, TextCompletionResponse from ... log_level import LogLevel @@ -18,27 +19,36 @@ default_ollama = 'http://localhost:11434' class Processor(ConsumerProducer): - def __init__( - self, - pulsar_host=None, - input_queue=default_input_queue, - output_queue=default_output_queue, - subscriber=default_subscriber, - log_level=LogLevel.INFO, - model=default_model, - ollama=default_ollama, - ): + def __init__(self, **params): + + input_queue = params.get("input_queue", default_input_queue) + output_queue = params.get("output_queue", default_output_queue) + subscriber = params.get("subscriber", default_subscriber) + model = params.get("model", default_model) + ollama = params.get("ollama", default_ollama) super(Processor, self).__init__( - pulsar_host=pulsar_host, - log_level=log_level, - input_queue=input_queue, - output_queue=output_queue, - subscriber=subscriber, - input_schema=TextCompletionRequest, - output_schema=TextCompletionResponse, + **params | { + "input_queue": input_queue, + "output_queue": output_queue, + "subscriber": subscriber, + "model": model, + "ollama": ollama, + "input_schema": TextCompletionRequest, + "output_schema": TextCompletionResponse, + } ) + if not hasattr(__class__, "model_metric"): + __class__.model_metric = Info( + 'model', 'Model information' + ) + + __class__.model_metric.info({ + "model": model, + "ollama": ollama, + }) + self.llm = Ollama(base_url=ollama, model=model) def handle(self, msg):