diff --git a/trustgraph-base/trustgraph/base/async_processor.py b/trustgraph-base/trustgraph/base/async_processor.py index 80b30356..8de3a231 100644 --- a/trustgraph-base/trustgraph/base/async_processor.py +++ b/trustgraph-base/trustgraph/base/async_processor.py @@ -48,9 +48,10 @@ class AsyncProcessor: self.config_handlers = [] self.config_sub_task = self.subscribe( - self.config_push_queue, config_subscriber_id, - schema=ConfigPush, handler=self.on_config_change, - start_of_messages=True + flow = None, + queue = self.config_push_queue, subscriber = config_subscriber_id, + schema = ConfigPush, handler = self.on_config_change, + start_of_messages = True ) self.running = True @@ -84,11 +85,12 @@ class AsyncProcessor: await asyncio.sleep(2) def subscribe( - self, queue, subscriber, schema, handler, metrics=None, + self, flow, queue, subscriber, schema, handler, metrics=None, start_of_messages=False ): return Consumer( + flow = flow, taskgroup = self.taskgroup, client = self.client, queue = queue, diff --git a/trustgraph-base/trustgraph/base/consumer.py b/trustgraph-base/trustgraph/base/consumer.py index 3a756a0f..adc33de6 100644 --- a/trustgraph-base/trustgraph/base/consumer.py +++ b/trustgraph-base/trustgraph/base/consumer.py @@ -1,4 +1,5 @@ +import _pulsar import asyncio import time @@ -7,13 +8,14 @@ from .. exceptions import TooManyRequests class Consumer: def __init__( - self, taskgroup, client, queue, subscriber, schema, + self, taskgroup, flow, client, queue, subscriber, schema, handler, rate_limit_retry_time = 10, rate_limit_timeout = 7200, metrics = None, start_of_messages=False, ): self.taskgroup = taskgroup + self.flow = flow self.client = client self.queue = queue self.subscriber = subscriber @@ -28,6 +30,12 @@ class Consumer: self.metrics = metrics + def __del__(self): + self.running = False + + if hasattr(self, "consumer"): + self.consumer.close() + async def stop(self): self.running = False @@ -58,7 +66,15 @@ class Consumer: while self.running: - msg = await asyncio.to_thread(self.consumer.receive) + try: + msg = await asyncio.to_thread( + self.consumer.receive, + timeout_millis=2000 + ) + except _pulsar.Timeout: + continue + except Exception as e: + raise e expiry = time.time() + self.rate_limit_timeout @@ -86,7 +102,7 @@ class Consumer: if self.metrics: with self.metrics.record_time(): - await self.handler(msg, self) + await self.handler(msg, self, self.flow) else: await self.handler(msg, self.consumer) @@ -132,3 +148,4 @@ class Consumer: if self.metrics: self.metrics.state("stopped") + diff --git a/trustgraph-base/trustgraph/base/flow_processor.py b/trustgraph-base/trustgraph/base/flow_processor.py index 8d7cec7a..a42b2164 100644 --- a/trustgraph-base/trustgraph/base/flow_processor.py +++ b/trustgraph-base/trustgraph/base/flow_processor.py @@ -10,6 +10,9 @@ from .. base import AsyncProcessor, Consumer, Producer from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics +class Flow: + pass + class FlowProcessor(AsyncProcessor): def __init__(self, **params): @@ -31,16 +34,12 @@ class FlowProcessor(AsyncProcessor): self.on_config(self.on_configuration) - self.consumers = {} - self.producers = {} + self.flows = {} # These can be overriden by a derived class self.consumer_spec = [] self.producer_spec = [] - -# "input", self.input_schema) -# "input_schema": Document, -# "output_schema": TextDocument, + self.config_spec = [] print("Service initialised.") @@ -50,9 +49,21 @@ class FlowProcessor(AsyncProcessor): def register_producer(self, name, schema): self.producer_spec.append((name, schema)) + def register_config(self, name): + self.config_spec.append((name,)) + async def start_flow(self, flow, defn): - producers = {} + flow_obj = Flow() + flow_obj.producer = {} + flow_obj.consumer = {} + flow_obj.config = {} + + for spec in self.config_spec: + name = spec[0] + flow_obj.config[name] = defn[name] + + setattr(flow_obj, name, defn[name]) for spec in self.producer_spec: @@ -68,9 +79,9 @@ class FlowProcessor(AsyncProcessor): metrics = producer_metrics, ) - producers[name] = producer + flow_obj.producer[name] = producer - consumers = {} + setattr(flow_obj, name, producer) for spec in self.consumer_spec: @@ -81,6 +92,7 @@ class FlowProcessor(AsyncProcessor): ) consumer = self.subscribe( + flow = flow_obj, queue = defn[name], subscriber = self.subscriber, schema = schema, @@ -92,25 +104,25 @@ class FlowProcessor(AsyncProcessor): # metadata consumer.id = self.id consumer.name = name - consumer.flow = flow - consumer.q = producers - - consumers[name] = consumer + consumer.flow = flow_obj + consumer.flow.name = flow await consumer.start() - self.consumers[flow] = consumers - self.producers[flow] = producers + flow_obj.consumer[name] = consumer + + setattr(flow_obj, name, consumer) + + self.flows[flow] = flow_obj print("Started flow: ", flow) async def stop_flow(self, flow): - for c in self.consumers[flow]: + for c in self.flows[flow].consumer: await c.stop() - del self.consumers[flow] - del self.producers[flow] + del self.flows[flow] print("Stopped flow: ", flow, flush=True) @@ -125,7 +137,7 @@ class FlowProcessor(AsyncProcessor): flow_config = json.loads(config["flows"][self.id]) wanted_keys = flow_config.keys() - current_keys = self.consumers.keys() + current_keys = self.flows.keys() for key in wanted_keys: if key not in current_keys: diff --git a/trustgraph-base/trustgraph/base/producer.py b/trustgraph-base/trustgraph/base/producer.py index 9cf62e0b..30bfabfc 100644 --- a/trustgraph-base/trustgraph/base/producer.py +++ b/trustgraph-base/trustgraph/base/producer.py @@ -12,7 +12,8 @@ class Producer: def __del__(self): - self.producer.close() + if hasattr(self, "producer"): + self.producer.close() async def send(self, msg, properties={}): @@ -21,4 +22,3 @@ class Producer: self.producer.send(msg, properties) - diff --git a/trustgraph-base/trustgraph/base/pubsub.py b/trustgraph-base/trustgraph/base/pubsub.py index 3c727a67..1ad08287 100644 --- a/trustgraph-base/trustgraph/base/pubsub.py +++ b/trustgraph-base/trustgraph/base/pubsub.py @@ -51,7 +51,8 @@ class PulsarClient: if self.client: self.client.close() - def subscribe(self, queue, subscriber, schema, start_of_messages=False): + def subscribe(self, queue, subscriber, schema, + start_of_messages=False): if start_of_messages: pos = pulsar.InitialPosition.Earliest @@ -59,10 +60,11 @@ class PulsarClient: pos = pulsar.InitialPosition.Latest return self.client.subscribe( - queue, subscriber, - consumer_type=pulsar.ConsumerType.Shared, - schema=JsonSchema(schema), - initial_position=pos, + topic = queue, + subscription_name = subscriber, + consumer_type = pulsar.ConsumerType.Shared, + schema = JsonSchema(schema), + initial_position = pos, ) def publish(self, queue, schema): diff --git a/trustgraph-base/trustgraph/base/request_response.py b/trustgraph-base/trustgraph/base/request_response.py index 6ccaed80..0f652603 100644 --- a/trustgraph-base/trustgraph/base/request_response.py +++ b/trustgraph-base/trustgraph/base/request_response.py @@ -1,4 +1,5 @@ +import json from pulsar.schema import JsonSchema from .. schema import Error @@ -8,109 +9,28 @@ from .. log_level import LogLevel from .. base import AsyncProcessor, Consumer, Producer from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics +from . flow_processor import FlowProcessor -class RequestResponseProcessor(AsyncProcessor): +class RequestResponseService(FlowProcessor): def __init__(self, **params): - - id = params.get("id") - self.subscriber = params.get("subscriber") - self.request_schema = params.get("request_schema") - self.response_schema = params.get("response_schema") - ProcessorMetrics(id=id).info( - { - "subscriber": self.subscriber, - "request_schema": self.request_schema.__name__, - "response_schema": self.response_schema.__name__, - } - ) + super(RequestResponseService, self).__init__(**params) - super(Processor, self).__init__( - **params | { - "request_schema": self.request_schema.__name__, - "response_schema": self.response_schema.__name__, - } - ) - - self.on_config(self.on_configuration) - - self.subs = {} - self.pubs = {} + # These can be overriden by a derived class + self.consumer_spec = [ + ("request", params.get("request_schema"), self.on_request) + ] + self.producer_spec = [ + ("response", params.get("response_schema")) + ] print("Service initialised.") - async def start_handler(self, flow, defn): - - request_metrics = ConsumerMetrics(self.id, flow) - response_metrics = ProducerMetrics(self.id, flow) - - self.subs[flow] = self.subscribe( - queue = defn["input"] - subscriber = subscriber, - schema = self.request_schema, - handler = self.on_message, - metrics = request_metrics, - ) - - self.pubs[flow] = self.publish( - queue = defn["output"], - schema = self.response_schema, - metrics = response_metrics, - ) - - self.subs[flow].start() - - print("Started flow for", flow) - - async def on_configuration(self, config, version): - - if "flows" not in config: return - - if self.id in config["flows"]: - - wanted_keys = config["flows"][self.id].keys() - current_keys = self.subs.keys() - - for key in wanted_keys: - if key not in current_keys: - self.start_handler(key, config["flows"][key]) - - for key in current_keys: - if key not in wanted_keys: - self.stop_handler(key, config["flows"][key]) - - print("Handled config update") - - async def start(self): - - await self.request_sub.start() - - async def on_message(self, msg, consumer): - - try: - - v = msg.value() - - # Sender-produced ID - id = msg.properties()["id"] - - print(f"Handling {id}...", flush=True) - - consumer.acknowledge(msg) - - except Exception as e: - - resp = "error" - - await self.response_pub.send(resp, properties={"id": id}) - - consumer.acknowledge(msg) - @staticmethod - def add_args(parser): + def add_args(parser, default_subscriber): - AsyncProcessor.add_args(parser) + FlowProcessor.add_args(parser) def run(): diff --git a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py index 963298d4..752a7787 100755 --- a/trustgraph-flow/trustgraph/chunking/recursive/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/recursive/chunker.py @@ -59,7 +59,7 @@ class Processor(FlowProcessor): print("Chunker initialised", flush=True) - async def on_message(self, msg, consumer): + async def on_message(self, msg, consumer, flow): v = msg.value() print(f"Chunking {v.metadata.id}...", flush=True) @@ -81,7 +81,7 @@ class Processor(FlowProcessor): id=consumer.id, flow=consumer.flow ).observe(len(chunk.page_content)) - await consumer.q["output"].send(r) + await flow.producer["output"].send(r) print("Done.", flush=True) diff --git a/trustgraph-flow/trustgraph/chunking/token/chunker.py b/trustgraph-flow/trustgraph/chunking/token/chunker.py index d35c003c..2ada07a0 100755 --- a/trustgraph-flow/trustgraph/chunking/token/chunker.py +++ b/trustgraph-flow/trustgraph/chunking/token/chunker.py @@ -58,7 +58,7 @@ class Processor(FlowProcessor): print("Chunker initialised", flush=True) - async def on_message(self, msg, consumer): + async def on_message(self, msg, consumer, flow): v = msg.value() print(f"Chunking {v.metadata.id}...", flush=True) @@ -80,7 +80,7 @@ class Processor(FlowProcessor): id=consumer.id, flow=consumer.flow ).observe(len(chunk.page_content)) - await consumer.q["output"].send(r) + await flow.producer["output"].send(r) print("Done.", flush=True) diff --git a/trustgraph-flow/trustgraph/config/service/service.py b/trustgraph-flow/trustgraph/config/service/service.py index d47d1feb..30cb3a0a 100644 --- a/trustgraph-flow/trustgraph/config/service/service.py +++ b/trustgraph-flow/trustgraph/config/service/service.py @@ -74,6 +74,7 @@ class Processor(AsyncProcessor): ) self.subs = self.subscribe( + flow = None, queue = request_queue, subscriber = subscriber, schema = request_schema, @@ -105,7 +106,7 @@ class Processor(AsyncProcessor): print("Pushed version ", self.config.version) - async def on_message(self, msg, consumer): + async def on_message(self, msg, consumer, flow): try: diff --git a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py index 8deb0394..ba47836a 100755 --- a/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py +++ b/trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py @@ -44,7 +44,7 @@ class Processor(FlowProcessor): print("PDF inited", flush=True) - async def on_message(self, msg, consumer): + async def on_message(self, msg, consumer, flow): print("PDF message received", flush=True) @@ -71,7 +71,7 @@ class Processor(FlowProcessor): text=page.page_content.encode("utf-8"), ) - await consumer.q["output"].send(r) + await flow.producer["output"].send(r) print("Done.", flush=True)