diff --git a/trustgraph-base/trustgraph/base/publisher.py b/trustgraph-base/trustgraph/base/publisher.py index 99a27c19..f40fbb30 100644 --- a/trustgraph-base/trustgraph/base/publisher.py +++ b/trustgraph-base/trustgraph/base/publisher.py @@ -1,4 +1,6 @@ +from pulsar.schema import JsonSchema + import asyncio import time import pulsar @@ -16,7 +18,6 @@ class Publisher: async def start(self): self.task = asyncio.create_task(self.run()) - await self.task.start() async def stop(self): self.running = False @@ -32,14 +33,14 @@ class Publisher: try: producer = self.client.create_producer( topic=self.topic, - schema=self.schema, + schema=JsonSchema(self.schema), chunking_enabled=self.chunking_enabled, ) while self.running: try: - id, item = asyncio.wait_for( + id, item = await asyncio.wait_for( self.q.get(), timeout=0.5 ) diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py index 6b4b4838..bbfb51a3 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_load.py @@ -1,6 +1,5 @@ import asyncio -from pulsar.schema import JsonSchema import uuid from aiohttp import WSMsgType @@ -26,12 +25,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): self.publisher = Publisher( self.pulsar_client, document_embeddings_store_queue, - schema=JsonSchema(DocumentEmbeddings) + schema=DocumentEmbeddings ) async def start(self): - self.publisher.start() + await self.publisher.start() async def listener(self, ws, running): @@ -59,6 +58,6 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint): ], ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py index 6d7db576..e59a0370 100644 --- a/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py @@ -1,7 +1,6 @@ import asyncio import queue -from pulsar.schema import JsonSchema import uuid from .. schema import DocumentEmbeddings @@ -27,7 +26,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): self.subscriber = Subscriber( self.pulsar_client, document_embeddings_store_queue, "api-gateway", "api-gateway", - schema=JsonSchema(DocumentEmbeddings), + schema=DocumentEmbeddings, ) async def listener(self, ws, running): @@ -44,17 +43,17 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): async def start(self): - self.subscriber.start() + await self.subscriber.start() async def async_thread(self, ws, running): id = str(uuid.uuid4()) - q = self.subscriber.subscribe_all(id) + q = await self.subscriber.subscribe_all(id) while running.get(): try: - resp = await asyncio.to_thread(q.get, timeout=0.5) + resp = await asyncio.wait_for(q.get(), timeout=0.5) await ws.send_json(serialize_document_embeddings(resp)) except TimeoutError: @@ -67,7 +66,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint): print(f"Exception: {str(e)}", flush=True) break - self.subscriber.unsubscribe_all(id) + await self.subscriber.unsubscribe_all(id) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/endpoint.py b/trustgraph-flow/trustgraph/gateway/endpoint.py index 5005463c..94980e8b 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint.py @@ -1,13 +1,9 @@ import asyncio -from pulsar.schema import JsonSchema from aiohttp import web import uuid import logging -from .. base import Publisher -from .. base import Subscriber - logger = logging.getLogger("endpoint") logger.setLevel(logging.INFO) diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py index c1354ce5..27e92a30 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py @@ -1,6 +1,5 @@ import asyncio -from pulsar.schema import JsonSchema import uuid from aiohttp import WSMsgType @@ -26,12 +25,12 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): self.publisher = Publisher( self.pulsar_client, graph_embeddings_store_queue, - schema=JsonSchema(GraphEmbeddings) + schema=GraphEmbeddings ) async def start(self): - self.publisher.start() + await self.publisher.start() async def listener(self, ws, running): @@ -60,6 +59,6 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint): ] ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py index 385eb9f4..37edc2bb 100644 --- a/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py +++ b/trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py @@ -1,7 +1,6 @@ import asyncio import queue -from pulsar.schema import JsonSchema import uuid from .. schema import GraphEmbeddings @@ -26,7 +25,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): self.subscriber = Subscriber( self.pulsar_client, graph_embeddings_store_queue, "api-gateway", "api-gateway", - schema=JsonSchema(GraphEmbeddings) + schema=GraphEmbeddings ) async def listener(self, ws, running): @@ -41,17 +40,17 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): async def start(self): - self.subscriber.start() + await self.subscriber.start() async def async_thread(self, ws, running): id = str(uuid.uuid4()) - q = self.subscriber.subscribe_all(id) + q = await self.subscriber.subscribe_all(id) while running.get(): try: - resp = await asyncio.to_thread(q.get, timeout=0.5) + resp = await asyncio.wait_for(q.get, timeout=0.5) await ws.send_json(serialize_graph_embeddings(resp)) except TimeoutError: @@ -64,7 +63,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint): print(f"Exception: {str(e)}", flush=True) break - self.subscriber.unsubscribe_all(id) + await self.subscriber.unsubscribe_all(id) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/metrics.py b/trustgraph-flow/trustgraph/gateway/metrics.py index 33c1fe3a..d8a1ef62 100644 --- a/trustgraph-flow/trustgraph/gateway/metrics.py +++ b/trustgraph-flow/trustgraph/gateway/metrics.py @@ -7,7 +7,6 @@ import aiohttp from aiohttp import web import asyncio -from pulsar.schema import JsonSchema import uuid import logging diff --git a/trustgraph-flow/trustgraph/gateway/mux.py b/trustgraph-flow/trustgraph/gateway/mux.py index 23b693ab..8195c542 100644 --- a/trustgraph-flow/trustgraph/gateway/mux.py +++ b/trustgraph-flow/trustgraph/gateway/mux.py @@ -1,7 +1,6 @@ import asyncio import queue -from pulsar.schema import JsonSchema import uuid from aiohttp import web, WSMsgType diff --git a/trustgraph-flow/trustgraph/gateway/requestor.py b/trustgraph-flow/trustgraph/gateway/requestor.py index dc74667d..63395203 100644 --- a/trustgraph-flow/trustgraph/gateway/requestor.py +++ b/trustgraph-flow/trustgraph/gateway/requestor.py @@ -1,6 +1,5 @@ import asyncio -from pulsar.schema import JsonSchema import uuid import logging @@ -23,21 +22,21 @@ class ServiceRequestor: self.pub = Publisher( pulsar_client, request_queue, - schema=JsonSchema(request_schema), + schema=request_schema, ) self.sub = Subscriber( pulsar_client, response_queue, subscription, consumer_name, - JsonSchema(response_schema) + response_schema ) self.timeout = timeout async def start(self): - self.pub.start() - self.sub.start() + await self.pub.start() + await self.sub.start() def to_request(self, request): raise RuntimeError("Not defined") @@ -51,18 +50,15 @@ class ServiceRequestor: try: - q = self.sub.subscribe(id) + q = await self.sub.subscribe(id) - await asyncio.to_thread( - self.pub.send, id, self.to_request(request) - ) + await self.pub.send(id, self.to_request(request)) while True: try: - resp = await asyncio.to_thread( - q.get, - timeout=self.timeout + resp = await asyncio.wait_for( + q.get(), timeout=self.timeout ) except Exception as e: raise RuntimeError("Timeout") @@ -99,5 +95,5 @@ class ServiceRequestor: return err finally: - self.sub.unsubscribe(id) + await self.sub.unsubscribe(id) diff --git a/trustgraph-flow/trustgraph/gateway/sender.py b/trustgraph-flow/trustgraph/gateway/sender.py index 32c586b1..81b64e6d 100644 --- a/trustgraph-flow/trustgraph/gateway/sender.py +++ b/trustgraph-flow/trustgraph/gateway/sender.py @@ -2,7 +2,6 @@ # Like ServiceRequestor, but just fire-and-forget instead of request/response import asyncio -from pulsar.schema import JsonSchema import uuid import logging @@ -21,12 +20,12 @@ class ServiceSender: self.pub = Publisher( pulsar_client, request_queue, - schema=JsonSchema(request_schema), + schema=request_schema, ) async def start(self): - self.pub.start() + await self.pub.start() def to_request(self, request): raise RuntimeError("Not defined") @@ -35,9 +34,7 @@ class ServiceSender: try: - await asyncio.to_thread( - self.pub.send, None, self.to_request(request) - ) + await self.pub.send(None, self.to_request(request)) if responder: await responder({}, True) diff --git a/trustgraph-flow/trustgraph/gateway/service.py b/trustgraph-flow/trustgraph/gateway/service.py index 3f9fac25..29b31483 100755 --- a/trustgraph-flow/trustgraph/gateway/service.py +++ b/trustgraph-flow/trustgraph/gateway/service.py @@ -19,7 +19,6 @@ import os import base64 import pulsar -from pulsar.schema import JsonSchema from prometheus_client import start_http_server from .. log_level import LogLevel diff --git a/trustgraph-flow/trustgraph/gateway/triples_load.py b/trustgraph-flow/trustgraph/gateway/triples_load.py index bc69975e..81c8ea82 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_load.py +++ b/trustgraph-flow/trustgraph/gateway/triples_load.py @@ -1,6 +1,5 @@ import asyncio -from pulsar.schema import JsonSchema import uuid from aiohttp import WSMsgType @@ -24,12 +23,12 @@ class TriplesLoadEndpoint(SocketEndpoint): self.publisher = Publisher( self.pulsar_client, triples_store_queue, - schema=JsonSchema(Triples) + schema=Triples ) async def start(self): - self.publisher.start() + await self.publisher.start() async def listener(self, ws, running): @@ -51,7 +50,7 @@ class TriplesLoadEndpoint(SocketEndpoint): triples=to_subgraph(data["triples"]), ) - self.publisher.send(None, elt) + await self.publisher.send(None, elt) running.stop() diff --git a/trustgraph-flow/trustgraph/gateway/triples_stream.py b/trustgraph-flow/trustgraph/gateway/triples_stream.py index a5d5ad0a..a660591e 100644 --- a/trustgraph-flow/trustgraph/gateway/triples_stream.py +++ b/trustgraph-flow/trustgraph/gateway/triples_stream.py @@ -1,7 +1,6 @@ import asyncio import queue -from pulsar.schema import JsonSchema import uuid from .. schema import Triples @@ -24,7 +23,7 @@ class TriplesStreamEndpoint(SocketEndpoint): self.subscriber = Subscriber( self.pulsar_client, triples_store_queue, "api-gateway", "api-gateway", - schema=JsonSchema(Triples) + schema=Triples ) async def listener(self, ws, running): @@ -39,7 +38,7 @@ class TriplesStreamEndpoint(SocketEndpoint): async def start(self): - self.subscriber.start() + await self.subscriber.start() async def async_thread(self, ws, running):