Fixed, working

This commit is contained in:
Cyber MacGeddon 2024-12-25 13:37:05 +00:00
parent 1ca46f43b0
commit cb639b7955
2 changed files with 35 additions and 40 deletions

View file

@ -8,6 +8,13 @@ from aiohttp import web, WSMsgType
from . socket import SocketEndpoint from . socket import SocketEndpoint
from . text_completion import TextCompletionRequestor from . text_completion import TextCompletionRequestor
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): class MuxEndpoint(SocketEndpoint):
def __init__( def __init__(
@ -20,31 +27,26 @@ class MuxEndpoint(SocketEndpoint):
endpoint_path=path, auth=auth, endpoint_path=path, auth=auth,
) )
# The outstanding request queue, max size is 10
self.q = asyncio.Queue(maxsize=10)
# Worker threads, servicing
self.workers = []
self.services = services self.services = services
async def start(self): async def start(self):
pass pass
async def maybe_tidy_workers(self): async def maybe_tidy_workers(self, workers):
while True: while True:
try: try:
await asyncio.wait_for( await asyncio.wait_for(
asyncio.shield(self.workers[0]), asyncio.shield(workers[0]),
0.05 WORKER_CLOSE_WAIT
) )
# worker[0] now stopped # worker[0] now stopped
self.workers = self.workers[1:] # FIXME: Delete reference???
del workers[0]
if len(self.workers) == 0: if len(workers) == 0:
break break
# Loop iterates to try the next worker # Loop iterates to try the next worker
@ -53,7 +55,7 @@ class MuxEndpoint(SocketEndpoint):
# worker[0] still running, move on # worker[0] still running, move on
break break
async def start_request_task(self, ws, id, svc, request): async def start_request_task(self, ws, id, svc, request, workers):
requestor = self.services[svc] requestor = self.services[svc]
@ -66,27 +68,30 @@ class MuxEndpoint(SocketEndpoint):
"complete": fin, "complete": fin,
}) })
# Wait for outstanding requests to go below 15 # Wait for outstanding requests to go below MAX_OUTSTANDING_REQUESTS
while len(self.workers) > 15: while len(workers) > MAX_OUTSTANDING_REQUESTS:
await asyncio.sleep(0.1) await asyncio.sleep(START_REQUEST_WAIT)
worker = asyncio.create_task( worker = asyncio.create_task(
requestor.process(request, responder) requestor.process(request, responder)
) )
self.workers.append(worker) workers.append(worker)
async def async_thread(self, ws, running): async def async_thread(self, ws, running, q):
# Worker threads, servicing
workers = []
while running.get(): while running.get():
try: try:
if len(self.workers) > 0: if len(workers) > 0:
await self.maybe_tidy_workers() await self.maybe_tidy_workers(workers)
# Get next request on queue # Get next request on queue
id, svc, request = await asyncio.wait_for(self.q.get(), 1) id, svc, request = await asyncio.wait_for(q.get(), 1)
except TimeoutError: except TimeoutError:
continue continue
@ -99,7 +104,7 @@ class MuxEndpoint(SocketEndpoint):
try: try:
print(id, svc, request) print(id, svc, request)
await self.start_request_task(ws, id, svc, request) await self.start_request_task(ws, id, svc, request, workers)
except Exception as e: except Exception as e:
print("Exception2:", e) print("Exception2:", e)
@ -108,6 +113,13 @@ class MuxEndpoint(SocketEndpoint):
running.stop() running.stop()
async def listener(self, ws, running): 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: async for msg in ws:
@ -127,7 +139,7 @@ class MuxEndpoint(SocketEndpoint):
if "id" not in data: if "id" not in data:
raise RuntimeError("Bad message") raise RuntimeError("Bad message")
await self.q.put( await q.put(
(data["id"], data["service"], data["request"]) (data["id"], data["service"], data["request"])
) )
@ -145,3 +157,4 @@ class MuxEndpoint(SocketEndpoint):
running.stop() running.stop()
await async_task

View file

@ -32,18 +32,6 @@ class SocketEndpoint:
break break
running.stop() running.stop()
async def async_thread(self, ws, running):
while running.get():
try:
await asyncio.sleep(1)
except TimeoutError:
continue
except Exception as e:
print(f"Exception: {str(e)}", flush=True)
async def handle(self, request): async def handle(self, request):
@ -59,12 +47,8 @@ class SocketEndpoint:
ws = web.WebSocketResponse() ws = web.WebSocketResponse()
await ws.prepare(request) await ws.prepare(request)
task = asyncio.create_task(self.async_thread(ws, running))
try: try:
await self.listener(ws, running) await self.listener(ws, running)
except Exception as e: except Exception as e:
print(e, flush=True) print(e, flush=True)
@ -72,8 +56,6 @@ class SocketEndpoint:
await ws.close() await ws.close()
await task
return ws return ws
async def start(self): async def start(self):