Comments, maybe broke something

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):
+    # Called when a new configuration message push occurs
-        await self.config_sub_task.start()
-
     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):
+    # Startup fabric.  This runs in 'async' mode, creates a taskgroup and
-
+    # runs the producer.
-        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)
-
     @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?
+                    # Create a processor instance, and include the taskgroup
-                p.module = ident
+                    # as a paramter.  A processor identity ident is used as
-                p.config_ident = args.get("ident", "FIXME")
+                    # - 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:
+                    # The taskgroup causes everything to wait until
-                print("Exception, closing taskgroup", flush=True)
+                    # all threads have stopped
-                raise e
 
+        # 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):
 

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(
+        # Initialise base class
-            {
+        super(FlowProcessor, self).__init__(**params)
-                "subscriber": self.subscriber,
-            }
-        )
 
-        super(FlowProcessor, self).__init__(
+        # Initialise metrics, records the parameters
-            **params | {
+        ProcessorMetrics(id=self.id).info(params)
-                "id": self.id,
-            }
-        )
 
-        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()
+            # Get list of flows which should be running and are currently
-            current_keys = self.flows.keys()
+            # running
+            wanted_flows = flow_config.keys()
+            current_flows = self.flows.keys()
 
-            for key in wanted_keys:
+            # Start all the flows which arent currently running
-                if key not in current_keys:
+            for flow in wanted_flows:
-                    await self.start_flow(key, flow_config[key])
+                if flow not in current_flows:
+                    await self.start_flow(flow, flow_config[flow])
 
-            for key in current_keys:
+            # Stop all the unwanted flows which are due to be stopped
-                if key not in wanted_keys:
+            for flow in current_flows:
-                    await self.stop_flow(key)
+                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,

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
 

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):
 

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_ident = "pdf-decoder"
-
-default_subscriber = module
 
 class Processor(FlowProcessor):
 
     def __init__(self, **params):
 
-        id = params.get("id")
+        id = params.get("id", default_ident)
-        subscriber = params.get("subscriber", default_subscriber)
 
         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_ident)
-        FlowProcessor.add_args(parser, default_subscriber)
 
 def run():
 
-    Processor.launch(module, __doc__)
+    Processor.launch(ident, __doc__)
 

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)
+        id = params.get("id")
-        output_queue = params.get("output_queue", default_output_queue)
         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,
+                "id": id,
-                "output_queue": output_queue,
                 "subscriber": subscriber,
-                "input_schema": PromptRequest,
+                "request_schema": PromptRequest,
-                "output_schema": PromptResponse,
+                "response_schema": PromptResponse,
-                "text_completion_request_queue": tc_request_queue,
-                "text_completion_response_queue": tc_response_queue,
             }
         )
 
@@ -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,
         )

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'