nodes/AustinMroz-ComfyUI-WorkflowCheckpointing
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

I attempt to keep socket open on exception

Browse Source 
Exceptions while executing a command are now caught and sent back to the
socket server
This commit is contained in:
Austin Mroz
2024-06-29 17:56:27 -05:00
parent 1089d6e68f
commit c3f11fa212
1 changed files with 26 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
orchestration.py
+26 -23
View File
@@ -44,28 +44,31 @@ async def websocket_loop():
async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"]) as ws:
print("connected to server")
async for msg in ws:
print("got command")
if msg.type == aiohttp.WSMsgType.TEXT:
breakpoint()
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")
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())