mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-22 03:31:02 +02:00
Working mux socket
This commit is contained in:
Cyber MacGeddon
parent
a70ae9793a
commit
bc33122e6e
4 changed files with 182 additions and 168 deletions
|
|
@ -27,6 +27,8 @@ from . triples_import import TriplesImport
|
|||
from . graph_embeddings_import import GraphEmbeddingsImport
|
||||
from . document_embeddings_import import DocumentEmbeddingsImport
|
||||
|
||||
from . mux import Mux
|
||||
|
||||
request_response_dispatchers = {
|
||||
"agent": AgentRequestor,
|
||||
"text-completion": TextCompletionRequestor,
|
||||
|
|
@ -35,7 +37,7 @@ request_response_dispatchers = {
|
|||
"document-rag": DocumentRagRequestor,
|
||||
"embeddings": EmbeddingsRequestor,
|
||||
"graph-embeddings": GraphEmbeddingsQueryRequestor,
|
||||
"triples-query": TriplesQueryRequestor,
|
||||
"triples": TriplesQueryRequestor,
|
||||
}
|
||||
|
||||
sender_dispatchers = {
|
||||
|
|
@ -120,6 +122,9 @@ class DispatcherManager:
|
|||
def dispatch_export(self):
|
||||
return self.invoke_export
|
||||
|
||||
def dispatch_socket(self):
|
||||
return self.invoke_socket
|
||||
|
||||
async def invoke_import(self, ws, running, params):
|
||||
|
||||
flow = params.get("flow")
|
||||
|
|
@ -184,11 +189,26 @@ class DispatcherManager:
|
|||
|
||||
return dispatcher
|
||||
|
||||
async def invoke_socket(self, ws, running, params):
|
||||
|
||||
flow = params.get("flow")
|
||||
|
||||
if flow not in self.flows:
|
||||
raise RuntimeError("Invalid flow")
|
||||
|
||||
dispatcher = Mux(self, ws, running)
|
||||
|
||||
return dispatcher
|
||||
|
||||
async def process(self, data, responder, params):
|
||||
|
||||
flow = params.get("flow")
|
||||
kind = params.get("kind")
|
||||
|
||||
return await self.invoke(data, responder, flow, kind)
|
||||
|
||||
async def invoke(self, data, responder, flow, kind):
|
||||
|
||||
if flow not in self.flows:
|
||||
raise RuntimeError("Invalid flow")
|
||||
|
||||
|
|
|
|||
156
trustgraph-flow/trustgraph/gateway/dispatch/mux.py
Normal file
156
trustgraph-flow/trustgraph/gateway/dispatch/mux.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
|
||||
MAX_OUTSTANDING_REQUESTS = 15
|
||||
WORKER_CLOSE_WAIT = 0.01
|
||||
START_REQUEST_WAIT = 0.1
|
||||
|
||||
# This buffers requests until task start, so short-lived
|
||||
MAX_QUEUE_SIZE = 10
|
||||
|
||||
class Mux:
|
||||
|
||||
def __init__(self, dispatcher_manager, ws, running):
|
||||
|
||||
self.dispatcher_manager = dispatcher_manager
|
||||
self.ws = ws
|
||||
self.running = running
|
||||
|
||||
self.q = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
|
||||
|
||||
async def destroy(self):
|
||||
|
||||
self.running.stop()
|
||||
|
||||
if self.ws:
|
||||
await self.ws.close()
|
||||
|
||||
async def receive(self, msg):
|
||||
|
||||
try:
|
||||
|
||||
data = msg.json()
|
||||
|
||||
# if data["service"] not in self.services:
|
||||
# raise RuntimeError("Bad service")
|
||||
|
||||
if "request" not in data:
|
||||
raise RuntimeError("Bad message")
|
||||
|
||||
if "id" not in data:
|
||||
raise RuntimeError("Bad message")
|
||||
|
||||
await self.q.put(
|
||||
(data["id"], data["service"], data["request"])
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print("receive exception:", str(e), flush=True)
|
||||
await self.ws.send_json({"error": str(e)})
|
||||
|
||||
async def maybe_tidy_workers(self, workers):
|
||||
|
||||
while True:
|
||||
|
||||
try:
|
||||
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(workers[0]),
|
||||
WORKER_CLOSE_WAIT
|
||||
)
|
||||
|
||||
# worker[0] now stopped
|
||||
# FIXME: Delete reference???
|
||||
|
||||
workers.pop(0)
|
||||
|
||||
if len(workers) == 0:
|
||||
break
|
||||
|
||||
# Loop iterates to try the next worker
|
||||
|
||||
except TimeoutError:
|
||||
# worker[0] still running, move on
|
||||
break
|
||||
|
||||
async def start_request_task(self, ws, id, svc, request, workers):
|
||||
|
||||
# Wait for outstanding requests to go below MAX_OUTSTANDING_REQUESTS
|
||||
while len(workers) > MAX_OUTSTANDING_REQUESTS:
|
||||
|
||||
# Fixes deadlock
|
||||
# FIXME: Put it in its own loop
|
||||
await asyncio.sleep(START_REQUEST_WAIT)
|
||||
|
||||
await self.maybe_tidy_workers(workers)
|
||||
|
||||
async def responder(resp, fin):
|
||||
await self.ws.send_json({
|
||||
"id": id,
|
||||
"response": resp,
|
||||
"complete": fin,
|
||||
})
|
||||
|
||||
worker = asyncio.create_task(
|
||||
self.request_task(request, responder, "0000", svc)
|
||||
)
|
||||
|
||||
workers.append(worker)
|
||||
|
||||
async def request_task(self, request, responder, flow, svc):
|
||||
|
||||
try:
|
||||
dispatcher = self.dispatcher_manager.dispatch_service()
|
||||
await dispatcher.invoke(request, responder, "0000", svc)
|
||||
except Exception as e:
|
||||
await self.ws.send_json({"error": str(e)})
|
||||
|
||||
async def run(self):
|
||||
|
||||
# Worker threads, servicing
|
||||
workers = []
|
||||
|
||||
while self.running.get():
|
||||
|
||||
try:
|
||||
|
||||
if len(workers) > 0:
|
||||
await self.maybe_tidy_workers(workers)
|
||||
|
||||
# Get next request on queue
|
||||
id, svc, request = await asyncio.wait_for(self.q.get(), 1)
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
# This is an internal working error, may not be recoverable
|
||||
print("run prepare exception:", e)
|
||||
await self.ws.send_json({"id": id, "error": str(e)})
|
||||
self.running.stop()
|
||||
|
||||
if self.ws:
|
||||
self.ws.close()
|
||||
self.ws = None
|
||||
|
||||
break
|
||||
|
||||
try:
|
||||
print(id, svc, request)
|
||||
|
||||
await self.start_request_task(
|
||||
self.ws, id, svc, request, workers
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print("Exception2:", e)
|
||||
await self.ws.send_json({"error": str(e)})
|
||||
|
||||
self.running.stop()
|
||||
|
||||
if self.ws:
|
||||
self.ws.close()
|
||||
self.ws = None
|
||||
|
||||
|
|
@ -55,6 +55,11 @@ class EndpointManager:
|
|||
auth = auth,
|
||||
dispatcher = dispatcher_manager.dispatch_export()
|
||||
),
|
||||
SocketEndpoint(
|
||||
endpoint_path = "/api/v1/flow/{flow}/socket",
|
||||
auth = auth,
|
||||
dispatcher = dispatcher_manager.dispatch_socket()
|
||||
),
|
||||
]
|
||||
|
||||
def add_routes(self, app):
|
||||
|
|
|
|||
|
|
@ -1,167 +0,0 @@
|
|||
|
||||
import asyncio
|
||||
import queue
|
||||
import uuid
|
||||
from aiohttp import web, WSMsgType
|
||||
|
||||
from . socket import SocketEndpoint
|
||||
|
||||
MAX_OUTSTANDING_REQUESTS = 15
|
||||
WORKER_CLOSE_WAIT = 0.01
|
||||
START_REQUEST_WAIT = 0.1
|
||||
|
||||
# This buffers requests until task start, so short-lived
|
||||
MAX_QUEUE_SIZE = 10
|
||||
|
||||
class MuxEndpoint(SocketEndpoint):
|
||||
|
||||
def __init__(
|
||||
self, pulsar_client, auth,
|
||||
services,
|
||||
path="/api/v1/socket",
|
||||
):
|
||||
|
||||
super(MuxEndpoint, self).__init__(
|
||||
endpoint_path=path, auth=auth,
|
||||
)
|
||||
|
||||
self.services = services
|
||||
|
||||
async def start(self):
|
||||
pass
|
||||
|
||||
async def maybe_tidy_workers(self, workers):
|
||||
|
||||
while True:
|
||||
|
||||
try:
|
||||
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(workers[0]),
|
||||
WORKER_CLOSE_WAIT
|
||||
)
|
||||
|
||||
# worker[0] now stopped
|
||||
# FIXME: Delete reference???
|
||||
|
||||
workers.pop(0)
|
||||
|
||||
if len(workers) == 0:
|
||||
break
|
||||
|
||||
# Loop iterates to try the next worker
|
||||
|
||||
except TimeoutError:
|
||||
# worker[0] still running, move on
|
||||
break
|
||||
|
||||
async def start_request_task(self, ws, id, svc, request, workers):
|
||||
|
||||
if svc not in self.services:
|
||||
await ws.send_json({"id": id, "error": "Service not recognised"})
|
||||
return
|
||||
|
||||
requestor = self.services[svc]
|
||||
|
||||
async def responder(resp, fin):
|
||||
await ws.send_json({
|
||||
"id": id,
|
||||
"response": resp,
|
||||
"complete": fin,
|
||||
})
|
||||
|
||||
# Wait for outstanding requests to go below MAX_OUTSTANDING_REQUESTS
|
||||
while len(workers) > MAX_OUTSTANDING_REQUESTS:
|
||||
|
||||
# Fixes deadlock
|
||||
# FIXME: Put it in its own loop
|
||||
await asyncio.sleep(START_REQUEST_WAIT)
|
||||
|
||||
await self.maybe_tidy_workers(workers)
|
||||
|
||||
worker = asyncio.create_task(
|
||||
requestor.process(request, responder)
|
||||
)
|
||||
|
||||
workers.append(worker)
|
||||
|
||||
async def async_thread(self, ws, running, q):
|
||||
|
||||
# Worker threads, servicing
|
||||
workers = []
|
||||
|
||||
while running.get():
|
||||
|
||||
try:
|
||||
|
||||
if len(workers) > 0:
|
||||
await self.maybe_tidy_workers(workers)
|
||||
|
||||
# Get next request on queue
|
||||
id, svc, request = await asyncio.wait_for(q.get(), 1)
|
||||
|
||||
except TimeoutError:
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
# This is an internal working error, may not be recoverable
|
||||
print("Exception:", e)
|
||||
await ws.send_json({"id": id, "error": str(e)})
|
||||
break
|
||||
|
||||
try:
|
||||
print(id, svc, request)
|
||||
await self.start_request_task(ws, id, svc, request, workers)
|
||||
|
||||
except Exception as e:
|
||||
print("Exception2:", e)
|
||||
await ws.send_json({"error": str(e)})
|
||||
|
||||
running.stop()
|
||||
|
||||
async def listener(self, ws, running):
|
||||
|
||||
# The outstanding request queue, max size is MAX_QUEUE_SIZE
|
||||
q = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
|
||||
|
||||
async_task = asyncio.create_task(self.async_thread(
|
||||
ws, running, q
|
||||
))
|
||||
|
||||
async for msg in ws:
|
||||
|
||||
# On error, finish
|
||||
if msg.type == WSMsgType.TEXT:
|
||||
|
||||
try:
|
||||
|
||||
data = msg.json()
|
||||
|
||||
if data["service"] not in self.services:
|
||||
raise RuntimeError("Bad service")
|
||||
|
||||
if "request" not in data:
|
||||
raise RuntimeError("Bad message")
|
||||
|
||||
if "id" not in data:
|
||||
raise RuntimeError("Bad message")
|
||||
|
||||
await q.put(
|
||||
(data["id"], data["service"], data["request"])
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
await ws.send_json({"error": str(e)})
|
||||
continue
|
||||
|
||||
elif msg.type == WSMsgType.ERROR:
|
||||
break
|
||||
elif msg.type == WSMsgType.CLOSE:
|
||||
break
|
||||
else:
|
||||
break
|
||||
|
||||
running.stop()
|
||||
|
||||
await async_task
|
||||
Loading…
Add table
Add a link
Reference in a new issue