Austin Mroz
c3f11fa212
Exceptions while executing a command are now caught and sent back to the socket server
75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
import json
|
|
import os
|
|
import server
|
|
#Check for availability of workflow checkpointing
|
|
import aiohttp
|
|
from .workflowcheckpointing import post_prompt_remote
|
|
|
|
web = server.web
|
|
ps = server.PromptServer.instance
|
|
|
|
finished_startup = False
|
|
original_server_start = ps.start
|
|
async def server_start(address, port, verbose=True, call_on_start=None):
|
|
original_on_start= call_on_start
|
|
def on_start(*args, **kwargs):
|
|
if original_on_start is not None:
|
|
original_on_start(*args, **kwargs)
|
|
global finished_startup
|
|
finished_startup= True
|
|
return await original_server_start(address, port, verbose, on_start)
|
|
ps.start = server_start
|
|
|
|
|
|
@ps.routes.get("/health")
|
|
async def heath(request):
|
|
#while any of the server endpoints could likely be used
|
|
return web.json_response([])
|
|
|
|
@ps.routes.get("/startup")
|
|
async def startup(request):
|
|
if finished_startup:
|
|
return web.json_response([])
|
|
return web.Response(status=503)
|
|
|
|
@ps.routes.get("/ready")
|
|
async def ready(request):
|
|
current_queue = ps.prompt_queue.get_current_queue()
|
|
if len(current_queue[0]) == 0 and len(current_queue[1]) == 0:
|
|
return web.json_response(current_queue)
|
|
return web.json_response(current_queue, status=503)
|
|
|
|
async def websocket_loop():
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"]) as ws:
|
|
print("connected to server")
|
|
async for msg in ws:
|
|
try:
|
|
print("got command")
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
js = msg.json()
|
|
match js['command']:
|
|
case 'prompt':
|
|
#wrap as mock request
|
|
class MockRequest:
|
|
async def json(self):
|
|
return js['data']
|
|
resp = await post_prompt_remote(MockRequest())
|
|
resp = json.loads(resp.body._value)
|
|
case "queue":
|
|
resp = ps.prompt_queue.get_current_queue()
|
|
case "files":
|
|
#Return a list of files, not yet implemented
|
|
resp = "Not yet implemented"
|
|
case _:
|
|
resp = "Unknown command"
|
|
print(resp)
|
|
await ws.send_json(resp)
|
|
elif msg.type == aiohttp.WSMsgType.ERROR:
|
|
await ws.send_json("Error")
|
|
except Exception as e:
|
|
#NOTE: this will reraise if error was socket closing
|
|
ws.send_json({"Error": str(e)})
|
|
|
|
process_loop = ps.loop.create_task(websocket_loop())
|