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

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

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

@ -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()

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

@ -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()

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

@ -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")

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

@ -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

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

@ -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]

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

@ -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()

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

@ -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()