diff --git a/trustgraph-flow/trustgraph/api/gateway/endpoint.py b/trustgraph-flow/trustgraph/api/gateway/endpoint.py index dc380f4b..26bbc364 100644 --- a/trustgraph-flow/trustgraph/api/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/api/gateway/endpoint.py @@ -43,8 +43,8 @@ class ServiceEndpoint: async def start(self): - self.pub_task = asyncio.create_task(self.pub.run()) - self.sub_task = asyncio.create_task(self.sub.run()) + self.pub.start() + self.sub.start() def add_routes(self, app): @@ -82,20 +82,16 @@ class ServiceEndpoint: print(data) - q = await self.sub.subscribe(id) + q = self.sub.subscribe(id) - await self.pub.send( - id, - self.to_request(data), + await asyncio.to_thread( + self.pub.send, id, self.to_request(data) ) - print("Request sent") try: - resp = await asyncio.wait_for(q.get(), self.timeout) - except: - raise RuntimeError("Timeout waiting for response") - - print("Response got") + resp = await asyncio.to_thread(q.get, timeout=self.timeout) + except Exception as e: + raise RuntimeError("Timeout") if resp.error: print("Error") @@ -103,8 +99,6 @@ class ServiceEndpoint: { "error": resp.error.message } ) - print("Send response") - return web.json_response( self.from_response(resp) ) @@ -117,7 +111,7 @@ class ServiceEndpoint: ) finally: - await self.sub.unsubscribe(id) + self.sub.unsubscribe(id) class MultiResponseServiceEndpoint(ServiceEndpoint): @@ -130,11 +124,10 @@ class MultiResponseServiceEndpoint(ServiceEndpoint): data = await request.json() - q = await self.sub.subscribe(id) + q = self.sub.subscribe(id) - await self.pub.send( - id, - self.to_request(data), + await asyncio.to_thread( + self.pub.send, id, self.to_request(data) ) # Keeps looking at responses... @@ -142,8 +135,8 @@ class MultiResponseServiceEndpoint(ServiceEndpoint): while True: try: - resp = await asyncio.wait_for(q.get(), self.timeout) - except: + resp = await asyncio.to_thread(q.get, timeout=self.timeout) + except Exception as e: raise RuntimeError("Timeout waiting for response") if resp.error: @@ -168,4 +161,4 @@ class MultiResponseServiceEndpoint(ServiceEndpoint): ) finally: - await self.sub.unsubscribe(id) + self.sub.unsubscribe(id) diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py index 15efdf5b..81fb6647 100644 --- a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py @@ -31,9 +31,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): async def start(self): - self.task = asyncio.create_task( - self.publisher.run() - ) + self.publisher.start() async def listener(self, ws, running): @@ -56,7 +54,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): vectors=data["vectors"], ) - await self.publisher.send(None, elt) + self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py index 7f3e5e18..7b3a9524 100644 --- a/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py @@ -30,19 +30,17 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): async def start(self): - self.task = asyncio.create_task( - self.subscriber.run() - ) + self.subscriber.start() async def async_thread(self, ws, running): id = str(uuid.uuid4()) - q = await self.subscriber.subscribe_all(id) + q = self.subscriber.subscribe_all(id) while running.get(): try: - resp = await asyncio.wait_for(q.get(), 0.5) + resp = await asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_graph_embeddings(resp)) except TimeoutError: @@ -52,7 +50,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): print(f"Exception: {str(e)}", flush=True) break - await self.subscriber.unsubscribe_all(id) + self.subscriber.unsubscribe_all(id) running.stop() diff --git a/trustgraph-flow/trustgraph/api/gateway/publisher.py b/trustgraph-flow/trustgraph/api/gateway/publisher.py index 1bff44dd..7e9c9677 100644 --- a/trustgraph-flow/trustgraph/api/gateway/publisher.py +++ b/trustgraph-flow/trustgraph/api/gateway/publisher.py @@ -1,6 +1,8 @@ -import asyncio -import aiopulsar +import queue +import time +import pulsar +import threading class Publisher: @@ -9,33 +11,46 @@ class Publisher: self.pulsar_host = pulsar_host self.topic = topic self.schema = schema - self.q = asyncio.Queue(maxsize=max_size) + self.q = queue.Queue(maxsize=max_size) self.chunking_enabled = chunking_enabled - async def run(self): + def start(self): + self.task = threading.Thread(target=self.run) + self.task.start() + + def run(self): while True: try: - async with aiopulsar.connect(self.pulsar_host) as client: - async with client.create_producer( - topic=self.topic, - schema=self.schema, - chunking_enabled=self.chunking_enabled, - ) as producer: - while True: - id, item = await self.q.get() - if id: - await producer.send(item, { "id": id }) - else: - await producer.send(item) + client = pulsar.Client( + self.pulsar_host, + ) + + producer = client.create_producer( + topic=self.topic, + schema=self.schema, + chunking_enabled=self.chunking_enabled, + ) + + while True: + + id, item = self.q.get() + + print("THING ON Q") + + if id: + producer.send(item, { "id": id }) + else: + producer.send(item) except Exception as e: print("Exception:", e, flush=True) # If handler drops out, sleep a retry - await asyncio.sleep(2) + time.sleep(2) - async def send(self, id, msg): - await self.q.put((id, msg)) + def send(self, id, msg): + self.q.put((id, msg)) + print("PUT ON Q") diff --git a/trustgraph-flow/trustgraph/api/gateway/service.py b/trustgraph-flow/trustgraph/api/gateway/service.py index a25dd9dc..38ff8291 100755 --- a/trustgraph-flow/trustgraph/api/gateway/service.py +++ b/trustgraph-flow/trustgraph/api/gateway/service.py @@ -166,7 +166,8 @@ class Api: # content is valid base64 doc = base64.b64decode(data["data"]) - resp = await self.document_out.send( + resp = await asyncio.to_thread( + self.document_out.send, None, Document( metadata=Metadata( @@ -211,7 +212,8 @@ class Api: # Text is base64 encoded text = base64.b64decode(data["text"]).decode(charset) - resp = await self.text_out.send( + resp = asyncio.to_thread( + self.text_out.send, None, TextDocument( metadata=Metadata( @@ -242,8 +244,8 @@ class Api: for ep in self.endpoints: await ep.start() - self.doc_ingest_pub_task = asyncio.create_task(self.document_out.run()) - self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) + self.document_out.start() + self.text_out.start() return self.app diff --git a/trustgraph-flow/trustgraph/api/gateway/subscriber.py b/trustgraph-flow/trustgraph/api/gateway/subscriber.py index 3d8840f6..ae947aff 100644 --- a/trustgraph-flow/trustgraph/api/gateway/subscriber.py +++ b/trustgraph-flow/trustgraph/api/gateway/subscriber.py @@ -1,6 +1,8 @@ -import asyncio -import aiopulsar +import queue +import pulsar +import threading +import time class Subscriber: @@ -13,56 +15,71 @@ class Subscriber: self.schema = schema self.q = {} self.full = {} + self.max_size = max_size + + def start(self): + self.task = threading.Thread(target=self.run) + self.task.start() + + def run(self): - async def run(self): while True: + try: - async with aiopulsar.connect(self.pulsar_host) as client: - async with client.subscribe( - topic=self.topic, - subscription_name=self.subscription, - consumer_name=self.consumer_name, - schema=self.schema, - ) as consumer: - while True: - msg = await consumer.receive() - # Acknowledge successful reception of the message - await consumer.acknowledge(msg) + client = pulsar.Client( + self.pulsar_host, + ) - try: - id = msg.properties()["id"] - except: - id = None + consumer = client.subscribe( + topic=self.topic, + subscription_name=self.subscription, + consumer_name=self.consumer_name, + schema=self.schema, + ) + + while True: - value = msg.value() - if id in self.q: - await self.q[id].put(value) + msg = consumer.receive() - for q in self.full.values(): - await q.put(value) + print("Received", msg) + + # Acknowledge successful reception of the message + consumer.acknowledge(msg) + + try: + id = msg.properties()["id"] + except: + id = None + + value = msg.value() + if id in self.q: + self.q[id].put(value) + + for q in self.full.values(): + q.put(value) except Exception as e: print("Exception:", e, flush=True) # If handler drops out, sleep a retry - await asyncio.sleep(2) + time.sleep(2) - async def subscribe(self, id): - q = asyncio.Queue() + def subscribe(self, id): + q = queue.Queue(maxsize=self.max_size) self.q[id] = q return q - async def unsubscribe(self, id): + def unsubscribe(self, id): if id in self.q: del self.q[id] - async def subscribe_all(self, id): - q = asyncio.Queue() + def subscribe_all(self, id): + q = queue.Queue(maxsize=self.max_size) self.full[id] = q return q - async def unsubscribe_all(self, id): + def unsubscribe_all(self, id): if id in self.full: del self.full[id] diff --git a/trustgraph-flow/trustgraph/api/gateway/triples_load.py b/trustgraph-flow/trustgraph/api/gateway/triples_load.py index 7f4561b1..dbb3e617 100644 --- a/trustgraph-flow/trustgraph/api/gateway/triples_load.py +++ b/trustgraph-flow/trustgraph/api/gateway/triples_load.py @@ -29,9 +29,7 @@ class TriplesLoadEndpoint(SocketEndpoint): async def start(self): - self.task = asyncio.create_task( - self.publisher.run() - ) + self.publisher.start() async def listener(self, ws, running): @@ -53,7 +51,7 @@ class TriplesLoadEndpoint(SocketEndpoint): triples=to_subgraph(data["triples"]), ) - await self.publisher.send(None, elt) + self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/api/gateway/triples_stream.py b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py index 6ecd2bdb..0875938e 100644 --- a/trustgraph-flow/trustgraph/api/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/api/gateway/triples_stream.py @@ -28,19 +28,17 @@ class TriplesStreamEndpoint(SocketEndpoint): async def start(self): - self.task = asyncio.create_task( - self.subscriber.run() - ) + self.subscriber.start() async def async_thread(self, ws, running): id = str(uuid.uuid4()) - q = await self.subscriber.subscribe_all(id) + q = self.subscriber.subscribe_all(id) while running.get(): try: - resp = await asyncio.wait_for(q.get(), 0.5) + resp = asyncio.to_thread(q.get, timeout=0.5) await ws.send_json(serialize_triples(resp)) except TimeoutError: @@ -50,7 +48,7 @@ class TriplesStreamEndpoint(SocketEndpoint): print(f"Exception: {str(e)}", flush=True) break - await self.subscriber.unsubscribe_all(id) + self.subscriber.unsubscribe_all(id) running.stop()