diff --git a/orchestration.py b/orchestration.py index c368654..4a97e0e 100644 --- a/orchestration.py +++ b/orchestration.py @@ -1,8 +1,10 @@ import json import os import server +import traceback #Check for availability of workflow checkpointing import aiohttp +import asyncio from .workflowcheckpointing import post_prompt_remote web = server.web @@ -41,6 +43,9 @@ async def ready(request): async def websocket_loop(): async with aiohttp.ClientSession() as session: + if 'ORCHESTRATION_SERVER' not in os.environ: + while True: + await asyncio.sleep(60) async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"]) as ws: print("connected to server") async for msg in ws: @@ -48,27 +53,43 @@ async def websocket_loop(): print("got command") if msg.type == aiohttp.WSMsgType.TEXT: js = msg.json() + resp = {"message_id": js.get('message_id', 0)} 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) + out = await post_prompt_remote(MockRequest()) + resp['data'] = json.loads(out.body._value) case "queue": - resp = ps.prompt_queue.get_current_queue() + resp['data'] = ps.prompt_queue.get_current_queue() case "files": - #Return a list of files, not yet implemented - resp = "Not yet implemented" + resp['data'] = [f.name for f in os.scandir('fetches') if f.is_file()] + case "info": + resp['data'] = {} + if 'SALAD_MACHINE_ID' in os.environ: + resp['data']['machine_id'] = os.environ['SALAD_MACHINE_ID'] + else: + resp['data']['machine_id'] = os.environ.get('HOSTNAME', 'local') case _: - resp = "Unknown command" + resp = {"error": "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)}) + await ws.send_json({"Error": str(e)}) +async def try_websocket(): + while True: + try: + await websocket_loop() + except aiohttp.client_exceptions.ClientConnectorError: + print("disconnected") + except: + print(traceback.format_exc()) + await asyncio.sleep(5) + print("Attempting re connection") -process_loop = ps.loop.create_task(websocket_loop()) +process_loop = ps.loop.create_task(try_websocket()) diff --git a/workflowcheckpointing.py b/workflowcheckpointing.py index ced5738..f669576 100644 --- a/workflowcheckpointing.py +++ b/workflowcheckpointing.py @@ -9,6 +9,7 @@ import threading import logging import itertools import hashlib +import time import comfy.samplers import execution @@ -348,8 +349,10 @@ async def post_prompt_remote(request): f = asyncio.Future() index = max(completion_futures.keys(),default=0)+1 completion_futures[index] = f + start_time = time.perf_counter() base_res = await original_post_prompt(request) outputs = await f + execution_time = time.perf_counter() - start_time completion_futures.pop(index) if "SALAD_ORGANIZATION" in os.environ: async with aiohttp.ClientSession('https://storage-api.salad.com') as s: @@ -358,13 +361,16 @@ async def post_prompt_remote(request): with open(outputs[i], 'rb') as f: data = f.read() #TODO support uploads > 100MB/ memory optimizations - fd = {'file': data} + fd = {'file': data, 'sign': True} url = '/'.join([base_url_path, uid, 'outputs', outputs[i]]) async with s.put(url, headers=headers, data=fd) as r: - await r.text() + url = await r.json()['url'] outputs[i] = url - res = base_res.text[:-1] + ', "outputs": ' + json.dumps(outputs) + '}' - return server.web.Response(body=res) + json_output = json.loads(base_res.text) + json_output['outputs'] = outputs + json_output['execution_time'] = execution_time + json_output['machineid'] = os.environ.get('SALAD_MACHINE_ID', "local") + return server.web.Response(body=json.dumps(json_output)) #Dangerous object.__setattr__(prompt_route, 'handler', post_prompt_remote)