From d2dbe3410f393f25730c6712d6fc64a73777de9f Mon Sep 17 00:00:00 2001 From: bennykok Date: Sat, 27 Jan 2024 18:07:10 +0800 Subject: [PATCH] feat(builder): update builder for updating time log and migrate to class objects --- builder/modal-builder/src/template/app.py | 155 ++++++++++-------- builder/modal-builder/src/template/config.py | 2 +- .../src/template/data/snapshot.json | 13 +- 3 files changed, 96 insertions(+), 74 deletions(-) diff --git a/builder/modal-builder/src/template/app.py b/builder/modal-builder/src/template/app.py index 64111d5..23b7b04 100644 --- a/builder/modal-builder/src/template/app.py +++ b/builder/modal-builder/src/template/app.py @@ -1,14 +1,14 @@ from config import config import modal -from modal import Image, Mount, web_endpoint, Stub, asgi_app +from modal import Image, Mount, web_endpoint, Stub, asgi_app, method, enter, exit import json import urllib.request import urllib.parse from pydantic import BaseModel -from fastapi import FastAPI, Request +from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse from volume_setup import volumes - +from datetime import datetime # deploy_test = False import os @@ -158,89 +158,106 @@ image = Image.debian_slim() target_image = image if deploy_test else dockerfile_image -@stub.function(image=target_image, gpu=config["gpu"] - ,volumes=volumes -) -def run(input: Input): - import subprocess - import time - # Make sure that the ComfyUI API is available - print(f"comfy-modal - check server") +@stub.cls(image=target_image, gpu=config["gpu"] ,volumes=volumes, timeout=60 * 10, container_idle_timeout=60 * 5) +class ComfyDeployRunner: - command = ["python", "main.py", - "--disable-auto-launch", "--disable-metadata"] + @enter() + def setup(self): + import subprocess + import time + # Make sure that the ComfyUI API is available + print(f"comfy-modal - check server") - server_process = subprocess.Popen(command, cwd="/comfyui") + command = ["python", "main.py", + "--disable-auto-launch", "--disable-metadata"] - check_server( - f"http://{COMFY_HOST}", - COMFY_API_AVAILABLE_MAX_RETRIES, - COMFY_API_AVAILABLE_INTERVAL_MS, - ) + self.server_process = subprocess.Popen(command, cwd="/comfyui") - job_input = input + check_server( + f"http://{COMFY_HOST}", + COMFY_API_AVAILABLE_MAX_RETRIES, + COMFY_API_AVAILABLE_INTERVAL_MS, + ) - # print(f"comfy-modal - got input {job_input}") + @exit() + def cleanup(self, exc_type, exc_value, traceback): + self.server_process.terminate() - # Queue the workflow - try: - # job_input is the json input - queued_workflow = queue_workflow_comfy_deploy( - job_input) # queue_workflow(workflow) - prompt_id = queued_workflow["prompt_id"] - print(f"comfy-modal - queued workflow with ID {prompt_id}") - except Exception as e: - import traceback - print(traceback.format_exc()) - return {"error": f"Error queuing workflow: {str(e)}"} + @method() + def run(self, input: Input): + data = json.dumps({ + "run_id": input.prompt_id, + "status": "started", + "time": datetime.now().isoformat() + }).encode('utf-8') + req = urllib.request.Request(input.status_endpoint, data=data, method='POST') + urllib.request.urlopen(req) - # Poll for completion - print(f"comfy-modal - wait until image generation is complete") - retries = 0 - status = "" - try: - print("getting request") - while retries < COMFY_POLLING_MAX_RETRIES: - status_result = check_status(prompt_id=prompt_id) - # history = get_history(prompt_id) + job_input = input - # Exit the loop if we have found the history - # if prompt_id in history and history[prompt_id].get("outputs"): - # break + try: + queued_workflow = queue_workflow_comfy_deploy(job_input) # queue_workflow(workflow) + prompt_id = queued_workflow["prompt_id"] + print(f"comfy-modal - queued workflow with ID {prompt_id}") + except Exception as e: + import traceback + print(traceback.format_exc()) + return {"error": f"Error queuing workflow: {str(e)}"} - # Exit the loop if we have found the status both success or failed - if 'status' in status_result and (status_result['status'] == 'success' or status_result['status'] == 'failed'): - status = status_result['status'] - print(status) - break + # Poll for completion + print(f"comfy-modal - wait until image generation is complete") + retries = 0 + status = "" + try: + print("getting request") + while retries < COMFY_POLLING_MAX_RETRIES: + status_result = check_status(prompt_id=prompt_id) + # history = get_history(prompt_id) + + # Exit the loop if we have found the history + # if prompt_id in history and history[prompt_id].get("outputs"): + # break + + # Exit the loop if we have found the status both success or failed + if 'status' in status_result and (status_result['status'] == 'success' or status_result['status'] == 'failed'): + status = status_result['status'] + print(status) + break + else: + # Wait before trying again + time.sleep(COMFY_POLLING_INTERVAL_MS / 1000) + retries += 1 else: - # Wait before trying again - time.sleep(COMFY_POLLING_INTERVAL_MS / 1000) - retries += 1 - else: - return {"error": "Max retries reached while waiting for image generation"} - except Exception as e: - return {"error": f"Error waiting for image generation: {str(e)}"} + return {"error": "Max retries reached while waiting for image generation"} + except Exception as e: + return {"error": f"Error waiting for image generation: {str(e)}"} - print(f"comfy-modal - Finished, turning off") - server_process.terminate() + print(f"comfy-modal - Finished, turning off") - # Get the generated image and return it as URL in an AWS bucket or as base64 - # images_result = process_output_images(history[prompt_id].get("outputs"), job["id"]) - # result = {**images_result, "refresh_worker": REFRESH_WORKER} - result = {"status": status} - - return result - print("Running remotely on Modal!") + result = {"status": status} + return result + @web_app.post("/run") -async def bar(request_input: RequestInput): - # print(request_input) +async def post_run(request_input: RequestInput): if not deploy_test: - return run.remote(request_input.input) - # pass + # print(request_input.input.prompt_id, request_input.input.status_endpoint) + data = json.dumps({ + "run_id": request_input.input.prompt_id, + "status": "queued", + "time": datetime.now().isoformat() + }).encode('utf-8') + req = urllib.request.Request(request_input.input.status_endpoint, data=data, method='POST') + urllib.request.urlopen(req) + model = ComfyDeployRunner() + call = model.run.spawn(request_input.input) + + # call = run.spawn() + return {"call_id": call.object_id} + + return {"call_id": None} @stub.function(image=image ,volumes=volumes diff --git a/builder/modal-builder/src/template/config.py b/builder/modal-builder/src/template/config.py index 2ecf253..62c21a2 100644 --- a/builder/modal-builder/src/template/config.py +++ b/builder/modal-builder/src/template/config.py @@ -1,6 +1,6 @@ config = { "name": "my-app", - "deploy_test": "True", + "deploy_test": "False", "gpu": "T4", "public_model_volume": "model-store", "private_model_volume": "private-model-store", diff --git a/builder/modal-builder/src/template/data/snapshot.json b/builder/modal-builder/src/template/data/snapshot.json index 3c42edb..e99a2e8 100644 --- a/builder/modal-builder/src/template/data/snapshot.json +++ b/builder/modal-builder/src/template/data/snapshot.json @@ -1,5 +1,10 @@ { - "comfyui": "d0165d819afe76bd4e6bdd710eb5f3e571b6a804", - "git_custom_nodes": {}, - "file_custom_nodes": [] -} \ No newline at end of file + "comfyui": "d0165d819afe76bd4e6bdd710eb5f3e571b6a804", + "git_custom_nodes": { + "https://github.com/BennyKok/comfyui-deploy.git": { + "hash": "a838cb7ad425e5652c3931fbafdc886b53c48a22", + "disabled": false + } + }, + "file_custom_nodes": [] +}