apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 01:01:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

To multithreads

Browse source
This commit is contained in:
Cyber MacGeddon 2024-12-03 17:21:26 +00:00
parent 0bd1a7265e
commit 2c3c508cc6
8 changed files with 114 additions and 95 deletions
Download patch file Download diff file Expand all files Collapse all files

37
trustgraph-flow/trustgraph/api/gateway/endpoint.py
View file

@ -43,8 +43,8 @@ class ServiceEndpoint:
async def start(self): async def start(self):
self.pub_task = asyncio.create_task(self.pub.run()) self.pub.start()
self.sub_task = asyncio.create_task(self.sub.run()) self.sub.start()
def add_routes(self, app): def add_routes(self, app):
@ -82,20 +82,16 @@ class ServiceEndpoint:
print(data) print(data)
q = await self.sub.subscribe(id) q = self.sub.subscribe(id)
await self.pub.send( await asyncio.to_thread(
id, self.pub.send, id, self.to_request(data)
self.to_request(data),
) )
print("Request sent")
try: try:
resp = await asyncio.wait_for(q.get(), self.timeout) resp = await asyncio.to_thread(q.get, timeout=self.timeout)
except: except Exception as e:
raise RuntimeError("Timeout waiting for response") raise RuntimeError("Timeout")
print("Response got")
if resp.error: if resp.error:
print("Error") print("Error")
@ -103,8 +99,6 @@ class ServiceEndpoint:
{ "error": resp.error.message } { "error": resp.error.message }
) )
print("Send response")
return web.json_response( return web.json_response(
self.from_response(resp) self.from_response(resp)
) )
@ -117,7 +111,7 @@ class ServiceEndpoint:
) )
finally: finally:
await self.sub.unsubscribe(id) self.sub.unsubscribe(id)
class MultiResponseServiceEndpoint(ServiceEndpoint): class MultiResponseServiceEndpoint(ServiceEndpoint):
@ -130,11 +124,10 @@ class MultiResponseServiceEndpoint(ServiceEndpoint):
data = await request.json() data = await request.json()
q = await self.sub.subscribe(id) q = self.sub.subscribe(id)
await self.pub.send( await asyncio.to_thread(
id, self.pub.send, id, self.to_request(data)
self.to_request(data),
) )
# Keeps looking at responses... # Keeps looking at responses...
@ -142,8 +135,8 @@ class MultiResponseServiceEndpoint(ServiceEndpoint):
while True: while True:
try: try:
resp = await asyncio.wait_for(q.get(), self.timeout) resp = await asyncio.to_thread(q.get, timeout=self.timeout)
except: except Exception as e:
raise RuntimeError("Timeout waiting for response") raise RuntimeError("Timeout waiting for response")
if resp.error: if resp.error:
@ -168,4 +161,4 @@ class MultiResponseServiceEndpoint(ServiceEndpoint):
) )
finally: finally:
await self.sub.unsubscribe(id) self.sub.unsubscribe(id)

6
trustgraph-flow/trustgraph/api/gateway/graph_embeddings_load.py
View file

@ -31,9 +31,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.task = asyncio.create_task( self.publisher.start()
self.publisher.run()
)
async def listener(self, ws, running): async def listener(self, ws, running):
@ -56,7 +54,7 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
vectors=data["vectors"], vectors=data["vectors"],
) )
await self.publisher.send(None, elt) self.publisher.send(None, elt)
running.stop() running.stop()

10
trustgraph-flow/trustgraph/api/gateway/graph_embeddings_stream.py
View file

@ -30,19 +30,17 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.task = asyncio.create_task( self.subscriber.start()
self.subscriber.run()
)
async def async_thread(self, ws, running): async def async_thread(self, ws, running):
id = str(uuid.uuid4()) id = str(uuid.uuid4())
q = await self.subscriber.subscribe_all(id) q = self.subscriber.subscribe_all(id)
while running.get(): while running.get():
try: 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)) await ws.send_json(serialize_graph_embeddings(resp))
except TimeoutError: except TimeoutError:
@ -52,7 +50,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True) print(f"Exception: {str(e)}", flush=True)
break break
await self.subscriber.unsubscribe_all(id) self.subscriber.unsubscribe_all(id)
running.stop() running.stop()

53
trustgraph-flow/trustgraph/api/gateway/publisher.py
View file

@ -1,6 +1,8 @@
import asyncio import queue
import aiopulsar import time
import pulsar
import threading
class Publisher: class Publisher:
@ -9,33 +11,46 @@ class Publisher:
self.pulsar_host = pulsar_host self.pulsar_host = pulsar_host
self.topic = topic self.topic = topic
self.schema = schema self.schema = schema
self.q = asyncio.Queue(maxsize=max_size) self.q = queue.Queue(maxsize=max_size)
self.chunking_enabled = chunking_enabled 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: while True:
try: 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: client = pulsar.Client(
await producer.send(item, { "id": id }) self.pulsar_host,
else: )
await producer.send(item)
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: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
# If handler drops out, sleep a retry # If handler drops out, sleep a retry
await asyncio.sleep(2) time.sleep(2)
async def send(self, id, msg): def send(self, id, msg):
await self.q.put((id, msg)) self.q.put((id, msg))
print("PUT ON Q")

10
trustgraph-flow/trustgraph/api/gateway/service.py
View file

@ -166,7 +166,8 @@ class Api:
# content is valid base64 # content is valid base64
doc = base64.b64decode(data["data"]) doc = base64.b64decode(data["data"])
resp = await self.document_out.send( resp = await asyncio.to_thread(
self.document_out.send,
None, None,
Document( Document(
metadata=Metadata( metadata=Metadata(
@ -211,7 +212,8 @@ class Api:
# Text is base64 encoded # Text is base64 encoded
text = base64.b64decode(data["text"]).decode(charset) text = base64.b64decode(data["text"]).decode(charset)
resp = await self.text_out.send( resp = asyncio.to_thread(
self.text_out.send,
None, None,
TextDocument( TextDocument(
metadata=Metadata( metadata=Metadata(
@ -242,8 +244,8 @@ class Api:
for ep in self.endpoints: for ep in self.endpoints:
await ep.start() await ep.start()
self.doc_ingest_pub_task = asyncio.create_task(self.document_out.run()) self.document_out.start()
self.text_ingest_pub_task = asyncio.create_task(self.text_out.run()) self.text_out.start()
return self.app return self.app

77
trustgraph-flow/trustgraph/api/gateway/subscriber.py
View file

@ -1,6 +1,8 @@
import asyncio import queue
import aiopulsar import pulsar
import threading
import time
class Subscriber: class Subscriber:
@ -13,56 +15,71 @@ class Subscriber:
self.schema = schema self.schema = schema
self.q = {} self.q = {}
self.full = {} 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: while True:
try: 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 client = pulsar.Client(
await consumer.acknowledge(msg) self.pulsar_host,
)
try: consumer = client.subscribe(
id = msg.properties()["id"] topic=self.topic,
except: subscription_name=self.subscription,
id = None consumer_name=self.consumer_name,
schema=self.schema,
)
while True:
value = msg.value() msg = consumer.receive()
if id in self.q:
await self.q[id].put(value)
for q in self.full.values(): print("Received", msg)
await q.put(value)
# 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: except Exception as e:
print("Exception:", e, flush=True) print("Exception:", e, flush=True)
# If handler drops out, sleep a retry # If handler drops out, sleep a retry
await asyncio.sleep(2) time.sleep(2)
async def subscribe(self, id): def subscribe(self, id):
q = asyncio.Queue() q = queue.Queue(maxsize=self.max_size)
self.q[id] = q self.q[id] = q
return q return q
async def unsubscribe(self, id): def unsubscribe(self, id):
if id in self.q: if id in self.q:
del self.q[id] del self.q[id]
async def subscribe_all(self, id): def subscribe_all(self, id):
q = asyncio.Queue() q = queue.Queue(maxsize=self.max_size)
self.full[id] = q self.full[id] = q
return q return q
async def unsubscribe_all(self, id): def unsubscribe_all(self, id):
if id in self.full: if id in self.full:
del self.full[id] del self.full[id]

6
trustgraph-flow/trustgraph/api/gateway/triples_load.py
View file

@ -29,9 +29,7 @@ class TriplesLoadEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.task = asyncio.create_task( self.publisher.start()
self.publisher.run()
)
async def listener(self, ws, running): async def listener(self, ws, running):
@ -53,7 +51,7 @@ class TriplesLoadEndpoint(SocketEndpoint):
triples=to_subgraph(data["triples"]), triples=to_subgraph(data["triples"]),
) )
await self.publisher.send(None, elt) self.publisher.send(None, elt)
running.stop() running.stop()

10
trustgraph-flow/trustgraph/api/gateway/triples_stream.py
View file

@ -28,19 +28,17 @@ class TriplesStreamEndpoint(SocketEndpoint):
async def start(self): async def start(self):
self.task = asyncio.create_task( self.subscriber.start()
self.subscriber.run()
)
async def async_thread(self, ws, running): async def async_thread(self, ws, running):
id = str(uuid.uuid4()) id = str(uuid.uuid4())
q = await self.subscriber.subscribe_all(id) q = self.subscriber.subscribe_all(id)
while running.get(): while running.get():
try: 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)) await ws.send_json(serialize_triples(resp))
except TimeoutError: except TimeoutError:
@ -50,7 +48,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True) print(f"Exception: {str(e)}", flush=True)
break break
await self.subscriber.unsubscribe_all(id) self.subscriber.unsubscribe_all(id)
running.stop() running.stop()