diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index f277e255..4fecdf56 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -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") diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/mux.py b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py new file mode 100644 index 00000000..786d31ff --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/mux.py @@ -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 + diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py index cee904a5..6e59f667 100644 --- a/trustgraph-flow/trustgraph/gateway/endpoint/manager.py +++ b/trustgraph-flow/trustgraph/gateway/endpoint/manager.py @@ -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): diff --git a/trustgraph-flow/trustgraph/gateway/endpoint/mux.py b/trustgraph-flow/trustgraph/gateway/endpoint/mux.py deleted file mode 100644 index 1afc3225..00000000 --- a/trustgraph-flow/trustgraph/gateway/endpoint/mux.py +++ /dev/null @@ -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