mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-30 02:46:23 +02:00
Reduc pulsar connections (#189)
This commit is contained in:
parent
df23e29971
commit
7e78aa6d91
9 changed files with 82 additions and 52 deletions
|
|
@ -41,10 +41,15 @@ class ServiceEndpoint:
|
|||
|
||||
self.operation = "service"
|
||||
|
||||
async def start(self):
|
||||
async def start(self, client):
|
||||
|
||||
self.pub_task = asyncio.create_task(self.pub.run())
|
||||
self.sub_task = asyncio.create_task(self.sub.run())
|
||||
self.pub_task = asyncio.create_task(self.pub.run(client))
|
||||
self.sub_task = asyncio.create_task(self.sub.run(client))
|
||||
|
||||
async def join(self):
|
||||
|
||||
await self.pub_task
|
||||
await self.sub_task
|
||||
|
||||
def add_routes(self, app):
|
||||
|
||||
|
|
|
|||
|
|
@ -29,10 +29,10 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
|
|||
schema=JsonSchema(GraphEmbeddings)
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
async def start(self, client):
|
||||
|
||||
self.task = asyncio.create_task(
|
||||
self.publisher.run()
|
||||
self.publisher.run(client)
|
||||
)
|
||||
|
||||
async def listener(self, ws, running):
|
||||
|
|
|
|||
|
|
@ -28,10 +28,10 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
|
|||
schema=JsonSchema(GraphEmbeddings)
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
async def start(self, client):
|
||||
|
||||
self.task = asyncio.create_task(
|
||||
self.subscriber.run()
|
||||
self.subscriber.run(client)
|
||||
)
|
||||
|
||||
async def async_thread(self, ws, running):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
|
||||
import asyncio
|
||||
import aiopulsar
|
||||
|
||||
class Publisher:
|
||||
|
||||
|
|
@ -12,24 +11,23 @@ class Publisher:
|
|||
self.q = asyncio.Queue(maxsize=max_size)
|
||||
self.chunking_enabled = chunking_enabled
|
||||
|
||||
async def run(self):
|
||||
async def run(self, client):
|
||||
|
||||
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()
|
||||
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)
|
||||
if id:
|
||||
await producer.send(item, { "id": id })
|
||||
else:
|
||||
await producer.send(item)
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from aiohttp import web
|
|||
import logging
|
||||
import os
|
||||
import base64
|
||||
import aiopulsar
|
||||
|
||||
import pulsar
|
||||
from pulsar.schema import JsonSchema
|
||||
|
|
@ -237,13 +238,35 @@ class Api:
|
|||
{ "error": str(e) }
|
||||
)
|
||||
|
||||
async def run_endpoints(self):
|
||||
|
||||
async with aiopulsar.connect(self.pulsar_host) as client:
|
||||
|
||||
for ep in self.endpoints:
|
||||
await ep.start(client)
|
||||
|
||||
self.doc_ingest_pub_task = asyncio.create_task(
|
||||
self.document_out.run(client)
|
||||
)
|
||||
|
||||
self.text_ingest_pub_task = asyncio.create_task(
|
||||
self.text_out.run(client)
|
||||
)
|
||||
|
||||
print("Endpoints are running...")
|
||||
|
||||
# They never exit
|
||||
for ep in self.endpoints:
|
||||
await ep.join()
|
||||
|
||||
await self.doc_ingest_pub_task
|
||||
await self.text_ingest_pub_task
|
||||
|
||||
print("Endpoints are stopped.")
|
||||
|
||||
async def app_factory(self):
|
||||
|
||||
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.endpoint_task = asyncio.create_task(self.run_endpoints())
|
||||
|
||||
return self.app
|
||||
|
||||
|
|
|
|||
|
|
@ -76,6 +76,12 @@ class SocketEndpoint:
|
|||
async def start(self):
|
||||
pass
|
||||
|
||||
async def join(self):
|
||||
|
||||
# Nothing to wait for
|
||||
while True:
|
||||
await asyncio.sleep(100)
|
||||
|
||||
def add_routes(self, app):
|
||||
|
||||
app.add_routes([
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
|
||||
import asyncio
|
||||
import aiopulsar
|
||||
|
||||
class Subscriber:
|
||||
|
||||
|
|
@ -14,33 +13,32 @@ class Subscriber:
|
|||
self.q = {}
|
||||
self.full = {}
|
||||
|
||||
async def run(self):
|
||||
async def run(self, client):
|
||||
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()
|
||||
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)
|
||||
# Acknowledge successful reception of the message
|
||||
await consumer.acknowledge(msg)
|
||||
|
||||
try:
|
||||
id = msg.properties()["id"]
|
||||
except:
|
||||
id = None
|
||||
try:
|
||||
id = msg.properties()["id"]
|
||||
except:
|
||||
id = None
|
||||
|
||||
value = msg.value()
|
||||
if id in self.q:
|
||||
await self.q[id].put(value)
|
||||
value = msg.value()
|
||||
if id in self.q:
|
||||
await self.q[id].put(value)
|
||||
|
||||
for q in self.full.values():
|
||||
await q.put(value)
|
||||
for q in self.full.values():
|
||||
await q.put(value)
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
|
|
|
|||
|
|
@ -27,10 +27,10 @@ class TriplesLoadEndpoint(SocketEndpoint):
|
|||
schema=JsonSchema(Triples)
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
async def start(self, client):
|
||||
|
||||
self.task = asyncio.create_task(
|
||||
self.publisher.run()
|
||||
self.publisher.run(client)
|
||||
)
|
||||
|
||||
async def listener(self, ws, running):
|
||||
|
|
|
|||
|
|
@ -26,10 +26,10 @@ class TriplesStreamEndpoint(SocketEndpoint):
|
|||
schema=JsonSchema(Triples)
|
||||
)
|
||||
|
||||
async def start(self):
|
||||
async def start(self, client):
|
||||
|
||||
self.task = asyncio.create_task(
|
||||
self.subscriber.run()
|
||||
self.subscriber.run(client)
|
||||
)
|
||||
|
||||
async def async_thread(self, ws, running):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue