diff --git a/trustgraph-base/trustgraph/base/async_processor.py b/trustgraph-base/trustgraph/base/async_processor.py index 8de3a231..59c84a10 100644 --- a/trustgraph-base/trustgraph/base/async_processor.py +++ b/trustgraph-base/trustgraph/base/async_processor.py @@ -1,13 +1,15 @@ +# Base class for processors. Implements: +# - Pulsar client, subscribe and consume basic +# - the async startup logic +# - Initialising metrics + import asyncio -import os import argparse -import pulsar -from pulsar.schema import JsonSchema import _pulsar import time import uuid -from prometheus_client import start_http_server, Info, Enum +from prometheus_client import start_http_server, Info from .. schema import ConfigPush, config_push_queue from .. log_level import LogLevel @@ -17,73 +19,104 @@ from . producer import Producer from . consumer import Consumer default_config_queue = config_push_queue -config_subscriber_id = str(uuid.uuid4()) +# Async processor class AsyncProcessor: - def __init__(self, **params): + # Store the identity + self.id = params.get("id") + + # Register a pulsar client self.client = PulsarClient(**params) + # The processor runs all activity in a taskgroup, it's mandatory + # that this is provded self.taskgroup = params.get("taskgroup") if self.taskgroup is None: raise RuntimeError("Essential taskgroup missing") + # Pubsub parameters passed in if not hasattr(__class__, "params_metric"): __class__.params_metric = Info( 'params', 'Parameters configuration' ) - # FIXME: Maybe outputs information it should not + # Record metrics for the processor __class__.params_metric.info({ k: str(params[k]) for k in params }) + # Get the configuration topic self.config_push_queue = params.get( "config_push_queue", default_config_queue ) + # This records registered configuration handlers self.config_handlers = [] + # Create a random ID for this subscription to the configuration + # service + config_subscriber_id = str(uuid.uuid4()) + + # Subscribe to config queue self.config_sub_task = self.subscribe( flow = None, queue = self.config_push_queue, subscriber = config_subscriber_id, schema = ConfigPush, handler = self.on_config_change, + + # This causes new subscriptions to view the entire history of + # configuration start_of_messages = True ) self.running = True + # This is called to start dynamic behaviour. An over-ride point for + # extra functionality + async def start(self): + await self.config_sub_task.start() + + # This is called to stop all threads. An over-ride point for extra + # functionality def stop(self): self.client.close() self.running = False + # Returns the pulsar host @property def pulsar_host(self): return self.client.pulsar_host - def on_config(self, handler): + # Register a new event handler for configuration change + def register_config_handler(self, handler): self.config_handlers.append(handler) - async def start(self): - await self.config_sub_task.start() - + # Called when a new configuration message push occurs async def on_config_change(self, message, consumer): + # Get configuration data and version number config = message.value().config version = message.value().version + # Acknowledge the message consumer.acknowledge(message) + # Invoke message handlers print("Config change event", config, version, flush=True) for ch in self.config_handlers: await ch(config, version) + # This is the 'main' body of the handler. It is a point to override + # if needed. By default does nothing. Processors are implemented + # by adding consumer/producer functionality so maybe nothing is needed + # in the run() body async def run(self): while self.running: await asyncio.sleep(2) + # Subscribe to a topic, returns a Consumer def subscribe( self, flow, queue, subscriber, schema, handler, metrics=None, start_of_messages=False @@ -101,8 +134,8 @@ class AsyncProcessor: start_of_messages=start_of_messages, ) + # Open a mechanism to publish messages to a topic. Returns a subscriber def publish(self, queue, schema, metrics=None): - return Producer( client = self.client, queue = queue, @@ -110,51 +143,45 @@ class AsyncProcessor: metrics = metrics, ) - def set_processor_state(self, flow, state): - - if not hasattr(__class__, "state_metric"): - __class__.state_metric = Enum( - 'processor_state', 'Processor state', - ["flow"], - states=['starting', 'running', 'stopped'] - ) - - __class__.state_metric.labels("flow").state(state) - - def set_pubsub_info(self, flow, info) : - - if not hasattr(__class__, "pubsub_metric"): - __class__.pubsub_metric = Info( - 'pubsub', 'Pub/sub configuration', - ["flow"] - ) - - __class__.pubsub_metric.labels(flow=flow).info(info) - + # Startup fabric. This runs in 'async' mode, creates a taskgroup and + # runs the producer. @classmethod async def launch_async(cls, args, ident): - async with asyncio.TaskGroup() as tg: + try: - try: + # Create a taskgroup. This seems complicated, when an exception + # occurs, unhandled it looks like it cancels all threads in the + # taskgroup, but I have observed this not working. I think + # it's fixed now that exceptions are caught in the right place. + async with asyncio.TaskGroup() as tg: - p = cls(**args | { "taskgroup": tg }) - # FIXME: Two sort of 'ident' things going on here? - p.module = ident - p.config_ident = args.get("ident", "FIXME") + # Create a processor instance, and include the taskgroup + # as a paramter. A processor identity ident is used as + # - subscriber name + # - an identifier for flow configuration + p = cls(**args | { "taskgroup": tg, "id", ident }) - await p.start() + # Start the processor + await p.start() - task2 = tg.create_task(p.run()) + # Run the processor + task = tg.create_task(p.run()) - except Exception as e: - print("Exception, closing taskgroup", flush=True) - raise e + # The taskgroup causes everything to wait until + # all threads have stopped + # This is here to output a debug message, shouldn't be needed. + except Exception as e: + print("Exception, closing taskgroup", flush=True) + raise e + + # Startup fabric. launch calls launch_async in async mode. @classmethod def launch(cls, ident, doc): + # Start assembling CLI arguments parser = argparse.ArgumentParser( prog=ident, description=doc @@ -166,22 +193,29 @@ class AsyncProcessor: help=f'Configuration identity (default: {ident})', ) + # Invoke the class-specific add_args, which manages adding all the + # command-line arguments cls.add_args(parser) + # Parse arguments args = parser.parse_args() args = vars(args) + # Debug print(args, flush=True) + # Start the Prometheus metrics service if needed if args["metrics"]: start_http_server(args["metrics_port"]) + # Loop forever, exception handler while True: print("Starting...", flush=True) try: + # Launch the processor in an asyncio handler asyncio.run(cls.launch_async( args, ident )) @@ -194,6 +228,7 @@ class AsyncProcessor: print("Pulsar Interrupted.", flush=True) return + # Exceptions from a taskgroup come in as an exception group except ExceptionGroup as e: print("Exception group:", flush=True) @@ -206,10 +241,13 @@ class AsyncProcessor: print("Type:", type(e), flush=True) print("Exception:", e, flush=True) + # Retry occurs here print("Will retry...", flush=True) time.sleep(4) print("Retrying...", flush=True) + # The command-line arguments are built using a stack of add_args + # invocations @staticmethod def add_args(parser): diff --git a/trustgraph-base/trustgraph/base/flow_processor.py b/trustgraph-base/trustgraph/base/flow_processor.py index a42b2164..2ade11cc 100644 --- a/trustgraph-base/trustgraph/base/flow_processor.py +++ b/trustgraph-base/trustgraph/base/flow_processor.py @@ -1,5 +1,10 @@ +# Base class for processor with management of flows in & out which are managed +# by configuration. This is probably all processor types, except for the +# configuration service which can't manage itself. + import json + from pulsar.schema import JsonSchema from .. schema import Error @@ -13,45 +18,51 @@ from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics class Flow: pass +# Parent class for configurable processors, configured with flows by +# the config service class FlowProcessor(AsyncProcessor): def __init__(self, **params): - - self.id = params.get("id") - self.subscriber = params.get("subscriber") - ProcessorMetrics(id=self.id).info( - { - "subscriber": self.subscriber, - } - ) + # Initialise base class + super(FlowProcessor, self).__init__(**params) - super(FlowProcessor, self).__init__( - **params | { - "id": self.id, - } - ) + # Initialise metrics, records the parameters + ProcessorMetrics(id=self.id).info(params) - self.on_config(self.on_configuration) + # Register configuration handler + self.register_config_handler(self.on_configuration) + # Initialise flow information state self.flows = {} - # These can be overriden by a derived class + # These can be overriden by a derived class: + + # Consumer specification, array of ("name", SchemaType, handler) self.consumer_spec = [] + + # Producer specification, array of ("name", SchemaType) self.producer_spec = [] + + # Configuration specification, collects some flow variables from + # config, array of "name" self.config_spec = [] print("Service initialised.") + # Register a new consumer name def register_consumer(self, name, schema, handler): self.consumer_spec.append((name, schema, handler)) + # Register a producer name def register_producer(self, name, schema): self.producer_spec.append((name, schema)) + # Register a configuration variable def register_config(self, name): - self.config_spec.append((name,)) + self.config_spec.append(name) + # Start processing for a new flow async def start_flow(self, flow, defn): flow_obj = Flow() @@ -117,6 +128,7 @@ class FlowProcessor(AsyncProcessor): print("Started flow: ", flow) + # Stop processing for a new flow async def stop_flow(self, flow): for c in self.flows[flow].consumer: @@ -126,43 +138,50 @@ class FlowProcessor(AsyncProcessor): print("Stopped flow: ", flow, flush=True) + # Event handler - called for a configuration change async def on_configuration(self, config, version): print("Got config version", version, flush=True) + # Skip over invalid data if "flows" not in config: return + # Check there's configuration information for me if self.id in config["flows"]: + # Get my flow config flow_config = json.loads(config["flows"][self.id]) - wanted_keys = flow_config.keys() - current_keys = self.flows.keys() + # Get list of flows which should be running and are currently + # running + wanted_flows = flow_config.keys() + current_flows = self.flows.keys() - for key in wanted_keys: - if key not in current_keys: - await self.start_flow(key, flow_config[key]) + # Start all the flows which arent currently running + for flow in wanted_flows: + if flow not in current_flows: + await self.start_flow(flow, flow_config[flow]) - for key in current_keys: - if key not in wanted_keys: - await self.stop_flow(key) + # Stop all the unwanted flows which are due to be stopped + for flow in current_flows: + if flow not in wanted_flows: + await self.stop_flow(flow) print("Handled config update") + else: + + print("No configuration settings for me!", flush=True) + + # Start threads, just call parent async def start(self): await super(FlowProcessor, self).start() @staticmethod - def add_args(parser, default_subscriber): + def add_args(parser): AsyncProcessor.add_args(parser) - parser.add_argument( - '-s', '--subscriber', - default=default_subscriber, - help=f'Queue subscriber name (default: {default_subscriber})' - ) - # parser.add_argument( # '--rate-limit-retry', # type=int, diff --git a/trustgraph-flow/trustgraph/config/service/config.py b/trustgraph-flow/trustgraph/config/service/config.py index 597fb443..46ade4c3 100644 --- a/trustgraph-flow/trustgraph/config/service/config.py +++ b/trustgraph-flow/trustgraph/config/service/config.py @@ -2,6 +2,8 @@ from trustgraph.schema import ConfigResponse from trustgraph.schema import ConfigValue, Error +# This behaves just like a dict, should be easier to add persistent storage +# later class ConfigurationItems(dict): pass diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index 30cb3a0a..c441e640 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -20,10 +20,6 @@ module = "config-svc" default_request_queue = config_request_queue default_response_queue = config_response_queue default_push_queue = config_push_queue -default_subscriber = module - -# This behaves just like a dict, should be easier to add persistent storage -# later class Processor(AsyncProcessor): diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index ba47836a..03e688c6 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -13,21 +13,17 @@ from ... schema import document_ingest_queue, text_ingest_queue from ... log_level import LogLevel from ... base import FlowProcessor -module = "pdf-decoder" - -default_subscriber = module +default_ident = "pdf-decoder" class Processor(FlowProcessor): def __init__(self, **params): - id = params.get("id") - subscriber = params.get("subscriber", default_subscriber) + id = params.get("id", default_ident) super(Processor, self).__init__( **params | { "id": id, - "subscriber": subscriber, } ) @@ -77,10 +73,9 @@ class Processor(FlowProcessor): @staticmethod def add_args(parser): - - FlowProcessor.add_args(parser, default_subscriber) + FlowProcessor.add_args(parser, default_ident) def run(): - Processor.launch(module, __doc__) + Processor.launch(ident, __doc__) diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index 24ed21a9..87b10a7f 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py @@ -13,42 +13,29 @@ from .... schema import TextCompletionRequest, TextCompletionResponse from .... schema import text_completion_request_queue from .... schema import text_completion_response_queue from .... schema import prompt_request_queue, prompt_response_queue -from .... base import ConsumerProducer from .... clients.llm_client import LlmClient +from .... base import RequestResponseService from . prompt_manager import PromptConfiguration, Prompt, PromptManager module = "prompt" - -default_input_queue = prompt_request_queue -default_output_queue = prompt_response_queue default_subscriber = module -class Processor(ConsumerProducer): +class Processor(RequestResponseService): def __init__(self, **params): - input_queue = params.get("input_queue", default_input_queue) - output_queue = params.get("output_queue", default_output_queue) + id = params.get("id") subscriber = params.get("subscriber", default_subscriber) - tc_request_queue = params.get( - "text_completion_request_queue", text_completion_request_queue - ) - tc_response_queue = params.get( - "text_completion_response_queue", text_completion_response_queue - ) self.config_key = params.get("config_type", "prompt") super(Processor, self).__init__( **params | { - "input_queue": input_queue, - "output_queue": output_queue, + "id": id, "subscriber": subscriber, - "input_schema": PromptRequest, - "output_schema": PromptResponse, - "text_completion_request_queue": tc_request_queue, - "text_completion_response_queue": tc_response_queue, + "request_schema": PromptRequest, + "response_schema": PromptResponse, } ) @@ -217,7 +204,7 @@ class Processor(ConsumerProducer): @staticmethod def add_args(parser): - ConsumerProducer.add_args( + RequestResponseService.add_args( parser, default_input_queue, default_subscriber, default_output_queue, ) diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index bc774997..a18c4f82 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -25,7 +25,6 @@ from .... exceptions import TooManyRequests from .... base import RequestResponseService module = "text-completion" - default_subscriber = module default_model = 'gemini-1.0-pro-001' default_region = 'us-central1'