Fruther orchestration implementation

Responses to orchestration are now wrapped and contain always contain a
message id to keep requests syncronized.

On connection failure, a reconnection is attempted.

Execution time is tracked and included in response data.

Response urls are now signed

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())

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)