From 8797b504af130f9b86c1ae9945f66cfc74d754c3 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Tue, 19 Aug 2025 17:59:20 +0100 Subject: [PATCH 01/14] Initial ComfyUI implementation with updated wrapper --- workers/comfyui-json/README.md | 179 +++++++++++++++++++++ workers/comfyui-json/__init__.py | 0 workers/comfyui-json/client.py | 98 +++++++++++ workers/comfyui-json/data_types.py | 82 ++++++++++ workers/comfyui-json/misc/test_prompts.txt | 34 ++++ workers/comfyui-json/server.py | 116 +++++++++++++ workers/comfyui-json/test_load.py | 8 + 7 files changed, 517 insertions(+) create mode 100644 workers/comfyui-json/README.md create mode 100644 workers/comfyui-json/__init__.py create mode 100644 workers/comfyui-json/client.py create mode 100644 workers/comfyui-json/data_types.py create mode 100644 workers/comfyui-json/misc/test_prompts.txt create mode 100644 workers/comfyui-json/server.py create mode 100644 workers/comfyui-json/test_load.py diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md new file mode 100644 index 0000000..472b1d1 --- /dev/null +++ b/workers/comfyui-json/README.md @@ -0,0 +1,179 @@ +# ComfyUI PyWorker + +This is the base PyWorker for ComfyUI. It provides a unified interface for running any ComfyUI workflow through a proxy-based architecture. + +## Endpoint + +The worker provides a single endpoint: + +- `/generate/sync`: Processes ComfyUI workflows using either predefined modifiers or custom workflow JSON + +## Request Format + +The worker accepts requests in the following format. Choose either modifier mode OR custom workflow mode: + +**Modifier Mode:** +```json +{ + "input": { + "request_id": "uuid-string", // optional - UUID generated if not provided + "modifier": "RawWorkflow", + "modifications": { + "prompt": "a beautiful landscape", + "width": 1024, + "height": 1024, + "steps": 20, + "seed": 123456789 + }, + "s3": { ... }, // optional + "webhook": { ... } // optional + }, + "expected_time": 30.0 +} +``` + +**Custom Workflow Mode:** +```json +{ + "input": { + "request_id": "uuid-string", // optional - UUID generated if not provided + "workflow_json": { + // Complete ComfyUI workflow JSON + }, + "s3": { ... }, // optional + "webhook": { ... } // optional + }, + "expected_time": 30.0 +} +``` + +## Request Fields + +### Required Fields + +- **`input`**: Contains the main workflow data +- **`input.request_id`**: Unique identifier for the request +- **`expected_time`**: Expected runtime in seconds on RTX4090 (defaults to 46.0 if not provided) + +### Workflow Mode (Choose One) + +You must provide either `modifier` OR `workflow_json`, but not both: + +#### Option 1: Modifier Mode +- **`input.modifier`**: Name of the predefined workflow modifier (e.g., "Text2Image") +- **`input.modifications`**: Parameters to pass to the modifier + +#### Option 2: Custom Workflow Mode +- **`input.workflow_json`**: Complete ComfyUI workflow JSON + +### Optional Fields + +- **`input.s3`**: S3 configuration for file storage +- **`input.webhook`**: Webhook configuration for notifications + +These configurations can be provided in the request JSON or via environment variables. Request-level configuration takes precedence over environment variables. + +#### S3 Configuration + +**Via Request JSON:** +```json +"s3": { + "access_key_id": "your-s3-access-key", + "secret_access_key": "your-s3-secret-access-key", + "endpoint_url": "https://my-endpoint.backblaze.com", + "bucket_name": "your-bucket", + "region": "us-east-1" +} +``` + +**Via Environment Variables:** +```bash +S3_ACCESS_KEY_ID=your-key +S3_SECRET_ACCESS_KEY=your-secret +S3_BUCKET_NAME=your-bucket +S3_ENDPOINT_URL=https://s3.amazonaws.com +S3_REGION=us-east-1 +``` + +#### Webhook Configuration + +**Via Request JSON:** +```json +"webhook": { + "url": "your-webhook-url", + "extra_params": { + "custom_field": "value" + } +} +``` + +**Via Environment Variables:** +```bash +WEBHOOK_URL=https://your-webhook.com # Default webhook URL +WEBHOOK_TIMEOUT=30 # Webhook timeout in seconds +``` + +## Examples + +### Basic Text-to-Image (Modifier Mode) + +```json +{ + "input": { + "modifier": "Text2Image", + "modifications": { + "prompt": "a cat sitting on a windowsill", + "width": 512, + "height": 512, + "steps": 20, + "seed": 42 + } + }, + "expected_time": 25.0 +} +``` + +### Custom Workflow Mode + +```json +{ + "input": { + "request_id": "67890", // optional - using custom ID for tracking + "workflow_json": { + "3": { + "inputs": { + "seed": 42, + "steps": 20, + "cfg": 8, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["5", 0] + }, + "class_type": "KSampler" + } + } + }, + "expected_time": 45.0 +} +``` + +## Expected Time Guidelines + +The `expected_time` field helps with resource planning and should reflect expected runtime on RTX4090: + +- **Simple text-to-image**: 15-30 seconds +- **Complex workflows with upscaling**: 60+ seconds +- **Video generation**: 180+ seconds +- **Default**: 46 seconds (if not specified) + +## Client Libraries + +See the test client examples for implementation details on how to integrate with the ComfyUI worker. + +--- + +See Vast's serverless documentation for more details on how to use ComfyUI with autoscaler. \ No newline at end of file diff --git a/workers/comfyui-json/__init__.py b/workers/comfyui-json/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py new file mode 100644 index 0000000..64289ec --- /dev/null +++ b/workers/comfyui-json/client.py @@ -0,0 +1,98 @@ +import logging +import uuid +import random +from urllib.parse import urljoin + +import requests + +from lib.test_utils import print_truncate_res +from utils.endpoint_util import Endpoint +from utils.ssl import get_cert_file_path + +logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s[%(levelname)-5s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +log = logging.getLogger(__file__) + + +def call_text2image_workflow( + endpoint_group_name: str, api_key: str, server_url: str +) -> None: + """Simple Text2Image using the new modifier-based approach""" + WORKER_ENDPOINT = "/generate/sync" + COST = 100 + + # Route to get worker URL + route_payload = { + "endpoint": endpoint_group_name, + "api_key": api_key, + "cost": COST, + } + response = requests.post( + urljoin(server_url, "/route/"), + json=route_payload, + timeout=4, + ) + response.raise_for_status() + message = response.json() + url = message["url"] + auth_data = dict( + signature=message["signature"], + cost=message["cost"], + endpoint=message["endpoint"], + reqnum=message["reqnum"], + url=message["url"], + ) + + # Build the new payload structure + payload = { + "input": { + "request_id": str(uuid.uuid4()), + "modifier": "RawWorkflow", # or whatever your Text2Image modifier is called + "modifications": { + "prompt": "a beautiful landscape with mountains and lakes", + "width": 1024, + "height": 1024, + "steps": 20, + "seed": random.randint(0, 2**32 - 1) + }, + "workflow_json": {} # Empty since using modifier approach + }, + "expected_time": 30.0 # Expected 30 seconds on RTX4090 + } + + req_data = dict(payload=payload, auth_data=auth_data) + url = urljoin(url, WORKER_ENDPOINT) + print(f"url: {url}") + + response = requests.post( + url, + json=req_data, + verify=get_cert_file_path(), + ) + response.raise_for_status() + print_truncate_res(str(response.json())) + + +if __name__ == "__main__": + from lib.test_utils import test_args + + args = test_args.parse_args() + endpoint_api_key = Endpoint.get_endpoint_api_key( + endpoint_name=args.endpoint_group_name, + account_api_key=args.api_key, + instance=args.instance, + ) + if endpoint_api_key: + try: + call_text2image_workflow( + api_key=endpoint_api_key, + endpoint_group_name=args.endpoint_group_name, + server_url=args.server_url, + ) + except Exception as e: + log.error(f"Error during API call: {e}") + else: + log.error(f"Failed to get API key for endpoint {args.endpoint_group_name}") \ No newline at end of file diff --git a/workers/comfyui-json/data_types.py b/workers/comfyui-json/data_types.py new file mode 100644 index 0000000..4f05a1a --- /dev/null +++ b/workers/comfyui-json/data_types.py @@ -0,0 +1,82 @@ +import sys +import json +import random +import dataclasses +import inspect +from typing import Dict, Any +from functools import cache +from math import ceil + +from lib.data_types import ApiPayload, JsonDataException + + +with open("workers/comfyui/misc/test_prompts.txt", "r") as f: + test_prompts = f.readlines() + +@dataclasses.dataclass +class ComfyWorkflowData(ApiPayload): + input: dict + expected_time: float = 46.0 # Default: 2x baseline (23s * 2) for RTX4090 + + @classmethod + def for_test(cls): + test_prompt = random.choice(test_prompts).rstrip() + return cls( + input={ + "request_id": f"test-{random.randint(1000, 9999)}", + "modifier": "RawWorkflow", + "modifications": { + "prompt": test_prompt, + "width": 1024, + "height": 1024, + "steps": 28, + "seed": random.randint(0, sys.maxsize), + } + }, + expected_time=25.0 # Test data: expect 25 seconds on RTX4090 (slightly above baseline) + ) + + def generate_payload_json(self) -> Dict[str, Any]: + # input is already a dict, just return it wrapped in the expected structure + return {"input": self.input} + + def count_workload(self) -> float: + """ + This needs review. We cannot reasonably predict the workload based on the inputs. We may be processing: + - Images + - Videos + - Audio... There may also be complex loops in the workflow. + + User will provide an expected time to complete and we will calculate equivalent cost + + Convert user-provided expected_time (RTX4090 seconds) to the old scoring system. + + The old system normalized to: 1024x1024, 28 steps = 200 tokens on RTX4090 + The old formula was: REQUEST_TIME_FOR_STANDARD_IMAGE * (time_ratio * 200) + + Now the user provides the expected request time directly. + Default expected_time is 46s (2x baseline) if not specified. + """ + # Baseline: standard image (1024x1024, 28 steps) = 23s = 200 tokens on RTX4090 + RTX4090_BASELINE_TIME = 23.0 # seconds for standard image on RTX4090 + BASELINE_TOKENS = 200 # tokens for standard image + + # Calculate time ratio compared to baseline + time_ratio = self.expected_time / RTX4090_BASELINE_TIME + + # Return workload score: time_ratio * baseline tokens + return time_ratio * BASELINE_TOKENS + + @classmethod + def from_json_msg(cls, json_msg: Dict[str, Any]) -> "ComfyWorkflowData": + # Extract required fields + if "input" not in json_msg: + raise JsonDataException({"input": "missing parameter"}) + + # expected_time is optional, uses default if not provided + expected_time = json_msg.get("expected_time", 46.0) # Default: 2x baseline + + return cls( + input=json_msg["input"], + expected_time=float(expected_time) + ) \ No newline at end of file diff --git a/workers/comfyui-json/misc/test_prompts.txt b/workers/comfyui-json/misc/test_prompts.txt new file mode 100644 index 0000000..cfb8f8c --- /dev/null +++ b/workers/comfyui-json/misc/test_prompts.txt @@ -0,0 +1,34 @@ +cartoon character of a person with a hoodie , in style of cytus and deemo, ork, gold chains, realistic anime cat, dripping black goo, lineage revolution style, thug life, cute anthropomorphic bunny, balrog, arknights, aliased, very buff, black and red and yellow paint, painting illustration collage style, character composition in vector with white background +stardew valley, fine details +2D Vector Illustration of a child with soccer ball Art for Sublimation, Design Art, Chrome Art, Painting and Stunning Artwork, Highly Detailed Digital Painting, Airbrush Art, Highly Detailed Digital Artwork, Dramatic Artwork, stained antique yellow copper paint, digital airbrush art, detailed by Mark Brooks, Chicano airbrush art, Swagger! snake Culture +realistic futuristic city-downtown with short buildings, sunset +seascape by Ray Collins and artgerm, front view of a perfect wave, sunny background, ultra detailed water +inspired by realflow-cinema4d editor features, create image of a transparent luxury cup with ice fruits and mint, connected with white, yellow and pink cream, Slow - High Speed MO Photography, YouTube Video Screenshot, Abstract Clay, Transparent Cup , molecular gastronomy, wheel, 3D fluid,Simulation rendering, still video, 4k polymer clay futras photography, very surreal, Houdini Fluid Simulation, hyperrealistic CGI and FLUIDS & MULTIPHYSICS SIMULATION effect, with Somali Stain Lurex, Metallic Jacquard, Gold Thread, Mulberry Silk, Toub Saree, Warm background, a fantastic image worthy of an award. +biker with backpack on his back riding a motorcycle, Style by Ade Santora, Oilpunk, Cover photo, craig mullins style, on the cover of a magazine, Outdoor Magazine, inspired by Alex Petruk APe, image of a male biker, Cover of an award-winning magazine, the man has a backpack, photo for magazine, with a backpack, magazine cover +generate a collage-style illustration inspired by the Procreate raster graphic editor, photographic illustration with the theme, 2D vector, art for textile sublimation, containing surrealistic cartoon cat wearing a baseball cap and jeans standing in front of a poster, inspired by Sadao Watanabe, Doraemon, Japanese cartoon style, Eichiro Oda, Iconic high detail character, Director: Nakahara Nantenbō, Kastuhiro Otomo, image detailed, by Miyamoto, Hidetaka Miyazaki, Katsuhiro illustration, 8k, masterpiece, Minimize noise and grain in photo quality without lose quality and increase brightness and lighting,Symmetry and Alignment, Avoid asymmetrical shapes and out-of-focus points. Focus and Sharpness: Make sure the image is focused and sharp and encourages the viewer to see it as a work of art printed on fabric. +fantasy medieval village world inside a glass sphere , high detail, fantasy, realistic, light effect, hyper detail, volumetric lighting, cinematic, macro, depth of field, blur, red light and clouds from the back, highly detailed epic cinematic concept art cg render made in maya, blender and photoshop, octane render, excellent composition, dynamic dramatic cinematic lighting, aesthetic, very inspirational, world inside a glass sphere by james gurney by artgerm with james jean, joe fenton and tristan eaton by ross tran, fine details +Iron Man, (Arnold Tsang, Toru Nakayama), Masterpiece, Studio Quality, 6k , toa, toaair, 1boy, glowing, axe, mecha, science_fiction, solo, weapon, jungle , green_background, nature, outdoors, solo, tree, weapon, mask, dynamic lighting, detailed shading, digital texture painting +(Pope Francis) wearing leather jacket is a DJ in a nightclub, mixing live on stage, giant mixing table, a masterpiece +Pope Francis wearing biker (leather jacket), a masterpiece +Luke Skywalker ordering a burger and fries from the Death Star canteen. +I want to generate a group avatar for a Feishu group chat. The role of this group is daily software technical communication. Now the subject technology stacks that members of this group discuss daily include: algorithms, data structures, optimization, functional programming, and the programming languages often discussed are: TypeScript, Java, python, etc. I hope this avatar has a simple aesthetic, this avatar is a single person avatar +portrait Anime black girl cute-fine-face, pretty face, realistic shaded Perfect face, fine details. Anime. realistic shaded lighting by Ilya Kuvshinov Giuseppe Dangelico Pino and Michael Garmash and Rob Rey, IAMAG premiere, WLOP matte print, cute freckles, masterpiece +young Disney socialite wearing a beige miniskirt, dark brown turtleneck sweater, small neckless, cute-fine-face, anime. illustration, realistic shaded perfect face, brown hair, grey eyes, fine details, realistic shaded lighting by ilya kuvshinov giuseppe dangelico pino and michael garmash and rob rey, iamag premiere, wlop matte print, a masterpiece +Cute small cat sitting in a movie theater eating chicken wiggs watching a movie ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render +Cute small dog sitting in a movie theater eating popcorn watching a movie ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render +fox bracelet made of buckskin with fox features, rich details, fine carvings, studio lighting +crane buckskin bracelet with crane features, rich details, fine carvings, studio lighting +london luxurious interior living-room, light walls +Parisian luxurious interior penthouse bedroom, dark walls, wooden panels +cute girl, crop-top, blond hair, black glasses, stretching, with background by greg rutkowski makoto shinkai kyoto animation key art feminine mid shot +houses in front, houses background, straight houses, digital art, smooth, sharp focus, gravity falls style, doraemon style, shinchan style, anime style +Simplified technical drawing, Leonardo da Vinci, Mechanical Dinosaur Skeleton, Minimalistic annotations, Hand-drawn illustrations, Basic design and engineering, Wonder and curiosity +High quality 8K painting impressionist style of a Japanese modern city street with a girl on the foreground wearing a traditional wedding dress with a fox mask, staring at the sky, daylight +a landscape from the Moon with the Earth setting on the horizon, realistic, detailed +Isometric Atlantis city,great architecture with columns, great details, ornaments,seaweed, blue ambiance, 3D cartoon style, soft light, 45° view +A hyper realistic avatar of a guy riding on a black honda cbr 650r in leather suit,high detail, high quality,8K,photo realism +the street of amedieval fantasy town, at dawn, dark, highly detailed +overwhelmingly beautiful eagle framed with vector flowers, long shiny wavy flowing hair, polished, ultra detailed vector floral illustration mixed with hyper realism, muted pastel colors, vector floral details in background, muted colors, hyper detailed ultra intricate overwhelming realism in detailed complex scene with magical fantasy atmosphere, no signature, no watermark +a highly detailed matte painting of a man on a hill watching a rocket launch in the distance by studio ghibli, makoto shinkai, by artgerm, by wlop, by greg rutkowski, volumetric lighting, octane render, 4 k resolution, trending on artstation, masterpiece | hyperrealism| highly detailed| insanely detailed| intricate| cinematic lighting| depth of field +electronik robot and ofice ,unreal engine, cozy indoor lighting, artstation, detailed, digital painting,cinematic,character design by mark ryden and pixar and hayao miyazaki, unreal 5, daz, hyperrealistic, octane render +exquisitely intricately detailed illustration, of a small world with a lake and a rainbow, inside a closed glass jar. diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py new file mode 100644 index 0000000..78aa804 --- /dev/null +++ b/workers/comfyui-json/server.py @@ -0,0 +1,116 @@ +import os +import logging +import dataclasses +import base64 +from typing import Optional, Union, Type + +from aiohttp import web, ClientResponse + +from lib.backend import Backend, LogAction +from lib.data_types import EndpointHandler +from lib.server import start_server +from .data_types import ComfyWorkflowData + + +MODEL_SERVER_URL = "http://127.0.0.1:18288" + +# This is the last log line that gets emitted once comfyui+extensions have been fully loaded +MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: http://127.0.0.1:18188" +MODEL_SERVER_ERROR_LOG_MSGS = [ + "MetadataIncompleteBuffer", # This error is emitted when the downloaded model is corrupted + "Value not in list: unet_name", # This error is emitted when the model file is not there at all +] + + +logging.basicConfig( + level=logging.DEBUG, + format="%(asctime)s[%(levelname)-5s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +log = logging.getLogger(__file__) + + +async def generate_client_response( + self, client_request: web.Request, model_response: ClientResponse + ) -> Union[web.Response, web.StreamResponse]: + # Check if the response is actually streaming based on response headers/content-type + is_streaming_response = ( + model_response.content_type == "text/event-stream" + or model_response.content_type == "application/x-ndjson" + or model_response.headers.get("Transfer-Encoding") == "chunked" + or "stream" in model_response.content_type.lower() + ) + + if is_streaming_response: + log.debug("Detected streaming response...") + res = web.StreamResponse() + res.content_type = model_response.content_type + await res.prepare(client_request) + async for chunk in model_response.content: + await res.write(chunk) + await res.write_eof() + log.debug("Done streaming response") + return res + else: + log.debug("Detected non-streaming response...") + content = await model_response.read() + return web.Response( + body=content, + status=model_response.status, + content_type=model_response.content_type + ) + + +@dataclasses.dataclass +class ComfyWorkflowHandler(EndpointHandler[ComfyWorkflowData]): + + @property + def endpoint(self) -> str: + return "/generate/sync" + + @property + def healthcheck_endpoint(self) -> Optional[str]: + return None + + @classmethod + def payload_cls(cls) -> Type[ComfyWorkflowData]: + return ComfyWorkflowData + + def make_benchmark_payload(self) -> ComfyWorkflowData: + return ComfyWorkflowData.for_test() + + async def generate_client_response( + self, client_request: web.Request, model_response: ClientResponse + ) -> Union[web.Response, web.StreamResponse]: + return await generate_client_response(client_request, model_response) + + +backend = Backend( + model_server_url=MODEL_SERVER_URL, + model_log_file=os.environ["MODEL_LOG"], + allow_parallel_requests=False, + benchmark_handler=ComfyWorkflowHandler( + benchmark_runs=3, benchmark_words=100 + ), + log_actions=[ + (LogAction.ModelLoaded, MODEL_SERVER_START_LOG_MSG), + (LogAction.Info, "Downloading:"), + *[ + (LogAction.ModelError, error_msg) + for error_msg in MODEL_SERVER_ERROR_LOG_MSGS + ], + ], +) + + +async def handle_ping(_): + return web.Response(body="pong") + + +routes = [ + web.post("/generate/sync", backend.create_handler(ComfyWorkflowHandler())), + web.get("/ping", handle_ping), +] + +if __name__ == "__main__": + start_server(backend, routes) diff --git a/workers/comfyui-json/test_load.py b/workers/comfyui-json/test_load.py new file mode 100644 index 0000000..c493f67 --- /dev/null +++ b/workers/comfyui-json/test_load.py @@ -0,0 +1,8 @@ +from lib.test_utils import test_load_cmd, test_args +from .data_types import ComfyWorkflowData + +WORKER_ENDPOINT = "/generate/sync" + + +if __name__ == "__main__": + test_load_cmd(ComfyWorkflowData, WORKER_ENDPOINT, arg_parser=test_args) From 08c88f75271ca86b347f7a0ca8f4659837685cd1 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Wed, 20 Aug 2025 09:34:09 +0100 Subject: [PATCH 02/14] Improve testability --- start_server.sh | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/start_server.sh b/start_server.sh index 8ef61a7..2d53525 100755 --- a/start_server.sh +++ b/start_server.sh @@ -41,20 +41,33 @@ echo_var DEBUG_LOG echo_var PYWORKER_LOG echo_var MODEL_LOG -env | grep _ >> /etc/environment; - +# Populate /etc/environment with quoted values +if ! grep -q "VAST" /etc/environment; then + env -0 | grep -zEv "^(HOME=|SHLVL=)|CONDA" | while IFS= read -r -d '' line; do + name=${line%%=*} + value=${line#*=} + printf '%s="%s"\n' "$name" "$value" + done > /etc/environment +fi if [ ! -d "$ENV_PATH" ] then echo "setting up venv" - curl -LsSf https://astral.sh/uv/install.sh | sh - source ~/.local/bin/env - git clone https://github.com/vast-ai/pyworker "$SERVER_DIR" + if ! which uv; then + curl -LsSf https://astral.sh/uv/install.sh | sh + source ~/.local/bin/env + fi - uv venv --managed-python "$WORKSPACE_DIR/worker-env" -p 3.10 - source "$WORKSPACE_DIR/worker-env/bin/activate" + # Fork testing + git clone "${PYWORKER_REPO:-https://github.com/vast-ai/pyworker}" "$SERVER_DIR" + if [[ -n ${PYWORKER_REF:-} ]]; then + (cd "$SERVER_DIR" && git checkout "$PYWORKER_REF") + fi - uv pip install -r vast-pyworker/requirements.txt + uv venv --managed-python "$ENV_PATH" -p 3.10 + source "$ENV_PATH/bin/activate" + + uv pip install -r "${SERVER_DIR}/requirements.txt" touch ~/.no_auto_tmux else From 636f17d27faaf07b5c5fbedd16aa7c10768ebf10 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Wed, 20 Aug 2025 09:57:07 +0100 Subject: [PATCH 03/14] Fix workflow modifier class --- workers/comfyui-json/data_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/comfyui-json/data_types.py b/workers/comfyui-json/data_types.py index 4f05a1a..5140d14 100644 --- a/workers/comfyui-json/data_types.py +++ b/workers/comfyui-json/data_types.py @@ -24,7 +24,7 @@ class ComfyWorkflowData(ApiPayload): return cls( input={ "request_id": f"test-{random.randint(1000, 9999)}", - "modifier": "RawWorkflow", + "modifier": "Text2Image", "modifications": { "prompt": test_prompt, "width": 1024, From f9fdf048844fd56ada0b8f22374237550534e766 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Wed, 20 Aug 2025 13:27:29 +0100 Subject: [PATCH 04/14] Fix signature --- workers/comfyui-json/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py index 78aa804..a2c0438 100644 --- a/workers/comfyui-json/server.py +++ b/workers/comfyui-json/server.py @@ -31,7 +31,7 @@ log = logging.getLogger(__file__) async def generate_client_response( - self, client_request: web.Request, model_response: ClientResponse + client_request: web.Request, model_response: ClientResponse ) -> Union[web.Response, web.StreamResponse]: # Check if the response is actually streaming based on response headers/content-type is_streaming_response = ( From 58b078f908a6de8b9bc323d0a7491a250850ffe4 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Wed, 20 Aug 2025 18:06:02 +0100 Subject: [PATCH 05/14] Fix modifier class --- workers/comfyui-json/client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index 64289ec..3d28e03 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -50,7 +50,7 @@ def call_text2image_workflow( payload = { "input": { "request_id": str(uuid.uuid4()), - "modifier": "RawWorkflow", # or whatever your Text2Image modifier is called + "modifier": "Text2Image", "modifications": { "prompt": "a beautiful landscape with mountains and lakes", "width": 1024, @@ -73,7 +73,7 @@ def call_text2image_workflow( verify=get_cert_file_path(), ) response.raise_for_status() - print_truncate_res(str(response.json())) + print(str(response.json())) if __name__ == "__main__": @@ -95,4 +95,4 @@ if __name__ == "__main__": except Exception as e: log.error(f"Error during API call: {e}") else: - log.error(f"Failed to get API key for endpoint {args.endpoint_group_name}") \ No newline at end of file + log.error(f"Failed to get API key for endpoint {args.endpoint_group_name}") From 3f4acb29faba10310d04c833d5c5893643306287 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Fri, 22 Aug 2025 15:20:15 +0100 Subject: [PATCH 06/14] Improved client exception handling --- workers/comfyui-json/client.py | 108 +++++++++++++++++++++++---------- 1 file changed, 77 insertions(+), 31 deletions(-) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index 3d28e03..acdb1f3 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -2,6 +2,7 @@ import logging import uuid import random from urllib.parse import urljoin +import json import requests @@ -21,6 +22,41 @@ def call_text2image_workflow( endpoint_group_name: str, api_key: str, server_url: str ) -> None: """Simple Text2Image using the new modifier-based approach""" + + def make_request(url: str, payload: dict, timeout: int = None, verify=True, context: str = "request"): + """Helper function for making requests with consistent error handling""" + try: + response = requests.post( + url, + json=payload, + timeout=timeout, + verify=verify + ) + response.raise_for_status() + return response.json() + + except requests.exceptions.HTTPError as http_err: + log.error(f"HTTP error occurred during {context}: {http_err}") + log.error(f"Status Code: {response.status_code}") + log.error("Response content:", response.text) + return None + except requests.exceptions.Timeout: + log.error(f"Timeout occurred during {context}: {url}") + return None + except requests.exceptions.ConnectionError: + log.error(f"Connection error occurred during {context}: {url}") + return None + except json.JSONDecodeError as json_err: + log.error(f"Failed to decode JSON response during {context}: {json_err}") + if 'response' in locals(): + print("Response content:", response.text) + return None + except Exception as err: + log.error(f"An unexpected error occurred during {context}: {err}") + if 'response' in locals(): + log.error("Response content (if available):", response.text) + return None + WORKER_ENDPOINT = "/generate/sync" COST = 100 @@ -30,24 +66,30 @@ def call_text2image_workflow( "api_key": api_key, "cost": COST, } - response = requests.post( - urljoin(server_url, "/route/"), - json=route_payload, + + # First request - get routing information + route_response = make_request( + url=urljoin(server_url, "/route/"), + payload=route_payload, timeout=4, - ) - response.raise_for_status() - message = response.json() - url = message["url"] - auth_data = dict( - signature=message["signature"], - cost=message["cost"], - endpoint=message["endpoint"], - reqnum=message["reqnum"], - url=message["url"], + context="route request" ) - # Build the new payload structure - payload = { + if route_response is None: + return None + + # Extract data from route response + url = route_response["url"] + auth_data = dict( + signature=route_response["signature"], + cost=route_response["cost"], + endpoint=route_response["endpoint"], + reqnum=route_response["reqnum"], + url=route_response["url"], + ) + + # Build the payload for the worker request + worker_payload = { "input": { "request_id": str(uuid.uuid4()), "modifier": "Text2Image", @@ -63,17 +105,19 @@ def call_text2image_workflow( "expected_time": 30.0 # Expected 30 seconds on RTX4090 } - req_data = dict(payload=payload, auth_data=auth_data) - url = urljoin(url, WORKER_ENDPOINT) - print(f"url: {url}") + req_data = dict(payload=worker_payload, auth_data=auth_data) + worker_url = urljoin(url, WORKER_ENDPOINT) + print(f"url: {worker_url}") - response = requests.post( - url, - json=req_data, + # Second request - call the worker endpoint + worker_response = make_request( + url=worker_url, + payload=req_data, verify=get_cert_file_path(), + context="worker request" ) - response.raise_for_status() - print(str(response.json())) + + return worker_response if __name__ == "__main__": @@ -85,14 +129,16 @@ if __name__ == "__main__": account_api_key=args.api_key, instance=args.instance, ) + if endpoint_api_key: - try: - call_text2image_workflow( - api_key=endpoint_api_key, - endpoint_group_name=args.endpoint_group_name, - server_url=args.server_url, - ) - except Exception as e: - log.error(f"Error during API call: {e}") + result = call_text2image_workflow( + api_key=endpoint_api_key, + endpoint_group_name=args.endpoint_group_name, + server_url=args.server_url, + ) + if result is None: + log.error("Text2Image workflow failed") + else: + print(result) else: log.error(f"Failed to get API key for endpoint {args.endpoint_group_name}") From b00bef547c4a5c265ae9c0665dc5a0938337b0e8 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Fri, 22 Aug 2025 17:08:42 +0100 Subject: [PATCH 07/14] Ensure uv env script is present before sourcing --- start_server.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/start_server.sh b/start_server.sh index 2d53525..fbeeeb1 100755 --- a/start_server.sh +++ b/start_server.sh @@ -71,7 +71,7 @@ then touch ~/.no_auto_tmux else - source ~/.local/bin/env + [[ -f ~/.local/bin/env ]] && source ~/.local/bin/env source "$WORKSPACE_DIR/worker-env/bin/activate" echo "environment activated" echo "venv: $VIRTUAL_ENV" From fc75a64684910b3f2d1a81eefd5d617322c0ada2 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Mon, 25 Aug 2025 17:56:27 +0100 Subject: [PATCH 08/14] Use MODEL_SERVER_URL environment variable --- workers/comfyui/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/comfyui/server.py b/workers/comfyui/server.py index 40ee389..4b5e025 100644 --- a/workers/comfyui/server.py +++ b/workers/comfyui/server.py @@ -13,7 +13,7 @@ from lib.server import start_server from .data_types import DefaultComfyWorkflowData, CustomComfyWorkflowData -MODEL_SERVER_URL = "http://0.0.0.0:38188" +MODEL_SERVER_URL = "http://127.0.0.1:18288" # API Wrapper Service # This is the last log line that gets emitted once comfyui+extensions have been fully loaded MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: http://127.0.0.1:18188" From 92ff412679bf4e2a138c073918a1950014ab0636 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Mon, 25 Aug 2025 17:57:32 +0100 Subject: [PATCH 09/14] Use MODEL_SERVER_URL environment variable --- workers/comfyui-json/server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py index a2c0438..45b3570 100644 --- a/workers/comfyui-json/server.py +++ b/workers/comfyui-json/server.py @@ -12,7 +12,7 @@ from lib.server import start_server from .data_types import ComfyWorkflowData -MODEL_SERVER_URL = "http://127.0.0.1:18288" +MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL", "http://127.0.0.1:18288") # This is the last log line that gets emitted once comfyui+extensions have been fully loaded MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: http://127.0.0.1:18188" From ba74ac8136e1e6f011fc86b6fd5fe0425d2f63d0 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Mon, 25 Aug 2025 17:58:22 +0100 Subject: [PATCH 10/14] Use cost value 1 for all jobs --- workers/comfyui-json/README.md | 44 ++++++++++++++---------- workers/comfyui-json/data_types.py | 55 +++++++++--------------------- 2 files changed, 42 insertions(+), 57 deletions(-) diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md index 472b1d1..8d8097a 100644 --- a/workers/comfyui-json/README.md +++ b/workers/comfyui-json/README.md @@ -2,6 +2,28 @@ This is the base PyWorker for ComfyUI. It provides a unified interface for running any ComfyUI workflow through a proxy-based architecture. +The cost for each request has a static value of `1`. ComfyUI does not handle concurrent workloads and there is no current provision to load multiple instances of ComfyUI per worker node. + +## Requirements + +This worker requires both [ComfyUI](https://github.com/comfyanonymous/ComfyUI) and [ComfyUI API Wrapper](https://github.com/ai-dock/comfyui-api-wrapper). + +A docker image is provided but you may use any if the above requirements are met. + +## Benchmarking + +A simple image generation benchmark is run when the worker is first loaded. + +This benchmark uses SD v1.5 in the default text to image workflow provided by ComfyUI. The following variables can be used to alter the complexity and running time of the benchmark: + +| Environment Variable | Default Value | +| -------------------- | ------------- | +| BENCHMARK_TEST_WIDTH | 512 | +| BENCHMARK_TEST_HEIGHT | 512 | +| BENCHMARK_TEST_STEPS | 20 | + +The prompt will be randomly selected from the file in misc/test_prompts.txt and a random seed used for every run of the benchmark. + ## Endpoint The worker provides a single endpoint: @@ -27,8 +49,7 @@ The worker accepts requests in the following format. Choose either modifier mode }, "s3": { ... }, // optional "webhook": { ... } // optional - }, - "expected_time": 30.0 + } } ``` @@ -42,8 +63,7 @@ The worker accepts requests in the following format. Choose either modifier mode }, "s3": { ... }, // optional "webhook": { ... } // optional - }, - "expected_time": 30.0 + } } ``` @@ -53,7 +73,6 @@ The worker accepts requests in the following format. Choose either modifier mode - **`input`**: Contains the main workflow data - **`input.request_id`**: Unique identifier for the request -- **`expected_time`**: Expected runtime in seconds on RTX4090 (defaults to 46.0 if not provided) ### Workflow Mode (Choose One) @@ -128,8 +147,7 @@ WEBHOOK_TIMEOUT=30 # Webhook timeout in seconds "steps": 20, "seed": 42 } - }, - "expected_time": 25.0 + } } ``` @@ -156,20 +174,10 @@ WEBHOOK_TIMEOUT=30 # Webhook timeout in seconds "class_type": "KSampler" } } - }, - "expected_time": 45.0 + } } ``` -## Expected Time Guidelines - -The `expected_time` field helps with resource planning and should reflect expected runtime on RTX4090: - -- **Simple text-to-image**: 15-30 seconds -- **Complex workflows with upscaling**: 60+ seconds -- **Video generation**: 180+ seconds -- **Default**: 46 seconds (if not specified) - ## Client Libraries See the test client examples for implementation details on how to integrate with the ComfyUI worker. diff --git a/workers/comfyui-json/data_types.py b/workers/comfyui-json/data_types.py index 5140d14..e73163d 100644 --- a/workers/comfyui-json/data_types.py +++ b/workers/comfyui-json/data_types.py @@ -1,8 +1,7 @@ +import os import sys -import json import random import dataclasses -import inspect from typing import Dict, Any from functools import cache from math import ceil @@ -13,27 +12,33 @@ from lib.data_types import ApiPayload, JsonDataException with open("workers/comfyui/misc/test_prompts.txt", "r") as f: test_prompts = f.readlines() +def count_workload() -> float: + # Always 1.0 where there is a single instance of ComfyUI handling requests + return 1.0 + @dataclasses.dataclass class ComfyWorkflowData(ApiPayload): input: dict - expected_time: float = 46.0 # Default: 2x baseline (23s * 2) for RTX4090 @classmethod def for_test(cls): + """ + Use the variables available to simulate workflows of the required running time + Example: SD1.5, simple image gen 10000 steps, 512px x 512px will run for approximately 9 minutes @ ~18 it/s (RTX 4090) + """ test_prompt = random.choice(test_prompts).rstrip() return cls( input={ - "request_id": f"test-{random.randint(1000, 9999)}", + "request_id": f"test-{random.randint(1000, 99999)}", "modifier": "Text2Image", "modifications": { "prompt": test_prompt, - "width": 1024, - "height": 1024, - "steps": 28, + "width": os.getenv('BENCHMARK_TEST_WIDTH', 512), + "height": os.getenv('BENCHMARK_TEST_HEIGHT', 512), + "steps": os.getenv('BENCHMARK_TEST_STEPS', 20), "seed": random.randint(0, sys.maxsize), } - }, - expected_time=25.0 # Test data: expect 25 seconds on RTX4090 (slightly above baseline) + } ) def generate_payload_json(self) -> Dict[str, Any]: @@ -41,31 +46,7 @@ class ComfyWorkflowData(ApiPayload): return {"input": self.input} def count_workload(self) -> float: - """ - This needs review. We cannot reasonably predict the workload based on the inputs. We may be processing: - - Images - - Videos - - Audio... There may also be complex loops in the workflow. - - User will provide an expected time to complete and we will calculate equivalent cost - - Convert user-provided expected_time (RTX4090 seconds) to the old scoring system. - - The old system normalized to: 1024x1024, 28 steps = 200 tokens on RTX4090 - The old formula was: REQUEST_TIME_FOR_STANDARD_IMAGE * (time_ratio * 200) - - Now the user provides the expected request time directly. - Default expected_time is 46s (2x baseline) if not specified. - """ - # Baseline: standard image (1024x1024, 28 steps) = 23s = 200 tokens on RTX4090 - RTX4090_BASELINE_TIME = 23.0 # seconds for standard image on RTX4090 - BASELINE_TOKENS = 200 # tokens for standard image - - # Calculate time ratio compared to baseline - time_ratio = self.expected_time / RTX4090_BASELINE_TIME - - # Return workload score: time_ratio * baseline tokens - return time_ratio * BASELINE_TOKENS + return count_workload() @classmethod def from_json_msg(cls, json_msg: Dict[str, Any]) -> "ComfyWorkflowData": @@ -73,10 +54,6 @@ class ComfyWorkflowData(ApiPayload): if "input" not in json_msg: raise JsonDataException({"input": "missing parameter"}) - # expected_time is optional, uses default if not provided - expected_time = json_msg.get("expected_time", 46.0) # Default: 2x baseline - return cls( - input=json_msg["input"], - expected_time=float(expected_time) + input=json_msg["input"] ) \ No newline at end of file From 16b414676ebd867cdd26bf134cbe51edb186d34c Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Mon, 25 Aug 2025 18:31:10 +0100 Subject: [PATCH 11/14] Use count_workload() function for cost --- workers/comfyui-json/client.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index acdb1f3..bed1b5d 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -9,6 +9,7 @@ import requests from lib.test_utils import print_truncate_res from utils.endpoint_util import Endpoint from utils.ssl import get_cert_file_path +from .data_types import count_workload logging.basicConfig( level=logging.DEBUG, @@ -58,7 +59,9 @@ def call_text2image_workflow( return None WORKER_ENDPOINT = "/generate/sync" - COST = 100 + + # This worker has concurrency = 1. All workloads have cost value 1.0 + COST = count_workload() # Route to get worker URL route_payload = { @@ -101,8 +104,7 @@ def call_text2image_workflow( "seed": random.randint(0, 2**32 - 1) }, "workflow_json": {} # Empty since using modifier approach - }, - "expected_time": 30.0 # Expected 30 seconds on RTX4090 + } } req_data = dict(payload=worker_payload, auth_data=auth_data) From 7c1a544b190fd0ed99a3b9b8d96942ea4445b617 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Tue, 26 Aug 2025 12:41:05 +0100 Subject: [PATCH 12/14] Improve error reporting when no ready workers --- workers/comfyui-json/client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/workers/comfyui-json/client.py b/workers/comfyui-json/client.py index bed1b5d..a7abee5 100644 --- a/workers/comfyui-json/client.py +++ b/workers/comfyui-json/client.py @@ -33,6 +33,7 @@ def call_text2image_workflow( timeout=timeout, verify=verify ) + response.raise_for_status() return response.json() @@ -81,6 +82,14 @@ def call_text2image_workflow( if route_response is None: return None + if "url" not in route_response or not route_response["url"]: + log.error("Error: No worker in 'Ready' state. Please wait while the serverless engine removes errored workers or finishes loading new workers.") + return None + + if "status" in route_response: + print(f"Autoscaler status: {route_response['status']}") + return None + # Extract data from route response url = route_response["url"] auth_data = dict( From 947fc5eea40cf03162268c757b4e3a2b1c6981cc Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Tue, 26 Aug 2025 12:41:30 +0100 Subject: [PATCH 13/14] Improve benchmarking explanation --- workers/comfyui-json/README.md | 39 +++++++++++++++++++++++++++------- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/workers/comfyui-json/README.md b/workers/comfyui-json/README.md index 8d8097a..bb07145 100644 --- a/workers/comfyui-json/README.md +++ b/workers/comfyui-json/README.md @@ -12,17 +12,40 @@ A docker image is provided but you may use any if the above requirements are met ## Benchmarking -A simple image generation benchmark is run when the worker is first loaded. +A simple image generation benchmark runs when each worker initializes to validate GPU performance and identify underperforming machines. -This benchmark uses SD v1.5 in the default text to image workflow provided by ComfyUI. The following variables can be used to alter the complexity and running time of the benchmark: +The benchmark uses Stable Diffusion v1.5 with ComfyUI's default text-to-image workflow. Configure the benchmark complexity and duration using these variables: -| Environment Variable | Default Value | -| -------------------- | ------------- | -| BENCHMARK_TEST_WIDTH | 512 | -| BENCHMARK_TEST_HEIGHT | 512 | -| BENCHMARK_TEST_STEPS | 20 | +| Environment Variable | Default Value | Description | +| -------------------- | ------------- | ----------- | +| BENCHMARK_TEST_WIDTH | 512 | Image width (pixels) | +| BENCHMARK_TEST_HEIGHT | 512 | Image height (pixels) | +| BENCHMARK_TEST_STEPS | 20 | Number of denoising steps | -The prompt will be randomly selected from the file in misc/test_prompts.txt and a random seed used for every run of the benchmark. +Each benchmark run uses a random prompt from `misc/test_prompts.txt` and a random seed to ensure consistent GPU load patterns. + +### Calibrating Benchmark Duration + +To screen for underperforming hardware, set `BENCHMARK_TEST_STEPS` to match your expected production workflow duration. This allows you to identify machines that won't meet performance requirements. + +**Example:** If your typical workflow should complete in 90 seconds on acceptable hardware: + +```bash +# 1. Measure it/sec on your reference machine +# RTX 4090 typically achieves ~43 it/sec with SD1.5 + +# 2. Calculate required steps +# 90 seconds × 43 it/sec = 3870 steps + +# 3. Configure benchmark +export BENCHMARK_TEST_STEPS=3870 + +# 4. Machines completing significantly slower than 90s indicate hardware issues +``` + +**Performance expectations:** +- Benchmark duration should remain consistent across identical GPU models +- Significant variation (>20%) may indicate thermal, power, or configuration issues ## Endpoint From 703435d10e5055d8daa87ecec4b4d293435c77a7 Mon Sep 17 00:00:00 2001 From: Rob Ballantyne Date: Tue, 26 Aug 2025 12:42:04 +0100 Subject: [PATCH 14/14] Improve MODEL_SERVER_START_* messages --- workers/comfyui-json/server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/workers/comfyui-json/server.py b/workers/comfyui-json/server.py index 45b3570..702330e 100644 --- a/workers/comfyui-json/server.py +++ b/workers/comfyui-json/server.py @@ -15,10 +15,10 @@ from .data_types import ComfyWorkflowData MODEL_SERVER_URL = os.getenv("MODEL_SERVER_URL", "http://127.0.0.1:18288") # This is the last log line that gets emitted once comfyui+extensions have been fully loaded -MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: http://127.0.0.1:18188" +MODEL_SERVER_START_LOG_MSG = "To see the GUI go to: " MODEL_SERVER_ERROR_LOG_MSGS = [ "MetadataIncompleteBuffer", # This error is emitted when the downloaded model is corrupted - "Value not in list: unet_name", # This error is emitted when the model file is not there at all + "Value not in list: ", # This error is emitted when the model file is not there at all ]