Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 robertvoy
|
1d4dcb985e |
@@ -1,3 +1,4 @@
|
||||
# These are supported funding model platforms
|
||||
|
||||
github: robertvoy
|
||||
buy_me_a_coffee: robertvoy
|
||||
|
||||
@@ -19,6 +19,6 @@ jobs:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
+1
-6
@@ -1,9 +1,4 @@
|
||||
bin/
|
||||
logs/
|
||||
gpu_config.json
|
||||
__pycache__/
|
||||
**/__pycache__/
|
||||
*.py[cod]
|
||||
node_modules/
|
||||
npm-debug.log*
|
||||
AGENTS.md
|
||||
__pycache__/
|
||||
@@ -51,6 +51,7 @@ ComfyUI Distributed supports three types of workers:
|
||||
## Requirements
|
||||
|
||||
- ComfyUI
|
||||
> Note: Desktop app not currently supported
|
||||
- Multiple NVIDIA GPUs
|
||||
> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
|
||||
- That's it
|
||||
@@ -91,9 +92,8 @@ Generate multiple images in the time it takes to generate one. Each worker uses
|
||||
1. Open your ComfyUI workflow
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
4. Enable workers in the UI
|
||||
5. Run the workflow!
|
||||
|
||||
### Parallel WAN Generation
|
||||
Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
|
||||
@@ -116,7 +116,7 @@ Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, w
|
||||
|
||||
|
||||
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
> [Download workflow](/workflows/distributed-txt2img.json)
|
||||
|
||||
1. Load your image
|
||||
2. Upscale with ESRGAN or similar
|
||||
@@ -147,40 +147,13 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
||||
Control your distributed cluster programmatically without opening the browser.
|
||||
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
|
||||
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
|
||||
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
||||
|
||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
|
||||
---
|
||||
|
||||
## FAQ
|
||||
|
||||
<details>
|
||||
@@ -195,12 +168,12 @@ No, it does not speed up the generation of a single image or video. Instead, it
|
||||
|
||||
<details>
|
||||
<summary>Does it work with the ComfyUI desktop app?</summary>
|
||||
Yes, it does now.
|
||||
Currently, it is not compatible with the ComfyUI desktop app.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I combine my RTX 5090 with a GTX 980 to get faster results?</summary>
|
||||
Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks.
|
||||
Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -208,11 +181,6 @@ Yes, you can combine different GPUs, but performance is optimized when using sim
|
||||
Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I use my main machine just to coordinate workers without rendering?</summary>
|
||||
Yes. Open the Distributed panel and uncheck the master toggle to run in orchestrator-only mode. The master will distribute work to workers but won't render locally. If all workers become unavailable, the master automatically re-enables to ensure your workflow still runs.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I make this work with my Docker setup?</summary>
|
||||
Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided.
|
||||
@@ -220,26 +188,6 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do
|
||||
|
||||
---
|
||||
|
||||
## Community projects
|
||||
|
||||
Community-maintained add-ons and forks that extend ComfyUI-Distributed.
|
||||
|
||||
### [ComfyUI-H3MC-Distributed](https://github.com/ymiyad/ComfyUI-H3MC-Distributed) by ymiyad
|
||||
|
||||
Connects H3 Motion Context with ComfyUI-Distributed to generate multiple candidate takes of the next video clip across your workers. Compare the results, pick your preferred take, and use it as context for the next clip.
|
||||
|
||||
Each saved candidate video includes its continuation latent, so you don't need a shared folder or separate latent files. This add-on speeds up iteration through parallel candidates; it does not split a single clip across GPUs.
|
||||
|
||||
### [ComfyUI-Deadline-Distributed](https://github.com/doubletwisted/ComfyUI-Deadline-Distributed) by doubletwisted
|
||||
|
||||
A fork of ComfyUI-Distributed with Thinkbox/AWS Deadline 10 integration for render farms. Claim Deadline workers from the Distributed panel, register them automatically, and collect their results back into your master workflow.
|
||||
|
||||
Requires Deadline 10 and the author's ComfyUI-Deadline-Plugin.
|
||||
|
||||
See each repository for installation instructions, requirements, and support.
|
||||
|
||||
---
|
||||
|
||||
## Disclaimer
|
||||
|
||||
This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats.
|
||||
@@ -256,6 +204,3 @@ Your support helps keep this project thriving.
|
||||
|
||||
Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
+56
-13
@@ -1,22 +1,65 @@
|
||||
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||
import os
|
||||
import sys
|
||||
|
||||
from .nodes.v3 import NODES
|
||||
from .runtime.bootstrap import initialize
|
||||
# Add the directory to Python path
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
|
||||
WEB_DIRECTORY = './web'
|
||||
# Patch ComfyUI execution validation for DistributedVideoCollector
|
||||
try:
|
||||
import execution
|
||||
from .distributed import ImageBatchDivider
|
||||
|
||||
# Store original validate function if it exists
|
||||
if hasattr(execution, 'validate_outputs'):
|
||||
original_validate_outputs = execution.validate_outputs
|
||||
|
||||
def patched_validate_outputs(executor, node_id, result, node_class):
|
||||
if node_class == ImageBatchDivider:
|
||||
return # Skip validation for our dynamic output node
|
||||
return original_validate_outputs(executor, node_id, result, node_class)
|
||||
|
||||
execution.validate_outputs = patched_validate_outputs
|
||||
else:
|
||||
# Fallback: patch the main execute method if validate_outputs doesn't exist
|
||||
if hasattr(execution.PromptExecutor, 'execute'):
|
||||
original_execute = execution.PromptExecutor.execute
|
||||
|
||||
def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
|
||||
# This is a more complex patch - for now just call original
|
||||
# The ByPassTypeTuple should handle most validation issues
|
||||
return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
|
||||
|
||||
execution.PromptExecutor.execute = patched_execute
|
||||
|
||||
except ImportError:
|
||||
pass # ComfyUI execution module not available during import
|
||||
|
||||
# Import everything needed from the main module
|
||||
from .distributed import (
|
||||
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
|
||||
class DistributedExtension(ComfyExtension):
|
||||
async def on_load(self) -> None:
|
||||
initialize()
|
||||
# Import utilities
|
||||
from .utils.config import ensure_config_exists, CONFIG_FILE
|
||||
from .utils.logging import debug_log
|
||||
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return list(NODES)
|
||||
# Import distributed upscale nodes
|
||||
from .distributed_upscale import (
|
||||
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
async def comfy_entrypoint() -> DistributedExtension:
|
||||
return DistributedExtension()
|
||||
ensure_config_exists()
|
||||
|
||||
# Merge node mappings
|
||||
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
|
||||
|
||||
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
from . import config_routes # noqa: F401
|
||||
from . import tunnel_routes # noqa: F401
|
||||
from . import worker_routes # noqa: F401
|
||||
from . import job_routes # noqa: F401
|
||||
from . import usdu_routes # noqa: F401
|
||||
@@ -1,277 +0,0 @@
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
try:
|
||||
from ..utils.config import config_transaction, load_config, save_config
|
||||
except ImportError:
|
||||
from ..utils.config import load_config
|
||||
|
||||
try:
|
||||
from ..utils.config import save_config
|
||||
except ImportError:
|
||||
def save_config(_config):
|
||||
return True
|
||||
|
||||
@asynccontextmanager
|
||||
async def config_transaction():
|
||||
config = load_config()
|
||||
original_snapshot = json.dumps(config, sort_keys=True)
|
||||
yield config
|
||||
if json.dumps(config, sort_keys=True) != original_snapshot:
|
||||
save_config(config)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host
|
||||
|
||||
|
||||
def _positive_int(value):
|
||||
return value > 0
|
||||
|
||||
|
||||
CONFIG_SCHEMA = {
|
||||
"workers": (list, None),
|
||||
"master": (dict, None),
|
||||
"settings": (dict, None),
|
||||
"tunnel": (dict, None),
|
||||
"managed_processes": (dict, None),
|
||||
"worker_timeout_seconds": (int, _positive_int),
|
||||
"debug": (bool, None),
|
||||
"auto_launch_workers": (bool, None),
|
||||
"stop_workers_on_master_exit": (bool, None),
|
||||
"master_delegate_only": (bool, None),
|
||||
"websocket_orchestration": (bool, None),
|
||||
"has_auto_populated_workers": (bool, None),
|
||||
}
|
||||
|
||||
_SETTINGS_FIELDS = {
|
||||
"worker_timeout_seconds",
|
||||
"debug",
|
||||
"auto_launch_workers",
|
||||
"stop_workers_on_master_exit",
|
||||
"master_delegate_only",
|
||||
"websocket_orchestration",
|
||||
"has_auto_populated_workers",
|
||||
}
|
||||
|
||||
_WORKER_FIELDS = [
|
||||
("enabled", None, False),
|
||||
("name", None, False),
|
||||
("port", None, False),
|
||||
("host", normalize_host, True),
|
||||
("cuda_device", None, True),
|
||||
("extra_args", None, True),
|
||||
("type", None, False),
|
||||
]
|
||||
|
||||
_MASTER_FIELDS = [
|
||||
("name", None, False),
|
||||
("host", normalize_host, True),
|
||||
("port", None, False),
|
||||
("cuda_device", None, False),
|
||||
("extra_args", None, False),
|
||||
]
|
||||
|
||||
|
||||
def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
|
||||
"""Apply a partial update to a target dict based on field rules."""
|
||||
for key, normalizer, remove_on_none in field_rules:
|
||||
if key not in data:
|
||||
continue
|
||||
value = data[key]
|
||||
if value is None and remove_on_none:
|
||||
target.pop(key, None)
|
||||
else:
|
||||
target[key] = normalizer(value) if (normalizer and value is not None) else value
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/config")
|
||||
async def get_config_endpoint(request):
|
||||
config = load_config()
|
||||
return web.json_response(config)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config")
|
||||
async def update_config_endpoint(request):
|
||||
"""Bulk config update with schema validation."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {e}", 400)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return await handle_api_error(request, "Config payload must be an object", 400)
|
||||
|
||||
validated_settings = {}
|
||||
validated_root = {}
|
||||
errors = []
|
||||
|
||||
for key, value in data.items():
|
||||
if key not in CONFIG_SCHEMA:
|
||||
errors.append(f"Unknown field: {key}")
|
||||
continue
|
||||
|
||||
expected_type, validator = CONFIG_SCHEMA[key]
|
||||
if not isinstance(value, expected_type):
|
||||
errors.append(f"{key}: expected {expected_type.__name__}")
|
||||
continue
|
||||
|
||||
if validator and not validator(value):
|
||||
errors.append(f"{key}: value {value!r} failed validation")
|
||||
continue
|
||||
|
||||
if key in _SETTINGS_FIELDS:
|
||||
validated_settings[key] = value
|
||||
else:
|
||||
validated_root[key] = value
|
||||
|
||||
if errors:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"error": errors,
|
||||
"message": "; ".join(errors),
|
||||
}, status=400)
|
||||
|
||||
try:
|
||||
async with config_transaction() as config:
|
||||
settings = config.setdefault("settings", {})
|
||||
settings.update(validated_settings)
|
||||
for key, value in validated_root.items():
|
||||
config[key] = value
|
||||
return web.json_response({"status": "success", "config": config})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
||||
async def queue_status_endpoint(request):
|
||||
"""Check if a job queue is initialized."""
|
||||
try:
|
||||
job_id = request.match_info['job_id']
|
||||
|
||||
# Import to ensure initialization
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
exists = job_id in prompt_server.distributed_pending_tile_jobs
|
||||
|
||||
debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
|
||||
return web.json_response({"exists": exists, "job_id": job_id})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
||||
async def update_worker_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
async with config_transaction() as config:
|
||||
worker_found = False
|
||||
workers = config.setdefault("workers", [])
|
||||
|
||||
for worker in workers:
|
||||
if worker["id"] == worker_id:
|
||||
_apply_field_patch(worker, data, _WORKER_FIELDS)
|
||||
worker_found = True
|
||||
break
|
||||
|
||||
if not worker_found:
|
||||
# If worker not found and all required fields are provided, create new worker
|
||||
if all(key in data for key in ["name", "port", "cuda_device"]):
|
||||
new_worker = {
|
||||
"id": worker_id,
|
||||
"name": data["name"],
|
||||
"host": normalize_host(data.get("host", "localhost")),
|
||||
"port": data["port"],
|
||||
"cuda_device": data["cuda_device"],
|
||||
"enabled": data.get("enabled", False),
|
||||
"extra_args": data.get("extra_args", ""),
|
||||
"type": data.get("type", "local")
|
||||
}
|
||||
workers.append(new_worker)
|
||||
else:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Worker {worker_id} not found and missing required fields for creation",
|
||||
404,
|
||||
)
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
||||
async def delete_worker_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
async with config_transaction() as config:
|
||||
workers = config.get("workers", [])
|
||||
|
||||
# Find and remove the worker
|
||||
worker_index = -1
|
||||
for i, worker in enumerate(workers):
|
||||
if worker["id"] == worker_id:
|
||||
worker_index = i
|
||||
break
|
||||
|
||||
if worker_index == -1:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Remove the worker
|
||||
removed_worker = workers.pop(worker_index)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"message": f"Worker {removed_worker.get('name', worker_id)} deleted"
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
|
||||
async def update_setting_endpoint(request):
|
||||
"""Updates a specific key in the settings object."""
|
||||
try:
|
||||
data = await request.json()
|
||||
key = data.get("key")
|
||||
value = data.get("value")
|
||||
|
||||
if not key or value is None:
|
||||
return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
|
||||
if key not in _SETTINGS_FIELDS:
|
||||
return await handle_api_error(request, f"Unknown setting: {key}", 400)
|
||||
|
||||
async with config_transaction() as config:
|
||||
if 'settings' not in config:
|
||||
config['settings'] = {}
|
||||
|
||||
config['settings'][key] = value
|
||||
|
||||
return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_master")
|
||||
async def update_master_endpoint(request):
|
||||
"""Updates master configuration."""
|
||||
try:
|
||||
data = await request.json()
|
||||
|
||||
async with config_transaction() as config:
|
||||
if 'master' not in config:
|
||||
config['master'] = {}
|
||||
_apply_field_patch(config['master'], data, _MASTER_FIELDS)
|
||||
|
||||
return web.json_response({"status": "success", "message": "Master configuration updated."})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
@@ -1,479 +0,0 @@
|
||||
import json
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import base64
|
||||
import binascii
|
||||
import time
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.image import pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import handle_api_error
|
||||
from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
|
||||
from ..utils import transfer_codec
|
||||
try:
|
||||
from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
|
||||
except ImportError:
|
||||
from .queue_orchestration import orchestrate_distributed_execution
|
||||
|
||||
def ensure_distributed_state():
|
||||
return None
|
||||
from .queue_request import parse_queue_request_payload
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
# Worker result envelopes accepted by POST /distributed/job_complete:
|
||||
#
|
||||
# application/json (legacy, one frame per request):
|
||||
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool,
|
||||
# "audio"?: {...} }
|
||||
#
|
||||
# multipart/form-data (compressed, many frames per request):
|
||||
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
|
||||
# "is_last", "audio"?: {...} }
|
||||
# media: one part per frame for still formats, or one video file for video formats
|
||||
|
||||
|
||||
def _decode_image_sync(image_path):
|
||||
"""Decode image/video file and compute hash in a threadpool worker."""
|
||||
import base64
|
||||
import hashlib
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(image_path)
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(image_path)
|
||||
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'}
|
||||
file_ext = os.path.splitext(full_path)[1].lower()
|
||||
|
||||
if file_ext in video_extensions:
|
||||
with open(full_path, 'rb') as f:
|
||||
file_data = f.read()
|
||||
mime_types = {
|
||||
'.mp4': 'video/mp4',
|
||||
'.avi': 'video/x-msvideo',
|
||||
'.mov': 'video/quicktime',
|
||||
'.mkv': 'video/x-matroska',
|
||||
'.webm': 'video/webm'
|
||||
}
|
||||
mime_type = mime_types.get(file_ext, 'video/mp4')
|
||||
image_data = f"data:{mime_type};base64,{base64.b64encode(file_data).decode('utf-8')}"
|
||||
else:
|
||||
with Image.open(full_path) as img:
|
||||
if img.mode not in ('RGB', 'RGBA'):
|
||||
img = img.convert('RGB')
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, format='PNG', compress_level=1)
|
||||
image_data = f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"image_data": image_data,
|
||||
"hash": file_hash,
|
||||
}
|
||||
|
||||
|
||||
def _check_file_sync(filename, expected_hash):
|
||||
"""Check file presence and hash in a threadpool worker."""
|
||||
import hashlib
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(filename)
|
||||
if not os.path.exists(full_path):
|
||||
return {
|
||||
"status": "success",
|
||||
"exists": False,
|
||||
}
|
||||
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"exists": True,
|
||||
"hash_matches": file_hash == expected_hash,
|
||||
}
|
||||
|
||||
|
||||
def _decode_canonical_png_tensor(image_payload):
|
||||
"""Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
|
||||
if not isinstance(image_payload, str) or not image_payload.strip():
|
||||
raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
|
||||
|
||||
encoded = image_payload.strip()
|
||||
if encoded.startswith("data:"):
|
||||
header, sep, data_part = encoded.partition(",")
|
||||
if not sep:
|
||||
raise ValueError("Field 'image' data URL is malformed.")
|
||||
if not header.lower().startswith("data:image/png;base64"):
|
||||
raise ValueError("Field 'image' must be a PNG data URL when using data:* format.")
|
||||
encoded = data_part
|
||||
|
||||
try:
|
||||
png_bytes = base64.b64decode(encoded, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError("Field 'image' is not valid base64 PNG data.") from exc
|
||||
|
||||
if not png_bytes:
|
||||
raise ValueError("Field 'image' decoded to empty PNG data.")
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(png_bytes)) as img:
|
||||
img = img.convert("RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
return ensure_contiguous(tensor)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
|
||||
|
||||
|
||||
def _decode_audio_payload(audio_payload):
|
||||
"""Decode canonical audio payload into an AUDIO dict."""
|
||||
from ..utils.audio_payload import decode_audio_payload
|
||||
|
||||
return decode_audio_payload(audio_payload)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/prepare_job")
|
||||
async def prepare_job_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id", 400)
|
||||
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
|
||||
debug_log(f"Prepared queue for job {multi_job_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/clear_memory")
|
||||
async def clear_memory_endpoint(request):
|
||||
debug_log("Received request to clear VRAM.")
|
||||
try:
|
||||
# Use ComfyUI's prompt server queue system like the /free endpoint does
|
||||
if hasattr(server.PromptServer.instance, 'prompt_queue'):
|
||||
server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
|
||||
server.PromptServer.instance.prompt_queue.set_flag("free_memory", True)
|
||||
debug_log("Set queue flags for memory clearing.")
|
||||
|
||||
# Wait a bit for the queue to process
|
||||
await asyncio.sleep(MEMORY_CLEAR_DELAY)
|
||||
|
||||
# Also do direct cleanup as backup, but with error handling
|
||||
import gc
|
||||
import comfy.model_management as mm
|
||||
|
||||
try:
|
||||
mm.unload_all_models()
|
||||
except AttributeError as e:
|
||||
debug_log(f"Warning during model unload: {e}")
|
||||
|
||||
try:
|
||||
mm.soft_empty_cache()
|
||||
except Exception as e:
|
||||
debug_log(f"Warning during cache clear: {e}")
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
debug_log("VRAM cleared successfully.")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared."})
|
||||
except Exception as e:
|
||||
# Even if there's an error, try to do basic cleanup
|
||||
import gc
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
debug_log(f"Partial VRAM clear completed with warning: {e}")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/queue")
|
||||
async def distributed_queue_endpoint(request):
|
||||
"""Queue a distributed workflow, mirroring the UI orchestration pipeline."""
|
||||
try:
|
||||
raw_payload = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
|
||||
try:
|
||||
payload = parse_queue_request_payload(raw_payload)
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
try:
|
||||
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
|
||||
payload.prompt,
|
||||
payload.workflow_meta,
|
||||
payload.client_id,
|
||||
enabled_worker_ids=payload.enabled_worker_ids,
|
||||
delegate_master=payload.delegate_master,
|
||||
trace_execution_id=payload.trace_execution_id,
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"number": prompt_number,
|
||||
"node_errors": node_errors,
|
||||
"worker_count": worker_count,
|
||||
"auto_prepare_supported": True,
|
||||
})
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/load_image")
|
||||
async def load_image_endpoint(request):
|
||||
"""Load an image or video file and return it as base64 data with hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
image_path = data.get("image_path")
|
||||
|
||||
if not image_path:
|
||||
return await handle_api_error(request, "Missing image_path", 400)
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _decode_image_sync, image_path)
|
||||
return web.json_response(payload)
|
||||
except FileNotFoundError:
|
||||
return await handle_api_error(request, f"File not found: {image_path}", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/check_file")
|
||||
async def check_file_endpoint(request):
|
||||
"""Check if a file exists and matches the given hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
filename = data.get("filename")
|
||||
expected_hash = data.get("hash")
|
||||
|
||||
if not filename or not expected_hash:
|
||||
return await handle_api_error(request, "Missing filename or hash", 400)
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _check_file_sync, filename, expected_hash)
|
||||
return web.json_response(payload)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
async def _enqueue_worker_items(multi_job_id, items):
|
||||
"""Put decoded items on the job queue, waiting briefly for the master to create it.
|
||||
|
||||
Returns the queue size after enqueueing, or None if the job never appeared.
|
||||
"""
|
||||
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
|
||||
while True:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
if pending is not None:
|
||||
for item in items:
|
||||
await pending.put(item)
|
||||
return pending.qsize()
|
||||
if time.monotonic() > deadline:
|
||||
return None
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
def _decode_compressed_chunk_sync(envelope, parts):
|
||||
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
|
||||
frames = transfer_codec.decode_chunk(envelope, parts)
|
||||
return [
|
||||
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
|
||||
for frame in frames
|
||||
]
|
||||
|
||||
|
||||
def _validate_compressed_envelope(envelope):
|
||||
errors = []
|
||||
if not isinstance(envelope, dict):
|
||||
return ["envelope: expected JSON object"]
|
||||
for key in ("job_id", "worker_id"):
|
||||
value = envelope.get(key)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
errors.append(f"{key}: expected non-empty string")
|
||||
start_index = envelope.get("start_index")
|
||||
if not isinstance(start_index, int) or isinstance(start_index, bool) or start_index < 0:
|
||||
errors.append("start_index: expected non-negative integer")
|
||||
if not isinstance(envelope.get("is_last"), bool):
|
||||
errors.append("is_last: expected boolean")
|
||||
try:
|
||||
if transfer_codec.normalize_format(envelope.get("format")) == transfer_codec.LEGACY_FORMAT:
|
||||
errors.append("format: legacy_png must be sent as JSON")
|
||||
except transfer_codec.TransferCodecError as exc:
|
||||
errors.append(f"format: {exc}")
|
||||
audio_payload = envelope.get("audio")
|
||||
if audio_payload is not None and not isinstance(audio_payload, dict):
|
||||
errors.append("audio: expected object when provided")
|
||||
return errors
|
||||
|
||||
|
||||
async def _job_complete_compressed(request):
|
||||
try:
|
||||
form = await request.post()
|
||||
except Exception as exc:
|
||||
too_large = getattr(web, "HTTPRequestEntityTooLarge", None)
|
||||
if too_large is not None and isinstance(exc, too_large):
|
||||
return await handle_api_error(
|
||||
request, f"Request body exceeds the master's upload limit (--max-upload-size): {exc}", 413)
|
||||
return await handle_api_error(request, f"Invalid multipart payload: {exc}", 400)
|
||||
|
||||
raw_envelope = form.get("envelope")
|
||||
if isinstance(raw_envelope, (bytes, bytearray)):
|
||||
raw_envelope = raw_envelope.decode("utf-8")
|
||||
elif hasattr(raw_envelope, "file"):
|
||||
raw_envelope = raw_envelope.file.read().decode("utf-8")
|
||||
try:
|
||||
envelope = json.loads(raw_envelope) if isinstance(raw_envelope, str) else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, f"envelope: invalid JSON ({exc})", 400)
|
||||
if envelope is None:
|
||||
return await handle_api_error(request, "envelope: missing field", 400)
|
||||
|
||||
errors = _validate_compressed_envelope(envelope)
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
|
||||
parts = []
|
||||
for field in form.getall("media", []):
|
||||
if hasattr(field, "file"):
|
||||
parts.append(field.file.read())
|
||||
elif isinstance(field, (bytes, bytearray)):
|
||||
parts.append(bytes(field))
|
||||
else:
|
||||
return await handle_api_error(request, "media: expected binary file parts", 400)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
tensors = await loop.run_in_executor(None, _decode_compressed_chunk_sync, envelope, parts)
|
||||
decoded_audio = _decode_audio_payload(envelope.get("audio")) if envelope.get("audio") is not None else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
multi_job_id = envelope["job_id"].strip()
|
||||
worker_id = envelope["worker_id"].strip()
|
||||
start_index = envelope["start_index"]
|
||||
is_last = envelope["is_last"]
|
||||
last_pos = len(tensors) - 1
|
||||
items = [
|
||||
{
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": start_index + pos,
|
||||
"is_last": bool(is_last and pos == last_pos),
|
||||
"audio": decoded_audio if pos == last_pos else None,
|
||||
}
|
||||
for pos, tensor in enumerate(tensors)
|
||||
]
|
||||
|
||||
queue_size = await _enqueue_worker_items(multi_job_id, items)
|
||||
if queue_size is None:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
|
||||
debug_log(
|
||||
f"job_complete received {envelope['format']} chunk - job_id: {multi_job_id}, worker: {worker_id}, "
|
||||
f"frames: {start_index}..{start_index + len(items) - 1}, "
|
||||
f"bytes: {sum(len(p) for p in parts)}, is_last: {is_last}, queue_size: {queue_size}"
|
||||
)
|
||||
return web.json_response({"status": "success", "frames": len(items)})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/job_complete")
|
||||
async def job_complete_endpoint(request):
|
||||
content_type = str(getattr(request, "content_type", "") or "").lower()
|
||||
if content_type.startswith("multipart/"):
|
||||
try:
|
||||
return await _job_complete_compressed(request)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return await handle_api_error(request, "Expected a JSON object body", 400)
|
||||
|
||||
try:
|
||||
job_id = data.get("job_id")
|
||||
worker_id = data.get("worker_id")
|
||||
batch_idx = data.get("batch_idx")
|
||||
image_payload = data.get("image")
|
||||
audio_payload = data.get("audio")
|
||||
is_last = data.get("is_last")
|
||||
|
||||
errors = []
|
||||
if not isinstance(job_id, str) or not job_id.strip():
|
||||
errors.append("job_id: expected non-empty string")
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
errors.append("worker_id: expected non-empty string")
|
||||
if not isinstance(batch_idx, int) or batch_idx < 0:
|
||||
errors.append("batch_idx: expected non-negative integer")
|
||||
image_provided = image_payload is not None
|
||||
audio_provided = audio_payload is not None
|
||||
if not image_provided and not audio_provided:
|
||||
errors.append("expected at least one of image or audio")
|
||||
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
|
||||
errors.append("image: expected non-empty base64 PNG string")
|
||||
if audio_provided and not isinstance(audio_payload, dict):
|
||||
errors.append("audio: expected object when provided")
|
||||
if not isinstance(is_last, bool):
|
||||
errors.append("is_last: expected boolean")
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
|
||||
try:
|
||||
tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
|
||||
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
multi_job_id = job_id.strip()
|
||||
worker_id = worker_id.strip()
|
||||
|
||||
queue_size = await _enqueue_worker_items(
|
||||
multi_job_id,
|
||||
[
|
||||
{
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
"audio": decoded_audio,
|
||||
}
|
||||
],
|
||||
)
|
||||
if queue_size is None:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
|
||||
debug_log(
|
||||
f"job_complete received canonical envelope - job_id: {multi_job_id}, "
|
||||
f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, "
|
||||
f"queue_size: {queue_size}"
|
||||
)
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
@@ -1,4 +0,0 @@
|
||||
# Orchestration helpers split out from queue_orchestration.py:
|
||||
# - prompt_transform.py: graph pruning + hidden input overrides
|
||||
# - media_sync.py: remote media/path normalization
|
||||
# - dispatch.py: worker probe + prompt dispatch
|
||||
@@ -1,268 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session, probe_worker
|
||||
try:
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
except ImportError:
|
||||
def trace_debug(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def trace_info(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
try:
|
||||
from ..schemas import parse_positive_int
|
||||
except ImportError:
|
||||
def parse_positive_int(value, default):
|
||||
try:
|
||||
parsed = int(value)
|
||||
return parsed if parsed > 0 else default
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
_least_busy_rr_index = 0
|
||||
|
||||
|
||||
async def worker_is_active(worker):
|
||||
"""Ping worker's /prompt endpoint to confirm it's reachable."""
|
||||
url = build_worker_url(worker)
|
||||
return await probe_worker(url, timeout=3.0) is not None
|
||||
|
||||
|
||||
async def worker_ws_is_active(worker):
|
||||
"""Ping worker's websocket endpoint to confirm it's reachable."""
|
||||
session = await get_client_session()
|
||||
url = build_worker_url(worker, "/distributed/worker_ws")
|
||||
try:
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
|
||||
await ws.close()
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
debug_log(f"[Distributed] Worker WS probe timed out: {url}")
|
||||
return False
|
||||
except aiohttp.ClientConnectorError:
|
||||
debug_log(f"[Distributed] Worker WS unreachable: {url}")
|
||||
return False
|
||||
except Exception as e:
|
||||
debug_log(f"[Distributed] Worker WS probe unexpected error: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def _probe_worker_active(worker, use_websocket, semaphore):
|
||||
async with semaphore:
|
||||
is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
|
||||
return worker, is_active
|
||||
|
||||
|
||||
async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
|
||||
"""Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
|
||||
request_id = uuid.uuid4().hex
|
||||
ws_payload = {
|
||||
"type": "dispatch_prompt",
|
||||
"request_id": request_id,
|
||||
"prompt": payload.get("prompt"),
|
||||
"workflow": payload.get("workflow"),
|
||||
"client_id": client_id,
|
||||
}
|
||||
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
|
||||
ws_url = f"{ws_url}/distributed/worker_ws"
|
||||
session = await get_client_session()
|
||||
|
||||
async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
|
||||
await ws.send_json(ws_payload)
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(msg.data or "{}")
|
||||
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
|
||||
if data.get("ok"):
|
||||
return
|
||||
error_text = data.get("error") or "Worker rejected websocket dispatch."
|
||||
validation_error = data.get("validation_error")
|
||||
node_errors = data.get("node_errors")
|
||||
if validation_error:
|
||||
error_text = f"{error_text} | validation_error={validation_error}"
|
||||
if node_errors:
|
||||
error_text = f"{error_text} | node_errors={node_errors}"
|
||||
raise RuntimeError(error_text)
|
||||
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
|
||||
raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
|
||||
|
||||
raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
|
||||
|
||||
|
||||
async def dispatch_worker_prompt(
|
||||
worker,
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id=None,
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
url = build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if extra_data:
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
if use_websocket:
|
||||
try:
|
||||
await _dispatch_via_websocket(
|
||||
worker_url,
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
},
|
||||
client_id,
|
||||
)
|
||||
return
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
raise
|
||||
|
||||
session = await get_client_session()
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
|
||||
|
||||
async def select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
):
|
||||
"""Probe workers and return (active_workers, updated_delegate_master)."""
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
|
||||
if trace_execution_id and workers:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
|
||||
)
|
||||
|
||||
probe_results = await asyncio.gather(
|
||||
*[
|
||||
_probe_worker_active(worker, use_websocket, probe_semaphore)
|
||||
for worker in workers
|
||||
]
|
||||
)
|
||||
|
||||
active_workers = []
|
||||
for worker, is_active in probe_results:
|
||||
if is_active:
|
||||
active_workers.append(worker)
|
||||
else:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
else:
|
||||
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
|
||||
if trace_execution_id and workers:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Worker probe complete: active={len(active_workers)}/{len(workers)}",
|
||||
)
|
||||
|
||||
if not active_workers and delegate_master:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.")
|
||||
else:
|
||||
debug_log("All workers offline while delegate-only requested; enabling master participation.")
|
||||
delegate_master = False
|
||||
|
||||
return active_workers, delegate_master
|
||||
|
||||
|
||||
def _extract_queue_remaining(payload):
|
||||
if not isinstance(payload, dict):
|
||||
return 0
|
||||
try:
|
||||
queue_remaining = int(payload.get("exec_info", {}).get("queue_remaining", 0))
|
||||
except (TypeError, ValueError):
|
||||
queue_remaining = 0
|
||||
return max(queue_remaining, 0)
|
||||
|
||||
|
||||
async def _probe_worker_queue(worker, semaphore, probe_timeout):
|
||||
async with semaphore:
|
||||
worker_url = build_worker_url(worker)
|
||||
payload = await probe_worker(worker_url, timeout=probe_timeout)
|
||||
if payload is None:
|
||||
return None
|
||||
return {
|
||||
"worker": worker,
|
||||
"queue_remaining": _extract_queue_remaining(payload),
|
||||
}
|
||||
|
||||
|
||||
def _select_idle_round_robin(statuses):
|
||||
global _least_busy_rr_index
|
||||
if not statuses:
|
||||
return None
|
||||
index = _least_busy_rr_index % len(statuses)
|
||||
_least_busy_rr_index += 1
|
||||
return statuses[index]
|
||||
|
||||
|
||||
async def select_least_busy_worker(
|
||||
workers,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
probe_timeout=3.0,
|
||||
):
|
||||
"""Select one worker by queue depth, round-robin among idle workers."""
|
||||
if not workers:
|
||||
return None
|
||||
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
statuses = await asyncio.gather(
|
||||
*[
|
||||
_probe_worker_queue(worker, probe_semaphore, probe_timeout)
|
||||
for worker in workers
|
||||
]
|
||||
)
|
||||
statuses = [status for status in statuses if status is not None]
|
||||
if not statuses:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
|
||||
else:
|
||||
log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
|
||||
return None
|
||||
|
||||
idle_statuses = [status for status in statuses if status["queue_remaining"] == 0]
|
||||
if idle_statuses:
|
||||
selected = _select_idle_round_robin(idle_statuses)
|
||||
else:
|
||||
selected = min(statuses, key=lambda status: status["queue_remaining"])
|
||||
|
||||
worker = selected["worker"]
|
||||
queue_remaining = selected["queue_remaining"]
|
||||
if trace_execution_id:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
|
||||
)
|
||||
else:
|
||||
debug_log(
|
||||
f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
|
||||
)
|
||||
return worker
|
||||
@@ -1,256 +0,0 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
|
||||
|
||||
LIKELY_FILENAME_RE = re.compile(
|
||||
r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|"
|
||||
r"wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma|latent|txt|vae|lora|embedding)"
|
||||
r"(\s*\[\w+\])?$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
MEDIA_FILE_RE = re.compile(
|
||||
r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma)(\s*\[\w+\])?$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_media_reference(value):
|
||||
"""Normalize one media string value to a path-like reference or None."""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
|
||||
if MEDIA_FILE_RE.search(cleaned):
|
||||
return cleaned
|
||||
return None
|
||||
|
||||
|
||||
def convert_paths_for_platform(obj, target_separator):
|
||||
"""Recursively normalize likely file paths for the worker platform separator."""
|
||||
if target_separator not in ("/", "\\"):
|
||||
return obj
|
||||
|
||||
def _convert(value):
|
||||
if isinstance(value, str):
|
||||
if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value):
|
||||
trimmed = value.strip()
|
||||
has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed))
|
||||
is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\")
|
||||
has_protocol = bool(re.match(r"^\w+://", trimmed))
|
||||
|
||||
# URLs are not local paths and should never be separator-normalized.
|
||||
if has_protocol:
|
||||
return trimmed
|
||||
|
||||
# Keep relative media-style paths in forward-slash form (Comfy-style annotated paths).
|
||||
if not has_drive and not is_absolute and not has_protocol and MEDIA_FILE_RE.search(trimmed):
|
||||
return re.sub(r"[\\]+", "/", trimmed)
|
||||
|
||||
if target_separator == "\\":
|
||||
return re.sub(r"[\\/]+", r"\\", trimmed)
|
||||
return re.sub(r"[\\/]+", "/", trimmed)
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return [_convert(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: _convert(item) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
return _convert(obj)
|
||||
|
||||
|
||||
def _find_media_references(prompt_obj):
|
||||
"""Find media file references in image/video/audio/file inputs used by worker prompts."""
|
||||
media_refs = set()
|
||||
for node in prompt_obj.values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
for key in ("image", "video", "audio", "file"):
|
||||
cleaned = _normalize_media_reference(inputs.get(key))
|
||||
if cleaned:
|
||||
media_refs.add(cleaned)
|
||||
return sorted(media_refs)
|
||||
|
||||
|
||||
def _rewrite_prompt_media_inputs(prompt_obj, worker_media_paths):
|
||||
"""Rewrite media string inputs to worker-local uploaded paths."""
|
||||
if not isinstance(worker_media_paths, dict) or not worker_media_paths:
|
||||
return
|
||||
|
||||
for node in prompt_obj.values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
if not isinstance(inputs, dict):
|
||||
continue
|
||||
for key in ("image", "video", "audio", "file"):
|
||||
value = inputs.get(key)
|
||||
cleaned = _normalize_media_reference(value)
|
||||
if not cleaned:
|
||||
continue
|
||||
worker_path = worker_media_paths.get(cleaned)
|
||||
if worker_path:
|
||||
inputs[key] = worker_path
|
||||
|
||||
|
||||
def _load_media_file_sync(filename):
|
||||
"""Load local media bytes and hash for worker upload sync."""
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(filename)
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(filename)
|
||||
|
||||
with open(full_path, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
file_hash = hashlib.md5(file_bytes).hexdigest()
|
||||
mime_type = mimetypes.guess_type(full_path)[0]
|
||||
if not mime_type:
|
||||
ext = os.path.splitext(full_path)[1].lower()
|
||||
if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}:
|
||||
mime_type = "video/mp4"
|
||||
else:
|
||||
mime_type = "image/png"
|
||||
return file_bytes, file_hash, mime_type
|
||||
|
||||
|
||||
async def fetch_worker_path_separator(worker, trace_execution_id=None):
|
||||
"""Best-effort fetch of a worker's path separator from /distributed/system_info."""
|
||||
url = build_worker_url(worker, "/distributed/system_info")
|
||||
session = await get_client_session()
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
|
||||
if resp.status != 200:
|
||||
return None
|
||||
payload = await resp.json()
|
||||
separator = ((payload or {}).get("platform") or {}).get("path_separator")
|
||||
return separator if separator in ("/", "\\") else None
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}")
|
||||
else:
|
||||
debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}")
|
||||
return None
|
||||
|
||||
|
||||
async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
|
||||
"""Upload one media file to worker iff missing or hash-mismatched."""
|
||||
session = await get_client_session()
|
||||
normalized = filename.replace("\\", "/")
|
||||
|
||||
check_url = build_worker_url(worker, "/distributed/check_file")
|
||||
try:
|
||||
async with session.post(
|
||||
check_url,
|
||||
json={"filename": normalized, "hash": file_hash},
|
||||
timeout=aiohttp.ClientTimeout(total=6),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
payload = await resp.json()
|
||||
if payload.get("exists") and payload.get("hash_matches"):
|
||||
return False, normalized
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
|
||||
else:
|
||||
debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
|
||||
|
||||
parts = normalized.split("/")
|
||||
clean_name = parts[-1]
|
||||
subfolder = "/".join(parts[:-1])
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type)
|
||||
form.add_field("type", "input")
|
||||
form.add_field("subfolder", subfolder)
|
||||
form.add_field("overwrite", "true")
|
||||
|
||||
upload_url = build_worker_url(worker, "/upload/image")
|
||||
async with session.post(
|
||||
upload_url,
|
||||
data=form,
|
||||
timeout=aiohttp.ClientTimeout(total=30),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
try:
|
||||
payload = await resp.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
|
||||
name = str((payload or {}).get("name") or clean_name).strip()
|
||||
subfolder = str((payload or {}).get("subfolder") or "").strip().replace("\\", "/").strip("/")
|
||||
worker_path = f"{subfolder}/{name}" if subfolder else name
|
||||
return True, worker_path
|
||||
|
||||
|
||||
async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
|
||||
"""Sync referenced media files from master to a remote worker before dispatch."""
|
||||
media_refs = _find_media_references(prompt_obj)
|
||||
if not media_refs:
|
||||
return
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
uploaded = 0
|
||||
skipped = 0
|
||||
missing = 0
|
||||
worker_media_paths = {}
|
||||
for filename in media_refs:
|
||||
try:
|
||||
file_bytes, file_hash, mime_type = await loop.run_in_executor(
|
||||
None, _load_media_file_sync, filename
|
||||
)
|
||||
except FileNotFoundError:
|
||||
missing += 1
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.")
|
||||
else:
|
||||
log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.")
|
||||
continue
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}")
|
||||
continue
|
||||
|
||||
try:
|
||||
changed, worker_path = await _upload_media_to_worker(
|
||||
worker,
|
||||
filename,
|
||||
file_bytes,
|
||||
file_hash,
|
||||
mime_type,
|
||||
trace_execution_id=trace_execution_id,
|
||||
)
|
||||
if worker_path:
|
||||
worker_media_paths[filename] = worker_path
|
||||
if changed:
|
||||
uploaded += 1
|
||||
else:
|
||||
skipped += 1
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
|
||||
|
||||
_rewrite_prompt_media_inputs(prompt_obj, worker_media_paths)
|
||||
|
||||
summary = (
|
||||
f"Media sync for worker {worker.get('id')}: "
|
||||
f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}"
|
||||
)
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, summary)
|
||||
else:
|
||||
debug_log(f"[Distributed] {summary}")
|
||||
@@ -1,579 +0,0 @@
|
||||
import json
|
||||
from collections import deque
|
||||
|
||||
from ...utils.logging import debug_log
|
||||
|
||||
|
||||
class PromptIndex:
|
||||
"""Cache prompt metadata for faster worker/master prompt preparation."""
|
||||
|
||||
def __init__(self, prompt_obj):
|
||||
self._prompt_json = json.dumps(prompt_obj)
|
||||
self.nodes_by_class = {}
|
||||
self.class_by_node = {}
|
||||
self.inputs_by_node = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
class_type = node.get("class_type")
|
||||
node_id_str = str(node_id)
|
||||
if class_type:
|
||||
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
|
||||
self.class_by_node[node_id_str] = class_type
|
||||
self.inputs_by_node[node_id_str] = node.get("inputs", {})
|
||||
self._upstream_cache = {}
|
||||
|
||||
def copy_prompt(self):
|
||||
return json.loads(self._prompt_json)
|
||||
|
||||
def nodes_for_class(self, class_name):
|
||||
return self.nodes_by_class.get(class_name, [])
|
||||
|
||||
def has_upstream(self, start_node_id, target_class):
|
||||
cache_key = (str(start_node_id), target_class)
|
||||
if cache_key in self._upstream_cache:
|
||||
return self._upstream_cache[cache_key]
|
||||
|
||||
visited = set()
|
||||
stack = [str(start_node_id)]
|
||||
while stack:
|
||||
node_id = stack.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
inputs = self.inputs_by_node.get(node_id, {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
upstream_id = str(value[0])
|
||||
if self.class_by_node.get(upstream_id) == target_class:
|
||||
self._upstream_cache[cache_key] = True
|
||||
return True
|
||||
if upstream_id in self.inputs_by_node:
|
||||
stack.append(upstream_id)
|
||||
|
||||
self._upstream_cache[cache_key] = False
|
||||
return False
|
||||
|
||||
|
||||
def _iter_prompt_nodes(prompt_obj):
|
||||
for node_id, node in prompt_obj.items():
|
||||
if isinstance(node, dict):
|
||||
yield str(node_id), node
|
||||
|
||||
|
||||
def find_nodes_by_class(prompt_obj, class_name):
|
||||
nodes = []
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") == class_name:
|
||||
nodes.append(node_id)
|
||||
return nodes
|
||||
|
||||
|
||||
def _find_downstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable downstream from the provided IDs."""
|
||||
adjacency = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
adjacency.setdefault(source_id, set()).add(str(node_id))
|
||||
|
||||
connected = set(start_ids)
|
||||
queue = deque(start_ids)
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
|
||||
if dependent not in connected:
|
||||
connected.add(dependent)
|
||||
queue.append(dependent)
|
||||
return connected
|
||||
|
||||
|
||||
def _create_numeric_id_generator(prompt_obj):
|
||||
"""Return a closure that yields new numeric string IDs."""
|
||||
max_id = 0
|
||||
for node_id in prompt_obj.keys():
|
||||
try:
|
||||
numeric = int(node_id)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
max_id = max(max_id, numeric)
|
||||
|
||||
counter = max_id
|
||||
|
||||
def _next_id():
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
return str(counter)
|
||||
|
||||
return _next_id
|
||||
|
||||
|
||||
def _find_upstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable upstream from start_ids, including start nodes."""
|
||||
connected = set(str(node_id) for node_id in start_ids)
|
||||
queue = deque(connected)
|
||||
while queue:
|
||||
node_id = queue.popleft()
|
||||
node = prompt_obj.get(node_id) or {}
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
if source_id in prompt_obj and source_id not in connected:
|
||||
connected.add(source_id)
|
||||
queue.append(source_id)
|
||||
return connected
|
||||
|
||||
|
||||
_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
|
||||
"PrimitiveBoolean",
|
||||
"PrimitiveFloat",
|
||||
"PrimitiveInt",
|
||||
"PrimitiveNode",
|
||||
"PrimitiveString",
|
||||
}
|
||||
|
||||
_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
|
||||
"LoadImage",
|
||||
}
|
||||
|
||||
_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
|
||||
_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
|
||||
|
||||
# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
|
||||
# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
|
||||
# only after its connected inputs recursively prove safe.
|
||||
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
|
||||
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
|
||||
"CreateList": ("inputs.",),
|
||||
}
|
||||
|
||||
# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
|
||||
_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
|
||||
|
||||
|
||||
def _get_delegate_master_node_class_mappings():
|
||||
"""Return ComfyUI node-class mappings when available."""
|
||||
if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
|
||||
return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
|
||||
try:
|
||||
import nodes as comfy_nodes # type: ignore
|
||||
except Exception: # pragma: no cover - depends on ComfyUI runtime imports
|
||||
return {}
|
||||
return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
|
||||
|
||||
|
||||
def _get_delegate_master_node_class(class_type):
|
||||
mappings = _get_delegate_master_node_class_mappings()
|
||||
return mappings.get(class_type) if isinstance(mappings, dict) else None
|
||||
|
||||
|
||||
def _normalize_delegate_master_return_type(return_type):
|
||||
if return_type is None:
|
||||
return ""
|
||||
return str(return_type).strip().upper()
|
||||
|
||||
|
||||
def _delegate_master_type_is_safe_scalar(type_name):
|
||||
return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
|
||||
|
||||
|
||||
def _delegate_master_type_is_safe_config(type_name):
|
||||
return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
|
||||
|
||||
|
||||
def _delegate_master_output_is_safe_scalar(class_type, output_index):
|
||||
"""Return True when a registered node output is lightweight config data."""
|
||||
if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
|
||||
return True
|
||||
node_class = _get_delegate_master_node_class(class_type)
|
||||
return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
|
||||
try:
|
||||
output_type = return_types[int(output_index)]
|
||||
except (IndexError, TypeError, ValueError):
|
||||
return False
|
||||
return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
|
||||
|
||||
|
||||
def _get_delegate_master_input_types(class_type):
|
||||
node_class = _get_delegate_master_node_class(class_type)
|
||||
input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
|
||||
if callable(input_types):
|
||||
try:
|
||||
input_types = input_types()
|
||||
except TypeError:
|
||||
return {}
|
||||
return input_types if isinstance(input_types, dict) else {}
|
||||
|
||||
|
||||
def _normalize_delegate_master_input_type(input_spec):
|
||||
if isinstance(input_spec, (list, tuple)) and input_spec:
|
||||
return _normalize_delegate_master_return_type(input_spec[0])
|
||||
return _normalize_delegate_master_return_type(input_spec)
|
||||
|
||||
|
||||
def _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
"""Return True when a registered downstream input expects config data."""
|
||||
for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
|
||||
if input_name.startswith(prefix):
|
||||
return True
|
||||
input_types = _get_delegate_master_input_types(class_type)
|
||||
for section_name in ("required", "optional"):
|
||||
section = input_types.get(section_name, {})
|
||||
if isinstance(section, dict) and input_name in section:
|
||||
input_type = _normalize_delegate_master_input_type(section[input_name])
|
||||
return _delegate_master_type_is_safe_config(input_type)
|
||||
return False
|
||||
|
||||
|
||||
def _is_delegate_master_always_retained_upstream_node(node):
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
class_type = node.get("class_type")
|
||||
return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
|
||||
|
||||
|
||||
def _is_delegate_master_retained_upstream_node(node, output_index=0):
|
||||
"""Return True for lightweight upstream nodes safe to keep on the master."""
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str):
|
||||
return False
|
||||
return (
|
||||
class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
|
||||
or class_type.startswith("Primitive")
|
||||
or _delegate_master_output_is_safe_scalar(class_type, output_index)
|
||||
)
|
||||
|
||||
|
||||
def _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
node_id,
|
||||
output_index,
|
||||
memo,
|
||||
visiting,
|
||||
):
|
||||
"""Return safe retained branch nodes, or None when the branch is not safe."""
|
||||
node_id = str(node_id)
|
||||
cache_key = (node_id, output_index)
|
||||
if cache_key in memo:
|
||||
cached = memo[cache_key]
|
||||
return None if cached is None else set(cached)
|
||||
if cache_key in visiting:
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
|
||||
node = prompt_obj.get(node_id)
|
||||
if not _is_delegate_master_retained_upstream_node(node, output_index):
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
|
||||
visiting.add(cache_key)
|
||||
retained = {node_id}
|
||||
inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
|
||||
class_type = node.get("class_type") if isinstance(node, dict) else None
|
||||
for input_name, value in inputs.items():
|
||||
if not (isinstance(value, list) and len(value) == 2):
|
||||
continue
|
||||
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
source_id = str(value[0])
|
||||
branch = _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
source_id,
|
||||
value[1],
|
||||
memo,
|
||||
visiting,
|
||||
)
|
||||
if branch is None:
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
retained.update(branch)
|
||||
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = frozenset(retained)
|
||||
return retained
|
||||
|
||||
|
||||
def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
|
||||
"""Return lightweight upstream nodes needed by kept delegate-master nodes."""
|
||||
connected = set()
|
||||
memo = {}
|
||||
for node_id in start_ids:
|
||||
node = prompt_obj.get(str(node_id)) or {}
|
||||
inputs = node.get("inputs", {})
|
||||
class_type = node.get("class_type") if isinstance(node, dict) else None
|
||||
for input_name, value in inputs.items():
|
||||
if not (isinstance(value, list) and len(value) == 2):
|
||||
continue
|
||||
source_node = prompt_obj.get(str(value[0]))
|
||||
if _is_delegate_master_always_retained_upstream_node(source_node):
|
||||
connected.add(str(value[0]))
|
||||
continue
|
||||
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
continue
|
||||
branch = _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
value[0],
|
||||
value[1],
|
||||
memo,
|
||||
set(),
|
||||
)
|
||||
if branch is not None:
|
||||
connected.update(branch)
|
||||
return connected
|
||||
|
||||
|
||||
def prune_prompt_for_worker(prompt_obj):
|
||||
"""Prune worker prompt to distributed nodes and their upstream dependencies."""
|
||||
collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
|
||||
upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
|
||||
distributed_ids = collector_ids + upscale_ids
|
||||
if not distributed_ids:
|
||||
return prompt_obj
|
||||
|
||||
connected = _find_upstream_nodes(prompt_obj, distributed_ids)
|
||||
pruned_prompt = {}
|
||||
for node_id in connected:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
# Generate IDs from the original prompt so we never reuse IDs from pruned downstream nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for dist_id in distributed_ids:
|
||||
if dist_id not in pruned_prompt:
|
||||
continue
|
||||
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
||||
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
||||
if has_removed_downstream:
|
||||
original_node = prompt_obj.get(str(dist_id), {})
|
||||
class_type = original_node.get("class_type")
|
||||
inputs = original_node.get("inputs", {})
|
||||
image_connected = class_type != "DistributedCollector" or (
|
||||
isinstance(inputs.get("images"), list)
|
||||
and len(inputs["images"]) == 2
|
||||
)
|
||||
audio_connected = (
|
||||
class_type == "DistributedCollector"
|
||||
and isinstance(inputs.get("audio"), list)
|
||||
and len(inputs["audio"]) == 2
|
||||
)
|
||||
|
||||
if image_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"images": [dist_id, 0]},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {"title": "Preview Image (auto-added)"},
|
||||
}
|
||||
elif audio_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"audio": [dist_id, 1]},
|
||||
"class_type": "PreviewAudio",
|
||||
"_meta": {"title": "Preview Audio (auto-added)"},
|
||||
}
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
nodes_to_keep.update(
|
||||
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
|
||||
)
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in nodes_to_keep:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
pruned_ids = set(pruned_prompt.keys())
|
||||
for node_id, node in pruned_prompt.items():
|
||||
inputs = node.get("inputs")
|
||||
if not inputs:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
if isinstance(input_value, list) and len(input_value) == 2:
|
||||
source_id = str(input_value[0])
|
||||
if source_id not in pruned_ids:
|
||||
inputs.pop(input_name, None)
|
||||
debug_log(
|
||||
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
||||
)
|
||||
|
||||
# Generate IDs from the original prompt to avoid ID collisions with pruned nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for collector_id in collector_ids:
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
|
||||
original_images = original_inputs.get("images")
|
||||
if not (isinstance(original_images, list) and len(original_images) == 2):
|
||||
continue
|
||||
placeholder_id = next_id()
|
||||
pruned_prompt[placeholder_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
"inputs": {
|
||||
"height": 64,
|
||||
"width": 64,
|
||||
"channels": 3,
|
||||
},
|
||||
"_meta": {
|
||||
"title": "Distributed Empty Image (auto-added)",
|
||||
},
|
||||
}
|
||||
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
|
||||
debug_log(
|
||||
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
|
||||
)
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
return job_map
|
||||
|
||||
|
||||
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedSeed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
):
|
||||
"""Configure DistributedCollector nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
||||
node.setdefault("inputs", {})["pass_through"] = True
|
||||
continue
|
||||
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
|
||||
def _override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
):
|
||||
"""Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
|
||||
|
||||
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedValue nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedValue"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
):
|
||||
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
||||
is_master = participant_id == "master"
|
||||
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
||||
enabled_json = json.dumps(enabled_worker_ids)
|
||||
|
||||
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
)
|
||||
_override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
@@ -1,418 +0,0 @@
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
from ..utils.config import load_config
|
||||
from ..utils.constants import (
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_master_url, build_master_callback_url
|
||||
from ..utils.trace_logger import trace_debug
|
||||
from .schemas import parse_positive_float, parse_positive_int
|
||||
from .orchestration.dispatch import (
|
||||
dispatch_worker_prompt,
|
||||
select_active_workers,
|
||||
select_least_busy_worker,
|
||||
)
|
||||
from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
|
||||
from .orchestration.prompt_transform import (
|
||||
PromptIndex,
|
||||
apply_participant_overrides,
|
||||
find_nodes_by_class,
|
||||
generate_job_id_map,
|
||||
prepare_delegate_master_prompt,
|
||||
prune_prompt_for_worker,
|
||||
)
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def _generate_execution_trace_id():
|
||||
return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
|
||||
|
||||
def ensure_distributed_state(server_instance=None):
|
||||
"""Ensure prompt_server has the state used by distributed queue orchestration."""
|
||||
ps = server_instance or prompt_server
|
||||
if not hasattr(ps, "distributed_pending_jobs"):
|
||||
ps.distributed_pending_jobs = {}
|
||||
if not hasattr(ps, "distributed_jobs_lock"):
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
|
||||
|
||||
# Initialize top-level distributed queue state at module import time.
|
||||
ensure_distributed_state()
|
||||
|
||||
|
||||
async def _ensure_distributed_queue(job_id):
|
||||
"""Ensure a queue exists for the given distributed job ID."""
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
|
||||
|
||||
def _resolve_enabled_workers(config, requested_ids=None):
|
||||
"""Return a list of worker configs that should participate."""
|
||||
workers = []
|
||||
for worker in config.get("workers", []):
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id:
|
||||
continue
|
||||
|
||||
if requested_ids is not None:
|
||||
if worker_id not in requested_ids:
|
||||
continue
|
||||
elif not worker.get("enabled", False):
|
||||
continue
|
||||
|
||||
raw_port = worker.get("port", worker.get("listen_port", 8188))
|
||||
try:
|
||||
port = int(raw_port or 8188)
|
||||
except (TypeError, ValueError):
|
||||
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
|
||||
port = 8188
|
||||
|
||||
workers.append(
|
||||
{
|
||||
"id": worker_id,
|
||||
"name": worker.get("name", worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
}
|
||||
)
|
||||
return workers
|
||||
|
||||
|
||||
def _resolve_orchestration_limits(config):
|
||||
"""Resolve bounded concurrency/timeouts for worker preparation pipeline."""
|
||||
settings = (config or {}).get("settings", {}) or {}
|
||||
worker_probe_concurrency = parse_positive_int(
|
||||
settings.get("worker_probe_concurrency"),
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
|
||||
)
|
||||
worker_prep_concurrency = parse_positive_int(
|
||||
settings.get("worker_prep_concurrency"),
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
media_sync_concurrency = parse_positive_int(
|
||||
settings.get("media_sync_concurrency"),
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
|
||||
)
|
||||
media_sync_timeout_seconds = parse_positive_float(
|
||||
settings.get("media_sync_timeout_seconds"),
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
|
||||
)
|
||||
return (
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
def _is_load_balance_enabled(value):
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return bool(value)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return False
|
||||
|
||||
|
||||
def _prompt_requests_load_balance(prompt_index):
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
inputs = prompt_index.inputs_by_node.get(node_id, {})
|
||||
if _is_load_balance_enabled(inputs.get("load_balance", False)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def _prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
trace_execution_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
):
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
worker_master_url = build_master_callback_url(
|
||||
worker,
|
||||
config=config,
|
||||
prompt_server_instance=prompt_server,
|
||||
)
|
||||
|
||||
worker_type = str(worker.get("type") or "local").strip().lower()
|
||||
is_remote_like = bool(worker.get("host")) and worker_type != "local"
|
||||
if is_remote_like:
|
||||
path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
|
||||
if path_separator:
|
||||
worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
|
||||
|
||||
worker_prompt = prune_prompt_for_worker(worker_prompt)
|
||||
worker_prompt = apply_participant_overrides(
|
||||
worker_prompt,
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
worker_master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if is_remote_like:
|
||||
async with media_sync_semaphore:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
sync_worker_media(worker, worker_prompt, trace_execution_id=trace_execution_id),
|
||||
timeout=media_sync_timeout_seconds,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
(
|
||||
f"Media sync timed out after {media_sync_timeout_seconds:.1f}s "
|
||||
f"for worker {worker.get('name')} ({worker.get('id')}); continuing dispatch."
|
||||
),
|
||||
)
|
||||
|
||||
return worker, worker_prompt
|
||||
|
||||
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=None,
|
||||
delegate_master=None,
|
||||
trace_execution_id=None,
|
||||
):
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
|
||||
(
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
(
|
||||
f"Orchestration start: requested_workers={len(workers)}, "
|
||||
f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
|
||||
f"websocket={use_websocket}, "
|
||||
f"probe_concurrency={worker_probe_concurrency}, "
|
||||
f"prep_concurrency={worker_prep_concurrency}, "
|
||||
f"media_sync_concurrency={media_sync_concurrency}, "
|
||||
f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
|
||||
f"load_balance={load_balance_requested}"
|
||||
),
|
||||
)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, delegate_master = await select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
|
||||
if load_balance_requested:
|
||||
candidate_workers = list(active_workers)
|
||||
if not delegate_master:
|
||||
# Include master in load balancing only when master participation is enabled.
|
||||
candidate_workers.append(
|
||||
{
|
||||
"id": "master",
|
||||
"name": "Master",
|
||||
"host": master_url,
|
||||
"type": "local",
|
||||
}
|
||||
)
|
||||
|
||||
selected_worker = None
|
||||
if candidate_workers:
|
||||
selected_worker = await select_least_busy_worker(
|
||||
candidate_workers,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
if selected_worker is None and candidate_workers:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance selection probe failed; using first available candidate.",
|
||||
)
|
||||
selected_worker = candidate_workers[0]
|
||||
|
||||
if selected_worker is not None and str(selected_worker.get("id")) == "master":
|
||||
# Master selected as least busy; run master workload only.
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance selected master for execution (workers skipped).",
|
||||
)
|
||||
elif selected_worker is not None:
|
||||
active_workers = [selected_worker]
|
||||
# Worker selected as least busy; keep master orchestrator-only for this run.
|
||||
delegate_master = True
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
|
||||
)
|
||||
else:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance requested but no execution candidates were available.",
|
||||
)
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
queue_result = await queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
queue_result["number"],
|
||||
0,
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = apply_participant_overrides(
|
||||
master_prompt,
|
||||
"master",
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
|
||||
)
|
||||
elif not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
|
||||
)
|
||||
else:
|
||||
master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
if active_workers:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Active distributed workers: "
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
|
||||
)
|
||||
worker_payloads = []
|
||||
if active_workers:
|
||||
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
|
||||
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
|
||||
worker_payloads = await asyncio.gather(
|
||||
*[
|
||||
_prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
)
|
||||
for worker in active_workers
|
||||
]
|
||||
)
|
||||
|
||||
if worker_payloads:
|
||||
await asyncio.gather(
|
||||
*[
|
||||
dispatch_worker_prompt(
|
||||
worker,
|
||||
wprompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
use_websocket=use_websocket,
|
||||
trace_execution_id=execution_trace_id,
|
||||
)
|
||||
for worker, wprompt in worker_payloads
|
||||
]
|
||||
)
|
||||
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
prompt_id = queue_result["prompt_id"]
|
||||
prompt_number = queue_result["number"]
|
||||
node_errors = queue_result.get("node_errors", {})
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
@@ -1,79 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueRequestPayload:
|
||||
prompt: Dict[str, Any]
|
||||
workflow_meta: Any
|
||||
client_id: str
|
||||
delegate_master: Optional[bool]
|
||||
enabled_worker_ids: List[str]
|
||||
auto_prepare: bool
|
||||
trace_execution_id: Optional[str]
|
||||
|
||||
|
||||
def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
"""Parse and validate /distributed/queue payload into a normalized shape."""
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
|
||||
auto_prepare_raw = data.get("auto_prepare", True)
|
||||
if not isinstance(auto_prepare_raw, bool):
|
||||
raise ValueError("auto_prepare must be a boolean when provided")
|
||||
auto_prepare = auto_prepare_raw
|
||||
|
||||
prompt = data.get("prompt")
|
||||
# Auto-prepare is always on server-side; keep the field for wire compatibility.
|
||||
if prompt is None:
|
||||
workflow_payload = data.get("workflow")
|
||||
if isinstance(workflow_payload, dict):
|
||||
candidate_prompt = workflow_payload.get("prompt")
|
||||
if isinstance(candidate_prompt, dict):
|
||||
prompt = candidate_prompt
|
||||
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
|
||||
enabled_ids_raw = data.get("enabled_worker_ids")
|
||||
workers_field = data.get("workers")
|
||||
if enabled_ids_raw is None and workers_field is not None:
|
||||
if not isinstance(workers_field, list):
|
||||
raise ValueError("Field 'workers' must be a list when provided")
|
||||
enabled_ids_raw = []
|
||||
for entry in workers_field:
|
||||
worker_id = entry.get("id") if isinstance(entry, dict) else entry
|
||||
if worker_id is not None:
|
||||
enabled_ids_raw.append(str(worker_id))
|
||||
|
||||
if enabled_ids_raw is None:
|
||||
raise ValueError("enabled_worker_ids required")
|
||||
else:
|
||||
if not isinstance(enabled_ids_raw, list):
|
||||
raise ValueError("enabled_worker_ids must be a list of worker IDs")
|
||||
enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
|
||||
|
||||
delegate_master = data.get("delegate_master")
|
||||
if delegate_master is not None and not isinstance(delegate_master, bool):
|
||||
raise ValueError("delegate_master must be a boolean when provided")
|
||||
|
||||
client_id = data.get("client_id")
|
||||
if not isinstance(client_id, str) or not client_id.strip():
|
||||
raise ValueError("client_id required")
|
||||
client_id = client_id.strip()
|
||||
|
||||
trace_execution_id = data.get("trace_execution_id")
|
||||
if trace_execution_id is not None:
|
||||
if not isinstance(trace_execution_id, str):
|
||||
raise ValueError("trace_execution_id must be a string when provided")
|
||||
trace_execution_id = trace_execution_id.strip() or None
|
||||
|
||||
return QueueRequestPayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=delegate_master,
|
||||
enabled_worker_ids=enabled_ids,
|
||||
auto_prepare=auto_prepare,
|
||||
trace_execution_id=trace_execution_id,
|
||||
)
|
||||
@@ -1,54 +0,0 @@
|
||||
def require_fields(data: dict, *fields) -> list[str]:
|
||||
"""Return field names that are missing or empty in a JSON object."""
|
||||
if not isinstance(data, dict):
|
||||
return list(fields)
|
||||
|
||||
missing = []
|
||||
for field in fields:
|
||||
if field not in data:
|
||||
missing.append(field)
|
||||
continue
|
||||
value = data.get(field)
|
||||
if value is None:
|
||||
missing.append(field)
|
||||
continue
|
||||
if isinstance(value, str) and not value.strip():
|
||||
missing.append(field)
|
||||
|
||||
return missing
|
||||
|
||||
|
||||
def validate_worker_id(worker_id: str, config: dict) -> bool:
|
||||
"""Return True when worker_id exists in config['workers']."""
|
||||
worker_id_str = str(worker_id)
|
||||
workers = (config or {}).get("workers", [])
|
||||
return any(str(worker.get("id")) == worker_id_str for worker in workers)
|
||||
|
||||
|
||||
def validate_positive_int(value, field_name: str) -> str | None:
|
||||
"""Validate positive integers and return an error string when invalid."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
if parsed <= 0:
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
return None
|
||||
|
||||
|
||||
def parse_positive_int(value, default: int) -> int:
|
||||
"""Parse value as positive int, returning default on failure."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return max(1, int(default))
|
||||
return max(1, parsed)
|
||||
|
||||
|
||||
def parse_positive_float(value, default: float) -> float:
|
||||
"""Parse value as positive float, returning default on failure."""
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return max(0.0, float(default))
|
||||
return max(0.0, parsed)
|
||||
@@ -1,51 +0,0 @@
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.cloudflare import cloudflare_tunnel_manager
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
|
||||
async def tunnel_status_endpoint(request):
|
||||
"""Return Cloudflare tunnel status and last known details."""
|
||||
try:
|
||||
status = cloudflare_tunnel_manager.get_status()
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host")
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": status,
|
||||
"master_host": master_host
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
|
||||
async def tunnel_start_endpoint(request):
|
||||
"""Start a Cloudflare tunnel pointing at the current ComfyUI server."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.start_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
|
||||
async def tunnel_stop_endpoint(request):
|
||||
"""Stop the managed Cloudflare tunnel if running."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.stop_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
@@ -1,228 +0,0 @@
|
||||
import asyncio
|
||||
import io
|
||||
import time
|
||||
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
import server
|
||||
|
||||
from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
|
||||
from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
|
||||
from ..upscale.payload_parsers import _parse_tiles_from_form
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if isinstance(job_data, BaseJobState):
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
debug_log(f"Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
return await handle_api_error(request, "Worker status tracking not available", 400)
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
||||
async def submit_tiles_endpoint(request):
|
||||
"""Endpoint for workers to submit processed tiles in static mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
|
||||
# Handle completion signal
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [],
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
try:
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
|
||||
q = job_data.queue
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last,
|
||||
})
|
||||
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [],
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
||||
async def submit_image_endpoint(request):
|
||||
"""Endpoint for workers to submit processed images in dynamic mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Handle image submission
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last,
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle completion signal (no image data)
|
||||
elif is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Missing image data or invalid request", 400)
|
||||
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
|
||||
mode = job_data.mode
|
||||
if isinstance(job_data, ImageJobState):
|
||||
pending_queue = job_data.pending_images
|
||||
elif isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid job configuration", 400)
|
||||
|
||||
try:
|
||||
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
remaining = pending_queue.qsize()
|
||||
|
||||
if mode == 'dynamic':
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
||||
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
||||
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
||||
return web.json_response({
|
||||
"tile_idx": task_idx,
|
||||
"estimated_remaining": remaining,
|
||||
"batched_static": job_data.batched_static,
|
||||
})
|
||||
except asyncio.TimeoutError:
|
||||
if mode == 'dynamic':
|
||||
return web.json_response({"image_idx": None})
|
||||
return web.json_response({"tile_idx": None})
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
|
||||
return web.json_response({"ready": ready})
|
||||
@@ -1,708 +0,0 @@
|
||||
import json
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import platform
|
||||
import subprocess
|
||||
import socket
|
||||
|
||||
import torch
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import (
|
||||
build_worker_url,
|
||||
get_client_session,
|
||||
get_server_port,
|
||||
handle_api_error,
|
||||
normalize_host,
|
||||
probe_worker,
|
||||
)
|
||||
from ..utils.constants import CHUNK_SIZE
|
||||
from ..workers import get_worker_manager
|
||||
from ..workers.ports import allocate_worker_ports
|
||||
from .schemas import require_fields, validate_worker_id
|
||||
from ..workers.detection import (
|
||||
get_machine_id,
|
||||
is_docker_environment,
|
||||
is_runpod_environment,
|
||||
)
|
||||
try:
|
||||
from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
|
||||
except ImportError:
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
def __init__(self, message, validation_error=None, node_errors=None):
|
||||
super().__init__(str(message))
|
||||
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_ws")
|
||||
async def worker_ws_endpoint(request):
|
||||
"""WebSocket endpoint for worker prompt dispatch."""
|
||||
ws = web.WebSocketResponse(heartbeat=30)
|
||||
await ws.prepare(request)
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
data = json.loads(msg.data or "{}")
|
||||
except json.JSONDecodeError:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": None,
|
||||
"ok": False,
|
||||
"error": "Invalid JSON payload.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Unsupported websocket message type.",
|
||||
})
|
||||
continue
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Field 'prompt' must be an object.",
|
||||
})
|
||||
continue
|
||||
|
||||
try:
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": True,
|
||||
"prompt_id": prompt_id,
|
||||
})
|
||||
except PromptValidationError as exc:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
"validation_error": exc.validation_error,
|
||||
"node_errors": exc.node_errors,
|
||||
})
|
||||
except Exception as exc:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
})
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
log(f"[Distributed] Worker websocket error: {ws.exception()}")
|
||||
|
||||
return ws
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
|
||||
async def clear_launching_state(request):
|
||||
"""Clear the launching flag when worker is confirmed running."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Clear launching flag in managed processes
|
||||
if worker_id in wm.processes:
|
||||
if 'launching' in wm.processes[worker_id]:
|
||||
del wm.processes[worker_id]['launching']
|
||||
wm.save_processes()
|
||||
debug_log(f"Cleared launching state for worker {worker_id}")
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
def get_network_ips():
|
||||
"""Get all network IPs, trying multiple methods."""
|
||||
ips = []
|
||||
hostname = socket.gethostname()
|
||||
|
||||
# Method 1: Try socket.getaddrinfo
|
||||
try:
|
||||
addr_info = socket.getaddrinfo(hostname, None)
|
||||
for info in addr_info:
|
||||
ip = info[4][0]
|
||||
if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now
|
||||
ips.append(ip)
|
||||
except (socket.gaierror, OSError):
|
||||
pass
|
||||
|
||||
# Method 2: Try to connect to external server and get local IP
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
s.connect(("8.8.8.8", 80)) # Google DNS
|
||||
local_ip = s.getsockname()[0]
|
||||
s.close()
|
||||
if local_ip not in ips:
|
||||
ips.append(local_ip)
|
||||
except (OSError, socket.error):
|
||||
pass
|
||||
|
||||
# Method 3: Platform-specific commands
|
||||
try:
|
||||
if platform.system() == "Windows":
|
||||
# Windows ipconfig
|
||||
result = subprocess.run(["ipconfig"], capture_output=True, text=True)
|
||||
lines = result.stdout.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
if 'IPv4' in line and i + 1 < len(lines):
|
||||
ip = lines[i].split(':')[-1].strip()
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
else:
|
||||
# Unix/Linux/Mac ifconfig or ip addr
|
||||
try:
|
||||
result = subprocess.run(["ip", "addr"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
try:
|
||||
result = subprocess.run(["ifconfig"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
result = None
|
||||
|
||||
import re
|
||||
ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)')
|
||||
if result is not None:
|
||||
for match in ip_pattern.finditer(result.stdout):
|
||||
ip = match.group(1)
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
pass
|
||||
|
||||
return ips
|
||||
|
||||
|
||||
def get_recommended_ip(ips):
|
||||
"""Choose the best IP for master-worker communication."""
|
||||
# Priority order:
|
||||
# 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
|
||||
# 2. Other non-localhost IPs
|
||||
# 3. Localhost as last resort
|
||||
|
||||
private_ips = []
|
||||
public_ips = []
|
||||
|
||||
for ip in ips:
|
||||
if ip.startswith('127.') or ip == 'localhost':
|
||||
continue
|
||||
elif (ip.startswith('192.168.')
|
||||
or ip.startswith('10.')
|
||||
or (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)):
|
||||
private_ips.append(ip)
|
||||
else:
|
||||
public_ips.append(ip)
|
||||
|
||||
# Prefer private IPs
|
||||
if private_ips:
|
||||
# Prefer 192.168 range as it's most common
|
||||
for ip in private_ips:
|
||||
if ip.startswith('192.168.'):
|
||||
return ip
|
||||
return private_ips[0]
|
||||
elif public_ips:
|
||||
return public_ips[0]
|
||||
elif ips:
|
||||
return ips[0]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _get_cuda_info():
|
||||
"""Detect CUDA device index and total physical GPU count.
|
||||
|
||||
Returns (cuda_device, cuda_device_count, physical_device_count).
|
||||
All three are 0/None if CUDA is unavailable.
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
return None, 0, 0
|
||||
try:
|
||||
cuda_device_count = torch.cuda.device_count()
|
||||
cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '')
|
||||
if cuda_visible and cuda_visible.strip():
|
||||
visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()]
|
||||
if visible_devices:
|
||||
cuda_device = visible_devices[0]
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
physical_device_count = (
|
||||
len(result.stdout.strip().split('\n'))
|
||||
if result.returncode == 0
|
||||
else max(visible_devices) + 1
|
||||
)
|
||||
except (FileNotFoundError, OSError, subprocess.SubprocessError):
|
||||
physical_device_count = max(visible_devices) + 1
|
||||
return cuda_device, cuda_device_count, physical_device_count
|
||||
else:
|
||||
return 0, cuda_device_count, cuda_device_count
|
||||
else:
|
||||
cuda_device = torch.cuda.current_device()
|
||||
return cuda_device, cuda_device_count, cuda_device_count
|
||||
except Exception as e:
|
||||
debug_log(f"CUDA detection error: {e}")
|
||||
return None, 0, 0
|
||||
|
||||
|
||||
def _collect_network_info_sync():
|
||||
"""Collect network/cuda info in a worker thread to avoid blocking route handlers."""
|
||||
cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
|
||||
device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
|
||||
master_port = get_server_port()
|
||||
config = load_config()
|
||||
worker_ports = []
|
||||
if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
|
||||
worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
|
||||
worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
|
||||
hostname = socket.gethostname()
|
||||
all_ips = get_network_ips()
|
||||
recommended_ip = get_recommended_ip(all_ips)
|
||||
return {
|
||||
"hostname": hostname,
|
||||
"all_ips": all_ips,
|
||||
"recommended_ip": recommended_ip,
|
||||
"cuda_device": cuda_device,
|
||||
"cuda_device_count": device_count,
|
||||
"master_port": master_port,
|
||||
"local_worker_ports": worker_ports,
|
||||
}
|
||||
|
||||
|
||||
def _read_worker_log_sync(log_file, lines_to_read):
|
||||
"""Read worker log content from disk in a threadpool worker."""
|
||||
file_size = os.path.getsize(log_file)
|
||||
|
||||
with open(log_file, 'r', encoding='utf-8', errors='replace') as f:
|
||||
if lines_to_read > 0 and file_size > 1024 * 1024:
|
||||
# Read last N lines efficiently from end of file.
|
||||
lines = []
|
||||
f.seek(0, 2)
|
||||
file_length = f.tell()
|
||||
chunk_size = min(CHUNK_SIZE, file_length)
|
||||
|
||||
while len(lines) < lines_to_read and f.tell() > 0:
|
||||
current_pos = max(0, f.tell() - chunk_size)
|
||||
f.seek(current_pos)
|
||||
chunk = f.read(chunk_size)
|
||||
chunk_lines = chunk.splitlines()
|
||||
if current_pos > 0:
|
||||
chunk_lines = chunk_lines[1:]
|
||||
lines = chunk_lines + lines
|
||||
f.seek(current_pos)
|
||||
|
||||
content = '\n'.join(lines[-lines_to_read:])
|
||||
truncated = len(lines) > lines_to_read
|
||||
else:
|
||||
content = f.read()
|
||||
truncated = False
|
||||
|
||||
return {
|
||||
"content": content,
|
||||
"file_size": file_size,
|
||||
"truncated": truncated,
|
||||
"lines_shown": lines_to_read if truncated else content.count('\n') + 1,
|
||||
}
|
||||
|
||||
|
||||
def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
|
||||
"""Parse bounded positive integer query params with sane fallback."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
parsed = max(minimum, parsed)
|
||||
if maximum is not None:
|
||||
parsed = min(maximum, parsed)
|
||||
return parsed
|
||||
|
||||
|
||||
def _find_worker_by_id(config, worker_id):
|
||||
worker_id_str = str(worker_id).strip()
|
||||
for worker in config.get("workers", []):
|
||||
if str(worker.get("id")).strip() == worker_id_str:
|
||||
return worker
|
||||
return None
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local_log")
|
||||
async def get_local_log_endpoint(request):
|
||||
"""Return this instance's in-memory ComfyUI log buffer."""
|
||||
try:
|
||||
from app.logger import get_logs
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to import app.logger: {e}", 500)
|
||||
|
||||
try:
|
||||
lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
|
||||
logs = get_logs()
|
||||
if logs is None:
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"content": "",
|
||||
"entries": 0,
|
||||
"source": "memory",
|
||||
"truncated": False,
|
||||
"lines_shown": 0,
|
||||
}
|
||||
)
|
||||
|
||||
entries = list(logs)
|
||||
selected_entries = entries[-lines_to_read:]
|
||||
content = "".join(
|
||||
entry.get("m", "") if isinstance(entry, dict) else str(entry)
|
||||
for entry in selected_entries
|
||||
)
|
||||
lines_shown = content.count("\n") + (1 if content else 0)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"content": content,
|
||||
"entries": len(selected_entries),
|
||||
"source": "memory",
|
||||
"truncated": len(entries) > len(selected_entries),
|
||||
"lines_shown": lines_shown,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/network_info")
|
||||
async def get_network_info_endpoint(request):
|
||||
"""Get network interfaces and recommend best IP for master."""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
info = await loop.run_in_executor(None, _collect_network_info_sync)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
**info,
|
||||
"message": "Auto-detected network configuration"
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/system_info")
|
||||
async def get_system_info_endpoint(request):
|
||||
"""Get system information including machine ID for local worker detection."""
|
||||
try:
|
||||
import socket
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"hostname": socket.gethostname(),
|
||||
"machine_id": get_machine_id(),
|
||||
"platform": {
|
||||
"system": platform.system(),
|
||||
"machine": platform.machine(),
|
||||
"node": platform.node(),
|
||||
"path_separator": os.sep, # Add path separator
|
||||
"os_name": os.name # Add OS name (posix, nt, etc.)
|
||||
},
|
||||
"is_docker": is_docker_environment(),
|
||||
"is_runpod": is_runpod_environment(),
|
||||
"runpod_pod_id": os.environ.get('RUNPOD_POD_ID')
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/launch_worker")
|
||||
async def launch_worker_endpoint(request):
|
||||
"""Launch a worker process from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
|
||||
# Find worker config
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = worker_id
|
||||
|
||||
# Check if already running (managed by this instance)
|
||||
if worker_id_str in wm.processes:
|
||||
proc_info = wm.processes[worker_id_str]
|
||||
process = proc_info.get('process')
|
||||
|
||||
# Check if still running
|
||||
is_running = False
|
||||
if process:
|
||||
is_running = process.poll() is None
|
||||
else:
|
||||
# Restored process without subprocess object
|
||||
is_running = wm._is_process_running(proc_info['pid'])
|
||||
|
||||
if is_running:
|
||||
return await handle_api_error(request, "Worker already running (managed by UI)", 409)
|
||||
else:
|
||||
# Process is dead, remove it
|
||||
del wm.processes[worker_id_str]
|
||||
wm.save_processes()
|
||||
|
||||
# Launch the worker
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
pid = await loop.run_in_executor(None, wm.launch_worker, worker)
|
||||
log_file = wm.processes[worker_id_str].get('log_file')
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"pid": pid,
|
||||
"message": f"Worker {worker['name']} launched",
|
||||
"log_file": log_file
|
||||
})
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/stop_worker")
|
||||
async def stop_worker_endpoint(request):
|
||||
"""Stop a worker process that was launched from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
success, message = wm.stop_worker(worker_id)
|
||||
|
||||
if success:
|
||||
return web.json_response({"status": "success", "message": message})
|
||||
else:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
message,
|
||||
404 if "not managed" in message else 409,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/managed_workers")
|
||||
async def get_managed_workers_endpoint(request):
|
||||
"""Get list of workers managed by this UI instance."""
|
||||
try:
|
||||
managed = get_worker_manager().get_managed_workers()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"managed_workers": managed
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
|
||||
async def get_local_worker_status_endpoint(request):
|
||||
"""Check status of all local workers (localhost/no host specified)."""
|
||||
try:
|
||||
config = load_config()
|
||||
worker_statuses = {}
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
# Only check local workers
|
||||
host = normalize_host(worker.get("host")) or ""
|
||||
if not host or host in ["localhost", "127.0.0.1"]:
|
||||
worker_id = worker["id"]
|
||||
port = worker["port"]
|
||||
|
||||
# Check if worker is enabled
|
||||
if not worker.get("enabled", False):
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": False,
|
||||
"processing": False,
|
||||
"queue_count": 0
|
||||
}
|
||||
continue
|
||||
|
||||
# Try to connect to worker
|
||||
try:
|
||||
worker_url = build_worker_url(worker)
|
||||
data = await probe_worker(worker_url, timeout=2.0)
|
||||
if data is None:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Unavailable",
|
||||
}
|
||||
continue
|
||||
queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
|
||||
worker_statuses[worker_id] = {
|
||||
"online": True,
|
||||
"enabled": True,
|
||||
"processing": queue_remaining > 0,
|
||||
"queue_count": queue_remaining
|
||||
}
|
||||
except asyncio.TimeoutError:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Timeout"
|
||||
}
|
||||
except Exception as e:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"worker_statuses": worker_statuses
|
||||
})
|
||||
except Exception as e:
|
||||
debug_log(f"Error checking local worker status: {e}")
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
|
||||
async def get_worker_log_endpoint(request):
|
||||
"""Get log content for a specific worker."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
worker_id = request.match_info['worker_id']
|
||||
|
||||
# Ensure worker_id is string
|
||||
worker_id = str(worker_id)
|
||||
|
||||
# Check if we manage this worker
|
||||
if worker_id not in wm.processes:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404)
|
||||
|
||||
proc_info = wm.processes[worker_id]
|
||||
log_file = proc_info.get('log_file')
|
||||
|
||||
if not log_file or not os.path.exists(log_file):
|
||||
return await handle_api_error(request, "Log file not found", 404)
|
||||
|
||||
# Read last N lines (or full file if small)
|
||||
lines_to_read = _parse_positive_int_query(request.query.get('lines'), default=1000)
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _read_worker_log_sync, log_file, lines_to_read)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"content": payload["content"],
|
||||
"log_file": log_file,
|
||||
"file_size": payload["file_size"],
|
||||
"truncated": payload["truncated"],
|
||||
"lines_shown": payload["lines_shown"],
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Error reading log file: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
|
||||
async def get_remote_worker_log_endpoint(request):
|
||||
"""Proxy a remote worker log request to the worker's local in-memory log endpoint."""
|
||||
try:
|
||||
worker_id = str(request.match_info["worker_id"]).strip()
|
||||
config = load_config()
|
||||
worker = _find_worker_by_id(config, worker_id)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Remote log proxy is only meaningful for remote/cloud workers.
|
||||
host = normalize_host(worker.get("host")) or ""
|
||||
if not host:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Worker {worker_id} is local; use /distributed/worker_log/{worker_id} instead.",
|
||||
400,
|
||||
)
|
||||
|
||||
lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
|
||||
worker_url = build_worker_url(worker, "/distributed/local_log")
|
||||
session = await get_client_session()
|
||||
async with session.get(
|
||||
worker_url,
|
||||
params={"lines": str(lines_to_read)},
|
||||
timeout=aiohttp.ClientTimeout(total=5),
|
||||
) as resp:
|
||||
if resp.status >= 400:
|
||||
body = await resp.text()
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Remote worker {worker_id} returned HTTP {resp.status}: {body[:400]}",
|
||||
resp.status,
|
||||
)
|
||||
|
||||
try:
|
||||
data = await resp.json()
|
||||
except Exception as e:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Remote worker {worker_id} returned invalid JSON: {e}",
|
||||
502,
|
||||
)
|
||||
|
||||
return web.json_response(data)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
-31
@@ -1,31 +0,0 @@
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
# All actual package context is provided by each test module via
|
||||
# importlib.util.spec_from_file_location with synthetic stub packages.
|
||||
|
||||
from _pytest.python import Package
|
||||
|
||||
_orig_pkg_setup = Package.setup
|
||||
|
||||
|
||||
def _patched_pkg_setup(self) -> None:
|
||||
# Skip the root package setup — its __init__.py uses relative imports
|
||||
# that require a parent package (ComfyUI's plugin loader) which is not
|
||||
# available in the test environment.
|
||||
if self.path == self.config.rootpath:
|
||||
return
|
||||
_orig_pkg_setup(self)
|
||||
|
||||
|
||||
Package.setup = _patched_pkg_setup
|
||||
|
||||
collect_ignore = [
|
||||
"__init__.py",
|
||||
"distributed.py",
|
||||
"distributed_upscale.py",
|
||||
]
|
||||
+2900
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,605 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
|
||||
import aiohttp
|
||||
|
||||
import execution
|
||||
import server
|
||||
|
||||
from .utils.config import load_config
|
||||
from .utils.logging import debug_log, log
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def ensure_distributed_state():
|
||||
"""Ensure prompt_server has the state used by distributed queue orchestration."""
|
||||
if not hasattr(prompt_server, "distributed_pending_jobs"):
|
||||
prompt_server.distributed_pending_jobs = {}
|
||||
if not hasattr(prompt_server, "distributed_jobs_lock"):
|
||||
prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
if not hasattr(prompt_server, "distributed_worker_ws"):
|
||||
prompt_server.distributed_worker_ws = {}
|
||||
|
||||
|
||||
async def _get_client_session():
|
||||
"""Get or create aiohttp client session (shared with distributed.py)."""
|
||||
if not hasattr(prompt_server, "_distributed_session"):
|
||||
prompt_server._distributed_session = aiohttp.ClientSession()
|
||||
return prompt_server._distributed_session
|
||||
|
||||
|
||||
class PromptIndex:
|
||||
"""Cache prompt metadata for faster worker/master prompt preparation."""
|
||||
|
||||
def __init__(self, prompt_obj):
|
||||
self._prompt_json = json.dumps(prompt_obj)
|
||||
self.nodes_by_class = {}
|
||||
self.class_by_node = {}
|
||||
self.inputs_by_node = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
class_type = node.get("class_type")
|
||||
node_id_str = str(node_id)
|
||||
if class_type:
|
||||
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
|
||||
self.class_by_node[node_id_str] = class_type
|
||||
self.inputs_by_node[node_id_str] = node.get("inputs", {})
|
||||
self._upstream_cache = {}
|
||||
|
||||
def copy_prompt(self):
|
||||
return json.loads(self._prompt_json)
|
||||
|
||||
def nodes_for_class(self, class_name):
|
||||
return self.nodes_by_class.get(class_name, [])
|
||||
|
||||
def has_upstream(self, start_node_id, target_class):
|
||||
cache_key = (str(start_node_id), target_class)
|
||||
if cache_key in self._upstream_cache:
|
||||
return self._upstream_cache[cache_key]
|
||||
|
||||
visited = set()
|
||||
stack = [str(start_node_id)]
|
||||
while stack:
|
||||
node_id = stack.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
inputs = self.inputs_by_node.get(node_id, {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
upstream_id = str(value[0])
|
||||
if self.class_by_node.get(upstream_id) == target_class:
|
||||
self._upstream_cache[cache_key] = True
|
||||
return True
|
||||
if upstream_id in self.inputs_by_node:
|
||||
stack.append(upstream_id)
|
||||
|
||||
self._upstream_cache[cache_key] = False
|
||||
return False
|
||||
|
||||
|
||||
def _iter_prompt_nodes(prompt_obj):
|
||||
for node_id, node in prompt_obj.items():
|
||||
if isinstance(node, dict):
|
||||
yield str(node_id), node
|
||||
|
||||
|
||||
def _find_nodes_by_class(prompt_obj, class_name):
|
||||
nodes = []
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") == class_name:
|
||||
nodes.append(node_id)
|
||||
return nodes
|
||||
|
||||
|
||||
def _find_downstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable downstream from the provided IDs."""
|
||||
adjacency = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
adjacency.setdefault(source_id, set()).add(str(node_id))
|
||||
|
||||
connected = set(start_ids)
|
||||
queue = deque(start_ids)
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
|
||||
if dependent not in connected:
|
||||
connected.add(dependent)
|
||||
queue.append(dependent)
|
||||
return connected
|
||||
|
||||
|
||||
def _create_numeric_id_generator(prompt_obj):
|
||||
"""Return a closure that yields new numeric string IDs."""
|
||||
max_id = 0
|
||||
for node_id in prompt_obj.keys():
|
||||
try:
|
||||
numeric = int(node_id)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
max_id = max(max_id, numeric)
|
||||
|
||||
counter = max_id
|
||||
|
||||
def _next_id():
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
return str(counter)
|
||||
|
||||
return _next_id
|
||||
|
||||
|
||||
def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in nodes_to_keep:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
pruned_ids = set(pruned_prompt.keys())
|
||||
for node_id, node in pruned_prompt.items():
|
||||
inputs = node.get("inputs")
|
||||
if not inputs:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
if isinstance(input_value, list) and len(input_value) == 2:
|
||||
source_id = str(input_value[0])
|
||||
if source_id not in pruned_ids:
|
||||
inputs.pop(input_name, None)
|
||||
debug_log(
|
||||
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
||||
)
|
||||
|
||||
next_id = _create_numeric_id_generator(pruned_prompt)
|
||||
for collector_id in collector_ids:
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
placeholder_id = next_id()
|
||||
pruned_prompt[placeholder_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
"inputs": {
|
||||
"height": 64,
|
||||
"width": 64,
|
||||
"channels": 3,
|
||||
},
|
||||
"_meta": {
|
||||
"title": "Distributed Empty Image (auto-added)",
|
||||
},
|
||||
}
|
||||
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
|
||||
debug_log(
|
||||
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
|
||||
)
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
async def _ensure_distributed_queue(job_id):
|
||||
"""Ensure a queue exists for the given distributed job ID."""
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
|
||||
|
||||
def _generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
return job_map
|
||||
|
||||
|
||||
def _resolve_enabled_workers(config, requested_ids=None):
|
||||
"""Return a list of worker configs that should participate."""
|
||||
workers = []
|
||||
for worker in config.get("workers", []):
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id:
|
||||
continue
|
||||
|
||||
if requested_ids is not None:
|
||||
if worker_id not in requested_ids:
|
||||
continue
|
||||
elif not worker.get("enabled", False):
|
||||
continue
|
||||
|
||||
raw_port = worker.get("port", worker.get("listen_port", 8188))
|
||||
try:
|
||||
port = int(raw_port or 8188)
|
||||
except (TypeError, ValueError):
|
||||
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
|
||||
port = 8188
|
||||
|
||||
workers.append(
|
||||
{
|
||||
"id": worker_id,
|
||||
"name": worker.get("name", worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
}
|
||||
)
|
||||
return workers
|
||||
|
||||
|
||||
def _resolve_master_url():
|
||||
"""Best-effort reconstruction of the master's public URL."""
|
||||
cfg = load_config()
|
||||
master_cfg = cfg.get("master", {}) or {}
|
||||
configured_host = (master_cfg.get("host") or "").strip()
|
||||
configured_port = master_cfg.get("port")
|
||||
default_port = getattr(prompt_server, "port", 8188) or 8188
|
||||
port = int(configured_port or default_port)
|
||||
|
||||
def _needs_https(hostname):
|
||||
hostname = hostname.lower()
|
||||
https_domains = (
|
||||
".proxy.runpod.net",
|
||||
".ngrok-free.app",
|
||||
".ngrok-free.dev",
|
||||
".ngrok.io",
|
||||
".trycloudflare.com",
|
||||
".cloudflare.dev",
|
||||
)
|
||||
return any(hostname.endswith(suffix) for suffix in https_domains)
|
||||
|
||||
if configured_host:
|
||||
if configured_host.startswith(("http://", "https://")):
|
||||
return configured_host.rstrip("/")
|
||||
|
||||
host = configured_host
|
||||
scheme = "https" if _needs_https(host) or port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
# For ngrok/cloud domains without explicit port, default to the scheme's default
|
||||
if configured_port is None and scheme == "https" and _needs_https(host):
|
||||
port = default_port_for_scheme
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{host}{port_part}"
|
||||
|
||||
address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
|
||||
if address in ("0.0.0.0", "::"):
|
||||
address = "127.0.0.1"
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{address}{port_part}"
|
||||
|
||||
|
||||
def _build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
host = (worker.get("host") or "").strip()
|
||||
port = int(worker.get("port", 8188) or 8188)
|
||||
|
||||
if not host:
|
||||
host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
|
||||
|
||||
if host.startswith(("http://", "https://")):
|
||||
base = host.rstrip("/")
|
||||
else:
|
||||
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
|
||||
scheme = "https" if is_cloud else "http"
|
||||
default_port = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port else f":{port}"
|
||||
base = f"{scheme}://{host}{port_part}"
|
||||
|
||||
return f"{base}{endpoint}"
|
||||
|
||||
|
||||
async def _worker_is_active(worker):
|
||||
"""Ping worker's /prompt endpoint to confirm it's reachable."""
|
||||
url = _build_worker_url(worker, "/prompt")
|
||||
session = await _get_client_session()
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
|
||||
return resp.status == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def _worker_ws_is_active(worker):
|
||||
"""Ping worker's websocket endpoint to confirm it's reachable."""
|
||||
session = await _get_client_session()
|
||||
url = _build_worker_url(worker, "/distributed/worker_ws")
|
||||
try:
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
|
||||
await ws.close()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def _get_worker_ws_connection(worker):
|
||||
"""Return a cached websocket connection to a worker (creating if needed)."""
|
||||
ensure_distributed_state()
|
||||
worker_id = worker.get("id")
|
||||
existing = prompt_server.distributed_worker_ws.get(worker_id)
|
||||
if existing and not existing.closed:
|
||||
return existing
|
||||
|
||||
session = await _get_client_session()
|
||||
url = _build_worker_url(worker, "/distributed/worker_ws")
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=10)
|
||||
prompt_server.distributed_worker_ws[worker_id] = ws
|
||||
return ws
|
||||
|
||||
|
||||
async def _close_worker_ws(worker_id):
|
||||
ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
|
||||
if ws and not ws.closed:
|
||||
await ws.close()
|
||||
|
||||
|
||||
async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
url = _build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if extra_data:
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
if use_websocket:
|
||||
request_id = uuid.uuid4().hex
|
||||
ws_payload = {
|
||||
"type": "dispatch_prompt",
|
||||
"request_id": request_id,
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
"client_id": client_id,
|
||||
}
|
||||
try:
|
||||
ws = await _get_worker_ws_connection(worker)
|
||||
await ws.send_json(ws_payload)
|
||||
response = await asyncio.wait_for(ws.receive(), timeout=60)
|
||||
if response.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(response.data or "{}")
|
||||
if (
|
||||
data.get("type") == "dispatch_ack"
|
||||
and data.get("request_id") == request_id
|
||||
and data.get("ok")
|
||||
):
|
||||
return
|
||||
error_msg = data.get("error") or "Worker rejected websocket dispatch."
|
||||
raise RuntimeError(error_msg)
|
||||
raise RuntimeError("Unexpected websocket response from worker.")
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
await _close_worker_ws(worker_id)
|
||||
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
|
||||
session = await _get_client_session()
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
|
||||
|
||||
async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
|
||||
"""Queue the master prompt through ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
if not valid[0]:
|
||||
raise RuntimeError(f"Invalid prompt: {valid[1]}")
|
||||
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
extra_data["client_id"] = client_id
|
||||
|
||||
sensitive = {}
|
||||
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
return prompt_id
|
||||
|
||||
|
||||
def _apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
):
|
||||
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
||||
is_master = participant_id == "master"
|
||||
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
||||
enabled_json = json.dumps(enabled_worker_ids)
|
||||
|
||||
# Distributed seed nodes
|
||||
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if not is_master:
|
||||
idx = worker_index_map.get(participant_id, 0)
|
||||
inputs["worker_id"] = f"worker_{idx}"
|
||||
else:
|
||||
inputs["worker_id"] = ""
|
||||
|
||||
# Distributed collectors
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
|
||||
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
||||
inputs["pass_through"] = True
|
||||
continue
|
||||
|
||||
unique_id = job_id_map.get(node_id, node_id)
|
||||
inputs["multi_job_id"] = unique_id
|
||||
inputs["is_worker"] = not is_master
|
||||
|
||||
if is_master:
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
# Distributed upscaler nodes
|
||||
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
unique_id = job_id_map.get(node_id, node_id)
|
||||
inputs["multi_job_id"] = unique_id
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
|
||||
return prompt_copy
|
||||
|
||||
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=None,
|
||||
delegate_master=None,
|
||||
):
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (prompt_id, worker_count)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
|
||||
config = load_config()
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
|
||||
delegate_master = False
|
||||
|
||||
# Filter to active workers
|
||||
active_workers = []
|
||||
for worker in workers:
|
||||
if use_websocket:
|
||||
is_active = await _worker_ws_is_active(worker)
|
||||
else:
|
||||
is_active = await _worker_is_active(worker)
|
||||
if is_active:
|
||||
active_workers.append(worker)
|
||||
else:
|
||||
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
|
||||
if not active_workers and delegate_master:
|
||||
debug_log("All workers offline while delegate-only requested; enabling master participation.")
|
||||
delegate_master = False
|
||||
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
|
||||
return prompt_id, 0
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
master_url = _resolve_master_url()
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = _apply_participant_overrides(
|
||||
master_prompt,
|
||||
"master",
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
|
||||
)
|
||||
elif not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
|
||||
)
|
||||
else:
|
||||
master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
if active_workers:
|
||||
debug_log(
|
||||
"Active distributed workers: "
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
|
||||
)
|
||||
worker_payloads = []
|
||||
for worker in active_workers:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
worker_prompt = _apply_participant_overrides(
|
||||
worker_prompt,
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
worker_payloads.append((worker, worker_prompt))
|
||||
|
||||
if worker_payloads:
|
||||
await asyncio.gather(
|
||||
*[
|
||||
_dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
|
||||
for worker, wprompt in worker_payloads
|
||||
]
|
||||
)
|
||||
|
||||
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
|
||||
return prompt_id, len(worker_payloads)
|
||||
File diff suppressed because it is too large
Load Diff
+17
-105
@@ -19,10 +19,9 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
|
||||
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
|
||||
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
|
||||
- Resolves enabled/selected workers.
|
||||
- Probes and dispatches workers through `/distributed/worker_ws` by default.
|
||||
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
|
||||
- Pings workers (`GET /prompt`) to include only reachable ones.
|
||||
- Dispatches the workflow to workers (`POST /prompt`).
|
||||
- Queues the master workflow in ComfyUI’s prompt queue.
|
||||
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
|
||||
|
||||
### What it does *not* add
|
||||
|
||||
@@ -49,39 +48,27 @@ Queue a workflow for distributed execution.
|
||||
{
|
||||
"prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
|
||||
"workflow": { },
|
||||
"client_id": "external-client",
|
||||
"client_id": "optional",
|
||||
"delegate_master": false,
|
||||
"enabled_worker_ids": ["1", "2"],
|
||||
"workers": ["1", "2"],
|
||||
"auto_prepare": true,
|
||||
"trace_execution_id": "exec_1700000000_ab12cd"
|
||||
"enabled_worker_ids": ["1", "2"]
|
||||
}
|
||||
```
|
||||
|
||||
#### Fields
|
||||
|
||||
- `prompt` (required unless `workflow.prompt` is present, object)
|
||||
- A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
|
||||
- This is not the normal visual workflow export from the ComfyUI editor.
|
||||
- `prompt` (required, object)
|
||||
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
|
||||
- `workflow` (optional, object)
|
||||
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
|
||||
- If you don’t care about UI metadata, you can omit it.
|
||||
- `client_id` (required, string)
|
||||
- `client_id` (optional, string)
|
||||
- Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
|
||||
- `delegate_master` (optional, boolean)
|
||||
- If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
|
||||
- Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
|
||||
- `enabled_worker_ids` (required, array of strings)
|
||||
- The explicit worker IDs to consider for this run.
|
||||
- `workers` (optional, array of strings or objects with `id`)
|
||||
- Transitional alias for `enabled_worker_ids` used by older clients.
|
||||
- `auto_prepare` (optional, boolean)
|
||||
- Kept for wire compatibility.
|
||||
- Backend orchestration always runs with auto-prepare semantics.
|
||||
- If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
|
||||
- `trace_execution_id` (optional, string)
|
||||
- Passed through to orchestration logs.
|
||||
- Server log lines include the marker as `[exec:<trace_execution_id>]`.
|
||||
- `enabled_worker_ids` (optional, array of strings)
|
||||
- If provided, only these worker IDs will be considered.
|
||||
- If omitted, the plugin uses workers marked as enabled in the UI config.
|
||||
|
||||
##### How to get `enabled_worker_ids`
|
||||
|
||||
@@ -106,8 +93,7 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
```json
|
||||
{
|
||||
"prompt_id": "<uuid>",
|
||||
"worker_count": 2,
|
||||
"auto_prepare_supported": true
|
||||
"worker_count": 2
|
||||
}
|
||||
```
|
||||
|
||||
@@ -124,89 +110,24 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
|
||||
## Worker requirements (important)
|
||||
|
||||
For a worker to participate, it must be reachable from the master. By default:
|
||||
|
||||
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
|
||||
|
||||
If `settings.websocket_orchestration=false`:
|
||||
For a worker to participate, it must be reachable from the master:
|
||||
|
||||
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
|
||||
|
||||
Also, for collector-based flows:
|
||||
|
||||
- Workers will send results back to the master via `POST /distributed/job_complete` (that route must be reachable from workers).
|
||||
|
||||
---
|
||||
|
||||
## Endpoint: `POST /distributed/job_complete`
|
||||
|
||||
Submit one completed worker image back to the master collector queue.
|
||||
|
||||
### URL
|
||||
|
||||
- `http://<master-host>:<master-port>/distributed/job_complete`
|
||||
|
||||
### Request Body
|
||||
|
||||
```json
|
||||
{
|
||||
"job_id": "exec_1234567890_17",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA...",
|
||||
"is_last": false
|
||||
}
|
||||
```
|
||||
|
||||
### Canonical envelope (required fields)
|
||||
|
||||
- `job_id` (string, required)
|
||||
- `worker_id` (string, required)
|
||||
- `batch_idx` (integer >= 0, required)
|
||||
- `image` (string, required)
|
||||
- PNG payload as either:
|
||||
- data URL: `data:image/png;base64,...`
|
||||
- raw base64 PNG bytes
|
||||
- `is_last` (boolean, required)
|
||||
|
||||
Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
|
||||
|
||||
### CORS note
|
||||
|
||||
If you call the API from a browser (not from a backend), ensure the master ComfyUI is started with `--enable-cors-header`.
|
||||
|
||||
---
|
||||
|
||||
## Log Endpoints
|
||||
|
||||
### `GET /distributed/worker_log/{worker_id}`
|
||||
|
||||
Read log files for workers launched locally by the master UI process manager.
|
||||
|
||||
- Intended for managed local workers.
|
||||
- Query param: `lines` (optional, default `1000`).
|
||||
|
||||
### `GET /distributed/local_log`
|
||||
|
||||
Read this ComfyUI instance's in-memory runtime log buffer.
|
||||
|
||||
- Available on any ComfyUI-Distributed instance (master or worker).
|
||||
- Query param: `lines` (optional, default `300`, max `3000`).
|
||||
|
||||
### `GET /distributed/remote_worker_log/{worker_id}`
|
||||
|
||||
Proxy endpoint on master that fetches logs from a configured remote/cloud worker's
|
||||
`/distributed/local_log`.
|
||||
|
||||
- Intended for remote/cloud workers in master config.
|
||||
- Query param: `lines` (optional, default `300`, max `3000`).
|
||||
|
||||
---
|
||||
|
||||
## Examples
|
||||
|
||||
### 1) Minimal `curl` request envelope
|
||||
### 1) Minimal `curl`
|
||||
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
@@ -214,23 +135,16 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
|
||||
Where `payload.json` contains at least:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": {
|
||||
"<node_id>": {
|
||||
"class_type": "<node class>",
|
||||
"inputs": {"<required_input>": "<value or connection>"}
|
||||
}
|
||||
},
|
||||
"enabled_worker_ids": [],
|
||||
"client_id": "external-client"
|
||||
"1": {"class_type": "KSampler", "inputs": {} }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
|
||||
|
||||
### 2) Python (`requests`)
|
||||
|
||||
```python
|
||||
@@ -240,7 +154,6 @@ url = "http://127.0.0.1:8188/distributed/queue"
|
||||
payload = {
|
||||
"prompt": {...},
|
||||
"workflow": {...},
|
||||
"client_id": "external-client",
|
||||
"delegate_master": False,
|
||||
"enabled_worker_ids": ["1", "2"],
|
||||
}
|
||||
@@ -258,7 +171,6 @@ const url = "http://127.0.0.1:8188/distributed/queue";
|
||||
const payload = {
|
||||
prompt: {/* ... */},
|
||||
workflow: {/* ... */},
|
||||
client_id: "external-client",
|
||||
delegate_master: false,
|
||||
enabled_worker_ids: ["1", "2"],
|
||||
};
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
# Compressed Result Transfer
|
||||
|
||||
**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
|
||||
|
||||
| `transfer_format` | Type | Notes |
|
||||
|---|---|---|
|
||||
| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
|
||||
| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
|
||||
| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
|
||||
| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
|
||||
| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
|
||||
|
||||
`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
|
||||
|
||||
Measured on a 121-frame 1280×720 batch (one worker's result):
|
||||
|
||||
| Format | Sent | vs legacy | PSNR |
|
||||
|---|---|---|---|
|
||||
| `legacy_png` | 446 MB | 1× | lossless |
|
||||
| `png` | 163 MB | 2.7× | lossless |
|
||||
| `ffv1` | 128 MB | 3.5× | lossless |
|
||||
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
|
||||
| `webp` q90 | 13 MB | 33× | 40.6 dB |
|
||||
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
|
||||
| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
|
||||
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
|
||||
|
||||
The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
|
||||
|
||||
Notes:
|
||||
- Choose the format on the master; the value is sent to workers with the workflow.
|
||||
- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
|
||||
- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
|
||||
- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
|
||||
- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
|
||||
- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
|
||||
- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
|
||||
- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.
|
||||
@@ -1,53 +0,0 @@
|
||||
# Native V3 node API
|
||||
|
||||
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
|
||||
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
|
||||
that API; there is no V1 registration fallback.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- All eight node IDs, display names, categories, visible input order, defaults
|
||||
and output order are retained. Existing execution algorithms remain in the
|
||||
private collector, utilities and upscale modules.
|
||||
- The image/audio dividers explicitly declare ten typed outputs. Their existing
|
||||
frontend extensions still show only the selected number of outputs. This
|
||||
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
|
||||
IMAGE/AUDIO socket indices or the ten returned values.
|
||||
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
|
||||
its existing prompt-input names and Python defaults via `accept_all_inputs`;
|
||||
it does not become visible widgets.
|
||||
- The collector retains list-input handling. Upscale retains its always-changing
|
||||
fingerprint and creates a private runtime helper for each V3 execution, so
|
||||
mutable helper state is not attached to sanitized V3 class clones.
|
||||
- Routes, distributed state and the existing worker startup/shutdown hooks are
|
||||
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
|
||||
Worker mode still suppresses automatic worker launch. `distributed.py` is no
|
||||
longer an entrypoint.
|
||||
|
||||
## Verification
|
||||
|
||||
Run ordinary unit tests from this repository:
|
||||
|
||||
```bash
|
||||
python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
|
||||
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
|
||||
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
|
||||
all node schemas with the V1 fixture, checks all five bundled workflows' node
|
||||
IDs and socket/link contracts, exercises injected worker metadata, image/audio
|
||||
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
|
||||
and decodes an actual preview PNG referenced by executor history.
|
||||
|
||||
The upscale GPU/model boundary is mocked to check argument forwarding and
|
||||
per-execution helper isolation. Full checkpoint inference, browser canvas
|
||||
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
|
||||
workers or model downloads are started; preview files use a temporary scratch
|
||||
directory. The test does not install this branch into a live custom-node folder.
|
||||
@@ -74,14 +74,12 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
|
||||
> ComfyUI instances running on a cloud service like Runpod.
|
||||
|
||||
For faster result transfers and lower bandwidth usage, see [Compressed Result Transfer](compressed-result-transfer.md).
|
||||
|
||||
### Deploy Cloud Worker on Runpod
|
||||
|
||||
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||
|
||||
**On Runpod:**
|
||||
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
|
||||
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
|
||||
|
||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
||||
@@ -94,7 +92,7 @@ For faster result transfers and lower bandwidth usage, see [Compressed Result Tr
|
||||
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
||||
5. Deploy your pod.
|
||||
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
||||
7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
```
|
||||
# Download from CivitAI
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||
@@ -1,731 +0,0 @@
|
||||
import torch
|
||||
import io
|
||||
import json
|
||||
import asyncio
|
||||
import time
|
||||
import base64
|
||||
|
||||
import aiohttp
|
||||
import server as _server
|
||||
import comfy.model_management
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
|
||||
from ..utils.constants import HEARTBEAT_INTERVAL, TRANSFER_MAX_CHUNK_BYTES, TRANSFER_MAX_REQUEST_BYTES
|
||||
from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import build_worker_url, get_client_session, probe_worker
|
||||
from ..utils.audio_payload import encode_audio_payload
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..utils import transfer_codec
|
||||
|
||||
prompt_server = _server.PromptServer.instance
|
||||
|
||||
|
||||
def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
|
||||
"""Materialize every transfer chunk (runs in an executor)."""
|
||||
return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
|
||||
|
||||
|
||||
class DistributedCollectorNode:
|
||||
INPUT_IS_LIST = True
|
||||
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"load_balance": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating).",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
"transfer_format": (
|
||||
list(transfer_codec.TRANSFER_FORMATS),
|
||||
{
|
||||
"default": transfer_codec.DEFAULT_TRANSFER_FORMAT,
|
||||
"tooltip": transfer_codec.TRANSFER_FORMAT_TOOLTIP,
|
||||
},
|
||||
),
|
||||
"transfer_quality": (
|
||||
"INT",
|
||||
{
|
||||
"default": transfer_codec.DEFAULT_TRANSFER_QUALITY,
|
||||
"min": 1,
|
||||
"max": 100,
|
||||
"step": 1,
|
||||
"tooltip": transfer_codec.TRANSFER_QUALITY_TOOLTIP,
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"pass_through": ("BOOLEAN", {"default": False}),
|
||||
"delegate_only": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "AUDIO")
|
||||
RETURN_NAMES = ("images", "audio")
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image"
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_list_input(value):
|
||||
"""Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
|
||||
if isinstance(value, (list, tuple)) and len(value) == 1:
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
def _normalize_images_input(self, images):
|
||||
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
|
||||
if isinstance(images, (list, tuple)):
|
||||
if not images:
|
||||
raise ValueError("Collector received an empty image list")
|
||||
if not all(isinstance(image, torch.Tensor) for image in images):
|
||||
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
|
||||
if len(images) == 1:
|
||||
return ensure_contiguous(images[0])
|
||||
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
|
||||
return ensure_contiguous(images)
|
||||
|
||||
def _normalize_audio_input(self, audio):
|
||||
"""Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
|
||||
if not isinstance(audio, (list, tuple)):
|
||||
return audio
|
||||
|
||||
audio_items = [item for item in audio if item is not None]
|
||||
if not audio_items:
|
||||
return None
|
||||
if len(audio_items) == 1:
|
||||
return audio_items[0]
|
||||
|
||||
waveforms = []
|
||||
sample_rate = 44100
|
||||
for item in audio_items:
|
||||
if not isinstance(item, dict):
|
||||
raise TypeError("Collector expected AUDIO list items to be dictionaries")
|
||||
waveform = item.get("waveform")
|
||||
if waveform is None or waveform.numel() == 0:
|
||||
continue
|
||||
waveforms.append(waveform)
|
||||
if sample_rate == 44100:
|
||||
sample_rate = item.get("sample_rate", 44100)
|
||||
|
||||
if not waveforms:
|
||||
return None
|
||||
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
||||
|
||||
def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
if images is not None:
|
||||
images = self._normalize_images_input(images)
|
||||
audio = self._normalize_audio_input(audio)
|
||||
load_balance = self._unwrap_list_input(load_balance)
|
||||
multi_job_id = self._unwrap_list_input(multi_job_id)
|
||||
is_worker = self._unwrap_list_input(is_worker)
|
||||
master_url = self._unwrap_list_input(master_url)
|
||||
enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
|
||||
worker_batch_size = self._unwrap_list_input(worker_batch_size)
|
||||
worker_id = self._unwrap_list_input(worker_id)
|
||||
pass_through = self._unwrap_list_input(pass_through)
|
||||
delegate_only = self._unwrap_list_input(delegate_only)
|
||||
transfer_format = self._unwrap_list_input(transfer_format)
|
||||
transfer_quality = self._unwrap_list_input(transfer_quality)
|
||||
|
||||
remote_only_master = (
|
||||
bool(multi_job_id)
|
||||
and not is_worker
|
||||
and (delegate_only or is_master_delegate_only())
|
||||
)
|
||||
if images is None and audio is None and not remote_only_master:
|
||||
raise ValueError("DistributedCollector requires at least one image or audio input")
|
||||
|
||||
# Create empty audio if not provided
|
||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
if not multi_job_id or pass_through:
|
||||
if pass_through:
|
||||
debug_log("Collector: pass-through mode enabled, returning images unchanged")
|
||||
return (images, audio if audio is not None else empty_audio)
|
||||
|
||||
# Use async helper to run in server loop
|
||||
result = run_async_in_server_loop(
|
||||
self.execute(
|
||||
images,
|
||||
audio,
|
||||
load_balance,
|
||||
multi_job_id,
|
||||
is_worker,
|
||||
master_url,
|
||||
enabled_worker_ids,
|
||||
worker_batch_size,
|
||||
worker_id,
|
||||
delegate_only,
|
||||
transfer_format,
|
||||
transfer_quality,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def send_batch_to_master(
|
||||
self,
|
||||
image_batch,
|
||||
audio,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
worker_id,
|
||||
transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT,
|
||||
transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY,
|
||||
):
|
||||
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||
fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
|
||||
if fmt != transfer_codec.LEGACY_FORMAT:
|
||||
sent = await self._send_compressed_batch(
|
||||
image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
|
||||
)
|
||||
if sent:
|
||||
return
|
||||
log(
|
||||
"Worker - Master did not accept compressed transfer "
|
||||
"(older ComfyUI-Distributed on the master?). Falling back to legacy_png."
|
||||
)
|
||||
await self._send_legacy_batch(image_batch, audio, multi_job_id, master_url, worker_id)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_transfer_format(transfer_format):
|
||||
try:
|
||||
fmt = transfer_codec.normalize_format(transfer_format)
|
||||
except transfer_codec.TransferCodecError as exc:
|
||||
log(f"Worker - {exc}; using legacy_png.")
|
||||
return transfer_codec.LEGACY_FORMAT
|
||||
if fmt not in transfer_codec.VIDEO_FORMATS or transfer_codec.video_codec_available(fmt):
|
||||
return fmt
|
||||
fallback = "h264" if fmt != "ffv1" and transfer_codec.video_codec_available("h264") else "png"
|
||||
log(f"Worker - Encoder for '{fmt}' is not available in this PyAV build; using '{fallback}'.")
|
||||
return fallback
|
||||
|
||||
@staticmethod
|
||||
def _build_multipart(envelope, parts):
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("envelope", json.dumps(envelope), content_type="application/json")
|
||||
for i, data in enumerate(parts):
|
||||
form.add_field("media", data, filename=f"{i}.bin", content_type="application/octet-stream")
|
||||
return form
|
||||
|
||||
# Room left in each request for the multipart envelope (ids, frame
|
||||
# metadata, part headers). Audio is budgeted separately.
|
||||
TRANSFER_ENVELOPE_RESERVE_BYTES = 256 * 1024
|
||||
# Body of a 400 from a master without compressed-transfer support: it
|
||||
# tries request.json() on the multipart body and fails.
|
||||
_OLD_MASTER_REJECTION = "Invalid JSON payload"
|
||||
|
||||
async def _send_compressed_batch(self, image_batch, audio, multi_job_id, master_url, worker_id, fmt, quality):
|
||||
"""Send the batch as compressed multipart chunks.
|
||||
|
||||
All chunks are encoded before anything is sent, so a frame that cannot
|
||||
fit in one request fails the transfer without delivering a partial
|
||||
batch. Every request stays within TRANSFER_MAX_REQUEST_BYTES.
|
||||
|
||||
Returns False (without having delivered anything) when the master
|
||||
rejects the first chunk because it predates compressed transfer, so
|
||||
the caller can retry with legacy_png. Any other rejection raises.
|
||||
"""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
frames = await loop.run_in_executor(None, transfer_codec.tensor_batch_to_uint8, image_batch)
|
||||
raw_bytes = int(frames.nbytes)
|
||||
total = int(frames.shape[0])
|
||||
media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
|
||||
started = time.time()
|
||||
# Encoding is CPU-bound (seconds for video); keep the server loop responsive.
|
||||
chunks = await loop.run_in_executor(
|
||||
None, _encode_all_chunks, frames, fmt, quality, min(TRANSFER_MAX_CHUNK_BYTES, media_limit), media_limit,
|
||||
)
|
||||
del frames
|
||||
|
||||
# Attach audio to the final chunk only if the request stays in bounds;
|
||||
# otherwise finish with a separate audio-only completion.
|
||||
audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
|
||||
last_media = sum(len(p) for p in chunks[-1]["parts"])
|
||||
audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
|
||||
audio_separate = encoded_audio is not None and not audio_inline
|
||||
|
||||
sent_bytes = 0
|
||||
for index, chunk in enumerate(chunks):
|
||||
parts = chunk.pop("parts")
|
||||
is_final_chunk = index == len(chunks) - 1
|
||||
envelope = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"is_last": bool(is_final_chunk and not audio_separate),
|
||||
**chunk,
|
||||
}
|
||||
if is_final_chunk and audio_inline:
|
||||
envelope["audio"] = encoded_audio
|
||||
chunk_bytes = sum(len(p) for p in parts)
|
||||
timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
data=self._build_multipart(envelope, parts),
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||
) as response:
|
||||
if response.status >= 400:
|
||||
body = await response.text()
|
||||
if index == 0 and response.status == 400 and self._OLD_MASTER_REJECTION in body:
|
||||
debug_log(f"Worker - Master rejected compressed chunk: {body[:300]}")
|
||||
return False
|
||||
raise RuntimeError(
|
||||
f"Master rejected {fmt} chunk {index + 1}/{len(chunks)} "
|
||||
f"({chunk_bytes} bytes) with HTTP {response.status}: {body[:300]}. "
|
||||
"If this is a size limit, make COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES "
|
||||
"on the worker smaller than the master's --max-upload-size."
|
||||
)
|
||||
except Exception as e:
|
||||
log(f"Worker - Failed to send {fmt} chunk to master: {e}")
|
||||
debug_log(f"Worker - Full error details: URL={url}")
|
||||
raise
|
||||
sent_bytes += chunk_bytes
|
||||
|
||||
requests = len(chunks)
|
||||
if audio_separate:
|
||||
await self._post_json_completion(session, url, {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": 0,
|
||||
"audio": encoded_audio,
|
||||
"is_last": True,
|
||||
})
|
||||
requests += 1
|
||||
|
||||
log(
|
||||
f"Worker - Sent {total} image(s) as {fmt} in {requests} request(s): "
|
||||
f"{sent_bytes / 1e6:.2f} MB (raw {raw_bytes / 1e6:.2f} MB, "
|
||||
f"{transfer_codec.describe_ratio(sent_bytes, raw_bytes)} smaller) in {time.time() - started:.1f}s"
|
||||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
async def _post_json_completion(session, url, payload):
|
||||
timeout_seconds = 60 if "image" in payload else 600
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
media_type = "image/audio" if "image" in payload else "audio-only"
|
||||
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
|
||||
debug_log(f"Worker - Full error details: URL={url}")
|
||||
raise
|
||||
|
||||
async def _send_legacy_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||
"""Original transfer: one JSON request per frame with a base64 PNG."""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
|
||||
payloads = []
|
||||
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||
if batch_size == 0:
|
||||
if encoded_audio is None:
|
||||
raise ValueError("Worker completion requires image or audio data")
|
||||
payloads.append(
|
||||
{
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": 0,
|
||||
"audio": encoded_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
else:
|
||||
for batch_idx in range(batch_size):
|
||||
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
||||
byte_io = io.BytesIO()
|
||||
img.save(byte_io, format='PNG', compress_level=0)
|
||||
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
||||
payload = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": int(batch_idx),
|
||||
"image": f"data:image/png;base64,{encoded_image}",
|
||||
"is_last": bool(batch_idx == batch_size - 1),
|
||||
}
|
||||
if payload["is_last"] and encoded_audio is not None:
|
||||
payload["audio"] = encoded_audio
|
||||
payloads.append(payload)
|
||||
|
||||
for payload in payloads:
|
||||
await self._post_json_completion(session, url, payload)
|
||||
|
||||
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
|
||||
"""Combine audio from master and workers into a single audio output.
|
||||
|
||||
Ordering: master first, then workers in `worker_order` (if provided),
|
||||
then any unexpected worker ids in sorted order.
|
||||
"""
|
||||
audio_pieces = []
|
||||
sample_rate = 44100
|
||||
|
||||
# Add master audio first if present
|
||||
if master_audio is not None:
|
||||
waveform = master_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
sample_rate = master_audio.get("sample_rate", 44100)
|
||||
|
||||
# Add worker audio in configured enabled-worker order first.
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
for worker_id_str in ordered_worker_ids:
|
||||
seen.add(worker_id_str)
|
||||
w_audio = worker_audio.get(worker_id_str)
|
||||
if w_audio is not None:
|
||||
waveform = w_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
# Use first available sample rate
|
||||
if sample_rate == 44100:
|
||||
sample_rate = w_audio.get("sample_rate", 44100)
|
||||
|
||||
# Append any audio from unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_audio.keys()):
|
||||
if worker_id_str in seen:
|
||||
continue
|
||||
w_audio = worker_audio[worker_id_str]
|
||||
if w_audio is not None:
|
||||
waveform = w_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
if sample_rate == 44100:
|
||||
sample_rate = w_audio.get("sample_rate", 44100)
|
||||
|
||||
if not audio_pieces:
|
||||
return empty_audio
|
||||
|
||||
try:
|
||||
# Concatenate along the samples dimension (dim=-1)
|
||||
# Ensure all pieces have same batch and channel dimensions
|
||||
combined_waveform = torch.cat(audio_pieces, dim=-1)
|
||||
debug_log(f"Master - Combined audio: {len(audio_pieces)} pieces, final shape={combined_waveform.shape}")
|
||||
return {"waveform": combined_waveform, "sample_rate": sample_rate}
|
||||
except Exception as e:
|
||||
log(f"[Distributed] Master - Audio combination failed, returning silence: {e}")
|
||||
return empty_audio
|
||||
|
||||
def _store_worker_result(self, worker_images: dict, item: dict) -> int:
|
||||
"""Store one canonical queue item in worker_images in-place.
|
||||
|
||||
Canonical format:
|
||||
- item has 'worker_id', 'image_index', and 'tensor'
|
||||
Returns 1 when stored, otherwise 0.
|
||||
"""
|
||||
worker_id = item['worker_id']
|
||||
tensor = item.get('tensor')
|
||||
image_index = item.get('image_index')
|
||||
if tensor is None or image_index is None:
|
||||
return 0
|
||||
|
||||
worker_images.setdefault(worker_id, {})
|
||||
worker_images[worker_id][image_index] = tensor
|
||||
return 1
|
||||
|
||||
def _reorder_and_combine_tensors(
|
||||
self,
|
||||
worker_images: dict,
|
||||
worker_order: list,
|
||||
master_batch_size: int,
|
||||
images_on_cpu,
|
||||
delegate_mode: bool,
|
||||
fallback_images,
|
||||
):
|
||||
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||
ordered_tensors = []
|
||||
if not delegate_mode and images_on_cpu is not None:
|
||||
for i in range(master_batch_size):
|
||||
ordered_tensors.append(images_on_cpu[i:i+1])
|
||||
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
for worker_id_str in ordered_worker_ids:
|
||||
seen.add(worker_id_str)
|
||||
if worker_id_str not in worker_images:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
|
||||
# Append any unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_images.keys()):
|
||||
if worker_id_str in seen:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
|
||||
cpu_tensors = []
|
||||
for t in ordered_tensors:
|
||||
if t.is_cuda:
|
||||
t = t.cpu()
|
||||
t = ensure_contiguous(t)
|
||||
cpu_tensors.append(t)
|
||||
|
||||
if cpu_tensors:
|
||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||
if fallback_images is not None:
|
||||
return ensure_contiguous(fallback_images)
|
||||
return None
|
||||
|
||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
if is_worker:
|
||||
# Worker mode: send images and audio to master in a single batch
|
||||
image_count = 0 if images is None else images.shape[0]
|
||||
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
|
||||
await self.send_batch_to_master(
|
||||
images, audio, multi_job_id, master_url, worker_id, transfer_format, transfer_quality
|
||||
)
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
else:
|
||||
delegate_mode = delegate_only or is_master_delegate_only()
|
||||
# Master mode: collect images and audio from workers
|
||||
enabled_workers_raw = json.loads(enabled_worker_ids)
|
||||
enabled_workers = []
|
||||
seen_enabled = set()
|
||||
for worker_id in enabled_workers_raw:
|
||||
worker_id_str = str(worker_id)
|
||||
if worker_id_str in seen_enabled:
|
||||
continue
|
||||
seen_enabled.add(worker_id_str)
|
||||
enabled_workers.append(worker_id_str)
|
||||
expected_workers = set(enabled_workers)
|
||||
num_workers = len(expected_workers)
|
||||
if num_workers == 0:
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
|
||||
# Create the queue before any expensive local work to avoid job_complete race.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
debug_log(f"Master - Initialized queue early for job {multi_job_id}")
|
||||
else:
|
||||
existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
|
||||
debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
|
||||
|
||||
if delegate_mode:
|
||||
master_batch_size = 0
|
||||
images_on_cpu = None
|
||||
master_audio = None
|
||||
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
|
||||
else:
|
||||
if images is None:
|
||||
images_on_cpu = None
|
||||
master_batch_size = 0
|
||||
else:
|
||||
images_on_cpu = ensure_contiguous(images.cpu())
|
||||
master_batch_size = images.shape[0]
|
||||
master_audio = audio # Keep master's audio for later
|
||||
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
|
||||
|
||||
|
||||
# Initialize storage for collected images and audio
|
||||
worker_images = {} # Dict to store images by worker_id and index
|
||||
worker_audio = {} # Dict to store audio by worker_id
|
||||
|
||||
# Collect images until all workers report they're done
|
||||
collected_count = 0
|
||||
workers_done = set()
|
||||
|
||||
# Use unified worker timeout from config/UI with simple sliced waits
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
|
||||
last_activity = time.time()
|
||||
|
||||
|
||||
# Get queue size before starting
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
initial_size = q.qsize()
|
||||
|
||||
# NEW: Initialize progress bar for workers (total = num_workers)
|
||||
p = ProgressBar(num_workers)
|
||||
|
||||
def mark_worker_done(done_worker_id):
|
||||
done_worker_id = str(done_worker_id)
|
||||
if done_worker_id not in expected_workers:
|
||||
debug_log(
|
||||
f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
|
||||
)
|
||||
return
|
||||
if done_worker_id in workers_done:
|
||||
debug_log(
|
||||
f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
|
||||
)
|
||||
return
|
||||
workers_done.add(done_worker_id)
|
||||
p.update(1) # +1 per completed expected worker
|
||||
|
||||
try:
|
||||
while len(workers_done) < num_workers:
|
||||
# Check for user interruption to abort collection promptly
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
try:
|
||||
# Get the queue again each time to ensure we have the right reference
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
current_size = q.qsize()
|
||||
|
||||
result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
count = self._store_worker_result(worker_images, result)
|
||||
collected_count += count
|
||||
debug_log(
|
||||
f"Master - Got canonical result from worker {worker_id}, "
|
||||
f"image {result.get('image_index', 0)}, is_last={is_last}"
|
||||
)
|
||||
|
||||
# Collect audio data if present
|
||||
result_audio = result.get('audio')
|
||||
if result_audio is not None:
|
||||
worker_audio[worker_id] = result_audio
|
||||
debug_log(f"Master - Got audio from worker {worker_id}")
|
||||
|
||||
# Record activity and refresh timeout baseline
|
||||
last_activity = time.time()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
|
||||
if is_last:
|
||||
mark_worker_done(worker_id)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# If we still have time, continue polling; otherwise handle timeout
|
||||
if (time.time() - last_activity) < base_timeout:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
continue
|
||||
# Re-check for user interruption after timeout expiry
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
missing_workers = set(str(w) for w in enabled_workers) - workers_done
|
||||
elapsed = time.time() - last_activity
|
||||
for missing_worker_id in sorted(missing_workers):
|
||||
log(
|
||||
"Master - Heartbeat timeout: "
|
||||
f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
log(
|
||||
f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
|
||||
f"(elapsed={elapsed:.1f}s)"
|
||||
)
|
||||
|
||||
# Probe missing workers' /prompt endpoints to check if they are actively processing
|
||||
any_busy = False
|
||||
try:
|
||||
cfg = load_config()
|
||||
cfg_workers = cfg.get('workers', [])
|
||||
for wid in list(missing_workers):
|
||||
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
|
||||
if not wrec:
|
||||
debug_log(f"Collector probe: worker {wid} not found in config")
|
||||
continue
|
||||
worker_url = build_worker_url(wrec)
|
||||
try:
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
queue_remaining = None
|
||||
if payload is not None:
|
||||
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
|
||||
debug_log(
|
||||
"Collector probe: worker "
|
||||
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
|
||||
)
|
||||
if payload is not None and queue_remaining and queue_remaining > 0:
|
||||
any_busy = True
|
||||
log(
|
||||
f"Master - Probe grace: worker {wid} appears busy "
|
||||
f"(queue_remaining={queue_remaining}). Continuing to wait."
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe failed for worker {wid}: {e}")
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe setup error: {e}")
|
||||
|
||||
if any_busy:
|
||||
# Refresh last_activity and continue waiting
|
||||
last_activity = time.time()
|
||||
# Refresh base timeout in case the user changed it in UI
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
continue
|
||||
|
||||
# Check queue size again with lock
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
final_q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
final_size = final_q.qsize()
|
||||
|
||||
# Try to drain any remaining items
|
||||
remaining_items = []
|
||||
while not final_q.empty():
|
||||
try:
|
||||
item = final_q.get_nowait()
|
||||
remaining_items.append(item)
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
if remaining_items:
|
||||
# Process them
|
||||
for item in remaining_items:
|
||||
worker_id = item['worker_id']
|
||||
is_last = item.get('is_last', False)
|
||||
|
||||
collected_count += self._store_worker_result(worker_images, item)
|
||||
|
||||
if is_last:
|
||||
mark_worker_done(worker_id)
|
||||
else:
|
||||
log(f"Master - Queue {multi_job_id} no longer exists!")
|
||||
break
|
||||
except comfy.model_management.InterruptProcessingException:
|
||||
# Cleanup queue on interruption and re-raise to abort prompt cleanly
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
raise
|
||||
|
||||
total_collected = sum(len(imgs) for imgs in worker_images.values())
|
||||
|
||||
# Clean up job queue
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
|
||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||
try:
|
||||
combined = self._reorder_and_combine_tensors(
|
||||
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
||||
)
|
||||
if combined is None:
|
||||
debug_log(f"Master - Job {multi_job_id} complete with audio only")
|
||||
else:
|
||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
|
||||
return (combined, combined_audio)
|
||||
except Exception as e:
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Preserve collected audio even when image assembly fails.
|
||||
return (images, combined_audio)
|
||||
@@ -1,270 +0,0 @@
|
||||
import json
|
||||
import math
|
||||
from functools import wraps
|
||||
|
||||
import comfy.samplers
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
|
||||
from ..upscale.tile_ops import TileOpsMixin
|
||||
from ..upscale.result_collector import ResultCollectorMixin
|
||||
from ..upscale.worker_comms import WorkerCommsMixin
|
||||
from ..upscale.job_state import JobStateMixin
|
||||
from ..upscale.modes.single_gpu import SingleGpuModeMixin
|
||||
from ..upscale.modes.static import StaticModeMixin
|
||||
from ..upscale.modes.dynamic import DynamicModeMixin
|
||||
|
||||
def sync_wrapper(async_func):
|
||||
"""Decorator to wrap async methods for synchronous execution."""
|
||||
@wraps(async_func)
|
||||
def sync_func(self, *args, **kwargs):
|
||||
# Use run_async_in_server_loop for ComfyUI compatibility
|
||||
return run_async_in_server_loop(
|
||||
async_func(self, *args, **kwargs),
|
||||
timeout=600.0 # 10 minute timeout for long operations
|
||||
)
|
||||
return sync_func
|
||||
|
||||
def _parse_enabled_worker_ids(enabled_worker_ids):
|
||||
"""Parse enabled worker IDs from either JSON or list input."""
|
||||
if isinstance(enabled_worker_ids, list):
|
||||
return [str(worker_id) for worker_id in enabled_worker_ids]
|
||||
if not enabled_worker_ids:
|
||||
return []
|
||||
if isinstance(enabled_worker_ids, str):
|
||||
try:
|
||||
parsed = json.loads(enabled_worker_ids)
|
||||
except json.JSONDecodeError:
|
||||
log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
|
||||
return []
|
||||
if isinstance(parsed, list):
|
||||
return [str(wid) for wid in parsed]
|
||||
return []
|
||||
|
||||
class UltimateSDUpscaleDistributed(
|
||||
DynamicModeMixin,
|
||||
StaticModeMixin,
|
||||
SingleGpuModeMixin,
|
||||
ResultCollectorMixin,
|
||||
WorkerCommsMixin,
|
||||
JobStateMixin,
|
||||
TileOpsMixin,
|
||||
):
|
||||
|
||||
"""
|
||||
Distributed version of Ultimate SD Upscale (No Upscale).
|
||||
|
||||
Supports two currently selected processing modes:
|
||||
1. Single GPU: No workers available, process everything locally
|
||||
2. Distributed tile queue: Workers pull tiles from a shared queue
|
||||
|
||||
Features:
|
||||
- Tile-based batch handling for video/image upscaling
|
||||
- Tiled VAE support for memory efficiency
|
||||
- Shared work queue so faster workers can process more tiles
|
||||
- Backward compatible with single-image workflows
|
||||
|
||||
Environment Variables:
|
||||
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
||||
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
||||
|
||||
The hidden dynamic_threshold input is retained for workflow compatibility but
|
||||
does not affect the current mode-selection policy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the node and ensure persistent storage exists."""
|
||||
# Pre-initialize the persistent storage on node creation
|
||||
ensure_tile_jobs_initialized()
|
||||
debug_log("UltimateSDUpscaleDistributed - Node initialized")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"upscaled_image": ("IMAGE",),
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
||||
"denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||||
"tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||||
"padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
|
||||
"mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}),
|
||||
"force_uniform_tiles": ("BOOLEAN", {"default": True}),
|
||||
"tiled_decode": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility
|
||||
"dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image/upscaling"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
"""Force re-execution."""
|
||||
return float("nan") # Always re-execute
|
||||
|
||||
def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode,
|
||||
multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]",
|
||||
worker_id="", tile_indices="", dynamic_threshold=8):
|
||||
"""Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
|
||||
# Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
|
||||
try:
|
||||
batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
|
||||
except Exception:
|
||||
batch_size = 1
|
||||
# Enforce 4n+1 batches globally for any model when batch > 1 (master only)
|
||||
if not is_worker and batch_size != 1 and (batch_size % 4 != 1):
|
||||
raise ValueError(
|
||||
f"Batch size {batch_size} is not of the form 4n+1. "
|
||||
"This node requires batch sizes of 1 or 4n+1 (1, 5, 9, 13, ...). "
|
||||
"Please adjust the batch size."
|
||||
)
|
||||
if not multi_job_id:
|
||||
# No distributed processing, run single GPU version
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
if is_worker:
|
||||
# Worker mode: process tiles synchronously
|
||||
return self.process_worker(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
else:
|
||||
# Master mode: distribute and collect synchronously
|
||||
return self.process_master(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold)
|
||||
|
||||
def process_worker(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold):
|
||||
"""Unified worker processing - handles both static and dynamic modes."""
|
||||
# Get batch size to determine mode
|
||||
batch_size = upscaled_image.shape[0]
|
||||
|
||||
# Ensure mode consistency across master/workers via shared threshold
|
||||
# Determine mode (must match master's logic)
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
# Compute number of tiles for this image to decide if tile distribution makes sense
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# For USDU-style processing, we want tile distribution whenever workers are available
|
||||
# and there is more than one tile to process, even if batch == 1.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
debug_log(f"USDU Dist Worker - Batch size {batch_size}")
|
||||
|
||||
if mode == "dynamic":
|
||||
return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
|
||||
# Static mode - enhanced with health monitoring and retry logic
|
||||
return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_workers)
|
||||
|
||||
def process_master(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold):
|
||||
"""Unified master processing with enhanced monitoring and failure handling."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get image dimensions and batch size
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
|
||||
# Calculate all tiles and grid
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
rows = math.ceil(height / tile_height)
|
||||
cols = math.ceil(width / tile_width)
|
||||
log(
|
||||
f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
|
||||
)
|
||||
|
||||
# Parse enabled workers
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
# Determine processing mode
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# Prefer tile-based static distribution when workers are available and there are multiple tiles,
|
||||
# even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
|
||||
|
||||
if mode == "single_gpu":
|
||||
# No workers, process all tiles locally
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
elif mode == "dynamic":
|
||||
# Dynamic mode for large batches
|
||||
return self.process_master_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers)
|
||||
|
||||
# Static mode - enhanced with unified job management
|
||||
return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
|
||||
all_tiles, num_tiles_per_image)
|
||||
|
||||
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
||||
"""Determines processing mode per requested policy:
|
||||
- any workers => prefer static (tile-based) for USDU
|
||||
- no workers => single_gpu
|
||||
"""
|
||||
if num_workers == 0:
|
||||
return "single_gpu"
|
||||
# Default to static when distributed; master/worker may still override if special cases arise
|
||||
return "static"
|
||||
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
@@ -1,354 +0,0 @@
|
||||
import torch
|
||||
import json
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
|
||||
|
||||
def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]:
|
||||
"""Return contiguous [start, end) bounds for n_splits chunks."""
|
||||
split_count = max(1, int(n_splits))
|
||||
total = max(0, int(total_items))
|
||||
base, remainder = divmod(total, split_count)
|
||||
|
||||
bounds: list[tuple[int, int]] = []
|
||||
start = 0
|
||||
for idx in range(split_count):
|
||||
size = base + (1 if idx < remainder else 0)
|
||||
end = start + size
|
||||
bounds.append((start, end))
|
||||
start = end
|
||||
return bounds
|
||||
|
||||
|
||||
class DistributedSeed:
|
||||
"""
|
||||
Distributes seed values across multiple GPUs.
|
||||
On master: passes through the original seed.
|
||||
On workers: adds offset based on worker ID.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"seed": ("INT", {
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": False # Widget by default, can be converted to input
|
||||
}),
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("seed",)
|
||||
FUNCTION = "distribute"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def distribute(self, seed, is_worker=False, worker_id=""):
|
||||
if not is_worker:
|
||||
# Master node: pass through original values
|
||||
debug_log(f"Distributor - Master: seed={seed}")
|
||||
return (seed,)
|
||||
else:
|
||||
# Worker node: apply offset based on worker index
|
||||
# Find worker index from enabled_worker_ids
|
||||
try:
|
||||
# Worker IDs are passed as "worker_0", "worker_1", etc.
|
||||
if worker_id.startswith("worker_"):
|
||||
worker_index = int(worker_id.split("_")[1])
|
||||
else:
|
||||
# Fallback: try to parse as direct index
|
||||
worker_index = int(worker_id)
|
||||
|
||||
offset = worker_index + 1
|
||||
new_seed = seed + offset
|
||||
debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
|
||||
return (new_seed,)
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
|
||||
# Fallback: return original seed
|
||||
return (seed,)
|
||||
|
||||
|
||||
# Define ByPassTypeTuple for flexible return types
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
any_type = AnyType("*")
|
||||
|
||||
|
||||
class DistributedValue:
|
||||
"""
|
||||
Outputs a different value per worker.
|
||||
On master: returns default_value.
|
||||
On workers: looks up the worker-specific value from a JSON map,
|
||||
falling back to default_value if not set.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"default_value": ("STRING", {"default": ""}),
|
||||
"worker_values": ("STRING", {"default": "{}"}),
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
RETURN_NAMES = ("value",)
|
||||
FUNCTION = "distribute"
|
||||
CATEGORY = "utils"
|
||||
|
||||
@staticmethod
|
||||
def _coerce(value, value_type):
|
||||
"""Convert a string value to the requested type."""
|
||||
if value_type == "INT":
|
||||
return int(float(value))
|
||||
if value_type == "FLOAT":
|
||||
return float(value)
|
||||
return value # STRING and COMBO stay as strings
|
||||
|
||||
@staticmethod
|
||||
def _coerce_safe(value, value_type):
|
||||
"""Best-effort coercion with graceful fallback to original value."""
|
||||
try:
|
||||
return DistributedValue._coerce(value, value_type)
|
||||
except (TypeError, ValueError):
|
||||
return value
|
||||
|
||||
def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
|
||||
values = {}
|
||||
value_type = "STRING"
|
||||
|
||||
try:
|
||||
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
|
||||
if not isinstance(values, dict):
|
||||
values = {}
|
||||
except json.JSONDecodeError as e:
|
||||
debug_log(f"DistributedValue - Error parsing worker_values: {e}")
|
||||
values = {}
|
||||
|
||||
value_type = values.get("_type", "STRING")
|
||||
coerced_default = self._coerce_safe(default_value, value_type)
|
||||
|
||||
if not is_worker:
|
||||
debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
try:
|
||||
if worker_id.startswith("worker_"):
|
||||
idx = int(worker_id.split("_")[1])
|
||||
else:
|
||||
idx = int(worker_id)
|
||||
key = str(idx + 1) # worker_0 → key "1" (1-indexed)
|
||||
raw = values.get(key, "")
|
||||
if raw:
|
||||
coerced = self._coerce(raw, value_type)
|
||||
debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
|
||||
return (coerced,)
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"DistributedValue - Error: {e}")
|
||||
debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
class DistributedModelName:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"default": ""}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
RETURN_NAMES = ("output",)
|
||||
FUNCTION = "log_input"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "utils"
|
||||
|
||||
def _stringify(self, value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (int, float, bool)):
|
||||
return str(value)
|
||||
try:
|
||||
return json.dumps(value, indent=4)
|
||||
except Exception:
|
||||
return str(value)
|
||||
|
||||
def _update_workflow(self, extra_pnginfo, unique_id, values):
|
||||
if not extra_pnginfo:
|
||||
return
|
||||
info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
|
||||
if not isinstance(info, dict) or "workflow" not in info:
|
||||
return
|
||||
node_id = None
|
||||
if isinstance(unique_id, list) and unique_id:
|
||||
node_id = str(unique_id[0])
|
||||
elif unique_id is not None:
|
||||
node_id = str(unique_id)
|
||||
if not node_id:
|
||||
return
|
||||
workflow = info["workflow"]
|
||||
node = next((x for x in workflow["nodes"] if str(x.get("id")) == node_id), None)
|
||||
if node:
|
||||
node["widgets_values"] = [values]
|
||||
|
||||
def log_input(self, text, unique_id=None, extra_pnginfo=None):
|
||||
values = []
|
||||
if isinstance(text, list):
|
||||
for val in text:
|
||||
values.append(self._stringify(val))
|
||||
else:
|
||||
values.append(self._stringify(text))
|
||||
|
||||
# Keep widget display in workflow metadata if available.
|
||||
self._update_workflow(extra_pnginfo, unique_id, values)
|
||||
|
||||
if isinstance(values, list) and len(values) == 1:
|
||||
return {"ui": {"text": values}, "result": (values[0],)}
|
||||
return {"ui": {"text": values}, "result": (values,)}
|
||||
|
||||
class ByPassTypeTuple(tuple):
|
||||
def __getitem__(self, index):
|
||||
if index > 0:
|
||||
index = 0
|
||||
item = super().__getitem__(index)
|
||||
if isinstance(item, str):
|
||||
return any_type
|
||||
return item
|
||||
|
||||
class ImageBatchDivider:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"divide_by": ("INT", {
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ByPassTypeTuple(("IMAGE", )) # Flexible for variable outputs
|
||||
RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)]))
|
||||
FUNCTION = "divide_batch"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "image"
|
||||
|
||||
def divide_batch(self, images, divide_by):
|
||||
total_splits = max(1, min(int(divide_by), 10))
|
||||
total_frames = images.shape[0]
|
||||
empty_tensor = images[:0]
|
||||
bounds = _chunk_bounds(total_frames, total_splits)
|
||||
outputs = [images[start:end] if end > start else empty_tensor for start, end in bounds]
|
||||
|
||||
while len(outputs) < 10:
|
||||
outputs.append(empty_tensor)
|
||||
|
||||
return tuple(outputs[:10])
|
||||
|
||||
|
||||
class AudioBatchDivider:
|
||||
"""Divides an audio waveform into sequential segments along the time dimension."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"divide_by": ("INT", {
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ByPassTypeTuple(("AUDIO",)) # Flexible for variable outputs
|
||||
RETURN_NAMES = ByPassTypeTuple(tuple([f"audio_{i+1}" for i in range(10)]))
|
||||
FUNCTION = "divide_audio"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "audio"
|
||||
|
||||
def divide_audio(self, audio, divide_by):
|
||||
import torch
|
||||
|
||||
waveform = audio.get("waveform")
|
||||
sample_rate = audio.get("sample_rate", 44100)
|
||||
|
||||
if waveform is None or waveform.numel() == 0:
|
||||
# Return empty audio for all outputs
|
||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": sample_rate}
|
||||
return tuple([empty_audio] * 10)
|
||||
|
||||
total_splits = max(1, min(int(divide_by), 10))
|
||||
total_samples = int(waveform.shape[-1])
|
||||
bounds = _chunk_bounds(total_samples, total_splits)
|
||||
|
||||
outputs = []
|
||||
empty_waveform = waveform[..., :0]
|
||||
for start, end in bounds:
|
||||
split_waveform = waveform[..., start:end] if end > start else empty_waveform
|
||||
outputs.append({
|
||||
"waveform": split_waveform,
|
||||
"sample_rate": sample_rate
|
||||
})
|
||||
|
||||
# Pad with empty audio up to max (10) to match RETURN_TYPES length
|
||||
empty_audio = {
|
||||
"waveform": empty_waveform,
|
||||
"sample_rate": sample_rate
|
||||
}
|
||||
|
||||
while len(outputs) < 10:
|
||||
outputs.append(empty_audio)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class DistributedEmptyImage:
|
||||
"""Produces an empty IMAGE batch used when the master delegates all work."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
|
||||
"width": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
|
||||
"channels": ("INT", {"default": 3, "min": 1, "max": 4, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "create"
|
||||
CATEGORY = "image"
|
||||
|
||||
def create(self, height, width, channels):
|
||||
import torch
|
||||
|
||||
shape = (0, height, width, channels)
|
||||
tensor = torch.zeros(shape, dtype=torch.float32)
|
||||
return (tensor,)
|
||||
-191
@@ -1,191 +0,0 @@
|
||||
"""Native V3 schemas with the existing execution algorithms kept intact.
|
||||
|
||||
Runtime objects are private, per-execution helpers, not registered V1 nodes.
|
||||
This avoids sharing mutable instance state through V3's sanitized class clones.
|
||||
"""
|
||||
import comfy.samplers
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from ..utils import transfer_codec as _transfer_codec
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
|
||||
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
|
||||
max=1125899906842624, force_input=False)],
|
||||
outputs=[io.Int.Output(display_name='seed')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed, is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedValue(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedValue', display_name='Distributed Value', category='utils',
|
||||
inputs=[io.String.Input('default_value', default=''),
|
||||
io.String.Input('worker_values', default='{}')],
|
||||
outputs=[io.AnyType.Output(display_name='value')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedValue().distribute(
|
||||
default_value, worker_values, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedModelName(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
|
||||
inputs=[io.String.Input('text', default='')],
|
||||
outputs=[io.AnyType.Output(display_name='output')],
|
||||
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text):
|
||||
result = _utilities.DistributedModelName().log_input(
|
||||
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
|
||||
return io.NodeOutput(*result['result'], ui=result['ui'])
|
||||
|
||||
|
||||
class ImageBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
|
||||
inputs=[io.Image.Input('images'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of parts to divide the batch into')],
|
||||
# The existing frontend still displays only divide_by sockets.
|
||||
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, divide_by):
|
||||
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
|
||||
|
||||
|
||||
class AudioBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
|
||||
inputs=[io.Audio.Input('audio'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of sequential time segments to create')],
|
||||
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, divide_by):
|
||||
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
|
||||
|
||||
|
||||
class DistributedEmptyImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
|
||||
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('width', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
|
||||
outputs=[io.Image.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, height, width, channels):
|
||||
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
|
||||
|
||||
|
||||
class DistributedCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False,
|
||||
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
|
||||
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
|
||||
io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
|
||||
default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,
|
||||
tooltip=_transfer_codec.TRANSFER_FORMAT_TOOLTIP),
|
||||
io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
|
||||
min=1, max=100, step=1, optional=True,
|
||||
tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
|
||||
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
|
||||
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
|
||||
worker_id='', pass_through=False, delegate_only=False,
|
||||
transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
|
||||
transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
return io.NodeOutput(*_CollectorRuntime().run(
|
||||
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
|
||||
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
|
||||
worker_batch_size=worker_batch_size, worker_id=worker_id,
|
||||
pass_through=pass_through, delegate_only=delegate_only,
|
||||
transfer_format=transfer_format, transfer_quality=transfer_quality))
|
||||
|
||||
|
||||
class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='UltimateSDUpscaleDistributed',
|
||||
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
|
||||
inputs=[
|
||||
io.Image.Input('upscaled_image'), io.Model.Input('model'),
|
||||
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
|
||||
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Int.Input('steps', default=20, min=1, max=10000),
|
||||
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
|
||||
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
|
||||
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
|
||||
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
|
||||
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('padding', default=32, min=0, max=256, step=8),
|
||||
io.Int.Input('mask_blur', default=8, min=0, max=256),
|
||||
io.Boolean.Input('force_uniform_tiles', default=True),
|
||||
io.Boolean.Input('tiled_decode', default=False),
|
||||
], outputs=[io.Image.Output()], accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
return _UpscaleRuntime.IS_CHANGED(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
|
||||
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
|
||||
return io.NodeOutput(*_UpscaleRuntime().run(
|
||||
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
|
||||
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
|
||||
Generated
-1419
File diff suppressed because it is too large
Load Diff
@@ -1,12 +0,0 @@
|
||||
{
|
||||
"name": "comfyui-distributed-web-tests",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"test:web": "bash ./scripts/test-web.sh",
|
||||
"test:web:watch": "bash ./scripts/test-web.sh --watch"
|
||||
},
|
||||
"devDependencies": {
|
||||
"vitest": "^2.1.9"
|
||||
}
|
||||
}
|
||||
+1
-5
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.5.0"
|
||||
version = "1.2.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
@@ -13,7 +13,3 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
PublisherId = "robertvoy"
|
||||
DisplayName = "ComfyUI-Distributed"
|
||||
Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
pythonpath = ["."]
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Internal extension lifecycle support (not a node provider)."""
|
||||
@@ -1,35 +0,0 @@
|
||||
"""Initialize routes, distributed state and the existing worker lifecycle."""
|
||||
import atexit
|
||||
import os
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.config import CONFIG_FILE, ensure_config_exists
|
||||
from ..utils.logging import debug_log
|
||||
from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
|
||||
_initialized = False
|
||||
|
||||
|
||||
def initialize():
|
||||
"""Called by ComfyExtension.on_load; initialize once per loaded package."""
|
||||
global _initialized
|
||||
if _initialized:
|
||||
return
|
||||
|
||||
ensure_config_exists()
|
||||
from .. import api # noqa: F401 - registers the existing @routes.* handlers
|
||||
from ..api.queue_orchestration import ensure_distributed_state
|
||||
|
||||
ensure_distributed_state(server.PromptServer.instance)
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
if not os.environ.get('COMFYUI_IS_WORKER'):
|
||||
atexit.register(sync_cleanup)
|
||||
delayed_auto_launch()
|
||||
register_async_signals()
|
||||
|
||||
_initialized = True
|
||||
debug_log('Loaded Distributed nodes.')
|
||||
debug_log(f'Config file: {CONFIG_FILE}')
|
||||
@@ -1,29 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
|
||||
if [[ -s "${NVM_DIR}/nvm.sh" ]]; then
|
||||
# shellcheck source=/dev/null
|
||||
. "${NVM_DIR}/nvm.sh"
|
||||
fi
|
||||
|
||||
if ! command -v node >/dev/null 2>&1 || ! command -v npm >/dev/null 2>&1; then
|
||||
echo "[test-web] node/npm are not available." >&2
|
||||
echo "[test-web] Install nvm and Node, or ensure node/npm are on PATH." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -f "${REPO_ROOT}/.nvmrc" ]] && command -v nvm >/dev/null 2>&1; then
|
||||
nvm use >/dev/null
|
||||
fi
|
||||
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
if [[ "${1:-}" == "--watch" ]]; then
|
||||
exec npx vitest web/tests
|
||||
fi
|
||||
|
||||
exec npx vitest run web/tests
|
||||
@@ -1,148 +0,0 @@
|
||||
import copy
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None):
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def _load_config_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
|
||||
package_name = "dist_api_config_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
default_config = {
|
||||
"workers": [],
|
||||
"master": {"host": ""},
|
||||
"settings": {"debug": False},
|
||||
"tunnel": {},
|
||||
}
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: copy.deepcopy(default_config)
|
||||
config_module.save_config = lambda _cfg: True
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
config_routes = _load_config_routes_module()
|
||||
|
||||
|
||||
class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_get_config_returns_core_sections(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIn("workers", response.payload)
|
||||
self.assertIn("master", response.payload)
|
||||
self.assertIn("settings", response.payload)
|
||||
|
||||
async def test_update_config_valid_field_persists(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
|
||||
config_routes, "save_config", return_value=True
|
||||
):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload["status"], "success")
|
||||
self.assertTrue(response.payload["config"]["settings"]["debug"])
|
||||
|
||||
async def test_update_config_unknown_field_returns_400(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"unknown_field": 1}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("unknown_field", " ".join(response.payload.get("error", [])).lower())
|
||||
|
||||
async def test_update_config_wrong_type_returns_400(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"debug": "true"}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("debug", " ".join(response.payload.get("error", [])).lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,551 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload):
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
||||
codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
||||
codec_module = importlib.util.module_from_spec(codec_spec)
|
||||
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
||||
codec_spec.loader.exec_module(codec_module)
|
||||
utils_pkg.transfer_codec = codec_module
|
||||
return codec_module
|
||||
|
||||
|
||||
def _load_job_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
||||
package_name = "dist_api_queue_testpkg"
|
||||
|
||||
# Reset package namespace to avoid stale module state across test runs.
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
_register_real_transfer_codec(package_name, utils_pkg)
|
||||
|
||||
# aiohttp.web stub
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
# server module stub with route decorators
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
prompt_server_instance = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
# torch stub (only needed to satisfy import)
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
# PIL stub (only needed to satisfy import)
|
||||
created_pil_stub = False
|
||||
if "PIL" not in sys.modules:
|
||||
created_pil_stub = True
|
||||
pil_module = types.ModuleType("PIL")
|
||||
image_module = types.ModuleType("PIL.Image")
|
||||
pil_module.Image = image_module
|
||||
sys.modules["PIL"] = pil_module
|
||||
sys.modules["PIL.Image"] = image_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
||||
image_module.ensure_contiguous = lambda tensor: tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
|
||||
def _decode_audio_payload(payload):
|
||||
if payload is None:
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Field 'audio' must be an object when provided.")
|
||||
|
||||
encoded = payload.get("data")
|
||||
shape = payload.get("shape")
|
||||
dtype = payload.get("dtype", "float32")
|
||||
sample_rate = payload.get("sample_rate", 44100)
|
||||
if not isinstance(encoded, str) or not encoded.strip():
|
||||
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
||||
if not isinstance(shape, list) or len(shape) != 3:
|
||||
raise ValueError("Field 'audio.shape' must be a 3-item list.")
|
||||
if dtype != "float32":
|
||||
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
||||
try:
|
||||
shape_tuple = tuple(int(dim) for dim in shape)
|
||||
except Exception as exc:
|
||||
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
||||
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError("Field 'audio.data' byte size mismatch.")
|
||||
|
||||
waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
|
||||
return {"waveform": waveform, "sample_rate": int(sample_rate)}
|
||||
|
||||
audio_payload_module.decode_audio_payload = _decode_audio_payload
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
||||
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _QueuePayload:
|
||||
prompt: dict
|
||||
workflow_meta: object
|
||||
client_id: str
|
||||
delegate_master: object
|
||||
enabled_worker_ids: list
|
||||
auto_prepare: bool
|
||||
trace_execution_id: object
|
||||
|
||||
def _parse_queue_request_payload(data):
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
enabled = data.get("enabled_worker_ids")
|
||||
if not isinstance(enabled, list):
|
||||
raise ValueError("enabled_worker_ids required")
|
||||
client_id = data.get("client_id")
|
||||
if not isinstance(client_id, str) or not client_id.strip():
|
||||
raise ValueError("client_id required")
|
||||
return _QueuePayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=data.get("delegate_master"),
|
||||
enabled_worker_ids=enabled,
|
||||
auto_prepare=bool(data.get("auto_prepare", True)),
|
||||
trace_execution_id=data.get("trace_execution_id"),
|
||||
)
|
||||
|
||||
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
||||
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
||||
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
if created_pil_stub:
|
||||
sys.modules.pop("PIL.Image", None)
|
||||
sys.modules.pop("PIL", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
job_routes = _load_job_routes_module()
|
||||
|
||||
|
||||
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
"auto_prepare": True,
|
||||
}
|
||||
)
|
||||
with patch.object(
|
||||
job_routes,
|
||||
"orchestrate_distributed_execution",
|
||||
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
||||
):
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
||||
self.assertEqual(response.payload.get("number"), 42)
|
||||
self.assertEqual(response.payload.get("node_errors"), {})
|
||||
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
||||
|
||||
async def test_distributed_queue_missing_prompt_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("prompt", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
||||
|
||||
|
||||
class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
def _encoded_audio_payload(self):
|
||||
waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
|
||||
return {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(waveform.tobytes()).decode("ascii"),
|
||||
}
|
||||
|
||||
async def test_job_complete_accepts_audio_payload(self):
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-1")
|
||||
self.assertTrue(queued["is_last"])
|
||||
self.assertIsNotNone(queued["audio"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_job_complete_accepts_audio_without_image(self):
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "audio-only-job",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
decode_image.assert_not_called()
|
||||
queued = await queue.get()
|
||||
self.assertIsNone(queued["tensor"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
|
||||
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": 123,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_returns_400_for_malformed_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_shape(self):
|
||||
bad = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(bad)
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_dtype(self):
|
||||
payload = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float16",
|
||||
"data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(payload)
|
||||
|
||||
|
||||
class _FakeFileField:
|
||||
def __init__(self, data):
|
||||
self.file = io.BytesIO(data)
|
||||
|
||||
|
||||
class _FakeMultipartRequest:
|
||||
content_type = "multipart/form-data"
|
||||
|
||||
def __init__(self, envelope, parts):
|
||||
from multidict import MultiDict
|
||||
|
||||
fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
|
||||
fields += [("media", _FakeFileField(part)) for part in parts]
|
||||
self._form = MultiDict(fields)
|
||||
|
||||
async def post(self):
|
||||
return self._form
|
||||
|
||||
async def json(self):
|
||||
raise ValueError("not JSON")
|
||||
|
||||
|
||||
class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
codec = job_routes.transfer_codec
|
||||
|
||||
def setUp(self):
|
||||
self.queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
|
||||
|
||||
def _frames(self, count=4, height=20, width=30):
|
||||
rng = np.random.default_rng(1)
|
||||
return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
|
||||
|
||||
def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
|
||||
chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
|
||||
parts = chunk.pop("parts")
|
||||
envelope = {
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"is_last": is_last,
|
||||
**chunk,
|
||||
"start_index": chunk["start_index"] + start_offset,
|
||||
}
|
||||
if audio is not None:
|
||||
envelope["audio"] = audio
|
||||
envelope.update(overrides)
|
||||
return _FakeMultipartRequest(envelope, parts)
|
||||
|
||||
async def _drain(self):
|
||||
items = []
|
||||
while not self.queue.empty():
|
||||
items.append(await self.queue.get())
|
||||
return items
|
||||
|
||||
async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
|
||||
frames = self._frames()
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
|
||||
self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
|
||||
self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
|
||||
for item, frame in zip(items, frames):
|
||||
self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
|
||||
self.assertEqual(item["tensor"].dtype, torch.float32)
|
||||
expected = frame.astype(np.float32) / 255.0
|
||||
self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
|
||||
|
||||
async def test_non_final_chunk_never_marks_worker_done(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "png", is_last=False)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertFalse(any(item["is_last"] for item in await self._drain()))
|
||||
|
||||
async def test_audio_is_attached_only_to_final_frame(self):
|
||||
audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=3), "jpeg", audio=audio)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
|
||||
self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_rejects_frame_count_mismatch(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "ffv1", count=5)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("frames", response.payload["message"])
|
||||
self.assertTrue(self.queue.empty())
|
||||
|
||||
async def test_rejects_bad_envelope(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("format", response.payload["message"])
|
||||
self.assertIn("is_last", response.payload["message"])
|
||||
|
||||
async def test_rejects_invalid_envelope_json(self):
|
||||
response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("envelope", response.payload["message"])
|
||||
|
||||
async def test_unknown_job_returns_404(self):
|
||||
job_routes.prompt_server.distributed_pending_jobs = {}
|
||||
with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
|
||||
|
||||
self.assertEqual(response.status, 404)
|
||||
|
||||
async def test_legacy_json_payload_still_accepted(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 3,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual(len(items), 1)
|
||||
self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,247 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_media_sync_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
|
||||
package_name = "dist_ms_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost{endpoint}"
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in pure-function tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
|
||||
trace_module.trace_debug = lambda *_args, **_kwargs: None
|
||||
trace_module.trace_info = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
class _FormData:
|
||||
def add_field(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.FormData = _FormData
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.media_sync",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
ms = _load_media_sync_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# convert_paths_for_platform
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ConvertPathsForPlatformTests(unittest.TestCase):
|
||||
def test_forward_slash_target_normalises_backslashes(self):
|
||||
obj = {"ckpt_name": "C:\\Models\\model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_backslash_target_normalises_forward_slashes(self):
|
||||
obj = {"ckpt_name": "/models/checkpoints/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("\\", result["ckpt_name"])
|
||||
self.assertNotIn("/", result["ckpt_name"])
|
||||
|
||||
def test_relative_media_paths_always_stay_forward_slash(self):
|
||||
"""Relative image/video/audio paths (Comfy annotated style) must not be backslash-ified."""
|
||||
obj = {"image": "subfolder/my_photo.png"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["image"], "subfolder/my_photo.png")
|
||||
|
||||
def test_relative_audio_paths_stay_forward_slash(self):
|
||||
obj = {"audio": "subfolder/my_track.wav"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["audio"], "subfolder/my_track.wav")
|
||||
|
||||
def test_annotated_relative_media_path_stays_forward_slash(self):
|
||||
obj = {"image": "input/frame.jpg [abc123]"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("/", result["image"])
|
||||
self.assertNotIn("\\", result["image"].split("[")[0])
|
||||
|
||||
def test_non_filename_strings_are_untouched(self):
|
||||
obj = {"prompt": "a beautiful cat", "count": 5}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["prompt"], "a beautiful cat")
|
||||
self.assertEqual(result["count"], 5)
|
||||
|
||||
def test_url_strings_are_untouched(self):
|
||||
obj = {"url": "https://example.com/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["url"], "https://example.com/model.safetensors")
|
||||
|
||||
def test_invalid_separator_returns_obj_unchanged(self):
|
||||
obj = {"ckpt_name": "/models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "|")
|
||||
self.assertEqual(result, obj)
|
||||
|
||||
def test_nested_dict_is_processed_recursively(self):
|
||||
obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_list_items_are_processed_recursively(self):
|
||||
obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"]
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
self.assertEqual(result[1], "plain string")
|
||||
|
||||
def test_non_string_scalar_values_are_untouched(self):
|
||||
obj = {"seed": 42, "enabled": True, "ratio": 1.5}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["seed"], 42)
|
||||
self.assertTrue(result["enabled"])
|
||||
|
||||
def test_absolute_unix_path_to_windows(self):
|
||||
obj = {"lora": "/home/user/loras/my_lora.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertNotIn("/", result["lora"])
|
||||
|
||||
def test_already_normalised_path_is_idempotent(self):
|
||||
obj = {"ckpt": "C:/Models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt"], "C:/Models/model.safetensors")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_media_references
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindMediaReferencesTests(unittest.TestCase):
|
||||
def test_finds_file_input_for_load_video(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("1 - Copy.mp4", refs)
|
||||
|
||||
def test_strips_annotation_suffix(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.jpg", refs)
|
||||
self.assertFalse(any("[" in r for r in refs))
|
||||
|
||||
def test_normalises_backslashes_in_path(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("sub/img.png", refs)
|
||||
|
||||
def test_ignores_non_media_text_inputs(self):
|
||||
prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_ignores_node_link_values(self):
|
||||
"""Inputs that are [node_id, slot] lists should be ignored."""
|
||||
prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_deduplicates_same_file_across_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(len(refs), 1)
|
||||
|
||||
def test_returns_sorted_list(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, sorted(refs))
|
||||
|
||||
def test_ignores_non_dict_nodes(self):
|
||||
prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("img.png", refs)
|
||||
|
||||
def test_empty_prompt_returns_empty_list(self):
|
||||
self.assertEqual(ms._find_media_references({}), [])
|
||||
|
||||
def test_multiple_media_types_all_found(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}},
|
||||
"2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}},
|
||||
"3": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("frame.png", refs)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
self.assertIn("track.wav", refs)
|
||||
|
||||
|
||||
class RewritePromptMediaInputsTests(unittest.TestCase):
|
||||
def test_rewrites_video_file_input_to_worker_path(self):
|
||||
prompt = {
|
||||
"79": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}},
|
||||
}
|
||||
ms._rewrite_prompt_media_inputs(prompt, {"1 - Copy.mp4": "videos/1 - Copy.mp4"})
|
||||
self.assertEqual(prompt["79"]["inputs"]["file"], "videos/1 - Copy.mp4")
|
||||
|
||||
def test_rewrites_audio_input_and_strips_annotation_when_matching(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadAudio", "inputs": {"audio": "song.wav [input]"}},
|
||||
}
|
||||
ms._rewrite_prompt_media_inputs(prompt, {"song.wav": "song.wav"})
|
||||
self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,216 +0,0 @@
|
||||
"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
_TESTS_DIR = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_from(path, name):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _load_modules():
|
||||
real_aiohttp = sys.modules["aiohttp"] # imported inside the isolated test
|
||||
queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
|
||||
collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
|
||||
job_routes = queue_tests._load_job_routes_module()
|
||||
collector = collector_tests._load_collector_module()
|
||||
|
||||
# Error responses must be real HTTP responses here (the unit-test harness
|
||||
# returns plain objects); mirror utils.network.handle_api_error.
|
||||
from aiohttp import web
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
if isinstance(error, list):
|
||||
return web.json_response({"errors": [str(item) for item in error]}, status=status)
|
||||
return web.json_response({"error": str(error)}, status=status)
|
||||
|
||||
job_routes.handle_api_error = handle_api_error
|
||||
# The collector harness stubs aiohttp; this test needs the real client.
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
collector.aiohttp = real_aiohttp
|
||||
return job_routes, collector
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_sys_modules():
|
||||
"""The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
|
||||
saved = dict(sys.modules)
|
||||
# Earlier test modules may have left a stub `aiohttp` behind; drop it so the
|
||||
# real package is imported for this test (restored afterwards).
|
||||
for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
|
||||
if getattr(sys.modules[name], "__file__", None) is None:
|
||||
del sys.modules[name]
|
||||
yield
|
||||
sys.modules.clear()
|
||||
sys.modules.update(saved)
|
||||
|
||||
|
||||
async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
|
||||
request_limit, chunk_target, audio=None):
|
||||
"""Send `images` from a real worker node to the real completion route.
|
||||
|
||||
Returns (queued items, sizes of every request body the server received,
|
||||
exception raised by the worker or None).
|
||||
"""
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
body_sizes = []
|
||||
|
||||
@web.middleware
|
||||
async def record_size(request, handler):
|
||||
body_sizes.append(request.content_length)
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
error = None
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
try:
|
||||
await node.send_batch_to_master(
|
||||
images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
except Exception as exc: # surfaced to the caller for assertions
|
||||
error = exc
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items, body_sizes, error
|
||||
|
||||
|
||||
def _noise_images(count, height=360, width=640, seed=97):
|
||||
return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
|
||||
def test_high_entropy_video_stays_under_server_upload_limit(fmt):
|
||||
"""Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(24)
|
||||
limit = 2 * 1024 * 1024
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, fmt,
|
||||
client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
|
||||
|
||||
assert error is None, error
|
||||
assert len(body_sizes) > 1 and max(body_sizes) <= limit
|
||||
assert [item["image_index"] for item in items] == list(range(24))
|
||||
assert [item["is_last"] for item in items] == [False] * 23 + [True]
|
||||
|
||||
|
||||
def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
|
||||
"""A worker configured above the master's limit fails loudly instead of
|
||||
silently resending the whole batch as uncompressed PNG."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(8)
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "ffv1",
|
||||
client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
|
||||
|
||||
assert isinstance(error, RuntimeError)
|
||||
assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
|
||||
assert len(body_sizes) == 1
|
||||
assert items == []
|
||||
|
||||
|
||||
def test_frame_that_cannot_fit_sends_nothing():
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(3)
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "png",
|
||||
client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
|
||||
|
||||
assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
|
||||
assert body_sizes == []
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
|
||||
def test_worker_to_master_over_real_http(fmt):
|
||||
# Import here, not at module level: other test harnesses only install their
|
||||
# aiohttp stubs when the real package has not been imported yet.
|
||||
aiohttp = pytest.importorskip("aiohttp")
|
||||
from aiohttp import web
|
||||
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = torch.rand(5, 24, 36, 3)
|
||||
|
||||
async def scenario():
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000 # several requests for still formats
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
await node.send_batch_to_master(
|
||||
images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items
|
||||
|
||||
items = asyncio.run(scenario())
|
||||
|
||||
assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
|
||||
assert [item["is_last"] for item in items] == [False, False, False, False, True]
|
||||
expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
|
||||
got = torch.cat([item["tensor"] for item in items]).numpy()
|
||||
assert got.shape == expected.shape
|
||||
if fmt in ("png", "ffv1"):
|
||||
assert np.array_equal(got, expected)
|
||||
else:
|
||||
assert np.abs(got - expected).mean() < 0.25 # random noise: only sanity-check lossy codecs
|
||||
@@ -1,304 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, json_payload=None, post_payload=None, headers=None, query=None):
|
||||
self._json_payload = json_payload
|
||||
self._post_payload = post_payload or {}
|
||||
self.headers = headers or {}
|
||||
self.query = query or {}
|
||||
|
||||
async def json(self):
|
||||
return self._json_payload
|
||||
|
||||
async def post(self):
|
||||
return self._post_payload
|
||||
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
def _load_usdu_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
|
||||
package_name = "dist_api_usdu_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.MAX_PAYLOAD_SIZE = 1024
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class BaseJobState:
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
|
||||
job_models_module.BaseJobState = BaseJobState
|
||||
job_models_module.TileJobState = TileJobState
|
||||
job_models_module.ImageJobState = ImageJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
|
||||
parsers_module._parse_tiles_from_form = lambda _data: []
|
||||
sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.usdu_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._TileJobState = TileJobState
|
||||
module._ImageJobState = ImageJobState
|
||||
return module
|
||||
|
||||
|
||||
usdu_routes = _load_usdu_routes_module()
|
||||
|
||||
|
||||
class _UploadField:
|
||||
def __init__(self, data):
|
||||
self.file = io.BytesIO(data)
|
||||
|
||||
|
||||
def _tiny_png_bytes():
|
||||
image = Image.new("RGB", (1, 1), (255, 0, 0))
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
usdu_routes._prompt_server_holder["value"] = types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
|
||||
async def test_heartbeat_updates_worker_status(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-1")
|
||||
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-1"})
|
||||
response = await usdu_routes.heartbeat_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertIn("worker-a", job_data.worker_status)
|
||||
|
||||
async def test_heartbeat_missing_fields_returns_400(self):
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a"})
|
||||
response = await usdu_routes.heartbeat_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("missing", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_request_image_dynamic_assigns_next_image(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._ImageJobState("job-2")
|
||||
await job_data.pending_images.put(7)
|
||||
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-2"})
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("image_idx"), 7)
|
||||
self.assertEqual(response.payload.get("estimated_remaining"), 0)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
|
||||
self.assertIn("worker-a", job_data.worker_status)
|
||||
|
||||
async def test_request_image_static_assigns_tile_and_batched_flag(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-3")
|
||||
job_data.batched_static = True
|
||||
await job_data.pending_tasks.put(4)
|
||||
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-3"})
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("tile_idx"), 4)
|
||||
self.assertTrue(response.payload.get("batched_static"))
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
|
||||
|
||||
async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-4")
|
||||
prompt_server.distributed_pending_tile_jobs["job-4"] = job_data
|
||||
|
||||
request = _FakeRequest(
|
||||
post_payload={
|
||||
"multi_job_id": "job-4",
|
||||
"worker_id": "worker-a",
|
||||
"batch_size": "0",
|
||||
"is_last": "true",
|
||||
},
|
||||
headers={"content-length": "128"},
|
||||
)
|
||||
|
||||
response = await usdu_routes.submit_tiles_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await job_data.queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-a")
|
||||
self.assertTrue(queued["is_last"])
|
||||
self.assertEqual(queued["tiles"], [])
|
||||
|
||||
async def test_submit_image_enqueues_processed_image_payload(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._ImageJobState("job-5")
|
||||
prompt_server.distributed_pending_tile_jobs["job-5"] = job_data
|
||||
|
||||
request = _FakeRequest(
|
||||
post_payload={
|
||||
"multi_job_id": "job-5",
|
||||
"worker_id": "worker-a",
|
||||
"image_idx": "2",
|
||||
"full_image": _UploadField(_tiny_png_bytes()),
|
||||
"is_last": "false",
|
||||
},
|
||||
headers={"content-length": "256"},
|
||||
)
|
||||
|
||||
response = await usdu_routes.submit_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await job_data.queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-a")
|
||||
self.assertEqual(queued["image_idx"], 2)
|
||||
self.assertIn("image", queued)
|
||||
|
||||
async def test_job_status_endpoint_reports_ready(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
prompt_server.distributed_pending_tile_jobs["job-6"] = usdu_routes._TileJobState("job-6")
|
||||
|
||||
request = _FakeRequest(query={"multi_job_id": "job-6"})
|
||||
response = await usdu_routes.job_status_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload.get("ready"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,435 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None, match_info=None, query=None):
|
||||
self._payload = payload
|
||||
self.match_info = match_info or {}
|
||||
self.query = query or {}
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
class _FakeHTTPClientResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self._payload = payload
|
||||
self.status = status
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
async def text(self):
|
||||
return str(self._payload)
|
||||
|
||||
|
||||
class _FakeHTTPClientSession:
|
||||
def __init__(self, payload, status=200):
|
||||
self._payload = payload
|
||||
self._status = status
|
||||
self.calls = []
|
||||
|
||||
def get(self, url, params=None, timeout=None):
|
||||
self.calls.append({"url": url, "params": params, "timeout": timeout})
|
||||
return _FakeHTTPClientResponse(self._payload, status=self._status)
|
||||
|
||||
|
||||
class _DummyWorkerManager:
|
||||
def __init__(self):
|
||||
self.processes = {}
|
||||
|
||||
def launch_worker(self, worker):
|
||||
worker_id = str(worker["id"])
|
||||
self.processes[worker_id] = {
|
||||
"pid": 12345,
|
||||
"log_file": f"/tmp/distributed_worker_{worker_id}.log",
|
||||
"process": None,
|
||||
}
|
||||
return 12345
|
||||
|
||||
def _is_process_running(self, _pid):
|
||||
return False
|
||||
|
||||
def save_processes(self):
|
||||
return None
|
||||
|
||||
def stop_worker(self, _worker_id):
|
||||
return True, "Stopped"
|
||||
|
||||
def get_managed_workers(self):
|
||||
return []
|
||||
|
||||
|
||||
class _ImmediateLoop:
|
||||
async def run_in_executor(self, _executor, func, *args):
|
||||
return func(*args)
|
||||
|
||||
|
||||
def _load_worker_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
|
||||
package_name = "dist_api_worker_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
|
||||
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
detection_module = types.ModuleType(f"{package_name}.workers.detection")
|
||||
detection_module.is_local_worker = lambda *_args, **_kwargs: True
|
||||
detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
|
||||
detection_module.get_machine_id = lambda: "machine-id"
|
||||
detection_module.is_docker_environment = lambda: False
|
||||
detection_module.is_runpod_environment = lambda: False
|
||||
detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
|
||||
sys.modules[f"{package_name}.workers.detection"] = detection_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _WebSocketResponse:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def prepare(self, _request):
|
||||
return None
|
||||
|
||||
async def send_json(self, _payload):
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
async def _empty():
|
||||
if False:
|
||||
yield None
|
||||
return _empty()
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
|
||||
WebSocketResponse=_WebSocketResponse,
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
current_device=lambda: 0,
|
||||
device_count=lambda: 0,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"
|
||||
|
||||
async def _probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
|
||||
async def _get_client_session():
|
||||
raise RuntimeError("not used in these tests")
|
||||
|
||||
network_module.get_client_session = _get_client_session
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.CHUNK_SIZE = 8192
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
async def _queue_prompt_payload(*_args, **_kwargs):
|
||||
return "prompt-id"
|
||||
|
||||
async_helpers_module.queue_prompt_payload = _queue_prompt_payload
|
||||
|
||||
class _PromptValidationError(RuntimeError):
|
||||
def __init__(self, message="invalid prompt", validation_error=None, node_errors=None):
|
||||
super().__init__(message)
|
||||
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
async_helpers_module.PromptValidationError = _PromptValidationError
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
|
||||
def _require_fields(data, *fields):
|
||||
missing = []
|
||||
for field in fields:
|
||||
value = data.get(field) if isinstance(data, dict) else None
|
||||
if value is None or (isinstance(value, str) and not value.strip()):
|
||||
missing.append(field)
|
||||
return missing
|
||||
|
||||
def _validate_worker_id(worker_id, config):
|
||||
return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
|
||||
|
||||
schemas_module.require_fields = _require_fields
|
||||
schemas_module.validate_worker_id = _validate_worker_id
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_network_info_reports_actual_master_and_available_ports(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": []}), \
|
||||
patch.object(worker_routes, "get_server_port", return_value=8189), \
|
||||
patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["master_port"], 8189)
|
||||
self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
|
||||
self.assertEqual(response.payload["cuda_device_count"], 4)
|
||||
allocate.assert_called_once_with(8189, [], 3)
|
||||
|
||||
async def test_network_info_does_not_reallocate_existing_configuration(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=[]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
|
||||
patch.object(worker_routes, "allocate_worker_ports") as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["local_worker_ports"], [])
|
||||
allocate.assert_not_called()
|
||||
|
||||
async def test_launch_conflict_is_a_clear_client_error(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
|
||||
patch.object(worker_routes, "load_config", return_value=config), \
|
||||
patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
|
||||
response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("master port 8189", response.payload["message"])
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
async def test_launch_worker_valid_id_returns_200(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
||||
request = _FakeRequest({"worker_id": "worker-a"})
|
||||
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes, "load_config", return_value=config
|
||||
), patch.object(
|
||||
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
||||
):
|
||||
response = await worker_routes.launch_worker_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertEqual(response.payload.get("pid"), 12345)
|
||||
|
||||
async def test_launch_worker_unknown_id_returns_404(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
||||
request = _FakeRequest({"worker_id": "missing-worker"})
|
||||
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes, "load_config", return_value=config
|
||||
):
|
||||
response = await worker_routes.launch_worker_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 404)
|
||||
self.assertIn("not found", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_worker_log_returns_content_json(self):
|
||||
manager = _DummyWorkerManager()
|
||||
with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
|
||||
handle.write("line-1\nline-2\nline-3\n")
|
||||
log_path = handle.name
|
||||
|
||||
manager.processes["worker-a"] = {
|
||||
"pid": 9999,
|
||||
"log_file": log_path,
|
||||
"process": None,
|
||||
}
|
||||
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
|
||||
try:
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
||||
):
|
||||
response = await worker_routes.get_worker_log_endpoint(request)
|
||||
finally:
|
||||
if os.path.exists(log_path):
|
||||
os.remove(log_path)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertIn("content", response.payload)
|
||||
self.assertIn("line-3", response.payload["content"])
|
||||
|
||||
async def test_local_log_reads_memory_buffer(self):
|
||||
request = _FakeRequest(query={"lines": "2"})
|
||||
fake_logs = deque(
|
||||
[
|
||||
{"m": "line-1\n"},
|
||||
{"m": "line-2\n"},
|
||||
{"m": "line-3\n"},
|
||||
],
|
||||
maxlen=300,
|
||||
)
|
||||
app_module = types.ModuleType("app")
|
||||
app_module.__path__ = []
|
||||
logger_module = types.ModuleType("app.logger")
|
||||
logger_module.get_logs = lambda: fake_logs
|
||||
app_module.logger = logger_module
|
||||
|
||||
with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}):
|
||||
response = await worker_routes.get_local_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertEqual(response.payload.get("source"), "memory")
|
||||
self.assertEqual(response.payload.get("entries"), 2)
|
||||
self.assertIn("line-3", response.payload.get("content", ""))
|
||||
|
||||
async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self):
|
||||
config = {
|
||||
"workers": [
|
||||
{
|
||||
"id": "worker-remote",
|
||||
"name": "Remote Worker",
|
||||
"host": "worker.example.com",
|
||||
"port": 8188,
|
||||
"type": "remote",
|
||||
}
|
||||
]
|
||||
}
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"})
|
||||
proxied_payload = {
|
||||
"status": "success",
|
||||
"content": "remote-log-content\n",
|
||||
"entries": 1,
|
||||
"source": "memory",
|
||||
"truncated": False,
|
||||
"lines_shown": 1,
|
||||
}
|
||||
fake_session = _FakeHTTPClientSession(proxied_payload)
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return fake_session
|
||||
|
||||
with patch.object(worker_routes, "load_config", return_value=config), patch.object(
|
||||
worker_routes, "get_client_session", side_effect=_fake_get_client_session
|
||||
):
|
||||
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("content"), "remote-log-content\n")
|
||||
self.assertEqual(len(fake_session.calls), 1)
|
||||
self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"})
|
||||
self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log"))
|
||||
|
||||
async def test_remote_worker_log_rejects_local_workers(self):
|
||||
config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]}
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-local"})
|
||||
|
||||
with patch.object(worker_routes, "load_config", return_value=config):
|
||||
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("local", response.payload.get("message", "").lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,2 +0,0 @@
|
||||
# This conftest.py marks the tests/ directory as the pytest collection root,
|
||||
# preventing pytest from traversing into the parent package's __init__.py.
|
||||
-691
@@ -1,691 +0,0 @@
|
||||
{
|
||||
"DistributedCollector": {
|
||||
"input": {
|
||||
"required": {
|
||||
"load_balance": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
|
||||
}
|
||||
]
|
||||
},
|
||||
"optional": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_batch_size": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1024
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"pass_through": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"delegate_only": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"load_balance"
|
||||
],
|
||||
"optional": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_batch_size",
|
||||
"worker_id",
|
||||
"pass_through",
|
||||
"delegate_only"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE",
|
||||
"AUDIO"
|
||||
],
|
||||
"output_name": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": true,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Collector"
|
||||
},
|
||||
"DistributedSeed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"seed"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"INT"
|
||||
],
|
||||
"output_name": [
|
||||
"seed"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Seed"
|
||||
},
|
||||
"DistributedModelName": {
|
||||
"input": {
|
||||
"required": {
|
||||
"text": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"text"
|
||||
],
|
||||
"hidden": [
|
||||
"unique_id",
|
||||
"extra_pnginfo"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"output"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Model Name"
|
||||
},
|
||||
"DistributedValue": {
|
||||
"input": {
|
||||
"required": {
|
||||
"default_value": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"worker_values": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "{}"
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"default_value",
|
||||
"worker_values"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"value"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Value"
|
||||
},
|
||||
"ImageBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"images",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"batch_1",
|
||||
"batch_2",
|
||||
"batch_3",
|
||||
"batch_4",
|
||||
"batch_5",
|
||||
"batch_6",
|
||||
"batch_7",
|
||||
"batch_8",
|
||||
"batch_9",
|
||||
"batch_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "image",
|
||||
"display_name": "Image Batch Divider"
|
||||
},
|
||||
"AudioBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"audio",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"audio_1",
|
||||
"audio_2",
|
||||
"audio_3",
|
||||
"audio_4",
|
||||
"audio_5",
|
||||
"audio_6",
|
||||
"audio_7",
|
||||
"audio_8",
|
||||
"audio_9",
|
||||
"audio_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "audio",
|
||||
"display_name": "Audio Segment Divider"
|
||||
},
|
||||
"DistributedEmptyImage": {
|
||||
"input": {
|
||||
"required": {
|
||||
"height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"channels": [
|
||||
"INT",
|
||||
{
|
||||
"default": 3,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"height",
|
||||
"width",
|
||||
"channels"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Empty Image"
|
||||
},
|
||||
"UltimateSDUpscaleDistributed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"upscaled_image": [
|
||||
"IMAGE"
|
||||
],
|
||||
"model": [
|
||||
"MODEL"
|
||||
],
|
||||
"positive": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"negative": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"vae": [
|
||||
"VAE"
|
||||
],
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 18446744073709551615
|
||||
}
|
||||
],
|
||||
"steps": [
|
||||
"INT",
|
||||
{
|
||||
"default": 20,
|
||||
"min": 1,
|
||||
"max": 10000
|
||||
}
|
||||
],
|
||||
"cfg": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0
|
||||
}
|
||||
],
|
||||
"sampler_name": [
|
||||
[
|
||||
"euler",
|
||||
"euler_cfg_pp",
|
||||
"euler_ancestral",
|
||||
"euler_ancestral_cfg_pp",
|
||||
"heun",
|
||||
"heunpp2",
|
||||
"exp_heun_2_x0",
|
||||
"exp_heun_2_x0_sde",
|
||||
"dpm_2",
|
||||
"dpm_2_ancestral",
|
||||
"lms",
|
||||
"dpm_fast",
|
||||
"dpm_adaptive",
|
||||
"dpmpp_2s_ancestral",
|
||||
"dpmpp_2s_ancestral_cfg_pp",
|
||||
"dpmpp_sde",
|
||||
"dpmpp_sde_gpu",
|
||||
"dpmpp_2m",
|
||||
"dpmpp_2m_cfg_pp",
|
||||
"dpmpp_2m_sde",
|
||||
"dpmpp_2m_sde_gpu",
|
||||
"dpmpp_2m_sde_heun",
|
||||
"dpmpp_2m_sde_heun_gpu",
|
||||
"dpmpp_3m_sde",
|
||||
"dpmpp_3m_sde_gpu",
|
||||
"ddpm",
|
||||
"lcm",
|
||||
"ipndm",
|
||||
"ipndm_v",
|
||||
"deis",
|
||||
"cfgpp_ud10_ab",
|
||||
"res_multistep",
|
||||
"res_multistep_cfg_pp",
|
||||
"res_multistep_ancestral",
|
||||
"res_multistep_ancestral_cfg_pp",
|
||||
"gradient_estimation",
|
||||
"gradient_estimation_cfg_pp",
|
||||
"er_sde",
|
||||
"seeds_2",
|
||||
"seeds_3",
|
||||
"sa_solver",
|
||||
"sa_solver_pece",
|
||||
"ddim",
|
||||
"uni_pc",
|
||||
"uni_pc_bh2"
|
||||
]
|
||||
],
|
||||
"scheduler": [
|
||||
[
|
||||
"simple",
|
||||
"sgm_uniform",
|
||||
"karras",
|
||||
"exponential",
|
||||
"ddim_uniform",
|
||||
"beta",
|
||||
"normal",
|
||||
"linear_quadratic",
|
||||
"kl_optimal"
|
||||
]
|
||||
],
|
||||
"denoise": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}
|
||||
],
|
||||
"tile_width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"tile_height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"padding": [
|
||||
"INT",
|
||||
{
|
||||
"default": 32,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"mask_blur": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 0,
|
||||
"max": 256
|
||||
}
|
||||
],
|
||||
"force_uniform_tiles": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": true
|
||||
}
|
||||
],
|
||||
"tiled_decode": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"tile_indices": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"dynamic_threshold": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 64
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"upscaled_image",
|
||||
"model",
|
||||
"positive",
|
||||
"negative",
|
||||
"vae",
|
||||
"seed",
|
||||
"steps",
|
||||
"cfg",
|
||||
"sampler_name",
|
||||
"scheduler",
|
||||
"denoise",
|
||||
"tile_width",
|
||||
"tile_height",
|
||||
"padding",
|
||||
"mask_blur",
|
||||
"force_uniform_tiles",
|
||||
"tiled_decode"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_id",
|
||||
"tile_indices",
|
||||
"dynamic_threshold"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image/upscaling",
|
||||
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
|
||||
}
|
||||
}
|
||||
@@ -1,289 +0,0 @@
|
||||
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
|
||||
|
||||
No HTTP listener, workers, model downloads or live installation changes.
|
||||
The original contracts were captured from 32ac027 using the same core.
|
||||
"""
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
|
||||
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
|
||||
|
||||
# Inputs intentionally added after the V1 baseline was captured. They must be
|
||||
# optional, appended after the baseline inputs, and keep legacy behaviour as
|
||||
# their default so saved workflows are unaffected.
|
||||
ADDED_OPTIONAL_INPUTS = {
|
||||
'DistributedCollector': {
|
||||
'transfer_format': [
|
||||
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
|
||||
{'default': 'legacy_png'},
|
||||
],
|
||||
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
kind = value[0]
|
||||
opts = dict(value[1]) if len(value) > 1 else {}
|
||||
if kind == 'STRING':
|
||||
# V3 serializes the same default single-line widget explicitly.
|
||||
opts.setdefault('multiline', False)
|
||||
if kind == 'COMBO':
|
||||
kind = opts['options']
|
||||
opts = {key: val for key, val in opts.items() if key != 'options'}
|
||||
if isinstance(kind, list):
|
||||
# Single-selection is the V1 dropdown default as well.
|
||||
opts.setdefault('multiselect', False)
|
||||
return [kind, opts]
|
||||
|
||||
|
||||
async def map_node(cls, values, extra=None):
|
||||
"""Use the actual input parser and sanitized V3 executor clones."""
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
|
||||
assert not missing, missing
|
||||
# Use a worker thread for synchronous nodes: collector/upscale bridge back
|
||||
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
|
||||
def run():
|
||||
return asyncio.run(execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
|
||||
return await asyncio.to_thread(run)
|
||||
|
||||
|
||||
def check_saved_workflows(mapping):
|
||||
checked = 0
|
||||
occurrences = 0
|
||||
for path in sorted((ROOT / 'workflows').glob('*.json')):
|
||||
workflow = json.loads(path.read_text())
|
||||
local_nodes = {node['id']: node for node in workflow['nodes']}
|
||||
for saved in workflow['nodes']:
|
||||
if saved['type'] not in mapping:
|
||||
continue
|
||||
occurrences += 1
|
||||
info = mapping[saved['type']].GET_NODE_INFO_V1()
|
||||
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
|
||||
for socket in saved.get('inputs', []):
|
||||
assert socket['name'] in inputs, (path.name, saved['id'], socket)
|
||||
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
|
||||
for index, output in enumerate(saved.get('outputs', [])):
|
||||
assert output['type'] == info['output'][index], (path.name, index)
|
||||
for link in workflow['links']:
|
||||
if link[1] == saved['id']:
|
||||
assert link[2] < len(info['output']), (path.name, link)
|
||||
target = local_nodes[link[3]]['inputs'][link[4]]
|
||||
assert target['type'] == info['output'][link[2]], (path.name, link)
|
||||
checked += 1
|
||||
assert checked == 5 and occurrences == 10, (checked, occurrences)
|
||||
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
|
||||
|
||||
|
||||
async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
result = await map_node(mapping['DistributedSeed'],
|
||||
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (126,)
|
||||
values = json.dumps({'_type': 'INT', '3': '17'})
|
||||
result = await map_node(mapping['DistributedValue'],
|
||||
{'default_value': '4', 'worker_values': values,
|
||||
'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (17,)
|
||||
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
|
||||
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
|
||||
{'extra_pnginfo': metadata})
|
||||
assert result[0].result == ('model.ckpt',)
|
||||
assert result[0].ui == {'text': ['model.ckpt']}
|
||||
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
|
||||
result = await map_node(mapping['DistributedEmptyImage'],
|
||||
{'height': 8, 'width': 8, 'channels': 3})
|
||||
empty = result[0].result[0]
|
||||
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
|
||||
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
|
||||
for count in (1, 3, 10):
|
||||
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat(result[0].result[:count]), images)
|
||||
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
|
||||
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
|
||||
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
|
||||
collector = mapping['DistributedCollector']
|
||||
prepared, missing, hidden = execution.get_input_data(
|
||||
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
|
||||
assert not missing
|
||||
prepared['images'] = [images[:2], images[2:]]
|
||||
prepared['audio'] = [audio, audio]
|
||||
result = await execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
|
||||
assert torch.equal(result[0].result[0], images)
|
||||
assert result[0].result[1]['waveform'].shape[-1] == 66
|
||||
# Exercise real async aggregation without an HTTP endpoint or worker.
|
||||
queue = asyncio.Queue()
|
||||
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
|
||||
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
|
||||
result = await map_node(collector,
|
||||
{'images': images[:2], 'load_balance': False,
|
||||
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
|
||||
assert torch.equal(result[0].result[0], images[:3])
|
||||
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
|
||||
|
||||
# The registered node must forward the transfer options to the runtime
|
||||
# helper, and default to the legacy transfer when they are not set.
|
||||
collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
|
||||
forwarded = []
|
||||
def fake_run(self, **kwargs):
|
||||
forwarded.append(kwargs)
|
||||
return (kwargs['images'], kwargs['audio'])
|
||||
with patch.object(collector_runtime, 'run', fake_run):
|
||||
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
|
||||
'transfer_format': 'av1', 'transfer_quality': 70})
|
||||
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
|
||||
# Set values arrive list-wrapped (is_input_list); unset ones use the
|
||||
# execute() defaults. The runtime helper unwraps both forms.
|
||||
assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
|
||||
(['av1'], [70]), ('legacy_png', 90)], forwarded
|
||||
|
||||
# Prove V3's per-execution class clones do not leak helper instance state;
|
||||
# replace only the GPU/model boundary, not the input parser or executor.
|
||||
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
|
||||
seen = []
|
||||
def fake_upscale(self, *args):
|
||||
assert not hasattr(self, 'acceptance_marker')
|
||||
self.acceptance_marker = True
|
||||
seen.append(args)
|
||||
return (args[0],)
|
||||
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
|
||||
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
|
||||
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
|
||||
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
|
||||
'force_uniform_tiles': True, 'tiled_decode': False,
|
||||
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
|
||||
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
|
||||
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
|
||||
with patch.object(upscale_runtime, 'run', fake_upscale):
|
||||
for _ in range(2):
|
||||
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
|
||||
assert result[0].result[0] is images
|
||||
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
|
||||
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
|
||||
'tile_indices', 'dynamic_threshold')), seen[0]
|
||||
import math
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
||||
'pass_through': True}},
|
||||
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
|
||||
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('v3-graph', graph, None)
|
||||
assert valid[0], valid
|
||||
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
|
||||
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
|
||||
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
|
||||
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
|
||||
assert not rejected[0]
|
||||
errors = [error for entry in rejected[3].values() for error in entry['errors']]
|
||||
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
|
||||
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
history = executor.history_result
|
||||
record = history['outputs']['4']['images'][0]
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
|
||||
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
|
||||
|
||||
|
||||
async def main():
|
||||
asset_manager = default_asset_manager()
|
||||
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
|
||||
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) == 8
|
||||
assert set(mapping) == set(BASELINE)
|
||||
for node_id, cls in mapping.items():
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
|
||||
old = dict(BASELINE[node_id])
|
||||
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
|
||||
# V1's ByPassTypeTuple advertises '*' when indexed, while its
|
||||
# underlying tuple and existing frontend declare IMAGE/AUDIO.
|
||||
# Native V3 declares all ten existing typed sockets explicitly.
|
||||
assert old['output'] == ['*'] * 10
|
||||
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
|
||||
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
|
||||
for group in ('required', 'optional'):
|
||||
old_inputs = old['input'].get(group, {})
|
||||
new_inputs = dict(new['input'].get(group, {}))
|
||||
if group == 'optional':
|
||||
for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
|
||||
assert name in new_inputs, (node_id, name, 'added input missing')
|
||||
actual = normalized_input(new_inputs.pop(name))
|
||||
kind, opts = normalized_input(expected)
|
||||
assert actual[0] == kind, (node_id, name, actual)
|
||||
for key, value in opts.items():
|
||||
assert actual[1].get(key) == value, (node_id, name, key, actual)
|
||||
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
|
||||
for name, original in old_inputs.items():
|
||||
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
|
||||
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
|
||||
assert old[key] == new[key], (node_id, key, old[key], new[key])
|
||||
# Standard context lives in cls.hidden; orchestrator metadata remains
|
||||
# accepted by its original kwarg name, without creating new widgets.
|
||||
signature = inspect.signature(cls.execute)
|
||||
for name, field in old['input'].get('hidden', {}).items():
|
||||
if isinstance(field, list):
|
||||
assert cls.GET_SCHEMA().accept_all_inputs, node_id
|
||||
assert name in signature.parameters, (node_id, name)
|
||||
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
|
||||
else:
|
||||
assert name in new['input']['hidden'], (node_id, name)
|
||||
expected_hidden = []
|
||||
if node_id == 'DistributedModelName':
|
||||
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
print('SCHEMA_PARITY_OK', len(mapping))
|
||||
check_saved_workflows(mapping)
|
||||
await check_execution(mapping, module, prompt_server, asset_manager)
|
||||
print('V3_ACCEPTANCE_OK')
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -1,91 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _load_async_helpers_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
|
||||
package_name = "dist_async_helpers_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
execution_module = types.ModuleType("execution")
|
||||
|
||||
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
|
||||
return (True, None, ["9"], {})
|
||||
|
||||
execution_module.validate_prompt = _validate_prompt
|
||||
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
|
||||
sys.modules["execution"] = execution_module
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
trigger_on_prompt=lambda payload: payload,
|
||||
number=12,
|
||||
prompt_queue=_PromptQueue(),
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_loop = lambda: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module, prompt_server
|
||||
|
||||
|
||||
async_helpers, prompt_server = _load_async_helpers_module()
|
||||
|
||||
|
||||
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
|
||||
result = await async_helpers.queue_prompt_payload(
|
||||
{"1": {"class_type": "Node"}},
|
||||
workflow_meta={"id": "workflow-1"},
|
||||
client_id="client-1",
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(result["prompt_id"], str)
|
||||
self.assertTrue(result["prompt_id"])
|
||||
self.assertEqual(result["number"], 12)
|
||||
self.assertEqual(result["node_errors"], {})
|
||||
|
||||
self.assertEqual(prompt_server.number, 13)
|
||||
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
|
||||
queued_item = prompt_server.prompt_queue.items[0]
|
||||
self.assertEqual(queued_item[0], 12)
|
||||
extra_data = queued_item[3]
|
||||
self.assertEqual(extra_data["client_id"], "client-1")
|
||||
self.assertIn("create_time", extra_data)
|
||||
self.assertIsInstance(extra_data["create_time"], int)
|
||||
self.assertGreater(extra_data["create_time"], 0)
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,102 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_utilities_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
package_name = "dist_divider_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.nodes.utilities",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
utils = _load_utilities_module()
|
||||
|
||||
|
||||
class ImageBatchDividerTests(unittest.TestCase):
|
||||
def test_divides_images_into_contiguous_chunks(self):
|
||||
divider = utils.ImageBatchDivider()
|
||||
images = torch.arange(10, dtype=torch.float32).reshape(10, 1, 1, 1)
|
||||
|
||||
outputs = divider.divide_batch(images, 3)
|
||||
|
||||
self.assertEqual(outputs[0].shape[0], 4)
|
||||
self.assertEqual(outputs[1].shape[0], 3)
|
||||
self.assertEqual(outputs[2].shape[0], 3)
|
||||
self.assertEqual(outputs[0][:, 0, 0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
|
||||
self.assertEqual(outputs[1][:, 0, 0, 0].tolist(), [4.0, 5.0, 6.0])
|
||||
self.assertEqual(outputs[2][:, 0, 0, 0].tolist(), [7.0, 8.0, 9.0])
|
||||
|
||||
def test_unused_image_outputs_are_empty(self):
|
||||
divider = utils.ImageBatchDivider()
|
||||
images = torch.arange(4, dtype=torch.float32).reshape(4, 1, 1, 1)
|
||||
|
||||
outputs = divider.divide_batch(images, 2)
|
||||
|
||||
self.assertEqual(len(outputs), 10)
|
||||
for idx in range(2, 10):
|
||||
self.assertEqual(outputs[idx].shape[0], 0)
|
||||
|
||||
|
||||
class AudioBatchDividerTests(unittest.TestCase):
|
||||
def test_divides_audio_samples_into_contiguous_chunks(self):
|
||||
divider = utils.AudioBatchDivider()
|
||||
audio = {
|
||||
"waveform": torch.arange(10, dtype=torch.float32).reshape(1, 1, 10),
|
||||
"sample_rate": 24000,
|
||||
}
|
||||
|
||||
outputs = divider.divide_audio(audio, 3)
|
||||
|
||||
self.assertEqual(outputs[0]["waveform"][0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
|
||||
self.assertEqual(outputs[1]["waveform"][0, 0].tolist(), [4.0, 5.0, 6.0])
|
||||
self.assertEqual(outputs[2]["waveform"][0, 0].tolist(), [7.0, 8.0, 9.0])
|
||||
|
||||
def test_unused_audio_outputs_are_empty(self):
|
||||
divider = utils.AudioBatchDivider()
|
||||
audio = {
|
||||
"waveform": torch.arange(8, dtype=torch.float32).reshape(1, 1, 8),
|
||||
"sample_rate": 24000,
|
||||
}
|
||||
|
||||
outputs = divider.divide_audio(audio, 2)
|
||||
|
||||
self.assertEqual(len(outputs), 10)
|
||||
for idx in range(2, 10):
|
||||
self.assertEqual(outputs[idx]["waveform"].shape[-1], 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,630 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
||||
codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
||||
codec_module = importlib.util.module_from_spec(codec_spec)
|
||||
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
||||
codec_spec.loader.exec_module(codec_module)
|
||||
utils_pkg.transfer_codec = codec_module
|
||||
return codec_module
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
_register_real_transfer_codec(package_name, utils_pkg)
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
def get(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
|
||||
class InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_management.InterruptProcessingException = InterruptProcessingException
|
||||
model_management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy_module.model_management = model_management
|
||||
|
||||
comfy_utils = types.ModuleType("comfy.utils")
|
||||
|
||||
class ProgressBar:
|
||||
def __init__(self, _total):
|
||||
self.total = _total
|
||||
self.updates = []
|
||||
|
||||
def update(self, value):
|
||||
self.updates.append(value)
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_management
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 0.1
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
config_module.is_master_delegate_only = lambda: False
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 1.0
|
||||
constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
|
||||
constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
def _ensure_contiguous(tensor):
|
||||
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
|
||||
|
||||
image_module.ensure_contiguous = _ensure_contiguous
|
||||
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
|
||||
image_module.pil_to_tensor = lambda value: value
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: "http://worker"
|
||||
network_module.get_client_session = lambda: None
|
||||
network_module.probe_worker = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
audio_payload_module.encode_audio_payload = lambda audio: audio
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_collector_opts_into_comfyui_list_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode
|
||||
|
||||
assert collector.INPUT_IS_LIST is True
|
||||
|
||||
|
||||
def test_collector_exposes_images_as_optional_input():
|
||||
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
|
||||
|
||||
assert "images" not in input_types["required"]
|
||||
assert input_types["optional"]["images"] == ("IMAGE",)
|
||||
|
||||
|
||||
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
|
||||
images, returned_audio = collector.run(images=None, audio=[audio])
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
|
||||
|
||||
def test_collector_rejects_missing_images_and_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
try:
|
||||
collector.run(images=None, audio=None)
|
||||
except ValueError as exc:
|
||||
assert "image or audio" in str(exc).lower()
|
||||
else:
|
||||
raise AssertionError("Expected collector to reject a run with no media input")
|
||||
|
||||
|
||||
def test_delegate_only_master_allows_no_local_media_input():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
images, audio = collector.run(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=["delegate-audio-job"],
|
||||
delegate_only=[True],
|
||||
enabled_worker_ids=["[]"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
load_balance=[False],
|
||||
audio=[None],
|
||||
multi_job_id=[""],
|
||||
is_worker=[False],
|
||||
master_url=[""],
|
||||
enabled_worker_ids=["[]"],
|
||||
worker_batch_size=[1],
|
||||
worker_id=[""],
|
||||
pass_through=[False],
|
||||
delegate_only=[False],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert torch.equal(images[0:1], first)
|
||||
assert torch.equal(images[1:2], second)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
posted_payloads = []
|
||||
|
||||
class _FakeImage:
|
||||
def save(self, fp, format=None, compress_level=None):
|
||||
fp.write(b"png-bytes")
|
||||
|
||||
class _FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, timeout):
|
||||
posted_payloads.append(json)
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda _audio: None
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
load_balance=[False],
|
||||
audio=[None],
|
||||
multi_job_id=["job-list-1"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
enabled_worker_ids=["[]"],
|
||||
worker_batch_size=[1],
|
||||
worker_id=["worker-a"],
|
||||
pass_through=[False],
|
||||
delegate_only=[False],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
assert len(posted_payloads) == 2
|
||||
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
|
||||
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
||||
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
||||
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
||||
|
||||
|
||||
def test_audio_only_worker_sends_one_completion_without_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
posted = []
|
||||
|
||||
class _FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, timeout):
|
||||
posted.append((url, json, timeout.total))
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda value: {"encoded": value is audio}
|
||||
|
||||
images, returned_audio = collector.run(
|
||||
images=None,
|
||||
audio=[audio],
|
||||
multi_job_id=["audio-job"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
worker_id=["worker-a"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
assert len(posted) == 1
|
||||
assert posted[0][0] == "http://master/distributed/job_complete"
|
||||
assert posted[0][2] == 600
|
||||
assert posted[0][1] == {
|
||||
"job_id": "audio-job",
|
||||
"worker_id": "worker-a",
|
||||
"batch_idx": 0,
|
||||
"audio": {"encoded": True},
|
||||
"is_last": True,
|
||||
}
|
||||
|
||||
|
||||
def test_audio_only_master_combines_local_and_worker_audio():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=master_audio,
|
||||
multi_job_id="audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
|
||||
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
|
||||
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
|
||||
|
||||
|
||||
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id="delegate-audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Compressed worker -> master transfer
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class _RecordingResponse:
|
||||
def __init__(self, status=200, text=""):
|
||||
self.status = status
|
||||
self._text = text
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def text(self):
|
||||
return self._text
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status >= 400:
|
||||
raise RuntimeError(f"HTTP {self.status}")
|
||||
|
||||
|
||||
class _RecordingSession:
|
||||
def __init__(self, statuses=None, text="Invalid JSON payload"):
|
||||
self.calls = []
|
||||
self._statuses = list(statuses or [])
|
||||
self._text = text
|
||||
|
||||
def post(self, url, timeout, json=None, data=None):
|
||||
self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
|
||||
status = self._statuses.pop(0) if self._statuses else 200
|
||||
return _RecordingResponse(status=status, text=self._text)
|
||||
|
||||
|
||||
def _install_fake_transport(module, session):
|
||||
async def _fake_get_client_session():
|
||||
return session
|
||||
|
||||
module.get_client_session = _fake_get_client_session
|
||||
# Record the multipart content instead of building a real aiohttp.FormData.
|
||||
module.DistributedCollectorNode._build_multipart = staticmethod(
|
||||
lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
|
||||
)
|
||||
|
||||
|
||||
def _run_worker(collector, images, audio=None, **kwargs):
|
||||
return collector.run(
|
||||
images=[images],
|
||||
audio=[audio],
|
||||
multi_job_id=["job-c"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
worker_id=["worker-a"],
|
||||
**{key: [value] for key, value in kwargs.items()},
|
||||
)
|
||||
|
||||
|
||||
def test_collector_exposes_transfer_widgets_with_legacy_default():
|
||||
module = _load_collector_module()
|
||||
optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
|
||||
|
||||
formats, options = optional["transfer_format"]
|
||||
assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
|
||||
assert options["default"] == "legacy_png"
|
||||
assert optional["transfer_quality"][0] == "INT"
|
||||
|
||||
|
||||
def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
|
||||
images = torch.rand(6, 32, 48, 3)
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
images,
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="ffv1",
|
||||
transfer_quality=90,
|
||||
)
|
||||
|
||||
assert len(session.calls) == 1
|
||||
call = session.calls[0]
|
||||
assert call["url"] == "http://master/distributed/job_complete"
|
||||
assert call["json"] is None
|
||||
envelope = call["data"]["envelope"]
|
||||
assert envelope["format"] == "ffv1"
|
||||
assert (envelope["start_index"], envelope["count"]) == (0, 6)
|
||||
assert (envelope["width"], envelope["height"]) == (48, 32)
|
||||
assert envelope["is_last"] is True
|
||||
assert envelope["audio"] == {"encoded": True}
|
||||
assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
|
||||
assert len(call["data"]["parts"]) == 1
|
||||
|
||||
|
||||
def test_worker_marks_only_final_chunk_as_last():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 1 # force one frame per request
|
||||
|
||||
_run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
|
||||
|
||||
envelopes = [call["data"]["envelope"] for call in session.calls]
|
||||
assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
|
||||
(0, 1, False),
|
||||
(1, 1, False),
|
||||
(2, 1, True),
|
||||
]
|
||||
assert "audio" not in envelopes[-1]
|
||||
|
||||
|
||||
def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession(statuses=[400])
|
||||
_install_fake_transport(module, session)
|
||||
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
||||
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
||||
)
|
||||
|
||||
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
|
||||
|
||||
assert session.calls[0]["data"] is not None # compressed attempt
|
||||
legacy = [call["json"] for call in session.calls[1:]]
|
||||
assert [payload["batch_idx"] for payload in legacy] == [0, 1]
|
||||
assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
|
||||
assert [payload["is_last"] for payload in legacy] == [False, True]
|
||||
|
||||
|
||||
def test_worker_uses_legacy_path_for_unknown_transfer_format():
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
||||
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
||||
)
|
||||
|
||||
_run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
|
||||
|
||||
assert len(session.calls) == 1
|
||||
assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
def test_worker_falls_back_when_video_encoder_missing():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
codec = module.transfer_codec
|
||||
original = codec.video_codec_available
|
||||
codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
|
||||
try:
|
||||
_run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
|
||||
finally:
|
||||
codec.video_codec_available = original
|
||||
|
||||
assert session.calls[0]["data"]["envelope"]["format"] == "h264"
|
||||
|
||||
|
||||
def test_worker_does_not_fall_back_on_non_legacy_rejection():
|
||||
import pytest
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
|
||||
_install_fake_transport(module, session)
|
||||
|
||||
with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
|
||||
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
|
||||
|
||||
# One compressed attempt, no legacy resend of the batch.
|
||||
assert len(session.calls) == 1
|
||||
assert session.calls[0]["json"] is None
|
||||
|
||||
|
||||
def test_worker_sends_nothing_when_a_frame_cannot_fit():
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
|
||||
# Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
|
||||
images = torch.zeros(3, 128, 128, 3)
|
||||
images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
|
||||
|
||||
with pytest.raises(module.transfer_codec.FrameTooLargeError):
|
||||
_run_worker(collector, images, transfer_format="png")
|
||||
|
||||
assert session.calls == []
|
||||
|
||||
|
||||
def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
torch.zeros(2, 16, 16, 3),
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="png",
|
||||
)
|
||||
|
||||
multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
|
||||
assert [e["is_last"] for e in multipart] == [False]
|
||||
assert all("audio" not in e for e in multipart)
|
||||
final = session.calls[-1]["json"]
|
||||
assert final["is_last"] is True and final["audio"]["data"].startswith("x")
|
||||
assert "image" not in final
|
||||
|
||||
|
||||
def test_worker_attaches_audio_to_last_chunk_when_it_fits():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
torch.zeros(2, 16, 16, 3),
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="png",
|
||||
)
|
||||
|
||||
assert len(session.calls) == 1
|
||||
envelope = session.calls[0]["data"]["envelope"]
|
||||
assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}
|
||||
@@ -1,293 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_config_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
|
||||
package_name = "dist_cfg_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[f"{package_name}.constants"] = constants_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
config = _load_config_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _merge_with_defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MergeWithDefaultsTests(unittest.TestCase):
|
||||
def test_non_dict_input_returns_defaults(self):
|
||||
result = config._merge_with_defaults("not a dict", {"key": "default"})
|
||||
self.assertEqual(result, {"key": "default"})
|
||||
|
||||
def test_fills_missing_keys_with_defaults(self):
|
||||
result = config._merge_with_defaults({}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result, {"a": 1, "b": 2})
|
||||
|
||||
def test_loaded_value_overrides_default(self):
|
||||
result = config._merge_with_defaults({"a": 99}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result["a"], 99)
|
||||
self.assertEqual(result["b"], 2)
|
||||
|
||||
def test_nested_dict_merges_recursively(self):
|
||||
defaults = {"settings": {"debug": False, "count": 5}}
|
||||
loaded = {"settings": {"debug": True}}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertTrue(result["settings"]["debug"])
|
||||
self.assertEqual(result["settings"]["count"], 5)
|
||||
|
||||
def test_preserves_unknown_keys_for_forward_compatibility(self):
|
||||
result = config._merge_with_defaults({"extra_key": "extra"}, {"a": 1})
|
||||
self.assertEqual(result["extra_key"], "extra")
|
||||
|
||||
def test_none_loaded_value_overrides_default(self):
|
||||
"""Explicitly set None in config should override non-None default."""
|
||||
result = config._merge_with_defaults({"a": None}, {"a": "default"})
|
||||
self.assertIsNone(result["a"])
|
||||
|
||||
def test_non_dict_nested_loaded_value_replaces_dict_default(self):
|
||||
"""If loaded has a scalar where default has a dict, use the scalar."""
|
||||
defaults = {"settings": {"debug": False}}
|
||||
loaded = {"settings": "flat_string"}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertEqual(result["settings"], "flat_string")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# load_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LoadConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_returns_defaults_when_file_missing(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent/path/config.json"):
|
||||
cfg = config.load_config()
|
||||
defaults = config.get_default_config()
|
||||
self.assertEqual(cfg["settings"]["debug"], defaults["settings"]["debug"])
|
||||
self.assertIn("workers", cfg)
|
||||
|
||||
def test_loads_valid_json_file(self):
|
||||
data = {
|
||||
"workers": [{"id": "w1"}],
|
||||
"master": {"host": "test.host"},
|
||||
"settings": {},
|
||||
"tunnel": {},
|
||||
}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertEqual(cfg["master"]["host"], "test.host")
|
||||
self.assertEqual(len(cfg["workers"]), 1)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_merges_loaded_file_with_defaults(self):
|
||||
"""Loaded file with partial settings should be filled in from defaults."""
|
||||
data = {"master": {"host": "h"}, "workers": [], "settings": {"debug": True}, "tunnel": {}}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
# debug was set to True
|
||||
self.assertTrue(cfg["settings"]["debug"])
|
||||
# auto_launch_workers is a default key and should be present
|
||||
self.assertIn("auto_launch_workers", cfg["settings"])
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_falls_back_to_defaults_on_invalid_json(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
f.write("{invalid json{{")
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertIn("settings", cfg)
|
||||
self.assertIn("workers", cfg)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_second_call_returns_cached_object(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(config.get_default_config(), f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
cfg2 = config.load_config()
|
||||
self.assertIs(cfg1, cfg2)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_invalidate_cache_forces_reload(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "first"
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
config.invalidate_config_cache()
|
||||
data["master"]["host"] = "second"
|
||||
with open(tmp_path, "w", encoding="utf-8") as fh:
|
||||
json.dump(data, fh)
|
||||
cfg2 = config.load_config()
|
||||
self.assertEqual(cfg1["master"]["host"], "first")
|
||||
self.assertEqual(cfg2["master"]["host"], "second")
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# save_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SaveConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_saves_and_reloads_correctly(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "saved.host"
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
result = config.save_config(data)
|
||||
self.assertTrue(result)
|
||||
loaded = config.load_config()
|
||||
self.assertEqual(loaded["master"]["host"], "saved.host")
|
||||
|
||||
def test_returns_false_when_path_unwritable(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent_dir/config.json"):
|
||||
result = config.save_config({})
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_save_invalidates_cache(self):
|
||||
"""After saving, the cache should be cleared so next load re-reads."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
data = config.get_default_config()
|
||||
config.save_config(data)
|
||||
# Cache is now None; load_config should re-read
|
||||
self.assertIsNone(config._config_cache)
|
||||
|
||||
def test_written_file_is_valid_json(self):
|
||||
data = config.get_default_config()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
config.save_config(data)
|
||||
with open(tmp_path, encoding="utf-8") as fh:
|
||||
parsed = json.load(fh)
|
||||
self.assertEqual(parsed["master"], data["master"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_worker_timeout_seconds
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetWorkerTimeoutSecondsTests(unittest.TestCase):
|
||||
def test_returns_configured_value(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 120
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 120)
|
||||
|
||||
def test_clamps_zero_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 0
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_clamps_negative_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = -10
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_falls_back_to_provided_default_when_key_missing(self):
|
||||
cfg = config.get_default_config()
|
||||
# worker_timeout_seconds is not present in default config
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=45)
|
||||
self.assertEqual(result, 45)
|
||||
|
||||
def test_fallback_also_clamped_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=0)
|
||||
self.assertEqual(result, 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_master_delegate_only
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsMasterDelegateOnlyTests(unittest.TestCase):
|
||||
def test_returns_false_by_default(self):
|
||||
cfg = config.get_default_config()
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_true_when_enabled(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["master_delegate_only"] = True
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertTrue(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_false_on_exception(self):
|
||||
def _raise():
|
||||
raise RuntimeError("config exploded")
|
||||
|
||||
with patch.object(config, "load_config", side_effect=RuntimeError("boom")):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,196 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_detection_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
|
||||
package_name = "dist_det_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.normalize_host = lambda value: value
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in these tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.detection",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
detection = _load_detection_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_docker_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsDockerEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_dockerenv_file_exists(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=True), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_docker_container_env_var_is_set(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": "1"}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_platform_node_contains_docker(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-docker-host"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_false_when_none_of_the_signals_are_present(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
def test_docker_node_name_is_case_insensitive(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="My-Docker-Box"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_docker_env_var_empty_string_is_falsy(self):
|
||||
"""An empty DOCKER_CONTAINER env var should NOT trigger docker detection."""
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": ""}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_runpod_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsRunpodEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_runpod_pod_id_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": "pod-abc"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_runpod_api_key_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_API_KEY": "key-xyz"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_both_vars_are_set(self):
|
||||
with patch.dict(
|
||||
detection.os.environ,
|
||||
{"RUNPOD_POD_ID": "pod-abc", "RUNPOD_API_KEY": "key-xyz"},
|
||||
clear=True,
|
||||
):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_false_when_neither_var_is_set(self):
|
||||
with patch.dict(detection.os.environ, {}, clear=True):
|
||||
self.assertFalse(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_pod_id_is_empty_string(self):
|
||||
"""is not None check means even empty string counts as detected."""
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": ""}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_local_worker (synchronous paths only)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_true_for_localhost_host(self):
|
||||
result = await detection.is_local_worker({"host": "localhost", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_127_0_0_1(self):
|
||||
result = await detection.is_local_worker({"host": "127.0.0.1", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_0_0_0_0(self):
|
||||
result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_when_type_is_local(self):
|
||||
result = await detection.is_local_worker({"type": "local", "host": "remote.example.com"})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_false_for_remote_host(self):
|
||||
result = await detection.is_local_worker({"host": "remote.example.com", "port": 8188})
|
||||
self.assertFalse(result)
|
||||
|
||||
async def test_true_when_no_host_key(self):
|
||||
"""Missing host defaults to 'localhost'."""
|
||||
result = await detection.is_local_worker({"port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_machine_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetMachineIdTests(unittest.TestCase):
|
||||
def test_returns_a_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertIsInstance(result, str)
|
||||
|
||||
def test_returns_non_empty_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertTrue(len(result) > 0)
|
||||
|
||||
def test_stable_across_calls(self):
|
||||
r1 = detection.get_machine_id()
|
||||
r2 = detection.get_machine_id()
|
||||
self.assertEqual(r1, r2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,254 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
|
||||
|
||||
package_name = "dist_dispatch_testpkg"
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
|
||||
|
||||
async def _probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("get_client_session should be mocked in these tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _ClientConnectorError(Exception):
|
||||
pass
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _TCPConnector:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _ClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.ClientConnectorError = _ClientConnectorError
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.TCPConnector = _TCPConnector
|
||||
aiohttp_module.ClientSession = _ClientSession
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: {"payload": payload, "status": status}
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.dispatch",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
return module
|
||||
|
||||
|
||||
dispatch = _load_dispatch_module()
|
||||
|
||||
|
||||
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_select_active_workers_filters_offline(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
|
||||
async def fake_probe(worker):
|
||||
return worker["id"] != "w2"
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=3,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_disables_delegate_when_all_offline(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker):
|
||||
return False
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=True,
|
||||
probe_concurrency=1,
|
||||
)
|
||||
|
||||
self.assertEqual(active_workers, [])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
|
||||
|
||||
async def fake_http_probe(_worker):
|
||||
return False
|
||||
|
||||
async def fake_ws_probe(_worker):
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
|
||||
dispatch,
|
||||
"worker_ws_is_active",
|
||||
side_effect=fake_ws_probe,
|
||||
) as ws_probe:
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=True,
|
||||
delegate_master=False,
|
||||
probe_concurrency=4,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
|
||||
self.assertEqual(ws_probe.call_count, 2)
|
||||
self.assertEqual(http_probe.call_count, 0)
|
||||
|
||||
async def test_probe_concurrency_is_bounded(self):
|
||||
workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
|
||||
state = {"in_flight": 0, "max_in_flight": 0}
|
||||
|
||||
async def fake_probe(_worker):
|
||||
state["in_flight"] += 1
|
||||
state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
|
||||
await asyncio.sleep(0.01)
|
||||
state["in_flight"] -= 1
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=2,
|
||||
)
|
||||
|
||||
self.assertEqual(len(active_workers), len(workers))
|
||||
self.assertLessEqual(state["max_in_flight"], 2)
|
||||
self.assertGreaterEqual(state["max_in_flight"], 2)
|
||||
|
||||
async def test_select_least_busy_worker_round_robins_idle_workers(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 0, "w2": 0, "w3": 2}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
dispatch._least_busy_rr_index = 0
|
||||
selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
|
||||
self.assertEqual(selected1["id"], "w1")
|
||||
self.assertEqual(selected2["id"], "w2")
|
||||
self.assertEqual(selected3["id"], "w1")
|
||||
|
||||
async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 5, "w2": 2, "w3": 4}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
|
||||
|
||||
self.assertEqual(selected["id"], "w2")
|
||||
|
||||
async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker_url, timeout=3.0):
|
||||
return None
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
|
||||
|
||||
self.assertIsNone(selected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,208 +0,0 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
|
||||
class DistributedValueTests(unittest.TestCase):
|
||||
"""Unit tests for the DistributedValue node's distribute() method."""
|
||||
|
||||
def _make_node(self):
|
||||
# Import inline to avoid plugin-level imports
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
pkg_name = "dv_test_pkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == pkg_name or mod_name.startswith(f"{pkg_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
# Mock torch if not available
|
||||
if "torch" not in sys.modules:
|
||||
sys.modules["torch"] = MagicMock()
|
||||
|
||||
root_pkg = types.ModuleType(pkg_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[pkg_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{pkg_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{pkg_name}.utils"] = utils_pkg
|
||||
|
||||
logging_mod = types.ModuleType(f"{pkg_name}.utils.logging")
|
||||
logging_mod.debug_log = lambda *_a, **_k: None
|
||||
logging_mod.log = lambda *_a, **_k: None
|
||||
sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{pkg_name}.nodes.utilities", module_path
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod.DistributedValue()
|
||||
|
||||
def setUp(self):
|
||||
self.node = self._make_node()
|
||||
|
||||
def test_master_returns_default(self):
|
||||
result = self.node.distribute(
|
||||
default_value="model_a",
|
||||
worker_values="{}",
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, ("model_a",))
|
||||
|
||||
def test_master_coerces_default_int(self):
|
||||
values = json.dumps({"_type": "INT"})
|
||||
result = self.node.distribute(
|
||||
default_value="42",
|
||||
worker_values=values,
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, (42,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_master_coerces_default_float(self):
|
||||
values = json.dumps({"_type": "FLOAT"})
|
||||
result = self.node.distribute(
|
||||
default_value="2.5",
|
||||
worker_values=values,
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, (2.5,))
|
||||
self.assertIsInstance(result[0], float)
|
||||
|
||||
def test_worker_returns_specific_value(self):
|
||||
values = json.dumps({"1": "model_x", "2": "model_y"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("model_x",))
|
||||
|
||||
def test_worker_second_index(self):
|
||||
values = json.dumps({"1": "model_x", "2": "model_y"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_1",
|
||||
)
|
||||
self.assertEqual(result, ("model_y",))
|
||||
|
||||
def test_worker_falls_back_to_default_when_key_missing(self):
|
||||
values = json.dumps({"_type": "INT", "1": "3"})
|
||||
result = self.node.distribute(
|
||||
default_value="9",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_5",
|
||||
)
|
||||
self.assertEqual(result, (9,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_worker_falls_back_to_default_on_empty_value(self):
|
||||
values = json.dumps({"1": ""})
|
||||
result = self.node.distribute(
|
||||
default_value="fallback",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("fallback",))
|
||||
|
||||
def test_worker_falls_back_on_invalid_json(self):
|
||||
result = self.node.distribute(
|
||||
default_value="safe",
|
||||
worker_values="not-json",
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("safe",))
|
||||
|
||||
def test_worker_falls_back_on_invalid_worker_id(self):
|
||||
values = json.dumps({"1": "model_x"})
|
||||
result = self.node.distribute(
|
||||
default_value="safe",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="bad_id",
|
||||
)
|
||||
self.assertEqual(result, ("safe",))
|
||||
|
||||
def test_worker_id_as_direct_integer(self):
|
||||
values = json.dumps({"1": "model_x"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="0",
|
||||
)
|
||||
self.assertEqual(result, ("model_x",))
|
||||
|
||||
def test_type_int_coerces_value(self):
|
||||
values = json.dumps({"_type": "INT", "1": "42"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, (42,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_type_float_coerces_value(self):
|
||||
values = json.dumps({"_type": "FLOAT", "1": "3.14"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertAlmostEqual(result[0], 3.14)
|
||||
self.assertIsInstance(result[0], float)
|
||||
|
||||
def test_type_combo_stays_string(self):
|
||||
values = json.dumps({"_type": "COMBO", "1": "model_v2"})
|
||||
result = self.node.distribute(
|
||||
default_value="model_v1",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("model_v2",))
|
||||
self.assertIsInstance(result[0], str)
|
||||
|
||||
def test_type_string_default_stays_string(self):
|
||||
values = json.dumps({"1": "hello"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("hello",))
|
||||
self.assertIsInstance(result[0], str)
|
||||
|
||||
def test_int_coerce_handles_float_string(self):
|
||||
"""INT coercion of '3.7' should truncate to 3."""
|
||||
values = json.dumps({"_type": "INT", "1": "3.7"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, (3,))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,197 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import unittest
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_job_timeout_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
|
||||
package_name = "dist_job_timeout_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
config_holder = {"value": {"settings": {}, "workers": []}}
|
||||
probe_holder = {"fn": None}
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: config_holder["value"]
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 60
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: f"http://{worker.get('host', '127.0.0.1')}:{worker.get('port', 8188)}"
|
||||
|
||||
async def _probe_worker(url, timeout=2.0):
|
||||
fn = probe_holder["fn"]
|
||||
if fn is None:
|
||||
return None
|
||||
return await fn(url, timeout)
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class BaseJobState:
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
job_models_module.BaseJobState = BaseJobState
|
||||
job_models_module.ImageJobState = ImageJobState
|
||||
job_models_module.TileJobState = TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.job_timeout", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
module._config_holder = config_holder
|
||||
module._probe_holder = probe_holder
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._ImageJobState = ImageJobState
|
||||
module._TileJobState = TileJobState
|
||||
return module
|
||||
|
||||
|
||||
jt = _load_job_timeout_module()
|
||||
|
||||
|
||||
class JobTimeoutRequeueTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
jt._prompt_server_holder["value"] = types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
jt._config_holder["value"] = {
|
||||
"settings": {"worker_timeout_seconds": 5},
|
||||
"workers": [{"id": "worker-1", "host": "worker.local", "port": 8188}],
|
||||
}
|
||||
|
||||
async def test_requeues_only_incomplete_dynamic_tasks_for_timed_out_worker(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-1")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [0, 1]
|
||||
job_data.completed_images[1] = "done"
|
||||
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-1", total_tasks=2)
|
||||
|
||||
self.assertEqual(requeued, 1)
|
||||
self.assertEqual(await job_data.pending_images.get(), 0)
|
||||
self.assertNotIn("worker-1", job_data.worker_status)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-1"], [])
|
||||
|
||||
async def test_busy_probe_graces_worker_and_skips_requeue(self):
|
||||
async def _busy_probe(_url, _timeout):
|
||||
return {"exec_info": {"queue_remaining": 3}}
|
||||
|
||||
jt._probe_holder["fn"] = _busy_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-2")
|
||||
old_heartbeat = time.time() - 60.0
|
||||
job_data.worker_status["worker-1"] = old_heartbeat
|
||||
job_data.assigned_to_workers["worker-1"] = [0]
|
||||
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-2", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertIn("worker-1", job_data.worker_status)
|
||||
self.assertGreaterEqual(job_data.worker_status["worker-1"], old_heartbeat)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
async def test_completed_dynamic_task_is_not_requeued(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-3")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [7]
|
||||
job_data.completed_images[7] = "complete"
|
||||
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-3", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,137 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_network_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py"
|
||||
|
||||
package_name = "dist_utils_testpkg"
|
||||
package_module = types.ModuleType(package_name)
|
||||
package_module.__path__ = [] # mark as package
|
||||
sys.modules[package_name] = package_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
|
||||
)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
if "aiohttp" not in sys.modules:
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _TCPConnector:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _ClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.TCPConnector = _TCPConnector
|
||||
aiohttp_module.ClientSession = _ClientSession
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: {"payload": payload, "status": status}
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.network", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
network = _load_network_module()
|
||||
|
||||
|
||||
class NetworkHelpersTests(unittest.TestCase):
|
||||
def test_normalize_host_strips_protocol_and_path(self):
|
||||
self.assertEqual(network.normalize_host(" https://example.com/a/b "), "example.com")
|
||||
|
||||
def test_normalize_host_keeps_none(self):
|
||||
self.assertIsNone(network.normalize_host(None))
|
||||
|
||||
def test_build_worker_url_defaults_to_server_address(self):
|
||||
worker = {"id": "w1", "port": 8189}
|
||||
self.assertEqual(network.build_worker_url(worker, "/prompt"), "http://127.0.0.1:8189/prompt")
|
||||
|
||||
def test_build_worker_url_cloud_defaults_to_https(self):
|
||||
worker = {"id": "w2", "host": "foo.proxy.runpod.net", "port": 443}
|
||||
self.assertEqual(network.build_worker_url(worker), "https://foo.proxy.runpod.net")
|
||||
|
||||
def test_build_worker_url_keeps_explicit_scheme(self):
|
||||
worker = {"id": "w3", "host": "https://worker.example.com", "port": 1234}
|
||||
self.assertEqual(network.build_worker_url(worker, "/prompt"), "https://worker.example.com/prompt")
|
||||
|
||||
def test_build_master_url_uses_https_for_cloud_host(self):
|
||||
cfg = {"master": {"host": "demo.proxy.runpod.net"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"https://demo.proxy.runpod.net",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_scheme(self):
|
||||
cfg = {"master": {"host": "https://master.example.com/"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"https://master.example.com",
|
||||
)
|
||||
|
||||
def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
|
||||
cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8188",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_port_in_host(self):
|
||||
cfg = {"master": {"host": "192.168.68.56:8001"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
def test_build_master_url_falls_back_to_server_address(self):
|
||||
cfg = {"master": {"host": "", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8190",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_uses_loopback_for_local_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8001",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,196 +0,0 @@
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
PIL_AVAILABLE = True
|
||||
except ImportError:
|
||||
PIL_AVAILABLE = False
|
||||
|
||||
|
||||
def _load_payload_parsers_module():
|
||||
# payload_parsers.py has no relative imports; only stdlib + PIL
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "payload_parsers.py"
|
||||
spec = importlib.util.spec_from_file_location("upscale_payload_parsers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
if PIL_AVAILABLE:
|
||||
pp = _load_payload_parsers_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_png_bytes(width=64, height=64, color=(128, 64, 32)):
|
||||
"""Return raw PNG bytes for a solid-colour image."""
|
||||
img = PILImage.new("RGB", (width, height), color=color)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class _MockFileField:
|
||||
"""Minimal multipart file-field stub."""
|
||||
|
||||
class _MockFile:
|
||||
def __init__(self, data: bytes):
|
||||
self._buf = io.BytesIO(data)
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._buf.read()
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self.file = self._MockFile(data)
|
||||
|
||||
|
||||
def _make_form(n_tiles, *, padding=None, extra_meta=None, image_color=(128, 64, 32)):
|
||||
"""Build a minimal form-data dict with `n_tiles` tile entries."""
|
||||
image_bytes = _make_png_bytes(color=image_color)
|
||||
metadata = []
|
||||
for i in range(n_tiles):
|
||||
entry = {
|
||||
"tile_idx": i,
|
||||
"x": i * 64,
|
||||
"y": 0,
|
||||
"extracted_width": 64,
|
||||
"extracted_height": 64,
|
||||
}
|
||||
if extra_meta and i < len(extra_meta):
|
||||
entry.update(extra_meta[i])
|
||||
metadata.append(entry)
|
||||
|
||||
form = {"tiles_metadata": json.dumps(metadata)}
|
||||
if padding is not None:
|
||||
form["padding"] = str(padding)
|
||||
for i in range(n_tiles):
|
||||
form[f"tile_{i}"] = _MockFileField(image_bytes)
|
||||
return form
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@unittest.skipUnless(PIL_AVAILABLE, "PIL not installed")
|
||||
class ParseTilesFromFormTests(unittest.TestCase):
|
||||
|
||||
# --- happy paths ---
|
||||
|
||||
def test_single_tile_returns_image_and_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertEqual(len(tiles), 1)
|
||||
tile = tiles[0]
|
||||
self.assertIsInstance(tile["image"], PILImage.Image)
|
||||
self.assertEqual(tile["tile_idx"], 0)
|
||||
self.assertEqual(tile["x"], 0)
|
||||
self.assertEqual(tile["y"], 0)
|
||||
self.assertEqual(tile["extracted_width"], 64)
|
||||
self.assertEqual(tile["extracted_height"], 64)
|
||||
|
||||
def test_multiple_tiles_preserve_count_order_and_coordinates(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(len(tiles), 3)
|
||||
for i, tile in enumerate(tiles):
|
||||
self.assertEqual(tile["tile_idx"], i)
|
||||
self.assertEqual(tiles[1]["x"], 64)
|
||||
self.assertEqual(tiles[2]["x"], 128)
|
||||
|
||||
def test_padding_is_parsed_from_form(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
|
||||
self.assertEqual(tiles[0]["padding"], 16)
|
||||
|
||||
def test_default_padding_is_zero(self):
|
||||
form = _make_form(1)
|
||||
form.pop("padding", None)
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_invalid_padding_string_falls_back_to_zero(self):
|
||||
form = _make_form(1)
|
||||
form["padding"] = "not_a_number"
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_optional_batch_idx_included_when_present(self):
|
||||
extra = [{"batch_idx": 2}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["batch_idx"], 2)
|
||||
|
||||
def test_optional_global_idx_included_when_present(self):
|
||||
extra = [{"global_idx": 5}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["global_idx"], 5)
|
||||
|
||||
def test_batch_idx_and_global_idx_absent_when_not_in_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertNotIn("batch_idx", tiles[0])
|
||||
self.assertNotIn("global_idx", tiles[0])
|
||||
|
||||
# --- error cases ---
|
||||
|
||||
def test_missing_tiles_metadata_raises_value_error(self):
|
||||
with self.assertRaises(ValueError, msg="Missing tiles_metadata"):
|
||||
pp._parse_tiles_from_form({})
|
||||
|
||||
def test_invalid_json_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": "{not valid json}"}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_list_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": json.dumps({"not": "a list"})}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_missing_tile_file_field_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
# tile_0 intentionally omitted
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_tile_field_without_file_attr_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": "plain string without .file",
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_image_bytes_raises_value_error(self):
|
||||
class _BadFileField:
|
||||
class _BadFile:
|
||||
def read(self):
|
||||
return b"this is definitely not image data"
|
||||
file = _BadFile()
|
||||
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": _BadFileField(),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_invalid_metadata_value_type_raises_value_error(self):
|
||||
"""Non-integer metadata fields (x, y, etc.) should raise ValueError."""
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": "not_int", "y": 0}]),
|
||||
"tile_0": _MockFileField(_make_png_bytes()),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,151 +0,0 @@
|
||||
import importlib.util
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_queue_request_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py"
|
||||
spec = importlib.util.spec_from_file_location("queue_request", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
queue_request = _load_queue_request_module()
|
||||
parse_queue_request_payload = queue_request.parse_queue_request_payload
|
||||
|
||||
|
||||
class QueueRequestPayloadTests(unittest.TestCase):
|
||||
def _base_payload(self):
|
||||
return {
|
||||
"prompt": {"1": {"class_type": "Anything"}},
|
||||
"enabled_worker_ids": ["worker-1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
|
||||
def test_normalizes_enabled_worker_ids(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["enabled_worker_ids"] = ["a", 2, 3]
|
||||
payload_data["delegate_master"] = True
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
|
||||
self.assertTrue(payload.delegate_master)
|
||||
|
||||
def test_supports_legacy_workers_field(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}]
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
|
||||
|
||||
def test_supports_auto_prepare_prompt_fallback(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("prompt", None)
|
||||
payload_data["auto_prepare"] = True
|
||||
payload_data["workflow"] = {
|
||||
"prompt": {
|
||||
"10": {"class_type": "DistributedCollector"},
|
||||
}
|
||||
}
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertIn("10", payload.prompt)
|
||||
self.assertTrue(payload.auto_prepare)
|
||||
|
||||
def test_normalizes_trace_execution_id(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = " exec_123 "
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.trace_execution_id, "exec_123")
|
||||
|
||||
def test_blank_trace_execution_id_normalizes_to_none(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = " "
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertIsNone(payload.trace_execution_id)
|
||||
|
||||
def test_auto_prepare_defaults_true(self):
|
||||
payload = parse_queue_request_payload(self._base_payload())
|
||||
self.assertTrue(payload.auto_prepare)
|
||||
|
||||
def test_workers_field_must_be_list(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
payload_data["workers"] = "worker-a"
|
||||
with self.assertRaisesRegex(ValueError, "Field 'workers' must be a list"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_trace_execution_id_must_be_string(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = 123
|
||||
with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("prompt", None)
|
||||
payload_data["auto_prepare"] = False
|
||||
payload_data["workflow"] = {
|
||||
"prompt": {"10": {"class_type": "DistributedCollector"}},
|
||||
}
|
||||
payload = parse_queue_request_payload(payload_data)
|
||||
self.assertIn("10", payload.prompt)
|
||||
self.assertFalse(payload.auto_prepare)
|
||||
|
||||
def test_auto_prepare_must_be_boolean(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["auto_prepare"] = "true"
|
||||
with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_delegate_master_type_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["delegate_master"] = "yes"
|
||||
with self.assertRaisesRegex(ValueError, "delegate_master must be a boolean"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_enabled_worker_ids_type_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["enabled_worker_ids"] = "worker-a"
|
||||
with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_top_level_payload_raises(self):
|
||||
with self.assertRaisesRegex(ValueError, "Expected a JSON object body"):
|
||||
parse_queue_request_payload(["not", "an", "object"])
|
||||
|
||||
def test_missing_prompt_raises(self):
|
||||
with self.assertRaisesRegex(ValueError, "Field 'prompt' must be an object"):
|
||||
parse_queue_request_payload(
|
||||
{
|
||||
"workflow": {},
|
||||
"enabled_worker_ids": ["worker-1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
|
||||
def test_missing_enabled_worker_ids_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
with self.assertRaisesRegex(ValueError, "enabled_worker_ids required"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_missing_client_id_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("client_id", None)
|
||||
with self.assertRaisesRegex(ValueError, "client_id required"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,260 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_static_mode_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
|
||||
package_name = "dist_static_mode_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
|
||||
modes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
created_comfy_stub = False
|
||||
if "comfy" not in sys.modules:
|
||||
created_comfy_stub = True
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_mgmt = types.ModuleType("comfy.model_management")
|
||||
|
||||
class _InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_mgmt.processing_interrupted = lambda: False
|
||||
model_mgmt.throw_exception_if_processing_interrupted = lambda: None
|
||||
model_mgmt.InterruptProcessingException = _InterruptProcessingException
|
||||
|
||||
comfy_module.model_management = model_mgmt
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_mgmt
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
from PIL import Image as PILImage
|
||||
import numpy as np
|
||||
|
||||
def _tensor_to_pil(img_tensor, batch_index=0):
|
||||
return PILImage.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
|
||||
|
||||
def _pil_to_tensor(image):
|
||||
arr = np.array(image).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
image_module.tensor_to_pil = _tensor_to_pil
|
||||
image_module.pil_to_tensor = _pil_to_tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
def _run_async_in_server_loop(coro, timeout=None):
|
||||
if timeout is not None:
|
||||
return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
|
||||
return asyncio.run(coro)
|
||||
|
||||
async_helpers_module.run_async_in_server_loop = _run_async_in_server_loop
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 60
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 10.0
|
||||
constants_module.JOB_POLL_INTERVAL = 0.0
|
||||
constants_module.JOB_POLL_MAX_ATTEMPTS = 3
|
||||
constants_module.MAX_BATCH = 20
|
||||
constants_module.TILE_SEND_TIMEOUT = 1.0
|
||||
constants_module.TILE_WAIT_TIMEOUT = 1.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
|
||||
async def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
job_store_module.init_static_job_batched = _noop
|
||||
job_store_module._mark_task_completed = _noop
|
||||
job_store_module._cleanup_job = _noop
|
||||
job_store_module._drain_results_queue = _noop
|
||||
job_store_module._get_completed_count = _noop
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class _TileJobState:
|
||||
pass
|
||||
|
||||
job_models_module.TileJobState = _TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_comfy_stub:
|
||||
sys.modules.pop("comfy.model_management", None)
|
||||
sys.modules.pop("comfy", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
static_mode = _load_static_mode_module()
|
||||
|
||||
|
||||
class _FakeStaticWorker(static_mode.StaticModeMixin):
|
||||
def __init__(self):
|
||||
self.sent_batches = []
|
||||
self.request_calls = 0
|
||||
self.heartbeat_calls = 0
|
||||
self.job_ready = True
|
||||
self.tile_sequence = [(0, 0, True), (None, 0, True)]
|
||||
|
||||
def round_to_multiple(self, value):
|
||||
return value
|
||||
|
||||
def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
|
||||
return [(0, 0)]
|
||||
|
||||
def _poll_job_ready(self, *_args, **_kwargs):
|
||||
return self.job_ready
|
||||
|
||||
async def _request_tile_from_master(self, *_args, **_kwargs):
|
||||
self.request_calls += 1
|
||||
return self.tile_sequence.pop(0)
|
||||
|
||||
async def _send_heartbeat_to_master(self, *_args, **_kwargs):
|
||||
self.heartbeat_calls += 1
|
||||
|
||||
async def send_tiles_batch_to_master(
|
||||
self,
|
||||
processed_tiles,
|
||||
_multi_job_id,
|
||||
_master_url,
|
||||
_padding,
|
||||
_worker_id,
|
||||
is_final_flush=False,
|
||||
):
|
||||
self.sent_batches.append(
|
||||
{
|
||||
"tiles": list(processed_tiles),
|
||||
"is_final_flush": bool(is_final_flush),
|
||||
}
|
||||
)
|
||||
|
||||
def _extract_and_process_tile(self, upscaled_image, *_args, **_kwargs):
|
||||
batch_size = upscaled_image.shape[0]
|
||||
processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32)
|
||||
return processed_batch, 0, 0, 2, 2
|
||||
|
||||
def create_tile_mask(self, *_args, **_kwargs):
|
||||
from PIL import Image
|
||||
return Image.new("L", (4, 4), 255)
|
||||
|
||||
def blend_tile(self, base_image, *_args, **_kwargs):
|
||||
return base_image
|
||||
|
||||
|
||||
def _call_worker_static(fake_worker):
|
||||
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
return fake_worker._process_worker_static_sync(
|
||||
image,
|
||||
model=None,
|
||||
positive=None,
|
||||
negative=None,
|
||||
vae=None,
|
||||
seed=1,
|
||||
steps=1,
|
||||
cfg=1.0,
|
||||
sampler_name="euler",
|
||||
scheduler="normal",
|
||||
denoise=0.5,
|
||||
tile_width=4,
|
||||
tile_height=4,
|
||||
padding=8,
|
||||
mask_blur=4,
|
||||
force_uniform_tiles=True,
|
||||
tiled_decode=False,
|
||||
multi_job_id="job-1",
|
||||
master_url="http://master:8188",
|
||||
worker_id="worker-1",
|
||||
enabled_workers=["worker-1"],
|
||||
)
|
||||
|
||||
|
||||
class StaticModeWorkerFlowTests(unittest.TestCase):
|
||||
def test_worker_static_aborts_when_job_not_ready(self):
|
||||
worker = _FakeStaticWorker()
|
||||
worker.job_ready = False
|
||||
|
||||
result = _call_worker_static(worker)
|
||||
|
||||
self.assertEqual(result[0].shape[0], 1)
|
||||
self.assertEqual(worker.request_calls, 0)
|
||||
self.assertEqual(worker.heartbeat_calls, 0)
|
||||
self.assertEqual(worker.sent_batches, [])
|
||||
|
||||
def test_worker_static_requests_tiles_and_flushes_final_batch(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
_call_worker_static(worker)
|
||||
|
||||
self.assertEqual(worker.request_calls, 2) # one tile, then sentinel
|
||||
self.assertEqual(worker.heartbeat_calls, 1)
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
tiles = worker.sent_batches[0]["tiles"]
|
||||
self.assertEqual(len(tiles), 1)
|
||||
self.assertEqual(tiles[0]["tile_idx"], 0)
|
||||
self.assertEqual(tiles[0]["global_idx"], 0)
|
||||
self.assertEqual(tiles[0]["batch_idx"], 0)
|
||||
|
||||
def test_flush_empty_final_still_sends_completion_signal(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
returned = worker._flush_tiles_to_master(
|
||||
[],
|
||||
"job-1",
|
||||
"http://master:8188",
|
||||
8,
|
||||
"worker-1",
|
||||
is_final_flush=True,
|
||||
)
|
||||
|
||||
self.assertEqual(returned, [])
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertEqual(worker.sent_batches[0]["tiles"], [])
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,219 +0,0 @@
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _load_codec():
|
||||
path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
|
||||
spec = importlib.util.spec_from_file_location("dist_transfer_codec_test", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules["dist_transfer_codec_test"] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
tc = _load_codec()
|
||||
|
||||
|
||||
def _smooth_frames(count=4, height=67, width=91):
|
||||
"""Smooth moving gradient, closer to generated content than random noise."""
|
||||
yy, xx = np.mgrid[0:height, 0:width].astype(np.float32)
|
||||
frames = []
|
||||
for i in range(count):
|
||||
r = 0.5 + 0.5 * np.sin((xx + 3 * i) / 13.0)
|
||||
g = 0.5 + 0.5 * np.cos((yy - 2 * i) / 11.0)
|
||||
b = (xx + yy + 5 * i) / (width + height + 5 * count)
|
||||
frames.append(np.stack([r, g, b], axis=-1))
|
||||
return torch.from_numpy(np.stack(frames)).float()
|
||||
|
||||
|
||||
def _legacy_master_tensor(frame_tensor):
|
||||
"""What the unmodified pipeline delivered: tensor_to_pil -> PNG -> pil_to_tensor."""
|
||||
img = Image.fromarray((255 * frame_tensor.numpy()).astype(np.uint8))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG", compress_level=0)
|
||||
with Image.open(io.BytesIO(buf.getvalue())) as decoded:
|
||||
return np.array(decoded.convert("RGB")).astype(np.float32) / 255.0
|
||||
|
||||
|
||||
def _round_trip(frames_u8, fmt, quality=90, max_chunk_bytes=10**9):
|
||||
out = []
|
||||
chunks = list(tc.encode_chunks(frames_u8, fmt, quality, max_chunk_bytes))
|
||||
for chunk in chunks:
|
||||
envelope = {k: v for k, v in chunk.items() if k != "parts"}
|
||||
out.extend(tc.decode_chunk(envelope, chunk["parts"]))
|
||||
return chunks, out
|
||||
|
||||
|
||||
def _psnr(a, b):
|
||||
mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
|
||||
return float("inf") if mse == 0 else 10 * np.log10(255.0**2 / mse)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "ffv1"])
|
||||
def test_lossless_formats_reproduce_legacy_master_tensors_exactly(fmt):
|
||||
images = _smooth_frames()
|
||||
frames = tc.tensor_batch_to_uint8(images)
|
||||
_, decoded = _round_trip(frames, fmt)
|
||||
|
||||
assert len(decoded) == images.shape[0]
|
||||
for i, frame in enumerate(decoded):
|
||||
got = tc.uint8_frame_to_tensor_array(frame)[0]
|
||||
assert np.array_equal(got, _legacy_master_tensor(images[i]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
|
||||
def test_lossy_formats_round_trip_odd_sizes_with_reasonable_quality(fmt):
|
||||
if fmt in tc.VIDEO_FORMATS and not tc.video_codec_available(fmt):
|
||||
pytest.skip(f"{fmt} encoder not available in this PyAV build")
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames())
|
||||
_, decoded = _round_trip(frames, fmt, quality=90)
|
||||
|
||||
assert len(decoded) == frames.shape[0]
|
||||
for i, frame in enumerate(decoded):
|
||||
assert frame.shape == frames[i].shape
|
||||
assert frame.dtype == np.uint8
|
||||
assert _psnr(frame, frames[i]) > 30
|
||||
|
||||
|
||||
def test_h265_handles_frames_smaller_than_ctu():
|
||||
if not tc.video_codec_available("h265"):
|
||||
pytest.skip("h265 encoder not available")
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=3, height=17, width=23))
|
||||
_, decoded = _round_trip(frames, "h265")
|
||||
assert [f.shape for f in decoded] == [(17, 23, 3)] * 3
|
||||
|
||||
|
||||
def test_image_formats_split_chunks_by_byte_cap_and_keep_order():
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
|
||||
one = len(tc.encode_image_frame(frames[0], "png", 90))
|
||||
chunks, decoded = _round_trip(frames, "png", max_chunk_bytes=one * 2 + 1)
|
||||
|
||||
assert len(chunks) >= 3
|
||||
assert [c["start_index"] for c in chunks] == sorted(c["start_index"] for c in chunks)
|
||||
assert sum(c["count"] for c in chunks) == 5
|
||||
assert all(len(c["parts"]) == c["count"] for c in chunks)
|
||||
assert np.array_equal(np.stack(decoded), frames)
|
||||
|
||||
|
||||
def test_video_formats_split_chunks_by_frame_cap(monkeypatch):
|
||||
monkeypatch.setattr(tc, "MAX_VIDEO_CHUNK_FRAMES", 2)
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
|
||||
chunks, decoded = _round_trip(frames, "ffv1")
|
||||
|
||||
assert [(c["start_index"], c["count"]) for c in chunks] == [(0, 2), (2, 2), (4, 1)]
|
||||
assert all(len(c["parts"]) == 1 for c in chunks)
|
||||
assert np.array_equal(np.stack(decoded), frames)
|
||||
|
||||
|
||||
def test_tensor_batch_to_uint8_clips_and_normalizes_channels():
|
||||
rgba = torch.tensor([[[[1.5, -0.2, 0.5, 0.3]]]])
|
||||
assert tc.tensor_batch_to_uint8(rgba).tolist() == [[[[255, 0, 127]]]]
|
||||
|
||||
gray = torch.full((1, 2, 2, 1), 0.5)
|
||||
out = tc.tensor_batch_to_uint8(gray)
|
||||
assert out.shape == (1, 2, 2, 3)
|
||||
assert (out == 127).all()
|
||||
|
||||
|
||||
def test_quality_to_crf_mapping():
|
||||
assert tc.quality_to_crf(100, 51) == 0
|
||||
assert tc.quality_to_crf(90, 51) == 13
|
||||
assert tc.quality_to_crf(80, 51) == 19
|
||||
assert tc.quality_to_crf(90, 63) == 16
|
||||
assert tc.quality_to_crf(1, 51) == 51
|
||||
assert tc.quality_to_crf("bogus", 51) == tc.quality_to_crf(tc.DEFAULT_TRANSFER_QUALITY, 51)
|
||||
|
||||
|
||||
def test_normalize_format_rejects_unknown_values():
|
||||
assert tc.normalize_format(" H264 ") == "h264"
|
||||
assert tc.normalize_format(None) == tc.DEFAULT_TRANSFER_FORMAT
|
||||
with pytest.raises(tc.TransferCodecError):
|
||||
tc.normalize_format("gif")
|
||||
|
||||
|
||||
def test_decode_chunk_validates_envelope_against_payload():
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=2))
|
||||
chunk = next(tc.encode_chunks(frames, "png", 90, 10**9))
|
||||
envelope = {k: v for k, v in chunk.items() if k != "parts"}
|
||||
|
||||
with pytest.raises(tc.TransferCodecError, match="image parts"):
|
||||
tc.decode_chunk({**envelope, "count": 3}, chunk["parts"])
|
||||
with pytest.raises(tc.TransferCodecError, match="expected"):
|
||||
tc.decode_chunk({**envelope, "width": 10}, chunk["parts"])
|
||||
with pytest.raises(tc.TransferCodecError, match="count"):
|
||||
tc.decode_chunk({**envelope, "count": 0}, chunk["parts"])
|
||||
with pytest.raises(tc.TransferCodecError):
|
||||
tc.decode_chunk(envelope, [b"not an image", b"nope"])
|
||||
|
||||
video = next(tc.encode_chunks(frames, "ffv1", 90, 10**9))
|
||||
video_env = {k: v for k, v in video.items() if k != "parts"}
|
||||
with pytest.raises(tc.TransferCodecError, match="frames"):
|
||||
tc.decode_chunk({**video_env, "count": 5}, video["parts"])
|
||||
|
||||
|
||||
def test_legacy_format_is_not_chunk_encoded():
|
||||
frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
|
||||
with pytest.raises(tc.TransferCodecError):
|
||||
list(tc.encode_chunks(frames, "legacy_png", 90, 10**9))
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Size bounds on finalized payloads
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _noise_frames(count=30, height=360, width=640, seed=97):
|
||||
"""High-entropy frames: worst case for every codec, and enough encoder
|
||||
lookahead that pre-flush buffer sizes badly underestimate final sizes."""
|
||||
return np.random.default_rng(seed).integers(0, 256, size=(count, height, width, 3), dtype=np.uint8)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
|
||||
def test_video_chunks_are_bounded_after_encoder_flush(fmt):
|
||||
frames = _noise_frames()
|
||||
limit = 1024 * 1024
|
||||
# Sanity: one unconstrained encode of this batch is far above the limit,
|
||||
# so the bound below is actually being enforced.
|
||||
unbounded = next(tc.encode_chunks(frames, fmt, 90, 10**12))
|
||||
assert len(unbounded["parts"][0]) > 3 * limit
|
||||
|
||||
chunks = list(tc.encode_chunks(frames, fmt, 90, limit))
|
||||
|
||||
assert all(len(c["parts"]) == 1 and len(c["parts"][0]) <= limit for c in chunks)
|
||||
assert [c["start_index"] for c in chunks] == list(np.cumsum([0] + [c["count"] for c in chunks[:-1]]))
|
||||
assert sum(c["count"] for c in chunks) == len(frames)
|
||||
decoded = [f for c in chunks for f in tc.decode_chunk(c, c["parts"])]
|
||||
assert len(decoded) == len(frames)
|
||||
if fmt == "ffv1":
|
||||
assert np.array_equal(np.stack(decoded), frames)
|
||||
|
||||
|
||||
def test_single_video_frame_may_exceed_chunk_target_up_to_frame_limit():
|
||||
frames = _noise_frames(count=3)
|
||||
one_frame = len(next(tc.encode_chunks(frames[:1], "ffv1", 90, 10**12))["parts"][0])
|
||||
|
||||
chunks = list(tc.encode_chunks(frames, "ffv1", 90, one_frame // 2, one_frame * 2))
|
||||
|
||||
assert [c["count"] for c in chunks] == [1, 1, 1]
|
||||
assert all(one_frame // 2 < len(c["parts"][0]) <= one_frame * 2 for c in chunks)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["ffv1", "h264", "png"])
|
||||
def test_frame_that_cannot_fit_alone_raises(fmt):
|
||||
frames = _noise_frames(count=2)
|
||||
with pytest.raises(tc.FrameTooLargeError, match="request limit"):
|
||||
list(tc.encode_chunks(frames, fmt, 90, 64 * 1024, 64 * 1024))
|
||||
|
||||
|
||||
def test_still_chunks_are_bounded():
|
||||
frames = _noise_frames(count=10)
|
||||
limit = 1024 * 1024
|
||||
chunks = list(tc.encode_chunks(frames, "jpeg", 90, limit))
|
||||
assert all(sum(len(p) for p in c["parts"]) <= limit for c in chunks)
|
||||
assert sum(c["count"] for c in chunks) == 10
|
||||
@@ -1,125 +0,0 @@
|
||||
"""Registered V3 DistributedCollector: transfer options in the public schema and forwarded to the runtime.
|
||||
|
||||
Runs without a ComfyUI checkout by loading nodes/v3.py against a recording
|
||||
stand-in for comfy_api's io module. tests/test_v3_migration.py covers the same
|
||||
boundary against the real framework when COMFYUI_SOURCE_ROOT is set.
|
||||
"""
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_v3_transfer_testpkg"
|
||||
|
||||
|
||||
class _Field:
|
||||
def __init__(self, kind, id=None, **kwargs):
|
||||
self.kind = kind
|
||||
self.id = id
|
||||
self.kwargs = kwargs
|
||||
|
||||
|
||||
def _io_type(kind):
|
||||
return types.SimpleNamespace(
|
||||
Input=lambda id, **kw: _Field(kind, id, **kw),
|
||||
Output=lambda id=None, **kw: _Field(kind, id, **kw),
|
||||
)
|
||||
|
||||
|
||||
def _fake_io():
|
||||
io = types.ModuleType("io")
|
||||
|
||||
class ComfyNode:
|
||||
pass
|
||||
|
||||
io.ComfyNode = ComfyNode
|
||||
io.Schema = lambda **kw: types.SimpleNamespace(**kw)
|
||||
io.NodeOutput = lambda *values, **kw: ("NodeOutput", values, kw)
|
||||
io.NumberDisplay = types.SimpleNamespace(number="number")
|
||||
io.Hidden = types.SimpleNamespace(unique_id="unique_id", extra_pnginfo="extra_pnginfo")
|
||||
for name in ("Int", "Float", "String", "Boolean", "Combo", "Image", "Audio",
|
||||
"AnyType", "Model", "Conditioning", "Vae"):
|
||||
setattr(io, name, _io_type(name))
|
||||
return io
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def v3(monkeypatch):
|
||||
calls = []
|
||||
|
||||
class CollectorRuntime:
|
||||
def run(self, **kwargs):
|
||||
calls.append(kwargs)
|
||||
return (kwargs["images"], kwargs["audio"])
|
||||
|
||||
def module(name, **attrs):
|
||||
mod = types.ModuleType(name)
|
||||
mod.__dict__.update(attrs)
|
||||
monkeypatch.setitem(sys.modules, name, mod)
|
||||
return mod
|
||||
|
||||
for name in [n for n in sys.modules if n == PKG or n.startswith(PKG + ".")]:
|
||||
monkeypatch.delitem(sys.modules, name)
|
||||
module(PKG, __path__=[])
|
||||
module(f"{PKG}.nodes", __path__=[])
|
||||
module(f"{PKG}.utils", __path__=[])
|
||||
module(f"{PKG}.nodes.utilities")
|
||||
module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
|
||||
module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
|
||||
module("comfy", __path__=[])
|
||||
module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
|
||||
module("comfy_api", __path__=[])
|
||||
module("comfy_api.v0_0_2", io=_fake_io())
|
||||
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py")
|
||||
codec = importlib.util.module_from_spec(codec_spec)
|
||||
monkeypatch.setitem(sys.modules, f"{PKG}.utils.transfer_codec", codec)
|
||||
codec_spec.loader.exec_module(codec)
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{PKG}.nodes.v3", ROOT / "nodes" / "v3.py")
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod, codec, calls
|
||||
|
||||
|
||||
def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
|
||||
mod, codec, _ = v3
|
||||
schema = mod.DistributedCollector.define_schema()
|
||||
inputs = {field.id: field for field in schema.inputs}
|
||||
|
||||
assert [field.id for field in schema.inputs] == [
|
||||
"load_balance", "images", "audio", "transfer_format", "transfer_quality"]
|
||||
fmt = inputs["transfer_format"]
|
||||
assert fmt.kind == "Combo"
|
||||
assert fmt.kwargs["options"] == list(codec.TRANSFER_FORMATS)
|
||||
assert fmt.kwargs["default"] == "legacy_png"
|
||||
assert fmt.kwargs["optional"] is True
|
||||
quality = inputs["transfer_quality"]
|
||||
assert quality.kind == "Int"
|
||||
assert (quality.kwargs["default"], quality.kwargs["min"], quality.kwargs["max"]) == (90, 1, 100)
|
||||
assert quality.kwargs["optional"] is True
|
||||
assert schema.is_input_list is True
|
||||
|
||||
|
||||
def test_registered_collector_forwards_transfer_options(v3):
|
||||
mod, _, calls = v3
|
||||
|
||||
mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=["av1"], transfer_quality=[70],
|
||||
multi_job_id=["job"], is_worker=[True])
|
||||
|
||||
assert calls[-1]["transfer_format"] == ["av1"]
|
||||
assert calls[-1]["transfer_quality"] == [70]
|
||||
assert calls[-1]["multi_job_id"] == ["job"]
|
||||
|
||||
|
||||
def test_registered_collector_defaults_to_legacy_transfer(v3):
|
||||
mod, _, calls = v3
|
||||
|
||||
mod.DistributedCollector.execute(images=["img"])
|
||||
|
||||
assert calls[-1]["transfer_format"] == "legacy_png"
|
||||
assert calls[-1]["transfer_quality"] == 90
|
||||
@@ -1,21 +0,0 @@
|
||||
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_native_v3_runtime():
|
||||
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not comfy_root:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
|
||||
helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
|
||||
capture_output=True, text=True, timeout=120,
|
||||
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'V3_ACCEPTANCE_OK' in result.stdout
|
||||
@@ -1,113 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import socket
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def load_ports():
|
||||
path = Path(__file__).resolve().parents[1] / "workers" / "ports.py"
|
||||
spec = importlib.util.spec_from_file_location("worker_ports_test", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class WorkerPortTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ports = load_ports()
|
||||
|
||||
def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self):
|
||||
workers = [
|
||||
{"id": "existing", "host": "localhost", "port": 8190, "enabled": False},
|
||||
{"id": "remote", "host": "remote.example", "port": 8192},
|
||||
]
|
||||
with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194])
|
||||
|
||||
def test_local_host_forms_reserve_ports(self):
|
||||
for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]:
|
||||
with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191])
|
||||
|
||||
def test_exhaustion_is_explicit_not_a_partial_allocation(self):
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
with self.assertRaisesRegex(ValueError, "available.*ports"):
|
||||
self.ports.allocate_worker_ports(65534, [], 2)
|
||||
|
||||
def test_launch_rejects_master_port_without_reassigning_configuration(self):
|
||||
worker = {"id": "manual", "port": 8189}
|
||||
with self.assertRaisesRegex(ValueError, "master.*8189"):
|
||||
self.ports.validate_worker_port(worker, 8189, [])
|
||||
self.assertEqual(worker["port"], 8189)
|
||||
|
||||
def test_launch_rejects_another_local_workers_reserved_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "other", "host": None, "port": 8190, "enabled": False}
|
||||
with self.assertRaisesRegex(ValueError, "other"):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_ignores_itself_and_remote_workers_with_same_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "remote", "host": "example.com", "port": 8190}
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_rejects_an_occupied_port(self):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen()
|
||||
port = listener.getsockname()[1]
|
||||
self.assertFalse(self.ports.is_port_available(port))
|
||||
with self.assertRaisesRegex(ValueError, "already in use"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "asyncio does not reuse addresses on Windows")
|
||||
def test_recently_closed_connection_does_not_block_worker_restart(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
port = listener.getsockname()[1]
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=2) as client:
|
||||
accepted, _ = listener.accept()
|
||||
with accepted:
|
||||
accepted.settimeout(2)
|
||||
accepted.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
client.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(accepted.recv(1), b"")
|
||||
self.assertTrue(self.ports.is_port_available(port))
|
||||
self.ports.validate_worker_port({"id": "restart", "port": port}, 1, [])
|
||||
with socket.socket() as restarted:
|
||||
restarted.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
restarted.bind(("127.0.0.1", port))
|
||||
restarted.listen(1)
|
||||
|
||||
def test_active_listener_with_reuseaddr_is_still_a_conflict(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
@unittest.skipUnless(socket.has_ipv6, "IPv6 unavailable")
|
||||
def test_ipv6_only_listener_is_still_a_conflict(self):
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as listener:
|
||||
listener.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
|
||||
try:
|
||||
listener.bind(("::1", 0))
|
||||
except OSError as exc:
|
||||
self.skipTest(f"IPv6 loopback unavailable: {exc}")
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
def test_invalid_ports_fail_clearly(self):
|
||||
for port in [0, 65536, "not-a-port", None, True, 8190.5]:
|
||||
with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,352 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import URL, make_url
|
||||
|
||||
|
||||
def _load_process_module(module_filename: str):
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
|
||||
package_name = "dist_proc_testpkg"
|
||||
module_name = module_filename[:-3]
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = [str(module_path.parents[1])]
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
process_pkg = types.ModuleType(f"{package_name}.workers.process")
|
||||
process_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers.process"] = process_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
process_module = types.ModuleType(f"{package_name}.utils.process")
|
||||
process_module.get_python_executable = lambda: "/usr/bin/test-python"
|
||||
process_module.is_process_alive = lambda _pid: False
|
||||
process_module.terminate_process = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.process"] = process_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}}
|
||||
config_module.save_config = lambda _config: None
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.PROCESS_TERMINATION_TIMEOUT = 1
|
||||
constants_module.PROCESS_WAIT_TIMEOUT = 1
|
||||
constants_module.WORKER_CHECK_INTERVAL = 0.01
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.process.{module_name}",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
root_discovery_module = _load_process_module("root_discovery.py")
|
||||
launch_builder_module = _load_process_module("launch_builder.py")
|
||||
|
||||
|
||||
class ComfyRootDiscoveryTests(unittest.TestCase):
|
||||
def test_prefers_loaded_comfyui_module_path(self):
|
||||
discovery = root_discovery_module.ComfyRootDiscovery()
|
||||
server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/opt/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(sys.modules, {"server": server_module}, clear=False), \
|
||||
patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
|
||||
patch.dict(root_discovery_module.os.environ, {}, clear=True):
|
||||
self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
|
||||
|
||||
|
||||
class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
def build_command(self, root, worker=None, runtime=None):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
worker = worker or {"id": "worker-a", "port": 8190}
|
||||
runtime = runtime if runtime is not None else Namespace(database_url=None)
|
||||
with patch.object(builder, "_get_runtime_args", return_value=runtime):
|
||||
return builder.build_launch_command(worker, str(root))
|
||||
|
||||
def test_database_is_unique_and_stable_by_worker_id(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
first = self.build_command(root, {"id": "worker-a", "port": 8190})
|
||||
renamed = self.build_command(root, {
|
||||
"id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3,
|
||||
})
|
||||
second = self.build_command(root, {"id": "worker-b", "port": 8191})
|
||||
database = first[first.index("--database-url") + 1]
|
||||
self.assertEqual(database, renamed[renamed.index("--database-url") + 1])
|
||||
self.assertNotEqual(database, second[second.index("--database-url") + 1])
|
||||
self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/"))
|
||||
self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir())
|
||||
self.assertNotIn("--disable-assets", first)
|
||||
|
||||
def test_database_uses_effective_user_or_base_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra_args, expected in [
|
||||
(f'--user-directory "{root / "custom user"}"', root / "custom user"),
|
||||
(f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"),
|
||||
]:
|
||||
with self.subTest(extra_args=extra_args):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra_args,
|
||||
})
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/"))
|
||||
cmd = self.build_command(root, runtime=Namespace(
|
||||
database_url="sqlite:///master.db", user_directory=str(root / "runtime user"),
|
||||
))
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/"))
|
||||
self.assertNotEqual(database, "sqlite:///master.db")
|
||||
|
||||
def test_relative_database_directories_use_worker_cwd(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory) / "ComfyUI"
|
||||
root.mkdir()
|
||||
(root / "main.py").touch()
|
||||
launcher = Path(directory) / "launcher"
|
||||
launcher.mkdir()
|
||||
cases = [
|
||||
(Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"),
|
||||
(Namespace(database_url=None, base_directory="data"), "", root / "data/user"),
|
||||
(Namespace(database_url=None), "--user-directory=profiles", root / "profiles"),
|
||||
(Namespace(database_url=None), "--base-directory data", root / "data/user"),
|
||||
]
|
||||
for runtime, extra, expected in cases:
|
||||
with self.subTest(runtime=runtime, extra=extra), \
|
||||
patch.object(os, "getcwd", return_value=str(launcher)):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra,
|
||||
}, runtime)
|
||||
database = make_url(cmd[cmd.index("--database-url") + 1]).database
|
||||
self.assertEqual(Path(database).parent, expected / "distributed/workers")
|
||||
self.assertFalse((launcher / "profiles").exists())
|
||||
self.assertFalse((launcher / "data").exists())
|
||||
|
||||
def test_database_url_round_trips_special_characters_without_collisions(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for name in ["user?profile", "user%3Fprofile"]:
|
||||
with self.subTest(name=name):
|
||||
user = root / name
|
||||
runtime = Namespace(database_url=None, user_directory=str(user))
|
||||
# SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail
|
||||
# clearly on those versions rather than silently sharing a DB.
|
||||
sample = (user / "test.db").as_posix()
|
||||
serialized = URL.create("sqlite", database=sample).render_as_string()
|
||||
if make_url(serialized).database != sample:
|
||||
with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"):
|
||||
self.build_command(root, runtime=runtime)
|
||||
self.assertFalse(user.exists())
|
||||
continue
|
||||
databases = []
|
||||
for worker_id in ["worker-a", "worker-b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime)
|
||||
url = cmd[cmd.index("--database-url") + 1]
|
||||
database = make_url(url).database
|
||||
self.assertEqual(Path(database).parent, user / "distributed/workers")
|
||||
engine = create_engine(url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
actual = connection.exec_driver_sql("PRAGMA database_list").one()[2]
|
||||
self.assertEqual(Path(actual), Path(database))
|
||||
finally:
|
||||
engine.dispose()
|
||||
databases.append(database)
|
||||
self.assertNotEqual(*databases)
|
||||
|
||||
def test_explicit_database_and_disabled_assets_are_preserved(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]:
|
||||
with self.subTest(extra=extra):
|
||||
cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_older_comfyui_without_database_flag_keeps_existing_command(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
cmd = self.build_command(root, runtime=Namespace())
|
||||
self.assertNotIn("--database-url", cmd)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_worker_id_cannot_escape_database_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
databases = []
|
||||
for worker_id in ["../outside", "a/b", "a_b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190})
|
||||
path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///"))
|
||||
self.assertTrue(path.is_relative_to(root / "user"))
|
||||
databases.append(path)
|
||||
self.assertEqual(len(set(databases)), 3)
|
||||
|
||||
def test_extra_args_cannot_silently_override_the_configured_port(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--port 8189", "--port=8189"]:
|
||||
with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"):
|
||||
self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
|
||||
def test_inherits_runtime_layout_args_for_desktop(self):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
runtime_args = Namespace(
|
||||
listen="127.0.0.1",
|
||||
base_directory="C:/Users/test/ComfyUI",
|
||||
temp_directory=None,
|
||||
input_directory="C:/Users/test/ComfyUI/input",
|
||||
output_directory="C:/Users/test/ComfyUI/output",
|
||||
user_directory="C:/Users/test/ComfyUI/user",
|
||||
front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
|
||||
extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
|
||||
enable_manager=True,
|
||||
disable_manager_ui=False,
|
||||
enable_manager_legacy_ui=False,
|
||||
windows_standalone_build=True,
|
||||
log_stdout=True,
|
||||
verbose="INFO",
|
||||
enable_cors_header="*",
|
||||
)
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
comfy_cli_args = types.ModuleType("comfy.cli_args")
|
||||
comfy_cli_args.args = runtime_args
|
||||
|
||||
worker_config = {
|
||||
"port": 9001,
|
||||
"extra_args": "--preview-method auto",
|
||||
}
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/desktop/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
|
||||
clear=False,
|
||||
), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
|
||||
cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
|
||||
|
||||
self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
|
||||
self.assertIn("--listen", cmd)
|
||||
self.assertIn("127.0.0.1", cmd)
|
||||
self.assertIn("--base-directory", cmd)
|
||||
self.assertIn("C:/Users/test/ComfyUI", cmd)
|
||||
self.assertIn("--input-directory", cmd)
|
||||
self.assertIn("--output-directory", cmd)
|
||||
self.assertIn("--user-directory", cmd)
|
||||
self.assertIn("--front-end-root", cmd)
|
||||
self.assertIn("--extra-model-paths-config", cmd)
|
||||
self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
|
||||
self.assertIn("--enable-manager", cmd)
|
||||
self.assertIn("--windows-standalone-build", cmd)
|
||||
self.assertIn("--log-stdout", cmd)
|
||||
self.assertIn("--disable-auto-launch", cmd)
|
||||
self.assertIn("--enable-cors-header", cmd)
|
||||
self.assertIn("*", cmd)
|
||||
self.assertIn("--port", cmd)
|
||||
self.assertIn("9001", cmd)
|
||||
self.assertNotIn("--auto-launch", cmd)
|
||||
|
||||
|
||||
class ProcessLaunchPortTests(unittest.TestCase):
|
||||
def test_launch_before_master_listens_uses_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
for missing in [AttributeError("PromptServer has no port"), None]:
|
||||
with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \
|
||||
patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
|
||||
def test_pre_listen_launch_still_rejects_the_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \
|
||||
patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
|
||||
def test_conflict_is_rejected_before_building_or_spawning(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
def test_available_port_reaches_existing_launch_path(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
self.assertEqual(manager.processes["manual"]["pid"], 1234)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,32 +0,0 @@
|
||||
import copy
|
||||
|
||||
|
||||
def clone_control_chain(control, clone_hint=True):
|
||||
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
|
||||
if control is None:
|
||||
return None
|
||||
new_control = copy.copy(control)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
hint = getattr(control, 'cond_hint_original', None)
|
||||
new_control.cond_hint_original = hint.clone() if hint is not None else None
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
|
||||
|
||||
def clone_conditioning(cond_list, clone_hints=True):
|
||||
"""Clone conditioning without duplicating ControlNet models."""
|
||||
new_cond = []
|
||||
for emb, cond_dict in cond_list:
|
||||
new_emb = emb.clone() if emb is not None else None
|
||||
new_dict = cond_dict.copy()
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict and new_dict['mask'] is not None:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
if 'pooled_output' in new_dict and new_dict['pooled_output'] is not None:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:]
|
||||
new_cond.append([new_emb, new_dict])
|
||||
return new_cond
|
||||
@@ -1,49 +0,0 @@
|
||||
from dataclasses import dataclass, field
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
|
||||
class BaseJobState:
|
||||
"""Marker base class for typed USDU job state containers."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
"""Typed state container for static (tile) USDU jobs."""
|
||||
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
"""Typed state container for dynamic (per-image) USDU jobs."""
|
||||
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
@@ -1,135 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from .job_store import ensure_tile_jobs_initialized
|
||||
from .job_timeout import _check_and_requeue_timed_out_workers as _requeue_usdu
|
||||
from .job_models import ImageJobState, TileJobState
|
||||
|
||||
|
||||
class JobStateMixin:
|
||||
async def _get_job_data(self, multi_job_id):
|
||||
"""Return current job data reference while holding lock briefly."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
return prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
|
||||
async def _get_all_completed_tasks(self, multi_job_id):
|
||||
"""Helper to retrieve all completed tasks from the job data."""
|
||||
job_data = await self._get_job_data(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
return dict(job_data.completed_tasks)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return dict(job_data.completed_images)
|
||||
return {}
|
||||
|
||||
async def _get_next_image_index(self, multi_job_id):
|
||||
"""Get next image index from pending queue for master."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
pending_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
pending_queue = job_data.pending_images
|
||||
|
||||
if pending_queue is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(pending_queue.get(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _get_next_tile_index(self, multi_job_id):
|
||||
"""Get next tile index from pending queue for master in static mode."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
pending_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
|
||||
if pending_queue is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _get_total_completed_count(self, multi_job_id):
|
||||
"""Get total count of all completed images (master + workers)."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return len(job_data.completed_images)
|
||||
if isinstance(job_data, TileJobState):
|
||||
return len(job_data.completed_tasks)
|
||||
return 0
|
||||
|
||||
async def _get_all_completed_images(self, multi_job_id):
|
||||
"""Get all completed images."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return job_data.completed_images.copy()
|
||||
return {}
|
||||
|
||||
async def _get_pending_count(self, multi_job_id):
|
||||
"""Get count of pending images in the queue."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return job_data.pending_images.qsize()
|
||||
if isinstance(job_data, TileJobState):
|
||||
return job_data.pending_tasks.qsize()
|
||||
return 0
|
||||
|
||||
async def _drain_worker_results_queue(self, multi_job_id):
|
||||
"""Drain pending worker results from queue and update completed images."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
worker_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
worker_queue = job_data.queue
|
||||
|
||||
if worker_queue is None:
|
||||
return 0
|
||||
|
||||
drained_results = []
|
||||
while True:
|
||||
try:
|
||||
drained_results.append(worker_queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
if not drained_results:
|
||||
return 0
|
||||
|
||||
collected = 0
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return 0
|
||||
|
||||
for result in drained_results:
|
||||
worker_id = result.get("worker_id")
|
||||
if "image_idx" in result and "image" in result:
|
||||
image_idx = result["image_idx"]
|
||||
image_pil = result["image"]
|
||||
if image_idx not in job_data.completed_images:
|
||||
job_data.completed_images[image_idx] = image_pil
|
||||
collected += 1
|
||||
debug_log(f"Drained image {image_idx} from worker {worker_id}")
|
||||
|
||||
if collected > 0:
|
||||
debug_log(f"Drained {collected} worker images during retry")
|
||||
|
||||
return collected
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
|
||||
"""Check for timed out workers and requeue their assigned images."""
|
||||
return await _requeue_usdu(multi_job_id, batch_size)
|
||||
@@ -1,180 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from .job_models import BaseJobState, ImageJobState, TileJobState
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
|
||||
def ensure_tile_jobs_initialized():
|
||||
"""Ensure tile job storage is initialized on the server instance."""
|
||||
prompt_server = server.PromptServer.instance
|
||||
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
|
||||
debug_log("Initializing persistent tile job queue on server instance.")
|
||||
prompt_server.distributed_pending_tile_jobs = {}
|
||||
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
else:
|
||||
invalid_job_ids = [
|
||||
job_id
|
||||
for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
|
||||
if not isinstance(job_data, BaseJobState)
|
||||
]
|
||||
for job_id in invalid_job_ids:
|
||||
debug_log(f"Removing invalid job state for {job_id}")
|
||||
del prompt_server.distributed_pending_tile_jobs[job_id]
|
||||
return prompt_server
|
||||
|
||||
|
||||
async def _init_job_queue(
|
||||
multi_job_id,
|
||||
mode,
|
||||
batch_size=None,
|
||||
num_tiles_per_image=None,
|
||||
all_indices=None,
|
||||
enabled_workers=None,
|
||||
batched_static: bool = False,
|
||||
):
|
||||
"""Unified initialization for job queues in static and dynamic modes."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
debug_log(f"Queue already exists for {multi_job_id}")
|
||||
return
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data = ImageJobState(multi_job_id=multi_job_id)
|
||||
elif mode == 'static':
|
||||
job_data = TileJobState(multi_job_id=multi_job_id)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
job_data.worker_status = {w: time.time() for w in enabled_workers or []}
|
||||
job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data.batch_size = int(batch_size or 0)
|
||||
pending_queue = job_data.pending_images
|
||||
for i in (all_indices or range(int(batch_size or 0))):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized image queue with {batch_size} pending items")
|
||||
elif mode == 'static':
|
||||
job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
|
||||
job_data.batch_size = int(batch_size or 0)
|
||||
job_data.batched_static = bool(batched_static)
|
||||
# For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
|
||||
pending_queue = job_data.pending_tasks
|
||||
if batched_static and num_tiles_per_image is not None:
|
||||
for i in range(num_tiles_per_image):
|
||||
await pending_queue.put(i)
|
||||
else:
|
||||
total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
|
||||
for i in range(total_tiles):
|
||||
await pending_queue.put(i)
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
|
||||
async def init_dynamic_job(
|
||||
multi_job_id: str,
|
||||
batch_size: int,
|
||||
enabled_workers: List[str],
|
||||
all_indices: Optional[List[int]] = None,
|
||||
):
|
||||
"""Initialize queue for dynamic mode (per-image), with collector fields."""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'dynamic',
|
||||
batch_size=batch_size,
|
||||
all_indices=all_indices or list(range(batch_size)),
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
|
||||
|
||||
|
||||
async def init_static_job_batched(
|
||||
multi_job_id: str,
|
||||
batch_size: int,
|
||||
num_tiles_per_image: int,
|
||||
enabled_workers: List[str],
|
||||
):
|
||||
"""Initialize queue for static mode (batched-per-tile)."""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
enabled_workers=enabled_workers,
|
||||
batched_static=True,
|
||||
)
|
||||
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
q = job_data.queue
|
||||
completed_tasks = job_data.completed_tasks
|
||||
|
||||
collected = 0
|
||||
while True:
|
||||
try:
|
||||
result = q.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
if is_last:
|
||||
if worker_id in job_data.worker_status:
|
||||
del job_data.worker_status[worker_id]
|
||||
|
||||
return collected
|
||||
|
||||
|
||||
async def _get_completed_count(multi_job_id):
|
||||
"""Get count of completed tasks."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, BaseJobState):
|
||||
return len(job_data.completed_tasks)
|
||||
return 0
|
||||
|
||||
|
||||
async def _mark_task_completed(multi_job_id, task_id, result):
|
||||
"""Mark a task as completed."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, BaseJobState):
|
||||
job_data.completed_tasks[task_id] = result
|
||||
|
||||
|
||||
async def _cleanup_job(multi_job_id):
|
||||
"""Cleanup the job data."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
debug_log(f"Cleaned up job {multi_job_id}")
|
||||
@@ -1,150 +0,0 @@
|
||||
import time
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.constants import HEARTBEAT_TIMEOUT
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_worker_url, probe_worker
|
||||
from .job_models import BaseJobState
|
||||
from .job_store import ensure_tile_jobs_initialized
|
||||
|
||||
|
||||
def _find_worker_record(worker_id):
|
||||
"""Return worker config entry by id, or None when missing."""
|
||||
workers = load_config().get("workers", [])
|
||||
return next((w for w in workers if str(w.get("id")) == str(worker_id)), None)
|
||||
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
"""Check timed out workers and requeue their tasks. Returns requeued count."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
current_time = time.time()
|
||||
|
||||
# Allow override via config setting 'worker_timeout_seconds'
|
||||
cfg = load_config()
|
||||
hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT))
|
||||
|
||||
# Snapshot timed-out workers and job details under lock.
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
|
||||
completed_tasks_snapshot = set(job_data.completed_tasks.keys())
|
||||
batched_static_snapshot = bool(job_data.batched_static)
|
||||
num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1)
|
||||
batch_size_snapshot = int(job_data.batch_size or 1)
|
||||
|
||||
timed_out_workers = []
|
||||
for worker, last_heartbeat in list(job_data.worker_status.items()):
|
||||
age = current_time - float(last_heartbeat)
|
||||
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
|
||||
if age > hb_timeout:
|
||||
timed_out_workers.append(
|
||||
{
|
||||
"worker_id": worker,
|
||||
"last_heartbeat": float(last_heartbeat),
|
||||
"assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])),
|
||||
}
|
||||
)
|
||||
|
||||
if not timed_out_workers:
|
||||
return 0
|
||||
|
||||
# Probe outside lock to avoid lock contention on network latency.
|
||||
workers_to_requeue = []
|
||||
workers_graced = []
|
||||
for worker_info in timed_out_workers:
|
||||
worker = worker_info["worker_id"]
|
||||
assigned = worker_info["assigned_tasks"]
|
||||
age = current_time - worker_info["last_heartbeat"]
|
||||
|
||||
incomplete_assigned = 0
|
||||
try:
|
||||
if assigned:
|
||||
if batched_static_snapshot:
|
||||
for task_id in assigned:
|
||||
for b in range(batch_size_snapshot):
|
||||
gidx = b * num_tiles_per_image_snapshot + task_id
|
||||
if gidx not in completed_tasks_snapshot:
|
||||
incomplete_assigned += 1
|
||||
break
|
||||
else:
|
||||
for task_id in assigned:
|
||||
if task_id not in completed_tasks_snapshot:
|
||||
incomplete_assigned += 1
|
||||
debug_log(
|
||||
f"Assigned diagnostics: total_assigned={len(assigned)} "
|
||||
f"incomplete_assigned={incomplete_assigned}"
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
|
||||
|
||||
busy = False
|
||||
probe_queue = None
|
||||
try:
|
||||
worker_record = _find_worker_record(worker)
|
||||
if worker_record:
|
||||
worker_url = build_worker_url(worker_record)
|
||||
debug_log(f"Probing worker {worker} at {worker_url}/prompt")
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
if payload is not None:
|
||||
probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0))
|
||||
busy = probe_queue is not None and probe_queue > 0
|
||||
else:
|
||||
debug_log(f"Probe skipped; worker {worker} not found in config")
|
||||
except Exception as e:
|
||||
debug_log(f"Probe failed for worker {worker}: {e}")
|
||||
finally:
|
||||
debug_log(
|
||||
f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
|
||||
)
|
||||
|
||||
if busy:
|
||||
workers_graced.append(worker)
|
||||
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
|
||||
continue
|
||||
|
||||
log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
|
||||
workers_to_requeue.append((worker, assigned))
|
||||
|
||||
# Re-acquire lock and apply requeue/cleanup decisions.
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
|
||||
# Refresh heartbeat for workers that we proved are still busy.
|
||||
for worker in workers_graced:
|
||||
if worker in job_data.worker_status:
|
||||
job_data.worker_status[worker] = current_time
|
||||
|
||||
requeued_count = 0
|
||||
completed_tasks = job_data.completed_tasks
|
||||
batched_static = bool(job_data.batched_static)
|
||||
num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
|
||||
batch_size = int(job_data.batch_size or 1)
|
||||
for worker, assigned_snapshot in workers_to_requeue:
|
||||
# Use current assignments if present, falling back to the snapshot.
|
||||
assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot))
|
||||
for task_id in assigned_tasks:
|
||||
# If batched_static, task_id is a tile_idx; consider it complete only if
|
||||
# all corresponding global_idx entries are present in completed_tasks.
|
||||
if batched_static:
|
||||
all_done = True
|
||||
for b in range(batch_size):
|
||||
gidx = b * num_tiles_per_image + task_id
|
||||
if gidx not in completed_tasks:
|
||||
all_done = False
|
||||
break
|
||||
if not all_done:
|
||||
await job_data.pending_tasks.put(task_id)
|
||||
requeued_count += 1
|
||||
else:
|
||||
if task_id not in completed_tasks:
|
||||
await job_data.pending_tasks.put(task_id)
|
||||
requeued_count += 1
|
||||
job_data.worker_status.pop(worker, None)
|
||||
if worker in job_data.assigned_to_workers:
|
||||
job_data.assigned_to_workers[worker] = []
|
||||
|
||||
return requeued_count
|
||||
@@ -1,313 +0,0 @@
|
||||
import asyncio, torch
|
||||
from PIL import Image
|
||||
import comfy.model_management
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ...utils.async_helpers import run_async_in_server_loop
|
||||
from ...utils.config import get_worker_timeout_seconds
|
||||
from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
|
||||
from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
|
||||
|
||||
|
||||
class DynamicModeMixin:
|
||||
"""
|
||||
Dynamic (per-image queue) USDU mode behaviors for master and worker roles.
|
||||
|
||||
Expected co-mixins on `self`:
|
||||
- TileOpsMixin (`calculate_tiles`, `_slice_conditioning`, `_process_and_blend_tile`).
|
||||
- JobStateMixin (image queue/task completion helpers).
|
||||
- WorkerCommsMixin (`_request_image_from_master`, `_send_full_image_to_master`, `_send_heartbeat_to_master`).
|
||||
"""
|
||||
|
||||
def process_master_dynamic(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers):
|
||||
"""Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
|
||||
# Get batch size and dimensions
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")
|
||||
|
||||
# No fixed share - all images are dynamic
|
||||
all_indices = list(range(batch_size))
|
||||
|
||||
debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
|
||||
|
||||
# Calculate tiles for processing
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
# Initialize job queue for communication
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
|
||||
timeout=2.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
|
||||
raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
|
||||
|
||||
# Convert batch to PIL list
|
||||
result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
|
||||
|
||||
# Process images dynamically with master participating
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
processed_count = 0
|
||||
consecutive_retries = 0
|
||||
max_consecutive_retries = 10
|
||||
|
||||
# Process loop - master pulls from queue and processes synchronously
|
||||
while processed_count < batch_size:
|
||||
# Try to get an image to process
|
||||
image_idx = run_async_in_server_loop(
|
||||
self._get_next_image_index(multi_job_id),
|
||||
timeout=5.0 # Short timeout to allow frequent checks
|
||||
)
|
||||
|
||||
if image_idx is not None:
|
||||
# Reset retry counter and process locally
|
||||
consecutive_retries = 0
|
||||
debug_log(f"Master processing image {image_idx} dynamically")
|
||||
processed_count += 1
|
||||
|
||||
# Process locally
|
||||
single_tensor = upscaled_image[image_idx:image_idx+1]
|
||||
local_image = result_images[image_idx]
|
||||
image_seed = seed
|
||||
|
||||
# Pre-slice conditioning once per image (not per tile)
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
|
||||
|
||||
for tile_idx, pos in enumerate(all_tiles):
|
||||
source_tensor = pil_to_tensor(local_image)
|
||||
if single_tensor.is_cuda:
|
||||
source_tensor = source_tensor.cuda()
|
||||
local_image = self._process_and_blend_tile(
|
||||
tile_idx, pos, source_tensor, local_image,
|
||||
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, width, height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx=image_idx
|
||||
)
|
||||
|
||||
# Yield after each tile to minimize worker downtime
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1)
|
||||
# Note: No per-tile drain here – that's what makes this "per-image"
|
||||
|
||||
result_images[image_idx] = local_image
|
||||
|
||||
# Mark as completed
|
||||
run_async_in_server_loop(
|
||||
self._mark_image_completed(multi_job_id, image_idx, local_image),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
|
||||
drained_count = run_async_in_server_loop(
|
||||
self._drain_worker_results_queue(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
if drained_count > 0:
|
||||
debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
|
||||
|
||||
# NEW: Log overall progress (includes master's image + any drained workers)
|
||||
completed_now = run_async_in_server_loop(
|
||||
self._get_total_completed_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
|
||||
|
||||
# Yield to allow workers to get new images after completing one
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1)
|
||||
else:
|
||||
# Queue empty: collect any queued worker results to update progress
|
||||
drained_count = run_async_in_server_loop(
|
||||
self._drain_worker_results_queue(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain
|
||||
|
||||
# Check for timed out workers and requeue their images
|
||||
requeued_count = run_async_in_server_loop(
|
||||
self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
|
||||
timeout=5.0
|
||||
)
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue
|
||||
|
||||
if requeued_count > 0:
|
||||
log(f"Requeued {requeued_count} images from timed out workers")
|
||||
consecutive_retries = 0 # Reset since we have work to do
|
||||
continue
|
||||
|
||||
# Now check total completed (includes newly collected)
|
||||
completed_now = run_async_in_server_loop(
|
||||
self._get_total_completed_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
|
||||
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
|
||||
|
||||
if completed_now >= batch_size:
|
||||
break
|
||||
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check
|
||||
|
||||
# Check if there are pending images in the queue (could be requeued)
|
||||
pending_count = run_async_in_server_loop(
|
||||
self._get_pending_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
|
||||
if pending_count > 0:
|
||||
consecutive_retries = 0 # Reset retries since there's work to do
|
||||
continue
|
||||
|
||||
consecutive_retries += 1
|
||||
if consecutive_retries >= max_consecutive_retries:
|
||||
log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
|
||||
break # Force exit to collection phase
|
||||
|
||||
debug_log("Waiting for workers")
|
||||
# Use async sleep to allow event loop to process worker requests
|
||||
run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
|
||||
|
||||
debug_log(f"Master processed {processed_count} images locally")
|
||||
|
||||
# Get all completed images to check what needs to be collected
|
||||
all_completed = run_async_in_server_loop(
|
||||
self._get_all_completed_images(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Calculate how many we still need to collect
|
||||
remaining_to_collect = batch_size - len(all_completed)
|
||||
|
||||
if remaining_to_collect > 0:
|
||||
debug_log(f"Waiting for {remaining_to_collect} more images from workers")
|
||||
# Use the unified worker timeout for the collection phase
|
||||
collection_timeout = float(get_worker_timeout_seconds())
|
||||
collected_images = run_async_in_server_loop(
|
||||
self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
|
||||
timeout=collection_timeout
|
||||
)
|
||||
|
||||
# Merge collected with already completed
|
||||
all_completed.update(collected_images)
|
||||
|
||||
# Update result images with all completed images
|
||||
for idx, processed_img in all_completed.items():
|
||||
if idx < batch_size:
|
||||
result_images[idx] = processed_img
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0])
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
|
||||
log(f"Completed processing all {batch_size} images")
|
||||
return (result_tensor,)
|
||||
|
||||
def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold):
|
||||
"""Worker processing in dynamic mode - processes whole images."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get dimensions and tile grid
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")
|
||||
|
||||
# Keep track of processed images for is_last detection
|
||||
processed_count = 0
|
||||
|
||||
# Poll for job readiness to avoid races during master init
|
||||
max_poll_attempts = 20 # ~20s at 1s sleep
|
||||
if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
|
||||
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
|
||||
return (upscaled_image,)
|
||||
|
||||
# Loop to request and process images
|
||||
while True:
|
||||
# Request an image to process
|
||||
image_idx, estimated_remaining = run_async_in_server_loop(
|
||||
self._request_image_from_master(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_WAIT_TIMEOUT
|
||||
)
|
||||
|
||||
if image_idx is None:
|
||||
debug_log(f"USDU Dist Worker - No more images to process")
|
||||
break
|
||||
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
|
||||
processed_count += 1
|
||||
|
||||
# Determine if this should be marked as last for this worker
|
||||
is_last_for_worker = (estimated_remaining == 0)
|
||||
|
||||
# Extract single image tensor
|
||||
single_tensor = upscaled_image[image_idx:image_idx+1]
|
||||
|
||||
# Convert to PIL for processing
|
||||
local_image = tensor_to_pil(single_tensor, 0).copy()
|
||||
|
||||
# Process all tiles for this image
|
||||
image_seed = seed
|
||||
|
||||
# Pre-slice conditioning once per image (not per tile)
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
|
||||
|
||||
for tile_idx, pos in enumerate(all_tiles):
|
||||
source_tensor = pil_to_tensor(local_image)
|
||||
if single_tensor.is_cuda:
|
||||
source_tensor = source_tensor.cuda()
|
||||
local_image = self._process_and_blend_tile(
|
||||
tile_idx, pos, source_tensor, local_image,
|
||||
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, width, height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx=image_idx
|
||||
)
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Send processed image back to master
|
||||
try:
|
||||
# Use the estimated remaining to determine if this is the last image
|
||||
is_last = is_last_for_worker
|
||||
run_async_in_server_loop(
|
||||
self._send_full_image_to_master(local_image, image_idx, multi_job_id,
|
||||
master_url, worker_id, is_last),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
# Send heartbeat after processing
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
if is_last:
|
||||
break
|
||||
except Exception as e:
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
|
||||
# Continue processing other images
|
||||
|
||||
# Send final is_last signal
|
||||
debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
self._send_worker_complete_signal(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
except Exception as e:
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")
|
||||
|
||||
return (upscaled_image,)
|
||||
@@ -1,72 +0,0 @@
|
||||
import math, torch
|
||||
from PIL import Image
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
|
||||
|
||||
class SingleGpuModeMixin:
|
||||
def process_single_gpu(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode):
|
||||
"""Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get image dimensions and batch size
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
|
||||
# Calculate all tiles
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
rows = math.ceil(height / tile_height)
|
||||
cols = math.ceil(width / tile_width)
|
||||
log(
|
||||
f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
|
||||
)
|
||||
|
||||
# Prepare result images list
|
||||
result_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
|
||||
result_images.append(image_pil.copy())
|
||||
|
||||
# Precompute tile masks once
|
||||
tile_masks = []
|
||||
for tx, ty in all_tiles:
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
# Process tiles batched across images
|
||||
for tile_idx, (tx, ty) in enumerate(all_tiles):
|
||||
# Progressive state parity: extract each tile from the current updated image batch.
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
|
||||
# Extract batched tile
|
||||
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
|
||||
source_batch, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
# Process batch
|
||||
region = (x1, y1, x1 + ew, y1 + eh)
|
||||
processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tiled_decode, region, (width, height))
|
||||
|
||||
# Blend results back into each image using cached mask
|
||||
tile_mask = tile_masks[tile_idx]
|
||||
for b in range(batch_size):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
# Resize back to extracted size
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensors = [pil_to_tensor(img) for img in result_images]
|
||||
result_tensor = torch.cat(result_tensors, dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
return (result_tensor,)
|
||||
@@ -1,570 +0,0 @@
|
||||
import asyncio, time, torch
|
||||
from PIL import Image
|
||||
import comfy.model_management
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ...utils.async_helpers import run_async_in_server_loop
|
||||
from ...utils.config import get_worker_timeout_seconds
|
||||
from ...utils.constants import (
|
||||
HEARTBEAT_INTERVAL,
|
||||
JOB_POLL_INTERVAL,
|
||||
JOB_POLL_MAX_ATTEMPTS,
|
||||
MAX_BATCH,
|
||||
TILE_SEND_TIMEOUT,
|
||||
TILE_WAIT_TIMEOUT,
|
||||
)
|
||||
from ..job_store import (
|
||||
ensure_tile_jobs_initialized, init_static_job_batched,
|
||||
_mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count,
|
||||
)
|
||||
from ..job_models import TileJobState
|
||||
|
||||
|
||||
class StaticModeMixin:
|
||||
"""
|
||||
Static (tile-queue) USDU mode behaviors for master and worker roles.
|
||||
|
||||
Expected co-mixins on `self`:
|
||||
- TileOpsMixin (`calculate_tiles`, tile extract/blend helpers).
|
||||
- JobStateMixin (`_get_next_tile_index`, `_get_all_completed_tasks`, requeue checks).
|
||||
- WorkerCommsMixin (`send_tiles_batch_to_master`, `_request_tile_from_master`, `_send_heartbeat_to_master`).
|
||||
"""
|
||||
|
||||
def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS):
|
||||
"""Poll master for job readiness to avoid worker/master initialization race."""
|
||||
for attempt in range(max_attempts):
|
||||
ready = run_async_in_server_loop(
|
||||
self._check_job_status(multi_job_id, master_url),
|
||||
timeout=5.0
|
||||
)
|
||||
if ready:
|
||||
if worker_id:
|
||||
debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
|
||||
else:
|
||||
debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
|
||||
return True
|
||||
time.sleep(JOB_POLL_INTERVAL)
|
||||
return False
|
||||
|
||||
def _extract_and_process_tile(
|
||||
self,
|
||||
upscaled_image,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
):
|
||||
"""Extract one tile position for the whole batch and process it."""
|
||||
tx, ty = all_tiles[tile_id]
|
||||
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
|
||||
upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
region = (x1, y1, x1 + ew, y1 + eh)
|
||||
processed_batch = self.process_tiles_batch(
|
||||
tile_batch, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
|
||||
region, (width, height)
|
||||
)
|
||||
return processed_batch, x1, y1, ew, eh
|
||||
|
||||
def _flush_tiles_to_master(
|
||||
self,
|
||||
processed_tiles,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=False,
|
||||
):
|
||||
"""Send accumulated tile payloads to master and return a fresh accumulator."""
|
||||
if not processed_tiles:
|
||||
if is_final_flush:
|
||||
run_async_in_server_loop(
|
||||
self.send_tiles_batch_to_master(
|
||||
[],
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=True,
|
||||
),
|
||||
timeout=TILE_SEND_TIMEOUT,
|
||||
)
|
||||
return processed_tiles
|
||||
run_async_in_server_loop(
|
||||
self.send_tiles_batch_to_master(
|
||||
processed_tiles,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=is_final_flush,
|
||||
),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
return []
|
||||
|
||||
def _master_process_one_tile(
|
||||
self,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
):
|
||||
"""Process one tile_id across the batch and blend into result_images."""
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
|
||||
source_batch,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
tile_mask = tile_masks[tile_id]
|
||||
out_bs = processed_batch.shape[0] if hasattr(processed_batch, "shape") else batch_size
|
||||
processed_items = min(batch_size, out_bs)
|
||||
for b in range(processed_items):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
|
||||
global_idx = b * num_tiles_per_image + tile_id
|
||||
run_async_in_server_loop(
|
||||
_mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
|
||||
timeout=5.0
|
||||
)
|
||||
return processed_items
|
||||
|
||||
def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_workers):
|
||||
"""Worker static mode processing with optional dynamic queue pulling."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get dimensions and calculate tiles
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
batch_size = upscaled_image.shape[0]
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
|
||||
processed_tiles = []
|
||||
working_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
|
||||
working_images.append(image_pil.copy())
|
||||
tile_masks = []
|
||||
for tx, ty in all_tiles:
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
# Dynamic queue mode (static processing): process batched-per-tile
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}")
|
||||
processed_count = 0
|
||||
|
||||
max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
|
||||
if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
|
||||
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
|
||||
return (upscaled_image,)
|
||||
|
||||
# Main processing loop - pull tile ids from queue
|
||||
while True:
|
||||
# Request a tile to process
|
||||
tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
|
||||
self._request_tile_from_master(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_WAIT_TIMEOUT
|
||||
)
|
||||
|
||||
if tile_idx is None:
|
||||
debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
|
||||
break
|
||||
|
||||
# Always batched-per-tile in static mode
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}")
|
||||
processed_count += batch_size
|
||||
tile_id = tile_idx
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in working_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
|
||||
source_batch,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
# Queue results
|
||||
for b in range(batch_size):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
working_images[b] = self.blend_tile(
|
||||
working_images[b],
|
||||
tile_pil,
|
||||
x1,
|
||||
y1,
|
||||
(ew, eh),
|
||||
tile_masks[tile_id],
|
||||
padding,
|
||||
)
|
||||
processed_tiles.append({
|
||||
'tile': processed_batch[b:b+1],
|
||||
'tile_idx': tile_id,
|
||||
'x': x1,
|
||||
'y': y1,
|
||||
'extracted_width': ew,
|
||||
'extracted_height': eh,
|
||||
'padding': padding,
|
||||
'batch_idx': b,
|
||||
'global_idx': b * num_tiles_per_image + tile_id
|
||||
})
|
||||
|
||||
# Send heartbeat
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
|
||||
|
||||
# Send tiles in batches within loop
|
||||
if len(processed_tiles) >= MAX_BATCH:
|
||||
processed_tiles = self._flush_tiles_to_master(
|
||||
processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False
|
||||
)
|
||||
|
||||
# Send any remaining tiles
|
||||
processed_tiles = self._flush_tiles_to_master(
|
||||
processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=True
|
||||
)
|
||||
|
||||
debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
|
||||
return (upscaled_image,)
|
||||
|
||||
async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total):
|
||||
"""Async helper for collection and monitoring in static mode.
|
||||
Returns collected tasks dict. Caller should check if all tasks are complete."""
|
||||
last_progress_log = time.time()
|
||||
progress_interval = 5.0
|
||||
last_heartbeat_check = time.time()
|
||||
last_completed_count = 0
|
||||
|
||||
while True:
|
||||
# Check for user interruption
|
||||
if comfy.model_management.processing_interrupted():
|
||||
log("Processing interrupted by user")
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
|
||||
# Drain any pending results
|
||||
collected_count = await _drain_results_queue(multi_job_id)
|
||||
|
||||
# Check and requeue timed-out workers periodically
|
||||
current_time = time.time()
|
||||
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
|
||||
requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
|
||||
if requeued_count > 0:
|
||||
log(f"Requeued {requeued_count} tasks from timed-out workers")
|
||||
last_heartbeat_check = current_time
|
||||
|
||||
# Get current completion count
|
||||
completed_count = await _get_completed_count(multi_job_id)
|
||||
|
||||
# Progress logging
|
||||
if current_time - last_progress_log >= progress_interval:
|
||||
log(f"Progress: {completed_count}/{expected_total} tasks completed")
|
||||
last_progress_log = current_time
|
||||
|
||||
# Check if all tasks are completed
|
||||
if completed_count >= expected_total:
|
||||
debug_log(f"All {expected_total} tasks completed")
|
||||
break
|
||||
|
||||
# If no active workers remain and there are pending tasks, return for local processing
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
active_workers = list(job_data.worker_status.keys())
|
||||
if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
|
||||
log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
|
||||
break
|
||||
|
||||
# Wait a bit before next check
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Get all completed tasks for return
|
||||
return await self._get_all_completed_tasks(multi_job_id)
|
||||
|
||||
def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
|
||||
all_tiles, num_tiles_per_image):
|
||||
"""Static mode master processing with optional dynamic queue pulling."""
|
||||
batch_size = upscaled_image.shape[0]
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
|
||||
# Convert batch to PIL list for processing
|
||||
result_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
|
||||
result_images.append(image_pil.copy())
|
||||
|
||||
# Initialize queue: pending queue holds tile ids (batched per tile)
|
||||
log("USDU Dist: Using tile queue distribution")
|
||||
run_async_in_server_loop(
|
||||
init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
|
||||
timeout=10.0
|
||||
)
|
||||
debug_log(
|
||||
f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
|
||||
)
|
||||
|
||||
# Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending
|
||||
tile_masks = []
|
||||
for idx, (tx, ty) in enumerate(all_tiles):
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
processed_count = 0
|
||||
consecutive_no_tile = 0
|
||||
max_consecutive_no_tile = 2
|
||||
|
||||
while processed_count < total_tiles:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
tile_idx = run_async_in_server_loop(
|
||||
self._get_next_tile_index(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
if tile_idx is not None:
|
||||
consecutive_no_tile = 0
|
||||
tile_id = tile_idx
|
||||
processed_count += self._master_process_one_tile(
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
|
||||
else:
|
||||
consecutive_no_tile += 1
|
||||
if consecutive_no_tile >= max_consecutive_no_tile:
|
||||
debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
|
||||
break
|
||||
time.sleep(0.1)
|
||||
master_processed_count = processed_count
|
||||
|
||||
# Continue processing any remaining tiles while collecting worker results
|
||||
remaining_tiles = total_tiles - master_processed_count
|
||||
if remaining_tiles > 0:
|
||||
debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
|
||||
|
||||
# Collect worker results using async operations
|
||||
try:
|
||||
# Wait until either all tasks are collected or there are no active workers left
|
||||
collected_tasks = run_async_in_server_loop(
|
||||
self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
|
||||
timeout=None
|
||||
)
|
||||
except comfy.model_management.InterruptProcessingException:
|
||||
# Clean up job on interruption
|
||||
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
|
||||
raise
|
||||
|
||||
# Check if we need to process any remaining tasks locally after collection
|
||||
completed_count = len(collected_tasks)
|
||||
if completed_count < total_tiles:
|
||||
log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
|
||||
|
||||
# Process any remaining pending tasks (batched-per-tile)
|
||||
while True:
|
||||
# Check for user interruption
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
# Get next tile_id from pending queue
|
||||
tile_id = run_async_in_server_loop(
|
||||
self._get_next_tile_index(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
if tile_id is None:
|
||||
break
|
||||
|
||||
self._master_process_one_tile(
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
else:
|
||||
# Master processed all tiles
|
||||
collected_tasks = run_async_in_server_loop(
|
||||
self._get_all_completed_tasks(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Blend worker tiles synchronously in deterministic tile order.
|
||||
def _sort_key(item):
|
||||
global_idx, tile_data = item
|
||||
batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
|
||||
tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
|
||||
return (tile_idx, batch_idx, global_idx)
|
||||
|
||||
for global_idx, tile_data in sorted(collected_tasks.items(), key=_sort_key):
|
||||
# Skip tiles that don't have tensor data (already processed)
|
||||
if 'tensor' not in tile_data and 'image' not in tile_data:
|
||||
continue
|
||||
|
||||
batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
|
||||
tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
|
||||
|
||||
if batch_idx >= batch_size:
|
||||
continue
|
||||
|
||||
# Blend tile synchronously
|
||||
x = tile_data.get('x', 0)
|
||||
y = tile_data.get('y', 0)
|
||||
# Prefer PIL image if present to avoid reconversion
|
||||
if 'image' in tile_data:
|
||||
tile_pil = tile_data['image']
|
||||
else:
|
||||
tile_tensor = tile_data['tensor']
|
||||
tile_pil = tensor_to_pil(tile_tensor, 0)
|
||||
orig_x, orig_y = all_tiles[tile_idx]
|
||||
tile_mask = tile_masks[tile_idx]
|
||||
extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
|
||||
extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
|
||||
result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil,
|
||||
x, y, (extracted_width, extracted_height), tile_mask, padding)
|
||||
|
||||
try:
|
||||
# Convert back to tensor
|
||||
if batch_size == 1:
|
||||
result_tensor = pil_to_tensor(result_images[0])
|
||||
else:
|
||||
result_tensors = [pil_to_tensor(img) for img in result_images]
|
||||
result_tensor = torch.cat(result_tensors, dim=0)
|
||||
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
|
||||
return (result_tensor,)
|
||||
finally:
|
||||
# Cleanup (async operation) - always execute
|
||||
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
|
||||
@@ -1,64 +0,0 @@
|
||||
import io
|
||||
import json
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _parse_tiles_from_form(data):
|
||||
"""Parse tiles submitted via multipart/form-data into a list of tile dicts."""
|
||||
try:
|
||||
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
|
||||
except Exception:
|
||||
padding = 0
|
||||
|
||||
meta_raw = data.get('tiles_metadata')
|
||||
if meta_raw is None:
|
||||
raise ValueError("Missing tiles_metadata")
|
||||
|
||||
try:
|
||||
metadata = json.loads(meta_raw)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
|
||||
|
||||
if not isinstance(metadata, list):
|
||||
raise ValueError("tiles_metadata must be a list")
|
||||
|
||||
tiles = []
|
||||
for i, meta in enumerate(metadata):
|
||||
file_field = data.get(f'tile_{i}')
|
||||
if file_field is None or not hasattr(file_field, 'file'):
|
||||
raise ValueError(f"Missing tile data for index {i}")
|
||||
|
||||
raw = file_field.file.read()
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid image data for tile {i}: {e}")
|
||||
|
||||
try:
|
||||
tile_info = {
|
||||
'image': img,
|
||||
'tile_idx': int(meta.get('tile_idx', i)),
|
||||
'x': int(meta.get('x', 0)),
|
||||
'y': int(meta.get('y', 0)),
|
||||
'extracted_width': int(meta.get('extracted_width', img.width)),
|
||||
'extracted_height': int(meta.get('extracted_height', img.height)),
|
||||
'padding': int(padding),
|
||||
}
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
|
||||
|
||||
if 'batch_idx' in meta:
|
||||
try:
|
||||
tile_info['batch_idx'] = int(meta['batch_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
if 'global_idx' in meta:
|
||||
try:
|
||||
tile_info['global_idx'] = int(meta['global_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tiles.append(tile_info)
|
||||
|
||||
return tiles
|
||||
@@ -1,198 +0,0 @@
|
||||
import asyncio, time
|
||||
import comfy.model_management
|
||||
import server
|
||||
from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.config import get_worker_timeout_seconds
|
||||
from .job_store import ensure_tile_jobs_initialized, _mark_task_completed
|
||||
from .job_timeout import _check_and_requeue_timed_out_workers
|
||||
from .job_models import BaseJobState, ImageJobState, TileJobState
|
||||
|
||||
|
||||
class ResultCollectorMixin:
|
||||
"""
|
||||
Mixin for master-side result collection in USDU distributed jobs.
|
||||
|
||||
Expected co-mixins/attributes:
|
||||
- JobStateMixin methods for queue/task access.
|
||||
- `self._check_and_requeue_timed_out_workers(...)` coroutine.
|
||||
- `self._async_yield(...)` optional helper from WorkerCommsMixin.
|
||||
"""
|
||||
|
||||
def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
|
||||
"""Log timeout elapsed seconds for each tracked worker and return worker ids."""
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return []
|
||||
|
||||
worker_status = dict(job_data.worker_status)
|
||||
for worker_id, last_seen in worker_status.items():
|
||||
elapsed = max(0.0, current_time - float(last_seen))
|
||||
log(
|
||||
"UltimateSDUpscale Master - Heartbeat timeout: "
|
||||
f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
return list(worker_status.keys())
|
||||
|
||||
async def _async_collect_results(self, multi_job_id, num_workers, mode='static',
|
||||
remaining_to_collect=None, batch_size=None):
|
||||
"""Unified async helper to collect results from workers (tiles or images)."""
|
||||
# Get the already initialized queue
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
|
||||
raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if mode == 'dynamic':
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
raise RuntimeError(
|
||||
f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
|
||||
)
|
||||
q = job_data.queue
|
||||
completed_images = job_data.completed_images
|
||||
expected_count = remaining_to_collect or batch_size
|
||||
elif mode == 'static':
|
||||
if not isinstance(job_data, TileJobState):
|
||||
raise RuntimeError(
|
||||
f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
|
||||
)
|
||||
q = job_data.queue
|
||||
expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported mode: {mode}")
|
||||
|
||||
item_type = "images" if mode == 'dynamic' else "tiles"
|
||||
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
|
||||
|
||||
collected_results = {}
|
||||
workers_done = set()
|
||||
# Unify collector/upscaler wait behavior with the UI worker timeout
|
||||
timeout = float(get_worker_timeout_seconds())
|
||||
last_heartbeat_check = time.time()
|
||||
wait_started_at = time.time()
|
||||
collected_count = 0
|
||||
|
||||
while len(workers_done) < num_workers:
|
||||
# Check for user interruption
|
||||
if comfy.model_management.processing_interrupted():
|
||||
log("Processing interrupted by user")
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
|
||||
# For dynamic mode with remaining_to_collect, check if we've collected enough
|
||||
if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
|
||||
break
|
||||
|
||||
job_data_snapshot = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(current_job_data, BaseJobState):
|
||||
job_data_snapshot = current_job_data
|
||||
|
||||
try:
|
||||
# Shorter poll for dynamic mode, but never exceed the configured timeout
|
||||
wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout)
|
||||
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if mode == 'static':
|
||||
# Handle tiles
|
||||
tiles = result.get('tiles', [])
|
||||
debug_log(
|
||||
f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
|
||||
f"'{worker_id}' (is_last={is_last})"
|
||||
)
|
||||
|
||||
for tile_data in tiles:
|
||||
if 'batch_idx' not in tile_data:
|
||||
log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
|
||||
continue
|
||||
|
||||
tile_idx = tile_data['tile_idx']
|
||||
key = tile_data.get('global_idx', tile_idx)
|
||||
entry = {
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_data['x'],
|
||||
'y': tile_data['y'],
|
||||
'extracted_width': tile_data['extracted_width'],
|
||||
'extracted_height': tile_data['extracted_height'],
|
||||
'padding': tile_data['padding'],
|
||||
'worker_id': worker_id,
|
||||
'batch_idx': tile_data.get('batch_idx', 0),
|
||||
'global_idx': tile_data.get('global_idx', tile_idx),
|
||||
}
|
||||
if 'image' in tile_data:
|
||||
entry['image'] = tile_data['image']
|
||||
elif 'tensor' in tile_data:
|
||||
entry['tensor'] = tile_data['tensor']
|
||||
collected_results[key] = entry
|
||||
|
||||
elif mode == 'dynamic':
|
||||
# Handle full images
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
image_idx = result['image_idx']
|
||||
image_pil = result['image']
|
||||
completed_images[image_idx] = image_pil
|
||||
collected_results[image_idx] = image_pil
|
||||
collected_count += 1
|
||||
debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
|
||||
|
||||
if is_last:
|
||||
workers_done.add(worker_id)
|
||||
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
current_time = time.time()
|
||||
waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
|
||||
if mode == 'dynamic':
|
||||
# Check for worker timeouts periodically
|
||||
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
|
||||
# Use the class method to check and requeue
|
||||
requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)
|
||||
if requeued > 0:
|
||||
log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
|
||||
last_heartbeat_check = current_time
|
||||
|
||||
# Check if we've been waiting too long overall
|
||||
if current_time - wait_started_at > timeout:
|
||||
elapsed = current_time - wait_started_at
|
||||
log(
|
||||
"UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
|
||||
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
break
|
||||
else:
|
||||
elapsed = current_time - wait_started_at
|
||||
log(
|
||||
f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; "
|
||||
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
break
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
|
||||
|
||||
# Clean up job queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
|
||||
return collected_results if mode == 'static' else completed_images
|
||||
|
||||
async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
|
||||
"""Async helper to collect tiles from workers."""
|
||||
return await self._async_collect_results(multi_job_id, num_workers, mode='static')
|
||||
|
||||
async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
|
||||
"""Mark an image as completed in the job data."""
|
||||
# Mark the image as completed with the image data
|
||||
await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
job_data.completed_images[image_idx] = image_pil
|
||||
|
||||
async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
|
||||
"""Collect remaining processed images from workers."""
|
||||
return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic',
|
||||
remaining_to_collect=remaining_to_collect,
|
||||
batch_size=batch_size)
|
||||
@@ -1,446 +0,0 @@
|
||||
import math, torch
|
||||
from contextlib import nullcontext
|
||||
from PIL import Image, ImageFilter, ImageDraw
|
||||
from typing import List, Tuple
|
||||
import comfy.samplers, comfy.model_management
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop
|
||||
from ..utils.crop_model_patch import crop_model_cond
|
||||
from .conditioning import clone_conditioning
|
||||
|
||||
|
||||
class TileOpsMixin:
|
||||
def round_to_multiple(self, value: int, multiple: int = 8) -> int:
|
||||
"""Round value to nearest multiple."""
|
||||
return round(value / multiple) * multiple
|
||||
|
||||
def calculate_tiles(self, image_width: int, image_height: int,
|
||||
tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]:
|
||||
"""Calculate tile positions to match Ultimate SD Upscale.
|
||||
|
||||
Positions are a simple grid starting at (0,0) with steps of
|
||||
`tile_width` and `tile_height`, using ceil(rows/cols) to cover edges.
|
||||
Uniform vs non-uniform affects only crop/resize, not positions.
|
||||
"""
|
||||
rows = math.ceil(image_height / tile_height)
|
||||
cols = math.ceil(image_width / tile_width)
|
||||
tiles: List[Tuple[int, int]] = []
|
||||
for yi in range(rows):
|
||||
for xi in range(cols):
|
||||
tiles.append((xi * tile_width, yi * tile_height))
|
||||
return tiles
|
||||
|
||||
def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
|
||||
tile_width: int, tile_height: int, padding: int,
|
||||
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
|
||||
"""Extract a tile region and resize to match USDU cropping logic.
|
||||
|
||||
Mirrors ComfyUI_UltimateSDUpscale processing:
|
||||
- Build a mask with a white rectangle at the tile rect
|
||||
- Compute crop_region via get_crop_region(mask, padding)
|
||||
- If force_uniform_tiles: expand by crop/aspect ratio, then resize to
|
||||
fixed processing size of round_to_multiple(tile + padding)
|
||||
- Else: target is ceil(crop_size/8)*8 per dimension
|
||||
- Extract the crop and resize to target tile_size
|
||||
Returns the resized tensor and crop origin/size for blending.
|
||||
"""
|
||||
_, h, w, _ = image.shape
|
||||
|
||||
# Create mask and compute initial padded crop region
|
||||
mask = Image.new('L', (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
x1, y1, x2, y2 = get_crop_region(mask, padding)
|
||||
|
||||
# Determine crop + processing size
|
||||
if force_uniform_tiles:
|
||||
process_w = self.round_to_multiple(tile_width + padding, 8)
|
||||
process_h = self.round_to_multiple(tile_height + padding, 8)
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
|
||||
process_ratio = process_w / process_h if process_h != 0 else 1.0
|
||||
if crop_ratio > process_ratio:
|
||||
target_w = crop_w
|
||||
target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
|
||||
else:
|
||||
target_w = round(crop_h * process_ratio)
|
||||
target_h = crop_h
|
||||
(x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
target_w = process_w
|
||||
target_h = process_h
|
||||
else:
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
target_w = max(8, math.ceil(crop_w / 8) * 8)
|
||||
target_h = max(8, math.ceil(crop_h / 8) * 8)
|
||||
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
|
||||
# Actual extracted size before resizing (for blending)
|
||||
extracted_width = x2 - x1
|
||||
extracted_height = y2 - y1
|
||||
|
||||
# Extract tile and resize to processing size
|
||||
tile = image[:, y1:y2, x1:x2, :]
|
||||
tile_pil = tensor_to_pil(tile, 0)
|
||||
if tile_pil.size != (target_w, target_h):
|
||||
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
|
||||
|
||||
tile_tensor = pil_to_tensor(tile_pil)
|
||||
if image.is_cuda:
|
||||
tile_tensor = tile_tensor.cuda()
|
||||
|
||||
return tile_tensor, x1, y1, extracted_width, extracted_height
|
||||
|
||||
def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int,
|
||||
tile_width: int, tile_height: int, padding: int,
|
||||
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
|
||||
"""Extract a tile region for the entire batch and resize to USDU logic.
|
||||
|
||||
- Computes a single crop region from a mask at (x,y,w,h) with padding
|
||||
- force_uniform_tiles controls target processing size logic
|
||||
- Returns a batched tensor [B,H',W',C] and crop origin/size for blending
|
||||
"""
|
||||
batch, h, w, _ = images.shape
|
||||
|
||||
# Create mask and compute initial padded crop region (same for all images)
|
||||
mask = Image.new('L', (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
x1, y1, x2, y2 = get_crop_region(mask, padding)
|
||||
|
||||
# Determine crop + processing size
|
||||
if force_uniform_tiles:
|
||||
process_w = self.round_to_multiple(tile_width + padding, 8)
|
||||
process_h = self.round_to_multiple(tile_height + padding, 8)
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
|
||||
process_ratio = process_w / process_h if process_h != 0 else 1.0
|
||||
if crop_ratio > process_ratio:
|
||||
target_w = crop_w
|
||||
target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
|
||||
else:
|
||||
target_w = round(crop_h * process_ratio)
|
||||
target_h = crop_h
|
||||
(x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
target_w = process_w
|
||||
target_h = process_h
|
||||
else:
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
target_w = max(8, math.ceil(crop_w / 8) * 8)
|
||||
target_h = max(8, math.ceil(crop_h / 8) * 8)
|
||||
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
|
||||
extracted_width = x2 - x1
|
||||
extracted_height = y2 - y1
|
||||
|
||||
# Slice batch region
|
||||
tiles = images[:, y1:y2, x1:x2, :]
|
||||
|
||||
# Resize each tile to target size
|
||||
resized_tiles = []
|
||||
for i in range(batch):
|
||||
tile_pil = tensor_to_pil(tiles, i)
|
||||
if tile_pil.size != (target_w, target_h):
|
||||
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
|
||||
resized_tiles.append(pil_to_tensor(tile_pil))
|
||||
tile_batch = torch.cat(resized_tiles, dim=0)
|
||||
|
||||
if images.is_cuda:
|
||||
tile_batch = tile_batch.cuda()
|
||||
|
||||
return tile_batch, x1, y1, extracted_width, extracted_height
|
||||
|
||||
def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae,
|
||||
seed: int, steps: int, cfg: float, sampler_name: str,
|
||||
scheduler: str, denoise: float, tiled_decode: bool = False,
|
||||
batch_idx: int = 0, region: Tuple[int, int, int, int] = None,
|
||||
image_size: Tuple[int, int] = None) -> torch.Tensor:
|
||||
"""Process a single tile through SD sampling.
|
||||
Note: positive and negative should already be pre-sliced for the current batch_idx."""
|
||||
debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}")
|
||||
|
||||
|
||||
# Import here to avoid circular dependencies
|
||||
from nodes import common_ksampler, VAEEncode, VAEDecode
|
||||
|
||||
# Try to import tiled VAE nodes if available
|
||||
try:
|
||||
from nodes import VAEEncodeTiled, VAEDecodeTiled
|
||||
tiled_vae_available = True
|
||||
except ImportError:
|
||||
tiled_vae_available = False
|
||||
if tiled_decode:
|
||||
debug_log("Tiled VAE nodes not available, falling back to standard VAE")
|
||||
|
||||
# Convert to PIL and back to ensure clean tensor without gradient tracking
|
||||
tile_pil = tensor_to_pil(tile_tensor, 0)
|
||||
clean_tensor = pil_to_tensor(tile_pil)
|
||||
|
||||
# Ensure tensor is detached and doesn't require gradients
|
||||
clean_tensor = clean_tensor.detach()
|
||||
if hasattr(clean_tensor, 'requires_grad_'):
|
||||
clean_tensor.requires_grad_(False)
|
||||
|
||||
# Move to correct device
|
||||
if tile_tensor.is_cuda:
|
||||
clean_tensor = clean_tensor.cuda()
|
||||
clean_tensor = clean_tensor.detach() # Detach again after device transfer
|
||||
|
||||
# Clone conditioning per tile (shares models, clones hints for cropping)
|
||||
positive_tile = clone_conditioning(positive, clone_hints=True)
|
||||
negative_tile = clone_conditioning(negative, clone_hints=True)
|
||||
|
||||
# Crop conditioning to tile region if provided (assumes hints at image resolution)
|
||||
if region is not None and image_size is not None:
|
||||
init_size = image_size # (width, height) of full image
|
||||
canvas_size = image_size
|
||||
tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height)
|
||||
w_pad = 0 # No extra pad needed; region already includes padding
|
||||
h_pad = 0
|
||||
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
else:
|
||||
# No region cropping needed, use cloned conditioning as-is
|
||||
positive_cropped = positive_tile
|
||||
negative_cropped = negative_tile
|
||||
|
||||
# Encode to latent (always non-tiled, matching original node)
|
||||
latent = VAEEncode().encode(vae, clean_tensor)[0]
|
||||
|
||||
# Sample with model patch cropping parity (ControlNet patch hints)
|
||||
if region is not None and image_size is not None:
|
||||
model_ctx = crop_model_cond(
|
||||
model,
|
||||
region,
|
||||
image_size,
|
||||
image_size,
|
||||
(clean_tensor.shape[2], clean_tensor.shape[1]),
|
||||
)
|
||||
else:
|
||||
model_ctx = nullcontext(model)
|
||||
with model_ctx as model_for_sampling:
|
||||
samples = common_ksampler(
|
||||
model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
|
||||
positive_cropped, negative_cropped, latent, denoise=denoise
|
||||
)[0]
|
||||
|
||||
# Decode back to image
|
||||
if tiled_decode and tiled_vae_available:
|
||||
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
|
||||
else:
|
||||
image = VAEDecode().decode(vae, samples)[0]
|
||||
|
||||
return image
|
||||
|
||||
def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae,
|
||||
seed: int, steps: int, cfg: float, sampler_name: str,
|
||||
scheduler: str, denoise: float, tiled_decode: bool,
|
||||
region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor:
|
||||
"""Process a batch of tiles together (USDU behavior).
|
||||
|
||||
tile_batch: [B, H, W, C]
|
||||
Returns image batch tensor [B, H, W, C]
|
||||
"""
|
||||
# Import locally to avoid circular deps
|
||||
from nodes import common_ksampler, VAEEncode, VAEDecode
|
||||
try:
|
||||
from nodes import VAEEncodeTiled, VAEDecodeTiled
|
||||
tiled_vae_available = True
|
||||
except ImportError:
|
||||
tiled_vae_available = False
|
||||
|
||||
# Detach and move device
|
||||
clean = tile_batch.detach()
|
||||
if hasattr(clean, 'requires_grad_'):
|
||||
clean.requires_grad_(False)
|
||||
if tile_batch.is_cuda:
|
||||
clean = clean.cuda().detach()
|
||||
|
||||
# Clone/crop conditioning once for the region
|
||||
positive_tile = clone_conditioning(positive, clone_hints=True)
|
||||
negative_tile = clone_conditioning(negative, clone_hints=True)
|
||||
|
||||
init_size = image_size
|
||||
canvas_size = image_size
|
||||
tile_size = (clean.shape[2], clean.shape[1]) # (W,H)
|
||||
w_pad = 0
|
||||
h_pad = 0
|
||||
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
|
||||
# Encode -> Sample -> Decode
|
||||
latent = VAEEncode().encode(vae, clean)[0]
|
||||
with crop_model_cond(model, region, image_size, image_size, tile_size) as model_for_sampling:
|
||||
samples = common_ksampler(
|
||||
model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
|
||||
positive_cropped, negative_cropped, latent, denoise=denoise
|
||||
)[0]
|
||||
if tiled_decode and tiled_vae_available:
|
||||
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
|
||||
else:
|
||||
image = VAEDecode().decode(vae, samples)[0]
|
||||
|
||||
return image
|
||||
|
||||
def create_tile_mask(self, image_width: int, image_height: int,
|
||||
x: int, y: int, tile_width: int, tile_height: int,
|
||||
mask_blur: int) -> Image.Image:
|
||||
"""Create a mask for blending tiles - matches Ultimate SD Upscale approach.
|
||||
|
||||
Creates a black image with a white rectangle at the tile position,
|
||||
then applies blur to create soft edges.
|
||||
"""
|
||||
# Create a full-size mask matching the image dimensions
|
||||
mask = Image.new('L', (image_width, image_height), 0) # Black background
|
||||
|
||||
# Draw white rectangle at tile position
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
|
||||
# Apply blur to soften edges
|
||||
if mask_blur > 0:
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
|
||||
|
||||
return mask
|
||||
|
||||
def blend_tile(self, base_image: Image.Image, tile_image: Image.Image,
|
||||
x: int, y: int, extracted_size: Tuple[int, int],
|
||||
mask: Image.Image, padding: int) -> Image.Image:
|
||||
"""Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach.
|
||||
|
||||
This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py
|
||||
"""
|
||||
extracted_width, extracted_height = extracted_size
|
||||
|
||||
# Debug logging (uncomment if needed)
|
||||
# debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}")
|
||||
|
||||
# Calculate the crop region that was used for extraction
|
||||
crop_region = (x, y, x + extracted_width, y + extracted_height)
|
||||
|
||||
# The mask is already full-size, no need to crop
|
||||
|
||||
# Resize the processed tile back to the extracted size
|
||||
if tile_image.size != (extracted_width, extracted_height):
|
||||
tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS)
|
||||
else:
|
||||
tile_resized = tile_image
|
||||
|
||||
# Follow Ultimate SD Upscale blending approach:
|
||||
# Put the tile into position
|
||||
image_tile_only = Image.new('RGBA', base_image.size)
|
||||
image_tile_only.paste(tile_resized, crop_region[:2])
|
||||
|
||||
# Add the mask as an alpha channel
|
||||
# Must make a copy due to the possibility of an edge becoming black
|
||||
temp = image_tile_only.copy()
|
||||
temp.putalpha(mask) # Use the full image mask
|
||||
image_tile_only.paste(temp, image_tile_only)
|
||||
|
||||
# Add back the tile to the initial image according to the mask in the alpha channel
|
||||
result = base_image.convert('RGBA')
|
||||
result.alpha_composite(image_tile_only)
|
||||
|
||||
# Convert back to RGB
|
||||
return result.convert('RGB')
|
||||
|
||||
def _slice_conditioning(self, positive, negative, batch_idx):
|
||||
"""Helper to slice conditioning for a specific batch index."""
|
||||
# Clone and slice conditioning properly, including ControlNet hints
|
||||
positive_sliced = clone_conditioning(positive)
|
||||
negative_sliced = clone_conditioning(negative)
|
||||
|
||||
for cond_list in [positive_sliced, negative_sliced]:
|
||||
for i in range(len(cond_list)):
|
||||
emb, cond_dict = cond_list[i]
|
||||
if emb.shape[0] > 1:
|
||||
cond_list[i][0] = emb[batch_idx:batch_idx+1]
|
||||
if 'control' in cond_dict:
|
||||
control = cond_dict['control']
|
||||
while control is not None:
|
||||
hint = control.cond_hint_original
|
||||
if hint.shape[0] > 1:
|
||||
control.cond_hint_original = hint[batch_idx:batch_idx+1]
|
||||
control = control.previous_controlnet
|
||||
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
|
||||
cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1]
|
||||
|
||||
return positive_sliced, negative_sliced
|
||||
|
||||
def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
|
||||
model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, image_width, image_height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx: int = 0):
|
||||
"""Process a single tile and blend it into the result image."""
|
||||
x, y = tile_pos
|
||||
|
||||
# Extract and process tile
|
||||
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
|
||||
upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, batch_idx=batch_idx,
|
||||
region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height))
|
||||
|
||||
# Convert and blend
|
||||
processed_pil = tensor_to_pil(processed_tile, 0)
|
||||
# Create mask for this specific tile (no cache here; only used in single-tile path)
|
||||
tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur)
|
||||
# Use extraction position and size for blending
|
||||
result_image = self.blend_tile(result_image, processed_pil,
|
||||
x1, y1, (ew, eh), tile_mask, padding)
|
||||
|
||||
return result_image
|
||||
|
||||
def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
|
||||
model, positive, negative, vae, seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
|
||||
width, height, force_uniform_tiles, sliced_conditioning_cache):
|
||||
"""Process a single tile."""
|
||||
# Calculate which image and tile this corresponds to
|
||||
batch_idx = global_idx // num_tiles_per_image
|
||||
tile_idx = global_idx % num_tiles_per_image
|
||||
|
||||
# Skip if batch_idx is out of range
|
||||
if batch_idx >= upscaled_image.shape[0]:
|
||||
debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}")
|
||||
return None
|
||||
|
||||
# Get or create sliced conditioning for this batch index
|
||||
if batch_idx not in sliced_conditioning_cache:
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx)
|
||||
sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced)
|
||||
else:
|
||||
positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx]
|
||||
|
||||
x, y = all_tiles[tile_idx]
|
||||
|
||||
# Extract tile from the specific image in the batch
|
||||
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
|
||||
upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
# Process tile through SD with the exact seed (USDU parity)
|
||||
image_seed = seed
|
||||
processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae,
|
||||
image_seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, batch_idx=batch_idx,
|
||||
region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height))
|
||||
|
||||
return {
|
||||
'tile': processed_tile,
|
||||
'global_idx': global_idx,
|
||||
'batch_idx': batch_idx,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x1,
|
||||
'y': y1,
|
||||
'extracted_width': ew,
|
||||
'extracted_height': eh
|
||||
}
|
||||
@@ -1,262 +0,0 @@
|
||||
import asyncio, io, json, time
|
||||
import aiohttp
|
||||
from PIL import Image
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import get_client_session
|
||||
from ..utils.constants import TILE_SEND_TIMEOUT
|
||||
from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master
|
||||
from ..utils.image import tensor_to_pil
|
||||
|
||||
|
||||
class WorkerCommsMixin:
|
||||
async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Proxy heartbeat helper used by worker processing mixins."""
|
||||
await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
|
||||
|
||||
async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url,
|
||||
padding, worker_id, is_final_flush=False):
|
||||
"""Send all processed tiles to master, chunked if large."""
|
||||
if not processed_tiles:
|
||||
if is_final_flush:
|
||||
await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id)
|
||||
return # Early exit if empty
|
||||
|
||||
total_tiles = len(processed_tiles)
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")
|
||||
|
||||
# Prepare encoded images and sizes to enable size-aware chunking
|
||||
encoded = []
|
||||
for idx, tile_data in enumerate(processed_tiles):
|
||||
img = tensor_to_pil(tile_data['tile'], 0)
|
||||
bio = io.BytesIO()
|
||||
# Keep compression low to balance speed and size; adjust if needed
|
||||
img.save(bio, format='PNG', compress_level=0)
|
||||
raw = bio.getvalue()
|
||||
encoded.append({
|
||||
'bytes': raw,
|
||||
'meta': {
|
||||
'tile_idx': tile_data['tile_idx'],
|
||||
'x': tile_data['x'],
|
||||
'y': tile_data['y'],
|
||||
'extracted_width': tile_data['extracted_width'],
|
||||
'extracted_height': tile_data['extracted_height'],
|
||||
**({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
|
||||
**({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
|
||||
}
|
||||
})
|
||||
|
||||
# Size-aware chunking
|
||||
max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom
|
||||
i = 0
|
||||
chunk_index = 0
|
||||
while i < total_tiles:
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('padding', str(padding))
|
||||
|
||||
metadata = []
|
||||
used = 0
|
||||
j = i
|
||||
while j < total_tiles:
|
||||
img_bytes = encoded[j]['bytes']
|
||||
meta = encoded[j]['meta']
|
||||
# Rough overhead for fields + JSON
|
||||
overhead = 1024
|
||||
if used + len(img_bytes) + overhead > max_bytes and j > i:
|
||||
break
|
||||
# Accept this tile in this chunk
|
||||
metadata.append(meta)
|
||||
data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
|
||||
used += len(img_bytes) + overhead
|
||||
j += 1
|
||||
|
||||
# Ensure at least one tile per chunk
|
||||
if j == i:
|
||||
# Single oversized tile, send anyway
|
||||
meta = encoded[j]['meta']
|
||||
metadata.append(meta)
|
||||
data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
|
||||
j += 1
|
||||
|
||||
chunk_size = j - i
|
||||
is_chunk_last = (j >= total_tiles)
|
||||
data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
|
||||
data.add_field('batch_size', str(chunk_size))
|
||||
data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
|
||||
|
||||
# Retry logic with exponential backoff
|
||||
max_retries = 5
|
||||
retry_delay = 0.5
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_tiles"
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay = min(retry_delay * 2, 5.0)
|
||||
else:
|
||||
log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
|
||||
raise
|
||||
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
|
||||
chunk_index += 1
|
||||
i = j
|
||||
|
||||
async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id):
|
||||
"""Send completion signal to master in static mode when no tiles are left."""
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('is_last', 'true')
|
||||
data.add_field('batch_size', '0')
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_tiles"
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Worker {worker_id} sent static completion signal")
|
||||
|
||||
async def _request_work_item_from_master(
|
||||
self,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
worker_id,
|
||||
endpoint="/distributed/request_image",
|
||||
):
|
||||
"""Request one work item from master with retry/backoff and total timeout."""
|
||||
max_retries = 10
|
||||
retry_delay = 0.5
|
||||
start_time = time.monotonic()
|
||||
url = f"{master_url}{endpoint}"
|
||||
|
||||
for attempt in range(max_retries):
|
||||
if time.monotonic() - start_time > 30:
|
||||
log(f"Total request timeout after 30s for worker {worker_id}")
|
||||
return None
|
||||
|
||||
try:
|
||||
session = await get_client_session()
|
||||
async with session.post(url, json={
|
||||
'worker_id': str(worker_id),
|
||||
'multi_job_id': multi_job_id
|
||||
}) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
if response.status == 404:
|
||||
text = await response.text()
|
||||
debug_log(f"Job not found (404), will retry: {text}")
|
||||
await asyncio.sleep(1.0)
|
||||
else:
|
||||
text = await response.text()
|
||||
debug_log(
|
||||
f"Request work item failed ({response.status}) for worker {worker_id}: {text}"
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
if attempt < max_retries - 1:
|
||||
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}")
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay = min(retry_delay * 2, 5.0)
|
||||
else:
|
||||
log(f"Failed to request work item after {max_retries} attempts: {exc}")
|
||||
raise
|
||||
|
||||
return None
|
||||
|
||||
async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Request an image index to process from master in dynamic mode."""
|
||||
data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
|
||||
if not data:
|
||||
return None, 0
|
||||
image_idx = data.get('image_idx')
|
||||
estimated_remaining = data.get('estimated_remaining', 0)
|
||||
return image_idx, estimated_remaining
|
||||
|
||||
async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
|
||||
data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
|
||||
if not data:
|
||||
return None, 0, False
|
||||
tile_idx = data.get('tile_idx')
|
||||
estimated_remaining = data.get('estimated_remaining', 0)
|
||||
batched_static = data.get('batched_static', False)
|
||||
return tile_idx, estimated_remaining, batched_static
|
||||
|
||||
async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id,
|
||||
master_url, worker_id, is_last):
|
||||
"""Send a processed full image back to master in dynamic mode."""
|
||||
# Serialize image to PNG
|
||||
byte_io = io.BytesIO()
|
||||
image_pil.save(byte_io, format='PNG', compress_level=0)
|
||||
byte_io.seek(0)
|
||||
|
||||
# Prepare form data
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('image_idx', str(image_idx))
|
||||
data.add_field('is_last', str(is_last))
|
||||
data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png',
|
||||
content_type='image/png')
|
||||
|
||||
# Retry logic
|
||||
max_retries = 5
|
||||
retry_delay = 0.5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_image"
|
||||
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Successfully sent image {image_idx} to master")
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay *= 2
|
||||
else:
|
||||
log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
|
||||
raise
|
||||
|
||||
async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
|
||||
"""Send completion signal to master in dynamic mode."""
|
||||
# Send a dummy request with is_last=True
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('is_last', 'true')
|
||||
# No image data - just completion signal
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_image"
|
||||
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Worker {worker_id} sent completion signal")
|
||||
|
||||
async def _check_job_status(self, multi_job_id, master_url):
|
||||
"""Check if job is ready on the master."""
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
|
||||
async with session.get(url) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
return data.get('ready', False)
|
||||
return False
|
||||
except Exception as e:
|
||||
debug_log(f"Job status check failed: {e}")
|
||||
return False
|
||||
|
||||
async def _async_yield(self):
|
||||
"""Simple async yield to allow event loop processing."""
|
||||
await asyncio.sleep(0)
|
||||
+1
-100
@@ -3,10 +3,6 @@ Async helper utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import execution
|
||||
import server
|
||||
from typing import Optional, Any, Coroutine
|
||||
from .network import get_server_loop
|
||||
|
||||
@@ -51,99 +47,4 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
|
||||
|
||||
if error:
|
||||
raise error
|
||||
return result
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def _summarize_node_errors(node_errors: dict) -> str:
|
||||
if not isinstance(node_errors, dict) or not node_errors:
|
||||
return ""
|
||||
|
||||
parts = []
|
||||
for node_id, entry in node_errors.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
class_type = str(entry.get("class_type") or "UnknownNode")
|
||||
for err in entry.get("errors", []):
|
||||
if not isinstance(err, dict):
|
||||
continue
|
||||
message = str(err.get("message") or "validation error")
|
||||
details = str(err.get("details") or "").strip()
|
||||
parts.append(
|
||||
f"{class_type}#{node_id}: {message}{f' ({details})' if details else ''}"
|
||||
)
|
||||
if len(parts) >= 5:
|
||||
return " | ".join(parts)
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
"""Raised when a prompt fails ComfyUI validation with structured context."""
|
||||
|
||||
def __init__(self, error_payload, node_errors=None):
|
||||
payload = error_payload if isinstance(error_payload, dict) else {
|
||||
"type": "prompt_validation_failed",
|
||||
"message": str(error_payload),
|
||||
"details": "",
|
||||
"extra_info": {},
|
||||
}
|
||||
self.validation_error = dict(payload)
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
if self.node_errors:
|
||||
details = str(self.validation_error.get("details") or "").strip()
|
||||
if not details:
|
||||
summary = _summarize_node_errors(self.node_errors)
|
||||
if summary:
|
||||
self.validation_error["details"] = summary
|
||||
|
||||
merged = dict(self.validation_error)
|
||||
if self.node_errors:
|
||||
merged["node_errors"] = self.node_errors
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
if not valid[0]:
|
||||
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
extra_data["client_id"] = client_id
|
||||
|
||||
sensitive = {}
|
||||
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
|
||||
if include_queue_metadata:
|
||||
return {
|
||||
"prompt_id": prompt_id,
|
||||
"number": number,
|
||||
"node_errors": {},
|
||||
}
|
||||
|
||||
return prompt_id
|
||||
return result
|
||||
@@ -1,103 +0,0 @@
|
||||
import base64
|
||||
import binascii
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .image import ensure_contiguous
|
||||
|
||||
|
||||
MAX_AUDIO_PAYLOAD_BYTES = int(
|
||||
os.environ.get("COMFYUI_MAX_AUDIO_PAYLOAD_BYTES", str(256 * 1024 * 1024))
|
||||
)
|
||||
|
||||
|
||||
def encode_audio_payload(audio_payload):
|
||||
"""Serialize an AUDIO dict into JSON-safe canonical envelope payload."""
|
||||
if not isinstance(audio_payload, dict):
|
||||
return None
|
||||
|
||||
waveform = audio_payload.get("waveform")
|
||||
if waveform is None or not isinstance(waveform, torch.Tensor) or waveform.numel() == 0:
|
||||
return None
|
||||
|
||||
sample_rate = audio_payload.get("sample_rate", 44100)
|
||||
try:
|
||||
sample_rate = int(sample_rate)
|
||||
except (TypeError, ValueError):
|
||||
sample_rate = 44100
|
||||
|
||||
waveform_cpu = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()
|
||||
data_bytes = waveform_cpu.numpy().tobytes()
|
||||
if len(data_bytes) > MAX_AUDIO_PAYLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Audio payload too large: {len(data_bytes)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
|
||||
)
|
||||
|
||||
return {
|
||||
"sample_rate": sample_rate,
|
||||
"shape": [int(dim) for dim in waveform_cpu.shape],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(data_bytes).decode("ascii"),
|
||||
}
|
||||
|
||||
|
||||
def decode_audio_payload(audio_payload):
|
||||
"""Decode canonical envelope audio payload into an AUDIO dict."""
|
||||
if audio_payload is None:
|
||||
return None
|
||||
if not isinstance(audio_payload, dict):
|
||||
raise ValueError("Field 'audio' must be an object when provided.")
|
||||
|
||||
encoded = audio_payload.get("data")
|
||||
shape = audio_payload.get("shape")
|
||||
sample_rate = audio_payload.get("sample_rate", 44100)
|
||||
dtype = audio_payload.get("dtype", "float32")
|
||||
|
||||
if not isinstance(encoded, str) or not encoded.strip():
|
||||
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
||||
if not isinstance(shape, list) or len(shape) != 3:
|
||||
raise ValueError("Field 'audio.shape' must be a 3-item list [batch, channels, samples].")
|
||||
if dtype != "float32":
|
||||
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
||||
|
||||
try:
|
||||
shape_tuple = tuple(int(dim) for dim in shape)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
||||
|
||||
if shape_tuple[0] <= 0 or shape_tuple[1] <= 0 or shape_tuple[2] < 0:
|
||||
raise ValueError(
|
||||
"Field 'audio.shape' must be [batch>0, channels>0, samples>=0]."
|
||||
)
|
||||
|
||||
try:
|
||||
sample_rate = int(sample_rate)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.sample_rate' must be an integer.") from exc
|
||||
if sample_rate <= 0:
|
||||
raise ValueError("Field 'audio.sample_rate' must be positive.")
|
||||
|
||||
try:
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.data' is not valid base64.") from exc
|
||||
|
||||
if len(raw) > MAX_AUDIO_PAYLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Field 'audio.data' too large: {len(raw)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
|
||||
)
|
||||
|
||||
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError(
|
||||
f"Field 'audio.data' byte size mismatch: expected {expected_bytes}, got {len(raw)}."
|
||||
)
|
||||
|
||||
array = np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple)
|
||||
waveform = torch.from_numpy(array.copy())
|
||||
return {
|
||||
"waveform": ensure_contiguous(waveform),
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
Lightweight Cloudflare Tunnel manager.
|
||||
|
||||
Responsibilities:
|
||||
- Locate or download the cloudflared binary for the current platform.
|
||||
- Start/stop a quick tunnel pointing to the local ComfyUI server.
|
||||
- Track status and persist minimal state in the shared config.
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import stat
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from urllib import request, error as urlerror
|
||||
|
||||
from .config import load_config, save_config
|
||||
from .logging import debug_log, log
|
||||
from .network import get_server_port
|
||||
from .process import is_process_alive, terminate_process
|
||||
|
||||
# Regex to capture the generated public URL from cloudflared output
|
||||
PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
|
||||
|
||||
# Default timeout (seconds) to wait for a tunnel URL before giving up
|
||||
TUNNEL_START_TIMEOUT = 25
|
||||
|
||||
|
||||
def _normalize_host(value):
|
||||
if not value or not isinstance(value, str):
|
||||
return ""
|
||||
value = value.strip()
|
||||
value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
|
||||
return value.split("/")[0]
|
||||
|
||||
|
||||
class CloudflareTunnelManager:
|
||||
def __init__(self):
|
||||
self.process = None
|
||||
self.pid = None
|
||||
self.public_url = None
|
||||
self.last_error = None
|
||||
self.log_file = None
|
||||
self.status = "stopped"
|
||||
self.previous_master_host = None
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
self._url_event = None
|
||||
self._reader_thread = None
|
||||
self._loop = None
|
||||
self._recent_logs = []
|
||||
|
||||
self.binary_path = None
|
||||
self._restore_state()
|
||||
|
||||
# --- Binary discovery & download ---
|
||||
@property
|
||||
def base_dir(self):
|
||||
return os.path.dirname(os.path.dirname(__file__))
|
||||
|
||||
@property
|
||||
def bin_dir(self):
|
||||
return os.path.join(self.base_dir, "bin")
|
||||
|
||||
def _get_asset_name(self):
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
|
||||
if system == "windows":
|
||||
if "arm" in machine:
|
||||
return "cloudflared-windows-arm64.exe"
|
||||
return "cloudflared-windows-amd64.exe"
|
||||
if system == "darwin":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-darwin-arm64"
|
||||
return "cloudflared-darwin-amd64"
|
||||
if system == "linux":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-linux-arm64"
|
||||
return "cloudflared-linux-amd64"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
|
||||
|
||||
def _download_cloudflared(self):
|
||||
asset = self._get_asset_name()
|
||||
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
|
||||
os.makedirs(self.bin_dir, exist_ok=True)
|
||||
target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
|
||||
|
||||
debug_log(f"Downloading cloudflared from {url}")
|
||||
try:
|
||||
with request.urlopen(url, timeout=30) as resp:
|
||||
with open(target_path, "wb") as f:
|
||||
shutil.copyfileobj(resp, f)
|
||||
except urlerror.URLError as exc:
|
||||
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
|
||||
|
||||
# Make executable
|
||||
st = os.stat(target_path)
|
||||
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
|
||||
debug_log(f"Downloaded cloudflared to {target_path}")
|
||||
return target_path
|
||||
|
||||
def ensure_binary(self):
|
||||
# Allow explicit override
|
||||
env_path = os.environ.get("CLOUDFLARED_PATH")
|
||||
if env_path and os.path.exists(env_path):
|
||||
self.binary_path = env_path
|
||||
return env_path
|
||||
|
||||
# Check cached path
|
||||
cached = self.binary_path
|
||||
if cached and os.path.exists(cached):
|
||||
return cached
|
||||
|
||||
# Local bin directory
|
||||
candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
|
||||
if os.path.exists(candidate):
|
||||
self.binary_path = candidate
|
||||
return candidate
|
||||
|
||||
# System PATH
|
||||
path_binary = shutil.which("cloudflared")
|
||||
if path_binary:
|
||||
self.binary_path = path_binary
|
||||
return path_binary
|
||||
|
||||
# Download on demand
|
||||
self.binary_path = self._download_cloudflared()
|
||||
return self.binary_path
|
||||
|
||||
# --- State helpers ---
|
||||
def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
|
||||
|
||||
if status is not None:
|
||||
tunnel_cfg["status"] = status
|
||||
if public_url is not None:
|
||||
tunnel_cfg["public_url"] = public_url
|
||||
if pid is not None:
|
||||
tunnel_cfg["pid"] = pid
|
||||
if log_file is not None:
|
||||
tunnel_cfg["log_file"] = log_file
|
||||
if previous_host is not None:
|
||||
tunnel_cfg["previous_master_host"] = previous_host
|
||||
if master_host is not None:
|
||||
cfg.setdefault("master", {})["host"] = master_host
|
||||
|
||||
cfg["tunnel"] = tunnel_cfg
|
||||
save_config(cfg)
|
||||
|
||||
def _restore_state(self):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
|
||||
|
||||
self.public_url = tunnel_cfg.get("public_url") or None
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host")
|
||||
self.log_file = tunnel_cfg.get("log_file")
|
||||
pid = tunnel_cfg.get("pid")
|
||||
|
||||
if pid and is_process_alive(pid):
|
||||
self.pid = pid
|
||||
self.status = tunnel_cfg.get("status", "running")
|
||||
debug_log(f"Detected existing cloudflared process (pid={pid})")
|
||||
else:
|
||||
# Clear stale info
|
||||
self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
|
||||
self.status = "stopped"
|
||||
self.pid = None
|
||||
|
||||
# --- Process output handling ---
|
||||
def _append_log(self, line):
|
||||
if self.log_file:
|
||||
try:
|
||||
with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
|
||||
f.write(line + "\n")
|
||||
except Exception as exc: # pragma: no cover - best effort logging
|
||||
debug_log(f"Failed to write tunnel log: {exc}")
|
||||
|
||||
self._recent_logs.append(line)
|
||||
# Keep last ~200 lines to avoid unbounded growth
|
||||
if len(self._recent_logs) > 200:
|
||||
self._recent_logs = self._recent_logs[-200:]
|
||||
|
||||
def _reader(self):
|
||||
assert self.process is not None
|
||||
loop = self._loop
|
||||
for raw_line in iter(self.process.stdout.readline, ""):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
self._append_log(line)
|
||||
match = PUBLIC_URL_PATTERN.search(line)
|
||||
if match and not self.public_url:
|
||||
self.public_url = match.group(1).rstrip("/")
|
||||
self.status = "running"
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
# Capture obvious error strings
|
||||
if "error" in line.lower() and not self.last_error:
|
||||
self.last_error = line
|
||||
|
||||
# Process exited
|
||||
if self.status == "starting" and not self.public_url:
|
||||
self.status = "error"
|
||||
if not self.last_error:
|
||||
self.last_error = "Cloudflare tunnel exited before becoming ready"
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
elif self.status == "running":
|
||||
# Treat unexpected exit as stopped so UI can retry
|
||||
self.status = "stopped"
|
||||
elif self.status == "stopped":
|
||||
# Normal stop path, nothing to do
|
||||
pass
|
||||
|
||||
# --- Public API ---
|
||||
async def start_tunnel(self):
|
||||
async with self._lock:
|
||||
if self.process and self.process.poll() is None:
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.process.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
# If we had a stale pid, ensure it's not running
|
||||
if self.pid and is_process_alive(self.pid):
|
||||
debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
|
||||
await self.stop_tunnel()
|
||||
|
||||
binary = await asyncio.to_thread(self.ensure_binary)
|
||||
port = get_server_port()
|
||||
self.status = "starting"
|
||||
self.last_error = None
|
||||
self.public_url = None
|
||||
self._recent_logs = []
|
||||
|
||||
# Remember current master host so we can restore it if the tunnel stops
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host") or ""
|
||||
tunnel_cfg = config.get("tunnel") or {}
|
||||
if tunnel_cfg.get("previous_master_host"):
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host")
|
||||
else:
|
||||
self.previous_master_host = master_host
|
||||
|
||||
os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
|
||||
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
||||
self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
|
||||
|
||||
cmd = [
|
||||
binary,
|
||||
"tunnel",
|
||||
"--no-autoupdate",
|
||||
"--url",
|
||||
f"http://127.0.0.1:{port}",
|
||||
]
|
||||
|
||||
debug_log(f"Starting cloudflared: {' '.join(cmd)}")
|
||||
try:
|
||||
self.process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
self.status = "error"
|
||||
raise RuntimeError("cloudflared binary not found")
|
||||
except Exception as exc:
|
||||
self.status = "error"
|
||||
raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
|
||||
|
||||
self.pid = self.process.pid
|
||||
self._persist_state(
|
||||
status="starting",
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host
|
||||
)
|
||||
|
||||
# Kick off background reader
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._url_event = asyncio.Event()
|
||||
self._reader_thread = threading.Thread(target=self._reader, daemon=True)
|
||||
self._reader_thread.start()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
self.last_error = "Timed out waiting for Cloudflare to assign a URL"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(self.last_error)
|
||||
|
||||
if self.status != "running" or not self.public_url:
|
||||
error_msg = self.last_error or "Cloudflare tunnel failed to start"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
|
||||
master_host = _normalize_host(self.public_url)
|
||||
self._persist_state(
|
||||
status="running",
|
||||
public_url=self.public_url,
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host or "",
|
||||
master_host=master_host
|
||||
)
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
async def stop_tunnel(self):
|
||||
async with self._lock:
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if not pid:
|
||||
self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
|
||||
self.status = "stopped"
|
||||
return {"status": "stopped"}
|
||||
|
||||
debug_log(f"Stopping cloudflared (pid={pid})")
|
||||
if self.process:
|
||||
terminate_process(self.process, timeout=5)
|
||||
else:
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
time.sleep(0.5)
|
||||
except Exception as exc: # pragma: no cover - best effort
|
||||
debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
|
||||
|
||||
config = load_config()
|
||||
tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
|
||||
active_url = tunnel_cfg.get("public_url")
|
||||
current_master_host = (config.get("master") or {}).get("host")
|
||||
restore_host = None
|
||||
if active_url:
|
||||
active_host = _normalize_host(active_url)
|
||||
current_host = _normalize_host(current_master_host)
|
||||
if current_host == active_host:
|
||||
restore_host = self.previous_master_host or ""
|
||||
|
||||
self.status = "stopped"
|
||||
self.public_url = None
|
||||
self.pid = None
|
||||
self.process = None
|
||||
self.last_error = None
|
||||
self._url_event = None
|
||||
if self._reader_thread and self._reader_thread.is_alive():
|
||||
self._reader_thread.join(timeout=1)
|
||||
self._reader_thread = None
|
||||
self._persist_state(
|
||||
status="stopped",
|
||||
public_url="",
|
||||
pid=None,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host,
|
||||
master_host=restore_host
|
||||
)
|
||||
return {"status": "stopped"}
|
||||
|
||||
def get_status(self):
|
||||
alive = False
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if pid:
|
||||
alive = is_process_alive(pid)
|
||||
if not alive and self.status == "running":
|
||||
self.status = "stopped"
|
||||
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": pid,
|
||||
"log_file": self.log_file,
|
||||
"last_error": self.last_error,
|
||||
"binary_path": self.binary_path or shutil.which("cloudflared"),
|
||||
"recent_logs": self._recent_logs[-20:],
|
||||
"previous_master_host": self.previous_master_host,
|
||||
}
|
||||
|
||||
|
||||
# Singleton tunnel manager
|
||||
cloudflare_tunnel_manager = CloudflareTunnelManager()
|
||||
@@ -1,5 +0,0 @@
|
||||
from .tunnel import CloudflareTunnelManager
|
||||
|
||||
cloudflare_tunnel_manager = CloudflareTunnelManager()
|
||||
|
||||
__all__ = ["CloudflareTunnelManager", "cloudflare_tunnel_manager"]
|
||||
@@ -1,83 +0,0 @@
|
||||
"""Cloudflared binary discovery and download helpers."""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import stat
|
||||
from urllib import error as urlerror
|
||||
from urllib import request
|
||||
|
||||
from ..logging import debug_log
|
||||
|
||||
|
||||
def _get_project_root():
|
||||
return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
|
||||
|
||||
def _get_cloudflared_dir():
|
||||
return os.path.join(_get_project_root(), "bin")
|
||||
|
||||
|
||||
def _get_platform_binary_name():
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
|
||||
if system == "windows":
|
||||
if "arm" in machine:
|
||||
return "cloudflared-windows-arm64.exe"
|
||||
return "cloudflared-windows-amd64.exe"
|
||||
if system == "darwin":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-darwin-arm64"
|
||||
return "cloudflared-darwin-amd64"
|
||||
if system == "linux":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-linux-arm64"
|
||||
return "cloudflared-linux-amd64"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
|
||||
|
||||
|
||||
def _get_binary_path(bin_dir=None):
|
||||
bin_dir = bin_dir or _get_cloudflared_dir()
|
||||
binary_name = "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared"
|
||||
return os.path.join(bin_dir, binary_name)
|
||||
|
||||
|
||||
def _download_cloudflared():
|
||||
asset = _get_platform_binary_name()
|
||||
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
|
||||
|
||||
bin_dir = _get_cloudflared_dir()
|
||||
os.makedirs(bin_dir, exist_ok=True)
|
||||
target_path = _get_binary_path(bin_dir)
|
||||
|
||||
debug_log(f"Downloading cloudflared from {url}")
|
||||
try:
|
||||
with request.urlopen(url, timeout=30) as resp:
|
||||
with open(target_path, "wb") as f:
|
||||
shutil.copyfileobj(resp, f)
|
||||
except urlerror.URLError as exc:
|
||||
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
|
||||
|
||||
st = os.stat(target_path)
|
||||
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
|
||||
debug_log(f"Downloaded cloudflared to {target_path}")
|
||||
return target_path
|
||||
|
||||
|
||||
def ensure_binary() -> str:
|
||||
"""Return a usable cloudflared binary path, downloading if necessary."""
|
||||
env_path = os.environ.get("CLOUDFLARED_PATH")
|
||||
if env_path and os.path.exists(env_path):
|
||||
return env_path
|
||||
|
||||
local_candidate = _get_binary_path()
|
||||
if os.path.exists(local_candidate):
|
||||
return local_candidate
|
||||
|
||||
path_binary = shutil.which("cloudflared")
|
||||
if path_binary:
|
||||
return path_binary
|
||||
|
||||
return _download_cloudflared()
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Background cloudflared process output reader."""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import threading
|
||||
|
||||
from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE
|
||||
from ..logging import debug_log
|
||||
|
||||
PUBLIC_URL_PATTERN = re.compile(
|
||||
r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
class ProcessReader:
|
||||
def __init__(self, log_file=None):
|
||||
self._process = None
|
||||
self._thread = None
|
||||
self._loop = None
|
||||
self._url_event = None
|
||||
self._public_url = None
|
||||
self._last_error = None
|
||||
self._recent_logs = []
|
||||
self._log_file = log_file
|
||||
|
||||
def set_log_file(self, log_file):
|
||||
self._log_file = log_file
|
||||
|
||||
def _append_log(self, line):
|
||||
if self._log_file:
|
||||
try:
|
||||
with open(self._log_file, "a", encoding="utf-8", errors="replace") as f:
|
||||
f.write(line + "\n")
|
||||
except Exception as exc: # pragma: no cover
|
||||
debug_log(f"Failed to write tunnel log: {exc}")
|
||||
|
||||
self._recent_logs.append(line)
|
||||
if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE:
|
||||
self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:]
|
||||
|
||||
def _reader(self):
|
||||
process = self._process
|
||||
if process is None:
|
||||
return
|
||||
|
||||
loop = self._loop
|
||||
for raw_line in iter(process.stdout.readline, ""):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
self._append_log(line)
|
||||
match = PUBLIC_URL_PATTERN.search(line)
|
||||
if match and not self._public_url:
|
||||
self._public_url = match.group(1).rstrip("/")
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
if "error" in line.lower() and not self._last_error:
|
||||
self._last_error = line
|
||||
|
||||
if self._url_event and loop:
|
||||
if not self._last_error and not self._public_url:
|
||||
self._last_error = "Cloudflare tunnel exited before becoming ready"
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
def start(self, process, loop):
|
||||
self._process = process
|
||||
self._loop = loop
|
||||
self._url_event = asyncio.Event()
|
||||
self._public_url = None
|
||||
self._last_error = None
|
||||
self._recent_logs = []
|
||||
self._thread = threading.Thread(target=self._reader, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
async def wait_for_url(self, timeout):
|
||||
if not self._url_event:
|
||||
return None
|
||||
await asyncio.wait_for(self._url_event.wait(), timeout=timeout)
|
||||
return self._public_url
|
||||
|
||||
def stop(self):
|
||||
if self._thread and self._thread.is_alive():
|
||||
self._thread.join(timeout=1)
|
||||
self._thread = None
|
||||
self._process = None
|
||||
self._loop = None
|
||||
self._url_event = None
|
||||
|
||||
def get_url(self):
|
||||
return self._public_url
|
||||
|
||||
def get_last_error(self):
|
||||
return self._last_error
|
||||
|
||||
def get_recent_logs(self):
|
||||
return list(self._recent_logs)
|
||||
@@ -1,81 +0,0 @@
|
||||
"""Cloudflare tunnel state persistence helpers."""
|
||||
|
||||
from ..config import load_config, save_config
|
||||
from ..network import normalize_host
|
||||
|
||||
|
||||
def _get_tunnel_config(cfg):
|
||||
tunnel_cfg = cfg.get("tunnel", {})
|
||||
if isinstance(tunnel_cfg, dict):
|
||||
return tunnel_cfg
|
||||
return {}
|
||||
|
||||
|
||||
def load_tunnel_state():
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {}
|
||||
return {
|
||||
"status": tunnel_cfg.get("status", "stopped"),
|
||||
"public_url": tunnel_cfg.get("public_url") or None,
|
||||
"pid": tunnel_cfg.get("pid"),
|
||||
"log_file": tunnel_cfg.get("log_file"),
|
||||
"previous_master_host": tunnel_cfg.get("previous_master_host"),
|
||||
"master_host": master_cfg.get("host"),
|
||||
}
|
||||
|
||||
|
||||
def persist_tunnel_state(
|
||||
status=None,
|
||||
public_url=None,
|
||||
pid=None,
|
||||
log_file=None,
|
||||
previous_host=None,
|
||||
master_host=None,
|
||||
):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
|
||||
if status is not None:
|
||||
tunnel_cfg["status"] = status
|
||||
if public_url is not None:
|
||||
tunnel_cfg["public_url"] = public_url
|
||||
if pid is not None:
|
||||
tunnel_cfg["pid"] = pid
|
||||
if log_file is not None:
|
||||
tunnel_cfg["log_file"] = log_file
|
||||
if previous_host is not None:
|
||||
tunnel_cfg["previous_master_host"] = previous_host
|
||||
if master_host is not None:
|
||||
cfg.setdefault("master", {})["host"] = master_host
|
||||
|
||||
cfg["tunnel"] = tunnel_cfg
|
||||
save_config(cfg)
|
||||
|
||||
|
||||
def clear_tunnel_state(log_file=None, previous_host=None, master_host=None):
|
||||
persist_tunnel_state(
|
||||
status="stopped",
|
||||
public_url="",
|
||||
pid=None,
|
||||
log_file=log_file,
|
||||
previous_host=previous_host,
|
||||
master_host=master_host,
|
||||
)
|
||||
|
||||
|
||||
def resolve_restore_master_host(previous_master_host):
|
||||
"""Determine whether master host should be restored after tunnel stop."""
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
active_url = tunnel_cfg.get("public_url")
|
||||
current_master_host = (cfg.get("master") or {}).get("host")
|
||||
|
||||
if not active_url:
|
||||
return None
|
||||
|
||||
active_host = normalize_host(active_url)
|
||||
current_host = normalize_host(current_master_host)
|
||||
if current_host == active_host:
|
||||
return previous_master_host or ""
|
||||
return None
|
||||
@@ -1,207 +0,0 @@
|
||||
"""Cloudflare tunnel lifecycle manager."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
from ..constants import TUNNEL_START_TIMEOUT
|
||||
from ..logging import debug_log
|
||||
from ..network import get_server_port, normalize_host
|
||||
from ..process import is_process_alive, terminate_process
|
||||
from .binary import ensure_binary
|
||||
from .process_reader import ProcessReader
|
||||
from .state import clear_tunnel_state, load_tunnel_state, persist_tunnel_state, resolve_restore_master_host
|
||||
|
||||
|
||||
class CloudflareTunnelManager:
|
||||
def __init__(self):
|
||||
self.process = None
|
||||
self.pid = None
|
||||
self.public_url = None
|
||||
self.last_error = None
|
||||
self.log_file = None
|
||||
self.status = "stopped"
|
||||
self.previous_master_host = None
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
self._reader = ProcessReader()
|
||||
self.binary_path = None
|
||||
|
||||
self._restore_state()
|
||||
|
||||
@property
|
||||
def base_dir(self):
|
||||
return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
|
||||
def _restore_state(self):
|
||||
state = load_tunnel_state()
|
||||
|
||||
self.public_url = state.get("public_url") or None
|
||||
self.previous_master_host = state.get("previous_master_host")
|
||||
self.log_file = state.get("log_file")
|
||||
pid = state.get("pid")
|
||||
|
||||
if pid and is_process_alive(pid):
|
||||
self.pid = pid
|
||||
self.status = state.get("status") or "running"
|
||||
debug_log(f"Detected existing cloudflared process (pid={pid})")
|
||||
else:
|
||||
clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
|
||||
self.status = "stopped"
|
||||
self.pid = None
|
||||
|
||||
async def start_tunnel(self):
|
||||
async with self._lock:
|
||||
if self.process and self.process.poll() is None:
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.process.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
if self.pid and is_process_alive(self.pid):
|
||||
debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
|
||||
await self.stop_tunnel()
|
||||
|
||||
binary = await asyncio.to_thread(ensure_binary)
|
||||
self.binary_path = binary
|
||||
port = get_server_port()
|
||||
self.status = "starting"
|
||||
self.last_error = None
|
||||
self.public_url = None
|
||||
|
||||
state = load_tunnel_state()
|
||||
master_host = state.get("master_host") or ""
|
||||
if state.get("previous_master_host"):
|
||||
self.previous_master_host = state.get("previous_master_host")
|
||||
else:
|
||||
self.previous_master_host = master_host
|
||||
|
||||
os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
|
||||
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
||||
self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
|
||||
|
||||
cmd = [
|
||||
binary,
|
||||
"tunnel",
|
||||
"--no-autoupdate",
|
||||
"--url",
|
||||
f"http://127.0.0.1:{port}",
|
||||
]
|
||||
|
||||
debug_log(f"Starting cloudflared: {' '.join(cmd)}")
|
||||
try:
|
||||
self.process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
self.status = "error"
|
||||
raise RuntimeError("cloudflared binary not found")
|
||||
except Exception as exc:
|
||||
self.status = "error"
|
||||
raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
|
||||
|
||||
self.pid = self.process.pid
|
||||
persist_tunnel_state(
|
||||
status="starting",
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host,
|
||||
)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
self._reader.set_log_file(self.log_file)
|
||||
self._reader.start(self.process, loop)
|
||||
|
||||
try:
|
||||
await self._reader.wait_for_url(timeout=TUNNEL_START_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
self.last_error = "Timed out waiting for Cloudflare to assign a URL"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(self.last_error)
|
||||
|
||||
public_url = self._reader.get_url()
|
||||
if not public_url:
|
||||
self.last_error = self._reader.get_last_error() or "Cloudflare tunnel failed to start"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(self.last_error)
|
||||
|
||||
self.public_url = public_url
|
||||
self.status = "running"
|
||||
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
|
||||
|
||||
persist_tunnel_state(
|
||||
status="running",
|
||||
public_url=self.public_url,
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host or "",
|
||||
master_host=normalize_host(self.public_url),
|
||||
)
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
async def stop_tunnel(self):
|
||||
async with self._lock:
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if not pid:
|
||||
clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
|
||||
self.status = "stopped"
|
||||
return {"status": "stopped"}
|
||||
|
||||
debug_log(f"Stopping cloudflared (pid={pid})")
|
||||
if self.process:
|
||||
terminate_process(self.process, timeout=5)
|
||||
else:
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
time.sleep(0.5)
|
||||
except Exception as exc: # pragma: no cover
|
||||
debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
|
||||
|
||||
restore_host = resolve_restore_master_host(self.previous_master_host)
|
||||
|
||||
self.status = "stopped"
|
||||
self.public_url = None
|
||||
self.pid = None
|
||||
self.process = None
|
||||
self.last_error = None
|
||||
self._reader.stop()
|
||||
|
||||
clear_tunnel_state(
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host,
|
||||
master_host=restore_host,
|
||||
)
|
||||
return {"status": "stopped"}
|
||||
|
||||
def get_status(self):
|
||||
alive = False
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if pid:
|
||||
alive = is_process_alive(pid)
|
||||
if not alive and self.status == "running":
|
||||
self.status = "stopped"
|
||||
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": pid,
|
||||
"log_file": self.log_file,
|
||||
"last_error": self.last_error or self._reader.get_last_error(),
|
||||
"binary_path": self.binary_path or shutil.which("cloudflared"),
|
||||
"recent_logs": self._reader.get_recent_logs()[-20:],
|
||||
"previous_master_host": self.previous_master_host,
|
||||
}
|
||||
+12
-84
@@ -1,23 +1,14 @@
|
||||
"""
|
||||
Configuration management for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from .logging import log
|
||||
|
||||
# Import defaults for timeout fallbacks
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
|
||||
CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
|
||||
_config_cache = None
|
||||
_config_mtime = 0.0
|
||||
_config_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _config_path():
|
||||
return CONFIG_FILE
|
||||
|
||||
def get_default_config():
|
||||
"""Returns the default configuration dictionary. Single source of truth."""
|
||||
@@ -29,11 +20,7 @@ def get_default_config():
|
||||
"auto_launch_workers": False,
|
||||
"stop_workers_on_master_exit": True,
|
||||
"master_delegate_only": False,
|
||||
"websocket_orchestration": True,
|
||||
"worker_probe_concurrency": 8,
|
||||
"worker_prep_concurrency": 4,
|
||||
"media_sync_concurrency": 2,
|
||||
"media_sync_timeout_seconds": 120
|
||||
"websocket_orchestration": True
|
||||
},
|
||||
"tunnel": {
|
||||
"status": "stopped",
|
||||
@@ -44,93 +31,34 @@ def get_default_config():
|
||||
}
|
||||
}
|
||||
|
||||
def _merge_with_defaults(data, defaults):
|
||||
"""Recursively merge loaded config data with default keys."""
|
||||
if not isinstance(data, dict):
|
||||
return defaults
|
||||
|
||||
merged = {}
|
||||
for key, default_value in defaults.items():
|
||||
loaded_value = data.get(key, default_value)
|
||||
if isinstance(default_value, dict) and isinstance(loaded_value, dict):
|
||||
merged[key] = _merge_with_defaults(loaded_value, default_value)
|
||||
else:
|
||||
merged[key] = loaded_value
|
||||
|
||||
# Preserve unknown keys for forward compatibility.
|
||||
for key, value in data.items():
|
||||
if key not in merged:
|
||||
merged[key] = value
|
||||
|
||||
return merged
|
||||
|
||||
|
||||
def invalidate_config_cache():
|
||||
"""Invalidate in-memory config cache so next load reads from disk."""
|
||||
global _config_cache, _config_mtime
|
||||
_config_cache = None
|
||||
_config_mtime = 0.0
|
||||
|
||||
|
||||
def load_config():
|
||||
"""Loads the config, falling back to defaults if the file is missing or invalid."""
|
||||
global _config_cache, _config_mtime
|
||||
path = _config_path()
|
||||
|
||||
try:
|
||||
mtime = os.path.getmtime(path)
|
||||
except OSError:
|
||||
if _config_cache is None:
|
||||
_config_cache = get_default_config()
|
||||
return _config_cache
|
||||
|
||||
if _config_cache is None or mtime != _config_mtime:
|
||||
if os.path.exists(CONFIG_FILE):
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
loaded = json.load(f)
|
||||
_config_cache = _merge_with_defaults(loaded, get_default_config())
|
||||
with open(CONFIG_FILE, 'r') as f:
|
||||
data = json.load(f)
|
||||
defaults = get_default_config()
|
||||
for key, value in defaults.items():
|
||||
if key not in data:
|
||||
data[key] = value
|
||||
return data
|
||||
except Exception as e:
|
||||
log(f"Error loading config, using defaults: {e}")
|
||||
_config_cache = get_default_config()
|
||||
_config_mtime = mtime
|
||||
|
||||
return _config_cache
|
||||
return get_default_config()
|
||||
|
||||
def save_config(config):
|
||||
"""Saves the configuration to file."""
|
||||
tmp_path = f"{_config_path()}.tmp"
|
||||
try:
|
||||
with open(tmp_path, 'w', encoding='utf-8') as f:
|
||||
with open(CONFIG_FILE, 'w') as f:
|
||||
json.dump(config, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, _config_path())
|
||||
invalidate_config_cache()
|
||||
return True
|
||||
except Exception as e:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
log(f"Error saving config: {e}")
|
||||
return False
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def config_transaction():
|
||||
"""Acquire config lock, yield loaded config, and save if changed."""
|
||||
async with _config_lock:
|
||||
config = load_config()
|
||||
original_snapshot = json.dumps(config, sort_keys=True)
|
||||
yield config
|
||||
updated_snapshot = json.dumps(config, sort_keys=True)
|
||||
if updated_snapshot != original_snapshot:
|
||||
if not save_config(config):
|
||||
raise RuntimeError("Failed to save config")
|
||||
|
||||
def ensure_config_exists():
|
||||
"""Creates default config file if it doesn't exist. Used by __init__.py"""
|
||||
if not os.path.exists(_config_path()):
|
||||
if not os.path.exists(CONFIG_FILE):
|
||||
default_config = get_default_config()
|
||||
if save_config(default_config):
|
||||
from .logging import debug_log
|
||||
|
||||
@@ -13,10 +13,6 @@ PROCESS_TERMINATION_TIMEOUT = 5.0
|
||||
WORKER_CHECK_INTERVAL = 2.0
|
||||
STATUS_CHECK_INTERVAL = 5.0
|
||||
|
||||
# Cloudflare tunnel
|
||||
TUNNEL_START_TIMEOUT = float(os.environ.get("TUNNEL_START_TIMEOUT", "25"))
|
||||
CLOUDFLARE_LOG_BUFFER_SIZE = 200
|
||||
|
||||
# Network
|
||||
CHUNK_SIZE = 8192
|
||||
LOG_TAIL_BYTES = 65536 # 64KB
|
||||
@@ -34,7 +30,6 @@ TILE_TRANSFER_TIMEOUT = 30.0
|
||||
PROCESS_WAIT_TIMEOUT = 3.0
|
||||
QUEUE_INIT_TIMEOUT = 5.0
|
||||
TILE_SEND_TIMEOUT = 60.0
|
||||
JOB_INIT_GRACE_PERIOD = 10.0
|
||||
|
||||
# Memory operations
|
||||
MEMORY_CLEAR_DELAY = 0.5
|
||||
@@ -42,38 +37,5 @@ MEMORY_CLEAR_DELAY = 0.5
|
||||
# Batch processing
|
||||
MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
|
||||
|
||||
# Compressed worker -> master result transfer (DistributedCollector)
|
||||
# Hard cap on one multipart request body. Must not exceed the master's
|
||||
# --max-upload-size (ComfyUI default: 100 MiB), which rejects larger requests.
|
||||
TRANSFER_MAX_REQUEST_BYTES = int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES', str(96 * 1024 * 1024)))
|
||||
# Target media bytes per request when batching several frames into one
|
||||
# request; a single frame may exceed this up to the hard cap above.
|
||||
TRANSFER_MAX_CHUNK_BYTES = min(
|
||||
int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES', str(32 * 1024 * 1024))),
|
||||
TRANSFER_MAX_REQUEST_BYTES,
|
||||
)
|
||||
|
||||
# Heartbeat monitoring
|
||||
HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10')) # Heartbeat/check interval in seconds
|
||||
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s)
|
||||
|
||||
# USDU result collection
|
||||
DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0
|
||||
|
||||
# Static mode job poll loop
|
||||
JOB_POLL_INTERVAL = 1.0
|
||||
JOB_POLL_MAX_ATTEMPTS = 20
|
||||
|
||||
# Orchestration pipeline
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int(
|
||||
os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8')
|
||||
)
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY = int(
|
||||
os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PREP_CONCURRENCY', '4')
|
||||
)
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = int(
|
||||
os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_CONCURRENCY', '2')
|
||||
)
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT = float(
|
||||
os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_TIMEOUT', '120')
|
||||
)
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
from contextlib import contextmanager
|
||||
|
||||
import torch
|
||||
|
||||
from .logging import debug_log
|
||||
from .usdu_utils import resize_region
|
||||
|
||||
|
||||
@contextmanager
|
||||
def crop_model_cond(model, crop_regions, init_size, canvas_size, tile_size, latent_crop=False):
|
||||
"""Clone model and crop compatible model patches for tile-local sampling."""
|
||||
try:
|
||||
patched_model = model.clone()
|
||||
except Exception:
|
||||
# Fallback to original model when clone/patch access is unavailable.
|
||||
yield model
|
||||
return
|
||||
|
||||
patches = (
|
||||
patched_model
|
||||
.model_options
|
||||
.get("transformer_options", {})
|
||||
.get("patches", {})
|
||||
)
|
||||
applied_croppers = {}
|
||||
for _module, module_patches in patches.items():
|
||||
for patch in module_patches:
|
||||
if id(patch) in applied_croppers:
|
||||
continue
|
||||
if type(patch).__name__ not in ("DiffSynthCnetPatch", "ZImageControlPatch"):
|
||||
continue
|
||||
try:
|
||||
cropper = ModelPatchCropper(patch).crop(crop_regions, canvas_size, latent_crop)
|
||||
applied_croppers[id(patch)] = cropper
|
||||
except Exception as exc:
|
||||
debug_log(f"crop_model_cond: patch crop skipped for {type(patch).__name__}: {exc}")
|
||||
try:
|
||||
yield patched_model
|
||||
finally:
|
||||
for cropper in applied_croppers.values():
|
||||
del cropper
|
||||
|
||||
|
||||
class ModelPatchCropper:
|
||||
"""Stateful crop helper that restores model patch tensors on cleanup."""
|
||||
|
||||
def __init__(self, patch):
|
||||
self.patch = patch
|
||||
self.original_state = {
|
||||
"image": patch.image.clone() if isinstance(patch.image, torch.Tensor) else patch.image,
|
||||
"encoded_image": patch.encoded_image.clone() if isinstance(patch.encoded_image, torch.Tensor) else patch.encoded_image,
|
||||
"encoded_image_size": patch.encoded_image_size,
|
||||
}
|
||||
self.patch_class = type(patch).__name__
|
||||
required_attrs = (
|
||||
"image",
|
||||
"model_patch",
|
||||
"vae",
|
||||
"strength",
|
||||
"encoded_image",
|
||||
"encoded_image_size",
|
||||
)
|
||||
missing = [attr for attr in required_attrs if not hasattr(patch, attr)]
|
||||
if missing:
|
||||
raise AttributeError(
|
||||
f"{self.patch_class} missing required attrs: {', '.join(missing)}"
|
||||
)
|
||||
|
||||
def __del__(self):
|
||||
self.patch.image = self.original_state["image"]
|
||||
self.patch.encoded_image = self.original_state["encoded_image"]
|
||||
self.patch.encoded_image_size = self.original_state["encoded_image_size"]
|
||||
|
||||
def crop(self, crop_regions, canvas_size, latent_crop=True):
|
||||
patch = self.patch
|
||||
|
||||
if not isinstance(crop_regions, list):
|
||||
crop_regions = [crop_regions]
|
||||
|
||||
image_size = (patch.image.shape[2], patch.image.shape[1]) # (W,H)
|
||||
|
||||
cropped_images = []
|
||||
for crop_region in crop_regions:
|
||||
resized_crop = resize_region(crop_region, canvas_size, image_size)
|
||||
x1, y1, x2, y2 = resized_crop
|
||||
cropped_image = patch.image[:, y1:y2, x1:x2, :]
|
||||
cropped_images.append(cropped_image)
|
||||
|
||||
concatenated_image = torch.cat(cropped_images, dim=0)
|
||||
patch.image = concatenated_image
|
||||
patch.encoded_image_size = (
|
||||
concatenated_image.shape[1],
|
||||
concatenated_image.shape[2],
|
||||
)
|
||||
|
||||
if latent_crop:
|
||||
downscale_ratio = patch.vae.spacial_compression_encode()
|
||||
cropped_latents = []
|
||||
for crop_region in crop_regions:
|
||||
resized_crop = resize_region(crop_region, canvas_size, image_size)
|
||||
x1, y1, x2, y2 = tuple(x // downscale_ratio for x in resized_crop)
|
||||
cropped_latent = patch.encoded_image[:, :, y1:y2, x1:x2]
|
||||
cropped_latents.append(cropped_latent)
|
||||
patch.encoded_image = torch.cat(cropped_latents, dim=0)
|
||||
else:
|
||||
patch.__init__(
|
||||
patch.model_patch,
|
||||
patch.vae,
|
||||
concatenated_image,
|
||||
patch.strength,
|
||||
inpaint_image=patch.inpaint_image,
|
||||
mask=patch.mask,
|
||||
)
|
||||
return self
|
||||
@@ -1,42 +0,0 @@
|
||||
"""Custom exceptions for ComfyUI-Distributed."""
|
||||
|
||||
|
||||
class DistributedError(Exception):
|
||||
"""Base exception for all ComfyUI-Distributed errors."""
|
||||
|
||||
|
||||
class WorkerError(DistributedError):
|
||||
"""Error related to a specific distributed worker."""
|
||||
|
||||
def __init__(self, message, worker_id=None):
|
||||
super().__init__(message)
|
||||
self.worker_id = worker_id
|
||||
|
||||
|
||||
class WorkerTimeoutError(WorkerError):
|
||||
"""Worker did not respond within the expected timeout."""
|
||||
|
||||
|
||||
class WorkerNotAvailableError(WorkerError):
|
||||
"""Worker is unreachable or not running."""
|
||||
|
||||
|
||||
class JobQueueError(DistributedError):
|
||||
"""Error in distributed job queue management."""
|
||||
|
||||
|
||||
class TileCollectionError(DistributedError):
|
||||
"""Error collecting processed tiles from workers."""
|
||||
|
||||
|
||||
class ProcessError(DistributedError):
|
||||
"""Error managing a worker subprocess."""
|
||||
|
||||
def __init__(self, message, pid=None, worker_id=None):
|
||||
super().__init__(message)
|
||||
self.pid = pid
|
||||
self.worker_id = worker_id
|
||||
|
||||
|
||||
class TunnelError(DistributedError):
|
||||
"""Error managing the Cloudflare tunnel."""
|
||||
+4
-19
@@ -3,35 +3,20 @@ Shared logging utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
|
||||
# Config file is in parent directory
|
||||
CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
|
||||
|
||||
_debug_cache: bool | None = None
|
||||
_debug_cache_time: float = 0.0
|
||||
_DEBUG_TTL: float = 5.0
|
||||
|
||||
def is_debug_enabled():
|
||||
"""Check if debug is enabled."""
|
||||
global _debug_cache, _debug_cache_time
|
||||
|
||||
now = time.monotonic()
|
||||
if _debug_cache is not None and (now - _debug_cache_time) < _DEBUG_TTL:
|
||||
return _debug_cache
|
||||
|
||||
enabled = False
|
||||
if os.path.exists(CONFIG_FILE):
|
||||
try:
|
||||
with open(CONFIG_FILE, 'r') as f:
|
||||
config = json.load(f)
|
||||
enabled = config.get("settings", {}).get("debug", False)
|
||||
except (OSError, json.JSONDecodeError, ValueError):
|
||||
return config.get("settings", {}).get("debug", False)
|
||||
except:
|
||||
pass
|
||||
|
||||
_debug_cache = enabled
|
||||
_debug_cache_time = now
|
||||
return enabled
|
||||
return False
|
||||
|
||||
def debug_log(message):
|
||||
"""Log debug messages only if debug is enabled in config."""
|
||||
@@ -40,4 +25,4 @@ def debug_log(message):
|
||||
|
||||
def log(message):
|
||||
"""Always log important messages."""
|
||||
print(f"[Distributed] {message}")
|
||||
print(f"[Distributed] {message}")
|
||||
+2
-162
@@ -1,10 +1,7 @@
|
||||
"""
|
||||
Network and API utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import re
|
||||
import server
|
||||
from aiohttp import web
|
||||
from .logging import debug_log
|
||||
|
||||
@@ -14,11 +11,6 @@ _client_session = None
|
||||
async def get_client_session():
|
||||
"""Get or create a shared aiohttp client session."""
|
||||
global _client_session
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError as exc:
|
||||
raise RuntimeError("get_client_session() requires an active asyncio event loop.") from exc
|
||||
|
||||
if _client_session is None or _client_session.closed:
|
||||
connector = aiohttp.TCPConnector(limit=100, limit_per_host=30)
|
||||
# Don't set timeout here - set it per request
|
||||
@@ -34,14 +26,8 @@ async def cleanup_client_session():
|
||||
|
||||
async def handle_api_error(request, error, status=500):
|
||||
"""Standardized error response handler."""
|
||||
if isinstance(error, list):
|
||||
messages = [str(item) for item in error]
|
||||
debug_log(f"API Error [{status}]: {messages}")
|
||||
return web.json_response({"errors": messages}, status=status)
|
||||
|
||||
message = str(error)
|
||||
debug_log(f"API Error [{status}]: {message}")
|
||||
return web.json_response({"error": message}, status=status)
|
||||
debug_log(f"API Error: {error}")
|
||||
return web.json_response({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
def get_server_port():
|
||||
"""Get the ComfyUI server port."""
|
||||
@@ -53,149 +39,3 @@ def get_server_loop():
|
||||
import server
|
||||
return server.PromptServer.instance.loop
|
||||
|
||||
|
||||
def normalize_host(value):
|
||||
if value is None:
|
||||
return None
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
host = value.strip()
|
||||
if not host:
|
||||
return host
|
||||
host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE)
|
||||
return host.split("/")[0]
|
||||
|
||||
|
||||
def _split_host_and_port(host):
|
||||
if not host:
|
||||
return host, None
|
||||
|
||||
if host.startswith("["):
|
||||
match = re.match(r"^(\[[^\]]+\])(?::(\d+))?$", host)
|
||||
if match:
|
||||
parsed_port = int(match.group(2)) if match.group(2) else None
|
||||
return match.group(1), parsed_port
|
||||
return host, None
|
||||
|
||||
if host.count(":") == 1:
|
||||
candidate_host, candidate_port = host.rsplit(":", 1)
|
||||
if candidate_port.isdigit():
|
||||
return candidate_host, int(candidate_port)
|
||||
|
||||
return host, None
|
||||
|
||||
|
||||
def build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
host = (worker.get("host") or "").strip()
|
||||
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
|
||||
|
||||
if not host:
|
||||
host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1"
|
||||
|
||||
if host.startswith(("http://", "https://")):
|
||||
base = host.rstrip("/")
|
||||
else:
|
||||
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
|
||||
scheme = "https" if is_cloud else "http"
|
||||
default_port = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port else f":{port}"
|
||||
base = f"{scheme}://{host}{port_part}"
|
||||
|
||||
return f"{base}{endpoint}"
|
||||
|
||||
|
||||
async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None:
|
||||
"""GET {worker_url}/prompt. Returns parsed JSON or None on any failure."""
|
||||
base_url = (worker_url or "").strip().rstrip("/")
|
||||
if not base_url:
|
||||
return None
|
||||
probe_url = base_url if base_url.endswith("/prompt") else f"{base_url}/prompt"
|
||||
session = await get_client_session()
|
||||
try:
|
||||
async with session.get(
|
||||
probe_url,
|
||||
timeout=aiohttp.ClientTimeout(total=float(timeout)),
|
||||
) as response:
|
||||
if response.status != 200:
|
||||
debug_log(f"[Distributed] Worker probe non-200 status: {response.status} ({probe_url})")
|
||||
return None
|
||||
payload = await response.json()
|
||||
if isinstance(payload, dict):
|
||||
return payload
|
||||
debug_log(f"[Distributed] Worker probe returned non-object JSON: {probe_url}")
|
||||
return None
|
||||
except asyncio.TimeoutError:
|
||||
debug_log(f"[Distributed] Worker probe timed out: {probe_url}")
|
||||
return None
|
||||
except aiohttp.ClientConnectorError:
|
||||
debug_log(f"[Distributed] Worker unreachable: {probe_url}")
|
||||
return None
|
||||
except Exception as exc:
|
||||
debug_log(f"[Distributed] Worker probe error ({probe_url}): {exc}")
|
||||
return None
|
||||
|
||||
|
||||
def build_master_url(config=None, prompt_server_instance=None):
|
||||
"""Build the best public URL workers should use to reach the master."""
|
||||
if config is None:
|
||||
from .config import load_config
|
||||
config = load_config()
|
||||
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
master_cfg = (config or {}).get("master", {}) or {}
|
||||
configured_host = (master_cfg.get("host") or "").strip()
|
||||
runtime_port = getattr(prompt_server_instance, "port", 8188) or 8188
|
||||
|
||||
def _needs_https(hostname):
|
||||
hostname = hostname.lower()
|
||||
https_domains = (
|
||||
".proxy.runpod.net",
|
||||
".ngrok-free.app",
|
||||
".ngrok-free.dev",
|
||||
".ngrok.io",
|
||||
".trycloudflare.com",
|
||||
".cloudflare.dev",
|
||||
)
|
||||
return any(hostname.endswith(suffix) for suffix in https_domains)
|
||||
|
||||
if configured_host:
|
||||
if configured_host.startswith(("http://", "https://")):
|
||||
return configured_host.rstrip("/")
|
||||
|
||||
host, explicit_port = _split_host_and_port(configured_host)
|
||||
port = explicit_port if explicit_port is not None else int(runtime_port)
|
||||
scheme = "https" if _needs_https(host) or port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
if explicit_port is None and scheme == "https" and _needs_https(host):
|
||||
port = default_port_for_scheme
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{host}{port_part}"
|
||||
|
||||
address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
|
||||
if address in ("0.0.0.0", "::"):
|
||||
address = "127.0.0.1"
|
||||
port = int(runtime_port)
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{address}{port_part}"
|
||||
|
||||
|
||||
def build_master_callback_url(worker, config=None, prompt_server_instance=None):
|
||||
"""Build the callback URL a specific worker should use to reach the master."""
|
||||
prompt_server_instance = prompt_server_instance or server.PromptServer.instance
|
||||
|
||||
worker_type = str((worker or {}).get("type") or "").strip().lower()
|
||||
worker_host = normalize_host((worker or {}).get("host"))
|
||||
local_hosts = {"", "localhost", "127.0.0.1", "::1", "[::1]", "0.0.0.0"}
|
||||
|
||||
is_local_worker = worker_type == "local" or worker_host in local_hosts
|
||||
if is_local_worker:
|
||||
port = int(getattr(prompt_server_instance, "port", 8188) or 8188)
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://127.0.0.1{port_part}"
|
||||
|
||||
return build_master_url(config=config, prompt_server_instance=prompt_server_instance)
|
||||
|
||||
@@ -1,13 +0,0 @@
|
||||
from .logging import debug_log, log
|
||||
|
||||
|
||||
def trace_prefix(trace_execution_id: str) -> str:
|
||||
return f"[Distributed][exec:{trace_execution_id}]"
|
||||
|
||||
|
||||
def trace_debug(trace_execution_id: str, message: str) -> None:
|
||||
debug_log(f"{trace_prefix(trace_execution_id)} {message}")
|
||||
|
||||
|
||||
def trace_info(trace_execution_id: str, message: str) -> None:
|
||||
log(f"{trace_prefix(trace_execution_id)} {message}")
|
||||
@@ -1,411 +0,0 @@
|
||||
"""
|
||||
Compressed worker -> master image transfer.
|
||||
|
||||
Workers encode their IMAGE batch with one of the formats below and POST it to
|
||||
the master as multipart/form-data chunks; the master decodes each chunk back
|
||||
into per-frame tensors. The goal is to cut network egress (e.g. cloud workers
|
||||
returning results to a local master), so formats range from bit-exact
|
||||
lossless to strongly compressed video.
|
||||
|
||||
Formats
|
||||
-------
|
||||
legacy_png Original behaviour: one JSON request per frame, uncompressed PNG
|
||||
in base64. Kept for compatibility with unmodified masters.
|
||||
png Lossless PNG per frame, sent as binary.
|
||||
ffv1 Lossless FFV1 video for the whole batch (good for frame sequences).
|
||||
jpeg, webp Lossy per-frame stills.
|
||||
h264, h265, av1
|
||||
Lossy video for the whole batch. Smallest by far for video
|
||||
frame batches, where consecutive frames are similar.
|
||||
|
||||
All formats quantize to 8-bit exactly like the legacy path did, so the
|
||||
lossless formats reproduce the legacy master tensors bit-for-bit.
|
||||
"""
|
||||
import io
|
||||
import math
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
LEGACY_FORMAT = "legacy_png"
|
||||
IMAGE_FORMATS = ("png", "jpeg", "webp")
|
||||
VIDEO_FORMATS = ("ffv1", "h264", "h265", "av1")
|
||||
LOSSLESS_FORMATS = (LEGACY_FORMAT, "png", "ffv1")
|
||||
TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
|
||||
|
||||
DEFAULT_TRANSFER_FORMAT = LEGACY_FORMAT
|
||||
DEFAULT_TRANSFER_QUALITY = 90
|
||||
|
||||
TRANSFER_FORMAT_TOOLTIP = (
|
||||
"How workers send images back to the master. "
|
||||
"legacy_png: original uncompressed PNG (largest). "
|
||||
"png / ffv1: lossless, identical result. "
|
||||
"jpeg / webp: lossy stills. "
|
||||
"h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences."
|
||||
)
|
||||
TRANSFER_QUALITY_TOOLTIP = (
|
||||
"Quality for lossy transfer formats (ignored by lossless ones). "
|
||||
"JPEG/WebP use it directly; video codecs map it to CRF "
|
||||
"(90 -> h264/h265 CRF 13, av1 CRF 16)."
|
||||
)
|
||||
|
||||
# encoder name, max CRF, container-friendly pixel format
|
||||
_VIDEO_CODECS = {
|
||||
"ffv1": ("ffv1", None, "bgr0"),
|
||||
"h264": ("libx264", 51, "yuv420p"),
|
||||
"h265": ("libx265", 51, "yuv420p"),
|
||||
"av1": ("libsvtav1", 63, "yuv420p"),
|
||||
}
|
||||
_VIDEO_RATE = 24
|
||||
# Upper bound on frames per video chunk, so one request never holds an
|
||||
# unbounded batch even when frames compress extremely well.
|
||||
MAX_VIDEO_CHUNK_FRAMES = 240
|
||||
# When sizing the next video chunk from observed bytes/frame, aim this far
|
||||
# below the limit so ordinary variation does not force a re-encode.
|
||||
VIDEO_CHUNK_SIZE_HEADROOM = 0.85
|
||||
|
||||
|
||||
class TransferCodecError(ValueError):
|
||||
"""Raised when a transfer payload cannot be encoded or decoded."""
|
||||
|
||||
|
||||
class FrameTooLargeError(TransferCodecError):
|
||||
"""A single encoded frame does not fit in one request.
|
||||
|
||||
Raised by encode_chunks before anything has been yielded for that frame;
|
||||
callers that materialize all chunks first can therefore fail without
|
||||
having delivered a partial batch.
|
||||
"""
|
||||
|
||||
|
||||
def normalize_format(value):
|
||||
fmt = str(value or DEFAULT_TRANSFER_FORMAT).strip().lower()
|
||||
if fmt not in TRANSFER_FORMATS:
|
||||
raise TransferCodecError(
|
||||
f"Unknown transfer format '{value}'. Expected one of: {', '.join(TRANSFER_FORMATS)}"
|
||||
)
|
||||
return fmt
|
||||
|
||||
|
||||
def normalize_quality(value):
|
||||
try:
|
||||
quality = int(value)
|
||||
except (TypeError, ValueError):
|
||||
quality = DEFAULT_TRANSFER_QUALITY
|
||||
return max(1, min(100, quality))
|
||||
|
||||
|
||||
def quality_to_crf(quality, max_crf):
|
||||
"""Map 1-100 quality onto a codec CRF scale (0 = lossless/best).
|
||||
|
||||
The curve is chosen so the same quality number looks roughly the same
|
||||
across JPEG/WebP and video codecs: 90 -> CRF 13 (h264/h265), 16 (av1);
|
||||
80 -> 19 / 24; 100 -> 0.
|
||||
"""
|
||||
quality = normalize_quality(quality)
|
||||
return int(round(max_crf * (1.0 - quality / 100.0) ** 0.6))
|
||||
|
||||
|
||||
def tensor_batch_to_uint8(images):
|
||||
"""IMAGE tensor [B,H,W,C] float 0..1 -> uint8 numpy [B,H,W,3].
|
||||
|
||||
Uses the same truncating quantization as utils.image.tensor_to_pil so that
|
||||
lossless formats reproduce the legacy transfer exactly, but clips first so
|
||||
out-of-range values no longer wrap around.
|
||||
"""
|
||||
arr = images.detach().cpu().float().numpy()
|
||||
if arr.ndim == 3:
|
||||
arr = arr[None, ...]
|
||||
if arr.ndim != 4:
|
||||
raise TransferCodecError(f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(arr.shape)}")
|
||||
channels = arr.shape[-1]
|
||||
if channels == 1:
|
||||
arr = np.repeat(arr, 3, axis=-1)
|
||||
elif channels >= 3:
|
||||
arr = arr[..., :3]
|
||||
else:
|
||||
raise TransferCodecError(f"Unsupported channel count {channels}")
|
||||
return np.ascontiguousarray(np.clip(255.0 * arr, 0, 255).astype(np.uint8))
|
||||
|
||||
|
||||
def uint8_frame_to_tensor_array(frame):
|
||||
"""uint8 [H,W,3] -> float32 [1,H,W,3] numpy, matching utils.image.pil_to_tensor."""
|
||||
return (frame.astype(np.float32) / 255.0)[None, ...]
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Still images
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def encode_image_frame(frame, fmt, quality):
|
||||
img = Image.fromarray(frame)
|
||||
quality = normalize_quality(quality)
|
||||
buf = io.BytesIO()
|
||||
if fmt == "png":
|
||||
img.save(buf, format="PNG", compress_level=4) # ~same size as 6, 2x faster
|
||||
elif fmt == "jpeg":
|
||||
img.save(buf, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
|
||||
elif fmt == "webp":
|
||||
img.save(buf, format="WEBP", quality=quality, method=4)
|
||||
else:
|
||||
raise TransferCodecError(f"'{fmt}' is not a still-image transfer format")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def decode_image_frame(data):
|
||||
try:
|
||||
with Image.open(io.BytesIO(data)) as img:
|
||||
return np.asarray(img.convert("RGB"), dtype=np.uint8)
|
||||
except Exception as exc:
|
||||
raise TransferCodecError(f"Failed to decode image frame: {exc}") from exc
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Video
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _import_av():
|
||||
# SVT-AV1 prints a banner per encoder instance unless told otherwise.
|
||||
os.environ.setdefault("SVT_LOG", "1")
|
||||
try:
|
||||
import av # PyAV ships with ComfyUI
|
||||
except ImportError as exc:
|
||||
raise TransferCodecError("PyAV ('av') is required for video transfer formats") from exc
|
||||
return av
|
||||
|
||||
|
||||
def video_codec_available(fmt):
|
||||
if fmt not in _VIDEO_CODECS:
|
||||
return False
|
||||
try:
|
||||
av = _import_av()
|
||||
av.codec.Codec(_VIDEO_CODECS[fmt][0], "w")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _padded_size(fmt, pix_fmt, width, height):
|
||||
"""Encoder-friendly frame size; the decoder crops back to width x height.
|
||||
|
||||
4:2:0 needs even dimensions, and libx265 crashes on frames smaller than
|
||||
its 64px CTU, so pad up to whichever is larger.
|
||||
"""
|
||||
if pix_fmt == "yuv420p":
|
||||
width, height = width + width % 2, height + height % 2
|
||||
if fmt == "h265":
|
||||
width, height = max(width, 64), max(height, 64)
|
||||
return width, height
|
||||
|
||||
|
||||
def _pad_to(frame, width, height):
|
||||
h, w = frame.shape[:2]
|
||||
if h == height and w == width:
|
||||
return frame
|
||||
return np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
|
||||
|
||||
|
||||
class _VideoChunkWriter:
|
||||
def __init__(self, fmt, quality, width, height):
|
||||
av = _import_av()
|
||||
self._av = av
|
||||
codec_name, max_crf, pix_fmt = _VIDEO_CODECS[fmt]
|
||||
self.fmt = fmt
|
||||
self.pix_fmt = pix_fmt
|
||||
self.width, self.height = _padded_size(fmt, pix_fmt, width, height)
|
||||
self.buffer = io.BytesIO()
|
||||
self.container = av.open(self.buffer, mode="w", format="matroska")
|
||||
self.stream = self.container.add_stream(codec_name, rate=_VIDEO_RATE)
|
||||
self.stream.width = self.width
|
||||
self.stream.height = self.height
|
||||
self.stream.pix_fmt = pix_fmt
|
||||
options = {}
|
||||
if max_crf is not None:
|
||||
options["crf"] = str(quality_to_crf(quality, max_crf))
|
||||
if fmt == "h264":
|
||||
options["preset"] = "medium"
|
||||
elif fmt == "h265":
|
||||
options["preset"] = "fast"
|
||||
options["x265-params"] = "log-level=error"
|
||||
elif fmt == "av1":
|
||||
options["preset"] = "8"
|
||||
elif fmt == "ffv1":
|
||||
options["level"] = "3"
|
||||
options["slicecrc"] = "0"
|
||||
self.stream.options = options
|
||||
self.count = 0
|
||||
|
||||
def add(self, frame):
|
||||
frame = _pad_to(frame, self.width, self.height)
|
||||
vf = self._av.VideoFrame.from_ndarray(frame, format="rgb24")
|
||||
for packet in self.stream.encode(vf):
|
||||
self.container.mux(packet)
|
||||
self.count += 1
|
||||
|
||||
def bytes_so_far(self):
|
||||
return self.buffer.tell()
|
||||
|
||||
def close(self):
|
||||
for packet in self.stream.encode():
|
||||
self.container.mux(packet)
|
||||
self.container.close()
|
||||
return self.buffer.getvalue()
|
||||
|
||||
|
||||
def decode_video(data, width, height, expected_frames=None):
|
||||
av = _import_av()
|
||||
frames = []
|
||||
try:
|
||||
with av.open(io.BytesIO(data), mode="r") as container:
|
||||
for vf in container.decode(video=0):
|
||||
arr = vf.to_ndarray(format="rgb24")
|
||||
frames.append(np.ascontiguousarray(arr[:height, :width]))
|
||||
except TransferCodecError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise TransferCodecError(f"Failed to decode video chunk: {exc}") from exc
|
||||
if expected_frames is not None and len(frames) != expected_frames:
|
||||
raise TransferCodecError(
|
||||
f"Video chunk decoded to {len(frames)} frames, expected {expected_frames}"
|
||||
)
|
||||
return frames
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Chunking
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes=None):
|
||||
"""Yield transfer chunks for a uint8 frame batch [B,H,W,3].
|
||||
|
||||
Each chunk is a dict:
|
||||
{"format", "start_index", "count", "width", "height", "parts": [bytes, ...]}
|
||||
Image formats carry one part per frame; video formats carry one part
|
||||
(a Matroska file) holding `count` frames.
|
||||
|
||||
Size guarantees, measured on the finalized bytes (after encoder flush and
|
||||
muxer finalization, not on encoder buffers):
|
||||
|
||||
* a chunk holding more than one frame never exceeds `max_chunk_bytes`;
|
||||
* a chunk holding a single frame may exceed `max_chunk_bytes` but never
|
||||
`max_frame_bytes` (defaults to `max_chunk_bytes`).
|
||||
|
||||
A frame that cannot fit in `max_frame_bytes` even on its own raises
|
||||
FrameTooLargeError. Video chunks that come out too large are re-encoded
|
||||
with fewer frames.
|
||||
"""
|
||||
fmt = normalize_format(fmt)
|
||||
if fmt == LEGACY_FORMAT:
|
||||
raise TransferCodecError("legacy_png is sent through the JSON path, not encode_chunks")
|
||||
total = int(frames.shape[0])
|
||||
if total == 0:
|
||||
return
|
||||
height, width = int(frames.shape[1]), int(frames.shape[2])
|
||||
max_chunk_bytes = max(1, int(max_chunk_bytes))
|
||||
max_frame_bytes = max_chunk_bytes if max_frame_bytes is None else max(max_chunk_bytes, int(max_frame_bytes))
|
||||
|
||||
def _chunk(start, parts, count):
|
||||
return {
|
||||
"format": fmt,
|
||||
"start_index": start,
|
||||
"count": count,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"parts": parts,
|
||||
}
|
||||
|
||||
def _too_large(index, size):
|
||||
return FrameTooLargeError(
|
||||
f"Frame {index} encodes to {size} bytes as {fmt}, above the {max_frame_bytes}-byte request limit. "
|
||||
"Use a lossy format or lower quality, or raise the master's --max-upload-size together with "
|
||||
"COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES."
|
||||
)
|
||||
|
||||
if fmt in IMAGE_FORMATS:
|
||||
start, parts, size = 0, [], 0
|
||||
for i in range(total):
|
||||
data = encode_image_frame(frames[i], fmt, quality)
|
||||
if len(data) > max_frame_bytes:
|
||||
raise _too_large(i, len(data))
|
||||
if parts and size + len(data) > max_chunk_bytes:
|
||||
yield _chunk(start, parts, len(parts))
|
||||
start, parts, size = i, [], 0
|
||||
parts.append(data)
|
||||
size += len(data)
|
||||
if parts:
|
||||
yield _chunk(start, parts, len(parts))
|
||||
return
|
||||
|
||||
start = 0
|
||||
budget = MAX_VIDEO_CHUNK_FRAMES
|
||||
while start < total:
|
||||
want = min(budget, total - start)
|
||||
data, count = _encode_video_range(frames, start, want, fmt, quality, width, height, max_chunk_bytes)
|
||||
if len(data) <= max_chunk_bytes or (count == 1 and len(data) <= max_frame_bytes):
|
||||
yield _chunk(start, [data], count)
|
||||
start += count
|
||||
per_frame = len(data) / count
|
||||
budget = max(1, min(MAX_VIDEO_CHUNK_FRAMES, int(max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / per_frame)))
|
||||
continue
|
||||
if count == 1:
|
||||
raise _too_large(start, len(data))
|
||||
# Finalized chunk overshot: retry this range with proportionally fewer
|
||||
# frames. Strictly fewer than `count`, so this always terminates.
|
||||
estimate = int(count * max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / len(data))
|
||||
budget = max(1, min(count - 1, estimate))
|
||||
|
||||
|
||||
def _encode_video_range(frames, start, count, fmt, quality, width, height, stop_after_bytes):
|
||||
"""Encode frames[start:start+count] as one finalized video.
|
||||
|
||||
Stops feeding frames early once the muxed output already exceeds
|
||||
`stop_after_bytes`, since the chunk is then known to be too large and will
|
||||
be re-encoded smaller; this bounds wasted work. Encoder lookahead means the
|
||||
pre-flush size underestimates the final size, so the caller must still
|
||||
check the returned (finalized) bytes. Returns (bytes, frames_encoded).
|
||||
"""
|
||||
writer = _VideoChunkWriter(fmt, quality, width, height)
|
||||
for i in range(start, start + count):
|
||||
writer.add(frames[i])
|
||||
if writer.bytes_so_far() > stop_after_bytes:
|
||||
break
|
||||
return writer.close(), writer.count
|
||||
|
||||
|
||||
def decode_chunk(envelope, parts):
|
||||
"""Decode one received chunk into a list of uint8 [H,W,3] frames."""
|
||||
fmt = normalize_format(envelope.get("format"))
|
||||
if fmt == LEGACY_FORMAT:
|
||||
raise TransferCodecError("legacy_png chunks are not sent as multipart")
|
||||
count = envelope.get("count")
|
||||
width = envelope.get("width")
|
||||
height = envelope.get("height")
|
||||
for name, value in (("count", count), ("width", width), ("height", height)):
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise TransferCodecError(f"envelope.{name}: expected positive integer")
|
||||
if not parts:
|
||||
raise TransferCodecError("chunk has no media parts")
|
||||
|
||||
if fmt in IMAGE_FORMATS:
|
||||
if len(parts) != count:
|
||||
raise TransferCodecError(f"expected {count} image parts, got {len(parts)}")
|
||||
frames = [decode_image_frame(p) for p in parts]
|
||||
else:
|
||||
if len(parts) != 1:
|
||||
raise TransferCodecError(f"expected 1 video part, got {len(parts)}")
|
||||
frames = decode_video(parts[0], width, height, expected_frames=count)
|
||||
|
||||
for frame in frames:
|
||||
if frame.shape[0] != height or frame.shape[1] != width:
|
||||
raise TransferCodecError(
|
||||
f"decoded frame is {frame.shape[1]}x{frame.shape[0]}, expected {width}x{height}"
|
||||
)
|
||||
return frames
|
||||
|
||||
|
||||
def describe_ratio(sent_bytes, raw_bytes):
|
||||
if sent_bytes <= 0:
|
||||
return "n/a"
|
||||
ratio = raw_bytes / sent_bytes
|
||||
return f"{ratio:.1f}x" if math.isfinite(ratio) else "n/a"
|
||||
+672
-40
@@ -1,32 +1,389 @@
|
||||
"""Backward-compatibility shim for USDU helpers.
|
||||
import asyncio
|
||||
import time
|
||||
import json
|
||||
import copy
|
||||
import os
|
||||
import io
|
||||
from aiohttp import web, ClientTimeout
|
||||
import server
|
||||
from PIL import Image
|
||||
|
||||
Route handlers and job logic now live in:
|
||||
- upscale.job_store
|
||||
- upscale.job_timeout
|
||||
- upscale.payload_parsers
|
||||
- api.usdu_routes
|
||||
"""
|
||||
# Import from other utilities
|
||||
from .logging import debug_log, log
|
||||
from .network import handle_api_error, get_client_session
|
||||
# We avoid converting to tensors on the master for tiles; blending uses PIL
|
||||
|
||||
from ..upscale.conditioning import clone_conditioning, clone_control_chain
|
||||
from ..upscale.job_store import (
|
||||
MAX_PAYLOAD_SIZE,
|
||||
_cleanup_job,
|
||||
_drain_results_queue,
|
||||
_get_completed_count,
|
||||
_init_job_queue,
|
||||
_mark_task_completed,
|
||||
ensure_tile_jobs_initialized,
|
||||
init_dynamic_job,
|
||||
init_static_job_batched,
|
||||
)
|
||||
from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
|
||||
from ..upscale.payload_parsers import _parse_tiles_from_form
|
||||
from .logging import debug_log
|
||||
from .network import get_client_session
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
# Import HEARTBEAT_TIMEOUT from constants
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
from .config import load_config
|
||||
|
||||
|
||||
def _parse_tiles_from_form(data):
|
||||
"""Parse tiles submitted via multipart/form-data into a list of tile dicts.
|
||||
|
||||
Expects the following fields in the aiohttp form data:
|
||||
- 'tiles_metadata': JSON list with per-tile metadata items containing at least
|
||||
'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
|
||||
'batch_idx' and 'global_idx' are included when available.
|
||||
- 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
|
||||
- 'padding': integer padding used during extraction (optional; defaults 0).
|
||||
|
||||
Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y',
|
||||
'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
|
||||
plus 'padding'.
|
||||
"""
|
||||
try:
|
||||
# Parse padding if present
|
||||
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
|
||||
except Exception:
|
||||
padding = 0
|
||||
|
||||
# Parse tiles metadata (JSON list)
|
||||
meta_raw = data.get('tiles_metadata')
|
||||
if meta_raw is None:
|
||||
raise ValueError("Missing tiles_metadata")
|
||||
|
||||
try:
|
||||
metadata = json.loads(meta_raw)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
|
||||
|
||||
if not isinstance(metadata, list):
|
||||
raise ValueError("tiles_metadata must be a list")
|
||||
|
||||
tiles = []
|
||||
# Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
|
||||
for i, meta in enumerate(metadata):
|
||||
file_field = data.get(f'tile_{i}')
|
||||
if file_field is None or not hasattr(file_field, 'file'):
|
||||
raise ValueError(f"Missing tile data for index {i}")
|
||||
|
||||
# Read image bytes and decode to PIL
|
||||
raw = file_field.file.read()
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid image data for tile {i}: {e}")
|
||||
|
||||
# Build tile dictionary (store PIL only; master blends via PIL)
|
||||
try:
|
||||
tile_info = {
|
||||
'image': img,
|
||||
'tile_idx': int(meta.get('tile_idx', i)),
|
||||
'x': int(meta.get('x', 0)),
|
||||
'y': int(meta.get('y', 0)),
|
||||
'extracted_width': int(meta.get('extracted_width', img.width)),
|
||||
'extracted_height': int(meta.get('extracted_height', img.height)),
|
||||
'padding': int(padding),
|
||||
}
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
|
||||
|
||||
# Optional fields
|
||||
if 'batch_idx' in meta:
|
||||
try:
|
||||
tile_info['batch_idx'] = int(meta['batch_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
if 'global_idx' in meta:
|
||||
try:
|
||||
tile_info['global_idx'] = int(meta['global_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tiles.append(tile_info)
|
||||
|
||||
return tiles
|
||||
|
||||
|
||||
# Unified Job Data Structure Keys
|
||||
JOB_QUEUE = 'queue'
|
||||
JOB_MODE = 'mode'
|
||||
JOB_COMPLETED_TASKS = 'completed_tasks'
|
||||
JOB_WORKER_STATUS = 'worker_status'
|
||||
JOB_ASSIGNED_TO_WORKERS = 'assigned_to_workers'
|
||||
JOB_PENDING_TASKS = 'pending_tasks'
|
||||
JOB_BATCH_SIZE = 'batch_size' # For dynamic
|
||||
JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
|
||||
|
||||
# Task Types
|
||||
TASK_TYPE_TILE = 'tile'
|
||||
TASK_TYPE_IMAGE = 'image'
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
|
||||
"""Initialize queue for dynamic mode (per-image), with collector fields.
|
||||
|
||||
- Creates JOB_PENDING_TASKS with image indices
|
||||
- Adds 'completed_images' dict and 'pending_images' alias used by collectors
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'dynamic',
|
||||
batch_size=batch_size,
|
||||
all_indices=all_indices or list(range(batch_size)),
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
job_data['completed_images'] = {}
|
||||
job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
|
||||
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
|
||||
|
||||
|
||||
async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
|
||||
"""Initialize queue for static mode (batched-per-tile).
|
||||
|
||||
- Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1]
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
enabled_workers=enabled_workers,
|
||||
batched_static=True,
|
||||
)
|
||||
# Initialization handled by master; avoid duplicate init logs here
|
||||
|
||||
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False):
|
||||
"""Unified initialization for job queues in static and dynamic modes."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
debug_log(f"Queue already exists for {multi_job_id}")
|
||||
return
|
||||
|
||||
job_data = {
|
||||
JOB_QUEUE: asyncio.Queue(),
|
||||
JOB_MODE: mode,
|
||||
JOB_COMPLETED_TASKS: {},
|
||||
JOB_WORKER_STATUS: {w: time.time() for w in enabled_workers or []},
|
||||
JOB_ASSIGNED_TO_WORKERS: {w: [] for w in enabled_workers or []},
|
||||
JOB_PENDING_TASKS: asyncio.Queue(),
|
||||
}
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data[JOB_BATCH_SIZE] = batch_size
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
for i in (all_indices or range(batch_size)):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized image queue with {batch_size} pending items")
|
||||
elif mode == 'static':
|
||||
job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
|
||||
job_data[JOB_BATCH_SIZE] = batch_size
|
||||
job_data['batched_static'] = bool(batched_static)
|
||||
# For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
if batched_static and num_tiles_per_image is not None:
|
||||
for i in range(num_tiles_per_image):
|
||||
await pending_queue.put(i)
|
||||
else:
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
for i in range(total_tiles):
|
||||
await pending_queue.put(i)
|
||||
|
||||
# Keep backward compatibility - if task assignments provided, still track them
|
||||
if task_assignments and enabled_workers:
|
||||
# task_assignments[0] is for master, 1+ are for workers
|
||||
for i, worker_id in enumerate(enabled_workers):
|
||||
if i + 1 < len(task_assignments):
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id] = task_assignments[i + 1]
|
||||
debug_log(f"Worker {worker_id} pre-assigned {len(task_assignments[i + 1])} tasks (legacy mode)")
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
# Note: legacy task distribution and queue pull helpers removed
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained.
|
||||
|
||||
Uses non-blocking get_nowait to avoid await timeouts and reduce latency.
|
||||
"""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not job_data or JOB_QUEUE not in job_data or JOB_COMPLETED_TASKS not in job_data:
|
||||
return 0
|
||||
q = job_data[JOB_QUEUE]
|
||||
completed_tasks = job_data[JOB_COMPLETED_TASKS]
|
||||
|
||||
collected = 0
|
||||
while True:
|
||||
try:
|
||||
result = q.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat
|
||||
task_id = result.get('global_idx', result['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = {
|
||||
'tensor': result['tensor'],
|
||||
'tile_idx': result['tile_idx'],
|
||||
'x': result['x'],
|
||||
'y': result['y'],
|
||||
'extracted_width': result['extracted_width'],
|
||||
'extracted_height': result['extracted_height'],
|
||||
'padding': result['padding'],
|
||||
'batch_idx': result.get('batch_idx', 0),
|
||||
'global_idx': task_id
|
||||
}
|
||||
collected += 1
|
||||
|
||||
if is_last:
|
||||
# Track worker completion
|
||||
if worker_id in job_data[JOB_WORKER_STATUS]:
|
||||
del job_data[JOB_WORKER_STATUS][worker_id]
|
||||
|
||||
return collected
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
"""Check timed out workers and requeue their tasks. Returns requeued count."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not job_data:
|
||||
return 0
|
||||
|
||||
current_time = time.time()
|
||||
requeued_count = 0
|
||||
completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
|
||||
|
||||
# Allow override via config setting 'worker_timeout_seconds'
|
||||
cfg = load_config()
|
||||
hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT))
|
||||
|
||||
for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
|
||||
age = current_time - last_heartbeat
|
||||
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
|
||||
if age > hb_timeout:
|
||||
# Busy-only grace policy: require positive signal from worker (/prompt)
|
||||
# We also log assignment state for diagnostics but do not grace on it alone.
|
||||
assigned = job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, [])
|
||||
incomplete_assigned = 0
|
||||
try:
|
||||
if assigned:
|
||||
batched_static = bool(job_data.get('batched_static', False))
|
||||
if batched_static:
|
||||
num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
|
||||
batch_size = job_data.get(JOB_BATCH_SIZE, 1)
|
||||
for task_id in assigned:
|
||||
for b in range(batch_size):
|
||||
gidx = b * num_tiles_per_image + task_id
|
||||
if gidx not in completed_tasks:
|
||||
incomplete_assigned += 1
|
||||
break
|
||||
else:
|
||||
for task_id in assigned:
|
||||
if task_id not in completed_tasks:
|
||||
incomplete_assigned += 1
|
||||
debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}")
|
||||
except Exception as e:
|
||||
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
|
||||
|
||||
busy = False
|
||||
probe_status = None
|
||||
probe_queue = None
|
||||
try:
|
||||
cfg_workers = load_config().get('workers', [])
|
||||
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None)
|
||||
if wrec:
|
||||
host = wrec.get('host') or 'localhost'
|
||||
port = int(wrec.get('port', 8188))
|
||||
url = f"http://{host}:{port}/prompt"
|
||||
debug_log(f"Probing worker {worker} at {url}")
|
||||
session = await get_client_session()
|
||||
async with session.get(url, timeout=ClientTimeout(total=2.0)) as resp:
|
||||
probe_status = resp.status
|
||||
if resp.status == 200:
|
||||
try:
|
||||
payload = await resp.json()
|
||||
probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0))
|
||||
except Exception:
|
||||
probe_queue = 0
|
||||
busy = probe_queue is not None and probe_queue > 0
|
||||
except Exception as e:
|
||||
debug_log(f"Probe failed for worker {worker}: {e}")
|
||||
finally:
|
||||
debug_log(f"Probe diagnostics: http_status={probe_status} queue_remaining={probe_queue}")
|
||||
|
||||
if busy:
|
||||
job_data[JOB_WORKER_STATUS][worker] = current_time
|
||||
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
|
||||
continue
|
||||
|
||||
log(f"Worker {worker} timed out")
|
||||
for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []):
|
||||
# If batched_static, task_id is a tile_idx; consider it complete only if
|
||||
# all corresponding global_idx entries are present in completed_tasks.
|
||||
batched_static = bool(job_data.get('batched_static', False))
|
||||
if batched_static:
|
||||
num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
|
||||
batch_size = job_data.get(JOB_BATCH_SIZE, 1)
|
||||
# Check all global indices for this tile across the batch
|
||||
all_done = True
|
||||
for b in range(batch_size):
|
||||
gidx = b * num_tiles_per_image + task_id
|
||||
if gidx not in completed_tasks:
|
||||
all_done = False
|
||||
break
|
||||
if not all_done:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
else:
|
||||
# Legacy/global-idx mode: task_id is a global index key
|
||||
if task_id not in completed_tasks:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
del job_data[JOB_WORKER_STATUS][worker]
|
||||
if JOB_ASSIGNED_TO_WORKERS in job_data:
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker] = []
|
||||
|
||||
return requeued_count
|
||||
|
||||
async def _get_completed_count(multi_job_id):
|
||||
"""Get count of completed tasks."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if job_data and JOB_COMPLETED_TASKS in job_data:
|
||||
return len(job_data[JOB_COMPLETED_TASKS])
|
||||
return 0
|
||||
|
||||
async def _mark_task_completed(multi_job_id, task_id, result):
|
||||
"""Mark a task as completed."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if job_data and JOB_COMPLETED_TASKS in job_data:
|
||||
job_data[JOB_COMPLETED_TASKS][task_id] = result
|
||||
|
||||
async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
|
||||
"""Send heartbeat to master from worker-side processing loops."""
|
||||
"""Send heartbeat to master."""
|
||||
try:
|
||||
data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
|
||||
session = await get_client_session()
|
||||
@@ -36,20 +393,295 @@ async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
|
||||
except Exception as e:
|
||||
debug_log(f"Heartbeat failed: {e}")
|
||||
|
||||
async def _cleanup_job(multi_job_id):
|
||||
"""Cleanup the job data."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
debug_log(f"Cleaned up job {multi_job_id}")
|
||||
|
||||
# API Endpoints (generalized)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
|
||||
debug_log(f"Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Worker status tracking not available", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
||||
async def submit_tiles_endpoint(request):
|
||||
"""Endpoint for workers to submit processed tiles in static mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
# Handle completion signal
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
try:
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
||||
async def submit_image_endpoint(request):
|
||||
"""Endpoint for workers to submit processed images in dynamic mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Handle image submission
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle completion signal (no image data)
|
||||
elif is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [] # For compatibility
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Missing image data or invalid request", 400)
|
||||
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.
|
||||
|
||||
|
||||
|
||||
# Helper functions for shallow copying conditioning without duplicating models
|
||||
def clone_control_chain(control, clone_hint=True):
|
||||
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
|
||||
if control is None:
|
||||
return None
|
||||
new_control = copy.copy(control) # Shallow copy (shares model)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
hint = getattr(control, 'cond_hint_original', None)
|
||||
new_control.cond_hint_original = hint.clone() if hint is not None else None
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
|
||||
def clone_conditioning(cond_list, clone_hints=True):
|
||||
"""Clone conditioning without duplicating ControlNet models."""
|
||||
new_cond = []
|
||||
for emb, cond_dict in cond_list:
|
||||
new_emb = emb.clone() if emb is not None else None
|
||||
new_dict = cond_dict.copy()
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict:
|
||||
if new_dict['mask'] is not None:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
# Handle other potential fields if needed
|
||||
if 'pooled_output' in new_dict:
|
||||
if new_dict['pooled_output'] is not None:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple
|
||||
new_cond.append([new_emb, new_dict])
|
||||
return new_cond
|
||||
|
||||
# Helper function to ensure persistent state is initialized
|
||||
def ensure_tile_jobs_initialized():
|
||||
"""Ensure tile job storage is initialized on the server instance."""
|
||||
prompt_server = server.PromptServer.instance
|
||||
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
|
||||
debug_log("Initializing persistent tile job queue on server instance.")
|
||||
prompt_server.distributed_pending_tile_jobs = {}
|
||||
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
else:
|
||||
# Clean up any legacy queue structures that don't have the 'mode' field
|
||||
# (Should be rare after fixes, but keep for safety)
|
||||
to_remove = [job_id for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data]
|
||||
for job_id in to_remove:
|
||||
debug_log(f"Removing legacy queue structure for job {job_id}")
|
||||
del prompt_server.distributed_pending_tile_jobs[job_id]
|
||||
return prompt_server
|
||||
|
||||
# API Endpoint for tile completion
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, dict) or 'mode' not in job_data:
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
|
||||
mode = job_data['mode']
|
||||
|
||||
# Handle both dynamic and static modes
|
||||
if mode == 'dynamic' and 'pending_images' in job_data:
|
||||
pending_queue = job_data['pending_images']
|
||||
elif mode == 'static' and JOB_PENDING_TASKS in job_data:
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid job configuration", 400)
|
||||
|
||||
try:
|
||||
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
# Track assigned task
|
||||
if 'assigned_to_workers' in job_data and worker_id in job_data['assigned_to_workers']:
|
||||
job_data['assigned_to_workers'][worker_id].append(task_idx)
|
||||
# Update worker heartbeat
|
||||
if 'worker_status' in job_data:
|
||||
job_data['worker_status'][worker_id] = time.time()
|
||||
# Get estimated remaining count
|
||||
remaining = pending_queue.qsize() # Approximate
|
||||
|
||||
# Return appropriate response based on mode
|
||||
if mode == 'dynamic':
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
||||
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
||||
else: # static
|
||||
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
||||
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.get('batched_static', False)})
|
||||
except asyncio.TimeoutError:
|
||||
if mode == 'dynamic':
|
||||
return web.json_response({"image_idx": None})
|
||||
else:
|
||||
return web.json_response({"tile_idx": None})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
ready = bool(job_data and isinstance(job_data, dict) and 'queue' in job_data)
|
||||
return web.json_response({"ready": ready})
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_PAYLOAD_SIZE",
|
||||
"_check_and_requeue_timed_out_workers",
|
||||
"_cleanup_job",
|
||||
"_drain_results_queue",
|
||||
"_get_completed_count",
|
||||
"_init_job_queue",
|
||||
"_mark_task_completed",
|
||||
"_parse_tiles_from_form",
|
||||
"_send_heartbeat_to_master",
|
||||
"clone_conditioning",
|
||||
"clone_control_chain",
|
||||
"ensure_tile_jobs_initialized",
|
||||
"init_dynamic_job",
|
||||
"init_static_job_batched",
|
||||
]
|
||||
|
||||
+3
-9
@@ -466,13 +466,9 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
|
||||
x1_px, y1_px, x2_px, y2_px = region
|
||||
|
||||
new_latents = []
|
||||
for t in latents: # (B,C,H_lat_in,W_lat_in) or (B,C,1,H_lat_in,W_lat_in)
|
||||
has_5d = False
|
||||
if t.ndim == 5:
|
||||
has_5d = True
|
||||
t = t.squeeze(2)
|
||||
for t in latents: # (B,C,H_lat_in,W_lat_in)
|
||||
if t.ndim != 4:
|
||||
raise ValueError(f"expected BCHW or BC1HW, got {t.shape}")
|
||||
raise ValueError(f"expected BCHW, got {t.shape}")
|
||||
|
||||
# 1. Resize to canvas resolution in latent units only if needed
|
||||
if t.shape[-2:] != (H_can_lat, W_can_lat):
|
||||
@@ -495,8 +491,6 @@ def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size,
|
||||
mode="bilinear",
|
||||
align_corners=False)
|
||||
|
||||
if has_5d:
|
||||
cropped = cropped.unsqueeze(2)
|
||||
new_latents.append(cropped)
|
||||
|
||||
cond_dict["reference_latents"] = new_latents
|
||||
@@ -514,4 +508,4 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
|
||||
crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
cropped.append(n)
|
||||
return cropped
|
||||
return cropped
|
||||
@@ -1,8 +0,0 @@
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ["web/tests/**/*.test.js"],
|
||||
environment: "node",
|
||||
},
|
||||
});
|
||||
+29
-139
@@ -1,43 +1,27 @@
|
||||
import { TIMEOUTS } from './constants.js';
|
||||
import { normalizeWorkerUrl } from './urlUtils.js';
|
||||
|
||||
export function createApiClient(baseUrl) {
|
||||
const normalizedBaseUrl = normalizeWorkerUrl(baseUrl);
|
||||
|
||||
const request = async (
|
||||
endpoint,
|
||||
options = {},
|
||||
{ retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
|
||||
) => {
|
||||
const maxAttempts = retry ? retries : 1;
|
||||
const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => {
|
||||
let lastError;
|
||||
let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff
|
||||
|
||||
for (let attempt = 0; attempt < maxAttempts; attempt++) {
|
||||
for (let attempt = 0; attempt < retries; attempt++) {
|
||||
try {
|
||||
const headers = {
|
||||
'Content-Type': 'application/json',
|
||||
...(options.headers || {}),
|
||||
};
|
||||
const response = await fetch(`${normalizedBaseUrl}${endpoint}`, {
|
||||
...options,
|
||||
headers,
|
||||
const response = await fetch(`${baseUrl}${endpoint}`, {
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
...options
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
const message = error.message
|
||||
|| error.error
|
||||
|| (Array.isArray(error.errors) ? error.errors.join('; ') : null)
|
||||
|| `HTTP ${response.status}`;
|
||||
throw new Error(message);
|
||||
const error = await response.json().catch(() => ({ message: 'Request failed' }));
|
||||
throw new Error(error.message || `HTTP ${response.status}`);
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
|
||||
if (attempt < maxAttempts - 1) {
|
||||
console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`);
|
||||
if (attempt < retries - 1) {
|
||||
await new Promise(resolve => setTimeout(resolve, delay));
|
||||
delay *= 2; // Exponential backoff
|
||||
}
|
||||
@@ -45,39 +29,6 @@ export function createApiClient(baseUrl) {
|
||||
}
|
||||
throw lastError;
|
||||
};
|
||||
|
||||
const requestUrl = async (
|
||||
url,
|
||||
options = {},
|
||||
{ retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
|
||||
) => {
|
||||
const maxAttempts = retry ? retries : 1;
|
||||
let lastError;
|
||||
let delay = TIMEOUTS.RETRY_DELAY;
|
||||
|
||||
for (let attempt = 0; attempt < maxAttempts; attempt++) {
|
||||
try {
|
||||
const response = await fetch(url, options);
|
||||
if (!response.ok) {
|
||||
const error = await response.json().catch(() => ({}));
|
||||
const message = error.message
|
||||
|| error.error
|
||||
|| (Array.isArray(error.errors) ? error.errors.join('; ') : null)
|
||||
|| `HTTP ${response.status}`;
|
||||
throw new Error(message);
|
||||
}
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${url} - ${error.message}`);
|
||||
if (attempt < maxAttempts - 1) {
|
||||
await new Promise(resolve => setTimeout(resolve, delay));
|
||||
delay *= 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
};
|
||||
|
||||
return {
|
||||
// Config endpoints
|
||||
@@ -89,28 +40,28 @@ export function createApiClient(baseUrl) {
|
||||
return request('/distributed/config/update_worker', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId, ...data })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async deleteWorker(workerId) {
|
||||
return request('/distributed/config/delete_worker', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async updateSetting(key, value) {
|
||||
return request('/distributed/config/update_setting', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ key, value })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async updateMaster(data) {
|
||||
return request('/distributed/config/update_master', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(data)
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
// Worker management endpoints
|
||||
@@ -118,14 +69,14 @@ export function createApiClient(baseUrl) {
|
||||
return request('/distributed/launch_worker', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async stopWorker(workerId) {
|
||||
return request('/distributed/stop_worker', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async getManagedWorkers() {
|
||||
@@ -135,89 +86,28 @@ export function createApiClient(baseUrl) {
|
||||
async getWorkerLog(workerId, lines = 1000) {
|
||||
return request(`/distributed/worker_log/${workerId}?lines=${lines}`);
|
||||
},
|
||||
|
||||
async getRemoteWorkerLog(workerId, lines = 300) {
|
||||
return request(`/distributed/remote_worker_log/${workerId}?lines=${lines}`);
|
||||
},
|
||||
|
||||
async clearLaunchingFlag(workerId) {
|
||||
return request('/distributed/worker/clear_launching', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId })
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async queueDistributed(payload) {
|
||||
return request('/distributed/queue', {
|
||||
// Job preparation
|
||||
async prepareJob(multiJobId) {
|
||||
return request('/distributed/prepare_job', {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
...(payload?.trace_execution_id
|
||||
? { 'X-Idempotency-Key': payload.trace_execution_id }
|
||||
: {}),
|
||||
},
|
||||
body: JSON.stringify(payload)
|
||||
}, { retry: false });
|
||||
body: JSON.stringify({ multi_job_id: multiJobId })
|
||||
});
|
||||
},
|
||||
|
||||
async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK, signal = null) {
|
||||
const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
|
||||
const controller = new AbortController();
|
||||
const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
|
||||
const effectiveSignal = signal
|
||||
? AbortSignal.any([controller.signal, signal])
|
||||
: controller.signal;
|
||||
try {
|
||||
const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
|
||||
method: 'GET',
|
||||
mode: 'cors',
|
||||
cache: 'no-store',
|
||||
signal: effectiveSignal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
return { ok: false, status: response.status, queueRemaining: null };
|
||||
}
|
||||
|
||||
let data;
|
||||
try {
|
||||
data = await response.json();
|
||||
} catch {
|
||||
return { ok: false, status: response.status, queueRemaining: null };
|
||||
}
|
||||
|
||||
if (!data || typeof data !== "object" || Array.isArray(data)) {
|
||||
return { ok: false, status: response.status, queueRemaining: null };
|
||||
}
|
||||
|
||||
const execInfo = data.exec_info;
|
||||
if (!execInfo || typeof execInfo !== "object" || Array.isArray(execInfo)) {
|
||||
return { ok: false, status: response.status, queueRemaining: null };
|
||||
}
|
||||
|
||||
const rawQueueRemaining = execInfo.queue_remaining;
|
||||
const queueRemaining = Number(rawQueueRemaining);
|
||||
if (!Number.isFinite(queueRemaining)) {
|
||||
return { ok: false, status: response.status, queueRemaining: null };
|
||||
}
|
||||
|
||||
return {
|
||||
ok: true,
|
||||
status: response.status,
|
||||
queueRemaining: Math.max(0, queueRemaining),
|
||||
};
|
||||
} finally {
|
||||
clearTimeout(timeoutId);
|
||||
}
|
||||
},
|
||||
|
||||
async dispatchToWorker(workerUrl, promptPayload) {
|
||||
const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
|
||||
return requestUrl(`${normalizedWorkerUrl}/prompt`, {
|
||||
|
||||
// Image loading
|
||||
async loadImage(imagePath) {
|
||||
return request('/distributed/load_image', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
mode: 'cors',
|
||||
body: JSON.stringify(promptPayload),
|
||||
}, { retry: false });
|
||||
body: JSON.stringify({ image_path: imagePath })
|
||||
});
|
||||
},
|
||||
|
||||
// Network info
|
||||
@@ -258,14 +148,14 @@ export function createApiClient(baseUrl) {
|
||||
return request('/distributed/tunnel/start', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({})
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async stopTunnel() {
|
||||
return request('/distributed/tunnel/stop', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({})
|
||||
}, { retry: false });
|
||||
});
|
||||
},
|
||||
|
||||
async getTunnelStatus() {
|
||||
|
||||
+14
-83
@@ -17,7 +17,6 @@ export const BUTTON_STYLES = {
|
||||
launch: "background-color: #4a7c4a;",
|
||||
stop: "background-color: #7c4a4a;",
|
||||
log: "background-color: #685434;",
|
||||
working: "background-color: #666;",
|
||||
clearMemory: "background-color: #555; padding: 6px 14px;",
|
||||
interrupt: "background-color: #555; padding: 6px 14px;",
|
||||
};
|
||||
@@ -63,22 +62,6 @@ export const PULSE_ANIMATION_CSS = `
|
||||
animation: pulse 1.2s ease-in-out infinite;
|
||||
transform-origin: center;
|
||||
}
|
||||
|
||||
.worker-status--online {
|
||||
background: var(--status-online, #3ca03c) !important;
|
||||
}
|
||||
|
||||
.worker-status--offline {
|
||||
background: var(--status-offline, #c04c4c) !important;
|
||||
}
|
||||
|
||||
.worker-status--unknown {
|
||||
background: var(--status-unknown, #888) !important;
|
||||
}
|
||||
|
||||
.worker-status--processing {
|
||||
background: var(--status-processing, #f0ad4e) !important;
|
||||
}
|
||||
|
||||
/* Button hover effects */
|
||||
.distributed-button:hover:not(:disabled) {
|
||||
@@ -106,9 +89,7 @@ export const PULSE_ANIMATION_CSS = `
|
||||
.worker-settings.expanded {
|
||||
max-height: 500px;
|
||||
opacity: 1;
|
||||
padding: 12px;
|
||||
margin-top: 8px;
|
||||
margin-bottom: 8px;
|
||||
padding: 12px 0;
|
||||
}
|
||||
|
||||
/* Cloudflare tunnel spinner */
|
||||
@@ -133,21 +114,20 @@ export const UI_STYLES = {
|
||||
statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
|
||||
controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;",
|
||||
formGroup: "display: flex; flex-direction: column; gap: 5px;",
|
||||
formLabel: "font-size: 12px; color: var(--dist-label-text, #ccc); font-weight: 500;",
|
||||
formInput:
|
||||
"padding: 6px 10px; color: var(--dist-input-text, white); background: var(--dist-input-bg, transparent); font-size: 12px; transition: border-color 0.2s;",
|
||||
formLabel: "font-size: 12px; color: #ccc; font-weight: 500;",
|
||||
formInput: "padding: 6px 10px; background: #2a2a2a; border: 1px solid #444; color: white; font-size: 12px; border-radius: 4px; transition: border-color 0.2s;",
|
||||
|
||||
// Card styles
|
||||
cardBase: "margin-bottom: 12px; overflow: hidden; display: flex;",
|
||||
workerCard: "margin-bottom: 12px; overflow: hidden; display: flex;",
|
||||
cardBlueprint: "cursor: pointer; transition: all 0.2s ease;",
|
||||
cardAdd: "cursor: pointer; transition: all 0.2s ease;",
|
||||
|
||||
cardBase: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex;",
|
||||
workerCard: "margin-bottom: 12px; border-radius: 6px; overflow: hidden; display: flex; background: #2a2a2a;",
|
||||
cardBlueprint: "border: 2px dashed #555; cursor: pointer; transition: all 0.2s ease; background: rgba(255, 255, 255, 0.02);",
|
||||
cardAdd: "border: 1px dashed #444; cursor: pointer; transition: all 0.2s ease; background: transparent;",
|
||||
|
||||
// Column styles
|
||||
columnBase: "display: flex; align-items: center; justify-content: center;",
|
||||
checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; cursor: default;",
|
||||
checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; border-right: 1px solid #3a3a3a; cursor: default; background: rgba(0,0,0,0.1);",
|
||||
contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
|
||||
iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: var(--dist-placeholder-add-color, #666);",
|
||||
iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: #666;",
|
||||
|
||||
// Row and content styles
|
||||
infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
|
||||
@@ -157,16 +137,14 @@ export const UI_STYLES = {
|
||||
buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
|
||||
settingsForm: "display: flex; flex-direction: column; gap: 10px;",
|
||||
checkboxGroup: "display: flex; align-items: center; gap: 8px; margin: 5px 0;",
|
||||
formLabelClickable: "font-size: 12px; color: var(--dist-label-text, #ccc); cursor: pointer;",
|
||||
formLabelClickable: "font-size: 12px; color: #ccc; cursor: pointer;",
|
||||
settingsToggle: "display: flex; align-items: center; gap: 6px; padding: 4px 0; cursor: pointer; user-select: none;",
|
||||
controlsWrapper: "display: flex; gap: 6px; align-items: stretch; width: 100%;",
|
||||
|
||||
// Existing styles
|
||||
settingsArrow:
|
||||
"font-size: 12px; color: var(--dist-settings-arrow, #888); transition: all 0.2s ease; margin-left: auto; padding: 4px;",
|
||||
infoBox:
|
||||
"color: var(--dist-info-box-text, #999); padding: 5px 14px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;",
|
||||
workerSettings: "margin: 0 12px; padding: 0 12px;"
|
||||
settingsArrow: "font-size: 12px; color: #888; transition: all 0.2s ease; margin-left: auto; padding: 4px;",
|
||||
infoBox: "background-color: #333; color: #999; padding: 5px 14px; border-radius: 4px; font-size: 11px; text-align: center; flex: 1; font-weight: 500;",
|
||||
workerSettings: "margin: 0 12px; padding: 0 12px; background: #1e1e1e; border-radius: 4px; border: 1px solid #2a2a2a;"
|
||||
};
|
||||
|
||||
export const TIMEOUTS = {
|
||||
@@ -190,50 +168,3 @@ export const TIMEOUTS = {
|
||||
LOG_REFRESH: 2000, // log auto-refresh interval
|
||||
IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache
|
||||
};
|
||||
|
||||
export const ENDPOINTS = {
|
||||
// ComfyUI core
|
||||
PROMPT: '/prompt',
|
||||
INTERRUPT: '/interrupt',
|
||||
UPLOAD_IMAGE: '/upload/image',
|
||||
SYSTEM_INFO: '/system_stats',
|
||||
|
||||
// Distributed API
|
||||
CONFIG: '/distributed/config',
|
||||
UPDATE_WORKER: '/distributed/config/update_worker',
|
||||
DELETE_WORKER: '/distributed/config/delete_worker',
|
||||
UPDATE_SETTING: '/distributed/config/update_setting',
|
||||
UPDATE_MASTER: '/distributed/config/update_master',
|
||||
LAUNCH_WORKER: '/distributed/launch_worker',
|
||||
STOP_WORKER: '/distributed/stop_worker',
|
||||
MANAGED_WORKERS: '/distributed/managed_workers',
|
||||
WORKER_LOG: '/distributed/worker_log',
|
||||
REMOTE_WORKER_LOG: '/distributed/remote_worker_log',
|
||||
LOCAL_LOG: '/distributed/local_log',
|
||||
CLEAR_LAUNCHING: '/distributed/worker/clear_launching',
|
||||
PREPARE_JOB: '/distributed/prepare_job',
|
||||
LOAD_IMAGE: '/distributed/load_image',
|
||||
NETWORK_INFO: '/distributed/network_info',
|
||||
CHECK_FILE: '/distributed/check_file',
|
||||
CLEAR_MEMORY: '/distributed/clear_memory',
|
||||
SYSTEM_INFO_DIST: '/distributed/system_info',
|
||||
TUNNEL_START: '/distributed/tunnel/start',
|
||||
TUNNEL_STOP: '/distributed/tunnel/stop',
|
||||
TUNNEL_STATUS: '/distributed/tunnel/status',
|
||||
};
|
||||
|
||||
export const NODE_CLASSES = {
|
||||
DISTRIBUTED_COLLECTOR: 'DistributedCollector',
|
||||
DISTRIBUTED_SEED: 'DistributedSeed',
|
||||
DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
|
||||
UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
|
||||
PREVIEW_IMAGE: 'PreviewImage',
|
||||
};
|
||||
|
||||
export function generateUUID() {
|
||||
if (crypto.randomUUID) return crypto.randomUUID();
|
||||
return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, c => {
|
||||
const r = Math.random() * 16 | 0;
|
||||
return (c === 'x' ? r : (r & 0x3 | 0x8)).toString(16);
|
||||
});
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user