Compare commits

...
Author SHA1 Message Date
Robert Wojciechowski a34969e6f3 fix(image-list): own job lifecycle and preserve worker preparation 2026-10-11 21:20:22 +00:00
Robert Wojciechowski 8e0ac3befe fix(media-sync): preserve output and temp source roots 2026-10-11 21:20:22 +00:00
Robert Wojciechowski eac29590bb fix(image-list): bound acknowledgements and configure run deadline 2026-10-11 21:07:19 +00:00
Robert Wojciechowski bac16afaad fix(image-list): isolate unavailable workers during preflight 2026-10-11 20:57:20 +00:00
Robert Wojciechowski 9fc8e3586a fix(image-list): make optional End text a connection-only input 2026-10-11 20:57:20 +00:00
Robert Wojciechowski 8cbf531448 fix(image-list): reject upstream data bypasses around Start 2026-10-11 20:52:12 +00:00
Robert Wojciechowski 3e280a0e29 feat: add dynamic per-image list queue with Start and End nodes 2026-10-11 08:49:23 +00:00
Robert Wojciechowski 87c0cda9f6 fix: preserve serverless capacity and recovery across job failures
Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.
2026-10-11 06:56:01 +00:00
Robert Wojciechowski 371517ae14 fix: match serverless endpoint box to master card 2026-10-11 06:10:44 +00:00
Robert Wojciechowski 6eb94a52fe fix: simplify serverless cards and keep status dots circular 2026-10-11 06:01:46 +00:00
Robert Wojciechowski a281d03c6d fix: validate complete serverless requests against 10 MB limit 2026-10-11 05:52:23 +00:00
Robert Wojciechowski d9485265ea fix: initialize worker settings before mounting the form 2026-10-11 05:24:39 +00:00
Robert Wojciechowski cd13521d24 fix: prevent worker settings clipping and remove preview copy 2026-10-11 05:06:01 +00:00
Robert Wojciechowski 1093818327 feat: add Runpod serverless image collector pools 2026-10-11 04:55:01 +00:00
Robert Wojciechowski 6f85129e45 Bump version from 1.4.7 to 1.5.0 2026-10-11 10:06:41 +11:00
Robert Wojciechowski a53d88f7fb test: remove redundant coverage and consolidate happy paths (#95) 2026-10-10 16:00:41 -07:00
Robert Wojciechowski 9be342f5fb feat: migrate distributed nodes to native ComfyUI V3 API (#94)
* feat: migrate distributed nodes to native ComfyUI V3 API

* fix: hide Distributed Value backing fields in Nodes 2.0

Use native widget visibility without converting backing fields into sockets. Preserve widget types, serializers and options identity; cover edits, reload, worker refresh and linked widgets.
2026-10-10 16:00:36 -07:00
Robert Wojciechowski a1faf7e916 fix: isolate local worker databases and avoid port conflicts (#93)
* fix: isolate local worker databases and avoid port conflicts

* fix: harden local worker startup preflight and database paths

Handle auto-launch before the master binds, allow TIME_WAIT restarts without accepting active listeners, resolve relative paths in the worker cwd, and round-trip SQLite URLs through SQLAlchemy. Add regressions and remove the README additions.
2026-10-10 16:00:31 -07:00
Robert Wojciechowski 32ac027e0e docs: align selected copy with current behavior 2026-07-12 11:13:22 +00:00
Robert Wojciechowski fec959e1ac chore: bump version to 1.4.7 2026-07-12 10:40:27 +00:00
Robert Wojciechowski cedc8d4591 feat: support audio-only distributed workflows (#87)
Support image, audio, or combined distributed collection with guarded audio-only worker, master, endpoint, preview, and delegate-only paths.
2026-07-12 20:36:43 +10:00
Robert Wojciechowski db01824d6e Update pyproject.toml 2026-06-06 09:10:43 +10:00
Robert Wojciechowski 6e7818bee6 fix: preserve delegate-only master config inputs (#86)
Fixes #85
2026-06-05 17:12:25 +10:00
Robert Wojciechowski c246d965e7 Update pyproject.toml 2026-05-28 09:00:17 +10:00
Robert Wojciechowski d6ce26f143 Merge pull request #84 from robertvoy/fix/issue-83
Fix collector handling of ComfyUI list inputs
2026-05-28 08:59:02 +10:00
101 changed files with 14014 additions and 608 deletions
+9 -3
View File
@@ -27,6 +27,11 @@
- Intelligent distribution
- Handles single images and videos
#### Dynamic Image-List Processing
- Schedule one image at a time to the next available direct worker; collect results in input order
- Use **Distributed Image List Start / End** around a per-image processing region
- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
- Convert any workflow to distributed with 2 nodes
@@ -38,11 +43,12 @@
<img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
ComfyUI Distributed supports three types of workers:
ComfyUI Distributed supports four types of workers:
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
- **Remote Workers** - GPUs on other computers in your network
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
- **Serverless (Runpod)** - Queue-based endpoint pools for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
@@ -147,7 +153,7 @@ 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 standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_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).
@@ -175,7 +181,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
| **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 Batch Divider** | Splits audio 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 |
+15 -22
View File
@@ -1,29 +1,22 @@
# 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
)
"""ComfyUI-Distributed's native V3 extension entrypoint."""
from comfy_api.v0_0_2 import ComfyExtension, io
# Import utilities
from .utils.config import ensure_config_exists, CONFIG_FILE
from .utils.logging import debug_log
from .nodes.v3 import NODES
from .runtime.bootstrap import initialize
# Import distributed upscale nodes
from .nodes.distributed_upscale import (
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
)
WEB_DIRECTORY = './web'
WEB_DIRECTORY = "./web"
ensure_config_exists()
class DistributedExtension(ComfyExtension):
async def on_load(self) -> None:
initialize()
# Merge node mappings
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return list(NODES)
__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())}")
async def comfy_entrypoint() -> DistributedExtension:
return DistributedExtension()
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
+2
View File
@@ -1,5 +1,7 @@
from . import config_routes # noqa: F401
from . import serverless_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
from . import image_list_routes # noqa: F401
+372 -69
View File
@@ -1,5 +1,6 @@
import json
from contextlib import asynccontextmanager
from typing import Any
from aiohttp import web
import server
@@ -24,6 +25,12 @@ except ImportError:
save_config(config)
from ..utils.logging import debug_log, log
from ..utils.network import handle_api_error, normalize_host
from ..utils.secret_fields import is_secret_field_name
from ..utils.serverless_config import (
ServerlessSecretChanges,
public_config,
validate_serverless_worker,
)
def _positive_int(value):
@@ -86,22 +93,160 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
target[key] = normalizer(value) if (normalizer and value is not None) else value
_SECRET_FIELDS = {"api_key", "clear_api_key"}
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
_SERVERLESS_ONLY_FIELDS = {
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
}
_SERVERLESS_CONFIG_FIELDS = {
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
}
class _ActivePoolCheckError(RuntimeError):
pass
def _error_response(message, status):
return web.json_response({"status": "error", "message": message}, status=status)
def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0):
if isinstance(value, dict):
for key, item in value.items():
field_name = key.casefold() if isinstance(key, str) else ""
if is_secret_field_name(field_name):
if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS):
return True
if _contains_secret_fields(
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
):
return True
elif isinstance(value, (list, tuple)):
return any(
_contains_secret_fields(
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
)
for item in value
)
return False
def _serverless_secret_store():
from .orchestration.serverless_runtime import get_secret_store
return get_secret_store()
def _serverless_pool_is_active(worker_id):
from .orchestration.serverless_runtime import get_manager
return get_manager().pool_is_active(worker_id)
def _ensure_pool_idle(worker_id):
try:
active = _serverless_pool_is_active(worker_id)
except Exception:
raise _ActivePoolCheckError from None
if active:
raise _ActivePoolCheckError
def _worker_by_id(workers, worker_id):
for worker in workers:
if isinstance(worker, dict) and worker.get("id") == worker_id:
return worker
return None
def _worker_type(worker):
return worker.get("type", "local") if isinstance(worker, dict) else None
def _strip_sensitive_fields(value: Any) -> Any:
if isinstance(value, dict):
return {
key: _strip_sensitive_fields(item)
for key, item in value.items()
if not is_secret_field_name(key)
}
if isinstance(value, list):
return [_strip_sensitive_fields(item) for item in value]
return value
def _clean_serverless_worker(worker):
clean = _strip_sensitive_fields(dict(worker))
for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}:
clean.pop(field, None)
return clean
def _validate_bulk_workers(workers):
validated = []
for worker in workers:
if isinstance(worker, dict) and worker.get("type") == "serverless":
clean = _clean_serverless_worker(worker)
validated.append(validate_serverless_worker(clean))
else:
validated.append(worker)
return validated
def _clean_bulk_type_conversions(old_workers, new_workers):
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
for index, worker in enumerate(new_workers):
if not isinstance(worker, dict) or worker.get("id") is None:
continue
old = old_by_id.get(str(worker["id"]))
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
continue
converted = dict(worker)
for field in _SERVERLESS_ONLY_FIELDS:
converted.pop(field, None)
new_workers[index] = converted
return new_workers
def _bulk_serverless_changes(old_workers, new_workers):
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
secret_ids = set()
for worker_id, old in old_by_id.items():
new = new_by_id.get(worker_id)
if old.get("type") == "serverless":
if new is None or new.get("type") != "serverless":
_ensure_pool_idle(worker_id)
secret_ids.add(worker_id)
elif (old.get("endpoint_id") != new.get("endpoint_id")
or old.get("concurrency", 1) != new.get("concurrency", 1)):
_ensure_pool_idle(worker_id)
elif new is not None and new.get("type") == "serverless":
secret_ids.add(worker_id)
return secret_ids
@server.PromptServer.instance.routes.get("/distributed/config")
async def get_config_endpoint(request):
config = load_config()
return web.json_response(config)
try:
config = load_config()
return web.json_response(public_config(config, _serverless_secret_store()))
except Exception:
return _error_response("Unable to read config", 500)
@server.PromptServer.instance.routes.post("/distributed/config")
async def update_config_endpoint(request):
"""Bulk config update with schema validation."""
"""Bulk config update with schema validation and credential redaction."""
try:
data = await request.json()
except Exception as e:
return await handle_api_error(request, f"Invalid JSON payload: {e}", 400)
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return await handle_api_error(request, "Config payload must be an object", 400)
return _error_response("Config payload must be an object", 400)
if _contains_secret_fields(data):
return _error_response("API keys cannot be set through bulk config", 400)
validated_settings = {}
validated_root = {}
@@ -118,9 +263,15 @@ async def update_config_endpoint(request):
continue
if validator and not validator(value):
errors.append(f"{key}: value {value!r} failed validation")
errors.append(f"{key}: failed validation")
continue
if key == "workers":
try:
value = _validate_bulk_workers(value)
except (TypeError, ValueError):
return _error_response("Invalid serverless worker configuration", 400)
if key in _SETTINGS_FIELDS:
validated_settings[key] = value
else:
@@ -133,15 +284,41 @@ async def update_config_endpoint(request):
"message": "; ".join(errors),
}, status=400)
store = None
response_config = None
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)
with ServerlessSecretChanges() as secret_changes:
async with config_transaction() as config:
if "workers" in validated_root:
old_workers = config.get("workers", [])
validated_root["workers"] = _clean_bulk_type_conversions(
old_workers, validated_root["workers"]
)
changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"])
for worker in validated_root["workers"]:
if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False):
store = store or _serverless_secret_store()
if worker["id"] in changed_ids or not store.has(worker["id"]):
return _error_response("Save an API key before enabling this pool", 400)
if changed_ids:
store = store or _serverless_secret_store()
for worker_id in sorted(changed_ids):
prior_secret = secret_changes.capture(store, worker_id)
if prior_secret is not None:
secret_changes.delete(store, worker_id)
settings = config.setdefault("settings", {})
settings.update(validated_settings)
for key, value in validated_root.items():
config[key] = value
response_config = public_config(config, store or _serverless_secret_store())
response = web.json_response({"status": "success", "config": response_config})
secret_changes.commit()
return response
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
except Exception:
return _error_response("Unable to update config", 500)
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
@@ -164,79 +341,205 @@ async def queue_status_endpoint(request):
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
async def update_worker_endpoint(request):
return await _run_with_serverless_secret_changes(_update_worker_endpoint, request)
async def _run_with_serverless_secret_changes(endpoint, request):
with ServerlessSecretChanges() as secret_changes:
response = await endpoint(request, secret_changes)
if response.status < 400:
secret_changes.commit()
return response
async def _update_worker_endpoint(request, secret_changes):
store = None
worker_id = None
try:
data = await request.json()
try:
data = await request.json()
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return _error_response("Worker payload must be an object", 400)
worker_id = data.get("worker_id")
if worker_id is None:
return await handle_api_error(request, "Missing worker_id", 400)
return _error_response("Missing worker_id", 400)
if _contains_secret_fields(data, allow_root_serverless_secrets=True):
return _error_response("Invalid worker credentials", 400)
result = None
async with config_transaction() as config:
worker_found = False
workers = config.setdefault("workers", [])
existing = _worker_by_id(workers, worker_id)
current_type = _worker_type(existing)
target_type = data.get("type", current_type or "local")
for worker in workers:
if worker["id"] == worker_id:
_apply_field_patch(worker, data, _WORKER_FIELDS)
worker_found = True
break
if target_type == "serverless":
unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS
if unknown:
return _error_response("Unknown serverless worker field", 400)
if not isinstance(worker_id, str) or not worker_id.strip():
return _error_response("worker_id must be a non-empty string", 400)
if "api_key" in data and not isinstance(data["api_key"], str):
return _error_response("Invalid worker credentials", 400)
if "clear_api_key" in data and type(data["clear_api_key"]) is not bool:
return _error_response("Invalid worker credentials", 400)
supplied_key = data.get("api_key", "")
clear_key = data.get("clear_api_key", False)
if supplied_key and clear_key:
return _error_response("Choose either api_key or clear_api_key", 400)
if "enabled" in data and type(data["enabled"]) is not bool:
return _error_response("enabled must be a boolean", 400)
if "name" in data and not isinstance(data["name"], str):
return _error_response("name must be a string", 400)
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)
candidate = dict(existing or {})
candidate["id"] = worker_id
candidate["type"] = "serverless"
candidate.setdefault("provider", "runpod")
candidate.setdefault("enabled", False)
for field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "enabled", "name"):
if field in data:
candidate[field] = data[field]
candidate = _clean_serverless_worker(candidate)
validated = validate_serverless_worker(candidate)
if "enabled" in validated and type(validated["enabled"]) is not bool:
return _error_response("enabled must be a boolean", 400)
if "name" in validated and not isinstance(validated["name"], str):
return _error_response("name must be a string", 400)
store = _serverless_secret_store()
prior_secret = secret_changes.capture(store, worker_id)
converting_to_serverless = current_type != "serverless"
next_secret = None if converting_to_serverless else prior_secret
if supplied_key:
next_secret = supplied_key
elif clear_key:
next_secret = None
endpoint_changed = (
current_type == "serverless"
and existing.get("endpoint_id") != validated.get("endpoint_id")
)
concurrency_changed = (
current_type == "serverless"
and existing.get("concurrency", 1) != validated.get("concurrency", 1)
)
if current_type == "serverless" and (
endpoint_changed or concurrency_changed or next_secret != prior_secret
):
_ensure_pool_idle(worker_id)
if validated.get("enabled", False) and next_secret is None:
return _error_response("Save an API key before enabling this pool; disable it before clearing the key", 400)
if next_secret != prior_secret:
if next_secret is None:
secret_changes.delete(store, worker_id)
else:
secret_changes.set(store, worker_id, next_secret)
if existing is None:
workers.append(validated)
else:
return await handle_api_error(
request,
f"Worker {worker_id} not found and missing required fields for creation",
404,
)
existing.clear()
existing.update(validated)
result = {"status": "success", "api_key_configured": next_secret is not None}
elif existing is not None:
cleaned_existing = _strip_sensitive_fields(existing)
existing.clear()
existing.update(cleaned_existing)
if current_type == "serverless":
if _contains_secret_fields(data):
return _error_response("Invalid worker credentials", 400)
_ensure_pool_idle(worker_id)
store = _serverless_secret_store()
prior_secret = secret_changes.capture(store, worker_id)
if prior_secret is not None:
secret_changes.delete(store, worker_id)
for field in _SERVERLESS_ONLY_FIELDS:
existing.pop(field, None)
elif _contains_secret_fields(data):
return _error_response("API keys are only accepted for serverless workers", 400)
_apply_field_patch(existing, data, _WORKER_FIELDS)
result = {"status": "success"}
elif all(key in data for key in ["name", "port", "cuda_device"]):
if _contains_secret_fields(data):
return _error_response("API keys are only accepted for serverless workers", 400)
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)
result = {"status": "success"}
else:
return _error_response(
"Worker not found and missing required fields for creation", 404
)
return web.json_response(result)
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
except ValueError:
return _error_response("Invalid worker configuration", 400)
except Exception:
return _error_response("Unable to update worker configuration", 500)
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):
return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
async def _delete_worker_endpoint(request, secret_changes):
store = None
worker_id = None
try:
data = await request.json()
try:
data = await request.json()
except Exception:
return _error_response("Invalid JSON payload", 400)
if not isinstance(data, dict):
return _error_response("Worker payload must be an object", 400)
worker_id = data.get("worker_id")
if worker_id is None:
return await handle_api_error(request, "Missing worker_id", 400)
return _error_response("Missing worker_id", 400)
result = None
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({
removed = _worker_by_id(workers, worker_id)
if removed is None:
return _error_response("Worker not found", 404)
if _worker_type(removed) == "serverless":
_ensure_pool_idle(worker_id)
store = _serverless_secret_store()
secret_worker_id = str(worker_id)
prior_secret = secret_changes.capture(store, secret_worker_id)
if prior_secret is not None:
secret_changes.delete(store, secret_worker_id)
workers.remove(removed)
result = {
"status": "success",
"message": f"Worker {removed_worker.get('name', worker_id)} deleted"
})
except Exception as e:
return await handle_api_error(request, e, 400)
"message": "Worker deleted" if _worker_type(removed) == "serverless"
else f"Worker {removed.get('name', worker_id)} deleted",
}
return web.json_response(result)
except _ActivePoolCheckError:
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
except Exception:
return _error_response("Unable to delete worker", 500)
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
async def update_setting_endpoint(request):
+88
View File
@@ -0,0 +1,88 @@
"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
import json
from aiohttp import web
import server
from ..image_list.runtime import (
MAX_REQUEST_BYTES,
decode_result,
delete_worker_item,
set_master_result,
store_worker_item,
)
prompt_server = server.PromptServer.instance
async def _json_body(request):
length = request.content_length
if length is not None and length > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
raw = bytearray()
async for chunk in request.content.iter_chunked(64 * 1024):
raw.extend(chunk)
if len(raw) > MAX_REQUEST_BYTES:
raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
try:
return json.loads(raw.decode("utf-8"))
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
@prompt_server.routes.get("/distributed/image-list/capabilities")
async def image_list_capabilities_endpoint(_request):
return web.json_response({
"supported": True,
"protocol_version": 1,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
"transport": "typed-float-tensor-json-v1",
})
@prompt_server.routes.post("/distributed/image-list/item")
async def image_list_item_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list item object")
context_id = data.get("context_id")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
item = data.get("item")
store_worker_item(data.get("token"), item)
return web.json_response({"status": "stored", "item_index": item["index"]})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
@prompt_server.routes.delete("/distributed/image-list/item/{token}")
async def image_list_delete_item_endpoint(request):
token = request.match_info.get("token", "")
return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
@prompt_server.routes.post("/distributed/image-list/result")
async def image_list_result_endpoint(request):
try:
data = await _json_body(request)
if not isinstance(data, dict):
raise ValueError("Expected an image-list result object")
context_id = data.get("context_id")
token = data.get("token")
item_index = data.get("item_index")
if not isinstance(context_id, str) or not context_id:
raise ValueError("context_id must be a non-empty string")
if not isinstance(token, str) or not token:
raise ValueError("token must be a non-empty string")
if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
raise ValueError("item_index must be a non-negative integer")
result = decode_result(data.get("result"))
set_master_result(context_id, item_index, token, result)
return web.json_response({"status": "accepted", "item_index": item_index})
except (ValueError, TypeError) as exc:
return web.json_response({"error": str(exc)}, status=400)
except Exception as exc:
return web.json_response({"error": str(exc)}, status=500)
+11 -4
View File
@@ -295,17 +295,24 @@ async def job_complete_endpoint(request):
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")
if not isinstance(image_payload, str) or not image_payload.strip():
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_payload is not None and not isinstance(audio_payload, dict):
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)
tensor = _decode_canonical_png_tensor(image_payload)
decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
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()
+140 -22
View File
@@ -55,6 +55,11 @@ async def worker_ws_is_active(worker):
async def _probe_worker_active(worker, use_websocket, semaphore):
async with semaphore:
if worker.get("type") == "serverless":
from .serverless_runtime import get_manager
# Authentication/configuration failures must surface, not silently skip paid work.
status = await get_manager().health(worker)
return worker, status["ok"]
is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
return worker, is_active
@@ -72,27 +77,44 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
ws_url = f"{ws_url}/distributed/worker_ws"
session = await get_client_session()
loop = asyncio.get_running_loop()
deadline = loop.time() + timeout
ws = None
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.")
try:
async with asyncio.timeout_at(deadline):
ws = await session.ws_connect(ws_url, heartbeat=20)
await ws.send_json(ws_payload)
while True:
msg = await ws.receive()
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 data.get("prompt_id")
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.CLOSE, aiohttp.WSMsgType.CLOSED):
raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
except asyncio.TimeoutError as exc:
raise RuntimeError(
"Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
) from exc
finally:
if ws is not None:
# Closing must not discard an already acknowledged prompt ID or
# replace a definite rejection with an ambiguous timeout.
try:
async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
await ws.close()
except (asyncio.TimeoutError, aiohttp.ClientError):
pass
async def dispatch_worker_prompt(
@@ -102,6 +124,7 @@ async def dispatch_worker_prompt(
client_id=None,
use_websocket=False,
trace_execution_id=None,
return_prompt_id=False,
):
"""Send the prepared prompt to a worker ComfyUI instance."""
worker_url = build_worker_url(worker)
@@ -115,7 +138,7 @@ async def dispatch_worker_prompt(
if use_websocket:
try:
await _dispatch_via_websocket(
prompt_id = await _dispatch_via_websocket(
worker_url,
{
"prompt": prompt_obj,
@@ -123,7 +146,9 @@ async def dispatch_worker_prompt(
},
client_id,
)
return
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
return prompt_id if return_prompt_id else None
except Exception as exc:
worker_id = worker.get("id")
if trace_execution_id:
@@ -139,6 +164,95 @@ async def dispatch_worker_prompt(
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError("Worker /prompt response omitted prompt_id")
return prompt_id if return_prompt_id else None
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
"""Require protocol/node support and report which direct workers are idle.
This performs read-only GETs only; it never submits a prompt or contacts a
serverless provider. Unsupported worker releases fail before master inputs run.
"""
if not workers:
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
session = await get_client_session()
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
semaphore = asyncio.Semaphore(8)
async def inspect(worker):
worker_id = str(worker.get("id") or "unknown")
check = "capability"
try:
async with semaphore:
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(
f"Image-list worker {worker_id} does not expose the capability endpoint "
f"(HTTP {response.status}); upgrade that worker before running this feature"
)
capability = await response.json()
if not isinstance(capability, dict):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
supported = capability.get("supported")
protocol_version = capability.get("protocol_version")
nodes = capability.get("node_ids")
if (
not isinstance(supported, bool)
or type(protocol_version) is not int
or not isinstance(nodes, list)
or any(not isinstance(node_id, str) for node_id in nodes)
):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
if (
supported is not True
or protocol_version != 1
or not required_nodes.issubset(nodes)
):
raise RuntimeError(
f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
f"{sorted(required_nodes)}; no work was submitted"
)
check = "queue"
queue_url = build_worker_url(worker, "/queue")
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
queue = await response.json()
if (
not isinstance(queue, dict)
or not isinstance(queue.get("queue_running"), list)
or not isinstance(queue.get("queue_pending"), list)
):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
idle = not queue["queue_running"] and not queue["queue_pending"]
return worker_id, idle
except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
log(
f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
f"skipping: {exc}"
)
return worker_id, None
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
idle_workers = {
worker_id: True
for worker_id, idle in statuses
if idle is True
}
if require_idle and not idle_workers:
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
return idle_workers
async def select_active_workers(
@@ -203,6 +317,10 @@ def _extract_queue_remaining(payload):
async def _probe_worker_queue(worker, semaphore, probe_timeout):
async with semaphore:
if worker.get("type") == "serverless":
from .serverless_runtime import get_manager
status = await get_manager().health(worker)
return {"worker": worker, "queue_remaining": status["queueRemaining"]}
worker_url = build_worker_url(worker)
payload = await probe_worker(worker_url, timeout=probe_timeout)
if payload is None:
+11
View File
@@ -29,6 +29,11 @@ def _normalize_media_reference(value):
return None
cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
if MEDIA_FILE_RE.search(cleaned):
# Output/temp annotations select a different master-side media root.
# Preserve them through discovery/loading and as distinct rewrite keys.
annotation = re.search(r"\s*\[(output|temp)\]$", value)
if annotation:
return f"{cleaned} [{annotation.group(1)}]"
return cleaned
return None
@@ -147,6 +152,12 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
"""Upload one media file to worker iff missing or hash-mismatched."""
session = await get_client_session()
normalized = filename.replace("\\", "/")
annotation = re.search(r"\s*\[(output|temp)\]$", normalized)
if annotation:
# Worker uploads live under input/. Do not collide across master roots,
# or overwrite a different reference with the same basename.
basename = normalized[:annotation.start()].rstrip().rsplit("/", 1)[-1]
normalized = f"distributed-media/{annotation.group(1)}/{file_hash}/{basename}"
check_url = build_worker_url(worker, "/distributed/check_file")
try:
+308 -12
View File
@@ -125,6 +125,209 @@ def _find_upstream_nodes(prompt_obj, start_ids):
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")
@@ -148,16 +351,33 @@ def prune_prompt_for_worker(prompt_obj):
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:
preview_id = next_id()
pruned_prompt[preview_id] = {
"inputs": {
"images": [dist_id, 0],
},
"class_type": "PreviewImage",
"_meta": {
"title": "Preview Image (auto-added)",
},
}
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
@@ -167,6 +387,9 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
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:
@@ -194,6 +417,10 @@ def prepare_delegate_master_prompt(prompt_obj, 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",
@@ -214,11 +441,55 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
return pruned_prompt
def _resolve_direct_workers_for_image_list(config, requested_ids):
"""Resolve enabled workers by default, or exact explicitly selected direct workers."""
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
configured_workers = (config or {}).get("workers", [])
configured_ids = {
str(worker.get("id") or "").strip()
for worker in configured_workers
if str(worker.get("id") or "").strip()
}
if requested is not None:
unknown_ids = requested - configured_ids
if unknown_ids:
raise RuntimeError(
"Unknown explicitly selected image-list worker ID(s): "
+ ", ".join(sorted(unknown_ids))
)
direct = []
for worker in configured_workers:
worker_id = str(worker.get("id") or "").strip()
if not worker_id:
continue
if requested is None:
if not worker.get("enabled", False):
continue
elif worker_id not in requested:
continue
if str(worker.get("type") or "local").strip().lower() == "serverless":
continue
try:
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
direct.append({
"id": worker_id,
"name": str(worker.get("name") or worker_id),
"host": worker.get("host"),
"port": port,
"type": worker.get("type", "local"),
})
return direct
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"
distributed_nodes = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+ prompt_index.nodes_for_class("DistributedImageListStart")
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -313,6 +584,21 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
job_id_map, master_url, workflow_meta=None):
"""Attach protocol context to Start without changing its visible V3 sockets."""
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
inputs["image_list_master_callback"] = str(master_url or "")
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
if workflow_meta is not None:
inputs["image_list_workflow_meta"] = workflow_meta
def apply_participant_overrides(
prompt_copy,
participant_id,
@@ -321,6 +607,7 @@ def apply_participant_overrides(
master_url,
delegate_master,
prompt_index,
workflow_meta=None,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
is_master = participant_id == "master"
@@ -348,5 +635,14 @@ def apply_participant_overrides(
master_url,
enabled_json,
)
_override_image_list_nodes(
prompt_copy,
prompt_index,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
workflow_meta=workflow_meta,
)
return prompt_copy
+123
View File
@@ -0,0 +1,123 @@
"""Small Runpod queue API client. Never retries a potentially billable POST."""
import asyncio
import json
import re
import aiohttp
API_BASE = "https://api.runpod.ai/v2"
MAX_REQUEST_BYTES = 10_000_000
_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z")
class RunpodError(RuntimeError):
"""Sanitized provider failure safe to display in the panel."""
class AmbiguousSubmission(RunpodError):
"""Submission may have been accepted: never retry it automatically."""
def serialize_request(workflow, images, timeout_seconds):
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
payload = {
"input": {"workflow": workflow, "images": images},
"policy": {
"executionTimeout": timeout_seconds * 1000,
"ttl": timeout_seconds * 1000,
},
}
body = json.dumps(
payload,
allow_nan=False,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
size = len(body)
if size > MAX_REQUEST_BYTES:
raise ValueError(
f"Serialized serverless request body is {size:,} bytes; Runpod allows at most "
f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw "
"image bytes; reduce image or workflow size."
)
return body
def checked_id(value):
if not isinstance(value, str) or not _ID.fullmatch(value):
raise ValueError("Invalid Runpod endpoint or job ID")
return value
class RunpodClient:
def __init__(self, session, endpoint_id, api_key):
self.session = session
self.endpoint_id = checked_id(endpoint_id)
if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"):
raise ValueError("A valid Runpod API key is required")
self._api_key = api_key
async def _request(self, method, path, payload=None, request_body=None):
# Base origin is code-owned, never a user-supplied URL. Do not follow redirects
# with credentials, or include response bodies/transport errors in exceptions.
url = f"{API_BASE}/{self.endpoint_id}/{path}"
attempts = 3 if method == "GET" else 1
for attempt in range(attempts):
try:
request_kwargs = {
"headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"},
"timeout": aiohttp.ClientTimeout(total=30),
"allow_redirects": False,
}
if request_body is None:
request_kwargs["json"] = payload
else:
request_kwargs["data"] = request_body
async with self.session.request(
method, url,
**request_kwargs,
) as response:
status = response.status
if status != 200:
if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts:
await asyncio.sleep(0.5 * (2 ** attempt))
continue
if method == "POST" and path == "run" and status >= 500:
raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.")
raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})")
try:
result = await response.json()
except (ValueError, aiohttp.ContentTypeError):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None
raise RunpodError("Runpod returned an invalid JSON response") from None
if not isinstance(result, dict):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.")
raise RunpodError("Runpod returned an invalid response")
return result
except (asyncio.TimeoutError, aiohttp.ClientError, OSError):
if method == "POST" and path == "run":
raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None
if method == "GET" and attempt + 1 < attempts:
await asyncio.sleep(0.5 * (2 ** attempt))
continue
raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None
raise RunpodError("Runpod request failed")
async def health(self):
return await self._request("GET", "health")
async def submit(self, workflow, images, timeout_seconds):
body = serialize_request(workflow, images, timeout_seconds)
result = await self._request("POST", "run", request_body=body)
try:
return checked_id(result.get("id"))
except ValueError:
raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None
async def status(self, job_id):
return await self._request("GET", f"status/{checked_id(job_id)}")
async def cancel(self, job_id):
return await self._request("POST", f"cancel/{checked_id(job_id)}")
+57
View File
@@ -0,0 +1,57 @@
"""Private, atomic job-ID journal. Contains identifiers/status only, never payloads."""
import json
import os
from pathlib import Path
import tempfile
import threading
from .serverless_state import TERMINAL
FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"}
class JobLedger:
def __init__(self, path):
self.path = Path(path)
self._lock = threading.RLock()
def _read(self):
if self.path.is_symlink():
raise RuntimeError("Unsafe serverless job journal path")
if not self.path.exists():
return {}
try:
data = json.loads(self.path.read_text(encoding="utf-8"))
if not isinstance(data, dict) or any(not isinstance(v, dict) or set(v) - FIELDS for v in data.values()):
raise ValueError()
return data
except (ValueError, OSError):
raise RuntimeError("Serverless job journal is unreadable; preserve it for recovery before submitting more work") from None
def all(self):
with self._lock:
return list(self._read().values())
def update(self, record_id, **fields):
if set(fields) - FIELDS:
raise ValueError("Unexpected serverless journal field")
with self._lock:
data = self._read()
data.setdefault(record_id, {"record_id": record_id}).update(fields)
# Keep every unresolved job, plus a bounded terminal history.
completed = sorted((k for k, v in data.items() if v.get("status") in TERMINAL), key=lambda k: data[k].get("updated_at", 0))
for key in completed[:-200]:
data.pop(key)
if self.path.parent.is_symlink():
raise RuntimeError("Unsafe serverless journal directory")
self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
fd, temporary = tempfile.mkstemp(prefix=".jobs-", dir=self.path.parent)
try:
with os.fdopen(fd, "w", encoding="utf-8") as stream:
json.dump(data, stream, allow_nan=False)
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, self.path)
finally:
if os.path.exists(temporary):
os.unlink(temporary)
return dict(data[record_id])
+601
View File
@@ -0,0 +1,601 @@
"""Bounded image transport for generic Runpod-backed ComfyUI jobs."""
import asyncio
import base64
import binascii
import hashlib
import io
import ipaddress
import os
import re
import socket
import stat
import warnings
from pathlib import Path, PurePosixPath
from urllib.parse import urljoin, urlsplit
import aiohttp
import numpy as np
from PIL import Image, UnidentifiedImageError
MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024
MAX_INPUT_TOTAL_BYTES = 10_000_000
MAX_OUTPUT_FILES = 1024
MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024
MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024
MAX_IMAGE_PIXELS = 40_000_000
MAX_OUTPUT_ARRAY_BYTES = 256 * 1024 * 1024
MAX_REDIRECTS = 3
_DOWNLOAD_TIMEOUT_SECONDS = 60
_CHUNK_SIZE = 64 * 1024
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
_UNSUPPORTED_IMAGE_EXTENSIONS = {".gif", ".tif", ".tiff", ".avif", ".heic", ".ico", ".svg"}
_AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus", ".aiff", ".aif", ".wma"}
_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
_MEDIA_EXTENSIONS = _IMAGE_EXTENSIONS | _UNSUPPORTED_IMAGE_EXTENSIONS | _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS
_MEDIA_ANNOTATION = re.compile(r"\s+\[([A-Za-z0-9_-]{1,64})\]$")
_PREFIX_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\Z", re.ASCII)
_SAFE_FILENAME_RE = re.compile(r"[A-Za-z0-9_.-]{1,255}\Z", re.ASCII)
class ServerlessMediaError(ValueError):
"""Sanitized input, artifact, or download validation failure."""
def _folder_roots(folder_paths):
roots = {}
for kind, getter_name in (
("input", "get_input_directory"),
("output", "get_output_directory"),
("temp", "get_temp_directory"),
):
getter = getattr(folder_paths, getter_name, None)
if not callable(getter):
raise ServerlessMediaError("ComfyUI media directories are unavailable")
try:
roots[kind] = Path(getter()).resolve(strict=True)
except Exception:
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
return roots
def _inside_roots(path, roots):
for root in roots.values():
try:
path.relative_to(root)
return True
except ValueError:
continue
return False
def _normalize_media_reference(value):
if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value:
raise ServerlessMediaError("Image input reference is invalid")
value = value.replace("\\", "/")
explicit_root = None
annotation = _MEDIA_ANNOTATION.search(value)
if annotation:
marker = annotation.group(1).lower()
if marker not in {"input", "output", "temp"}:
raise ServerlessMediaError("Image input annotation is unsupported")
explicit_root = marker
value = value[: annotation.start()]
if "://" in value:
raise ServerlessMediaError("Remote media references are not supported as image inputs")
if re.match(r"^[A-Za-z]:", value) or value.startswith("//"):
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
path = PurePosixPath(value)
if any(part == ".." for part in path.parts):
raise ServerlessMediaError("Image input path traversal is not allowed")
if not path.parts or value in {".", ""}:
raise ServerlessMediaError("Image input reference is invalid")
if explicit_root:
value = f"{value} [{explicit_root}]"
suffix = Path(path.name).suffix.lower()
return value, suffix
def _resolve_input_path(value, suffix, folder_paths, roots):
if suffix not in _IMAGE_EXTENSIONS:
if suffix in _MEDIA_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
raise ServerlessMediaError("Unsupported serverless image file type")
try:
source = Path(value)
if source.is_absolute():
annotation = _MEDIA_ANNOTATION.search(value)
unannotated = value[:annotation.start()] if annotation else value
resolved = Path(unannotated).resolve(strict=True)
else:
resolver = getattr(folder_paths, "get_annotated_filepath", None)
if not callable(resolver):
raise ServerlessMediaError("ComfyUI media path resolver is unavailable")
resolved = Path(resolver(value)).resolve(strict=True)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if not _inside_roots(resolved, roots):
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
try:
mode = resolved.stat().st_mode
except OSError:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if not stat.S_ISREG(mode):
raise ServerlessMediaError("Image input is not a regular file")
return resolved
def _iter_image_references(workflow):
if not isinstance(workflow, dict):
raise ServerlessMediaError("Serverless workflow must be an API graph")
references = []
for node_id, node in workflow.items():
if not isinstance(node, dict):
continue
class_type = node.get("class_type")
inputs = node.get("inputs", {})
if class_type == "LoadImage" and not isinstance(inputs, dict):
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
if not isinstance(inputs, dict):
continue
if class_type == "LoadImage":
value = inputs.get("image")
if not isinstance(value, str):
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
references.append((node_id, inputs, "image", value, True))
for key in ("image", "file"):
if key == "image" and class_type == "LoadImage":
continue
value = inputs.get(key)
if not isinstance(value, str):
continue
normalized, suffix = _normalize_media_reference(value)
if suffix in _IMAGE_EXTENSIONS:
references.append((node_id, inputs, key, normalized, False))
elif suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
elif suffix in _UNSUPPORTED_IMAGE_EXTENSIONS or (key == "image" and suffix):
raise ServerlessMediaError("Unsupported serverless image file type")
for key in ("audio", "video"):
value = inputs.get(key)
if not isinstance(value, str):
continue
_normalized, suffix = _normalize_media_reference(value)
if suffix in _MEDIA_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
return references
def _read_limited(path, expected_size):
try:
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
with os.fdopen(descriptor, "rb") as handle:
info = os.fstat(handle.fileno())
if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES:
raise ServerlessMediaError("Image input exceeds the permitted size")
if info.st_size > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
data = handle.read(MAX_INPUT_TOTAL_BYTES + 1)
except ServerlessMediaError:
raise
except OSError:
raise ServerlessMediaError("Image input file could not be read safely") from None
if len(data) > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
if len(data) != expected_size:
raise ServerlessMediaError("Image input changed while being read")
return data
async def prepare_images(workflow: dict) -> list[dict]:
"""Pack referenced local images and rewrite only their graph input references.
``workflow`` is expected to be a private API-graph copy owned by the caller.
Returned ``name`` values are flat generated names suitable for ``input.images``.
"""
references = _iter_image_references(workflow)
if not references:
return []
# Import only in ComfyUI, keeping this module importable in the serverless client.
try:
import folder_paths
except Exception:
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
roots = _folder_roots(folder_paths)
resolved_refs = []
unique_paths = {}
total_source_bytes = 0
for node_id, inputs, key, value, _is_load_image in references:
normalized, suffix = _normalize_media_reference(value)
if key == "image" and suffix not in _IMAGE_EXTENSIONS:
if suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
raise ServerlessMediaError("Unsupported serverless image file type")
path = _resolve_input_path(normalized, suffix, folder_paths, roots)
if path not in unique_paths:
try:
size = path.stat().st_size
except OSError:
raise ServerlessMediaError("Image input file is missing or invalid") from None
if size < 0 or size > MAX_INPUT_IMAGE_BYTES:
raise ServerlessMediaError("Image input exceeds the permitted size")
total_source_bytes += size
if total_source_bytes > MAX_INPUT_TOTAL_BYTES:
raise ServerlessMediaError(
f"Total raw serverless image input size is {total_source_bytes:,} bytes; "
f"the input read guard is {MAX_INPUT_TOTAL_BYTES:,} bytes. Base64 encoding "
"adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
)
unique_paths[path] = {"size": size, "suffix": suffix}
resolved_refs.append((inputs, key, path))
packed_by_name = {}
path_to_name = {}
for path, metadata in unique_paths.items():
data = _read_limited(path, metadata["size"])
digest = hashlib.sha256(data).hexdigest()
name = f"serverless_{digest}{metadata['suffix']}"
path_to_name[path] = name
packed_by_name.setdefault(name, {"name": name, "image": base64.b64encode(data).decode("ascii")})
for inputs, key, path in resolved_refs:
inputs[key] = path_to_name[path]
return [packed_by_name[name] for name in sorted(packed_by_name)]
def _validate_artifact_url(url):
if not isinstance(url, str) or not url or any(ord(char) < 0x20 or char.isspace() for char in url):
raise ServerlessMediaError("Serverless image artifact URL is invalid")
try:
parsed = urlsplit(url)
port = parsed.port
except ValueError:
raise ServerlessMediaError("Serverless image artifact URL is invalid") from None
if (
parsed.scheme.lower() != "https"
or not parsed.hostname
or parsed.username is not None
or parsed.password is not None
or port not in (None, 443)
or parsed.fragment
):
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
host = parsed.hostname.rstrip(".").lower()
if host in {"localhost", "metadata", "metadata.google.internal"} or host.endswith((".localhost", ".local", ".internal")):
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
try:
address = ipaddress.ip_address(host)
except ValueError:
return url
if not address.is_global:
raise ServerlessMediaError("Serverless image artifact URL must resolve to a public address")
return url
class _PublicResolver(aiohttp.abc.AbstractResolver):
"""Resolve once, reject non-public answers, and pin the checked address."""
async def _getaddrinfo(self, host, port, family):
return await asyncio.get_running_loop().getaddrinfo(
host, port, family=family, type=socket.SOCK_STREAM
)
async def resolve(self, host, port=0, family=socket.AF_INET):
try:
results = await self._getaddrinfo(host, port, family)
except (OSError, asyncio.TimeoutError):
raise OSError("Public artifact host could not be resolved") from None
if not results:
raise OSError("Public artifact host could not be resolved")
resolved = []
seen = set()
for result_family, _socktype, proto, _canonname, sockaddr in results:
try:
address = ipaddress.ip_address(sockaddr[0])
except (ValueError, TypeError, IndexError):
raise OSError("Public artifact host resolved to an invalid address") from None
if not address.is_global:
raise OSError("Public artifact host resolved to a non-public address")
key = (str(address), result_family)
if key in seen:
continue
seen.add(key)
resolved.append({
"hostname": host,
"host": str(address),
"port": port,
"family": result_family,
"proto": proto,
"flags": 0,
})
if not resolved:
raise OSError("Public artifact host could not be resolved")
return resolved
async def close(self):
return None
async def download_public_artifact(url: str, max_bytes: int = MAX_OUTPUT_IMAGE_BYTES) -> bytes:
"""Download a signed public HTTPS artifact without provider credentials."""
if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or not 0 < max_bytes <= MAX_OUTPUT_IMAGE_BYTES:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
current_url = _validate_artifact_url(url)
resolver = _PublicResolver()
connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False, limit=4)
timeout = aiohttp.ClientTimeout(
total=_DOWNLOAD_TIMEOUT_SECONDS,
connect=10,
sock_read=15,
)
try:
async with aiohttp.ClientSession(
connector=connector,
timeout=timeout,
trust_env=False,
cookie_jar=aiohttp.DummyCookieJar(),
) as session:
for redirect_count in range(MAX_REDIRECTS + 1):
_validate_artifact_url(current_url)
async with session.get(current_url, allow_redirects=False) as response:
if response.status in {301, 302, 303, 307, 308}:
location = response.headers.get("Location")
if not location or redirect_count >= MAX_REDIRECTS:
raise ServerlessMediaError("Serverless image artifact redirect is invalid")
current_url = _validate_artifact_url(urljoin(current_url, location))
continue
if response.status != 200:
raise ServerlessMediaError("Serverless image artifact download failed")
payload = bytearray()
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
if not isinstance(chunk, bytes):
raise ServerlessMediaError("Serverless image artifact response is invalid")
if len(payload) + len(chunk) > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
payload.extend(chunk)
if not payload:
raise ServerlessMediaError("Serverless image artifact is empty")
return bytes(payload)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Serverless image artifact download failed") from None
raise ServerlessMediaError("Serverless image artifact download failed")
def _validated_manifest(manifest):
if not isinstance(manifest, list) or len(manifest) > MAX_OUTPUT_FILES:
raise ServerlessMediaError("Serverless image manifest is invalid")
normalized = []
seen_prefixes = set()
prefixes = []
for item in manifest:
if not isinstance(item, dict):
raise ServerlessMediaError("Serverless image manifest is invalid")
if item.get("kind") != "image":
raise ServerlessMediaError("Only image collectors are supported by serverless media transport")
for key in ("collector_id", "job_id", "node_id"):
value = item.get(key)
if not isinstance(value, (str, int)) or isinstance(value, bool) or not str(value):
raise ServerlessMediaError("Serverless image manifest is invalid")
prefix = item.get("prefix")
if not isinstance(prefix, str) or not _PREFIX_RE.fullmatch(prefix) or prefix in {".", ".."}:
raise ServerlessMediaError("Serverless image manifest contains an unsafe filename prefix")
folded = prefix.casefold()
if folded in seen_prefixes:
raise ServerlessMediaError("Serverless image manifest prefixes must be unique")
seen_prefixes.add(folded)
prefixes.append(prefix)
normalized.append({
"collector_id": str(item["collector_id"]),
"job_id": str(item["job_id"]),
"node_id": str(item["node_id"]),
"prefix": prefix,
})
for first in prefixes:
for second in prefixes:
if first != second and second.casefold().startswith(first.casefold() + "_"):
raise ServerlessMediaError("Serverless image manifest prefixes are ambiguous")
return normalized
def _artifact_filename_matches(filename, prefixes):
if not isinstance(filename, str):
return None
basename = filename.replace("\\", "/").rsplit("/", 1)[-1]
matches = [item for item in prefixes if basename.startswith(item["prefix"] + "_")]
if len(matches) > 1:
raise ServerlessMediaError("Serverless image artifact matches multiple collectors")
return matches[0] if matches else None
def _canonical_artifact_identity(entry, filename):
artifact_path = entry.get("artifact_path")
artifact_type = entry.get("artifact_type", "output")
if artifact_path is not None:
if not isinstance(artifact_path, str) or not artifact_path or "\x00" in artifact_path:
raise ServerlessMediaError("Serverless image artifact identity is invalid")
normalized_path = os.path.normcase(os.path.normpath(artifact_path.replace("\\", "/")))
else:
subfolder = entry.get("subfolder", "")
if not isinstance(subfolder, str) or "\x00" in subfolder:
raise ServerlessMediaError("Serverless image artifact identity is invalid")
normalized_path = os.path.normcase(os.path.normpath(f"{subfolder}/{filename}"))
return str(artifact_type), normalized_path
def _decode_base64_data(value, max_bytes=MAX_OUTPUT_IMAGE_BYTES):
if not isinstance(value, str) or not value:
raise ServerlessMediaError("Serverless image artifact data is invalid")
max_bytes = min(max_bytes, MAX_OUTPUT_IMAGE_BYTES)
max_encoded = ((max_bytes + 2) // 3) * 4
if len(value) > max_encoded:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
padding = len(value) - len(value.rstrip("="))
estimated_size = (len(value) // 4) * 3 - padding
if estimated_size > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
try:
payload = base64.b64decode(value, validate=True)
except (binascii.Error, ValueError):
raise ServerlessMediaError("Serverless image artifact data is invalid") from None
if not payload:
raise ServerlessMediaError("Serverless image artifact is empty")
if len(payload) > max_bytes:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
return payload
def _decode_image_array(payload, total_array_bytes):
try:
with warnings.catch_warnings():
warnings.simplefilter("error", Image.DecompressionBombWarning)
image = Image.open(io.BytesIO(payload))
width, height = image.size
if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
raise ServerlessMediaError("Serverless image dimensions exceed the permitted limit")
has_alpha = image.mode in {"RGBA", "LA"} or "transparency" in image.info
mode = "RGBA" if has_alpha else "RGB"
output_bytes = width * height * (4 if has_alpha else 3) * 4
if total_array_bytes + output_bytes > MAX_OUTPUT_ARRAY_BYTES:
raise ServerlessMediaError("Decoded serverless images exceed the permitted memory limit")
converted = image.convert(mode)
converted.load()
array = np.array(converted, dtype=np.float32, copy=True)
np.multiply(array, 1.0 / 255.0, out=array)
return array, output_bytes
except ServerlessMediaError:
raise
except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
except Exception:
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
def _image_output(output):
"""Unwrap handler output and reject explicit execution errors."""
if not isinstance(output, dict):
raise ServerlessMediaError("Serverless job output is invalid")
nested = output.get("output")
data = nested if isinstance(nested, dict) else output
for candidate in (output, data):
if candidate.get("error") or candidate.get("errors"):
raise ServerlessMediaError("Serverless job reported an execution error")
return data
async def _select_image_payloads(data, expected, download):
"""Match artifact identities and fetch bounded bytes, without image decoding."""
artifacts = data.get("images", [])
if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES:
raise ServerlessMediaError("Serverless job returned too many image files")
selected = {item["prefix"]: [] for item in expected}
identities = set()
total_payload_bytes = 0
for entry in artifacts:
if not isinstance(entry, dict):
continue
filename = entry.get("filename")
manifest_entry = _artifact_filename_matches(filename, expected)
if manifest_entry is None:
continue
normalized_filename = filename.replace("\\", "/")
if "/" in normalized_filename or not _SAFE_FILENAME_RE.fullmatch(filename):
raise ServerlessMediaError("Serverless image artifact filename is invalid")
if Path(filename).suffix.lower() not in _IMAGE_EXTENSIONS:
raise ServerlessMediaError("Serverless image artifact has an unsupported file type")
if "artifact_type" in entry and entry["artifact_type"] != "output":
raise ServerlessMediaError("Expected serverless image artifact was not an output image")
identity = _canonical_artifact_identity(entry, filename)
if identity in identities:
raise ServerlessMediaError("Duplicate serverless image artifact identity")
identities.add(identity)
match = re.match(rf"^{re.escape(manifest_entry['prefix'])}_(\d+)(?:_|\.)", filename)
if match is None:
raise ServerlessMediaError("Expected serverless image filename counter is invalid")
artifact_type = entry.get("type")
remaining_bytes = MAX_OUTPUT_TOTAL_BYTES - total_payload_bytes
if remaining_bytes <= 0:
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
if artifact_type == "base64":
payload = _decode_base64_data(entry.get("data"), remaining_bytes)
elif artifact_type in {"gcs_url", "s3_url", "url"}:
url = entry.get("data") or entry.get("url")
url = _validate_artifact_url(url)
try:
if download is None:
payload = await download_public_artifact(
url, max_bytes=min(MAX_OUTPUT_IMAGE_BYTES, remaining_bytes)
)
else:
payload = await download(url)
except ServerlessMediaError:
raise
except Exception:
raise ServerlessMediaError("Serverless image artifact download failed") from None
if not isinstance(payload, bytes) or not payload:
raise ServerlessMediaError("Serverless image artifact response is invalid")
if len(payload) > MAX_OUTPUT_IMAGE_BYTES:
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
else:
raise ServerlessMediaError("Serverless image artifact transport type is unsupported")
total_payload_bytes += len(payload)
if total_payload_bytes > MAX_OUTPUT_TOTAL_BYTES:
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
selected[manifest_entry["prefix"]].append((int(match.group(1)), filename, identity, payload))
if any(not selected[item["prefix"]] for item in expected):
raise ServerlessMediaError("Serverless job is missing an expected collector image")
return selected
def _decode_image_groups(selected, expected):
"""Decode in stable counter order, bounding array memory across collectors."""
result = {}
total_array_bytes = 0
for manifest_entry in expected:
prefix = manifest_entry["prefix"]
rows = sorted(selected[prefix], key=lambda row: (row[0], row[1], row[2]))
arrays = []
shape = None
for _counter, _filename, _identity, payload in rows:
array, added_bytes = _decode_image_array(payload, total_array_bytes)
total_array_bytes += added_bytes
if shape is None:
shape = array.shape
elif array.shape != shape:
raise ServerlessMediaError("Images from one serverless collector have inconsistent dimensions")
arrays.append(array)
result.setdefault(manifest_entry["job_id"], []).extend(arrays)
return result
async def collect_images(output: dict, manifest: list, download=None) -> dict[str, list[np.ndarray]]:
"""Validate and decode expected image artifacts, grouped by serverless job ID."""
data = _image_output(output)
expected = _validated_manifest(manifest)
if not expected:
return {}
selected = await _select_image_payloads(data, expected, download)
return _decode_image_groups(selected, expected)
+420
View File
@@ -0,0 +1,420 @@
"""Compile prepared distributed prompts into worker-side ComfyUI image jobs.
This module is deliberately pure: it only copies and rewrites an API prompt
mapping and never imports or executes ComfyUI nodes.
"""
from __future__ import annotations
import copy
import hashlib
import json
import re
_SUPPORTED_DISTRIBUTED = {
"DistributedCollector",
"DistributedSeed",
"DistributedValue",
"DistributedModelName",
}
_REMOVED_SCALAR_NODES = {
"DistributedSeed",
"DistributedValue",
"DistributedModelName",
}
_PRIMITIVE_NODES = {
"PrimitiveBoolean": False,
"PrimitiveFloat": 0.0,
"PrimitiveInt": 0,
"PrimitiveString": "",
}
_REMOVED_NODE_OUTPUTS = {
"DistributedCollector": {0, 1},
"DistributedSeed": {0},
"DistributedValue": {0},
"DistributedModelName": {0},
}
_SEED_DEFAULT = 1125899906842
def _is_link(value):
return isinstance(value, list) and len(value) == 2
def _is_unsupported_distributed(class_type):
if not isinstance(class_type, str) or class_type in _SUPPORTED_DISTRIBUTED:
return False
lowered = class_type.lower()
return "distributed" in lowered or "batchdivider" in lowered
def _safe_component(value):
cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_-")
return cleaned[:24] or "empty"
def _filename_prefix(job_id, worker_id, collector_id, kind):
identity = json.dumps(
[job_id, worker_id, collector_id, kind],
ensure_ascii=False,
separators=(",", ":"),
)
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
parts = (
"serverless",
_safe_component(job_id),
_safe_component(worker_id),
_safe_component(collector_id),
kind,
digest,
)
return "_".join(parts)
def _worker_index(worker_id):
"""Mirror nodes.utilities worker-index parsing; None means fallback."""
try:
if worker_id.startswith("worker_"):
return int(worker_id.split("_")[1])
return int(worker_id)
except (ValueError, IndexError):
return None
def _coerce(value, value_type):
if value_type == "INT":
return int(float(value))
if value_type == "FLOAT":
return float(value)
return value
def _coerce_safe(value, value_type):
try:
return _coerce(value, value_type)
except (TypeError, ValueError):
return value
def _stringify_model_name(value):
if isinstance(value, str):
return value
if isinstance(value, (int, float, bool)):
return str(value)
try:
return json.dumps(value, indent=4)
except (TypeError, ValueError):
return str(value)
def _model_name_result(text):
if isinstance(text, list):
values = [_stringify_model_name(value) for value in text]
return values[0] if len(values) == 1 else values
return _stringify_model_name(text)
def _validate_prepared_prompt(prepared_prompt):
"""Validate graph shape, links, collectors, and supported topology."""
if not isinstance(prepared_prompt, dict):
raise ValueError("prepared prompt must be an API graph mapping")
source_nodes = {}
for raw_id, node in prepared_prompt.items():
node_id = str(raw_id)
if node_id in source_nodes:
raise ValueError("prompt has duplicate node IDs after string normalization")
if not isinstance(node, dict):
raise ValueError(f"prompt node {node_id!r} must be an object")
class_type = node.get("class_type")
if not isinstance(class_type, str):
raise ValueError(f"prompt node {node_id!r} is missing a class_type")
if _is_unsupported_distributed(class_type):
raise ValueError(f"unsupported distributed or batch node: {class_type}")
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
source_nodes[node_id] = node
if not source_nodes:
raise ValueError("prepared prompt contains no nodes")
# Validate every API link before any graph rewriting. Slots on the nodes
# removed by this compiler have a known schema and are checked exactly.
for node_id, node in source_nodes.items():
for input_name, value in node.get("inputs", {}).items():
if not _is_link(value):
continue
upstream_id, output_slot = value
upstream_id = str(upstream_id)
if upstream_id not in source_nodes:
raise ValueError(
f"invalid link at node {node_id!r} input {input_name!r}: "
f"source {upstream_id!r} does not exist"
)
if isinstance(output_slot, bool) or not isinstance(output_slot, int) or output_slot < 0:
raise ValueError(
f"invalid link at node {node_id!r} input {input_name!r}: "
"output slot must be a nonnegative integer"
)
upstream_class = source_nodes[upstream_id]["class_type"]
valid_slots = _REMOVED_NODE_OUTPUTS.get(upstream_class)
if valid_slots is not None and output_slot not in valid_slots:
raise ValueError(
f"invalid output slot {output_slot} for removed node "
f"{upstream_class} {upstream_id!r}"
)
collector_ids = sorted(
(node_id for node_id, node in source_nodes.items() if node["class_type"] == "DistributedCollector"),
key=str,
)
if not collector_ids:
raise ValueError("prepared prompt contains no DistributedCollector nodes")
# Only image collectors are supported by the current generic handler.
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
if "audio" in inputs:
raise ValueError("collector audio input is unsupported by this image-only compiler")
if not _is_link(inputs.get("images")):
raise ValueError(f"collector {collector_id!r} is missing a linked image input")
source_id = str(inputs["images"][0])
source_class = source_nodes[source_id]["class_type"]
if source_class in _REMOVED_SCALAR_NODES:
raise ValueError(f"collector {collector_id!r} image input must link to an image node")
# A collector cannot be used, directly or through processing, as another
# collector's input: this first compiler handles disconnected jobs only.
for collector_id in collector_ids:
pending = [str(source_nodes[collector_id]["inputs"]["images"][0])]
seen = set()
while pending:
node_id = pending.pop()
if node_id in seen:
continue
seen.add(node_id)
if source_nodes[node_id]["class_type"] == "DistributedCollector":
raise ValueError("collector-to-collector dependence is unsupported")
for value in source_nodes[node_id].get("inputs", {}).values():
if _is_link(value):
pending.append(str(value[0]))
return source_nodes, collector_ids
def _make_scalar_resolver(source_nodes):
"""Build pure resolvers for distributed scalar nodes and stock primitives."""
# Resolve only stock primitive nodes. A visited set makes malformed
# primitive-to-primitive cycles fail clearly instead of recursing forever.
def resolve_primitive(node_id, visiting=None):
node_id = str(node_id)
node = source_nodes[node_id]
class_type = node["class_type"]
if class_type not in _PRIMITIVE_NODES:
raise ValueError("linked value must be literal or supported PrimitiveInt/Float/String/Boolean")
if visiting is None:
visiting = set()
if node_id in visiting:
raise ValueError(f"cycle while resolving primitive value at node {node_id!r}")
visiting.add(node_id)
value = node.get("inputs", {}).get("value", _PRIMITIVE_NODES[class_type])
if _is_link(value):
parent_id, output_slot = str(value[0]), value[1]
if output_slot != 0:
raise ValueError(f"invalid output slot {output_slot} for primitive node {parent_id!r}")
if source_nodes[parent_id]["class_type"] not in _PRIMITIVE_NODES:
raise ValueError("linked primitive value must come from a stock Primitive node")
result = resolve_primitive(parent_id, visiting)
else:
result = copy.deepcopy(value)
visiting.remove(node_id)
return result
def resolve_literal_or_primitive(value, context):
if _is_link(value):
source_id, output_slot = str(value[0]), value[1]
if output_slot != 0 or source_nodes[source_id]["class_type"] not in _PRIMITIVE_NODES:
raise ValueError(f"{context} must be a literal or supported Primitive link")
try:
return resolve_primitive(source_id)
except ValueError as exc:
if "cycle" in str(exc):
raise
raise ValueError(f"{context} must be a literal or supported Primitive link") from None
return copy.deepcopy(value)
def resolve_hidden(node_id, name, default):
value = source_nodes[node_id].get("inputs", {}).get(name, default)
if _is_link(value):
raise ValueError(f"{name} on distributed node {node_id!r} must be a literal")
return copy.deepcopy(value)
resolved_distributed = {}
def resolve_distributed(node_id, output_slot):
node_id = str(node_id)
class_type = source_nodes[node_id]["class_type"]
if output_slot != 0:
raise ValueError(f"invalid output slot {output_slot} for {class_type} {node_id!r}")
cache_key = (node_id, output_slot)
if cache_key in resolved_distributed:
return copy.deepcopy(resolved_distributed[cache_key])
inputs = source_nodes[node_id].get("inputs", {})
if class_type == "DistributedSeed":
seed = resolve_literal_or_primitive(inputs.get("seed", _SEED_DEFAULT), f"seed on {node_id!r}")
is_worker = resolve_hidden(node_id, "is_worker", False)
worker_id = resolve_hidden(node_id, "worker_id", "")
if is_worker:
worker_index = _worker_index(worker_id)
if worker_index is not None:
seed = seed + worker_index + 1
result = seed
elif class_type == "DistributedValue":
default_value = resolve_literal_or_primitive(
inputs.get("default_value", ""), f"default_value on {node_id!r}"
)
worker_values = resolve_literal_or_primitive(
inputs.get("worker_values", "{}"), f"worker_values on {node_id!r}"
)
is_worker = resolve_hidden(node_id, "is_worker", False)
worker_id = resolve_hidden(node_id, "worker_id", "")
try:
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
except json.JSONDecodeError:
values = {}
if not isinstance(values, dict):
values = {}
value_type = values.get("_type", "STRING")
fallback = _coerce_safe(default_value, value_type)
result = fallback
if is_worker:
worker_index = _worker_index(worker_id)
if worker_index is not None:
raw_value = values.get(str(worker_index + 1), "")
if raw_value:
try:
result = _coerce(raw_value, value_type)
except (ValueError, IndexError):
result = fallback
elif class_type == "DistributedModelName":
text = resolve_literal_or_primitive(
inputs.get("text", ""), f"DistributedModelName text on {node_id!r}"
)
result = _model_name_result(text)
else:
raise ValueError(f"unsupported distributed node: {class_type}")
resolved_distributed[cache_key] = copy.deepcopy(result)
return result
return resolve_distributed
def _build_savers(source_nodes, collector_ids):
"""Construct stable SaveImage nodes and their collector metadata."""
# The prepared collector values define the artifact identity. The fallback
# for absent hidden values is stable and retains the collector's job scope.
collector_specs = []
max_numeric_id = 0
for node_id in source_nodes:
try:
max_numeric_id = max(max_numeric_id, int(node_id))
except ValueError:
pass
next_numeric_id = max_numeric_id + 1
for collector_id in collector_ids:
inputs = source_nodes[collector_id].get("inputs", {})
raw_job_id = inputs.get("multi_job_id", collector_id)
if _is_link(raw_job_id):
raise ValueError(f"collector {collector_id!r} multi_job_id must be a literal")
job_id = str(collector_id if raw_job_id in (None, "") else raw_job_id)
raw_worker_id = inputs.get("worker_id", "")
if _is_link(raw_worker_id):
raise ValueError(f"collector {collector_id!r} worker_id must be a literal")
worker_id = "" if raw_worker_id is None else str(raw_worker_id)
saver_id = str(next_numeric_id)
next_numeric_id += 1
prefix = _filename_prefix(job_id, worker_id, collector_id, "image")
collector_specs.append(
{
"collector_id": collector_id,
"job_id": job_id,
"node_id": saver_id,
"kind": "image",
"prefix": prefix,
"source_link": copy.deepcopy(inputs["images"]),
}
)
saver_nodes = {}
for spec in collector_specs:
saver_nodes[spec["node_id"]] = {
"class_type": "SaveImage",
"inputs": {
"images": spec["source_link"],
"filename_prefix": spec["prefix"],
},
"_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"},
}
return collector_specs, saver_nodes
def _prune_graph(source_nodes, collector_specs, saver_nodes, resolve_distributed):
"""Keep saver ancestors, replacing distributed scalar links with literals."""
all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()}
all_nodes.update(copy.deepcopy(saver_nodes))
needed = set()
pending = [spec["node_id"] for spec in collector_specs]
while pending:
node_id = pending.pop()
if node_id in needed:
continue
needed.add(node_id)
node = all_nodes[node_id]
if node["class_type"] == "DistributedCollector":
raise ValueError("collector-to-collector dependence is unsupported")
for input_name, value in list(node.get("inputs", {}).items()):
if not _is_link(value):
continue
upstream_id, output_slot = str(value[0]), value[1]
upstream_class = source_nodes[upstream_id]["class_type"]
if upstream_class in _REMOVED_SCALAR_NODES:
node["inputs"][input_name] = resolve_distributed(upstream_id, output_slot)
else:
pending.append(upstream_id)
compiled = {
node_id: copy.deepcopy(node)
for node_id, node in all_nodes.items()
if node_id in needed
}
if any("distributed" in node["class_type"].lower() for node in compiled.values()):
raise ValueError("compiled prompt still contains a distributed dependency")
return compiled
def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]:
"""Compile a prepared worker prompt to an image-only serverless graph.
Every DistributedCollector's original ``images`` input is connected to a
new SaveImage node. Only those savers and their upstream dependencies are
retained. DistributedSeed/Value/ModelName links are resolved as scalar
literals; unsupported node execution is never attempted.
"""
source_nodes, collector_ids = _validate_prepared_prompt(prepared_prompt)
resolve_distributed = _make_scalar_resolver(source_nodes)
collector_specs, saver_nodes = _build_savers(source_nodes, collector_ids)
compiled = _prune_graph(
source_nodes,
collector_specs,
saver_nodes,
resolve_distributed,
)
manifest = [
{key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")}
for spec in collector_specs
]
return compiled, manifest
+375
View File
@@ -0,0 +1,375 @@
"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
import asyncio
import time
import uuid
from pathlib import Path
from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
from .serverless_ledger import JobLedger
from .serverless_state import (
PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
occupies_capacity, terminal_provider_status,
)
async def decode_results(output, manifest):
from .serverless_media import collect_images
import torch
images = await collect_images(output, manifest)
return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
def state_directory():
import folder_paths
return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
def get_secret_store():
from ...utils.serverless_config import SecretStore
return SecretStore(state_directory() / "credentials.json")
def get_manager():
import server
from ...utils.network import get_client_session
ps = server.PromptServer.instance
manager = getattr(ps, "distributed_serverless", None)
if manager is None:
manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
ps.distributed_serverless = manager
return manager
class ServerlessManager:
def __init__(self, prompt_server, secrets, ledger_path, session_factory):
self.server = prompt_server
self.secrets = secrets
self.ledger = JobLedger(ledger_path)
self.session_factory = session_factory
self.tasks = {}
self.runs = {}
self.pool_limits = {}
self.health_cache = {}
self.health_locks = {}
self.poll_seconds = 2.0
self.capacity_changed = asyncio.Condition()
self.recovery_lock = asyncio.Lock()
self._recover_local_records()
def _recover_local_records(self):
# WAITING is committed before admission; it cannot have reached Runpod.
# SUBMITTING is an intent record, so a lost ID must remain reserved.
for record in self.ledger.all():
if record.get("status") == "WAITING" and not record.get("job_id"):
self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time())
elif record.get("status") == "SUBMITTING" and not record.get("job_id"):
self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time())
def _task_is_active(self, record_id):
task = self.tasks.get(record_id)
return task is not None and not task.done()
def pool_is_active(self, worker_id):
"""Protect both live local work and unresolved pre-restart provider work."""
for rid, run in self.runs.items():
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
if str(pool_id) == str(worker_id) and self._task_is_active(rid):
return True
return any(belongs_to_pool(record, worker_id) and is_unresolved(record)
for record in self.ledger.all())
async def _reserve_provider_capacity(self, run):
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
async with self.capacity_changed:
while True:
self._check_master(run)
outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record)
for record in self.ledger.all())
if outstanding < run["worker"]["concurrency"]:
# Claim atomically before any await that could submit a billed job.
self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time())
return
await self.capacity_changed.wait()
async def _notify_capacity(self):
async with self.capacity_changed:
self.capacity_changed.notify_all()
async def _client(self, worker):
pool_id = worker.get("pool_id", worker["id"])
key = self.secrets.get(pool_id)
return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
async def health(self, worker, *, fresh=False):
pool_id = worker.get("pool_id", worker["id"])
cache_key = (pool_id, worker["endpoint_id"])
lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
async with lock:
cached = self.health_cache.get(cache_key)
if cached and not fresh and time.monotonic() - cached[0] < 10:
return cached[1]
client = await self._client(worker)
data = await client.health()
jobs = data.get("jobs", {})
if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
raise RunpodError("Runpod returned an invalid endpoint health record")
try:
pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
except (ValueError, TypeError):
raise RunpodError("Runpod returned invalid queue counts") from None
result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
self.health_cache[cache_key] = (time.monotonic(), result)
return result
def _semaphore(self, worker):
pool_id = worker.get("pool_id", worker["id"])
limit = worker["concurrency"]
current = self.pool_limits.get(pool_id)
if current is None:
current = (limit, asyncio.Semaphore(limit))
self.pool_limits[pool_id] = current
elif current[0] != limit:
if self.pool_is_active(pool_id):
raise ValueError("Wait for this serverless pool's current jobs before changing concurrency")
current = (limit, asyncio.Semaphore(limit))
self.pool_limits[pool_id] = current
return current[1]
def ensure_capacity(self, count):
if sum(not task.done() for task in self.tasks.values()) + count > 64:
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
async def start(self, execution_id, prompt_id, prepared_jobs):
self.ensure_capacity(len(prepared_jobs))
new_runs = []
# Journal every local slot before starting any billed submission.
for prepared in prepared_jobs:
worker = prepared["worker"]
semaphore = self._semaphore(worker)
record_id = uuid.uuid4().hex
now = time.time()
self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
self.runs[record_id] = run
new_runs.append((record_id, run, semaphore))
for record_id, run, semaphore in new_runs:
self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
# Bound in-memory completed tracking as well as the on-disk journal.
done = [rid for rid, task in self.tasks.items() if task.done()]
for rid in done[:-200]:
self.tasks.pop(rid, None)
self.runs.pop(rid, None)
def is_pending(self, collector_job_id, participant_id):
return any(not task.done() and run["worker"]["id"] == participant_id and
any(m["job_id"] == collector_job_id for m in run["manifest"])
for rid, task in self.tasks.items() for run in [self.runs[rid]])
def owns(self, collector_job_id, participant_id):
return any(run["worker"]["id"] == participant_id and
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
async def _deliver(self, run, images=None, error=None):
async with self.server.distributed_jobs_lock:
for binding in run["manifest"]:
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
if queue is None:
continue # A cancelled/finished collector must never be resurrected.
if error:
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
else:
batch = images[binding["job_id"]]
for index, tensor in enumerate(batch):
queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
"tensor": tensor, "is_last": index == len(batch) - 1})
async def _cancel_provider(self, run):
if run["provider_job_id"] and run["client"]:
try:
result = await run["client"].cancel(run["provider_job_id"])
return terminal_provider_status(result)
except Exception:
return None
return None
def _check_master(self, run):
queue = getattr(self.server, "prompt_queue", None)
if queue is None:
return
history = queue.get_history(run["prompt_id"])
record = history.get(run["prompt_id"], {})
if record.get("status", {}).get("status_str") == "error":
raise RunpodError("Master workflow failed; serverless work stopped")
if not record and hasattr(queue, "get_current_queue"):
running, pending = queue.get_current_queue()
if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
async def _run(self, run, semaphore):
rid = run["record_id"]
submitted = False
provider_terminal = False
try:
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
async with semaphore:
await self._reserve_provider_capacity(run)
client = await self._client(run["worker"])
run["client"] = client
submitted = True
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
run["provider_job_id"] = job_id
# Persist immediately: all subsequent actions use this exact ID.
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
while True:
self._check_master(run)
response = await client.status(job_id)
status = response.get("status")
if not isinstance(status, str) or status not in PROVIDER_STATES:
raise RunpodError("Runpod returned an unknown job state")
provider_terminal = status in PROVIDER_TERMINAL
if status == "COMPLETED":
output = response.get("output")
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
images = await decode_results(output, run["manifest"])
await self._deliver(run, images=images)
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
return
self.ledger.update(rid, status=status, updated_at=time.time())
if provider_terminal:
raise RunpodError(f"Serverless job ended with {status}")
await asyncio.sleep(self.poll_seconds)
except AmbiguousSubmission as exc:
self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time())
await self._deliver(run, error=str(exc))
except (asyncio.TimeoutError, asyncio.CancelledError) as exc:
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
if submitted and not run["provider_job_id"]:
status = "UNKNOWN"
error = "Submission interrupted before its job ID was received; check Runpod before retrying"
elif isinstance(exc, asyncio.TimeoutError):
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
error = "Serverless job exceeded its queue/execution deadline"
else:
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED")
error = "Serverless work stopped" if terminal or not submitted else "Serverless cancellation requested; provider confirmation is pending"
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
await self._deliver(run, error=error)
except Exception as exc:
# If polling/decoding failed, stop only our own provider job, never the pool.
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract"
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED")
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
await self._deliver(run, error=error)
finally:
# Completed tracking retains IDs/bindings only, not potentially large
# base64 inputs, workflow metadata, or a client holding credentials.
run.pop("workflow", None)
run.pop("images", None)
run["client"] = None
# A released local semaphore alone cannot release a durable reservation.
await self._notify_capacity()
async def wait_for_execution(self, execution_id):
tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id]
await asyncio.gather(*tasks, return_exceptions=False)
def _detached_records(self, worker_id=None, execution_id=None):
return [record for record in self.ledger.all()
if is_unresolved(record) and not self._task_is_active(record["record_id"])
and (worker_id is None or belongs_to_pool(record, worker_id))
and (execution_id is None or record.get("execution_id") == execution_id)]
async def _recover_provider_record(self, record, *, cancel=False):
"""Reconcile an exact saved ID, never submit or resume a workflow."""
if not record.get("job_id"):
return False
status = None
try:
client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]})
response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"]))
status = terminal_provider_status(response)
except Exception:
pass # Unreachable/expired is not proof of provider termination.
if status is not None or cancel:
self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time())
return status is not None
async def reconcile(self, worker_id=None):
"""Read terminal states for detached jobs and release only proven capacity."""
resolved = 0
async with self.recovery_lock:
for record in self._detached_records(worker_id):
resolved += await self._recover_provider_record(record)
await self._notify_capacity()
return resolved
async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False):
"""Explicit operator escape hatch for lost IDs/expired provider records.
This records a human's verification, not a cancellation or provider
status. It is never called automatically or allowed for live tasks.
"""
if provider_checked is not True:
raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it")
async with self.recovery_lock:
record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None)
if record is None:
raise ValueError("Unresolved inactive job not found for this pool")
self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(),
error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed")
await self._notify_capacity()
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
targets = {}
for rid, task in self.tasks.items():
run = self.runs[rid]
if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
continue
if execution_id and run["execution_id"] != execution_id:
continue
if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
continue
if not task.done():
# Repeated requests must not interrupt the first cancellation's cleanup.
if not task.cancelling():
task.cancel()
targets[rid] = task
if targets:
await asyncio.gather(*targets.values(), return_exceptions=True)
recovered = 0
async with self.recovery_lock:
for record in self._detached_records(worker_id, execution_id):
rid = record["record_id"]
# A task cancelled before its first step never enters _run's handlers.
# A failed start() may also leave a journalled but unstarted local slot.
if record.get("status") == "WAITING" and not record.get("job_id"):
run = self.runs.get(rid)
if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])):
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
run.pop("workflow", None)
run.pop("images", None)
await self._deliver(run, error="Serverless job cancelled before submission")
continue
# Do not send two cancellation requests for the same live job in
# one action. Its unresolved reservation survives for later recovery.
if rid in targets or collector_job_id or not record.get("job_id"):
continue
await self._recover_provider_record(record, cancel=True)
recovered += 1
await self._notify_capacity()
return len(targets) + recovered
def public_jobs(self, worker_id):
records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)]
# Never hide unresolved IDs behind the bounded completed history.
visible = [record for record in records if is_unresolved(record)]
visible.extend(record for record in records if not is_unresolved(record))
unresolved_count = sum(is_unresolved(record) for record in records)
return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")}
for record in visible[:max(20, unresolved_count)]]
+29
View File
@@ -0,0 +1,29 @@
"""Shared rules for provider capacity and durable recovery protection.
A local task ending is not evidence that provider work ended. RESOLVED is an
explicit operator acknowledgment, never a provider-reported state.
"""
PROVIDER_TERMINAL = frozenset({"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"})
TERMINAL = PROVIDER_TERMINAL | {"RESOLVED"}
PROVIDER_STATES = PROVIDER_TERMINAL | {"IN_QUEUE", "IN_PROGRESS"}
def terminal_provider_status(response):
status = response.get("status") if isinstance(response, dict) else None
return status if isinstance(status, str) and status in PROVIDER_TERMINAL else None
def is_unresolved(record):
return record.get("status") not in TERMINAL
def occupies_capacity(record):
# WAITING is journalled before admission and cannot have been submitted.
# Every other unresolved state is conservatively reserved, including a
# submission whose provider ID was lost and an unconfirmed cancellation.
return is_unresolved(record) and (record.get("status") != "WAITING" or bool(record.get("job_id")))
def belongs_to_pool(record, worker_id):
return str(record.get("worker_id")) == str(worker_id)
+121 -15
View File
@@ -29,7 +29,9 @@ from .orchestration.prompt_transform import (
generate_job_id_map,
prepare_delegate_master_prompt,
prune_prompt_for_worker,
_resolve_direct_workers_for_image_list,
)
from ..image_list.runtime import validate_image_list_regions
prompt_server = server.PromptServer.instance
@@ -74,6 +76,14 @@ def _resolve_enabled_workers(config, requested_ids=None):
elif not worker.get("enabled", False):
continue
if worker.get("type") == "serverless":
from ..utils.serverless_config import validate_serverless_worker
pool = validate_serverless_worker(worker)
for slot in range(pool["jobs_per_run"]):
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
continue
raw_port = worker.get("port", worker.get("listen_port", 8188))
try:
port = int(raw_port or 8188)
@@ -154,6 +164,21 @@ async def _prepare_worker_payload(
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
async with worker_prep_semaphore:
worker_prompt = prompt_index.copy_prompt()
if worker.get("type") == "serverless":
from .orchestration.serverless_prompt import compile_serverless_prompt
from .orchestration.serverless_media import prepare_images
from .orchestration.runpod_client import serialize_request
worker_prompt = apply_participant_overrides(
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
job_id_map, "", delegate_master, prompt_index,
)
workflow, manifest = compile_serverless_prompt(worker_prompt)
images = await prepare_images(workflow)
# Fail during preparation, before regular workers or the master are
# dispatched. submit() uses this same serializer for the sent bytes.
serialize_request(workflow, images, worker["timeout_seconds"])
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images}
worker_master_url = build_master_callback_url(
worker,
config=config,
@@ -223,8 +248,42 @@ async def orchestrate_distributed_execution(
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)
image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
image_list_mode = bool(image_list_start_ids)
if image_list_mode:
# Validate before probing or queueing anything so malformed boundaries and
# old worker releases fail before any upstream image is materialized.
validate_image_list_regions(prompt_obj)
other_distributed = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
)
if other_distributed:
raise ValueError(
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
"Ultimate SD Upscale nodes in one prompt"
)
configured_serverless = [
str(worker.get("id")) for worker in config.get("workers", [])
if worker.get("enabled", False)
and str(worker.get("type") or "local").lower() == "serverless"
and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
]
if configured_serverless:
trace_debug(
execution_trace_id,
"Image List v1 explicitly excludes unsupported serverless pools: "
+ ", ".join(configured_serverless),
)
workers = _resolve_direct_workers_for_image_list(config, requested_ids)
if not workers:
raise RuntimeError(
"Distributed Image List v1 requires an enabled selected direct local/remote worker; "
"serverless pools and master-only execution are unsupported"
)
else:
workers = _resolve_enabled_workers(config, requested_ids)
load_balance_requested = _prompt_requests_load_balance(prompt_index)
trace_debug(
execution_trace_id,
@@ -243,6 +302,10 @@ async def orchestrate_distributed_execution(
# Respect master delegate-only configuration
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
if image_list_mode:
# The master must execute each upstream dependency once to obtain the
# full typed lists; only the enclosed branch is delegated.
delegate_master = False
if not workers and delegate_master:
trace_debug(
@@ -258,8 +321,19 @@ async def orchestrate_distributed_execution(
trace_execution_id=execution_trace_id,
probe_concurrency=worker_probe_concurrency,
)
if image_list_mode:
from .orchestration.dispatch import preflight_image_list_workers
if not active_workers:
raise RuntimeError(
"No selected direct worker is reachable for Distributed Image List; "
"the master will not duplicate per-image processing"
)
idle_ids = await preflight_image_list_workers(active_workers, require_idle=True)
active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
if not active_workers:
raise RuntimeError("No feature-compatible direct worker is idle; no master job was queued")
if load_balance_requested:
if load_balance_requested and not image_list_mode:
candidate_workers = list(active_workers)
if not delegate_master:
# Include master in load balancing only when master participation is enabled.
@@ -330,9 +404,6 @@ async def orchestrate_distributed_execution(
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,
@@ -342,6 +413,7 @@ async def orchestrate_distributed_execution(
master_url,
delegate_master,
prompt_index,
workflow_meta=workflow_meta,
)
if delegate_master:
@@ -365,7 +437,7 @@ async def orchestrate_distributed_execution(
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
)
worker_payloads = []
if active_workers:
if active_workers and not image_list_mode:
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
worker_payloads = await asyncio.gather(
@@ -387,7 +459,20 @@ async def orchestrate_distributed_execution(
]
)
if worker_payloads:
serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"]
regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"]
manager = None
if serverless_jobs:
from .orchestration.serverless_runtime import get_manager
manager = get_manager()
manager.ensure_capacity(len(serverless_jobs))
# Do not leave collector queues behind when offline compilation/media validation fails.
if not image_list_mode:
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
if regular_payloads:
await asyncio.gather(
*[
dispatch_worker_prompt(
@@ -398,21 +483,42 @@ async def orchestrate_distributed_execution(
use_websocket=use_websocket,
trace_execution_id=execution_trace_id,
)
for worker, wprompt in worker_payloads
for worker, wprompt in regular_payloads
]
)
queue_result = await queue_prompt_payload(
master_prompt,
workflow_meta,
client_id,
include_queue_metadata=True,
)
try:
queue_result = await queue_prompt_payload(
master_prompt,
workflow_meta,
client_id,
include_queue_metadata=True,
)
except Exception:
if serverless_jobs:
async with prompt_server.distributed_jobs_lock:
for job_id in job_id_map.values():
prompt_server.distributed_pending_jobs.pop(job_id, None)
raise
prompt_id = queue_result["prompt_id"]
if manager is not None:
try:
await manager.start(execution_trace_id, prompt_id, serverless_jobs)
except Exception:
# Wake queued collectors with an explicit error if journal/capacity setup failed.
await manager.cancel(execution_id=execution_trace_id)
async with prompt_server.distributed_jobs_lock:
for prepared in serverless_jobs:
for binding in prepared["manifest"]:
q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
if q is not None:
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
raise
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
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
+98
View File
@@ -0,0 +1,98 @@
"""Master-only Runpod pool monitoring. Health checks never submit inference."""
from aiohttp import web
import server
from ..utils.config import load_config
from .orchestration.serverless_runtime import get_manager
from .orchestration.runpod_client import RunpodError
def _pool(worker_id):
return next((w for w in load_config().get("workers", [])
if w.get("id") == worker_id and w.get("type") == "serverless"), None)
async def _status(worker_id, fresh=False):
worker = _pool(worker_id)
if worker is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
jobs = []
try:
manager = get_manager()
jobs = manager.public_jobs(worker_id)
result = await manager.health(worker, fresh=fresh)
return web.json_response(dict(result, jobs=jobs))
except Exception as exc:
message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
@server.PromptServer.instance.routes.post("/distributed/serverless/test")
async def test_serverless_endpoint(request):
try:
data = await request.json()
worker_id = data.get("worker_id")
except Exception:
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
return await _status(worker_id, fresh=True)
@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
async def serverless_status_endpoint(request):
return await _status(request.match_info["worker_id"])
@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
async def cancel_serverless_endpoint(request):
try:
data = await request.json()
worker_id = data.get("worker_id")
except Exception:
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
if _pool(worker_id) is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
try:
manager = get_manager()
count = await manager.cancel(worker_id=worker_id)
return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
"message": "Cancellation requested for this master's jobs only"})
except Exception:
return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
async def _recover(request, *, acknowledge=False):
try:
data = await request.json()
worker_id = data.get("worker_id")
if not isinstance(worker_id, str) or not worker_id.strip():
raise ValueError()
record_id = data.get("record_id")
if acknowledge:
if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
raise ValueError()
except Exception:
return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
if _pool(worker_id) is None:
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
try:
manager = get_manager()
if acknowledge:
await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
else:
result = {"resolved": await manager.reconcile(worker_id=worker_id)}
return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
except ValueError:
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this pool"}, status=409)
except Exception:
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
async def reconcile_serverless_endpoint(request):
return await _recover(request)
@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
async def resolve_serverless_endpoint(request):
return await _recover(request, acknowledge=True)
+20 -1
View File
@@ -16,12 +16,14 @@ 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,
@@ -277,6 +279,13 @@ def _get_cuda_info():
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)
@@ -285,7 +294,9 @@ def _collect_network_info_sync():
"all_ips": all_ips,
"recommended_ip": recommended_ip,
"cuda_device": cuda_device,
"cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
"cuda_device_count": device_count,
"master_port": master_port,
"local_worker_ports": worker_ports,
}
@@ -449,6 +460,9 @@ async def launch_worker_endpoint(request):
if not worker:
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
if worker.get("type") == "serverless":
return await handle_api_error(request, "Serverless pools are managed by Runpod, not local processes", 400)
# Ensure consistent string ID
worker_id_str = worker_id
@@ -483,6 +497,8 @@ async def launch_worker_endpoint(request):
"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)
@@ -541,6 +557,9 @@ async def get_local_worker_status_endpoint(request):
worker_statuses = {}
for worker in config.get("workers", []):
# Serverless pools have no host, port, or local process.
if worker.get("type") == "serverless":
continue
# Only check local workers
host = normalize_host(worker.get("host")) or ""
if not host or host in ["localhost", "127.0.0.1"]:
+1 -1
View File
@@ -1,7 +1,7 @@
# conftest.py — project-level pytest configuration.
#
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
# (from .distributed import ...) that fail when pytest tries to import it as a
# (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.
-51
View File
@@ -1,51 +0,0 @@
"""
ComfyUI-Distributed: thin entry point.
All implementation lives in workers/, nodes/, api/.
"""
import atexit
import os
import server
from .utils.config import ensure_config_exists
from .utils.logging import debug_log
from .utils.network import cleanup_client_session
from .workers import get_worker_manager
from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
from .upscale.job_store import ensure_tile_jobs_initialized
from .nodes import (
NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS,
ImageBatchDivider,
DistributedCollectorNode,
DistributedSeed,
DistributedModelName,
DistributedValue,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from . import api # noqa: F401 - triggers all @routes.* registrations
from .api.queue_orchestration import ensure_distributed_state
ensure_config_exists()
# Aiohttp session cleanup
async def _cleanup_session():
await cleanup_client_session()
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
# Initialize distributed job state on prompt_server
prompt_server = server.PromptServer.instance
ensure_distributed_state(prompt_server)
ensure_tile_jobs_initialized()
# Worker startup
if not os.environ.get('COMFYUI_IS_WORKER'):
atexit.register(sync_cleanup)
delayed_auto_launch()
register_async_signals()
+18 -8
View File
@@ -19,8 +19,8 @@ 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.
- Pings workers (`GET /prompt`) to include only reachable ones.
- Dispatches the workflow to workers (`POST /prompt`).
- 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.
- Queues the master workflow in ComfyUI’s prompt queue.
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
@@ -61,7 +61,8 @@ Queue a workflow for distributed execution.
#### Fields
- `prompt` (required unless `workflow.prompt` is present, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
- 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.
- `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.
@@ -123,10 +124,14 @@ $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:
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`:
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
Also, for collector-based flows:
@@ -201,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
## Examples
### 1) Minimal `curl`
### 1) Minimal `curl` request envelope
```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -209,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-d @payload.json
```
Where `payload.json` contains at least:
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
```json
{
"prompt": {
"1": {"class_type": "KSampler", "inputs": {} }
"<node_id>": {
"class_type": "<node class>",
"inputs": {"<required_input>": "<value or connection>"}
}
},
"enabled_worker_ids": [],
"client_id": "external-client"
}
```
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
+88
View File
@@ -0,0 +1,88 @@
# Dynamic image-list queue
Process a list of images across **direct local or remote ComfyUI workers** without assigning fixed chunks to each worker.
The master evaluates the input list once. Each selected idle worker receives one image job; when it returns a result, it gets the next unassigned image. Faster workers can therefore process more images. The master restores input order before running downstream nodes.
## Wiring
```text
Image list or IMAGE batch
|
Distributed Image List Start -------- context ---------+
| image |
v |
Per-image processing |
| IMAGE |
v v
Distributed Image List End
| image_list
v
Preview / Save / next stage
```
1. Install the same feature-capable version of Distributed on the master and every selected worker. Restart those instances and refresh the frontend after installation.
2. Enable Distributed and select at least one reachable, idle **direct** worker in the panel. The master must be reachable from remote workers at its configured callback address.
3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; on a worker, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
## Inputs and outputs
Start accepts:
- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames.
- Optional `reference_image`, `mask`, `text`, `strength`, and `seed`.
- A `per_item` / `shared` mode for each optional input. `per_item` requires exactly one value per image; `shared` requires exactly one value, reused for all images. Inputs are never silently padded by repeating the last element.
Reference images and masks default to `per_item`. Text, strength and seed default to `shared`; strength defaults to `1.0`, seed to `0`.
Start outputs, in order: `image`, `reference_image`, `mask`, `text`, `strength`, `seed`, `item_index` (zero-based), `context`. Connect only optional values that your processing region actually provides and uses. Seeds are tied to input values, not worker assignments; use a per-item seed list for deterministic variation.
End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
Models are **loaded and cached on workers**, not serialized from the master. Connect loader outputs directly to the processing nodes. Each worker needs the region's node packs and model files. **The master still validates the prompt:** it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
End's optional `texts` is a connection-only socket. For workflows saved before this fix, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations while being read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens.
## Scope and failure behavior
- This version supports direct workers only, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
- Only the enclosed processing is distributed. List creation and downstream saving run on the master, even when master-delegate-only is enabled. There is no master fallback when workers are unavailable.
- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/reference/mask combinations can hit the aggregate request limit first.
- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
- Worker errors fail the region. The worker timeout measures loss of evidence for the **owned prompt ID**, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; transport and cleanup have finite time limits. Lost acknowledgements are **ambiguous**, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
### Time and size limits
`settings.image_list_job_timeout_seconds` defaults to `0` (no arbitrary total-run cap). A positive finite number opts into an overall deadline with owned-work cleanup; negative, non-finite or malformed settings fail explicitly. There is no fixed one-hour region deadline. The setting is separate from the existing worker liveness timeout.
Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each.
**4K float32 RGB outputs are not supported by these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. 1080p-to-4K upscaling therefore cannot use this configuration. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/reference/mask combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
## Verification
The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
```bash
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
bash scripts/test-web.sh
```
Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.
+163
View File
@@ -0,0 +1,163 @@
# Runpod serverless pools (preview)
A **Serverless (Runpod)** worker is a queue-based endpoint/pool, not a permanent
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
normal ComfyUI API workflow per logical slot and collects the returned images.
The endpoint workers **do not need ComfyUI-Distributed installed**.
## Setup
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
2. Enter its name, Runpod **endpoint ID** (not a URL), and API key.
3. Start with **Jobs per run = 1**, **Concurrent jobs = 1**, and **Timeout = 1800**.
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
endpoint health only; it does not start GPU inference or prove model/node compatibility.
5. Enable the pool when you want queued Distributed workflows to use it.
6. Use an image **Distributed Collector** and queue through Distributed as usual.
Enable **Delegate Only** if the master should collect without generating locally.
Keys are write-only. Leave the key field blank to retain a saved key, or select
**Clear saved API key** to remove it (disable the pool first). A saved key is
required before enabling a pool. Use HTTPS or a trusted local connection when
entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
## Endpoint contract
The endpoint must accept this generic input shape; named application modes are
not used:
```json
{
"input": {
"workflow": {
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}}
},
"images": []
}
}
```
For image inputs, `images` contains `{ "name": "input.png", "image": "<base64>" }`
entries and the handler must rewrite `LoadImage` references to the staged files.
The adapter transfers local input images; it does not copy models or install nodes.
All referenced models and custom nodes must already exist in the endpoint image.
On completion, Runpod's `output.images` must contain artifact records such as:
```json
{
"filename": "example_00001_.png",
"subfolder": "",
"artifact_type": "output",
"artifact_path": "example_00001_.png",
"type": "base64",
"data": "<base64>"
}
```
Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"`
with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without
Runpod credentials; private-network destinations and unsafe redirects are refused.
The handler must preserve injected `SaveImage.filename_prefix` values so results
can be matched to the right collector and slot. Duplicate basenames in different
subfolders remain different artifact identities.
A Runpod job marked `COMPLETED` is still rejected if `output.error` or
`output.errors` is nonempty, required collector outputs are missing, or image
artifacts are invalid. Failures are surfaced rather than silently returning an
incomplete serverless batch.
## Scheduling and cancellation
- **Jobs per run** adds logical collector participants. Each receives its own
Distributed Seed/Value overrides. It is not a physical GPU count.
- **Concurrent jobs** bounds this master's outstanding submissions per pool,
including retained jobs after a restart. Ambiguous submissions and unconfirmed
cancellations keep their capacity reservation until reconciled. Other
applications using the same endpoint are unaffected.
- **Load balance** selects one participant, so it does not fan out every pool slot.
- **Timeout** includes waiting for a local slot, Runpod queue/cold-start time,
execution, and output retrieval. Provider policy limits are submitted as well.
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
- Interrupting a collector or using the panel cancellation action requests
cancellation only for jobs submitted by this master. It never purges the endpoint.
- A network failure during submission is **not retried automatically**: the job
may already be billable. Check the Runpod job list before queuing again.
## Persistence and recovery
Credentials live outside `gpu_config.json` and Git in the master ComfyUI user
folder under `distributed/serverless/credentials.json`, with restrictive POSIX
permissions. They are omitted from public configuration, workflow exports, and
browser configuration state. The first version requires a POSIX (Linux/macOS)
master for serverless secret storage; existing non-serverless workers remain
available on Windows.
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
not image data, workflow payloads, or API keys. A restarted master cannot resume
an interrupted collector. It does **not** automatically resubmit jobs. Retained
IDs are available through `/distributed/serverless/status/<worker_id>`; the pool's
cancel action can request cancellation of known pre-restart jobs. If no job ID
was received, check Runpod manually. Cancellation is best-effort and may race
with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation.
Unresolved records protect the pool's endpoint, concurrency, API key, worker type,
and deletion even when no local task remains. Disabling a pool is still allowed;
it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are
closed locally, while interrupted `SUBMITTING` records without IDs become
`UNKNOWN` and remain protected.
### Reconcile or acknowledge recovery
Recovery uses the same trusted-master API as the panel; it never resubmits a job.
1. Read `GET /distributed/serverless/status/<worker_id>`. The `jobs` list includes
every unresolved record with its local `record_id`, provider `job_id` (if
received), and state. These records remain visible if endpoint health fails.
2. Request cancellation from the panel, or send
`POST /distributed/serverless/reconcile` with the body below to **read** the
retained provider IDs. Only a confirmed terminal state releases protection.
Network errors, expired IDs, or ambiguous responses do not.
```json
{"worker_id": "YOUR_POOL_ID"}
```
3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
manually first. **Only after verifying that the job is no longer running or
queued**, send `POST /distributed/serverless/resolve` for that one record:
```json
{"worker_id": "YOUR_POOL_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
```
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
provider-confirmed cancellation. It cannot resolve a live local task and does
not stop remote compute. It releases the record's reservation and may let
already-waiting jobs proceed. Never automate this acknowledgment as a retry
mechanism; an unverified job can still be billable.
Pool changes become available once all its local tasks and unresolved records
are finished. No credentials or ledger files need to be edited by hand.
## First-version limits
- Image collectors only. Audio/video transfer and audio collectors are rejected.
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
node types are rejected before submission.
- Downstream post-collector processing remains local; the remote graph contains
only the upstream compute required for image saves.
- Inputs must be supported images under ComfyUI's input/output/temp roots.
- Raw image input reads are guarded at 10,000,000 bytes total, and the complete
serialized `/run` JSON request body (workflow, base64 images, and policy) must
also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image
budget is lower and depends on workflow size. Large-media/object-storage upload
support is not included.
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
scaling changes, or inbound master callback/tunnel is needed.
This feature is verified with offline provider/transport tests. A successful
health test is not a paid end-to-end generation test; test your deployed handler
and model inventory with one slot before increasing concurrency.
+53
View File
@@ -0,0 +1,53 @@
# 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.
+2 -2
View File
@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
**On Runpod:**
> 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!**
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and 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).
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
- 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 /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
7. Download models into `/workspace/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
View File
@@ -0,0 +1 @@
"""Dynamic per-image image-list execution internals."""
+142
View File
@@ -0,0 +1,142 @@
"""Server-owned background cleanup for image-list runtime state."""
from __future__ import annotations
import asyncio
import logging
import math
import time
from collections.abc import Iterable
from typing import Any, cast
LOGGER = logging.getLogger(__name__)
DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0
DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
_LIFECYCLE_MARKER = object()
def _queue_prompt_ids(prompt_server: Any) -> set[str]:
"""Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
queue = getattr(prompt_server, "prompt_queue", None)
readers = (
getattr(queue, "get_current_queue_volatile", None),
getattr(queue, "get_current_queue", None),
)
errors: list[Exception] = []
for reader in readers:
if not callable(reader):
continue
try:
snapshot = reader()
if not isinstance(snapshot, (tuple, list)) or len(snapshot) < 2:
raise TypeError("prompt queue snapshot must contain running and pending entries")
running, pending = snapshot[0], snapshot[1]
prompt_ids: set[str] = set()
for entries in (running, pending):
if not isinstance(entries, Iterable):
raise TypeError("prompt queue entries must be iterable")
for entry in entries:
if isinstance(entry, (tuple, list)) and len(entry) > 1:
prompt_id = entry[1]
if isinstance(prompt_id, str):
prompt_ids.add(prompt_id)
return prompt_ids
except Exception as exc:
errors.append(exc)
if errors:
raise RuntimeError("unable to read the prompt queue") from errors[-1]
raise RuntimeError("prompt queue has no supported queue reader")
def _sweep_orphan_contexts(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
queue_read_ok = True
try:
queued_prompt_ids = _queue_prompt_ids(prompt_server)
except Exception:
# A failed queue read must not make a submitted context look orphaned.
LOGGER.warning("Could not read prompt queue during image-list lifecycle sweep", exc_info=True)
queued_prompt_ids = set()
queue_read_ok = False
now = time.monotonic()
with runtime._LOCK:
for context_id, state in list(runtime._MASTER_JOBS.items()):
if getattr(state, "active", False):
continue
prompt_id = getattr(state, "prompt_id", None)
if prompt_id is not None:
if not queue_read_ok or prompt_id in queued_prompt_ids:
continue
created_at = getattr(state, "created_at", None)
expired = (
isinstance(created_at, (int, float))
and not isinstance(created_at, bool)
and math.isfinite(created_at)
and now - created_at >= orphan_grace_seconds
)
if expired and runtime._MASTER_JOBS.get(context_id) is state:
runtime.remove_master_job(context_id)
def _sweep_once(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
try:
runtime._sweep_worker_items()
except Exception:
LOGGER.exception("Could not sweep expired image-list worker staging")
_sweep_orphan_contexts(prompt_server, runtime, orphan_grace_seconds)
async def _sweep_loop(
prompt_server: Any,
runtime: Any,
interval_seconds: float,
orphan_grace_seconds: float,
) -> None:
while True:
try:
_sweep_once(prompt_server, runtime, orphan_grace_seconds)
except Exception:
# Keep the server-owned task alive if one pass encounters unexpected state.
LOGGER.exception("Image-list lifecycle sweep failed")
await asyncio.sleep(interval_seconds)
def install_lifecycle(
prompt_server: Any,
runtime: Any,
*,
interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS,
orphan_grace_seconds: float = DEFAULT_ORPHAN_GRACE_SECONDS,
) -> None:
"""Register one cancellable periodic sweep in the prompt server's app context."""
interval_seconds = float(interval_seconds)
orphan_grace_seconds = float(orphan_grace_seconds)
if not math.isfinite(interval_seconds) or interval_seconds <= 0:
raise ValueError("interval_seconds must be a positive finite value")
if not math.isfinite(orphan_grace_seconds) or orphan_grace_seconds < 0:
raise ValueError("orphan_grace_seconds must be a non-negative finite value")
app = getattr(prompt_server, "app", None)
cleanup_ctx = getattr(app, "cleanup_ctx", None)
if cleanup_ctx is None:
raise TypeError("prompt_server must expose an app with cleanup_ctx")
if any(getattr(factory, "_image_list_lifecycle_marker", None) is _LIFECYCLE_MARKER
for factory in cleanup_ctx):
return
async def lifecycle_context(_app):
task = asyncio.create_task(
_sweep_loop(prompt_server, runtime, interval_seconds, orphan_grace_seconds),
name="image-list-lifecycle-sweeper",
)
try:
yield
finally:
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
cast(Any, lifecycle_context)._image_list_lifecycle_marker = _LIFECYCLE_MARKER
cleanup_ctx.append(lifecycle_context)
+921
View File
@@ -0,0 +1,921 @@
"""Runtime primitives for the V3 dynamic image-list feature.
The network protocol transports typed float tensors as bounded JSON/base64 payloads;
no Python object deserialization or arbitrary-path access is used.
"""
from __future__ import annotations
import asyncio
import base64
import json
import math
import secrets
import threading
import time
from dataclasses import dataclass, field
from typing import Any
PROTOCOL_VERSION = 1
START_CLASS = "DistributedImageListStart"
END_CLASS = "DistributedImageListEnd"
START_NODE_ID = "DistributedImageListStart"
END_NODE_ID = "DistributedImageListEnd"
# Start output indexes are a stable internal prompt protocol.
IMAGE_OUT, REFERENCE_OUT, MASK_OUT, TEXT_OUT, STRENGTH_OUT, SEED_OUT, INDEX_OUT, CONTEXT_OUT = range(8)
MAX_PAYLOAD_BYTES = 50 * 1024 * 1024
MAX_REQUEST_BYTES = MAX_PAYLOAD_BYTES * 2 + 1024 * 1024
MAX_ITEMS = 4096
SUBMISSION_TIMEOUT_SECONDS = 190.0
SUBMISSION_CANCEL_GRACE_SECONDS = 1.0
CLEANUP_TIMEOUT_SECONDS = 16.0
_ALLOWED_FLOAT_DTYPES = {"float16", "bfloat16", "float32", "float64"}
_LOCK = threading.RLock()
_MASTER_JOBS: dict[str, "ImageListJob"] = {}
_WORKER_ITEMS: dict[str, tuple[float, dict[str, Any]]] = {}
_WORKER_ITEM_TTL_SECONDS = 3600.0
class ImageListCancelledError(RuntimeError):
"""Raised after owned worker prompts have been cleaned up on cancellation."""
@dataclass
class ImageListJob:
context_id: str
start_id: str
end_id: str
original_prompt: dict[str, Any]
items: list[dict[str, Any]]
worker_ids: list[str]
master_callback_url: str
workflow_meta: dict[str, Any] | None = None
results: dict[int, dict[str, Any]] = field(default_factory=dict)
received_indexes: set[int] = field(default_factory=set)
expected_tokens: dict[int, str] = field(default_factory=dict)
expected_workers: dict[int, str] = field(default_factory=dict)
event: asyncio.Event | None = None
created_at: float = field(default_factory=time.monotonic)
failed: str | None = None
prompt_id: str | None = None
active: bool = False
def _torch():
import torch
return torch
def unwrap_single(value: Any) -> Any:
"""Unwrap the one-element wrappers V3/INPUT_IS_LIST uses for scalar inputs."""
while isinstance(value, (list, tuple)) and len(value) == 1:
value = value[0]
return value
def _image_frames(value: Any, name: str, *, mask: bool = False) -> list[Any]:
torch = _torch()
if value is None:
return []
if isinstance(value, (list, tuple)):
frames: list[Any] = []
for item in value:
frames.extend(_image_frames(item, name, mask=mask))
return frames
if not isinstance(value, torch.Tensor):
raise TypeError(f"{name} must be a ComfyUI tensor or a list of tensors")
if not value.is_floating_point():
raise TypeError(f"{name} tensor must have a floating-point dtype")
if mask:
if value.ndim == 2:
value = value.unsqueeze(0)
elif value.ndim == 4 and value.shape[-1] == 1:
value = value[..., 0]
if value.ndim != 3:
raise ValueError(f"{name} MASK must have shape [B,H,W] or [B,H,W,1]")
if min(value.shape) < 1:
raise ValueError(f"{name} MASK dimensions must be positive")
return [value[i:i + 1] for i in range(value.shape[0])]
if value.ndim == 3:
value = value.unsqueeze(0)
if value.ndim != 4:
raise ValueError(f"{name} IMAGE must have shape [B,H,W,C] or [H,W,C]")
if min(value.shape[:3]) < 1 or value.shape[-1] not in (3, 4):
raise ValueError(f"{name} IMAGE must have positive dimensions and 3 or 4 channels")
return [value[i:i + 1] for i in range(value.shape[0])]
def _scalar_items(value: Any, name: str) -> list[Any]:
if value is None:
return []
if isinstance(value, (list, tuple)):
return list(value)
return [value]
def _normalize_aligned(value: Any, mode: str, count: int, name: str, kind: str) -> list[Any]:
if value is None:
return [None] * count
if mode not in ("per_item", "shared"):
raise ValueError(f"{name}_mode must be 'per_item' or 'shared'")
if kind == "image":
items = _image_frames(value, name)
elif kind == "mask":
items = _image_frames(value, name, mask=True)
else:
items = _scalar_items(value, name)
if mode == "shared":
if len(items) != 1:
raise ValueError(f"{name} is shared but has {len(items)} values; shared inputs require exactly one value")
items = [items[0] for _ in range(count)]
if len(items) != count:
raise ValueError(f"{name} has {len(items)} per-item values; expected exactly {count}")
if kind == "text" and any(item is not None and not isinstance(item, str) for item in items):
raise TypeError(f"{name} per-item values must be strings")
if kind == "int" and any(item is not None and (isinstance(item, bool) or not isinstance(item, int)) for item in items):
raise TypeError(f"{name} per-item values must be integers")
if kind == "float" and any(item is not None and (isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item))) for item in items):
raise TypeError(f"{name} per-item values must be finite numbers")
return items
def normalize_image_list_inputs(images: Any, reference_image: Any = None, mask: Any = None,
text: Any = None, strength: Any = None, seed: Any = None,
reference_image_mode: Any = "per_item", mask_mode: Any = "per_item",
text_mode: Any = "per_item", strength_mode: Any = "per_item",
seed_mode: Any = "per_item") -> list[dict[str, Any]]:
"""Normalize batches and explicitly shared/per-item optional inputs."""
frames = _image_frames(images, "images")
if not frames:
raise ValueError("Distributed Image List Start requires at least one image")
if len(frames) > MAX_ITEMS:
raise ValueError(f"Image list contains {len(frames)} items; maximum is {MAX_ITEMS}")
count = len(frames)
refs = _normalize_aligned(reference_image, str(unwrap_single(reference_image_mode)), count,
"reference_image", "image")
masks = _normalize_aligned(mask, str(unwrap_single(mask_mode)), count, "mask", "mask")
texts = _normalize_aligned(text, str(unwrap_single(text_mode)), count, "text", "text")
strengths = _normalize_aligned(strength, str(unwrap_single(strength_mode)), count, "strength", "float")
seeds = _normalize_aligned(seed, str(unwrap_single(seed_mode)), count, "seed", "int")
return [
{"image": frames[i], "reference_image": refs[i], "mask": masks[i],
"text": texts[i], "strength": strengths[i], "seed": seeds[i], "index": i}
for i in range(count)
]
def _tensor_dtype_name(tensor: Any) -> str:
name = str(tensor.dtype).removeprefix("torch.")
if name not in _ALLOWED_FLOAT_DTYPES:
raise TypeError(f"Tensor dtype {name!r} is unsupported; use float16, bfloat16, float32, or float64")
return name
def encode_value(value: Any, kind: str) -> dict[str, Any] | None:
"""Encode only the explicit IMAGE/MASK/STRING/FLOAT/INT protocol types."""
if value is None:
return None
if kind in ("IMAGE", "MASK"):
torch = _torch()
if not isinstance(value, torch.Tensor) or not value.is_floating_point():
raise TypeError(f"{kind} result must be a floating-point tensor")
if kind == "IMAGE":
if value.ndim != 4 or value.shape[0] != 1 or value.shape[-1] not in (3, 4):
raise ValueError("IMAGE item must have shape [1,H,W,3|4]")
else:
if value.ndim == 2:
value = value.unsqueeze(0)
if value.ndim == 4 and value.shape[-1] == 1:
value = value[..., 0]
if value.ndim != 3 or value.shape[0] != 1:
raise ValueError("MASK item must have shape [1,H,W] or [1,H,W,1]")
if value.numel() * value.element_size() > MAX_PAYLOAD_BYTES:
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
cpu = value.detach().to("cpu").contiguous()
raw = cpu.view(torch.uint8).numpy().tobytes()
if len(raw) > MAX_PAYLOAD_BYTES:
raise ValueError(f"{kind} item exceeds the {MAX_PAYLOAD_BYTES}-byte transport limit")
return {"type": kind, "dtype": _tensor_dtype_name(cpu),
"shape": list(cpu.shape), "data": base64.b64encode(raw).decode("ascii")}
if kind == "STRING":
if not isinstance(value, str):
raise TypeError("STRING transport value must be a string")
return {"type": kind, "value": value}
if kind == "INT":
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError("INT transport value must be an integer")
return {"type": kind, "value": value}
if kind == "FLOAT":
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)):
raise TypeError("FLOAT transport value must be finite")
return {"type": kind, "value": float(value)}
raise ValueError(f"Unsupported image-list transport type {kind!r}")
def decode_value(data: Any, expected_kind: str) -> Any:
if data is None:
return None
if not isinstance(data, dict) or data.get("type") != expected_kind:
raise ValueError(f"Expected {expected_kind} transport value")
if expected_kind in ("IMAGE", "MASK"):
torch = _torch()
dtype_name = data.get("dtype")
if dtype_name not in _ALLOWED_FLOAT_DTYPES:
raise ValueError(f"Unsupported tensor dtype {dtype_name!r}")
shape = data.get("shape")
if not isinstance(shape, list) or not shape or len(shape) > 4:
raise ValueError("Tensor shape must be a list of 1 to 4 positive dimensions")
if any(isinstance(dim, bool) or not isinstance(dim, int) or dim < 1 for dim in shape):
raise ValueError("Tensor shape dimensions must be positive integers")
expected_count = math.prod(shape)
element_size = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}[dtype_name]
expected_bytes = expected_count * element_size
if expected_bytes > MAX_PAYLOAD_BYTES:
raise ValueError("Tensor exceeds the image-list transport limit")
encoded = data.get("data")
if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4:
raise ValueError("Tensor data is missing or exceeds the transport limit")
try:
raw = base64.b64decode(encoded, validate=True)
except Exception as exc:
raise ValueError("Tensor data must be valid base64") from exc
if len(raw) != expected_bytes:
raise ValueError("Tensor byte length does not match shape and dtype")
dtype = getattr(torch, dtype_name)
return torch.frombuffer(bytearray(raw), dtype=torch.uint8).view(dtype).reshape(shape).clone()
value = data.get("value")
return encode_value(value, expected_kind) and value
def encode_request_payload(payload: dict[str, Any]) -> bytes:
"""Budget the final JSON envelope, including base64, before HTTP upload."""
body = json.dumps(payload, separators=(",", ":"), allow_nan=False).encode("utf-8")
if len(body) > MAX_REQUEST_BYTES:
raise ValueError(
f"Image-list complete JSON/base64 request is {len(body)} bytes; "
f"limit is {MAX_REQUEST_BYTES} bytes (image/reference/mask combined)"
)
return body
def encode_item(item: dict[str, Any]) -> dict[str, Any]:
return {"image": encode_value(item["image"], "IMAGE"),
"reference_image": encode_value(item.get("reference_image"), "IMAGE"),
"mask": encode_value(item.get("mask"), "MASK"),
"text": encode_value(item.get("text"), "STRING"),
"strength": encode_value(item.get("strength"), "FLOAT"),
"seed": encode_value(item.get("seed"), "INT"),
"index": int(item["index"])}
def decode_item(data: Any) -> dict[str, Any]:
if not isinstance(data, dict):
raise ValueError("item must be an object")
index = data.get("index")
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
raise ValueError("item.index must be a non-negative integer")
image = decode_value(data.get("image"), "IMAGE")
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
raise ValueError("item.image must contain exactly one RGB/RGBA image frame")
reference = decode_value(data.get("reference_image"), "IMAGE")
if reference is not None and (reference.ndim != 4 or reference.shape[0] != 1 or reference.shape[-1] not in (3, 4)):
raise ValueError("item.reference_image must contain exactly one RGB/RGBA image frame")
mask = decode_value(data.get("mask"), "MASK")
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("item.mask must contain exactly one MASK frame")
return {"image": image,
"reference_image": reference,
"mask": mask,
"text": decode_value(data.get("text"), "STRING"),
"strength": decode_value(data.get("strength"), "FLOAT"),
"seed": decode_value(data.get("seed"), "INT"), "index": index}
def _sweep_worker_items() -> None:
now = time.monotonic()
with _LOCK:
for token, (expires, _) in list(_WORKER_ITEMS.items()):
if expires <= now:
del _WORKER_ITEMS[token]
def store_worker_item(token: str, item: dict[str, Any]) -> None:
if not isinstance(token, str) or len(token) < 32 or len(token) > 128:
raise ValueError("token must be a 32-128 character capability")
clean = decode_item(item)
_sweep_worker_items()
with _LOCK:
if token in _WORKER_ITEMS:
raise ValueError("duplicate item token")
_WORKER_ITEMS[token] = (time.monotonic() + _WORKER_ITEM_TTL_SECONDS, clean)
def consume_worker_item(token: str) -> dict[str, Any]:
if not isinstance(token, str) or not token:
raise ValueError("worker item token is required")
_sweep_worker_items()
with _LOCK:
entry = _WORKER_ITEMS.pop(token, None)
if entry is None:
raise RuntimeError("Image-list item is missing, expired, or already consumed")
return entry[1]
def delete_worker_item(token: str) -> bool:
with _LOCK:
return _WORKER_ITEMS.pop(str(token), None) is not None
def create_master_job(start_id: str, end_id: str, prompt: dict[str, Any], items: list[dict[str, Any]],
worker_ids: list[str], master_callback_url: str,
workflow_meta: dict[str, Any] | None = None,
prompt_id: str | None = None) -> str:
if not isinstance(prompt, dict):
raise ValueError("Image List Start requires the original ComfyUI prompt for region extraction")
if not worker_ids:
raise RuntimeError("Distributed Image List requires at least one preflighted direct worker; master fallback is disabled")
context_id = secrets.token_urlsafe(36)
state = ImageListJob(context_id, str(start_id), str(end_id), prompt, items,
[str(worker_id) for worker_id in worker_ids], master_callback_url,
workflow_meta, prompt_id=prompt_id)
with _LOCK:
_MASTER_JOBS[context_id] = state
return context_id
def get_master_job(context_id: str) -> ImageListJob | None:
with _LOCK:
return _MASTER_JOBS.get(context_id)
def remove_master_job(context_id: str) -> None:
with _LOCK:
_MASTER_JOBS.pop(context_id, None)
def _prompt_links(node: Any, prompt: dict[str, Any]):
if not isinstance(node, dict):
return
inputs = node.get("inputs", {})
if not isinstance(inputs, dict):
return
for input_name, value in inputs.items():
if isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt:
yield input_name, str(value[0]), value[1]
def _direct_link(value: Any, prompt: dict[str, Any], node_id: str, input_name: str) -> tuple[str, int]:
if not (isinstance(value, (list, tuple)) and len(value) == 2 and str(value[0]) in prompt):
raise ValueError(f"{node_id}.{input_name} must be connected to an in-graph node")
try:
return str(value[0]), int(value[1])
except (TypeError, ValueError) as exc:
raise ValueError(f"{node_id}.{input_name} has an invalid link") from exc
def _ancestors_for_inputs(prompt: dict[str, Any], end_id: str, start_id: str,
input_names: tuple[str, ...]) -> set[str]:
result = {end_id}
pending = []
end_inputs = prompt[end_id].get("inputs", {})
for name in input_names:
if name in end_inputs:
source, _ = _direct_link(end_inputs[name], prompt, end_id, name)
pending.append(source)
while pending:
node_id = pending.pop()
if node_id in result:
continue
result.add(node_id)
if node_id == start_id:
continue
node = prompt[node_id]
for _, source, _ in _prompt_links(node, prompt):
if source not in result:
pending.append(source)
return result
def _downstream_until(prompt: dict[str, Any], start_id: str, end_id: str) -> set[str]:
downstream: dict[str, set[str]] = {str(node_id): set() for node_id in prompt}
for node_id, node in prompt.items():
for _, source, _ in _prompt_links(node, prompt):
downstream[source].add(str(node_id))
visited = set()
pending = [start_id]
while pending:
node_id = pending.pop()
for target in downstream.get(node_id, ()):
if target == end_id or target in visited:
continue
visited.add(target)
pending.append(target)
return visited
def _start_ancestors(prompt: dict[str, Any], start_id: str) -> set[str]:
"""All dependencies materialized on the master, without crossing Start again."""
ancestors = set()
pending = [source for _, source, _ in _prompt_links(prompt[start_id], prompt)]
while pending:
node_id = pending.pop()
if node_id == start_id or node_id in ancestors:
continue
ancestors.add(node_id)
pending.extend(source for _, source, _ in _prompt_links(prompt[node_id], prompt))
return ancestors
def _safe_shared_dependency(node: dict[str, Any], prompt: dict[str, Any], registry) -> bool:
"""Only typed model dependencies or scalar constants may straddle Start.
A leaf alone is not sufficient: LoadImage/EmptyImage are leaves too. Unknown
classes and list-producing nodes fail closed. All shared ancestors are checked,
so a model-producing node cannot hide a shared image/latent dependency.
"""
node_class = registry.get(node.get("class_type"))
if node_class is None:
return False
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
outputs = info.get("output", getattr(node_class, "RETURN_TYPES", ()))
if not outputs or any(info.get("output_is_list", getattr(node_class, "OUTPUT_IS_LIST", ()))):
return False
types = {str(output) for output in outputs}
model_types = {"MODEL", "CLIP", "VAE", "CLIP_VISION", "CONTROL_NET"}
scalar_types = {"INT", "FLOAT", "STRING", "BOOLEAN"}
if types <= model_types:
return True
return types <= scalar_types and not any(_prompt_links(node, prompt))
def validate_image_list_regions(prompt: dict[str, Any]) -> list[dict[str, Any]]:
"""Validate pairing, non-nesting, and closed per-image graph boundaries."""
if not isinstance(prompt, dict):
raise ValueError("Prompt must be an object")
starts = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == START_CLASS)
ends = sorted(str(k) for k, v in prompt.items() if isinstance(v, dict) and v.get("class_type") == END_CLASS)
if len(starts) != len(ends) or not starts:
raise ValueError("Distributed Image List needs one paired Start and End per region")
regions = []
for start_id in starts:
candidates = []
for end_id in ends:
context = prompt[end_id].get("inputs", {}).get("context")
if isinstance(context, (list, tuple)) and len(context) == 2 and str(context[0]) == start_id and context[1] == CONTEXT_OUT:
candidates.append(end_id)
if len(candidates) != 1:
raise ValueError(f"Start node {start_id} must connect its typed context output directly to exactly one End")
end_id = candidates[0]
if "images" not in prompt[end_id].get("inputs", {}):
raise ValueError(f"End node {end_id}.images must be connected to the per-item result branch")
end_inputs = prompt[end_id].get("inputs", {})
for name in ("images", "masks", "texts"):
if name in end_inputs:
_direct_link(end_inputs[name], prompt, end_id, name)
region_nodes = _ancestors_for_inputs(prompt, end_id, start_id, ("images", "masks", "texts"))
if start_id not in region_nodes:
raise ValueError(f"Start node {start_id} is not an ancestor of End node {end_id}")
forbidden = {n for n in region_nodes if n != start_id and n != end_id and
isinstance(prompt[n], dict) and prompt[n].get("class_type") in
{START_CLASS, END_CLASS, "DistributedCollector", "UltimateSDUpscaleDistributed"}}
if forbidden:
raise ValueError("Nested/ambiguous regions and other distributed nodes are unsupported in an image-list region: " + ", ".join(sorted(forbidden)))
# Output nodes execute independently of the lazy End, so they would
# evaluate the processing branch on the master as well as the worker.
import sys
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
shared = region_nodes & _start_ancestors(prompt, start_id)
unsafe = {node_id for node_id in shared
if not _safe_shared_dependency(prompt[node_id], prompt, registry)}
if unsafe:
links = [f"{node_id}.{name} <- {source}:{slot}"
for node_id in sorted(region_nodes - {start_id})
for name, source, slot in _prompt_links(prompt[node_id], prompt)
if source in unsafe]
raise ValueError(
f"Image-list region {start_id}->{end_id} bypasses Start through upstream nodes "
f"{sorted(unsafe)} ({'; '.join(links)}); route per-item data through Start"
)
for node_id in region_nodes - {start_id, end_id}:
node_class = registry.get(prompt[node_id].get("class_type"))
if node_class is None:
continue
info = node_class.GET_NODE_INFO_V1() if hasattr(node_class, "GET_NODE_INFO_V1") else {}
if info.get("output_node", getattr(node_class, "OUTPUT_NODE", False)):
raise ValueError(f"Output node {node_id} must be placed after Distributed Image List End")
downstream = _downstream_until(prompt, start_id, end_id)
outside = sorted(downstream - region_nodes)
if outside:
raise ValueError(f"Image-list region {start_id}->{end_id} has unsafe cross-boundary nodes: {outside}")
# Model loaders and constants can be shared; item-dependent processing cannot.
regions.append({"start_id": start_id, "end_id": end_id,
"node_ids": region_nodes, "branch_ids": (region_nodes & downstream) - {start_id}})
regions.sort(key=lambda region: (0, int(region["start_id"])) if region["start_id"].isdigit()
else (1, region["start_id"]))
internal_sets = [r["branch_ids"] for r in regions]
for i in range(len(regions)):
for j in range(i + 1, len(regions)):
overlap = internal_sets[i] & internal_sets[j]
if overlap:
raise ValueError(f"Image-list regions share ambiguous processing nodes: {sorted(overlap)}")
return regions
def build_worker_prompt(prompt: dict[str, Any], region: dict[str, Any], item_token: str,
context_id: str, item_index: int, callback_url: str) -> dict[str, Any]:
"""Copy only the per-image branch and its worker-side dependencies."""
start_id, end_id = region["start_id"], region["end_id"]
_direct_link(prompt[end_id].get("inputs", {}).get("context"), prompt, end_id, "context")
kept = {node_id: json.loads(json.dumps(prompt[node_id])) for node_id in region["node_ids"]}
start_inputs = kept[start_id].setdefault("inputs", {})
for name in ("images", "reference_image", "mask", "text", "strength", "seed"):
start_inputs.pop(name, None)
start_inputs.update({"worker_item_token": item_token, "worker_context_id": context_id})
end_inputs = kept[end_id].setdefault("inputs", {})
end_inputs.update({"worker_result_token": item_token,
"worker_item_index": int(item_index),
"master_callback_url": callback_url})
return kept
def create_worker_task_token() -> str:
return secrets.token_urlsafe(36)
async def dynamic_schedule(item_count: int, workers: list[str], dispatch, wait_next,
cleanup, interrupted=None):
"""One in-flight job per worker; completed workers immediately take next item."""
if item_count <= 0:
raise ValueError("Dynamic image-list queue has no items")
if not workers:
raise RuntimeError("No free, reachable, feature-compatible direct workers are available")
pending_index = 0
inflight: dict[str, Any] = {}
results: dict[int, Any] = {}
def check_interrupted():
if interrupted is not None and interrupted():
raise ImageListCancelledError("Image-list job interrupted by user")
async def next_result():
# Keep one waiter alive while checking ComfyUI's thread-safe interrupt flag.
# Restarting wait_next on a timeout can lose a result it already consumed.
waiter = asyncio.create_task(wait_next(inflight))
try:
while True:
check_interrupted()
done, _ = await asyncio.wait({waiter}, timeout=0.1)
check_interrupted()
if done:
return waiter.result()
finally:
if not waiter.done():
waiter.cancel()
await asyncio.gather(waiter, return_exceptions=True)
async def submit(worker, index):
# Bound the shield itself, and inspect the real Comfy interrupt flag while
# awaiting an acknowledgement. Briefly retain a late acknowledged handle.
task = asyncio.create_task(asyncio.wait_for(
dispatch(worker, index), SUBMISSION_TIMEOUT_SECONDS))
try:
while True:
check_interrupted()
done, _ = await asyncio.wait({task}, timeout=0.1)
if done:
handle = task.result()
inflight[worker] = handle
check_interrupted()
return handle
except (asyncio.CancelledError, ImageListCancelledError):
try:
inflight[worker] = await asyncio.wait_for(
asyncio.shield(task), SUBMISSION_CANCEL_GRACE_SECONDS)
except (Exception, asyncio.CancelledError):
pass
raise
finally:
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
try:
for worker in workers:
check_interrupted()
if pending_index >= item_count:
break
inflight[worker] = await submit(worker, pending_index)
pending_index += 1
while inflight:
check_interrupted()
worker, index, result = await next_result()
if worker not in inflight:
raise RuntimeError(f"Unexpected completion from unassigned worker {worker}")
if isinstance(index, bool) or not isinstance(index, int) or index != inflight[worker].get("index"):
raise RuntimeError(f"Worker {worker} returned an unexpected image index {index!r}")
if index in results:
raise RuntimeError(f"Duplicate image-list result for item {index}")
results[index] = result
del inflight[worker]
if pending_index < item_count:
inflight[worker] = await submit(worker, pending_index)
pending_index += 1
missing = sorted(set(range(item_count)) - set(results))
if missing:
raise RuntimeError(f"Image-list results are incomplete; missing indexes {missing}")
return [results[i] for i in range(item_count)]
except BaseException:
if inflight:
await asyncio.gather(*(asyncio.wait_for(cleanup(worker, handle), CLEANUP_TIMEOUT_SECONDS)
for worker, handle in inflight.items()), return_exceptions=True)
raise
def set_master_result(context_id: str, item_index: int, token: str, result: dict[str, Any]) -> None:
with _LOCK:
state = _MASTER_JOBS.get(context_id)
if state is None:
raise ValueError("Unknown or expired image-list context")
expected = state.expected_tokens.get(item_index)
if expected is None or not secrets.compare_digest(expected, token):
raise ValueError("Image-list result token does not match an assigned item")
if item_index in state.received_indexes:
raise ValueError(f"Duplicate image-list result for item {item_index}")
state.received_indexes.add(item_index)
state.results[item_index] = result
event = state.event
if event is not None:
event.set()
def register_assignment(context_id: str, item_index: int, worker_id: str, token: str) -> None:
state = get_master_job(context_id)
if state is None:
raise RuntimeError("Image-list context expired before worker dispatch")
with _LOCK:
if item_index in state.expected_tokens:
raise RuntimeError(f"Image-list item {item_index} was already assigned")
state.expected_tokens[item_index] = token
state.expected_workers[item_index] = worker_id
def decode_result(result: Any) -> dict[str, Any]:
if not isinstance(result, dict):
raise ValueError("result must be an object")
image = decode_value(result.get("image"), "IMAGE")
if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4):
raise ValueError("result.image must contain exactly one RGB/RGBA image")
mask = decode_value(result.get("mask"), "MASK")
if mask is not None and (mask.ndim != 3 or mask.shape[0] != 1):
raise ValueError("result.mask must contain exactly one MASK frame")
text = decode_value(result.get("text"), "STRING")
return {"image": image, "mask": mask, "text": text}
def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) -> dict[str, Any]:
def one(value, name, kind):
if value is None:
return None
if kind in ("IMAGE", "MASK"):
values = _image_frames(value, name, mask=(kind == "MASK"))
if len(values) != 1:
raise ValueError(f"Worker End expects exactly one {kind} result, received {len(values)}")
return values[0]
values = _scalar_items(value, name)
values = [v for v in values if v is not None]
if not values:
return None
if len(values) != 1:
raise ValueError(f"Worker End expects at most one {kind} result, received {len(values)}")
return values[0]
image = one(images, "images", "IMAGE")
if image is None:
raise ValueError("Worker End received no image result")
mask = one(masks, "masks", "MASK")
text = one(texts, "texts", "STRING")
if text is not None and not isinstance(text, str):
raise TypeError("Worker End text result must be a string")
return {"image": image, "mask": mask, "text": text}
def encode_result(result: dict[str, Any]) -> dict[str, Any]:
normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text"))
return {"image": encode_value(normalized["image"], "IMAGE"),
"mask": encode_value(normalized["mask"], "MASK"),
"text": encode_value(normalized["text"], "STRING")}
def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]:
if not results:
raise RuntimeError("Image-list execution returned no results")
images, masks, texts = [], [], []
for expected_index, result in enumerate(results):
if not isinstance(result, dict) or "image" not in result:
raise RuntimeError(f"Image-list result {expected_index} is missing its image")
images.append(result["image"])
masks.append(result.get("mask"))
texts.append(result.get("text"))
return images, masks, texts
async def run_image_list_job(context_id: str) -> tuple[list[Any], list[Any], list[Any]]:
"""Resolve compatible idle direct workers and schedule one image at a time."""
state = get_master_job(context_id)
if state is None:
raise RuntimeError("Distributed Image List context is missing or expired")
with _LOCK:
state.active = True
try:
from ..utils.config import get_image_list_job_timeout_seconds
deadline = get_image_list_job_timeout_seconds()
timer = asyncio.timeout(deadline if deadline > 0 else None)
try:
async with timer:
return await _run_image_list_job(state)
except TimeoutError:
if timer.expired():
raise TimeoutError(f"Image-list total deadline exceeded ({deadline:g}s); cleanup attempted") from None
raise
finally:
remove_master_job(context_id)
async def _run_image_list_job(state: ImageListJob):
context_id = state.context_id
state.event = asyncio.Event()
from ..api.orchestration.dispatch import (
preflight_image_list_workers, dispatch_worker_prompt,
)
from ..api.orchestration.prompt_transform import _resolve_direct_workers_for_image_list
from ..utils.config import load_config, get_worker_timeout_seconds
from ..utils.network import build_master_callback_url, build_worker_url, get_client_session
from ..utils.logging import log
import aiohttp
config = load_config()
workers = _resolve_direct_workers_for_image_list(config, state.worker_ids)
if not workers:
raise RuntimeError("Distributed Image List requires enabled direct workers; serverless pools are unsupported")
# Recheck worker capability and idle state immediately before the first item.
worker_statuses = await preflight_image_list_workers(workers, require_idle=True)
workers = [worker for worker in workers if worker["id"] in worker_statuses]
if not workers:
raise RuntimeError("No free, reachable, feature-compatible direct workers are available; master execution is disabled")
session = await get_client_session()
regions = validate_image_list_regions(state.original_prompt)
region = next((entry for entry in regions if entry["start_id"] == state.start_id and entry["end_id"] == state.end_id), None)
if region is None:
raise RuntimeError("Image-list region no longer matches the submitted prompt")
timeout_seconds = max(5.0, float(get_worker_timeout_seconds()))
use_websocket = bool((config.get("settings", {}) or {}).get("websocket_orchestration", False))
prepared_worker_prompts: dict[str, dict[str, Any]] = {}
owned_uploads: dict[str, dict[str, Any]] = {}
async def dispatch(worker_id: str, index: int) -> dict[str, Any]:
worker = next((entry for entry in workers if entry["id"] == worker_id), None)
if worker is None:
raise RuntimeError(f"Worker {worker_id} was removed before dispatch")
token = create_worker_task_token()
register_assignment(context_id, index, worker_id, token)
item_payload = {"context_id": context_id, "item": encode_item(state.items[index])}
body = encode_request_payload({"token": token, **item_payload})
handle = {"prompt_id": None, "token": token, "index": index, "worker": worker,
"submission_pending": False}
# Own even a lost upload response; DELETE is idempotent.
owned_uploads[token] = handle
input_url = build_worker_url(worker, "/distributed/image-list/item")
async with session.post(input_url, data=body, headers={"Content-Type": "application/json"},
timeout=aiohttp.ClientTimeout(total=60)) as response:
if response.status >= 400:
raise RuntimeError(f"Worker {worker_id} rejected image-list item {index}: HTTP {response.status}")
callback_url = build_master_callback_url(worker, config=config)
if worker_id not in prepared_worker_prompts:
worker_prompt = build_worker_prompt(state.original_prompt, region, token, context_id,
index, callback_url)
if worker.get("host") and str(worker.get("type") or "local").lower() != "local":
from ..api.orchestration.media_sync import fetch_worker_path_separator, convert_paths_for_platform
separator = await fetch_worker_path_separator(worker)
if separator:
worker_prompt = convert_paths_for_platform(worker_prompt, separator)
from ..api.orchestration.media_sync import sync_worker_media
await asyncio.wait_for(sync_worker_media(worker, worker_prompt), timeout=60)
prepared_worker_prompts[worker_id] = worker_prompt
# Reuse media rewrites, never assignment identity or a mutable submitted graph.
worker_prompt = build_worker_prompt(prepared_worker_prompts[worker_id], region, token,
context_id, index, callback_url)
handle["submission_pending"] = True
prompt_id = await dispatch_worker_prompt(
worker, worker_prompt, state.workflow_meta,
client_id=f"image-list-{context_id[:20]}", use_websocket=use_websocket,
return_prompt_id=True,
)
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError(f"Worker {worker_id} did not return a prompt_id; submission status is ambiguous")
handle.update(prompt_id=prompt_id, last_seen=time.monotonic(), submission_pending=False)
return handle
async def wait_next(inflight: dict[str, Any]):
while True:
if state.failed:
raise RuntimeError(state.failed)
# Treat worker-side execution errors as failures even when End never posts a result.
async def history_status(worker_id, handle):
url = build_worker_url(handle["worker"], f"/history/{handle['prompt_id']}")
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as response:
if response.status >= 400:
return None
data = await response.json()
entry = data.get(handle["prompt_id"])
if not isinstance(entry, dict):
async with session.get(build_worker_url(handle["worker"], "/queue"),
timeout=aiohttp.ClientTimeout(total=3)) as response:
response.raise_for_status()
queue = await response.json()
if isinstance(queue, dict) and any(
isinstance(row, (list, tuple)) and len(row) > 1 and row[1] == handle["prompt_id"]
for key in ("queue_running", "queue_pending")
for row in (queue.get(key) if isinstance(queue.get(key), list) else [])
):
handle["last_seen"] = time.monotonic()
return None
status = entry.get("status", {}) or {}
if status.get("status_str") == "error":
return f"Worker {worker_id} failed image-list item {handle['index']}: {status.get('messages', status)}"
if status.get("completed") and status.get("status_str") not in (None, "success"):
return f"Worker {worker_id} failed image-list item {handle['index']}: {status}"
if status.get("completed"):
with _LOCK:
result = state.results.get(handle["index"])
if result is None:
return f"Worker {worker_id} completed item {handle['index']} without posting a result"
return worker_id, handle["index"], result
except (asyncio.TimeoutError, aiohttp.ClientError):
return None
return None
errors = await asyncio.gather(*(history_status(wid, handle) for wid, handle in inflight.items()))
failure = next((error for error in errors if isinstance(error, str)), None)
if failure:
raise RuntimeError(failure)
completed = next((entry for entry in errors if isinstance(entry, tuple)), None)
if completed is not None:
# The callback precedes the worker prompt finishing. Wait for its
# successful history before releasing the slot / next region.
with _LOCK:
state.results.pop(completed[1], None)
owned_uploads.pop(inflight[completed[0]]["token"], None)
return completed
now = time.monotonic()
expired = [f"{wid}/item {handle['index']}" for wid, handle in inflight.items()
if now - handle["last_seen"] > timeout_seconds]
if expired:
raise TimeoutError(f"Image-list owned-prompt liveness timeout after {timeout_seconds:.1f}s: {', '.join(expired)}")
await asyncio.sleep(0.25)
async def cleanup(worker_id: str, handle: dict[str, Any]) -> None:
worker = handle["worker"]
prompt_id = handle.get("prompt_id")
if prompt_id:
# Try both independently: failure of queue deletion must not suppress
# the targeted interrupt. No global interrupt or queue clearing.
for route, body in (("/queue", {"delete": [prompt_id]}),
("/interrupt", {"prompt_id": prompt_id})):
try:
async with session.post(build_worker_url(worker, route), json=body,
timeout=aiohttp.ClientTimeout(total=5)) as response:
response.raise_for_status()
except Exception as exc:
log(f"[Distributed Image List] Owned prompt cleanup unconfirmed for worker {worker_id} ({route}): {exc}")
elif handle.get("submission_pending"):
log(f"[Distributed Image List] Worker {worker_id} submission status is ambiguous; "
"no prompt ID to target. Check worker manually; not retrying.")
try:
async with session.delete(build_worker_url(worker, f"/distributed/image-list/item/{handle['token']}"),
timeout=aiohttp.ClientTimeout(total=5)) as response:
response.raise_for_status()
except Exception as exc:
log(f"[Distributed Image List] Staged input cleanup unconfirmed for worker {worker_id}: {exc}")
finally:
owned_uploads.pop(handle["token"], None)
def interrupted() -> bool:
try:
import comfy.model_management
return bool(comfy.model_management.processing_interrupted())
except Exception:
return False
try:
results = await dynamic_schedule(len(state.items), [worker["id"] for worker in workers],
dispatch, wait_next, cleanup, interrupted)
return _collect_output_lists(results)
finally:
# Upload ownership exists before dynamic_schedule has an acknowledged
# handle. Cover rejection/lost upload responses as well as cancellation.
pending = list(owned_uploads.values())
outcomes = await asyncio.gather(*(
asyncio.wait_for(cleanup(handle["worker"]["id"], handle), CLEANUP_TIMEOUT_SECONDS)
for handle in pending), return_exceptions=True)
for handle, outcome in zip(pending, outcomes):
if isinstance(outcome, BaseException):
log(f"[Distributed Image List] Cleanup unconfirmed for worker {handle['worker']['id']}: {outcome}")
+1 -31
View File
@@ -1,31 +1 @@
from .utilities import (
DistributedSeed,
DistributedModelName,
DistributedValue,
ImageBatchDivider,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
ByPassTypeTuple,
any_type,
)
from .collector import DistributedCollectorNode
NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"DistributedValue": DistributedValue,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed",
"DistributedModelName": "Distributed Model Name",
"DistributedValue": "Distributed Value",
"ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider",
"DistributedEmptyImage": "Distributed Empty Image",
}
"""Private execution helpers; public registration lives in nodes.v3."""
+221 -148
View File
@@ -29,7 +29,6 @@ class DistributedCollectorNode:
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"load_balance": (
"BOOLEAN",
{
@@ -38,7 +37,10 @@ class DistributedCollectorNode:
},
),
},
"optional": { "audio": ("AUDIO",) },
"optional": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
},
"hidden": {
"multi_job_id": ("STRING", {"default": ""}),
"is_worker": ("BOOLEAN", {"default": False}),
@@ -102,8 +104,9 @@ class DistributedCollectorNode:
return None
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
def run(self, images, 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):
images = self._normalize_images_input(images)
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):
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)
@@ -115,6 +118,14 @@ class DistributedCollectorNode:
pass_through = self._unwrap_list_input(pass_through)
delegate_only = self._unwrap_list_input(delegate_only)
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}
@@ -140,42 +151,66 @@ class DistributedCollectorNode:
)
return result
@staticmethod
def _get_serverless_manager():
"""Return the serverless manager only for collector jobs with serverless slots."""
manager = getattr(prompt_server, "distributed_serverless", None)
if manager is not None:
return manager
from ..api.orchestration.serverless_runtime import get_manager
return get_manager()
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
"""Send image batch to master via canonical JSON envelopes."""
batch_size = image_batch.shape[0]
if batch_size == 0:
return
"""Send an image batch, optionally with audio, or an audio-only completion."""
encoded_audio = encode_audio_payload(audio)
session = await get_client_session()
url = f"{master_url}/distributed/job_complete"
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 = []
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:
timeout_seconds = 60 if "image" in payload else 600
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
) as response:
response.raise_for_status()
except Exception as e:
log(f"Worker - Failed to send canonical image envelope to master: {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 # Re-raise to handle at caller level
raise
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
"""Combine audio from master and workers into a single audio output.
@@ -257,8 +292,8 @@ class DistributedCollectorNode:
images_on_cpu,
delegate_mode: bool,
fallback_images,
) -> torch.Tensor:
"""Assemble final tensor: master first, then workers in enabled order."""
):
"""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):
@@ -289,15 +324,15 @@ class DistributedCollectorNode:
if cpu_tensors:
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
elif fallback_images is not None:
if fallback_images is not None:
return ensure_contiguous(fallback_images)
else:
raise ValueError("No image data collected from master or workers")
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):
if is_worker:
# Worker mode: send images and audio to master in a single batch
debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")
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)
return (images, audio if audio is not None else self.EMPTY_AUDIO)
else:
@@ -313,6 +348,17 @@ class DistributedCollectorNode:
seen_enabled.add(worker_id_str)
enabled_workers.append(worker_id_str)
expected_workers = set(enabled_workers)
serverless_workers = {
worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
}
serverless_manager = None
def get_serverless_manager():
nonlocal serverless_manager
if serverless_workers and serverless_manager is None:
serverless_manager = self._get_serverless_manager()
return serverless_manager
num_workers = len(expected_workers)
if num_workers == 0:
return (images, audio if audio is not None else self.EMPTY_AUDIO)
@@ -332,14 +378,15 @@ class DistributedCollectorNode:
master_audio = None
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
else:
images_on_cpu = images.cpu()
master_batch_size = images.shape[0]
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...")
# Ensure master images are contiguous
images_on_cpu = ensure_contiguous(images_on_cpu)
# Initialize storage for collected images and audio
worker_images = {} # Dict to store images by worker_id and index
@@ -378,151 +425,177 @@ class DistributedCollectorNode:
workers_done.add(done_worker_id)
p.update(1) # +1 per completed expected worker
def consume_result(result):
nonlocal collected_count
if "serverless_error" in result:
error = result.get("serverless_error")
message = str(error).strip() if error is not None else "unspecified serverless failure"
participant = result.get("worker_id", "unknown")
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
result_worker_id = str(result["worker_id"])
is_last = result.get("is_last", False)
collected_count += self._store_worker_result(worker_images, result)
debug_log(
f"Master - Got canonical result from worker {result_worker_id}, "
f"image {result.get('image_index', 0)}, is_last={is_last}"
)
result_audio = result.get("audio")
if result_audio is not None:
worker_audio[result_worker_id] = result_audio
debug_log(f"Master - Got audio from worker {result_worker_id}")
if is_last:
mark_worker_done(result_worker_id)
try:
while len(workers_done) < num_workers:
# Check for user interruption to abort collection promptly
# 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}"
)
consume_result(result)
# 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
# Real results refresh the regular-worker 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
missing_workers = expected_workers - workers_done
# A queue delivery can race the timed wait. Process it before
# declaring an owned provider slot stopped without a result.
async with prompt_server.distributed_jobs_lock:
current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
queue_has_results = current_q is not None and not current_q.empty()
if queue_has_results:
continue
missing_serverless = missing_workers & serverless_workers
if missing_serverless:
manager = get_serverless_manager()
if manager is None:
raise RuntimeError("Serverless manager is unavailable for pending collector slots")
pending_serverless = []
for participant_id in sorted(missing_serverless):
if manager.is_pending(multi_job_id, participant_id):
pending_serverless.append(participant_id)
elif manager.owns(multi_job_id, participant_id):
raise RuntimeError(
f"Serverless participant {participant_id} stopped without delivering a result"
)
# Provider deadlines are managed by the serverless runtime;
# do not apply the regular worker heartbeat timeout to it.
if pending_serverless:
continue
missing_regular = missing_workers - serverless_workers
elapsed = time.time() - last_activity
for missing_worker_id in sorted(missing_workers):
for missing_worker_id in sorted(missing_regular):
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)"
)
if missing_regular:
log(
f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
f"(elapsed={elapsed:.1f}s)"
)
# Probe missing workers' /prompt endpoints to check if they are actively processing
# Probe only regular worker /prompt endpoints.
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."
if missing_regular:
try:
cfg = load_config()
cfg_workers = cfg.get('workers', [])
for wid in sorted(missing_regular):
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}"
)
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 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 = []
final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
remaining_items = []
if final_q is not None:
while not final_q.empty():
try:
item = final_q.get_nowait()
remaining_items.append(item)
remaining_items.append(final_q.get_nowait())
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!")
for item in remaining_items:
consume_result(item)
break
except comfy.model_management.InterruptProcessingException:
# Cleanup queue on interruption and re-raise to abort prompt cleanly
total_collected = sum(len(imgs) for imgs in worker_images.values())
# Remove the queue before combining so a late provider callback cannot
# resurrect results for a collector that has completed.
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]
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
try:
combined = self._reorder_and_combine_tensors(
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
)
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})")
# Combine audio from master and workers
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
return (combined, combined_audio)
except Exception as e:
log(f"Master - Error combining images: {e}")
# Return just the master images as fallback
return (images, audio if audio is not None else self.EMPTY_AUDIO)
serverless_results = set(worker_images) & serverless_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:
if serverless_results:
raise RuntimeError(
"Serverless result batch shape mismatch while combining participants "
f"{sorted(serverless_results)}: {e}"
) from e
log(f"Master - Error combining images: {e}")
# Preserve collected audio for regular workers when image assembly fails.
return (images, combined_audio)
except BaseException:
# Remove the queue first, preventing cancelled work from delivering into
# a collector that is already unwinding.
async with prompt_server.distributed_jobs_lock:
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
if serverless_workers:
try:
manager = get_serverless_manager()
if manager is not None:
await manager.cancel(collector_job_id=multi_job_id)
except Exception as cancel_error:
debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
raise
+6 -15
View File
@@ -56,22 +56,22 @@ class UltimateSDUpscaleDistributed(
"""
Distributed version of Ultimate SD Upscale (No Upscale).
Supports three processing modes:
Supports two currently selected processing modes:
1. Single GPU: No workers available, process everything locally
2. Static Mode: Small batches, distributes tiles across workers (flattened)
3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
2. Distributed tile queue: Workers pull tiles from a shared queue
Features:
- Multi-mode batch handling for efficient video/image upscaling
- Tile-based batch handling for video/image upscaling
- Tiled VAE support for memory efficiency
- Dynamic load balancing for large batches
- 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)
Threshold: dynamic_threshold input controls mode switch (default 8)
The hidden dynamic_threshold input is retained for workflow compatibility but
does not affect the current mode-selection policy.
"""
def __init__(self):
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
# Ensure initialization before registering routes
ensure_tile_jobs_initialized()
# Node registration
NODE_CLASS_MAPPINGS = {
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
}
+230
View File
@@ -0,0 +1,230 @@
"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
from __future__ import annotations
import json
from comfy_api.v0_0_2 import io
from ..utils.async_helpers import run_async_in_server_loop
from ..image_list.runtime import (
create_master_job,
normalize_image_list_inputs,
normalize_worker_result,
run_image_list_job,
encode_result,
encode_request_payload,
unwrap_single,
)
class DistributedImageListStart(io.ComfyNode):
"""Materialize a batch once, then expose ordered per-item typed values."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListStart",
display_name="Distributed Image List Start",
category="image/distributed",
description=("Starts one per-image processing region. Optional inputs explicitly choose shared "
"(exactly one value) or per-item (exactly one value per image) alignment."),
inputs=[
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one worker job."),
io.Image.Input("reference_image", optional=True),
io.Combo.Input("reference_image_mode", options=["per_item", "shared"], default="per_item"),
io.Mask.Input("mask", optional=True),
io.Combo.Input("mask_mode", options=["per_item", "shared"], default="per_item"),
io.String.Input("text", optional=True),
io.Combo.Input("text_mode", options=["per_item", "shared"], default="shared"),
io.Float.Input("strength", optional=True, default=1.0),
io.Combo.Input("strength_mode", options=["per_item", "shared"], default="shared"),
io.Int.Input("seed", optional=True, default=0, min=0, max=0xffffffffffffffff),
io.Combo.Input("seed_mode", options=["per_item", "shared"], default="shared"),
],
outputs=[
io.Image.Output(display_name="image", is_output_list=True),
io.Image.Output(display_name="reference_image", is_output_list=True),
io.Mask.Output(display_name="mask", is_output_list=True),
io.String.Output(display_name="text", is_output_list=True),
io.Float.Output(display_name="strength", is_output_list=True),
io.Int.Output(display_name="seed", is_output_list=True),
io.Int.Output(display_name="item_index", is_output_list=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
],
is_input_list=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
)
@classmethod
def execute(cls, images=None, reference_image=None, reference_image_mode="per_item",
mask=None, mask_mode="per_item", text=None, text_mode="shared",
strength=None, strength_mode="shared", seed=None, seed_mode="shared",
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
image_list_master_callback="", image_list_execution_id="",
image_list_workflow_meta=None):
worker_item_token = unwrap_single(worker_item_token)
worker_context_id = unwrap_single(worker_context_id)
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
if worker_item_token:
item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
context_id = worker_context_id
if not isinstance(context_id, str) or not context_id:
raise RuntimeError("Worker image-list prompt is missing its paired context")
return io.NodeOutput(
[item["image"]], [item.get("reference_image")], [item.get("mask")],
[item.get("text")], [item.get("strength")], [item.get("seed")],
[item["index"]], [context_id],
)
items = normalize_image_list_inputs(
images, reference_image, mask, text, strength, seed,
reference_image_mode, mask_mode, text_mode, strength_mode, seed_mode,
)
prompt = getattr(cls.hidden, "prompt", None)
start_id = str(getattr(cls.hidden, "unique_id", ""))
if not prompt or not start_id:
raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
try:
worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
except (TypeError, ValueError) as exc:
raise RuntimeError("Invalid preflighted image-list worker list") from exc
if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
raise RuntimeError("Invalid preflighted image-list worker list")
from ..image_list.runtime import validate_image_list_regions
regions = validate_image_list_regions(prompt)
region = next((row for row in regions if row["start_id"] == start_id), None)
if region is None:
raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
from comfy_execution.utils import get_executing_context
executing = get_executing_context()
context_id = create_master_job(
start_id, region["end_id"], prompt, items, worker_ids,
str(unwrap_single(image_list_master_callback or "")),
image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
prompt_id=executing.prompt_id if executing is not None else None,
)
return io.NodeOutput(
[item["image"] for item in items],
[item.get("reference_image") for item in items],
[item.get("mask") for item in items],
[item.get("text") for item in items],
[item.get("strength") for item in items],
[item.get("seed") for item in items],
[item["index"] for item in items],
[context_id for _ in items],
)
async def _consume_worker_item(token):
# Item state is local to each worker process; source nodes consume it exactly once.
from ..image_list.runtime import consume_worker_item
return consume_worker_item(token)
class DistributedImageListEnd(io.ComfyNode):
"""Collect branch results in source order without evaluating the branch on master."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="DistributedImageListEnd",
display_name="Distributed Image List End",
category="image/distributed",
description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
inputs=[
io.Image.Input("images", lazy=True),
io.Mask.Input("masks", optional=True, lazy=True),
io.String.Input("texts", optional=True, lazy=True, force_input=True),
io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
],
outputs=[
io.Image.Output(display_name="image_list", is_output_list=True),
io.Mask.Output(display_name="masks", is_output_list=True),
io.String.Output(display_name="texts", is_output_list=True),
],
is_input_list=True,
is_output_node=True,
accept_all_inputs=True,
not_idempotent=True,
hidden=[io.Hidden.prompt, io.Hidden.unique_id],
)
@classmethod
def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", **kwargs):
worker_result_token = unwrap_single(worker_result_token)
if worker_result_token:
# Decide which optional sockets are connected from the actual worker prompt;
# absent optional inputs must not be demanded from Comfy's lazy executor.
prompt = getattr(cls.hidden, "prompt", {}) or {}
end_inputs = next((node.get("inputs", {}) for node in prompt.values()
if isinstance(node, dict)
and node.get("class_type") == "DistributedImageListEnd"
and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
required = []
for name, value in (("images", images), ("masks", masks), ("texts", texts)):
connected = name == "images" or name in end_inputs
missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
if connected and missing:
required.append(name)
return required
# Master only needs Start's context. Never request branch outputs here.
if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
return ["context"]
return []
@classmethod
def execute(cls, images=None, masks=None, texts=None, context=None,
worker_result_token="", worker_item_index=-1, master_callback_url=""):
worker_result_token = unwrap_single(worker_result_token)
master_callback_url = unwrap_single(master_callback_url)
if worker_result_token:
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Worker Image List End is missing its paired context")
result = normalize_worker_result(images, masks, texts)
payload = encode_result(result)
run_async_in_server_loop(
submit_worker_result(str(master_callback_url), str(worker_result_token),
context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
)
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
if isinstance(item, str) and item]
if not context_values or len(set(context_values)) != 1:
raise RuntimeError("Distributed Image List End requires the paired Start context wire")
context_id = context_values[0]
from ..image_list.runtime import ImageListCancelledError
try:
# The runtime owns the configurable deadline and bounded remote cleanup.
result_lists = run_async_in_server_loop(run_image_list_job(context_id))
except ImageListCancelledError:
import comfy.model_management
raise comfy.model_management.InterruptProcessingException()
images_out, masks_out, texts_out = result_lists
if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
return io.NodeOutput(images_out, masks_out, texts_out)
async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
if not callback_url.startswith(("http://", "https://")):
raise RuntimeError("Worker image-list callback URL is missing or invalid")
import aiohttp
from ..utils.network import get_client_session
session = await get_client_session()
payload = {"context_id": context_id, "token": item_token,
"item_index": item_index, "result": result}
async with session.post(
f"{callback_url.rstrip('/')}/distributed/image-list/result",
data=encode_request_payload(payload),
headers={"Content-Type": "application/json"},
timeout=aiohttp.ClientTimeout(total=60),
) as response:
if response.status >= 400:
body = (await response.text())[:500]
raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")
+2 -2
View File
@@ -269,7 +269,7 @@ class ImageBatchDivider:
class AudioBatchDivider:
"""Divides an audio waveform into multiple parts along the time/samples dimension."""
"""Divides an audio waveform into sequential segments along the time dimension."""
@classmethod
def INPUT_TYPES(s):
@@ -282,7 +282,7 @@ class AudioBatchDivider:
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the audio into"
"tooltip": "Number of sequential time segments to create"
}),
}
}
+183
View File
@@ -0,0 +1,183 @@
"""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 .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
from .image_list import DistributedImageListEnd, DistributedImageListStart
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)],
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):
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))
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,
DistributedImageListStart, DistributedImageListEnd]
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.4.4"
version = "1.5.0"
license = {file = "LICENSE"}
dependencies = []
+1
View File
@@ -0,0 +1 @@
"""Internal extension lifecycle support (not a node provider)."""
+38
View File
@@ -0,0 +1,38 @@
"""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()
from ..image_list import runtime as image_list_runtime
from ..image_list.lifecycle import install_lifecycle
install_lifecycle(server.PromptServer.instance, image_list_runtime)
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}')
+10 -11
View File
@@ -38,17 +38,10 @@ def _load_config_routes_module():
sys.modules[f"{package_name}.api"] = api_pkg
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = []
utils_pkg.__path__ = [str(module_path.parents[1] / "utils")]
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
import aiohttp # Load the real dependency; isolate the fake response at the tested module.
class _Routes:
def get(self, _path):
@@ -96,8 +89,14 @@ def _load_config_routes_module():
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)
class _EmptySecretStore:
def has(self, _worker_id):
return False
module.web = types.SimpleNamespace(
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
)
module._serverless_secret_store = lambda: _EmptySecretStore()
return module
+75
View File
@@ -205,6 +205,10 @@ def _load_job_routes_module():
module = importlib.util.module_from_spec(spec)
assert spec is not None and spec.loader is not None
spec.loader.exec_module(module)
# Keep response assertions independent of real aiohttp import ordering.
module.web = types.SimpleNamespace(
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
)
if created_aiohttp_stub:
sys.modules.pop("aiohttp", None)
@@ -302,6 +306,77 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
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,
+49 -15
View File
@@ -164,26 +164,11 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
# ---------------------------------------------------------------------------
class FindMediaReferencesTests(unittest.TestCase):
def test_finds_image_input(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
refs = ms._find_media_references(prompt)
self.assertIn("photo.png", refs)
def test_finds_video_input(self):
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
refs = ms._find_media_references(prompt)
self.assertIn("clip.mp4", refs)
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_finds_audio_input(self):
prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
refs = ms._find_media_references(prompt)
self.assertIn("track.wav", refs)
def test_strips_annotation_suffix(self):
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
refs = ms._find_media_references(prompt)
@@ -258,5 +243,54 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase):
async def test_output_temp_and_input_uploads_do_not_collide(self):
from unittest.mock import AsyncMock, patch
class Response:
status = 200
def __init__(self, data):
self.data = data
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.data
def raise_for_status(self):
pass
class Form:
def __init__(self):
self.fields = {}
def add_field(self, key, value, **kwargs):
self.fields[key] = value
if key == 'image':
self.fields['name'] = kwargs['filename']
checked = []
def post(url, *, json=None, data=None, **kwargs):
if url.endswith('/check_file'):
checked.append(json['filename'])
return Response({'exists': False})
self.assertEqual(data.fields['type'], 'input')
self.assertEqual(data.fields['image'], b'real file bytes')
return Response({'name': data.fields['name'], 'subfolder': data.fields['subfolder']})
refs = ['shared/ref.png', 'shared/ref.png [output]', 'shared/ref.png [temp]']
prompt = {str(i): {'inputs': {'image': ref}} for i, ref in enumerate(refs)}
self.assertEqual(ms._find_media_references(prompt), refs)
with patch.object(ms, 'get_client_session', AsyncMock(return_value=types.SimpleNamespace(post=post))), \
patch.object(ms, 'aiohttp', types.SimpleNamespace(FormData=Form, ClientTimeout=lambda **kw: kw)):
paths = [await ms._upload_media_to_worker({}, ref, b'real file bytes', 'fixture-hash', 'image/png')
for ref in refs]
expected = ['shared/ref.png', 'distributed-media/output/fixture-hash/ref.png',
'distributed-media/temp/fixture-hash/ref.png']
self.assertEqual(checked, expected)
self.assertEqual(paths, [(True, path) for path in expected])
ms._rewrite_prompt_media_inputs(prompt, dict(zip(refs, expected)))
self.assertEqual([node['inputs']['image'] for node in prompt.values()], expected)
if __name__ == "__main__":
unittest.main()
+36 -1
View File
@@ -106,7 +106,7 @@ def _load_worker_routes_module():
sys.modules[f"{package_name}.utils"] = utils_pkg
workers_pkg = types.ModuleType(f"{package_name}.workers")
workers_pkg.__path__ = []
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
sys.modules[f"{package_name}.workers"] = workers_pkg
@@ -213,6 +213,7 @@ def _load_worker_routes_module():
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")
@@ -256,6 +257,8 @@ def _load_worker_routes_module():
module = importlib.util.module_from_spec(spec)
assert spec is not None and spec.loader is not None
spec.loader.exec_module(module)
# Keep this module's response double independent of aiohttp import order.
module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
if created_aiohttp_stub:
sys.modules.pop("aiohttp", None)
@@ -269,6 +272,38 @@ 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}]}
+691
View File
@@ -0,0 +1,691 @@
{
"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)"
}
}
+397
View File
@@ -0,0 +1,397 @@
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
No managed ComfyUI service, user workflow, model, or remote worker is touched.
"""
import asyncio
from collections import Counter
import importlib
import json
import os
from pathlib import Path
import sys
import tempfile
import time
import uuid
from unittest.mock import patch
ROOT = Path(__file__).resolve().parents[2]
COMFY_ROOT = Path(sys.argv[1]).resolve()
WORKER = '--worker' in sys.argv
DELAY = float(sys.argv[3]) if WORKER else 0
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
import comfy.model_management
import execution
import nodes
import server
import torch
from aiohttp import web
import aiohttp
torch.set_num_threads(1)
SOURCE_CALLS, BRANCH_CALLS, CAPTURED = [], [], []
COUNT = 7
def source_rows():
return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
for i in range(COUNT)]
class TestImageSource:
@classmethod
def INPUT_TYPES(cls):
return {'required': {}}
RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
OUTPUT_IS_LIST = (True,) * 6
FUNCTION = 'run'
def run(self):
assert not WORKER, 'upstream image source leaked into worker prompt'
SOURCE_CALLS.append(True)
rows = source_rows()
return tuple([row[key] for row in rows] for key in
('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
class TestImageOperation:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
'text': ('STRING', {'forceInput': True}),
'strength': ('FLOAT', {'forceInput': True}),
'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True})}}
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
FUNCTION = 'run'
def run(self, image, reference_image, mask, text, strength, seed, index):
assert WORKER, 'per-image branch executed on master'
BRANCH_CALLS.append(index)
# Sleep in short chunks so the real Comfy interrupt signal can be tested.
until = time.monotonic() + DELAY
while time.monotonic() < until:
comfy.model_management.throw_exception_if_processing_interrupted()
time.sleep(0.01)
assert seed == 100 + index
assert text == f'item-{index}'
assert strength == (index + 1) / 16
assert torch.all(reference_image == strength)
return image + strength, mask * 0.5, f'{text}:seed-{seed}'
class TestCapture:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
INPUT_IS_LIST = True
OUTPUT_NODE = True
RETURN_TYPES = ()
FUNCTION = 'run'
def run(self, images, masks, texts):
CAPTURED.append((images, masks, texts))
return ()
async def setup():
assets = default_asset_manager()
ps = server.PromptServer(asyncio.get_running_loop(), assets)
assert await nodes.load_custom_node(str(ROOT))
module = sys.modules[str(ROOT).replace('.', '_x_')]
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
TestImageOperation=TestImageOperation, TestCapture=TestCapture)
# Match the CORS-enabled direct-worker launch contract for canvas probes.
app = web.Application(middlewares=[server.create_cors_middleware('*')])
# Register the actual feature routes, not test replacements.
app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
return ps, assets, module, app
async def listen(app):
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(runner, '127.0.0.1', 0)
await site.start()
return runner, site._server.sockets[0].getsockname()[1]
async def worker_main():
ps, assets, module, app = await setup()
pending, running, history = [], {}, {}
wake = asyncio.Event()
interrupted_ids = []
async def prompt(request):
data = await request.json()
prompt_id = str(uuid.uuid4())
graph = data['prompt']
valid = await execution.validate_prompt(prompt_id, graph, None)
if not valid[0]:
return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
pending.append((prompt_id, graph, valid[2]))
wake.set()
return web.json_response({'prompt_id': prompt_id})
async def queue(request):
if request.method == 'POST':
deleted = (await request.json()).get('delete', [])
pending[:] = [entry for entry in pending if entry[0] not in deleted]
return web.json_response({})
return web.json_response({'queue_running': [[0, key] for key in running],
'queue_pending': [[0, entry[0]] for entry in pending]})
async def get_history(request):
pid = request.match_info['prompt_id']
return web.json_response({pid: history[pid]} if pid in history else {})
async def interrupt(request):
pid = (await request.json())['prompt_id']
interrupted_ids.append(pid)
if pid in running:
comfy.model_management.interrupt_current_processing(True)
return web.json_response({})
async def stats(request):
return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
'history': history})
async def consume():
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
while True:
await wake.wait()
while pending:
pid, graph, outputs = pending.pop(0)
running[pid] = True
await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
history[pid] = {'status': {'completed': True,
'status_str': 'success' if executor.success else 'error',
'messages': str(executor.status_messages)}}
running.pop(pid, None)
wake.clear()
async def probe(_request):
return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
app.router.add_get('/prompt', probe)
app.router.add_post('/prompt', prompt)
app.router.add_get('/queue', queue)
app.router.add_post('/queue', queue)
app.router.add_get('/history/{prompt_id}', get_history)
app.router.add_post('/interrupt', interrupt)
app.router.add_get('/test-stats', stats)
runner, port = await listen(app)
consumer = asyncio.create_task(consume())
print(f'READY:{port}', flush=True)
try:
await asyncio.Event().wait()
finally:
consumer.cancel()
await asyncio.gather(consumer, return_exceptions=True)
await runner.cleanup()
def graph(callback, worker_ids):
return {
'1': {'class_type': 'TestImageSource', 'inputs': {}},
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'reference_image': ['1', 1], 'mask': ['1', 2], 'text': ['1', 3],
'strength': ['1', 4], 'seed': ['1', 5],
'reference_image_mode': 'per_item', 'mask_mode': 'per_item', 'text_mode': 'per_item',
'strength_mode': 'per_item', 'seed_mode': 'per_item',
'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
'3': {'class_type': 'TestImageOperation', 'inputs': {
key: ['2', index] for key, index in [('image', 0), ('reference_image', 1), ('mask', 2),
('text', 3), ('strength', 4), ('seed', 5), ('index', 6)]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 7]}},
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
}
async def coordinator():
ps, assets, module, app = await setup()
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
config_mod = importlib.import_module(module.__name__ + '.utils.config')
orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
runner, port = await listen(app)
ps.port = port # Match the listener used by local-worker callback resolution.
processes, stderr_files = [], []
# Old worker capability refusal is read-only and precedes input evaluation.
legacy_app = web.Application()
legacy_requests = []
async def legacy(request):
legacy_requests.append(request.method)
return web.json_response({'error': 'not found'}, status=404)
legacy_app.router.add_route('*', '/{path:.*}', legacy)
legacy_runner, legacy_port = await listen(legacy_app)
dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
try:
try:
await dispatch_module.preflight_image_list_workers([
{'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
raise AssertionError('legacy worker accepted')
except RuntimeError as exc:
assert 'upgrade' in str(exc), str(exc)
assert legacy_requests == ['GET'] and not SOURCE_CALLS
finally:
await legacy_runner.cleanup()
print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
scratch.mkdir(parents=True, exist_ok=True)
try:
with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
workers = []
for name, delay in [('fast', 0.02), ('slow', 0.5)]:
err = open(Path(tmp) / f'{name}.log', 'w+')
stderr_files.append(err)
proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
env={**os.environ, 'TMPDIR': tmp})
processes.append(proc)
async def ready():
while True:
line = await proc.stdout.readline()
if not line:
err.seek(0)
raise AssertionError(f'Worker exited: {err.read()}')
if line.startswith(b'READY:'):
return int(line.split(b':')[1])
worker_port = await asyncio.wait_for(ready(), 60)
workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
'port': worker_port, 'enabled': True})
config = {'workers': workers, 'settings': {'websocket_orchestration': False},
'master': {'host': f'http://127.0.0.1:{port}'}}
prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
valid = await execution.validate_prompt('network-acceptance', prompt, None)
assert valid[0], valid
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
# Exercise production queue preparation; only master enqueue is redirected
# to this isolated executor, not the managed server's live queue.
prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
prompt['2']['inputs'].pop('image_list_master_callback')
async def enqueue(prepared, workflow, client_id, **kwargs):
checked = await execution.validate_prompt('network-acceptance', prepared, None)
assert checked[0], checked
await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
prompt, None, 'test', [w['id'] for w in workers]), 45)
assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
if not executor.success:
for err in stderr_files:
err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
assert executor.success, executor.status_messages
assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
assert len(CAPTURED) == 1
images, masks, texts = CAPTURED[0]
assert len(images) == len(masks) == len(texts) == COUNT
for i, row in enumerate(source_rows()):
assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
assert torch.equal(masks[i], row['mask'] * 0.5), i
assert texts[i] == f'item-{i}:seed-{100 + i}', texts
async with aiohttp.ClientSession() as session:
stats = []
for worker in workers:
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
stats.append(await response.json())
counts = [len(row['calls']) for row in stats]
assert counts[0] > counts[1] > 0, counts
assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
assert not runtime._MASTER_JOBS
# Three sequential regions on one worker, including optional outputs
# left unconnected. Region N+1 must not race the prior worker End.
chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
previous = '1'
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
'images': [previous, 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
'text_mode': 'shared', 'strength_mode': 'shared', 'seed_mode': 'shared'}}
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': [start, 0], 'context': [start, 7]}}
previous = end
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(orchestration, 'load_config', return_value=config), \
patch.object(orchestration, 'queue_prompt_payload', enqueue), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
chained, None, 'test', ['slow']), 45)
assert executor.success, executor.status_messages
assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
assert CAPTURED[1][1:] == ([None, None], [None, None])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
# Real Comfy user-interrupt flag while a real worker executor is active.
slow = workers[1]
async with aiohttp.ClientSession() as session:
async def slow_stats():
async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
return await response.json()
async def until(predicate):
for _ in range(150):
current = await slow_stats()
if predicate(current):
return current
await asyncio.sleep(0.02)
raise AssertionError('worker state transition timed out')
before = await until(lambda state: not state['running'] and not state['pending'])
context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
task = asyncio.create_task(runtime.run_image_list_job(context))
active = await until(lambda state: bool(state['running']))
owned_id = active['running'][0]
comfy.model_management.interrupt_current_processing(True)
try:
await asyncio.wait_for(task, 5)
raise AssertionError('interrupt was ignored')
except runtime.ImageListCancelledError:
pass
finally:
comfy.model_management.interrupt_current_processing(False)
after = await until(lambda state: not state['running'] and not state['pending'])
assert after['interrupted'] == [owned_id], after
assert len(after['calls']) <= len(before['calls']) + 1
assert set(before['history']) <= set(after['history'])
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
# A worker exception before End must fail, not hang or return a partial list.
bad_rows = source_rows()
bad_rows[0]['seed'] = -1
context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
f'http://127.0.0.1:{port}')
with patch.object(config_mod, 'load_config', return_value=config), \
patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
try:
await asyncio.wait_for(runtime.run_image_list_job(context), 5)
raise AssertionError('worker exception was ignored')
except RuntimeError as exc:
assert 'failed image-list item' in str(exc), str(exc)
assert not runtime._MASTER_JOBS
print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
finally:
for proc in processes:
if proc.returncode is None:
proc.terminate()
await asyncio.gather(*(proc.wait() for proc in processes))
for err in stderr_files:
err.close()
await runner.cleanup()
network = importlib.import_module(module.__name__ + '.utils.network')
session = await network.get_client_session()
await session.close()
if __name__ == '__main__':
asyncio.run(worker_main() if WORKER else coordinator())
+376
View File
@@ -0,0 +1,376 @@
"""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())
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
# 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 check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
"""Execute real native V3 list/lazy flow; master must skip the processing branch."""
branch_calls = []
class ImageListProbe:
@classmethod
def INPUT_TYPES(cls):
return {'required': {'image': ('IMAGE',)}}
RETURN_TYPES = ('IMAGE',)
FUNCTION = 'run'
CATEGORY = 'tests'
def run(self, image):
branch_calls.append(image)
return (image,)
nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
runtime = sys.modules[module.__name__ + '.image_list.runtime']
original_contexts = set(runtime._MASTER_JOBS)
image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
async def fake_worker_results(_context_id):
state = runtime.get_master_job(_context_id)
assert state.prompt_id == 'image-list-lazy', state.prompt_id
assert state.workflow_meta == workflow_meta, state.workflow_meta
assert len(prompt_server.app.cleanup_ctx) >= 1, 'lifecycle cleanup was not registered'
return (
[torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
[None, None, None], [None, None, None],
)
graph = {
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
'images': ['1', 0], 'reference_image_mode': 'per_item', 'mask_mode': 'per_item',
'text_mode': 'shared', 'strength': 1.0, 'strength_mode': 'shared',
'seed': 17, 'seed_mode': 'shared', 'image_list_enabled_worker_ids': '["test-worker"]',
'image_list_master_callback': 'http://127.0.0.1:1',
'image_list_workflow_meta': workflow_meta,
}},
'3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
'images': ['3', 0], 'context': ['2', 7],
}},
'5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
}
valid = await execution.validate_prompt('image-list-lazy', graph, None)
assert valid[0], valid
validator = runtime.validate_image_list_regions
assert 'texts' not in graph['4']['inputs']
assert len(validator(graph)) == 1
connected_text_graph = json.loads(json.dumps(graph))
connected_text_graph['6'] = {
'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 3]},
}
connected_text_graph['4']['inputs']['texts'] = ['6', 0]
assert len(validator(connected_text_graph)) == 1
for literal_text in ('', 'per-item text'):
literal_text_graph = json.loads(json.dumps(graph))
literal_text_graph['4']['inputs']['texts'] = literal_text
try:
validator(literal_text_graph)
except ValueError:
pass
else:
raise AssertionError(f'End accepted unconnected literal text: {literal_text!r}')
print('IMAGE_LIST_END_TEXT_INPUT_OK native forceInput; unconnected omitted; linked accepted; literals rejected')
import folder_paths
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
try:
# Resolve genuine Comfy output/temp annotations against the pinned core,
# not only the controlled folder_paths seam in unit tests.
media = sys.modules[module.__name__ + '.api.orchestration.media_sync']
with tempfile.TemporaryDirectory(prefix='media-roots-', dir=scratch) as media_root:
for kind in ('input', 'output', 'temp'):
source = Path(media_root) / kind / 'shared/ref.png'
source.parent.mkdir(parents=True)
source.write_bytes(kind.encode())
with patch.object(folder_paths, 'input_directory', str(Path(media_root) / 'input')), \
patch.object(folder_paths, 'output_directory', str(Path(media_root) / 'output')), \
patch.object(folder_paths, 'temp_directory', str(Path(media_root) / 'temp')):
for kind in ('output', 'temp'):
refs = media._find_media_references({'1': {'inputs': {'image': f'shared/ref.png [{kind}]'}}})
assert len(refs) == 1
assert media._load_media_file_sync(refs[0])[0] == kind.encode()
with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
with patch.object(folder_paths, 'temp_directory', temp), \
patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
executor = execution.PromptExecutor(
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
assert executor.success, executor.status_messages
assert branch_calls == [], 'master evaluated the lazy per-image branch'
records = executor.history_result['outputs']['5']['images']
assert len(records) == 3, records
for record in records:
preview = Path(temp) / record.get('subfolder', '') / record['filename']
with Image.open(preview) as image:
assert image.size == (8, 8) and image.mode == 'RGB'
finally:
nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
for context_id in set(runtime._MASTER_JOBS) - original_contexts:
runtime.remove_master_job(context_id)
print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews 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) + 2 == 10
assert set(mapping) == set(BASELINE) | {'DistributedImageListStart', 'DistributedImageListEnd'}
for node_id in BASELINE:
cls = mapping[node_id]
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 = new['input'].get(group, {})
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'))
start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
assert start_schema.node_id == 'DistributedImageListStart'
assert start_schema.display_name == 'Distributed Image List Start'
assert end_schema.node_id == 'DistributedImageListEnd'
assert end_schema.display_name == 'Distributed Image List End'
texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
texts_schema = texts_input.as_dict()
assert texts_schema.get('forceInput') is True, texts_schema
assert 'default' not in texts_schema, texts_schema
assert [out.io_type for out in start_schema.outputs] == [
'IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT']
assert all(out.is_output_list for out in start_schema.outputs)
assert start_schema.is_input_list and start_schema.accept_all_inputs
assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
assert all(out.is_output_list for out in end_schema.outputs)
assert end_schema.is_input_list and end_schema.is_output_node
print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
print('SCHEMA_PARITY_OK', len(mapping))
check_saved_workflows(mapping)
await check_execution(mapping, module, prompt_server, asset_manager)
await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
print('V3_ACCEPTANCE_OK')
asyncio.run(main())
+159
View File
@@ -124,6 +124,49 @@ def test_collector_opts_into_comfyui_list_inputs():
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)
@@ -203,3 +246,119 @@ def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final
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"])
+1 -1
View File
@@ -18,7 +18,7 @@ def _load_config_module():
del sys.modules[mod_name]
root_pkg = types.ModuleType(package_name)
root_pkg.__path__ = []
root_pkg.__path__ = [str(module_path.parent)]
sys.modules[package_name] = root_pkg
logging_module = types.ModuleType(f"{package_name}.logging")
+266 -1
View File
@@ -55,7 +55,13 @@ def _load_dispatch_module():
def __init__(self, total=None):
self.total = total
class _ClientConnectorError(Exception):
class _ClientError(Exception):
pass
class _ClientConnectorError(_ClientError):
pass
class _ClientConnectionError(_ClientError):
pass
class _WSMsgType:
@@ -76,7 +82,9 @@ def _load_dispatch_module():
self.closed = True
aiohttp_module.ClientTimeout = _ClientTimeout
setattr(aiohttp_module, "ClientError", _ClientError)
aiohttp_module.ClientConnectorError = _ClientConnectorError
setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
aiohttp_module.WSMsgType = _WSMsgType
aiohttp_module.TCPConnector = _TCPConnector
aiohttp_module.ClientSession = _ClientSession
@@ -100,6 +108,75 @@ def _load_dispatch_module():
dispatch = _load_dispatch_module()
def _load_prompt_transform_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
package_name = "prompt_transform_worker_selection_testpkg"
for suffix in ("", ".api", ".api.orchestration", ".utils"):
package = types.ModuleType(package_name + suffix)
package.__path__ = []
sys.modules[package_name + suffix] = package
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
sys.modules[f"{package_name}.utils.logging"] = logging_module
spec = importlib.util.spec_from_file_location(
f"{package_name}.api.orchestration.prompt_transform",
module_path,
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
prompt_transform = _load_prompt_transform_module()
class _FakeResponse:
def __init__(self, status=200, payload=None, error=None):
self.status = status
self.payload = payload
self.error = error
async def __aenter__(self):
if self.error is not None:
raise self.error
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.payload
class _FakeSession:
def __init__(self, responses):
self.responses = responses
def get(self, url, timeout=None):
del timeout
worker_id = url.removeprefix("http://").split("/", 1)[0]
endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
response = self.responses[(worker_id, endpoint)]
if isinstance(response, BaseException):
return _FakeResponse(error=response)
status, payload = response
return _FakeResponse(status=status, payload=payload)
def _valid_capability():
return {
"supported": True,
"protocol_version": 1,
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
}
def _valid_queue(running=None, pending=None):
return {"queue_running": running or [], "queue_pending": pending or []}
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
async def test_select_active_workers_filters_offline(self):
workers = [
@@ -249,6 +326,194 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
self.assertIsNone(selected)
async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "offline"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "offline")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "offline")
})
responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
workers = [{"id": "healthy"}, {"id": "slow"}]
responses: dict = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("healthy", "slow")
}
responses.update({
(worker_id, "queue"): (200, _valid_queue())
for worker_id in ("healthy", "slow")
})
responses[("slow", "queue")] = asyncio.TimeoutError()
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
), patch.object(dispatch, "log") as log_mock:
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"healthy": True})
self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
async def test_image_list_preflight_excludes_busy_workers(self):
workers = [{"id": "idle"}, {"id": "busy"}]
responses = {
(worker_id, "capabilities"): (200, _valid_capability())
for worker_id in ("idle", "busy")
}
responses.update({
("idle", "queue"): (200, _valid_queue()),
("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
})
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
result = await dispatch.preflight_image_list_workers(workers)
self.assertEqual(result, {"idle": True})
async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
workers = [{"id": "offline-a"}, {"id": "offline-b"}]
responses = {
(worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
for worker_id in ("offline-a", "offline-b")
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
self.assertEqual(
await dispatch.preflight_image_list_workers(workers, require_idle=False),
{},
)
with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
await dispatch.preflight_image_list_workers(workers, require_idle=True)
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
malformed = {"supported": True, "protocol_version": 1, "node_ids": None}
session = _FakeSession({("worker", "capabilities"): (200, malformed)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
incompatible = _valid_capability()
incompatible["protocol_version"] = 2
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
incompatible = _valid_capability()
incompatible["node_ids"] = ["DistributedImageListStart"]
session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "incompatible"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
responses = {
("worker", "capabilities"): (200, _valid_capability()),
("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
}
session = _FakeSession(responses)
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
async def fake_session():
return session
with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
):
with self.assertRaisesRegex(ValueError, "programming error"):
await dispatch.preflight_image_list_workers([{"id": "worker"}])
class ImageListDirectWorkerResolutionTests(unittest.TestCase):
def setUp(self):
self.config = {"workers": [
{"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
{"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
{"id": "serverless", "enabled": True, "type": "serverless"},
{"id": "other", "enabled": True, "host": "other-host", "port": 8190},
]}
def test_implicit_selection_uses_only_enabled_direct_workers(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
def test_explicit_selection_includes_exact_disabled_direct_worker(self):
workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
self.assertEqual([worker["id"] for worker in workers], ["disabled"])
def test_explicit_empty_selection_returns_no_workers(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
def test_unknown_explicit_id_raises_instead_of_falling_back(self):
with self.assertRaisesRegex(RuntimeError, "missing"):
prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
def test_explicit_serverless_selection_is_excluded(self):
self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
if __name__ == "__main__":
unittest.main()
+291
View File
@@ -0,0 +1,291 @@
"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
import asyncio
import unittest
import torch
from image_list.runtime import (
END_CLASS,
IMAGE_OUT,
CONTEXT_OUT,
START_CLASS,
ImageListCancelledError,
build_worker_prompt,
create_master_job,
decode_item,
decode_result,
dynamic_schedule,
encode_item,
encode_result,
normalize_image_list_inputs,
normalize_worker_result,
remove_master_job,
set_master_result,
validate_image_list_regions,
)
def start_node(image_link):
return {"class_type": START_CLASS, "inputs": {
"images": image_link,
"reference_image_mode": "per_item", "mask_mode": "per_item",
"text_mode": "shared", "strength_mode": "shared", "seed_mode": "shared",
}}
def graph_one_region():
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
"3": start_node(["2", 0]),
"4": {"class_type": "GenericFrameOperation", "inputs": {
"image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", 5],
}},
"5": {"class_type": END_CLASS, "inputs": {
"images": ["4", 0], "context": ["3", CONTEXT_OUT],
}},
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
}
def graph_two_sequential_regions():
return {
"1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
"2": start_node(["1", 0]),
"3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
"4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
"5": start_node(["4", 0]),
"6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
"7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
}
class ImageListNormalizationTests(unittest.TestCase):
def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
images = [
torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
]
refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
rows = normalize_image_list_inputs(
images, refs, masks, ["first", "second"], [0.2], [37],
"per_item", "per_item", "per_item", "shared", "shared",
)
self.assertEqual(len(rows), 2)
self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
self.assertTrue(torch.equal(rows[0]["reference_image"], refs[:1]))
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
self.assertEqual([row["index"] for row in rows], [0, 1])
self.assertEqual([row["seed"] for row in rows], [37, 37])
self.assertEqual([row["strength"] for row in rows], [0.2, 0.2])
def test_explicit_per_item_and_shared_scalar_contracts(self):
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
rows = normalize_image_list_inputs(image, text="same", strength=0.4, seed=9,
text_mode="shared", strength_mode="shared", seed_mode="shared")
self.assertEqual([row["text"] for row in rows], ["same"] * 3)
self.assertEqual([row["strength"] for row in rows], [0.4] * 3)
self.assertEqual([row["seed"] for row in rows], [9] * 3)
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
normalize_image_list_inputs(image, text=["one", "two"], text_mode="per_item")
with self.assertRaisesRegex(ValueError, "shared inputs require exactly one"):
normalize_image_list_inputs(image, seed=[1, 2], seed_mode="shared")
with self.assertRaisesRegex(ValueError, "expected exactly 3"):
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)), mask_mode="per_item")
def test_no_implicit_tail_repetition_and_no_empty_primary_batch(self):
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
with self.assertRaisesRegex(ValueError, "expected exactly 2"):
normalize_image_list_inputs(images, reference_image=torch.zeros((1, 2, 2, 3)),
reference_image_mode="per_item")
with self.assertRaisesRegex(ValueError, "at least one image"):
normalize_image_list_inputs(None)
def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
"reference_image": torch.ones((1, 1, 3, 4), dtype=torch.float16),
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
"text": "safe text", "strength": 0.625, "seed": 42, "index": 3}
restored = decode_item(encode_item(item))
self.assertTrue(torch.equal(restored["image"], item["image"]))
self.assertEqual(restored["image"].dtype, torch.float64)
self.assertTrue(torch.equal(restored["reference_image"], item["reference_image"]))
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
self.assertEqual(restored["text"], "safe text")
self.assertEqual(restored["seed"], 42)
with self.assertRaisesRegex(TypeError, "floating-point tensor"):
encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
from image_list.runtime import encode_value
encode_value("/untrusted/path", "PATH")
def test_result_transport_preserves_ragged_result_shapes(self):
result = normalize_worker_result(
[torch.ones((1, 3, 5, 4), dtype=torch.float32)],
[torch.ones((1, 3, 5), dtype=torch.float16)],
["done"],
)
encoded = encode_result(result)
decoded = decode_result(encoded)
self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
self.assertEqual(decoded["image"].dtype, torch.float32)
self.assertEqual(decoded["mask"].dtype, torch.float16)
self.assertEqual(decoded["text"], "done")
with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
normalize_worker_result(torch.zeros((2, 3, 4, 3)))
class ImageListGraphBoundaryTests(unittest.TestCase):
def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
prompt = graph_one_region()
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
"ctx", 1, "http://master.invalid")
self.assertEqual(set(worker), {"1", "3", "4", "5"})
self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
self.assertNotIn("images", worker["3"]["inputs"])
self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
self.assertNotIn("6", worker)
def test_rejects_upstream_processing_and_leaf_image_bypasses(self):
# Both a leaf source and a generated batch must cross Start exactly once.
for generated in (False, True):
with self.subTest(generated=generated):
prompt = graph_one_region()
if generated:
prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
prompt["4"]["inputs"]["reference"] = ["2", 0]
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
validate_image_list_regions(prompt)
def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
import sys
import types
from unittest.mock import patch
prompt = graph_one_region()
prompt["7"] = {"class_type": "SeedConstant", "inputs": {"value": 9}}
prompt["2"]["inputs"].update({"model": ["1", 0], "seed": ["7", 0]})
prompt["4"]["inputs"]["shared_seed"] = ["7", 0]
registry = {
"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL", "CLIP", "VAE")}),
"SeedConstant": type("Constant", (), {"RETURN_TYPES": ("INT",)}),
}
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
region = validate_image_list_regions(prompt)[0]
worker = build_worker_prompt(prompt, region, "token", "ctx", 0, "http://master")
self.assertNotIn("2", worker)
self.assertTrue({"1", "7"}.issubset(worker))
def test_unknown_or_list_producing_shared_dependency_is_not_exempt(self):
import sys
import types
from unittest.mock import patch
for node_class in (None, type("ListLoader", (), {
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,)})):
prompt = graph_one_region()
prompt["2"]["inputs"]["model"] = ["1", 0]
registry = {} if node_class is None else {"CheckpointLoaderSimple": node_class}
with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
with self.assertRaisesRegex(ValueError, "bypasses Start.*1"):
validate_image_list_regions(prompt)
def test_two_sequential_regions_are_valid_and_isolated(self):
regions = validate_image_list_regions(graph_two_sequential_regions())
self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
p = graph_two_sequential_regions()
second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
"http://master.invalid")
self.assertEqual(set(second), {"5", "6", "7"})
def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
prompt = graph_one_region()
prompt["5"]["inputs"]["context"] = ["2", 0]
with self.assertRaisesRegex(ValueError, "typed context output"):
validate_image_list_regions(prompt)
nested = graph_one_region()
nested["7"] = start_node(["4", 0])
nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
nested["5"]["inputs"]["images"] = ["8", 0]
nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
validate_image_list_regions(nested)
outside = graph_one_region()
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
with self.assertRaisesRegex(ValueError, "cross-boundary"):
validate_image_list_regions(outside)
def test_rejects_duplicate_or_missing_results(self):
prompt = graph_one_region()
prompt["5"]["inputs"].pop("images")
with self.assertRaisesRegex(ValueError, "images must be connected"):
validate_image_list_regions(prompt)
class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
async def test_fast_workers_get_more_items_and_order_is_stable(self):
assignments = {"fast": [], "slow": []}
async def dispatch(worker, index):
assignments[worker].append(index)
delay = 0.002 if worker == "fast" else 0.025
task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
return {"index": index, "task": task}
async def wait_next(inflight):
done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
task = next(iter(done))
worker = next(w for w, handle in inflight.items() if handle["task"] is task)
index, result = task.result()
return worker, index, result
async def cleanup(_worker, handle):
handle["task"].cancel()
ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
async def dispatch(worker, index):
return {"index": index, "prompt_id": f"{worker}-{index}"}
cancelled = []
async def cleanup(worker, handle):
cancelled.append((worker, handle["prompt_id"]))
async def timeout(_inflight):
raise TimeoutError("worker timed out")
with self.assertRaisesRegex(TimeoutError, "timed out"):
await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
cancelled.clear()
async def wrong_index(inflight):
worker, handle = next(iter(inflight.items()))
return worker, handle["index"] + 99, "bad"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
self.assertEqual(len(cancelled), 1)
cancelled.clear()
with self.assertRaises(ImageListCancelledError):
await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
interrupted=lambda: True)
self.assertEqual(len(cancelled), 0)
async def test_no_workers_and_duplicate_results_fail_closed(self):
async def unused(*_args):
self.fail("must not dispatch without workers")
with self.assertRaisesRegex(RuntimeError, "No free"):
await dynamic_schedule(2, [], unused, unused, unused)
async def dispatch(worker, index):
return {"index": index}
async def duplicate(inflight):
worker, _handle = next(iter(inflight.items()))
return worker, 0, "same"
with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
if __name__ == "__main__":
unittest.main()
+182
View File
@@ -0,0 +1,182 @@
"""Regression coverage at the scheduler and real HTTP transport boundaries."""
import asyncio
import importlib.util
import json
from pathlib import Path
import sys
import types
import unittest
from unittest.mock import patch
import torch
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from image_list import runtime
from test_image_list import graph_two_sequential_regions, start_node
def load_routes():
root = Path(__file__).resolve().parents[1]
name = '_image_list_transport_tests'
routes = web.RouteTableDef()
modules = {
name: types.ModuleType(name),
name + '.api': types.ModuleType(name + '.api'),
name + '.image_list': types.ModuleType(name + '.image_list'),
name + '.image_list.runtime': runtime,
'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
instance=types.SimpleNamespace(routes=routes))),
}
spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', root / 'api/image_list_routes.py')
module = importlib.util.module_from_spec(spec)
with patch.dict(sys.modules, modules):
spec.loader.exec_module(module)
return module, routes
class TransportTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.module, routes = load_routes()
app = web.Application()
app.add_routes(routes)
self.client = TestClient(TestServer(app))
await self.client.start_server()
async def asyncTearDown(self):
await self.client.close()
runtime._WORKER_ITEMS.clear()
async def test_fragmented_chunked_item_upload_is_consumable(self):
token = runtime.create_worker_task_token()
item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
async def chunks():
yield data[:8]
await asyncio.sleep(0.02)
yield data[8:]
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 200, await response.text())
result = runtime.consume_worker_item(token)
self.assertTrue(torch.equal(result['image'], item['image']))
self.assertEqual(result['index'], 2)
async def test_complete_item_upload_and_replay_rejection(self):
token = runtime.create_worker_task_token()
data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
{'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 200, await response.text())
response = await self.client.post('/distributed/image-list/item', json=data)
self.assertEqual(response.status, 400)
runtime.consume_worker_item(token)
with self.assertRaisesRegex(RuntimeError, 'already consumed'):
runtime.consume_worker_item(token)
async def test_chunked_body_limit_is_aggregate(self):
async def chunks():
yield b'{"context_id":"'
await asyncio.sleep(0.01)
yield b'x' * 60
yield b'"}'
with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
response = await self.client.post('/distributed/image-list/item', data=chunks())
self.assertEqual(response.status, 400)
self.assertIn('exceeds', await response.text())
async def test_invalid_result_shapes_are_rejected(self):
for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
'mask': torch.zeros((1, 2, 3))})
result[kind]['shape'] = list(shape)
with self.assertRaises(ValueError):
runtime.decode_result(result)
class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
entered = asyncio.Event()
cancelled = asyncio.Event()
wait_cancelled = asyncio.Event()
dispatches, cleanups = [], []
async def dispatch(worker, index):
dispatches.append(index)
return {'index': index}
async def wait_next(_inflight):
entered.set()
try:
await asyncio.Event().wait()
finally:
wait_cancelled.set()
async def cleanup(worker, handle):
cleanups.append((worker, handle['index']))
task = asyncio.create_task(runtime.dynamic_schedule(
3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
await entered.wait()
cancelled.set()
with self.assertRaises(runtime.ImageListCancelledError):
await asyncio.wait_for(task, timeout=1)
self.assertTrue(wait_cancelled.is_set())
self.assertEqual(dispatches, [0])
self.assertEqual(cleanups, [('one', 0)])
async def test_already_interrupted_does_not_submit_any_work(self):
async def unused(*args):
self.fail('pre-cancelled job must not dispatch or wait')
with self.assertRaises(runtime.ImageListCancelledError):
await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
async def test_task_cancel_during_submit_waits_for_owned_handle(self):
entered, release = asyncio.Event(), asyncio.Event()
cleanups = []
async def dispatch(worker, index):
entered.set()
await release.wait()
return {'index': index, 'prompt_id': 'owned'}
async def unused(*args):
self.fail('cancelled dispatch must not enter result wait')
async def cleanup(worker, handle):
cleanups.append(handle['prompt_id'])
task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
await entered.wait()
task.cancel()
release.set()
with self.assertRaises(asyncio.CancelledError):
await task
self.assertEqual(cleanups, ['owned'])
def test_output_nodes_inside_region_are_rejected(self):
prompt = graph_two_sequential_regions()
output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
with patch.dict(sys.modules, {'nodes': registry}):
with self.assertRaisesRegex(ValueError, 'must be placed after'):
runtime.validate_image_list_regions(prompt)
def test_shared_scalars_are_validated_before_dispatch(self):
for kwargs in ({'text': 4, 'text_mode': 'shared'},
{'strength': float('nan'), 'strength_mode': 'shared'},
{'seed': True, 'seed_mode': 'shared'}):
with self.assertRaises(TypeError):
runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), **kwargs)
def test_three_sequential_regions_can_share_worker_loader(self):
prompt = graph_two_sequential_regions()
prompt.pop('8')
prompt.update({
'8': start_node(['7', 0]),
'9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
'10': {'class_type': runtime.END_CLASS,
'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
'11': {'class_type': 'GenericModelLoader', 'inputs': {}},
})
for node_id in ['3', '6', '9']:
prompt[node_id]['inputs']['model'] = ['11', 0]
# Shared dependencies are verified against the registered output schema.
loader = type('GenericModelLoader', (), {'RETURN_TYPES': ('MODEL',)})
registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'GenericModelLoader': loader})
with patch.dict(sys.modules, {'nodes': registry}):
regions = runtime.validate_image_list_regions(prompt)
self.assertEqual(len(regions), 3)
for region in regions:
self.assertIn('11', region['node_ids'])
+177
View File
@@ -0,0 +1,177 @@
"""Tests for periodic image-list server lifecycle cleanup."""
import asyncio
import time
import unittest
from types import SimpleNamespace
from typing import Any, cast
from unittest.mock import patch
from aiohttp import web
from image_list import lifecycle, runtime
class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
def setUp(self):
with runtime._LOCK:
runtime._MASTER_JOBS.clear()
runtime._WORKER_ITEMS.clear()
def tearDown(self):
with runtime._LOCK:
runtime._MASTER_JOBS.clear()
runtime._WORKER_ITEMS.clear()
@staticmethod
def _state(prompt_id=None, *, active=False, age=10_000.0):
return SimpleNamespace(
prompt_id=prompt_id,
active=active,
created_at=time.monotonic() - age,
)
@staticmethod
def _server(queue):
return SimpleNamespace(app=web.Application(), prompt_queue=queue)
async def _start(self, server, *, interval=0.01, grace=0.01):
lifecycle.install_lifecycle(
server,
runtime,
interval_seconds=interval,
orphan_grace_seconds=grace,
)
runner = web.AppRunner(server.app)
await runner.setup()
return runner
async def test_unattended_sweep_removes_orphan_and_expired_staging(self):
context_id = "orphan"
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)[context_id] = self._state()
runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"})
queue = SimpleNamespace(
get_current_queue_volatile=lambda: ([], []),
)
server = self._server(queue)
swept = asyncio.Event()
sweep_worker_items = runtime._sweep_worker_items
def track_worker_sweep():
swept.set()
sweep_worker_items()
with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep):
runner = await self._start(server)
try:
await asyncio.wait_for(swept.wait(), timeout=0.5)
self.assertNotIn(context_id, runtime._MASTER_JOBS)
self.assertNotIn("expired-token", runtime._WORKER_ITEMS)
with self.assertRaisesRegex(ValueError, "Unknown or expired"):
runtime.set_master_result(context_id, 0, "late-token", {})
finally:
await runner.cleanup()
async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self):
states = {
"active": self._state("active-prompt", active=True),
"queued": self._state("queued-prompt"),
"orphan": self._state("missing-prompt"),
"unbound": self._state(None),
}
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS).update(states)
queue = SimpleNamespace(
get_current_queue_volatile=lambda: (
[[1, "active-prompt", {}]],
[(2, "queued-prompt", {}, {}, None, None)],
),
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"])
self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"])
self.assertNotIn("orphan", runtime._MASTER_JOBS)
self.assertNotIn("unbound", runtime._MASTER_JOBS)
finally:
await runner.cleanup()
async def test_falls_back_to_nonvolatile_queue_reader(self):
state = self._state("queued-prompt")
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["queued"] = state
fallback_reads = 0
def volatile_failure():
raise RuntimeError("volatile view unavailable")
def current_queue():
nonlocal fallback_reads
fallback_reads += 1
return ([], [[1, "queued-prompt", {}]])
queue = SimpleNamespace(
get_current_queue_volatile=volatile_failure,
get_current_queue=current_queue,
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertGreaterEqual(fallback_reads, 1)
self.assertIs(runtime._MASTER_JOBS.get("queued"), state)
finally:
await runner.cleanup()
async def test_queue_read_error_protects_bound_contexts(self):
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt")
cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None)
def fail_read():
raise RuntimeError("queue unavailable")
queue = SimpleNamespace(
get_current_queue_volatile=fail_read,
get_current_queue=fail_read,
)
runner = await self._start(self._server(queue))
try:
await asyncio.sleep(0.04)
self.assertIn("bound", runtime._MASTER_JOBS)
self.assertNotIn("unbound", runtime._MASTER_JOBS)
finally:
await runner.cleanup()
async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self):
queue_reads = 0
def read_queue():
nonlocal queue_reads
queue_reads += 1
return ([], [])
server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue))
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
self.assertEqual(len(server.app.cleanup_ctx), 1)
with runtime._LOCK:
cast(Any, runtime._MASTER_JOBS)["removed"] = self._state()
runtime.remove_master_job("removed")
runtime.remove_master_job("removed")
self.assertNotIn("removed", runtime._MASTER_JOBS)
runner = web.AppRunner(server.app)
await runner.setup()
await asyncio.sleep(0.035)
self.assertGreaterEqual(queue_reads, 1)
await runner.cleanup()
reads_after_cleanup = queue_reads
await asyncio.sleep(0.035)
self.assertEqual(queue_reads, reads_after_cleanup)
if __name__ == "__main__":
unittest.main()
+340
View File
@@ -0,0 +1,340 @@
"""Exercise production per-item orchestration through controlled provider seams."""
import asyncio
import importlib.util
import json
from pathlib import Path
import sys
import types
import unittest
from unittest.mock import AsyncMock, Mock, patch
import torch
from test_image_list import graph_one_region
ROOT = Path(__file__).resolve().parents[1]
class Response:
def __init__(self, data=None, status=200):
self.data, self.status = data, status
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return None
async def json(self):
return self.data
def raise_for_status(self):
if self.status >= 400:
raise RuntimeError(f'HTTP {self.status}')
class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
name = '_image_list_boundaries'
self.modules = {}
for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
module = types.ModuleType(name + suffix)
module.__path__ = []
self.modules[name + suffix] = module
self.patch_modules = patch.dict(sys.modules, self.modules)
self.patch_modules.start()
self.addCleanup(self.patch_modules.stop)
spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
self.rt = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = self.rt
spec.loader.exec_module(self.rt)
self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188}
self.config = {'workers': [self.worker], 'settings': {}}
self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], []
self.reject = False
self.interrupted = False
self.completed = True
self.queue = {'queue_running': [], 'queue_pending': []}
self.callback = Mock(side_effect=lambda worker, **_kwargs:
'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master')
self.sync = AsyncMock(side_effect=self.sync_media)
self.dispatch = AsyncMock(side_effect=self.dispatch_prompt)
self.logs = []
dependencies = {
'.api.orchestration.dispatch': types.SimpleNamespace(
preflight_image_list_workers=AsyncMock(return_value={'worker': True}),
dispatch_worker_prompt=self.dispatch),
'.api.orchestration.prompt_transform': types.SimpleNamespace(
_resolve_direct_workers_for_image_list=lambda config, ids: config['workers']),
'.api.orchestration.media_sync': types.SimpleNamespace(
sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'),
convert_paths_for_platform=lambda prompt, separator: prompt),
'.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
get_worker_timeout_seconds=lambda: 60,
get_image_list_job_timeout_seconds=lambda config=None: 0),
'.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master',
build_master_callback_url=self.callback,
build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}",
get_client_session=AsyncMock(return_value=self)),
'.utils.logging': types.SimpleNamespace(log=self.logs.append),
}
for suffix, module in dependencies.items():
sys.modules[name + suffix] = module
self.config_module = dependencies['.utils.config']
# Other legacy test modules replace aiohttp during collection. This
# controlled provider fixture needs only its explicit timeout/error seam.
sys.modules['aiohttp'] = types.SimpleNamespace(
ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError)
comfy = types.ModuleType('comfy')
comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
sys.modules['comfy'] = comfy
sys.modules['comfy.model_management'] = comfy.model_management
async def sync_media(self, worker, prompt):
prompt['7']['inputs']['image'] = 'shared/ref.png'
async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
if self.reject:
raise RuntimeError('prompt rejected')
self.prompts.append((json.loads(json.dumps(prompt)), meta))
inputs = prompt['5']['inputs']
token = inputs['worker_result_token']
item = self.rt.consume_worker_item(token)
self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']})
return f'prompt-{len(self.prompts)}'
def post(self, url, *, json=None, data=None, **kwargs):
if data is not None:
json = __import__('json').loads(data)
if url.endswith('/distributed/image-list/item'):
self.uploads.append(json['token'])
self.rt.store_worker_item(json['token'], json['item'])
else:
self.cancel_requests.append((url, json))
return Response({})
def delete(self, url, **kwargs):
token = url.rsplit('/', 1)[-1]
self.deletions.append(token)
self.rt.delete_worker_item(token)
return Response({})
def get(self, url, **kwargs):
if url.endswith('/queue'):
return Response(self.queue)
pid = url.rsplit('/', 1)[-1]
return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}}
if self.completed else {})
def create(self, count=2):
prompt = graph_one_region()
prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}}
prompt['4']['inputs']['reference'] = ['7', 0]
self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}}
self.context = self.rt.create_master_job('3', '5', prompt,
[{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)],
['worker'], 'http://public-master', self.metadata)
return self.context
async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
context = self.create()
await self.rt.run_image_list_job(context)
self.assertEqual(self.sync.await_count, 1)
self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts],
['shared/ref.png', 'shared/ref.png'])
tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts]
self.assertEqual(len(set(tokens)), 2)
self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1])
self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata])
async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self):
import tempfile
context = self.create(2)
state = self.rt.get_master_job(context)
state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}}
state.original_prompt['4']['inputs']['second_reference'] = ['8', 0]
name = '_image_list_boundaries.api.orchestration.media_sync'
spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py')
assert spec is not None and spec.loader is not None
media = importlib.util.module_from_spec(spec)
sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append
sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace(
trace_debug=lambda *args: None, trace_info=lambda *args: None)
spec.loader.exec_module(media)
media.fetch_worker_path_separator = AsyncMock(return_value='\\')
converted = Mock(wraps=media.convert_paths_for_platform)
media.convert_paths_for_platform = converted
uploads = []
async def upload(worker, filename, content, digest, mime, **kwargs):
uploads.append((filename, content))
return True, f'worker-input/{len(uploads)}.png'
media._upload_media_to_worker = AsyncMock(side_effect=upload)
media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media)
with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory:
root = Path(directory)
for kind in ('input', 'output', 'temp'):
path = root / kind / 'shared/ref.png'
path.parent.mkdir(parents=True)
path.write_bytes(kind.encode())
def resolve_annotated(filename):
for kind in ('output', 'temp', 'input'):
suffix = f' [{kind}]'
if filename.endswith(suffix):
return str(root / kind / filename[:-len(suffix)])
return str(root / 'input' / filename)
with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace(
get_annotated_filepath=resolve_annotated)}):
await self.rt.run_image_list_job(context)
self.assertEqual(uploads, [('shared/ref.png [output]', b'output'),
('shared/ref.png [temp]', b'temp')])
media.sync_worker_media.assert_awaited_once()
media.fetch_worker_path_separator.assert_awaited_once()
converted.assert_called_once()
self.assertEqual(converted.call_args.args[1], '\\')
for prompt, _ in self.prompts:
self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png')
self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png')
self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]')
async def test_local_worker_callback_does_not_use_public_master_address(self):
self.worker['type'] = 'local'
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188')
self.callback.assert_called()
async def test_remote_worker_uses_topology_callback_helper(self):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master')
self.callback.assert_called()
def test_metadata_is_forwarded_into_master_start_prompt(self):
from test_prompt_transform import pt
prompt = graph_one_region()
metadata = {'nodes': [{'id': 1}]}
index = pt.PromptIndex(prompt)
transformed = pt.apply_participant_overrides(
prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False,
index, workflow_meta=metadata)
self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata)
async def test_aggregate_input_limit_fails_before_upload(self):
self.rt.MAX_REQUEST_BYTES = 64
with self.assertRaisesRegex(ValueError, 'complete.*request'):
await self.rt.run_image_list_job(self.create(1))
self.assertFalse(self.uploads)
self.assertFalse(self.rt._WORKER_ITEMS)
def test_tensor_and_combined_request_limits_preserve_dtype(self):
item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0}
encoded = self.rt.encode_item(item)
payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded}
single = self.rt.encode_request_payload(payload)
self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32')
self.rt.MAX_REQUEST_BYTES = len(single)
self.assertEqual(self.rt.encode_request_payload(payload), single)
item['reference_image'] = item['image']
item['mask'] = torch.zeros((1, 4, 4))
payload['item'] = self.rt.encode_item(item)
with self.assertRaisesRegex(ValueError, 'complete.*request'):
self.rt.encode_request_payload(payload)
# Size is checked before allocating/copying a 4K tensor onto the CPU.
with self.assertRaisesRegex(ValueError, 'transport limit'):
self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE')
async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self):
self.reject = True
with self.assertRaisesRegex(RuntimeError, 'prompt rejected'):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._WORKER_ITEMS)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertFalse(self.cancel_requests)
async def test_user_interrupt_during_pending_submission_is_bounded(self):
entered = asyncio.Event()
async def never_ack(*args, **kwargs):
entered.set()
await asyncio.Event().wait()
self.dispatch.side_effect = never_ack
self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02
task = asyncio.create_task(self.rt.run_image_list_job(self.create(2)))
await asyncio.wait_for(entered.wait(), 1)
self.interrupted = True
with self.assertRaises(self.rt.ImageListCancelledError):
await asyncio.wait_for(task, 0.8)
self.assertEqual(len(self.uploads), 1)
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._MASTER_JOBS)
self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs)
async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self):
self.completed = False
context = self.create(2)
clock = [0.0]
calls = [0]
original_get = self.get
def get(url, **kwargs):
if '/history/' in url:
calls[0] += 1
clock[0] += 3700
self.completed = calls[0] >= 3
if url.endswith('/queue'):
return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [],
'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []})
return original_get(url, **kwargs)
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
with patch.object(self, 'get', side_effect=get):
result = await self.rt.run_image_list_job(context)
self.assertEqual(len(result[0]), 2)
self.assertGreaterEqual(clock[0], 7200)
self.assertFalse(self.cancel_requests)
async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self):
self.completed = False
clock = [0.0]
original_get = self.get
def get(url, **kwargs):
if '/history/' in url:
clock[0] += 70
if url.endswith('/queue'):
return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []})
return original_get(url, **kwargs)
self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
with patch.object(self, 'get', side_effect=get):
with self.assertRaisesRegex(TimeoutError, 'liveness'):
await self.rt.run_image_list_job(self.create(1))
self.assertEqual(self.cancel_requests, [
('http://worker/queue', {'delete': ['prompt-1']}),
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self):
self.completed = False
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
original_post = self.post
def post(url, **kwargs):
result = original_post(url, **kwargs)
return Response(status=500) if url.endswith('/queue') else result
with patch.object(self, 'post', side_effect=post):
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
await self.rt.run_image_list_job(self.create(2))
self.assertEqual(self.cancel_requests, [
('http://worker/queue', {'delete': ['prompt-1']}),
('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
self.assertEqual(self.deletions, self.uploads)
self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs)
self.assertEqual(self.dispatch.await_count, 1)
async def test_configured_total_deadline_cleans_owned_work(self):
self.completed = False
self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []}
with self.assertRaisesRegex(TimeoutError, 'total deadline'):
await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1)
self.assertEqual(len(self.uploads), 1)
self.assertEqual(self.deletions, self.uploads)
self.assertFalse(self.rt._MASTER_JOBS)
+106
View File
@@ -0,0 +1,106 @@
import importlib.util
import sys
import types
from pathlib import Path
import pytest
def _load_config_module():
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
package_name = "image_list_timeout_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__ = [str(module_path.parent)]
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_module = _load_config_module()
SETTING_KEY = "image_list_job_timeout_seconds"
def _config_with_timeout(value):
return {"settings": {SETTING_KEY: value}}
def test_default_config_sets_image_list_job_timeout_to_zero():
assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0
def test_missing_value_defaults_to_disabled_timeout():
assert config_module.get_image_list_job_timeout_seconds({}) == 0.0
assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0
def test_none_config_loads_current_config(monkeypatch):
cfg = _config_with_timeout("12.5")
monkeypatch.setattr(config_module, "load_config", lambda: cfg)
assert config_module.get_image_list_job_timeout_seconds() == 12.5
@pytest.mark.parametrize(
("value", "expected"),
[
(0, 0.0),
(0.0, 0.0),
("0", 0.0),
(" 0.0 ", 0.0),
(1, 1.0),
(2.5, 2.5),
("3", 3.0),
(" 4.25 ", 4.25),
],
)
def test_accepts_zero_and_positive_numeric_values(value, expected):
result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
assert result == expected
assert isinstance(result, float)
@pytest.mark.parametrize(
"value",
[
True,
False,
-1,
-0.01,
float("nan"),
float("inf"),
float("-inf"),
"not-a-number",
"",
" ",
"-1",
"nan",
"inf",
None,
[],
{},
],
)
def test_rejects_invalid_values_with_setting_name(value):
with pytest.raises(ValueError, match=SETTING_KEY):
config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
+212
View File
@@ -0,0 +1,212 @@
import asyncio
import importlib.util
import sys
import types
import unittest
import uuid
from pathlib import Path
from typing import Any, cast
from unittest.mock import patch
import aiohttp
from aiohttp import web
def _load_dispatch_module():
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}"
package_modules = {}
for suffix in ("", ".api", ".api.orchestration", ".utils"):
package = types.ModuleType(package_name + suffix)
package.__path__ = []
package_modules[package_name + suffix] = package
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
setattr(logging_module, "log", lambda *_args, **_kwargs: None)
package_modules[logging_module.__name__] = logging_module
network_module = types.ModuleType(f"{package_name}.utils.network")
setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid")
async def _unused_probe_worker(*_args, **_kwargs):
return None
async def _unused_get_client_session(*_args, **_kwargs):
raise RuntimeError("get_client_session must be patched by this test")
setattr(network_module, "probe_worker", _unused_probe_worker)
setattr(network_module, "get_client_session", _unused_get_client_session)
package_modules[network_module.__name__] = network_module
module_name = f"{package_name}.api.orchestration.dispatch"
spec = importlib.util.spec_from_file_location(module_name, module_path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
package_modules[module_name] = module
previous_modules = {name: sys.modules.get(name) for name in package_modules}
try:
sys.modules.update(package_modules)
spec.loader.exec_module(module)
finally:
for name, previous in previous_modules.items():
if previous is None:
sys.modules.pop(name, None)
else:
sys.modules[name] = previous
return module
dispatch = _load_dispatch_module()
class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.mode = None
self.received_payload = None
self.dispatch_received = asyncio.Event()
self.peer_finished = asyncio.Event()
self.irrelevant_messages_sent = 0
app = web.Application()
app.router.add_get("/distributed/worker_ws", self._handle_websocket)
self.runner = web.AppRunner(app)
await self.runner.setup()
self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
await self.site.start()
server = cast(Any, self.site._server)
port = server.sockets[0].getsockname()[1]
self.worker_url = f"http://127.0.0.1:{port}"
self.session = aiohttp.ClientSession()
async def asyncTearDown(self):
await self.session.close()
await self.runner.cleanup()
async def _get_client_session(self):
return self.session
async def _handle_websocket(self, request):
ws = web.WebSocketResponse()
await ws.prepare(request)
try:
payload = await ws.receive_json()
self.received_payload = payload
self.dispatch_received.set()
if self.mode in ("success", "slow-close"):
await ws.send_json({
"type": "dispatch_ack",
"request_id": "unrelated-request",
"ok": True,
"prompt_id": "wrong-prompt",
})
await ws.send_json({"type": "heartbeat"})
await ws.send_json({
"type": "dispatch_ack",
"request_id": payload["request_id"],
"ok": True,
"prompt_id": "prompt-123",
})
if self.mode == "slow-close":
await asyncio.sleep(0.15)
elif self.mode == "rejection":
await ws.send_json({
"type": "dispatch_ack",
"request_id": payload["request_id"],
"ok": False,
"error": "Invalid prompt",
"validation_error": {"node": "bad"},
"node_errors": {"bad": "missing input"},
})
elif self.mode in ("no-ack", "wait-for-cancel"):
async def send_irrelevant_messages():
while True:
try:
await ws.send_json({"type": "heartbeat"})
self.irrelevant_messages_sent += 1
await ws.send_json({
"type": "dispatch_ack",
"request_id": "unrelated-request",
"ok": True,
"prompt_id": "not-this-request",
})
self.irrelevant_messages_sent += 1
except (ConnectionError, RuntimeError):
return
await asyncio.sleep(0.01)
sender = asyncio.create_task(send_irrelevant_messages())
try:
async for _message in ws:
pass
finally:
sender.cancel()
try:
await sender
except asyncio.CancelledError:
pass
finally:
self.peer_finished.set()
return ws
async def _dispatch(self, timeout):
with patch.object(dispatch, "get_client_session", self._get_client_session):
return await dispatch._dispatch_via_websocket(
self.worker_url,
{"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}},
"client-456",
timeout=timeout,
)
async def test_returns_matching_acknowledgement_prompt_id(self):
self.mode = "success"
prompt_id = await self._dispatch(timeout=0.5)
self.assertEqual(prompt_id, "prompt-123")
received_payload = self.received_payload
assert received_payload is not None
self.assertEqual(received_payload["type"], "dispatch_prompt")
self.assertTrue(received_payload["request_id"])
self.assertEqual(received_payload["client_id"], "client-456")
self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}})
self.assertEqual(received_payload["workflow"], {"name": "test"})
async def test_acknowledged_prompt_id_survives_slow_close(self):
self.mode = "slow-close"
self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123")
async def test_matching_rejection_preserves_worker_error_details(self):
self.mode = "rejection"
with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="):
await self._dispatch(timeout=0.5)
async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self):
self.mode = "no-ack"
started = asyncio.get_running_loop().time()
with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"):
await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5)
elapsed = asyncio.get_running_loop().time() - started
self.assertLess(elapsed, 0.5)
self.assertGreaterEqual(self.irrelevant_messages_sent, 4)
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
async def test_outer_cancellation_propagates_and_closes_peer(self):
self.mode = "wait-for-cancel"
task = asyncio.create_task(self._dispatch(timeout=2.0))
await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5)
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=0.5)
await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
if __name__ == "__main__":
unittest.main()
+10 -22
View File
@@ -86,27 +86,25 @@ class ParseTilesFromFormTests(unittest.TestCase):
# --- happy paths ---
def test_single_tile_returns_one_entry(self):
def test_single_tile_returns_image_and_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertEqual(len(tiles), 1)
def test_multiple_tiles_all_returned(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(len(tiles), 3)
def test_tile_image_is_pil_image(self):
tiles = pp._parse_tiles_from_form(_make_form(1))
self.assertIsInstance(tiles[0]["image"], PILImage.Image)
def test_tile_metadata_fields_are_parsed(self):
tiles = pp._parse_tiles_from_form(_make_form(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)
@@ -138,16 +136,6 @@ class ParseTilesFromFormTests(unittest.TestCase):
self.assertNotIn("batch_idx", tiles[0])
self.assertNotIn("global_idx", tiles[0])
def test_tile_indices_match_metadata_order(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
for i, tile in enumerate(tiles):
self.assertEqual(tile["tile_idx"], i)
def test_x_coordinates_reflect_metadata(self):
tiles = pp._parse_tiles_from_form(_make_form(3))
self.assertEqual(tiles[1]["x"], 64)
self.assertEqual(tiles[2]["x"], 128)
# --- error cases ---
def test_missing_tiles_metadata_raises_value_error(self):
+541 -9
View File
@@ -71,6 +71,15 @@ def _collector_only_prompt():
}
def _audio_only_prompt():
"""1(Audio source) → 2(DistributedCollector) → 3(Audio sink)."""
return {
"1": {"class_type": "LoadAudio", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
"3": {"class_type": "SaveAudio", "inputs": {"audio": ["2", 1]}},
}
def _delegate_prompt():
"""1 → 2 → 3(DistributedCollector) → 4(SaveImage)"""
return {
@@ -217,6 +226,35 @@ class PrunePromptForWorkerTests(unittest.TestCase):
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["4", 0])
def test_injects_preview_audio_for_audio_only_collector(self):
result = pt.prune_prompt_for_worker(_audio_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewAudio"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["audio"], ["2", 1])
self.assertFalse(any(n.get("class_type") == "PreviewImage" for n in result.values()))
def test_prefers_image_preview_when_collector_has_images_and_audio(self):
prompt = _linear_prompt()
prompt["6"] = {"class_type": "LoadAudio", "inputs": {}}
prompt["4"]["inputs"]["audio"] = ["6", 0]
result = pt.prune_prompt_for_worker(prompt)
self.assertEqual(
len([n for n in result.values() if n.get("class_type") == "PreviewImage"]),
1,
)
self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
def test_preserves_image_preview_for_distributed_upscale(self):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
self.assertEqual(len(preview_nodes), 1)
self.assertEqual(preview_nodes[0]["inputs"]["images"], ["2", 0])
def test_no_preview_image_when_no_downstream(self):
result = pt.prune_prompt_for_worker(_collector_only_prompt())
preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
@@ -242,7 +280,7 @@ class PrunePromptForWorkerTests(unittest.TestCase):
def test_upscale_node_is_treated_as_distributed(self):
prompt = {
"1": {"class_type": "KSampler", "inputs": {}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"image": ["1", 0]}},
"2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
}
result = pt.prune_prompt_for_worker(prompt)
@@ -264,8 +302,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotIn("1", result)
self.assertNotIn("2", result)
def test_removes_dangling_upstream_refs(self):
"""Collector must not retain dangling refs to pruned upstream nodes."""
def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
"""Collector must point to exactly one valid placeholder after pruning."""
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
collector_inputs = result["3"].get("inputs", {})
@@ -276,26 +314,520 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
self.assertNotEqual(source_id, "2")
self.assertIn(source_id, result)
self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
def test_injects_empty_image_placeholder(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 1)
placeholder_id = empty_nodes[0][0]
self.assertEqual(result["3"]["inputs"]["images"], [placeholder_id, 0])
def test_audio_only_collector_does_not_get_image_placeholder(self):
prompt = _audio_only_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(empty_nodes, [])
self.assertNotIn("images", result["2"].get("inputs", {}))
self.assertNotIn("audio", result["2"].get("inputs", {}))
def test_one_placeholder_per_collector(self):
"""Two collectors → two placeholders."""
prompt = {
"1": {"class_type": "DistributedCollector", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {}},
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0]}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["11", 0]}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
"10": {"class_type": "LoadImage", "inputs": {}},
"11": {"class_type": "LoadImage", "inputs": {}},
}
result = pt.prepare_delegate_master_prompt(prompt, ["1", "2"])
empty_nodes = [n for n in result.values() if n.get("class_type") == "DistributedEmptyImage"]
self.assertEqual(len(empty_nodes), 2)
def test_preserves_primitive_string_for_downstream_required_input(self):
"""Delegate-only master keeps primitive inputs needed by SaveImage."""
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {"SaveImage": save_image}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test/input_bug"}},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["15", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertEqual(result["15"], prompt["15"])
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
def test_preserves_load_image_for_switch_alternate_required_input(self):
"""Delegate-only master keeps LoadImage inputs needed by switches."""
prompt = {
"848": {"class_type": "LoadImage", "inputs": {"image": "input_bug2_00003_.png"}},
"862": {"class_type": "VAEDecode", "inputs": {}},
"854": {"class_type": "DistributedCollector", "inputs": {"images": ["862", 0]}},
"850": {
"class_type": "ComfySwitchNode",
"inputs": {
"on_false": ["848", 0],
"on_true": ["854", 0],
},
},
"851": {"class_type": "PreviewImage", "inputs": {"images": ["850", 0]}},
}
result = pt.prepare_delegate_master_prompt(prompt, ["854"])
self.assertIn("848", result)
self.assertEqual(result["850"]["inputs"]["on_false"], ["848", 0])
self.assertEqual(result["850"]["inputs"]["on_true"], ["854", 0])
self.assertNotIn("862", result)
def test_preserves_registered_string_utility_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps scalar utility chains used by SaveImage."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug3"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["16", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["17", 0])
self.assertEqual(result["17"]["inputs"]["string_a"], ["15", 0])
self.assertEqual(result["17"]["inputs"]["string_b"], ["16", 0])
def test_preserves_registered_multi_string_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps multi-input scalar utility chains."""
join_string_multi = type(
"JoinStringMulti",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"string_1": ("STRING",)},
"optional": {
"string_2": ("STRING",),
"string_3": ("STRING",),
},
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"JoinStringMulti": join_string_multi,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"16": {"class_type": "PrimitiveString", "inputs": {"value": "input"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "bug4"}},
"18": {
"class_type": "JoinStringMulti",
"inputs": {
"string_1": ["15", 0],
"string_2": ["16", 0],
"string_3": ["17", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["18", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("16", result)
self.assertIn("17", result)
self.assertIn("18", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["18", 0])
def test_does_not_preserve_scalar_utility_with_heavy_upstream_dependency(self):
"""Scalar utility nodes are retained only when their full input branch is safe."""
string_concat = type(
"StringConcatenate",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"string_a": ("STRING",),
"string_b": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"SaveImage": save_image,
"StringConcatenate": string_concat,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {
"class_type": "StringConcatenate",
"inputs": {
"string_a": ["15", 0],
"string_b": ["8", 0],
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("8", result)
self.assertNotIn("17", result)
self.assertNotIn("filename_prefix", result["9"]["inputs"])
def test_does_not_preserve_scalar_output_for_non_scalar_downstream_input(self):
"""Scalar outputs are retained only for scalar/config downstream inputs."""
string_provider = type("StringProvider", (), {"RETURN_TYPES": ("STRING",)})
image_consumer = type(
"ImageConsumer",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"mask": ("IMAGE",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"ImageConsumer": image_consumer,
"StringProvider": string_provider,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"17": {"class_type": "StringProvider", "inputs": {}},
"9": {
"class_type": "ImageConsumer",
"inputs": {
"images": ["11", 0],
"mask": ["17", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertNotIn("17", result)
self.assertNotIn("mask", result["9"]["inputs"])
def test_preserves_scalar_list_join_subgraph_for_downstream_required_input(self):
"""Delegate-only master keeps list-shaped scalar config chains."""
create_list = type(
"CreateList",
(),
{
"RETURN_TYPES": ("LIST",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {"inputs.input0": ("STRING",)},
"optional": {
"inputs.input1": ("STRING",),
"inputs.input2": ("STRING",),
},
}
),
},
)
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("LIST",),
"delimiter": ("STRING",),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"CreateList": create_list,
"SaveImage": save_image,
"StringDataListJoin": string_data_list_join,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["15", 0],
"inputs.input1": ["17", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"delimiter": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
self.assertIn("15", result)
self.assertIn("17", result)
self.assertIn("28", result)
self.assertIn("29", result)
self.assertIn("32", result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["15", 0])
self.assertNotIn("8", result)
def test_preserves_builtin_create_list_string_data_list_join_subgraph(self):
"""Delegate-only master handles ComfyUI 0.23 CreateList data-list output."""
string_data_list_join = type(
"StringDataListJoin",
(),
{
"RETURN_TYPES": ("STRING",),
"INPUT_IS_LIST": True,
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"strings": ("STRING", {"forceInput": True}),
"sep": ("STRING", {"default": " "}),
}
}
),
},
)
save_image = type(
"SaveImage",
(),
{
"INPUT_TYPES": classmethod(
lambda cls: {
"required": {
"images": ("IMAGE",),
"filename_prefix": ("STRING",),
}
}
)
},
)
previous_mappings = getattr(pt, "_DELEGATE_MASTER_NODE_CLASS_MAPPINGS", None)
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = {
"Basic data handling: StringDataListJoin": string_data_list_join,
"SaveImage": save_image,
}
try:
prompt = {
"8": {"class_type": "KSampler", "inputs": {}},
"11": {"class_type": "DistributedCollector", "inputs": {"images": ["8", 0]}},
"15": {"class_type": "PrimitiveString", "inputs": {"value": "input_bug"}},
"17": {"class_type": "PrimitiveString", "inputs": {"value": "test"}},
"29": {"class_type": "PrimitiveString", "inputs": {"value": "new"}},
"28": {
"class_type": "CreateList",
"inputs": {
"inputs.input0": ["17", 0],
"inputs.input1": ["15", 0],
"inputs.input2": ["29", 0],
},
},
"32": {
"class_type": "Basic data handling: StringDataListJoin",
"inputs": {
"strings": ["28", 0],
"sep": "/",
},
},
"9": {
"class_type": "SaveImage",
"inputs": {
"images": ["11", 0],
"filename_prefix": ["32", 0],
},
},
}
result = pt.prepare_delegate_master_prompt(prompt, ["11"])
finally:
pt._DELEGATE_MASTER_NODE_CLASS_MAPPINGS = previous_mappings
for node_id in ("15", "17", "28", "29", "32"):
self.assertIn(node_id, result)
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["32", 0])
self.assertEqual(result["32"]["inputs"]["strings"], ["28", 0])
self.assertEqual(result["28"]["inputs"]["inputs.input0"], ["17", 0])
self.assertNotIn("8", result)
def test_result_is_independent_copy(self):
prompt = _delegate_prompt()
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+167
View File
@@ -0,0 +1,167 @@
import importlib.util
import json
import sys
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
def load_module():
spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py")
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
client = load_module()
def workflow_for_request_size(target_size, timeout_seconds):
"""Build an ASCII workflow whose compact UTF-8 request body has exact size."""
workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}}
def body_size(candidate):
payload = {
"input": {"workflow": candidate, "images": []},
"policy": {
"executionTimeout": timeout_seconds * 1000,
"ttl": timeout_seconds * 1000,
},
}
return len(json.dumps(
payload, allow_nan=False, ensure_ascii=False, separators=(",", ":")
).encode("utf-8"))
padding = target_size - body_size(workflow)
if padding < 0:
raise AssertionError("target is smaller than the request fixture")
workflow["1"]["inputs"]["value"] = "x" * padding
if body_size(workflow) != target_size:
raise AssertionError("request fixture did not reach the exact target size")
return workflow
class Response:
def __init__(self, status=200, payload=None):
self.status, self.payload = status, payload
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return False
async def json(self):
return self.payload
class Session:
def __init__(self, responses):
self.responses, self.calls = list(responses), []
def request(self, method, url, **kwargs):
self.calls.append((method, url, kwargs))
response = self.responses.pop(0)
if isinstance(response, Exception):
raise response
return response
class ClientTests(unittest.IsolatedAsyncioTestCase):
async def test_submit_exact_generic_contract_and_once(self):
session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})])
api = client.RunpodClient(session, "endpoint123", "secret-value")
result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800)
self.assertEqual(result, "job-1")
method, url, kwargs = session.calls[0]
self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run"))
self.assertIsInstance(kwargs["data"], bytes)
sent_payload = json.loads(kwargs["data"])
self.assertEqual(set(sent_payload["input"]), {"workflow", "images"})
self.assertEqual(sent_payload["policy"]["ttl"], 1800000)
self.assertFalse(kwargs["allow_redirects"])
self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value")
async def test_submit_timeout_is_ambiguous_and_never_retried(self):
session = Session([TimeoutError("secret-value in unsafe error")])
api = client.RunpodClient(session, "endpoint123", "secret-value")
with self.assertRaises(client.AmbiguousSubmission) as ctx:
await api.submit({}, [], 30)
self.assertNotIn("secret-value", str(ctx.exception))
self.assertEqual(len(session.calls), 1)
async def test_http_submit_failure_not_retried_or_leaked(self):
for status in [401, 429, 500, 302]:
session = Session([Response(status, {"error": "secret-value"})])
api = client.RunpodClient(session, "endpoint123", "secret-value")
with self.assertRaises(client.RunpodError) as ctx:
await api.submit({}, [], 30)
self.assertEqual(len(session.calls), 1)
self.assertNotIn("secret-value", str(ctx.exception))
async def test_health_zero_workers_is_valid(self):
session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})])
api = client.RunpodClient(session, "endpoint123", "key")
self.assertEqual((await api.health())["workers"]["idle"], 0)
async def test_invalid_provider_id_no_network(self):
session = Session([])
with self.assertRaises(ValueError):
client.RunpodClient(session, "https://attacker.invalid", "key")
api = client.RunpodClient(session, "endpoint123", "key")
with self.assertRaises(ValueError):
await api.status("../other")
self.assertFalse(session.calls)
async def test_status_and_cancel_only_target_exact_job(self):
session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})])
api = client.RunpodClient(session, "endpoint123", "key")
await api.status("job-1")
await api.cancel("job-1")
self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")])
async def test_missing_id_is_ambiguous(self):
api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key")
with self.assertRaises(client.AmbiguousSubmission):
await api.submit({}, [], 30)
async def test_large_payload_rejected_before_submit(self):
session = Session([])
api = client.RunpodClient(session, "endpoint123", "key")
with self.assertRaises(ValueError):
await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30)
self.assertFalse(session.calls)
async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self):
session = Session([Response(payload={"id": "boundary-job"})])
api = client.RunpodClient(session, "endpoint123", "key")
timeout_seconds = 37
workflow = workflow_for_request_size(10_000_000, timeout_seconds)
result = await api.submit(workflow, [], timeout_seconds)
self.assertEqual(result, "boundary-job")
sent = session.calls[0][2]["data"]
self.assertIsInstance(sent, bytes)
self.assertEqual(len(sent), 10_000_000)
self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000)
async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self):
session = Session([Response(payload={"id": "must-not-be-used"})])
api = client.RunpodClient(session, "endpoint123", "key")
workflow = workflow_for_request_size(10_000_001, 37)
workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:]
with self.assertRaisesRegex(ValueError, "10,000,001") as ctx:
await api.submit(workflow, [], 37)
self.assertNotIn("private graph value", str(ctx.exception))
self.assertFalse(session.calls)
async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self):
session = Session([Response(payload={"id": "unicode-job"})])
api = client.RunpodClient(session, "endpoint123", "key")
workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}}
await api.submit(workflow, [], 41)
sent = session.calls[0][2]["data"]
expected = json.dumps({
"input": {"workflow": workflow, "images": []},
"policy": {"executionTimeout": 41_000, "ttl": 41_000},
}, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
self.assertEqual(sent, expected)
self.assertGreater(len(sent), len(sent.decode("utf-8")))
+202
View File
@@ -0,0 +1,202 @@
import asyncio
import pytest
import torch
from test_collector_list_inputs import _load_collector_module
class _FakeManager:
def __init__(self, *, pending=False, owned=True):
self.pending = pending
self.owned = owned
self.pending_calls = []
self.owned_calls = []
self.cancel_calls = []
def is_pending(self, collector_job_id, participant_id):
self.pending_calls.append((collector_job_id, participant_id))
return self.pending
def owns(self, collector_job_id, participant_id):
self.owned_calls.append((collector_job_id, participant_id))
return self.owned
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
self.cancel_calls.append({
"worker_id": worker_id,
"execution_id": execution_id,
"collector_job_id": collector_job_id,
})
return 1
def _set_up_collector(module, job_id, queue, manager=None):
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
module.prompt_server.distributed_pending_jobs = {job_id: queue}
if manager is not None:
module.prompt_server.distributed_serverless = manager
return module.DistributedCollectorNode()
def test_serverless_pending_waits_past_normal_worker_timeout():
module = _load_collector_module()
job_id = "collector-pending"
participant = "serverless:pool:1"
queue = asyncio.Queue()
manager = _FakeManager(pending=True)
collector = _set_up_collector(module, job_id, queue, manager)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
probe_calls = []
module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
async def deliver_after_provider_work():
await asyncio.sleep(0.14)
queue.put_nowait({
"worker_id": participant,
"image_index": 0,
"tensor": torch.ones(1, 2, 2, 2),
"is_last": True,
})
async def run_collection():
producer = asyncio.create_task(deliver_after_provider_work())
result = await collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids=f'["{participant}"]',
delegate_only=True,
)
await producer
return result
images, _audio = asyncio.run(run_collection())
assert tuple(images.shape) == (1, 2, 2, 2)
assert manager.pending_calls
assert probe_calls == []
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
module = _load_collector_module()
job_id = "collector-error"
queue = asyncio.Queue()
queue.put_nowait({
"worker_id": "serverless:pool:1",
"serverless_error": "Provider execution failed safely",
"is_last": True,
})
manager = _FakeManager()
collector = _set_up_collector(module, job_id, queue, manager)
with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.cancel_calls == [{
"worker_id": None,
"execution_id": None,
"collector_job_id": job_id,
}]
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
module = _load_collector_module()
job_id = "collector-stopped"
queue = asyncio.Queue()
manager = _FakeManager(pending=False, owned=True)
collector = _set_up_collector(module, job_id, queue, manager)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.owned_calls
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_interruption_cancels_serverless_job_and_removes_collector_queue():
module = _load_collector_module()
job_id = "collector-interrupted"
queue = asyncio.Queue()
manager = _FakeManager(pending=True)
collector = _set_up_collector(module, job_id, queue, manager)
interrupted = module.comfy.model_management.InterruptProcessingException("stop")
module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
asyncio.run(collector.execute(
images=None,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
delegate_only=True,
))
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
module = _load_collector_module()
job_id = "collector-shape-mismatch"
queue = asyncio.Queue()
queue.put_nowait({
"worker_id": "serverless:pool:1",
"image_index": 0,
"tensor": torch.ones(1, 3, 2, 2),
"is_last": True,
})
manager = _FakeManager()
collector = _set_up_collector(module, job_id, queue, manager)
local_images = torch.zeros(1, 2, 2, 2)
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
asyncio.run(collector.execute(
images=local_images,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["serverless:pool:1"]',
))
assert manager.cancel_calls[0]["collector_job_id"] == job_id
assert job_id not in module.prompt_server.distributed_pending_jobs
def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
module = _load_collector_module()
job_id = "collector-regular-timeout"
queue = asyncio.Queue()
collector = _set_up_collector(module, job_id, queue)
module.get_worker_timeout_seconds = lambda: 0.01
module.HEARTBEAT_INTERVAL = 0.1
local_images = torch.zeros(1, 2, 2, 2)
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
images, _audio = asyncio.run(collector.execute(
images=local_images,
audio=None,
multi_job_id=job_id,
enabled_worker_ids='["worker-a"]',
))
assert torch.equal(images, local_images)
assert job_id not in module.prompt_server.distributed_pending_jobs
+398
View File
@@ -0,0 +1,398 @@
import importlib.util
import json
import os
import stat
import sys
import tempfile
import threading
import types
import unittest
from pathlib import Path
from unittest import mock
def _load_serverless_config_module(*, without_fcntl=False):
module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py"
module_name = (
"dist_serverless_config_no_fcntl_test_module"
if without_fcntl
else "dist_serverless_config_test_module"
)
package_name = module_name + "_pkg"
package = types.ModuleType(package_name)
package.__path__ = [str(module_path.parent)]
sys.modules[package_name] = package
module_name = package_name + ".serverless_config"
sys.modules.pop(module_name, None)
spec = importlib.util.spec_from_file_location(module_name, module_path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
if without_fcntl:
with mock.patch.dict(sys.modules, {"fcntl": None}):
spec.loader.exec_module(module)
else:
spec.loader.exec_module(module)
return module
serverless_config = _load_serverless_config_module()
class ValidateServerlessWorkerTests(unittest.TestCase):
def test_returns_copy_with_runpod_defaults(self):
worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}
result = serverless_config.validate_serverless_worker(worker)
self.assertEqual(
result,
{
"id": "worker-1",
"endpoint_id": "ep_123-abc",
"name": "GPU",
"type": "serverless",
"provider": "runpod",
"jobs_per_run": 1,
"concurrency": 1,
"timeout_seconds": 1800,
},
)
self.assertIsNot(result, worker)
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
def test_accepts_inclusive_numeric_bounds(self):
for jobs, concurrency, timeout in ((1, 1, 30), (16, 16, 86400)):
with self.subTest(jobs=jobs, concurrency=concurrency, timeout=timeout):
result = serverless_config.validate_serverless_worker(
{
"id": "worker-1",
"endpoint_id": "ep-123",
"jobs_per_run": jobs,
"concurrency": concurrency,
"timeout_seconds": timeout,
}
)
self.assertEqual(result["jobs_per_run"], jobs)
self.assertEqual(result["concurrency"], concurrency)
self.assertEqual(result["timeout_seconds"], timeout)
def test_rejects_missing_or_empty_id_and_endpoint(self):
base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
with self.subTest(field=field, value=value):
worker = dict(base)
worker[field] = value
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
worker = dict(base)
del worker["id"]
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
def test_endpoint_id_must_be_ascii_opaque_identifier(self):
for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"):
with self.subTest(endpoint_id=endpoint_id):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": endpoint_id}
)
def test_endpoint_id_matches_runpod_client_length_limit(self):
valid_id = "e" * 200
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": valid_id}
)
self.assertEqual(result["endpoint_id"], valid_id)
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "e" * 201}
)
def test_enabled_must_be_boolean_when_supplied(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
for enabled in (True, False):
with self.subTest(enabled=enabled):
self.assertIs(
serverless_config.validate_serverless_worker(
{**base, "enabled": enabled}
)["enabled"],
enabled,
)
for enabled in ("false", 0, 1, None):
with self.subTest(enabled=enabled):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(
{**base, "enabled": enabled}
)
def test_rejects_unknown_provider_and_non_serverless_type(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
for extra in ({"provider": "other"}, {"type": "local"}):
with self.subTest(extra=extra):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker({**base, **extra})
def test_numeric_fields_reject_bool_float_and_out_of_range(self):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
invalid_values = {
"jobs_per_run": (True, 1.0, 0, 17),
"concurrency": (False, 2.5, 0, 17),
"timeout_seconds": (True, 30.0, 29, 86401),
}
for field, values in invalid_values.items():
for value in values:
with self.subTest(field=field, value=value):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker({**base, field: value})
def test_requires_worker_object(self):
for worker in (None, [], "worker"):
with self.subTest(worker=worker):
with self.assertRaises(ValueError):
serverless_config.validate_serverless_worker(worker)
class SecretStoreTests(unittest.TestCase):
def setUp(self):
self.tempdir = tempfile.TemporaryDirectory()
self.root = Path(self.tempdir.name)
self.path = self.root / "private" / "runpod-secrets.json"
def tearDown(self):
self.tempdir.cleanup()
def test_set_get_has_delete_and_persists_as_json_mapping(self):
store = serverless_config.SecretStore(self.path)
store.set("worker-1", "secret-value")
self.assertEqual(store.get("worker-1"), "secret-value")
self.assertTrue(store.has("worker-1"))
self.assertIsNone(store.get("missing"))
self.assertFalse(store.has("missing"))
self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"})
self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value")
store.delete("worker-1")
self.assertIsNone(store.get("worker-1"))
self.assertFalse(store.has("worker-1"))
def test_rejects_api_keys_with_newlines_or_over_4096_characters(self):
store = serverless_config.SecretStore(self.path)
for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097):
with self.subTest(key_length=len(key)):
with self.assertRaises(serverless_config.SecretStoreError):
store.set("worker-1", key)
store.set("worker-1", "k" * 4096)
self.assertEqual(store.get("worker-1"), "k" * 4096)
def test_new_secret_file_and_directory_have_restrictive_modes(self):
store = serverless_config.SecretStore(self.path)
store.set("worker-1", "secret-value")
self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600)
self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700)
def test_missing_store_reads_as_empty_without_creating_files(self):
store = serverless_config.SecretStore(self.path)
self.assertIsNone(store.get("worker-1"))
self.assertFalse(store.has("worker-1"))
self.assertFalse(self.path.parent.exists())
def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self):
self.path.parent.mkdir()
corrupt = b'{"partial": "NEVER-LEAK-THIS-secret'
self.path.write_bytes(corrupt)
store = serverless_config.SecretStore(self.path)
for operation in (
lambda: store.get("worker-1"),
lambda: store.set("worker-1", "replacement-secret"),
lambda: store.delete("worker-1"),
):
with self.subTest(operation=operation):
with self.assertRaises(Exception) as raised:
operation()
self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception))
self.assertNotIn("replacement-secret", str(raised.exception))
self.assertEqual(self.path.read_bytes(), corrupt)
def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self):
for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'):
with self.subTest(corrupt=corrupt):
self.path.parent.mkdir(exist_ok=True)
self.path.write_bytes(corrupt)
with self.assertRaises(Exception) as raised:
serverless_config.SecretStore(self.path).get("worker")
self.assertNotIn("opaque-stored-value", str(raised.exception))
def test_rejects_symlink_file_and_preserves_link_target(self):
target = self.root / "target.json"
target.write_text('{"worker-1": "target-secret"}', encoding="utf-8")
self.path.parent.mkdir()
self.path.symlink_to(target)
store = serverless_config.SecretStore(self.path)
with self.assertRaises(Exception):
store.get("worker-1")
with self.assertRaises(Exception):
store.set("worker-1", "replacement-secret")
self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}')
def test_rejects_symlinked_parent_directory(self):
real_dir = self.root / "real"
real_dir.mkdir()
link_dir = self.root / "linked"
link_dir.symlink_to(real_dir, target_is_directory=True)
store = serverless_config.SecretStore(link_dir / "secrets.json")
with self.assertRaises(Exception):
store.set("worker-1", "secret-value")
self.assertFalse((real_dir / "secrets.json").exists())
def test_rejects_parent_traversal_in_store_path(self):
store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json")
with self.assertRaises(Exception):
store.set("worker-1", "secret-value")
self.assertFalse((self.root / "secrets.json").exists())
def test_threaded_updates_from_independent_instances_are_not_lost(self):
workers = 12
barrier = threading.Barrier(workers)
errors = []
def write_one(index):
try:
barrier.wait(timeout=5)
serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}")
except Exception as exc: # collected to fail in the test thread
errors.append(exc)
threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=10)
self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish")
self.assertEqual(errors, [])
saved = json.loads(self.path.read_text(encoding="utf-8"))
self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)})
class PublicConfigTests(unittest.TestCase):
def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self):
with tempfile.TemporaryDirectory() as tempdir:
store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json")
store.set("serverless-1", "stored-api-key")
config = {
"workers": [
{
"id": "serverless-1",
"type": "serverless",
"provider": "runpod",
"endpoint_id": "ep-123",
"api_key": "inline-secret",
"clear_api_key": "another-secret",
"credential": "credential-secret",
"nested": {"runpod_credential": "nested-secret", "safe": "kept"},
"name": "remote GPU",
},
{"id": "serverless-2", "type": "serverless", "name": "not configured"},
{"id": "local-1", "type": "local", "api_key": "local-secret"},
],
"settings": {"mode": "preserved", "credentials_note": "remove-me"},
}
original = json.loads(json.dumps(config))
result = serverless_config.public_config(config, store)
first, second, local = result["workers"]
self.assertEqual(first["api_key_configured"], True)
self.assertEqual(second["api_key_configured"], False)
self.assertNotIn("api_key_configured", local)
self.assertNotIn("api_key", first)
self.assertNotIn("clear_api_key", first)
self.assertNotIn("credential", first)
self.assertNotIn("runpod_credential", first["nested"])
self.assertEqual(first["nested"]["safe"], "kept")
self.assertNotIn("api_key", local)
self.assertNotIn("credentials_note", result["settings"])
self.assertEqual(first["name"], "remote GPU")
self.assertEqual(config, original)
self.assertIsNot(result, config)
self.assertIsNot(result["workers"], config["workers"])
def test_redacts_normalized_secret_field_aliases(self):
config = {
"workers": [{
"id": "serverless-1",
"type": "serverless",
"apiKey": "camel-api-key",
"APIKEY": "upper-api-key",
"clearApiKey": "camel-clear-key",
"CLEAR_API_KEY": "upper-clear-key",
"runpod_CREDENTIAL": "provider-credential",
"nested": {"credential_note": "nested-credential"},
"safe": "preserved",
}]
}
result = serverless_config.public_config(
config, serverless_config.SecretStore(self.root_path())
)
worker = result["workers"][0]
self.assertEqual(worker, {
"id": "serverless-1",
"type": "serverless",
"nested": {},
"safe": "preserved",
"api_key_configured": False,
})
def test_handles_non_dict_config_without_mutating_it(self):
config = ["safe", {"api_key": "secret"}]
result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path()))
self.assertEqual(result, ["safe", {}])
self.assertEqual(config, ["safe", {"api_key": "secret"}])
def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self):
module = _load_serverless_config_module(without_fcntl=True)
store = module.SecretStore(self.root_path())
config = {
"workers": [
{"id": "local-1", "type": "local", "api_key": "inline-secret"}
]
}
result = module.public_config(config, store)
self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]})
self.assertEqual(config["workers"][0]["api_key"], "inline-secret")
def test_secret_operations_fail_clearly_without_fcntl(self):
module = _load_serverless_config_module(without_fcntl=True)
store = module.SecretStore(self.root_path())
unsupported_message = "Serverless credential storage is currently supported on POSIX masters only"
for operation in (
lambda: store.get("worker-1"),
lambda: store.set("worker-1", "api-key"),
lambda: store.delete("worker-1"),
lambda: store.has("worker-1"),
):
with self.subTest(operation=operation):
with self.assertRaisesRegex(module.SecretStoreError, unsupported_message):
operation()
@staticmethod
def root_path():
return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json"
if __name__ == "__main__":
unittest.main()
+183
View File
@@ -0,0 +1,183 @@
"""Offline integration: real graph preparation + orchestration, fake provider."""
import asyncio
import aiohttp as real_aiohttp
import importlib.util
import json
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
ROOT = Path(__file__).resolve().parents[1]
PKG = "dist_serverless_integration"
def module(name, **attrs):
value = types.ModuleType(name)
value.__dict__.update(attrs)
return value
def load_modules():
modules = {}
for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]:
modules[PKG + part] = module(PKG + part, __path__=[str(directory)])
ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn))
modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps))
modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None)
modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None)
modules[PKG + ".utils.config"] = module("config", load_config=lambda: {})
modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock())
modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master")
modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock())
modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), sync_worker_media=AsyncMock())
with patch.dict(sys.modules, modules):
spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py")
q = importlib.util.module_from_spec(spec)
spec.loader.exec_module(q)
# Keep imported package namespace only; restore shared Comfy runtime stub.
modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)})
return q, modules, ps
class IntegrationTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.q, modules, self.ps = load_modules()
self.patch = patch.dict(sys.modules, modules)
self.patch.start()
self.addCleanup(self.patch.stop)
self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30}
self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}}
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None)
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager)
sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[]))
self.prompt = {
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}},
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
}
async def test_real_compiler_slots_to_master_collector_without_callbacks(self):
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
self.assertEqual(result[2], 2)
self.q.dispatch_worker_prompt.assert_not_awaited()
self.manager.start.assert_awaited_once()
jobs = self.manager.start.call_args.args[2]
self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"])
self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2)
for job in jobs:
graph = job["workflow"]
self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values()))
self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values()))
self.assertNotIn("http://master", json.dumps(graph))
master_graph = self.q.queue_prompt_payload.call_args.args[0]
self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"])
async def test_compile_pack_provider_decode_and_collect_real_components(self):
"""Exercise the whole image path, mocking only the provider and Comfy queue."""
import base64
import io
import tempfile
import torch
from PIL import Image
from test_collector_list_inputs import _load_collector_module
sys.modules["aiohttp"] = real_aiohttp
for leaf in ("serverless_media", "serverless_runtime"):
name = PKG + ".api.orchestration." + leaf
spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py"))
real_module = importlib.util.module_from_spec(spec)
sys.modules[name] = real_module
spec.loader.exec_module(real_module)
runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"]
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
image_buffer = io.BytesIO()
Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG")
image_bytes = image_buffer.getvalue()
(root / "source.png").write_bytes(image_bytes)
paths = module("folder_paths", get_input_directory=lambda: directory,
get_output_directory=lambda: directory, get_temp_directory=lambda: directory,
get_annotated_filepath=lambda name: str(root / name))
submissions = []
class FakeProvider:
async def submit(self, graph, images, timeout_seconds):
self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage")
submissions.append((graph, images))
self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage")
assert self_image == images[0]["name"]
assert base64.b64decode(images[0]["image"]) == image_bytes
return "provider-job-" + str(len(submissions))
async def status(self, _job_id):
filename = self.prefix + "_00001_.png"
return {"status": "COMPLETED", "output": {"images": [{
"filename": filename, "artifact_path": filename, "artifact_type": "output",
"type": "base64", "data": base64.b64encode(image_bytes).decode(),
}]}}
manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None)
manager._client = AsyncMock(side_effect=lambda _: FakeProvider())
self.ps.distributed_serverless = manager
self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}}
collector_module = _load_collector_module()
collector_module.prompt_server = self.ps
sys.modules["aiohttp"] = real_aiohttp
with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}):
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
master_graph = self.q.queue_prompt_payload.call_args.args[0]
inputs = master_graph["2"]["inputs"]
images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute(
images=None, audio=None, multi_job_id=inputs["multi_job_id"],
enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True,
), timeout=3)
await asyncio.gather(*manager.tasks.values())
self.assertEqual(result[2], 2)
self.assertEqual(len(submissions), 2)
self.assertEqual(tuple(images.shape), (2, 3, 2, 3))
self.assertEqual(images.dtype, torch.float32)
self.assertTrue(torch.equal(images[0], images[1]))
self.assertEqual(float(images[0, 0, 0, 0]), 1.0)
self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"])
self.assertEqual(self.ps.distributed_pending_jobs, {})
self.q.dispatch_worker_prompt.assert_not_awaited()
async def test_audio_rejected_before_any_dispatch_or_master_queue(self):
self.prompt["2"]["inputs"]["audio"] = ["9", 0]
self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}}
with self.assertRaises(ValueError):
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
self.manager.start.assert_not_awaited()
self.q.dispatch_worker_prompt.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
async def test_disabled_pool_excluded(self):
self.pool["enabled"] = False
self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), [])
async def test_master_validation_failure_submits_no_paid_job(self):
self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph")
with self.assertRaisesRegex(ValueError, "invalid local graph"):
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
self.manager.start.assert_not_awaited()
async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self):
media_module = sys.modules[PKG + ".api.orchestration.serverless_media"]
media_module.prepare_images.return_value = [{
"name": "private-image-name.png",
"image": "x" * 10_000_001,
}]
with self.assertRaisesRegex(ValueError, "10,000,") as ctx:
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
self.assertNotIn("private-image-name.png", str(ctx.exception))
self.assertNotIn("x" * 100, str(ctx.exception))
self.q.dispatch_worker_prompt.assert_not_awaited()
self.q.queue_prompt_payload.assert_not_awaited()
self.manager.start.assert_not_awaited()
self.assertEqual(self.ps.distributed_pending_jobs, {})
+192
View File
@@ -0,0 +1,192 @@
"""Billing-safe lifecycle regressions: local completion is not remote completion."""
import asyncio
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
from test_serverless_runtime import runtime, Secrets
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
self.path = Path(self.temp.name) / "jobs.json"
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]}
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
self.client_patch.start()
self.addCleanup(self.client_patch.stop)
self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
self.decode_patch.start()
self.addCleanup(self.decode_patch.stop)
async def asyncTearDown(self):
await self.manager.cancel()
def second(self, **changes):
return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes))
def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"):
self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id,
status=status, created_at=1, updated_at=1)
async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self):
cancel_started, release_cancel = asyncio.Event(), asyncio.Event()
self.api.submit.side_effect = ["job-1", "job-2"]
self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}]
async def cancel(_job_id):
cancel_started.set()
await release_cancel.wait()
return {"status": "CANCELLED"}
self.api.cancel.side_effect = cancel
await self.manager.start("exec", "prompt", [self.prepared, self.second()])
try:
await asyncio.wait_for(cancel_started.wait(), 1)
await asyncio.sleep(0.01)
self.assertEqual(self.api.submit.await_count, 1)
finally:
release_cancel.set()
await self.manager.wait_for_execution("exec")
self.assertEqual(self.api.submit.await_count, 2)
async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self):
self.api.status.side_effect = runtime.RunpodError("poll failed")
self.api.cancel.return_value = {"status": "IN_PROGRESS"}
await self.manager.start("first", "prompt", [self.prepared])
await self.manager.wait_for_execution("first")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
self.assertTrue(self.manager.pool_is_active("pool"))
await self.manager.start("second", "prompt", [self.second()])
try:
await asyncio.sleep(0.01)
self.assertEqual(self.api.submit.await_count, 1)
self.api.status.side_effect = None
self.api.status.return_value = {"status": "COMPLETED", "output": {}}
await self.manager.reconcile(worker_id="pool")
await asyncio.wait_for(self.manager.wait_for_execution("second"), 1)
self.assertEqual(self.api.submit.await_count, 2)
finally:
self.api.cancel.return_value = {"status": "CANCELLED"}
await self.manager.cancel()
async def test_restarted_manager_reserves_known_unresolved_jobs(self):
self.retained("IN_PROGRESS")
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertFalse(self.manager.pool_is_active("other"))
await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)])
await self.manager.wait_for_execution("exec")
self.api.submit.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT")
async def test_recovery_closes_never_submitted_waiting_records(self):
self.retained("WAITING", job_id=None)
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertFalse(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
self.api.submit.assert_not_awaited()
async def test_recovery_keeps_interrupted_submission_unresolved(self):
self.retained("SUBMITTING", job_id=None)
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
await self.manager.reconcile(worker_id="pool")
self.api.status.assert_not_awaited()
self.api.submit.assert_not_awaited()
async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self):
self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown")
await self.manager.start("exec", "prompt", [self.prepared])
await self.manager.wait_for_execution("exec")
record = self.manager.ledger.all()[0]
self.assertTrue(self.manager.pool_is_active("pool"))
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False)
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True)
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
self.assertFalse(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED")
self.api.submit.assert_awaited_once()
async def test_cancel_already_completed_provider_job_is_terminal(self):
self.retained()
self.api.cancel.return_value = {"status": "COMPLETED"}
await self.manager.cancel(worker_id="pool")
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
self.assertFalse(self.manager.pool_is_active("pool"))
async def test_reconciliation_failure_never_releases_capacity(self):
self.retained()
self.api.status.side_effect = runtime.RunpodError("not reachable")
await self.manager.reconcile(worker_id="pool")
self.assertTrue(self.manager.pool_is_active("pool"))
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
async def test_cancel_before_first_task_step_finalizes_local_record(self):
await self.manager.start("exec", "prompt", [self.prepared])
await self.manager.cancel(worker_id="pool")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
self.assertFalse(self.server.distributed_pending_jobs["collector"].empty())
self.assertFalse(self.manager.pool_is_active("pool"))
self.api.submit.assert_not_awaited()
async def test_live_record_cannot_be_acknowledged_or_reconciled(self):
polling = asyncio.Event()
async def status(_job_id):
polling.set()
await asyncio.Event().wait()
self.api.status.side_effect = status
await self.manager.start("exec", "prompt", [self.prepared])
await asyncio.wait_for(polling.wait(), 1)
record = self.manager.ledger.all()[0]
with self.assertRaises(ValueError):
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
await self.manager.reconcile(worker_id="pool")
self.assertEqual(self.api.status.await_count, 1)
self.assertTrue(self.manager.pool_is_active("pool"))
async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self):
polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
async def status(_job_id):
polling.set()
await asyncio.Event().wait()
async def cancel(_job_id):
cancelling.set()
await release.wait()
return {"status": "CANCELLED"}
self.api.status.side_effect = status
self.api.cancel.side_effect = cancel
await self.manager.start("exec", "prompt", [self.prepared])
await polling.wait()
first = asyncio.create_task(self.manager.cancel(worker_id="pool"))
await cancelling.wait()
second = asyncio.create_task(self.manager.cancel(worker_id="pool"))
await asyncio.sleep(0)
release.set()
await asyncio.gather(first, second)
self.api.cancel.assert_awaited_once_with("job")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
async def test_partial_start_cleanup_does_not_leave_waiting_protection(self):
with self.assertRaises(ValueError):
await self.manager.start("exec", "prompt", [self.prepared, self.second(concurrency=2)])
await self.manager.cancel(execution_id="exec")
self.assertFalse(self.manager.pool_is_active("pool"))
self.api.submit.assert_not_awaited()
async def test_public_jobs_never_hides_unresolved_recovery_ids(self):
for index in range(22):
self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index)
for index in range(5):
self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index)
records = self.manager.public_jobs("pool")
self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)})
self.assertTrue(all(record["status"] == "UNKNOWN" for record in records))
+447
View File
@@ -0,0 +1,447 @@
import base64
import importlib.util
import io
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
from PIL import Image
ROOT = Path(__file__).resolve().parents[1]
MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py"
def load_module():
spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
media = load_module()
def load_runpod_module():
spec = importlib.util.spec_from_file_location(
"dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py"
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
runpod = load_runpod_module()
def png_bytes(mode="RGB", size=(2, 1), color=None):
if color is None:
color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30)
buffer = io.BytesIO()
Image.new(mode, size, color).save(buffer, format="PNG")
return buffer.getvalue()
class FolderPathsFixture:
def __init__(self, root):
self.input = root / "input"
self.output = root / "output"
self.temp = root / "temp"
for directory in (self.input, self.output, self.temp):
directory.mkdir(parents=True, exist_ok=True)
def get_input_directory(self):
return str(self.input)
def get_output_directory(self):
return str(self.output)
def get_temp_directory(self):
return str(self.temp)
def get_annotated_filepath(self, filename):
value = filename.replace("\\", "/")
# Match real ComfyUI folder_paths: root annotations are suffixes,
# not directory prefixes (output/foo.png is an input subdirectory).
for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)):
suffix = f" [{marker}]"
if value.endswith(suffix):
return str(root / value[:-len(suffix)])
return str(self.input / value)
class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.tempdir = tempfile.TemporaryDirectory()
self.paths = FolderPathsFixture(Path(self.tempdir.name))
self.folder_paths = types.ModuleType("folder_paths")
self.folder_paths.get_input_directory = self.paths.get_input_directory
self.folder_paths.get_output_directory = self.paths.get_output_directory
self.folder_paths.get_temp_directory = self.paths.get_temp_directory
self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath
self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths})
self.folder_patch.start()
self.addCleanup(self.folder_patch.stop)
self.addCleanup(self.tempdir.cleanup)
async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self):
payload = png_bytes()
(self.paths.input / "cat.png").write_bytes(payload)
workflow = {
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}},
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}},
"3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}},
}
packed = await media.prepare_images(workflow)
self.assertEqual(len(packed), 1)
self.assertEqual(set(packed[0]), {"name", "image"})
self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$")
self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload)
self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"])
self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"])
self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged")
self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable")
async def test_resolves_annotated_input_output_and_temp_sources(self):
sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp"))
workflow = {}
for index, (directory, filename) in enumerate(sources, 1):
(directory / filename).write_bytes(png_bytes())
root = directory.name
workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}}
packed = await media.prepare_images(workflow)
self.assertEqual(len(packed), 3)
self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()})
self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed})
async def test_directory_prefix_is_not_a_root_annotation(self):
(self.paths.input / "output").mkdir()
payload = png_bytes(color=(100, 0, 0))
(self.paths.input / "output" / "same.png").write_bytes(payload)
(self.paths.output / "same.png").write_bytes(png_bytes())
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}}
packed = await media.prepare_images(workflow)
self.assertEqual(base64.b64decode(packed[0]["image"]), payload)
async def test_rejects_unknown_root_annotation(self):
(self.paths.input / "cat.png").write_bytes(png_bytes())
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}}
with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"):
await media.prepare_images(workflow)
async def test_recognizes_custom_file_image_reference(self):
(self.paths.input / "source.bmp").write_bytes(png_bytes())
workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}}
packed = await media.prepare_images(workflow)
self.assertEqual(len(packed), 1)
self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"])
async def test_rejects_missing_and_dynamic_loadimage_values(self):
cases = (
{"1": {"class_type": "LoadImage", "inputs": {}}},
{"1": {"class_type": "LoadImage"}},
{"1": {"class_type": "LoadImage", "inputs": None}},
{"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}},
)
for workflow in cases:
with self.subTest(workflow=workflow):
with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"):
await media.prepare_images(workflow)
async def test_rejects_missing_source_and_unapproved_extension(self):
for value in ("absent.png", "animation.gif", "clip.mp4"):
with self.subTest(value=value):
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
with self.assertRaises(media.ServerlessMediaError):
await media.prepare_images(workflow)
async def test_rejects_unsupported_audio_or_video_references_without_reading(self):
for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")):
with self.subTest(key=key):
workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}}
with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"):
await media.prepare_images(workflow)
async def test_rejects_parent_traversal_and_absolute_outside_roots(self):
outside = Path(self.tempdir.name).parent / "outside.png"
outside.write_bytes(png_bytes())
self.addCleanup(lambda: outside.unlink(missing_ok=True))
for value in ("../outside.png", str(outside)):
with self.subTest(value=value):
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
with self.assertRaises(media.ServerlessMediaError):
await media.prepare_images(workflow)
async def test_rejects_symlink_escape_from_allowed_root(self):
outside = Path(self.tempdir.name).parent / "linked-outside.png"
outside.write_bytes(png_bytes())
self.addCleanup(lambda: outside.unlink(missing_ok=True))
(self.paths.input / "link.png").symlink_to(outside)
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}}
with self.assertRaises(media.ServerlessMediaError):
await media.prepare_images(workflow)
async def test_rejects_individual_and_aggregate_input_size_overages(self):
huge = self.paths.input / "large.jpg"
with huge.open("wb") as handle:
handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1)
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}}
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
await media.prepare_images(workflow)
huge.unlink()
first = self.paths.input / "first.png"
second = self.paths.input / "second.png"
first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
workflow = {
"1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}},
"2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}},
}
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
await media.prepare_images(workflow)
async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self):
raw_size = 7 * 1024 * 1024 + 1
source = self.paths.input / "large-input.png"
source.write_bytes(b"x" * raw_size)
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}}
packed = await media.prepare_images(workflow)
request_body = runpod.serialize_request(workflow, packed, 30)
self.assertGreater(raw_size, 7 * 1024 * 1024)
self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size)
self.assertLessEqual(len(request_body), 10_000_000)
class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase):
def manifest(self, prefix="render"):
return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}]
def artifact(self, filename, payload=None, **extra):
entry = {
"filename": filename,
"artifact_type": "output",
"artifact_path": f"/safe/output/{filename}",
"subfolder": "",
"type": "base64",
"data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"),
}
entry.update(extra)
return entry
async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self):
output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]}
result = await media.collect_images(output, self.manifest())
self.assertEqual(list(result), ["job-1"])
self.assertEqual(len(result["job-1"]), 2)
self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"]))
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255)
async def test_groups_images_under_manifest_job_ids(self):
manifest = [
self.manifest("first")[0],
{"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"},
]
output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]}
result = await media.collect_images(output, manifest)
self.assertEqual(set(result), {"job-1", "job-2"})
self.assertEqual(len(result["job-1"]), 1)
self.assertEqual(len(result["job-2"]), 1)
async def test_preserves_rgba_alpha(self):
output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]}
result = await media.collect_images(output, self.manifest())
self.assertEqual(result["job-1"][0].shape, (1, 2, 4))
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255)
async def test_signed_url_is_downloaded_only_through_injected_downloader(self):
output = {"images": [{
"filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url",
"url": "https://storage.example.test/signed",
}]}
fetch = AsyncMock(return_value=png_bytes())
result = await media.collect_images(output, self.manifest(), download=fetch)
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
fetch.assert_awaited_once_with("https://storage.example.test/signed")
async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self):
output = {"images": [
self.artifact("preview_00001_.png"),
self.artifact("render_00001_.png"),
{"filename": "mesh.obj", "type": "url", "url": "not a fetch target"},
]}
result = await media.collect_images(output, self.manifest())
self.assertEqual(len(result["job-1"]), 1)
with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"):
await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest())
async def test_accepts_compiler_length_flat_safe_prefix(self):
prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64
manifest = self.manifest(prefix)
output = {"images": [self.artifact(f"{prefix}_00001_.png")]}
result = await media.collect_images(output, manifest)
self.assertEqual(len(result["job-1"]), 1)
async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self):
duplicate = self.manifest() * 2
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images({"images": []}, duplicate)
unsafe = self.manifest("../render")
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images({"images": []}, unsafe)
async def test_rejects_duplicate_canonical_artifact_identity(self):
item = self.artifact("render_00001_.png")
with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"):
await media.collect_images({"images": [item, dict(item)]}, self.manifest())
async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self):
wrong_type = self.artifact("render_00001_.png", artifact_type="temp")
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images({"images": [wrong_type]}, self.manifest())
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest())
async def test_rejects_provider_errors_even_when_job_completed(self):
for output in (
{"error": "raw provider secret", "images": []},
{"errors": ["raw provider secret"], "images": []},
{"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}},
):
with self.subTest(output=output):
with self.assertRaises(media.ServerlessMediaError) as ctx:
await media.collect_images(output, self.manifest())
self.assertNotIn("raw provider secret", str(ctx.exception))
async def test_rejects_invalid_base64_and_malformed_image(self):
for item in (
self.artifact("render_00001_.png", data="!not base64!"),
self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()),
):
with self.subTest(item=item):
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images({"images": [item]}, self.manifest())
async def test_validates_shape_consistency_within_collector(self):
output = {"images": [
self.artifact("render_00001_.png", png_bytes(size=(2, 1))),
self.artifact("render_00002_.png", png_bytes(size=(3, 1))),
]}
with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"):
await media.collect_images(output, self.manifest())
async def test_bounds_output_file_count_and_bytes(self):
too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)}
with self.assertRaisesRegex(media.ServerlessMediaError, "files"):
await media.collect_images(too_many, self.manifest())
with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4):
output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]}
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
await media.collect_images(output, self.manifest())
async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self):
output = {"images": [
self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()),
self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()),
]}
with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3):
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
await media.collect_images(output, self.manifest())
async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self):
for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"):
with self.subTest(url=url):
fetch = AsyncMock(return_value=png_bytes())
output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]}
with self.assertRaises(media.ServerlessMediaError):
await media.collect_images(output, self.manifest(), download=fetch)
fetch.assert_not_awaited()
async def test_public_resolver_rejects_non_global_dns_answers(self):
resolver = media._PublicResolver()
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))])
self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1)
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))])
with self.assertRaises(OSError):
await resolver.resolve("cdn.example.test", 443)
async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self):
class Content:
def __init__(self, chunks):
self.chunks = iter(chunks)
async def iter_chunked(self, _size):
for chunk in self.chunks:
yield chunk
class Response:
def __init__(self, status, headers=None, chunks=()):
self.status, self.headers = status, headers or {}
self.content = Content(chunks)
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return False
class Session:
def __init__(self, responses):
self.responses, self.calls = list(responses), []
async def __aenter__(self):
return self
async def __aexit__(self, *_args):
return False
def get(self, url, **kwargs):
self.calls.append((url, kwargs))
return self.responses.pop(0)
session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])])
with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session:
payload = await media.download_public_artifact("https://storage.example.test/signed")
self.assertEqual(payload, b"abcdef")
self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"])
self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls))
self.assertFalse(create_session.call_args.kwargs["trust_env"])
self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver)
self.assertNotIn("Authorization", session.calls[0][1].get("headers", {}))
await create_session.call_args.kwargs["connector"].close()
blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})])
with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session:
with self.assertRaises(media.ServerlessMediaError):
await media.download_public_artifact("https://storage.example.test/signed")
self.assertEqual(len(blocked.calls), 1)
await blocked_session.call_args.kwargs["connector"].close()
if __name__ == "__main__":
unittest.main()
+392
View File
@@ -0,0 +1,392 @@
import importlib.util
import json
import sys
import types
from pathlib import Path
import pytest
MODULE_PATH = (
Path(__file__).resolve().parents[1]
/ "api"
/ "orchestration"
/ "serverless_prompt.py"
)
_spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH)
assert _spec is not None and _spec.loader is not None
_serverless_prompt = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_serverless_prompt)
compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt
def _prepared_prompt():
return {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}},
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
"3": {
"class_type": "DistributedSeed",
"inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"},
},
"4": {
"class_type": "DistributedCollector",
"inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True},
},
"5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
"7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}},
"8": {"class_type": "UnrelatedNode", "inputs": {}},
}
def _run(prompt):
return compile_serverless_prompt(prompt)
def _load_prompt_transform():
module_path = MODULE_PATH.with_name("prompt_transform.py")
package_name = "serverless_prompt_transform_testpkg"
root_pkg = types.ModuleType(package_name)
root_pkg.__path__ = []
sys.modules[package_name] = root_pkg
for suffix in ("api", "api.orchestration", "utils"):
pkg = types.ModuleType(f"{package_name}.{suffix}")
pkg.__path__ = []
sys.modules[f"{package_name}.{suffix}"] = pkg
logging_module = types.ModuleType(f"{package_name}.utils.logging")
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
setattr(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}.api.orchestration.prompt_transform", module_path
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def test_compiles_prompt_prepared_by_existing_worker_helpers():
prompt_transform = _load_prompt_transform()
original = {
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
"3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
"5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
}
pruned = prompt_transform.prune_prompt_for_worker(original)
prompt_index = prompt_transform.PromptIndex(pruned)
prepared = prompt_transform.apply_participant_overrides(
pruned,
participant_id="worker-a",
enabled_worker_ids=["worker-a"],
job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"),
master_url="http://master.invalid",
delegate_master=False,
prompt_index=prompt_index,
)
graph, manifest = _run(prepared)
assert graph["2"]["inputs"]["seed"] == 101
assert set(graph) == {"1", "2", "7"}
assert graph["7"]["class_type"] == "SaveImage"
assert manifest[0]["job_id"] == "run_4"
assert manifest[0]["collector_id"] == "4"
def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream():
original = _prepared_prompt()
graph, manifest = _run(original)
assert manifest == [
{
"collector_id": "4",
"job_id": "job/42",
"node_id": "9",
"kind": "image",
"prefix": graph["9"]["inputs"]["filename_prefix"],
}
]
prefix = manifest[0]["prefix"]
assert graph["9"] == {
"class_type": "SaveImage",
"inputs": {"images": ["2", 0], "filename_prefix": prefix},
"_meta": {"title": "Serverless output: collector 4 (image)"},
}
assert "/" not in prefix and "\\" not in prefix
assert graph["2"]["inputs"]["seed"] == 41
assert set(graph) == {"1", "2", "9"}
assert all("Distributed" not in node["class_type"] for node in graph.values())
assert original["3"]["inputs"]["seed"] == 40
assert original["4"]["inputs"]["images"] == ["2", 0]
def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts():
prompt = _prepared_prompt()
prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe")
first_graph, first_manifest = _run(prompt)
second_graph, second_manifest = _run(prompt)
assert first_manifest == second_manifest
assert first_graph == second_graph
prefix = first_manifest[0]["prefix"]
assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"]
assert not any(char in prefix for char in ("/", "\\", ":"))
assert first_manifest[0]["job_id"] == "../run/a"
def test_compiles_multiple_disconnected_collectors_in_stable_order():
prompt = {
"10": {"class_type": "LoadImage", "inputs": {}},
"20": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}},
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}},
"30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}},
}
graph, manifest = _run(prompt)
assert [item["collector_id"] for item in manifest] == ["1", "2"]
assert [item["node_id"] for item in manifest] == ["31", "32"]
assert graph["31"]["inputs"]["images"] == ["10", 0]
assert graph["32"]["inputs"]["images"] == ["20", 0]
assert set(graph) == {"10", "20", "31", "32"}
def test_rejects_audio_input_instead_of_emitting_audio_saver():
prompt = {
"1": {"class_type": "LoadAudio", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
}
with pytest.raises(ValueError, match="audio.*unsupported"):
_run(prompt)
def test_rejects_collector_without_image_input():
with pytest.raises(ValueError, match="collector.*image"):
_run({"1": {"class_type": "DistributedCollector", "inputs": {}}})
@pytest.mark.parametrize(
"class_type",
["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"],
)
def test_rejects_unsupported_distributed_and_batch_nodes(class_type):
prompt = {
"1": {"class_type": class_type, "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
"3": {"class_type": "LoadImage", "inputs": {}},
}
with pytest.raises(ValueError, match="unsupported"):
_run(prompt)
def test_rejects_collector_to_collector_dependency_through_processing_node():
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
"3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
}
with pytest.raises(ValueError, match="collector-to-collector"):
_run(prompt)
@pytest.mark.parametrize(
"link",
[["missing", 0], ["1", -1], ["1", "0"], ["1", True]],
)
def test_rejects_malformed_api_links(link):
prompt = {
"1": {"class_type": "LoadImage", "inputs": {}},
"2": {"class_type": "DistributedCollector", "inputs": {"images": link}},
}
with pytest.raises(ValueError, match="link"):
_run(prompt)
def test_rejects_unexpected_output_slot_from_removed_distributed_node():
prompt = {
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}},
"2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}},
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
}
with pytest.raises(ValueError, match="output slot"):
_run(prompt)
def test_distributed_seed_matches_worker_offset_and_master_passthrough():
prompt = {
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}},
"2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}},
"3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
}
graph, _manifest = _run(prompt)
assert graph["3"]["inputs"] == {"seed": 12, "steps": 7}
assert "1" not in graph and "2" not in graph
def test_distributed_value_uses_one_based_worker_values_and_type_coercion():
prompt = {
"1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}},
"2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}},
"3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
}
graph, _manifest = _run(prompt)
assert graph["3"]["inputs"] == {"steps": 8, "seed": 3}
assert "1" not in graph and "2" not in graph
def test_resolves_stock_primitive_links_for_distributed_seed_and_value():
prompt = {
"1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}},
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
"3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}},
"4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}},
"5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}},
"6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}},
"7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}},
}
graph, _manifest = _run(prompt)
assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75}
assert not {"1", "2", "3", "4", "5"}.intersection(graph)
def test_rejects_dynamic_linked_seed_or_value_without_executing_it():
prompt = {
"1": {"class_type": "ArbitraryNode", "inputs": {}},
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
"3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
}
with pytest.raises(ValueError, match="literal or supported Primitive"):
_run(prompt)
def test_primitive_resolution_detects_cycles():
prompt = {
"1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}},
"2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}},
"3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}},
"4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}},
"5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}},
}
with pytest.raises(ValueError, match="cycle"):
_run(prompt)
def test_model_name_trivial_text_passthrough_is_inlined():
prompt = {
"1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}},
"2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}},
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
}
graph, _manifest = _run(prompt)
assert graph["2"]["inputs"]["model_name"] == "model.safetensors"
assert "1" not in graph
def test_rejects_dynamic_linked_model_name_text():
prompt = {
"1": {"class_type": "ArbitraryNode", "inputs": {}},
"2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}},
"3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}},
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
}
with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"):
_run(prompt)
def test_resolves_boolean_and_float_primitive_links():
prompt = {
"1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}},
"2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}},
"3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}},
"4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}},
"5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}},
"6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}},
}
graph, _manifest = _run(prompt)
assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5}
assert not {"1", "2", "3", "4"}.intersection(graph)
def test_compiled_graph_matches_golden_prompt_and_manifest():
prompt = _prepared_prompt()
prompt["3"]["inputs"]["seed"] = ["10", 0]
prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}}
prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}
prompt["12"] = {
"class_type": "DistributedValue",
"inputs": {
"default_value": ["11", 0],
"worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}),
"is_worker": True,
"worker_id": "worker_0",
},
}
prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}}
prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}}
prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0])
graph, manifest = _run(prompt)
expected_prefix = (
"serverless_job_42_worker_0_4_image_"
"1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244"
)
assert list(graph) == ["1", "2", "15"]
assert graph == {
"1": {
"class_type": "CheckpointLoaderSimple",
"inputs": {},
"_meta": {"title": "source"},
},
"2": {
"class_type": "KSampler",
"inputs": {
"model": ["1", 0],
"seed": 41,
"cfg": 2.5,
"model_name": "model.safetensors",
},
},
"15": {
"class_type": "SaveImage",
"inputs": {"images": ["2", 0], "filename_prefix": expected_prefix},
"_meta": {"title": "Serverless output: collector 4 (image)"},
},
}
assert manifest == [
{
"collector_id": "4",
"job_id": "job/42",
"node_id": "15",
"kind": "image",
"prefix": expected_prefix,
}
]
+608
View File
@@ -0,0 +1,608 @@
import copy
import importlib.util
import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import patch
class _FakeResponse:
def __init__(self, payload, status=200):
self.payload = copy.deepcopy(payload)
self.status = status
class _FakeRequest:
def __init__(self, payload=None, error=None):
self.payload = payload
self.error = error
async def json(self):
if self.error is not None:
raise self.error
return self.payload
class _Routes:
def get(self, _path):
return lambda handler: handler
def post(self, _path):
return lambda handler: handler
def _load_modules():
repo = Path(__file__).resolve().parents[1]
package_name = "dist_serverless_routes_testpkg"
for name in list(sys.modules):
if name == package_name or name.startswith(f"{package_name}."):
del sys.modules[name]
root = types.ModuleType(package_name)
root.__path__ = [str(repo)]
sys.modules[package_name] = root
api_pkg = types.ModuleType(f"{package_name}.api")
api_pkg.__path__ = [str(repo / "api")]
sys.modules[f"{package_name}.api"] = api_pkg
utils_pkg = types.ModuleType(f"{package_name}.utils")
utils_pkg.__path__ = [str(repo / "utils")]
sys.modules[f"{package_name}.utils"] = utils_pkg
orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")]
sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
logging_module = types.ModuleType(f"{package_name}.utils.logging")
logging_module.log = lambda *_args, **_kwargs: None
logging_module.debug_log = lambda *_args, **_kwargs: None
sys.modules[logging_module.__name__] = logging_module
constants_module = types.ModuleType(f"{package_name}.utils.constants")
constants_module.HEARTBEAT_TIMEOUT = 30
sys.modules[constants_module.__name__] = constants_module
def load_file(module_name, path):
spec = importlib.util.spec_from_file_location(module_name, path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
spec.loader.exec_module(module)
return module
config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py")
serverless_config = load_file(
f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py"
)
runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime")
runtime.secret_store = None
runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
runtime.get_secret_store = lambda: runtime.secret_store
runtime.get_manager = lambda: runtime.manager
sys.modules[runtime.__name__] = runtime
import aiohttp # Ensure the real dependency is loaded; never replace global aiohttp.
previous_server = sys.modules.get("server")
server = types.ModuleType("server")
server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
sys.modules["server"] = server
network = types.ModuleType(f"{package_name}.utils.network")
async def handle_api_error(_request, error, status=500):
return _FakeResponse({"status": "error", "message": str(error)}, status)
network.handle_api_error = handle_api_error
network.normalize_host = lambda value: value
sys.modules[network.__name__] = network
routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py")
routes.web = types.SimpleNamespace(
json_response=lambda payload, status=200: _FakeResponse(payload, status)
)
if previous_server is None:
sys.modules.pop("server", None)
else:
sys.modules["server"] = previous_server
return config_module, serverless_config, runtime, routes
config_module, serverless_config, runtime_module, config_routes = _load_modules()
class _RunningTask:
def done(self):
return False
class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.tempdir = tempfile.TemporaryDirectory()
self.root = Path(self.tempdir.name)
config_module.CONFIG_FILE = str(self.root / "gpu_config.json")
config_module.invalidate_config_cache()
self.config = config_module.get_default_config()
self.config["workers"] = []
self.assertTrue(config_module.save_config(self.config))
self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json")
runtime_module.secret_store = self.secrets
runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
def tearDown(self):
config_module.invalidate_config_cache()
self.tempdir.cleanup()
def read_config_file(self):
return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8"))
class ServerlessConfigRouteTests(_ConfigRouteFixture):
async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self):
self.secrets.set("pool-1", "private-api-key")
self.config["workers"] = [
{"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1",
"api_key": "inline-api-key", "name": "GPU"}
]
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
config_module.invalidate_config_cache()
response = await config_routes.get_config_endpoint(_FakeRequest())
self.assertEqual(response.status, 200)
worker = response.payload["workers"][0]
self.assertTrue(worker["api_key_configured"])
self.assertNotIn("api_key", worker)
self.assertNotIn("private-api-key", json.dumps(response.payload))
self.assertNotIn("inline-api-key", json.dumps(response.payload))
async def test_get_config_redacts_normalized_secret_field_aliases(self):
aliases = {
"apiKey": "camel-api-key",
"APIKEY": "upper-api-key",
"clearApiKey": "camel-clear-key",
"CLEAR_API_KEY": "upper-clear-key",
"runpod_CREDENTIAL": "provider-credential",
}
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
**aliases,
"nested": {"credential_note": "nested-credential"},
}]
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
config_module.invalidate_config_cache()
response = await config_routes.get_config_endpoint(_FakeRequest())
self.assertEqual(response.status, 200)
worker = response.payload["workers"][0]
for alias, value in aliases.items():
self.assertNotIn(alias, worker)
self.assertNotIn(value, json.dumps(response.payload))
self.assertNotIn("credential_note", worker["nested"])
self.assertNotIn("nested-credential", json.dumps(response.payload))
async def test_bulk_config_rejects_normalized_secret_field_aliases(self):
for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"):
with self.subTest(secret_field=secret_field):
secret = f"private-{secret_field}"
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{"id": "pool-1", "type": "local", secret_field: secret}]
}))
self.assertEqual(response.status, 400)
self.assertNotIn(secret, json.dumps(response.payload))
self.assertEqual(self.read_config_file()["workers"], [])
async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self):
secret = "runpod-secret-never-returned"
payload = {
"worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2,
"timeout_seconds": 600, "enabled": True, "api_key": secret,
}
response = await config_routes.update_worker_endpoint(_FakeRequest(payload))
self.assertEqual(response.status, 200)
self.assertTrue(response.payload["api_key_configured"])
self.assertNotIn(secret, json.dumps(response.payload))
saved = self.read_config_file()["workers"][0]
self.assertEqual(saved["endpoint_id"], "ep-123")
self.assertEqual(saved["jobs_per_run"], 3)
self.assertEqual(saved["concurrency"], 2)
self.assertEqual(saved["timeout_seconds"], 600)
for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"):
self.assertNotIn(obsolete, saved)
self.assertEqual(self.secrets.get("pool-1"), secret)
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
self.secrets.set("pool-1", "saved-secret")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool",
"endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3,
"timeout_seconds": 900, "enabled": False,
}]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_worker_endpoint(
_FakeRequest({"worker_id": "pool-1", "enabled": True})
)
worker = self.read_config_file()["workers"][0]
self.assertTrue(worker["enabled"])
self.assertEqual(worker["endpoint_id"], "ep-123")
self.assertEqual(worker["jobs_per_run"], 4)
self.assertEqual(worker["concurrency"], 3)
self.assertEqual(worker["timeout_seconds"], 900)
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
self.assertTrue(response.payload["api_key_configured"])
async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self):
secret = "must-not-leak"
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "concurrency": True, "api_key": secret,
}))
self.assertEqual(response.status, 400)
self.assertNotIn(secret, json.dumps(response.payload))
self.assertEqual(self.read_config_file()["workers"], [])
self.assertIsNone(self.secrets.get("pool-1"))
async def test_json_decode_error_is_generic_and_does_not_echo_body(self):
secret = "body-secret-value"
response = await config_routes.update_worker_endpoint(
_FakeRequest(error=ValueError(f"bad json near {secret}"))
)
self.assertEqual(response.status, 400)
self.assertNotIn(secret, json.dumps(response.payload))
async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self):
secret = "bulk-secret-never-write"
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{
"id": "pool-1", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "nested": {"api_key": secret},
}]
}))
self.assertEqual(response.status, 400)
self.assertNotIn(secret, json.dumps(response.payload))
self.assertEqual(self.read_config_file()["workers"], [])
self.assertFalse(self.secrets.path.exists())
async def test_bulk_update_validates_and_returns_redacted_serverless_config(self):
self.secrets.set("pool-1", "saved-secret")
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{
"id": "pool-1", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "name": "Pool", "enabled": True,
"host": "obsolete-host", "port": 8188, "cuda_device": 0,
}]
}))
self.assertEqual(response.status, 200)
self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"])
saved = self.read_config_file()["workers"][0]
for obsolete in ("host", "port", "cuda_device"):
self.assertNotIn(obsolete, saved)
self.assertNotIn("api_key", response.payload["config"]["workers"][0])
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self):
self.secrets.set("pool-1", "bulk-clear-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{
"id": "pool-1", "type": "local", "port": 8188,
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
}]
}))
self.assertEqual(response.status, 200)
self.assertIsNone(self.secrets.get("pool-1"))
saved = self.read_config_file()["workers"][0]
for field in ("provider", "endpoint_id", "api_key_configured"):
self.assertNotIn(field, saved)
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
existing = {"id": "local-1", "type": "local", "port": 8188}
self.config["workers"] = [existing]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
}))
self.assertEqual(response.status, 400)
self.assertEqual(self.read_config_file()["workers"], [existing])
async def test_delete_serverless_worker_clears_saved_key(self):
self.secrets.set("pool-1", "delete-me")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123"
}]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"], [])
self.assertIsNone(self.secrets.get("pool-1"))
async def test_type_conversion_clears_secret_and_removes_serverless_fields(self):
self.secrets.set("pool-1", "old-key")
self.config["workers"] = [{
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
}]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188,
"cuda_device": 0, "enabled": False,
}))
self.assertEqual(response.status, 200)
self.assertIsNone(self.secrets.get("pool-1"))
worker = self.read_config_file()["workers"][0]
self.assertEqual(worker["type"], "local")
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
self.assertNotIn(serverless_field, worker)
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
self.secrets.set("pool-1", "saved-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
"enabled": False,
}]
self.assertTrue(config_module.save_config(self.config))
runtime_module.manager = types.SimpleNamespace(
tasks={"run-1": _RunningTask()},
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
pool_is_active=lambda worker_id: worker_id == "pool-1",
)
enabled = await config_routes.update_worker_endpoint(
_FakeRequest({"worker_id": "pool-1", "enabled": True})
)
blocked = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "endpoint_id": "ep-456",
}))
self.assertEqual(enabled.status, 200)
self.assertEqual(blocked.status, 409)
self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self):
self.secrets.set("pool-1", "saved-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
keep = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "api_key": "",
}))
clear = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "clear_api_key": True,
}))
self.assertTrue(keep.payload["api_key_configured"])
self.assertFalse(clear.payload["api_key_configured"])
self.assertIsNone(self.secrets.get("pool-1"))
async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self):
with self.subTest(case="inline key"):
config = config_module.get_default_config()
config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}]
self.assertFalse(config_module.save_config(config))
with self.subTest(case="invalid serverless record"):
config = config_module.get_default_config()
config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
self.assertFalse(config_module.save_config(config))
async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self):
self.secrets.set("pool-1", "saved-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
runtime_module.manager = types.SimpleNamespace(
tasks={"run-1": _RunningTask()},
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
pool_is_active=lambda worker_id: worker_id == "pool-1",
)
credential_change = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "api_key": "new-key",
}))
type_change = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0,
}))
deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(credential_change.status, 409)
self.assertEqual(type_change.status, 409)
self.assertEqual(deletion.status, 409)
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
async def test_secret_write_is_rolled_back_when_config_persistence_fails(self):
secret = "write-then-rollback"
with patch.object(config_module, "save_config", return_value=False):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
"api_key": secret,
}))
self.assertEqual(response.status, 500)
self.assertIsNone(self.secrets.get("pool-1"))
self.assertEqual(self.read_config_file()["workers"], [])
self.assertEqual(config_module.load_config()["workers"], [])
async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self):
self.secrets.set("pool-1", "bulk-restore-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
with patch.object(config_module, "save_config", return_value=False):
response = await config_routes.update_config_endpoint(_FakeRequest({
"workers": [{"id": "pool-1", "type": "local", "port": 8188}]
}))
self.assertEqual(response.status, 500)
self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key")
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self):
self.secrets.set("pool-1", "prior-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
with patch.object(config_module, "save_config", return_value=False):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "api_key": "replacement-key",
}))
self.assertEqual(response.status, 500)
self.assertEqual(self.secrets.get("pool-1"), "prior-key")
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self):
self.secrets.set("pool-1", "delete-restore-key")
self.config["workers"] = [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
}]
self.assertTrue(config_module.save_config(self.config))
with patch.object(config_module, "save_config", return_value=False):
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(response.status, 500)
self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key")
self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1")
async def test_regular_worker_delete_does_not_require_posix_secret_storage(self):
self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}]
self.assertTrue(config_module.save_config(self.config))
with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")):
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"], [])
async def test_enabling_serverless_without_key_is_rejected(self):
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True,
}))
self.assertEqual(response.status, 400)
self.assertEqual(self.read_config_file()["workers"], [])
async def test_bulk_concurrency_change_is_blocked_during_active_run(self):
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
self.assertTrue(config_module.save_config(self.config))
runtime_module.manager = types.SimpleNamespace(
tasks={"run-1": _RunningTask()},
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
pool_is_active=lambda worker_id: worker_id == "pool-1",
)
updated = copy.deepcopy(self.config["workers"])
updated[0]["concurrency"] = 2
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
self.assertEqual(response.status, 409)
self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1)
async def test_non_serverless_update_behavior_is_preserved(self):
self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}]
self.assertTrue(config_module.save_config(self.config))
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "local-1", "name": "Renamed", "enabled": True,
}))
self.assertEqual(response.status, 200)
worker = self.read_config_file()["workers"][0]
self.assertEqual(worker["name"], "Renamed")
self.assertTrue(worker["enabled"])
async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self):
raw = {
"workers": [{
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
"host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old",
}]
}
normalized = config_module._merge_with_defaults(raw, config_module.get_default_config())
worker = normalized["workers"][0]
for obsolete in ("host", "port", "cuda_device", "extra_args"):
self.assertNotIn(obsolete, worker)
class PersistedPoolGuardTests(_ConfigRouteFixture):
"""Use the actual manager/ledger across restart, not a mocked idle result."""
def setUp(self):
super().setUp()
from test_serverless_runtime import runtime
from unittest.mock import AsyncMock
self.runtime = runtime
self.session_factory = AsyncMock()
self.ledger_path = self.root / "jobs.json"
self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
runtime_module.manager = self.manager
self.secrets.set("pool-1", "saved-key")
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
self.assertTrue(config_module.save_config(self.config))
async def test_retained_jobs_block_destructive_changes_after_restart(self):
original = self.read_config_file()
for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"):
for restart in (False, True):
with self.subTest(status=status, restart=restart):
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id")
if restart:
self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
runtime_module.manager = self.manager
else:
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"concurrency": 2}, {"type": "local", "port": 8188}):
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
self.assertEqual(response.status, 409, fields)
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(response.status, 409)
worker = original["workers"][0]
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, concurrency=2)]):
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
self.assertEqual(response.status, 409, workers)
self.assertEqual(self.read_config_file(), original)
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
self.session_factory.assert_not_awaited()
async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self):
self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN")
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused")
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True}))
self.assertEqual(response.status, 409)
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
async def test_terminal_or_other_pool_records_do_not_block_deletion(self):
self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED")
self.manager.ledger.update("other", worker_id="other", status="UNKNOWN")
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
self.assertEqual(response.status, 200)
self.assertIsNone(self.secrets.get("pool-1"))
if __name__ == "__main__":
unittest.main()
+143
View File
@@ -0,0 +1,143 @@
import asyncio
import importlib.util
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
ROOT = Path(__file__).resolve().parents[1]
PACKAGE = "dist_serverless_runtime_tests"
for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]:
module = types.ModuleType(name)
module.__path__ = [str(path)]
sys.modules[name] = module
def load(name):
spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py"))
assert spec and spec.loader
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
runtime = load("serverless_runtime")
class Secrets:
def get(self, worker_id):
return "private-test-key"
class RuntimeTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.temp = tempfile.TemporaryDirectory()
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
self.manager.poll_seconds = 0
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
async def asyncTearDown(self):
await self.manager.cancel()
self.temp.cleanup()
async def test_job_id_persisted_before_status_and_results_injected(self):
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock())
async def status(job_id):
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
return {"status": "COMPLETED", "output": {"images": []}}
api.status = status
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
self.assertEqual(api.submit.await_count, 1)
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
self.assertEqual(item["worker_id"], "pool:1")
self.assertEqual(item["tensor"], "tensor")
self.assertTrue(item["is_last"])
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
async def test_output_error_marks_failed_and_delivers_explicit_error(self):
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock())
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
record = self.manager.ledger.all()[0]
self.assertEqual(record["status"], "FAILED")
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
self.assertIn("serverless_error", item)
self.assertNotIn("private-test-key", str(record) + str(item))
async def test_ambiguous_submit_never_retries(self):
api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock())
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
self.assertEqual(api.submit.await_count, 1)
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
async def test_cancel_targets_owned_job_only(self):
polled = asyncio.Event()
async def status(_job):
polled.set()
await asyncio.Event().wait()
api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"}))
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await polled.wait()
await self.manager.cancel(worker_id="other-pool")
api.cancel.assert_not_awaited()
await self.manager.cancel(worker_id="pool")
api.cancel.assert_awaited_once_with("owned-job")
async def test_per_pool_concurrency_bounds_submission(self):
gate = asyncio.Event()
reached = asyncio.Event()
count = 0
async def submit(*args):
nonlocal count
count += 1
reached.set()
return f"provider-{count}"
async def status(_job):
await gate.wait()
return {"status": "COMPLETED", "output": {"images": []}}
api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
second = dict(self.prepared, worker=dict(self.worker, id="pool:2"))
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
await self.manager.start("execution", "prompt", [self.prepared, second])
await reached.wait()
await asyncio.sleep(0)
self.assertEqual(count, 1)
gate.set()
await self.manager.wait_for_execution("execution")
self.assertEqual(count, 2)
async def test_cancel_persisted_job_after_master_restart(self):
self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1)
api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"}))
with patch.object(runtime, "RunpodClient", return_value=api) as client:
await self.manager.cancel(worker_id="pool")
api.cancel.assert_awaited_once_with("old-job")
self.assertEqual(client.call_args.args[1], "old-endpoint")
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
async def test_master_failure_cancels_job_without_partial_success(self):
self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}})
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
api.submit.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
async def test_deadline_includes_local_slot_wait(self):
self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
# A blocked slot expires without a billable submission.
self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0))
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
api.submit.assert_not_awaited()
self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT")
+98
View File
@@ -0,0 +1,98 @@
import importlib.util
import json
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch
ROOT = Path(__file__).resolve().parents[1]
PKG = "dist_serverless_status_tests"
class StatusTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
def mod(name, **kwargs):
result = types.ModuleType(name)
result.__dict__.update(kwargs)
return result
self.manager = types.SimpleNamespace(health=AsyncMock(return_value={"ok": True, "queueRemaining": 0, "message": "Ready (scales from zero)"}), public_jobs=lambda _: [], cancel=AsyncMock(return_value=1))
self.pool = {"id": "pool", "type": "serverless", "endpoint_id": "ep"}
modules = {PKG: mod(PKG, __path__=[str(ROOT)]), PKG + ".api": mod("api", __path__=[str(ROOT/'api')]),
PKG + ".utils": mod("utils", __path__=[]),
PKG + ".utils.config": mod("config", load_config=lambda: {"workers": [self.pool]}),
PKG + ".api.orchestration": mod("orchestration", __path__=[str(ROOT/'api/orchestration')]),
PKG + ".api.orchestration.serverless_runtime": mod("runtime", get_manager=lambda: self.manager),
"server": mod("server", PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda f: f, post=lambda _: lambda f: f))))}
with patch.dict(sys.modules, modules):
spec = importlib.util.spec_from_file_location(PKG+'.api.serverless_routes', ROOT/'api/serverless_routes.py')
self.routes = importlib.util.module_from_spec(spec)
spec.loader.exec_module(self.routes)
async def test_health_only_never_submits_and_supports_scale_to_zero(self):
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
result = await self.routes.test_serverless_endpoint(request)
self.assertEqual(result.status, 200)
self.assertTrue(json.loads(result.text)["ok"])
self.manager.health.assert_awaited_once_with(self.pool, fresh=True)
self.manager.cancel.assert_not_awaited()
async def test_unknown_pool_cannot_cancel_other_jobs(self):
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "other"}))
result = await self.routes.cancel_serverless_endpoint(request)
self.assertEqual(result.status, 404)
self.manager.cancel.assert_not_awaited()
async def test_transport_error_does_not_leak_credentials(self):
self.manager.health.side_effect = RuntimeError("Bearer secret-token https://signed.invalid/?token=value")
request = types.SimpleNamespace(match_info={"worker_id": "pool"})
result = await self.routes.serverless_status_endpoint(request)
self.assertFalse(json.loads(result.text)["ok"])
self.assertNotIn("secret-token", result.text)
self.assertNotIn("signed.invalid", result.text)
async def test_cancel_uses_only_requested_pool(self):
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
result = await self.routes.cancel_serverless_endpoint(request)
self.assertEqual(result.status, 200)
self.manager.cancel.assert_awaited_once_with(worker_id="pool")
async def test_reconcile_only_reads_the_requested_pools_jobs(self):
self.manager.reconcile = AsyncMock(return_value=2)
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
result = await self.routes.reconcile_serverless_endpoint(request)
self.assertEqual(result.status, 200)
self.assertEqual(json.loads(result.text)["resolved"], 2)
self.manager.reconcile.assert_awaited_once_with(worker_id="pool")
self.manager.cancel.assert_not_awaited()
async def test_resolve_requires_explicit_boolean_confirmation_and_record_id(self):
self.manager.acknowledge_recovery = AsyncMock()
for payload in ({"worker_id": "pool", "record_id": "record"},
{"worker_id": "pool", "record_id": "record", "provider_checked": "true"},
{"worker_id": "pool", "provider_checked": True}):
with self.subTest(payload=payload):
result = await self.routes.resolve_serverless_endpoint(types.SimpleNamespace(json=AsyncMock(return_value=payload)))
self.assertEqual(result.status, 400)
self.manager.acknowledge_recovery.assert_not_awaited()
async def test_resolve_is_scoped_and_does_not_claim_provider_cancellation(self):
self.manager.acknowledge_recovery = AsyncMock()
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
result = await self.routes.resolve_serverless_endpoint(request)
self.assertEqual(result.status, 200)
self.manager.acknowledge_recovery.assert_awaited_once_with("pool", "record", provider_checked=True)
self.assertIn("not provider-confirmed", json.loads(result.text)["message"])
self.manager.cancel.assert_not_awaited()
async def test_live_or_mismatched_record_refuses_manual_resolution(self):
self.manager.acknowledge_recovery = AsyncMock(side_effect=ValueError("Unresolved inactive job not found for this pool"))
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
result = await self.routes.resolve_serverless_endpoint(request)
self.assertEqual(result.status, 409)
async def test_failed_health_still_exposes_recovery_ids(self):
self.manager.health.side_effect = RuntimeError("private transport error")
self.manager.public_jobs = lambda _: [{"record_id": "record", "status": "UNKNOWN"}]
result = await self.routes.serverless_status_endpoint(types.SimpleNamespace(match_info={"worker_id": "pool"}))
self.assertEqual(json.loads(result.text)["jobs"][0]["record_id"], "record")
+36
View File
@@ -0,0 +1,36 @@
"""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
def test_image_list_network_runtime():
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
if not comfy_root:
pytest.skip('Set COMFYUI_SOURCE_ROOT for model-free multi-process HTTP acceptance')
helper = Path(__file__).parent / 'helpers' / 'image_list_network_check.py'
result = subprocess.run(
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
capture_output=True, text=True, timeout=180,
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
)
assert result.returncode == 0, result.stdout + '\n' + result.stderr
assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
+113
View File
@@ -0,0 +1,113 @@
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()
+219 -1
View File
@@ -1,11 +1,16 @@
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
@@ -21,7 +26,7 @@ def _load_process_module(module_filename: str):
sys.modules[package_name] = root_pkg
workers_pkg = types.ModuleType(f"{package_name}.workers")
workers_pkg.__path__ = []
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")
@@ -39,8 +44,23 @@ def _load_process_module(module_filename: str):
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,
@@ -70,6 +90,145 @@ class ComfyRootDiscoveryTests(unittest.TestCase):
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(
@@ -130,5 +289,64 @@ class LaunchCommandBuilderTests(unittest.TestCase):
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()
+99 -8
View File
@@ -2,10 +2,13 @@
Configuration management for ComfyUI-Distributed.
"""
import asyncio
import copy
import os
import json
import math
from contextlib import asynccontextmanager
from .logging import log
from .secret_fields import is_secret_field_name as _is_secret_field_name
# Import defaults for timeout fallbacks
from .constants import HEARTBEAT_TIMEOUT
@@ -15,6 +18,52 @@ _config_cache = None
_config_mtime = 0.0
_config_lock = asyncio.Lock()
_SERVERLESS_OBSOLETE_WORKER_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
def _contains_serverless_secret(value):
if isinstance(value, dict):
return any(
(isinstance(key, str) and _is_secret_field_name(key))
or _contains_serverless_secret(item)
for key, item in value.items()
)
if isinstance(value, (list, tuple)):
return any(_contains_serverless_secret(item) for item in value)
return False
def _normalize_serverless_workers(config):
"""Remove obsolete local-worker fields without adding local defaults."""
if not isinstance(config, dict):
return config
workers = config.get("workers")
if isinstance(workers, list):
for worker in workers:
if isinstance(worker, dict) and worker.get("type") == "serverless":
for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
worker.pop(field, None)
return config
def _validated_config_for_save(config):
"""Return a safe normalized copy ready for gpu_config.json."""
if _contains_serverless_secret(config):
raise ValueError("API keys must be stored outside the config file")
saved = _normalize_serverless_workers(copy.deepcopy(config))
workers = saved.get("workers") if isinstance(saved, dict) else None
if isinstance(workers, list) and any(
isinstance(worker, dict) and worker.get("type") == "serverless" for worker in workers
):
from .serverless_config import validate_serverless_worker
for index, worker in enumerate(workers):
if isinstance(worker, dict) and worker.get("type") == "serverless":
saved_worker = validate_serverless_worker(worker)
for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
saved_worker.pop(field, None)
workers[index] = saved_worker
return saved
def _config_path():
return CONFIG_FILE
@@ -33,7 +82,8 @@ def get_default_config():
"worker_probe_concurrency": 8,
"worker_prep_concurrency": 4,
"media_sync_concurrency": 2,
"media_sync_timeout_seconds": 120
"media_sync_timeout_seconds": 120,
"image_list_job_timeout_seconds": 0.0
},
"tunnel": {
"status": "stopped",
@@ -62,7 +112,7 @@ def _merge_with_defaults(data, defaults):
if key not in merged:
merged[key] = value
return merged
return _normalize_serverless_workers(merged)
def invalidate_config_cache():
@@ -97,11 +147,12 @@ def load_config():
return _config_cache
def save_config(config):
"""Saves the configuration to file."""
"""Saves normalized configuration without allowing API keys in the config file."""
tmp_path = f"{_config_path()}.tmp"
try:
saved_config = _validated_config_for_save(config)
with open(tmp_path, 'w', encoding='utf-8') as f:
json.dump(config, f, indent=2)
json.dump(saved_config, f, indent=2)
f.flush()
os.fsync(f.fileno())
os.replace(tmp_path, _config_path())
@@ -121,12 +172,21 @@ async def config_transaction():
"""Acquire config lock, yield loaded config, and save if changed."""
async with _config_lock:
config = load_config()
original_config = copy.deepcopy(config)
original_snapshot = json.dumps(config, sort_keys=True)
yield config
try:
yield config
except BaseException:
if isinstance(config, dict):
config.clear()
config.update(original_config)
raise
updated_snapshot = json.dumps(config, sort_keys=True)
if updated_snapshot != original_snapshot:
if not save_config(config):
raise RuntimeError("Failed to save config")
if updated_snapshot != original_snapshot and not save_config(config):
if isinstance(config, dict):
config.clear()
config.update(original_config)
raise RuntimeError("Failed to save config")
def ensure_config_exists():
"""Creates default config file if it doesn't exist. Used by __init__.py"""
@@ -157,6 +217,37 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
return max(1, int(default))
def get_image_list_job_timeout_seconds(config=None) -> float:
"""Return the validated image-list total-job timeout in seconds.
A value of zero disables the total-job deadline. Missing settings default
to zero; configured values must be finite, non-negative numbers.
"""
setting_key = "settings.image_list_job_timeout_seconds"
cfg = load_config() if config is None else config
if not isinstance(cfg, dict):
raise ValueError(f"{setting_key} must be a finite non-negative number")
settings = cfg.get("settings", {})
if not isinstance(settings, dict):
raise ValueError(f"{setting_key} must be a finite non-negative number")
value = settings.get("image_list_job_timeout_seconds", 0.0)
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
raise ValueError(f"{setting_key} must be a finite non-negative number")
if isinstance(value, str) and not value.strip():
raise ValueError(f"{setting_key} must be a finite non-negative number")
try:
timeout = float(value)
except (OverflowError, ValueError):
raise ValueError(f"{setting_key} must be a finite non-negative number") from None
if not math.isfinite(timeout) or timeout < 0:
raise ValueError(f"{setting_key} must be a finite non-negative number")
return timeout
def is_master_delegate_only() -> bool:
"""Returns True when master should skip local workload and act as orchestrator only."""
try:
+2
View File
@@ -87,6 +87,8 @@ def _split_host_and_port(host):
def build_worker_url(worker, endpoint=""):
"""Construct the worker base URL with optional endpoint."""
if worker.get("type") == "serverless":
raise ValueError("Serverless pools have no ComfyUI worker URL")
host = (worker.get("host") or "").strip()
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
+9
View File
@@ -0,0 +1,9 @@
"""Shared classification for credential-bearing configuration fields."""
def is_secret_field_name(key: object) -> bool:
"""Return whether a field name identifies an API key or credential."""
if not isinstance(key, str):
return False
normalized = "".join(character for character in key.casefold() if character.isalnum())
return normalized in {"apikey", "clearapikey"} or "credential" in normalized
+441
View File
@@ -0,0 +1,441 @@
"""Validation and private storage for Runpod serverless worker settings.
Secret storage is supported only on POSIX masters. The module remains
importable elsewhere so ordinary worker configuration and redaction continue
to work, but secret-store operations fail explicitly without POSIX locking.
"""
from __future__ import annotations
import copy
import json
import os
import re
import secrets
import stat
import threading
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator
from .secret_fields import is_secret_field_name
try:
import fcntl
except ImportError: # Windows and other platforms without POSIX file locking.
fcntl = None
class ServerlessConfigError(ValueError):
"""A serverless worker configuration is invalid."""
class SecretStoreError(RuntimeError):
"""The secret store could not be accessed safely."""
_ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII)
_UNSUPPORTED_SECRET_STORE_MESSAGE = (
"Serverless credential storage is currently supported on POSIX masters only"
)
_NUMERIC_FIELDS = {
"jobs_per_run": (1, 16, 1),
"concurrency": (1, 16, 1),
"timeout_seconds": (30, 86400, 1800),
}
_THREAD_LOCKS: dict[str, threading.Lock] = {}
_THREAD_LOCKS_GUARD = threading.Lock()
def validate_serverless_worker(worker: dict) -> dict:
"""Return a validated copy of a Runpod serverless worker record.
Provider credentials are intentionally not required here; they are stored
separately by :class:`SecretStore`.
"""
if not isinstance(worker, dict):
raise ServerlessConfigError("worker must be an object")
try:
result = copy.deepcopy(worker)
except Exception:
raise ServerlessConfigError("worker must be copyable") from None
worker_id = result.get("id")
if not isinstance(worker_id, str) or not worker_id.strip():
raise ServerlessConfigError("id must be a non-empty string")
endpoint_id = result.get("endpoint_id")
if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id):
raise ServerlessConfigError("endpoint_id must be an opaque Runpod identifier")
if result.get("type", "serverless") != "serverless":
raise ServerlessConfigError("type must be serverless")
if result.get("provider", "runpod") != "runpod":
raise ServerlessConfigError("provider is not supported")
if "enabled" in result and type(result["enabled"]) is not bool:
raise ServerlessConfigError("enabled must be a boolean")
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
value = result.get(field, default)
if type(value) is not int or not minimum <= value <= maximum:
raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}")
result[field] = value
result["type"] = "serverless"
result["provider"] = "runpod"
return result
def _thread_lock_for(path: str) -> threading.Lock:
with _THREAD_LOCKS_GUARD:
lock = _THREAD_LOCKS.get(path)
if lock is None:
lock = threading.Lock()
_THREAD_LOCKS[path] = lock
return lock
class SecretStore:
"""Atomically persist worker-id to API-key mappings in a private JSON file."""
def __init__(self, path: Path):
self.path = Path(path)
def get(self, worker_id: str) -> str | None:
self._require_supported()
worker_id = self._validate_worker_id(worker_id)
try:
with self._locked(create_parent=False) as (parent_fd, filename):
data = self._read_mapping(parent_fd, filename)
return data.get(worker_id)
except FileNotFoundError:
return None
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
def set(self, worker_id: str, key: str) -> None:
self._require_supported()
worker_id = self._validate_worker_id(worker_id)
if (
not isinstance(key, str)
or not key
or len(key) > 4096
or "\r" in key
or "\n" in key
):
raise SecretStoreError("secret must be 1 to 4096 characters and contain no CR or LF")
try:
with self._locked(create_parent=True) as (parent_fd, filename):
data = self._read_mapping(parent_fd, filename)
data[worker_id] = key
self._write_mapping(parent_fd, filename, data)
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
def delete(self, worker_id: str) -> None:
self._require_supported()
worker_id = self._validate_worker_id(worker_id)
try:
with self._locked(create_parent=False) as (parent_fd, filename):
data = self._read_mapping(parent_fd, filename)
if worker_id in data:
del data[worker_id]
self._write_mapping(parent_fd, filename, data)
except FileNotFoundError:
return
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
def has(self, worker_id: str) -> bool:
return self.get(worker_id) is not None
@staticmethod
def _require_supported() -> None:
# Never substitute a weaker Windows store: all operations retain the
# dir_fd/O_NOFOLLOW and flock protections implemented below.
if os.name != "posix" or fcntl is None:
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
@staticmethod
def _validate_worker_id(worker_id: str) -> str:
if not isinstance(worker_id, str) or not worker_id.strip():
raise SecretStoreError("worker_id must be a non-empty string")
return worker_id
def _absolute_path(self) -> Path:
path = self.path
if ".." in path.parts:
raise SecretStoreError("unsafe secret store path")
if not path.is_absolute():
path = Path.cwd() / path
if len(path.parts) < 2 or path.name in ("", ".", ".."):
raise SecretStoreError("unsafe secret store path")
return path
@contextmanager
def _locked(self, create_parent: bool) -> Iterator[tuple[int, str]]:
lock_api = fcntl
if os.name != "posix" or lock_api is None:
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
path = self._absolute_path()
filename = path.name
lock_name = f".{filename}.lock"
thread_lock = _thread_lock_for(str(path))
with thread_lock:
parent_fd = self._open_parent(path, create=create_parent)
lock_fd = None
try:
flags = os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW
lock_fd = os.open(lock_name, flags, 0o600, dir_fd=parent_fd)
lock_stat = os.fstat(lock_fd)
if not stat.S_ISREG(lock_stat.st_mode):
raise SecretStoreError("unsafe secret store lock")
os.fchmod(lock_fd, 0o600)
lock_api.flock(lock_fd, lock_api.LOCK_EX)
self._check_target(parent_fd, filename)
yield parent_fd, filename
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
finally:
if lock_fd is not None:
try:
lock_api.flock(lock_fd, lock_api.LOCK_UN)
except OSError:
pass
os.close(lock_fd)
os.close(parent_fd)
@staticmethod
def _open_parent(path: Path, create: bool) -> int:
parts = path.parts
current_fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
try:
for part in parts[1:-1]:
created = False
try:
next_fd = os.open(
part,
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
dir_fd=current_fd,
)
except FileNotFoundError:
if not create:
raise
try:
os.mkdir(part, 0o700, dir_fd=current_fd)
created = True
except FileExistsError:
pass
next_fd = os.open(
part,
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
dir_fd=current_fd,
)
if created:
os.fchmod(next_fd, 0o700)
os.close(current_fd)
current_fd = next_fd
return current_fd
except Exception:
os.close(current_fd)
raise
@staticmethod
def _check_target(parent_fd: int, filename: str) -> None:
try:
target_stat = os.stat(filename, dir_fd=parent_fd, follow_symlinks=False)
except FileNotFoundError:
return
if not stat.S_ISREG(target_stat.st_mode):
raise SecretStoreError("unsafe secret store file")
@classmethod
def _read_mapping(cls, parent_fd: int, filename: str) -> dict[str, str]:
cls._check_target(parent_fd, filename)
try:
fd = os.open(
filename,
os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
dir_fd=parent_fd,
)
except FileNotFoundError:
return {}
except Exception:
raise SecretStoreError("Unable to access secret store") from None
try:
file_stat = os.fstat(fd)
if not stat.S_ISREG(file_stat.st_mode):
raise SecretStoreError("unsafe secret store file")
os.fchmod(fd, 0o600)
with os.fdopen(fd, "rb") as stream:
raw = stream.read()
try:
mapping = json.loads(raw.decode("utf-8"), object_pairs_hook=cls._unique_object)
except Exception:
raise SecretStoreError("secret store is corrupt") from None
if not isinstance(mapping, dict) or any(
not isinstance(worker_id, str)
or not worker_id.strip()
or not isinstance(key, str)
or not key
for worker_id, key in mapping.items()
):
raise SecretStoreError("secret store is corrupt")
return mapping
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
finally:
try:
os.close(fd)
except OSError:
pass
@staticmethod
def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
result: dict[str, object] = {}
for key, value in pairs:
if key in result:
raise ValueError("duplicate key")
result[key] = value
return result
@staticmethod
def _write_mapping(parent_fd: int, filename: str, mapping: dict[str, str]) -> None:
payload = json.dumps(mapping, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
temp_name = f".{filename}.{secrets.token_hex(12)}.tmp"
temp_fd = None
try:
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW
temp_fd = os.open(temp_name, flags, 0o600, dir_fd=parent_fd)
os.fchmod(temp_fd, 0o600)
with os.fdopen(temp_fd, "wb") as stream:
temp_fd = None
stream.write(payload)
stream.flush()
os.fsync(stream.fileno())
SecretStore._check_target(parent_fd, filename)
os.replace(temp_name, filename, src_dir_fd=parent_fd, dst_dir_fd=parent_fd)
os.fsync(parent_fd)
except SecretStoreError:
raise
except Exception:
raise SecretStoreError("Unable to access secret store") from None
finally:
if temp_fd is not None:
try:
os.close(temp_fd)
except OSError:
pass
try:
os.unlink(temp_name, dir_fd=parent_fd)
except OSError:
pass
class ServerlessSecretChanges:
"""Apply serverless secret changes and restore them unless committed.
Use this around the associated config transaction so a failed config save
compensates any secret-store writes or deletions made along the way.
"""
def __init__(self):
self._originals: dict[tuple[int, str], tuple[SecretStore, str, str | None]] = {}
self._mutations: list[tuple[int, str]] = []
self._committed = False
def __enter__(self) -> "ServerlessSecretChanges":
return self
def __exit__(self, _exc_type, _exc_value, _traceback) -> bool:
if not self._committed:
self.rollback()
return False
def capture(self, store: SecretStore, worker_id: str) -> str | None:
"""Capture and return the original key once per store/worker pair."""
mutation_id = (id(store), worker_id)
if mutation_id not in self._originals:
self._originals[mutation_id] = (store, worker_id, store.get(worker_id))
return self._originals[mutation_id][2]
def set(self, store: SecretStore, worker_id: str, key: str) -> None:
self._prepare_mutation(store, worker_id)
store.set(worker_id, key)
def delete(self, store: SecretStore, worker_id: str) -> None:
self._prepare_mutation(store, worker_id)
store.delete(worker_id)
def commit(self) -> None:
"""Keep the secret changes after the corresponding config save."""
self._committed = True
def rollback(self) -> None:
"""Best-effort restoration without masking the original route error."""
for mutation_id in reversed(self._mutations):
store, worker_id, prior_secret = self._originals[mutation_id]
try:
if prior_secret is None:
store.delete(worker_id)
else:
store.set(worker_id, prior_secret)
except Exception:
pass
self._mutations.clear()
def _prepare_mutation(self, store: SecretStore, worker_id: str) -> None:
self.capture(store, worker_id)
mutation_id = (id(store), worker_id)
if mutation_id not in self._mutations:
# Record before applying: a store operation may fail after replacing
# its file, so the exception path must still attempt compensation.
self._mutations.append(mutation_id)
def _redact_credentials(value):
if isinstance(value, dict):
redacted = {}
for key, item in value.items():
if is_secret_field_name(key):
continue
redacted[key] = _redact_credentials(item)
return redacted
if isinstance(value, list):
return [_redact_credentials(item) for item in value]
if isinstance(value, tuple):
return tuple(_redact_credentials(item) for item in value)
return value
def public_config(config, store: SecretStore):
"""Return a deep-copied public configuration with credentials removed."""
result = _redact_credentials(copy.deepcopy(config))
if not isinstance(result, dict):
return result
workers = result.get("workers")
if isinstance(workers, list):
for worker in workers:
if isinstance(worker, dict) and worker.get("type") == "serverless":
worker_id = worker.get("id")
configured = (
store.has(worker_id)
if isinstance(worker_id, str) and worker_id.strip()
else False
)
worker["api_key_configured"] = configured
return result
+30 -4
View File
@@ -7,7 +7,7 @@ export function createApiClient(baseUrl) {
const request = async (
endpoint,
options = {},
{ retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
{ retries = TIMEOUTS.MAX_RETRIES, retry = true, redactErrors = false } = {},
) => {
const maxAttempts = retry ? retries : 1;
let lastError;
@@ -36,7 +36,7 @@ export function createApiClient(baseUrl) {
return await response.json();
} catch (error) {
lastError = error;
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${redactErrors ? "[redacted]" : error.message}`);
if (attempt < maxAttempts - 1) {
await new Promise(resolve => setTimeout(resolve, delay));
delay *= 2; // Exponential backoff
@@ -86,9 +86,35 @@ export function createApiClient(baseUrl) {
},
async updateWorker(workerId, data) {
return request('/distributed/config/update_worker', {
const apiKey = data?.api_key;
try {
return await request('/distributed/config/update_worker', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId, ...data })
}, { retry: false, redactErrors: Boolean(apiKey) });
} catch (error) {
if (apiKey && typeof error?.message === "string") {
error.message = error.message.split(apiKey).join("[redacted]");
}
throw error;
}
},
async testServerlessWorker(workerId) {
return request('/distributed/serverless/test', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId, ...data })
body: JSON.stringify({ worker_id: workerId }),
}, { retry: false });
},
async getServerlessStatus(workerId) {
return request(`/distributed/serverless/status/${encodeURIComponent(workerId)}`);
},
async cancelServerlessWorker(workerId) {
return request('/distributed/serverless/cancel', {
method: 'POST',
body: JSON.stringify({ worker_id: workerId }),
}, { retry: false });
},
+6 -4
View File
@@ -104,7 +104,8 @@ export const PULSE_ANIMATION_CSS = `
transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease;
}
.worker-settings.expanded {
max-height: 500px;
/* Let long forms and wrapped labels fit; the sidebar owns scrolling. */
max-height: none;
opacity: 1;
padding: 12px;
margin-top: 8px;
@@ -130,7 +131,7 @@ export const PULSE_ANIMATION_CSS = `
`;
export const UI_STYLES = {
statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
statusDot: "display: inline-block; flex: 0 0 10px; 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;",
@@ -146,12 +147,12 @@ export const UI_STYLES = {
// 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;",
contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
contentColumn: "flex: 1; min-width: 0; 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);",
// Row and content styles
infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;",
workerContent: "display: flex; align-items: center; gap: 10px; flex: 1; min-width: 0;",
// Form and controls styles
buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
@@ -224,6 +225,7 @@ export const ENDPOINTS = {
export const NODE_CLASSES = {
DISTRIBUTED_COLLECTOR: 'DistributedCollector',
IMAGE_LIST_START: 'DistributedImageListStart',
DISTRIBUTED_SEED: 'DistributedSeed',
DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',
+11
View File
@@ -235,6 +235,17 @@
border-radius: 4px;
}
.dist-worker-info {
min-width: 0;
overflow-wrap: anywhere;
}
.dist-serverless-status {
display: block;
margin-top: 4px;
line-height: 1.4;
}
.dist-worker-info__title {
color: var(--dist-card-title) !important;
font-size: 1.03em;
+9 -12
View File
@@ -2,7 +2,6 @@ import { app } from "/scripts/app.js";
import { ENDPOINTS } from "./constants.js";
const NODE_CLASS = "DistributedValue";
const CONVERTED_WIDGET = "converted-widget";
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
@@ -42,15 +41,13 @@ function getDynamicWorkerWidgets(node) {
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
}
function hideWidgetForGood(node, widget, suffix = "") {
function hideWidgetForGood(widget) {
if (!widget) return;
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
widget.origType = widget.type;
widget.origComputeSize = widget.computeSize;
widget.origSerializeValue = widget.serializeValue;
widget.computeSize = () => [0, -4];
widget.type = `${CONVERTED_WIDGET}${suffix}`;
// These are serialized backing fields, not widgets converted into sockets.
// Keep their native type/serializer and mutate the shared options in place
// so both the classic canvas and Nodes 2.0 suppress the whole widget row.
widget.options ??= {};
widget.options.hidden = true;
// Hide any attached DOM element (multiline widgets).
if (widget.element) widget.element.style.display = "none";
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
if (widget.linkedWidgets) {
for (const linked of widget.linkedWidgets) {
hideWidgetForGood(node, linked, `:${widget.name}`);
hideWidgetForGood(linked);
}
}
}
function hideRawWidgets(node) {
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
hideWidgetForGood(getRawDefaultWidget(node));
hideWidgetForGood(getRawWorkerValuesWidget(node));
}
function removeDynamicDefaultWidget(node) {
+27 -2
View File
@@ -9,7 +9,9 @@ export function setupInterceptor(extension) {
const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
if (hasCollector || hasDistUpscale) {
const hasImageList = findNodesByClass(prompt.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
if (hasCollector || hasDistUpscale || hasImageList) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
checkAllWorkerStatuses(extension);
@@ -25,12 +27,20 @@ export function setupInterceptor(extension) {
export async function executeParallelDistributed(extension, promptWrapper) {
const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
try {
const enabledWorkers = extension.enabledWorkers;
const hasImageList = findNodesByClass(promptWrapper.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
// Image-list v1 uses direct workers only; do not contact paid provider pools.
const enabledWorkers = hasImageList
? extension.enabledWorkers.filter(worker => worker.type !== "serverless")
: extension.enabledWorkers;
extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
// Pre-flight health check on all enabled workers
const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
if (hasImageList && activeWorkers.length === 0) {
throw new Error("Distributed Image List requires an active direct worker; master fallback is disabled.");
}
// Case: Enabled workers but all offline
if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
extension.log("No active workers found. All enabled workers are offline.");
@@ -112,6 +122,21 @@ export async function performPreflightCheck(extension, workers) {
const startTime = Date.now();
const checkPromises = workers.map(async (worker) => {
if (worker.type === "serverless") {
try {
const providerStatus = await extension.api.getServerlessStatus(worker.id);
if (providerStatus?.ok) {
extension.log(`Serverless worker ${worker.name} is ready to scale`, "debug");
return { worker, active: true };
}
extension.log(`Serverless worker ${worker.name} is not ready: ${providerStatus?.message || "unavailable"}`, "debug");
return { worker, active: false };
} catch (error) {
extension.log(`Serverless worker ${worker.name} status check failed: ${error.message}`, "debug");
return { worker, active: false };
}
}
const workerUrl = getWorkerUrl(extension, worker);
extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug");
+10 -1
View File
@@ -42,6 +42,15 @@ export async function detectMasterIP(extension) {
if (shouldAutoPopulate) {
extension.log(`Auto-populating workers based on ${extension.cudaDeviceCount} CUDA devices (excluding master on CUDA ${extension.masterCudaDevice})`, "info");
const workerCount = extension.cudaDeviceCount - (
Number.isInteger(extension.masterCudaDevice) &&
extension.masterCudaDevice >= 0 &&
extension.masterCudaDevice < extension.cudaDeviceCount ? 1 : 0
);
const workerPorts = data.local_worker_ports;
if (!Array.isArray(workerPorts) || workerPorts.length !== workerCount) {
throw new Error("Could not allocate enough available local worker ports; check the server log");
}
const newWorkers = [];
let workerNum = 1;
let portOffset = 0;
@@ -56,7 +65,7 @@ export async function detectMasterIP(extension) {
id: generateUUID(),
name: `Worker ${workerNum}`,
host: isRunpod ? null : "localhost",
port: 8189 + portOffset,
port: workerPorts[portOffset],
cuda_device: i,
enabled: true,
extra_args: isRunpod ? "--listen" : "",
+2 -2
View File
@@ -72,7 +72,7 @@ export function renderSettingsSection(extension) {
createCheckboxSetting(
"setting-debug",
"Debug Mode",
"Enable verbose logging in the browser console.",
"Enable verbose logging in the browser console and ComfyUI server output.",
extension.config?.settings?.debug || false,
(event) => extension._updateSetting("debug", event.target.checked)
)
@@ -100,7 +100,7 @@ export function renderSettingsSection(extension) {
createNumberSetting(
"setting-worker-timeout",
"Worker Timeout",
"Seconds without a heartbeat before a worker is considered timed out. Default 60.",
"Maximum result-wait and heartbeat inactivity period before recovery begins. Busy workers may receive additional grace. Default: 60 seconds.",
extension.config?.settings?.worker_timeout_seconds ?? 60,
10,
1,
+13
View File
@@ -18,6 +18,19 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) {
const workerState = extension.state.getWorker(workerId);
const isLaunching = Boolean(workerState?.launching);
if (worker.type === "serverless") {
const providerStatus = workerState?.serverlessStatus || {};
if (newStatus.online && newStatus.processing) {
extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, `Runpod active (${newStatus.queueCount || 0} queued)`, false);
} else if (newStatus.online) {
extension.ui.updateStatusDot(workerId, STATUS_COLORS.ONLINE_GREEN, providerStatus.message || "Runpod ready to scale", false);
} else if (worker.enabled) {
extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, providerStatus.message || "Runpod endpoint unavailable", false);
}
updateWorkerControls(extension, workerId);
return true;
}
if (isLaunching && !newStatus.online) {
extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true);
} else if (newStatus.online && newStatus.processing) {
+120
View File
@@ -0,0 +1,120 @@
import { readFileSync } from "node:fs";
import vm from "node:vm";
import { describe, expect, it, vi } from "vitest";
// Evaluate this repository's browser entrypoint without needing a ComfyUI server.
// The real-renderer acceptance additionally checks the input rows in Chromium.
const source = readFileSync(new URL("../distributedValue.js", import.meta.url), "utf8")
.replace(/^import .*;\r?\n/gm, "");
function setup({ connected = false, targetType = "STRING" } = {}) {
let extension;
const listeners = new Map();
const tasks = [];
const workers = [
{ id: "gpu-a", name: "GPU A", enabled: true },
{ id: "gpu-b", name: "GPU B", enabled: true },
{ id: "off", name: "Disabled", enabled: false },
];
const target = {
inputs: [{ name: "value" }],
widgets: [{ name: "value", type: targetType === "INT" ? "number" : "string", options: { step: 1, precision: 0 } }],
};
const graph = { links: { 1: { target_id: 2, target_slot: 0 } }, getNodeById: () => target };
const raw = ["default_value", "worker_values"].map((name, index) => ({
name,
type: "text",
value: index === 0 ? "saved default" : '{"1":"saved worker"}',
options: { multiline: false },
computeSize: vi.fn(() => [200, 20]),
serializeValue: vi.fn(function () { return this.value; }),
}));
const node = {
comfyClass: "DistributedValue", graph, widgets: [...raw], size: [200, 100],
inputs: raw.map(w => ({ name: w.name, widget: { name: w.name }, link: null })),
outputs: [{ links: connected ? [1] : [] }],
computeSize: () => [200, 100], setSize: vi.fn(), setDirtyCanvas: vi.fn(),
addWidget(type, name, value, callback, options) {
const widget = { type, name, value, callback, options };
this.widgets.push(widget);
return widget;
},
configure(data) {
data.widgets_values.forEach((value, i) => { this.widgets[i].value = value; });
return "configured";
},
};
vm.runInNewContext(source, {
app: { graph, registerExtension: ext => { extension = ext; } },
ENDPOINTS: { CONFIG: "/distributed/config" },
fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
window: { addEventListener: (name, handler) => listeners.set(name, handler) },
setTimeout: callback => { tasks.push(callback); }, console,
}, { filename: "distributedValue.js" });
return { node, raw, extension, listeners, tasks, flush: () => { while (tasks.length) tasks.shift()(); } };
}
describe("Distributed Value internal widget visibility", () => {
it("hides raw fields natively without converting their types or replacing their options", async () => {
const h = setup();
const options = h.raw.map(w => w.options);
const sizes = h.raw.map(w => w.computeSize);
const serializers = h.raw.map(w => w.serializeValue);
const slots = h.node.inputs;
await h.extension.nodeCreated(h.node);
for (const [i, widget] of h.raw.entries()) {
expect(widget.type).toBe("text");
expect(widget.options).toBe(options[i]);
expect(widget.options.hidden).toBe(true);
expect(widget.computeSize).toBe(sizes[i]);
expect(widget.serializeValue).toBe(serializers[i]);
expect(h.node.widgets[i]).toBe(widget);
}
expect(h.node.inputs).toBe(slots);
expect(h.node.widgets.slice(2).map(w => w.label)).toEqual(["default_value", "GPU A", "GPU B"]);
expect(h.node.widgets.slice(2).every(w => !w.options.hidden)).toBe(true);
});
it("keeps typed edits in the original serialized fields across rebuild/configure", async () => {
const h = setup({ connected: true, targetType: "INT" });
await h.extension.nodeCreated(h.node);
h.node.widgets.find(w => w.name === "_dv_default").callback(42);
h.node.widgets.find(w => w.name === "_dv_worker_1").callback(7);
h.node.widgets.find(w => w.name === "_dv_worker_2").callback(9);
const values = h.raw.map(w => w.serializeValue());
expect(values[0]).toBe(42);
expect(JSON.parse(values[1])).toMatchObject({
_type: "INT", 1: "7", 2: "9", _by_worker_id: { "gpu-a": "7", "gpu-b": "9" },
});
expect(h.node.configure({ widgets_values: values })).toBe("configured");
h.flush();
expect(h.raw.map(w => w.serializeValue())).toEqual(values);
expect(h.raw.every(w => w.options.hidden && w.type === "text")).toBe(true);
expect(h.node.widgets.slice(2).map(w => w.value)).toEqual([42, 7, 9]);
});
it("reapplies visibility on worker refresh without duplicating raw or dynamic fields", async () => {
const h = setup();
await h.extension.nodeCreated(h.node);
h.raw[0].options.hidden = false;
await h.listeners.get("distributed:workers-changed")({ detail: { workers: [
{ id: "gpu-b", name: "Renamed B", enabled: true },
] } });
expect(h.raw.every(w => w.options.hidden)).toBe(true);
expect(h.node.widgets.map(w => w.name)).toEqual(["default_value", "worker_values", "_dv_default", "_dv_worker_1"]);
expect(h.node.widgets[3].label).toBe("Renamed B");
});
it("hides linked backing widgets and handles absent options without altering serializers", async () => {
const h = setup();
delete h.raw[1].options;
const linked = { type: "text", options: {}, serializeValue: vi.fn(() => "linked") };
h.raw[0].linkedWidgets = [linked];
await h.extension.nodeCreated(h.node);
expect(h.raw[1].options.hidden).toBe(true);
expect(linked.options.hidden).toBe(true);
expect(linked.type).toBe("text");
expect(linked.serializeValue()).toBe("linked");
});
});
-15
View File
@@ -1,15 +0,0 @@
import { describe, expect, it } from "vitest";
import { buildWorkerWebSocketUrl } from "../urlUtils.js";
describe("execution decision helpers", () => {
it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
"ws://worker.local:8188/distributed/worker_ws"
);
expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
"wss://worker.example.com/distributed/worker_ws"
);
});
});
+56
View File
@@ -0,0 +1,56 @@
import { afterEach, describe, expect, it, vi } from 'vitest';
vi.mock('../../scripts/api.js', () => ({ api: { clientId: 'test-client' } }), { virtual: true });
vi.mock('../workerLifecycle.js', () => ({
checkAllWorkerStatuses: vi.fn(), getWorkerUrl: () => 'http://worker.invalid',
}));
vi.mock('../workerUtils.js', () => ({
applyProbeResultToWorkerDot: vi.fn(),
findNodesByClass: (prompt, cls) => Object.values(prompt).filter(node => node.class_type === cls),
}));
import { api } from '../../scripts/api.js';
import { setupInterceptor } from '../executionUtils.js';
const prompt = { output: { '1': { class_type: 'DistributedImageListStart', inputs: {} } } };
function extension(workers, online = true) {
return {
isEnabled: true, enabledWorkers: workers, log: vi.fn(), config: { settings: {} },
originalQueuePrompt: vi.fn().mockResolvedValue({ prompt_id: 'local' }),
api: { probeWorker: vi.fn().mockResolvedValue({ ok: online }),
getServerlessStatus: vi.fn(),
queueDistributed: vi.fn().mockResolvedValue({ prompt_id: 'distributed' }) },
ui: { showToast: vi.fn() },
};
}
afterEach(() => vi.useRealTimers());
describe('image-list queue interception', () => {
it('routes image-list graphs through backend orchestration', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'one', type: 'local' }]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: ['one'] }));
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
});
it('never falls back to master when all image-list workers are offline', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'one', type: 'local' }], false);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toBeNull();
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
});
it('does not probe serverless providers for an image-list job', async () => {
vi.useFakeTimers();
const ext = extension([{ id: 'pool', type: 'serverless' }]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, prompt)).toBeNull();
expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
});
it('preserves ordinary prompt handling', async () => {
const ext = extension([]);
setupInterceptor(ext);
expect(await api.queuePrompt(0, { output: {} })).toEqual({ prompt_id: 'local' });
expect(ext.api.queueDistributed).not.toHaveBeenCalled();
});
});
+65
View File
@@ -0,0 +1,65 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { detectMasterIP } from "../masterDetection.js";
describe("automatic local worker ports", () => {
let originalWindow;
beforeEach(() => {
originalWindow = globalThis.window;
globalThis.window = { location: { hostname: "localhost", port: "443" } };
});
afterEach(() => { globalThis.window = originalWindow; });
function makeExtension(overrides = {}) {
return {
config: { workers: [], settings: {}, master: {} },
log: vi.fn(),
api: {
getNetworkInfo: vi.fn().mockResolvedValue({
cuda_device: 0, cuda_device_count: 3,
master_port: 8189, local_worker_ports: [8191, 8193],
...overrides,
}),
updateMaster: vi.fn().mockResolvedValue({}),
updateWorker: vi.fn().mockResolvedValue({}),
updateSetting: vi.fn().mockResolvedValue({}),
},
ui: { updateMasterDisplay: vi.fn() },
app: {},
loadConfig: vi.fn().mockResolvedValue(),
};
}
it("uses server-allocated ports, not hardcoded ports or the browser proxy port", async () => {
const extension = makeExtension();
await detectMasterIP(extension);
expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
expect(extension.config.workers.map(worker => worker.cuda_device)).toEqual([1, 2]);
expect(extension.api.updateWorker).toHaveBeenCalledTimes(2);
});
it("does not overwrite an existing manually configured worker", async () => {
const extension = makeExtension();
const worker = { id: "manual", host: "localhost", port: 8189 };
extension.config.workers = [worker];
await detectMasterIP(extension);
expect(extension.config.workers).toEqual([worker]);
expect(extension.api.updateWorker).not.toHaveBeenCalled();
});
it("fails without partially saving workers if not enough ports are available", async () => {
const extension = makeExtension({ local_worker_ports: [8191] });
await detectMasterIP(extension);
expect(extension.config.workers).toEqual([]);
expect(extension.api.updateWorker).not.toHaveBeenCalled();
expect(extension.api.updateSetting).not.toHaveBeenCalled();
expect(extension.log).toHaveBeenCalledWith(expect.stringContaining("port"), "error");
});
it("preserves Runpod worker host and launch settings", async () => {
globalThis.window.location.hostname = "pod-8189.proxy.runpod.net";
const extension = makeExtension();
await detectMasterIP(extension);
expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
expect(extension.config.workers.every(worker => worker.host === null && worker.extra_args === "--listen")).toBe(true);
});
});
+507
View File
@@ -0,0 +1,507 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client" } }), { virtual: true });
import { createApiClient } from "../apiClient.js";
import { createWorkerSettingsForm } from "../ui/settingsForm.js";
import { DistributedUI } from "../ui.js";
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
import { isRemoteWorker, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
import { performPreflightCheck } from "../executionUtils.js";
function createDocumentMock() {
const elements = new Map();
const makeElement = (tagName) => {
const element = {
tagName,
children: [],
dataset: {},
style: {},
setAttribute(name, value) { this[name] = value; },
classList: {
values: new Set(),
add(...names) { names.forEach((name) => this.values.add(name)); },
remove(...names) { names.forEach((name) => this.values.delete(name)); },
toggle(name, force) {
const enabled = force ?? !this.values.has(name);
if (enabled) this.values.add(name);
else this.values.delete(name);
return enabled;
},
},
appendChild(child) { this.children.push(child); child.parentNode = this; return child; },
dispatchEvent() { this.onchange?.({ target: this }); },
};
Object.defineProperty(element, "id", {
get() { return this._id || ""; },
set(value) { this._id = value; elements.set(value, this); },
});
return element;
};
return {
elements,
document: {
createElement: makeElement,
getElementById: (id) => elements.get(id) || null,
},
};
}
function createFormUi() {
return {
createFormGroup(label, value, id, type = "text", placeholder = "") {
const group = document.createElement("div");
const input = document.createElement("input");
input.id = id;
input.value = String(value ?? "");
input.type = type;
input.placeholder = placeholder;
group.appendChild(input);
return { group, input, label };
},
createButton(text, onclick) {
const button = document.createElement("button");
button.textContent = text;
button.onclick = onclick;
return button;
},
createButtonGroup(buttons) {
const group = document.createElement("div");
buttons.forEach((button) => group.appendChild(button));
return group;
},
};
}
function makeSaveExtension(worker, updateWorker = vi.fn().mockResolvedValue({})) {
return {
config: { workers: [worker] },
api: { updateWorker },
state: { updateWorker: vi.fn(), setWorkerExpanded: vi.fn() },
app: { extensionManager: { toast: { add: vi.fn() } } },
panelElement: null,
};
}
describe("Runpod serverless worker panel", () => {
let originalDocument;
let originalWindow;
let elements;
beforeEach(() => {
originalDocument = globalThis.document;
originalWindow = globalThis.window;
const mock = createDocumentMock();
elements = mock.elements;
globalThis.document = mock.document;
globalThis.window = {
location: { hostname: "localhost", protocol: "http:", port: "8188", origin: "http://localhost:8188" },
};
});
afterEach(() => {
globalThis.document = originalDocument;
globalThis.window = originalWindow;
vi.restoreAllMocks();
});
it("classifies Runpod serverless as remote for worker allocation", () => {
expect(isRemoteWorker({}, { type: "serverless" })).toBe(true);
});
it("initializes a saved serverless form before it is attached to the document", () => {
// Real document lookups cannot see a form while its card is being built.
vi.spyOn(document, "getElementById").mockReturnValue(null);
const worker = {
id: "saved", name: "Runpod", type: "serverless", endpoint_id: "ep_123",
jobs_per_run: 2, concurrency: 1, timeout_seconds: 900, api_key_configured: true,
};
const form = createWorkerSettingsForm(createFormUi(), {}, worker);
expect(form.parentNode).toBeUndefined();
expect(elements.get("worker-type-saved").value).toBe("serverless");
for (const field of ["host", "port", "cuda", "args"]) {
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
}
expect(elements.get("runpod-text-saved").style.display).toBe("none");
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "concurrency", "timeout-seconds", "serverless-test"]) {
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
}
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
expect(elements.get("jobs-per-run-saved").value).toBe("2");
expect(elements.get("concurrency-saved").value).toBe("1");
expect(elements.get("timeout-seconds-saved").value).toBe("900");
expect(elements.get("api-key-saved").value).toBe("");
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
});
it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
vi.spyOn(document, "getElementById").mockReturnValue(null);
createWorkerSettingsForm(createFormUi(), {}, {
id: "cloud", name: "Cloud", type: "cloud", host: "worker.example", port: 8443,
});
const type = elements.get("worker-type-cloud");
expect(type.value).toBe("cloud");
expect(elements.get("port-cloud").value).toBe("8443");
expect(elements.get("host-cloud").value).toBe("worker.example");
expect(elements.get("runpod-text-cloud").style.display).toBe("block");
type.value = "remote";
type.dispatchEvent();
type.value = "cloud";
type.dispatchEvent();
expect(elements.get("port-cloud").value).toBe("443");
});
it("shows only serverless settings and keeps credentials write-only", () => {
const worker = {
id: "w1", name: "Runpod", type: "remote", host: "worker.example", port: 8188,
cuda_device: 0, extra_args: "--test", api_key_configured: true,
};
const form = createWorkerSettingsForm(createFormUi(), {}, worker);
const type = elements.get("worker-type-w1");
type.value = "serverless";
type.dispatchEvent();
expect(type.children.map((option) => [option.value, option.textContent])).toContainEqual([
"serverless", "Serverless (Runpod)",
]);
for (const id of ["host-group-w1", "port-group-w1", "cuda-group-w1", "args-group-w1"]) {
expect(elements.get(id).style.display).toBe("none");
}
for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "concurrency-w1", "timeout-seconds-w1"]) {
expect(elements.has(id)).toBe(true);
}
for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "concurrency-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) {
expect(elements.get(id).style.display).toBe("flex");
}
expect(elements.has("serverless-help-w1")).toBe(false);
expect(elements.get("api-key-w1").type).toBe("password");
expect(elements.get("api-key-w1").value).toBe("");
expect(elements.get("api-key-w1").placeholder).toMatch(/leave blank/i);
expect(elements.get("jobs-per-run-w1").value).toBe("1");
expect(elements.get("jobs-per-run-w1").min).toBe("1");
expect(elements.get("jobs-per-run-w1").max).toBe("16");
expect(elements.get("concurrency-w1").value).toBe("1");
expect(elements.get("concurrency-w1").min).toBe("1");
expect(elements.get("concurrency-w1").max).toBe("16");
expect(elements.get("timeout-seconds-w1").value).toBe("1800");
expect(elements.get("timeout-seconds-w1").min).toBe("30");
expect(elements.get("timeout-seconds-w1").max).toBe("86400");
expect(form.children.length).toBeGreaterThan(0);
});
it("saves fixed Runpod settings without enabling or retaining the API key", async () => {
const worker = { id: "w2", name: "Old", type: "remote", enabled: false, api_key_configured: false };
const extension = makeSaveExtension(worker);
elements.set("name-w2", { value: "Runpod worker" });
elements.set("worker-type-w2", { value: "serverless" });
elements.set("endpoint-id-w2", { value: "ep_123" });
elements.set("api-key-w2", { value: "secret-token" });
elements.set("clear-api-key-w2", { checked: false });
elements.set("jobs-per-run-w2", { value: "4" });
elements.set("concurrency-w2", { value: "3" });
elements.set("timeout-seconds-w2", { value: "2400" });
await saveWorkerSettings(extension, "w2");
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
api_key: "secret-token", jobs_per_run: 4, concurrency: 3, timeout_seconds: 2400, enabled: false,
});
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
expect(worker).not.toHaveProperty("api_key");
expect(JSON.stringify(extension.config)).not.toContain("secret-token");
expect(elements.get("api-key-w2").value).toBe("");
expect(extension.state.updateWorker).toHaveBeenCalledWith("w2", { enabled: false });
expect(JSON.stringify(extension.state.updateWorker.mock.calls)).not.toContain("secret-token");
});
it("omits a blank API key so the saved key is preserved", async () => {
const worker = { id: "w-blank", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
const extension = makeSaveExtension(worker);
elements.set("name-w-blank", { value: "Runpod" });
elements.set("worker-type-w-blank", { value: "serverless" });
elements.set("endpoint-id-w-blank", { value: "ep_123" });
elements.set("api-key-w-blank", { value: "" });
elements.set("clear-api-key-w-blank", { checked: false });
elements.set("jobs-per-run-w-blank", { value: "1" });
elements.set("concurrency-w-blank", { value: "1" });
elements.set("timeout-seconds-w-blank", { value: "1800" });
await saveWorkerSettings(extension, "w-blank");
const payload = extension.api.updateWorker.mock.calls[0][1];
expect(payload).not.toHaveProperty("api_key");
expect(payload).not.toHaveProperty("clear_api_key");
expect(worker.api_key_configured).toBe(true);
expect(worker).not.toHaveProperty("api_key");
});
it("omits a blank key to preserve it and sends an explicit clear flag when requested", async () => {
const worker = { id: "w3", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
const extension = makeSaveExtension(worker);
elements.set("name-w3", { value: "Runpod" });
elements.set("worker-type-w3", { value: "serverless" });
elements.set("endpoint-id-w3", { value: "ep_123" });
elements.set("api-key-w3", { value: "" });
elements.set("clear-api-key-w3", { checked: true });
elements.set("jobs-per-run-w3", { value: "1" });
elements.set("concurrency-w3", { value: "1" });
elements.set("timeout-seconds-w3", { value: "1800" });
await saveWorkerSettings(extension, "w3");
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
clear_api_key: true, jobs_per_run: 1, concurrency: 1, timeout_seconds: 1800, enabled: true,
});
expect(worker.api_key_configured).toBe(false);
expect(worker).not.toHaveProperty("api_key");
});
it.each([
[undefined, "Checking status…"],
[{ ok: true, queueRemaining: 0, jobs: [] }, "Ready"],
[{ ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] }, "Queued"],
[{ ok: true, jobs: [{ status: "IN_PROGRESS" }] }, "Running"],
[{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_QUEUE" }] }, "Running"],
[{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_PROGRESS" }] }, "Running · 2 jobs"],
[{ ok: true, jobs: [{ status: "DOWNLOADING" }] }, "Downloading result"],
[{ ok: true, jobs: [{ status: "CANCEL_REQUESTED" }] }, "Cancelling"],
[{ ok: true, queueRemaining: 1 }, "Active · 1 job"],
[{ ok: true, queueRemaining: 3 }, "Active · 3 jobs"],
[{ ok: true, jobs: [{ status: "COMPLETED" }] }, "Ready"],
[{ ok: true, jobs: [{ status: "FAILED" }] }, "Last job failed"],
[{ ok: false, message: "Check the saved endpoint ID and API key" }, "Unavailable"],
])("formats provider state as one plain-language summary: %j", (status, expected) => {
expect(formatServerlessStatus(status)).toBe(expected);
});
it("keeps status above a consistent endpoint info box and retains expanded settings", () => {
const ui = new DistributedUI();
const worker = {
id: "compact", name: "Used Car Inventory", type: "serverless", enabled: true,
endpoint_id: "ep_private_id", jobs_per_run: 1, concurrency: 1,
};
const extension = {
config: { workers: [worker] },
state: {
isWorkerExpanded: () => false,
getWorker: () => ({ serverlessStatus: { ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] } }),
},
};
const card = ui.renderEntityCard("worker", worker, extension);
const collect = (element) => [element, ...(element.children || []).flatMap(collect)];
const info = collect(card).find((element) => element.classList.values.has("dist-worker-info"));
expect(info.innerHTML).not.toContain("Runpod Serverless");
expect(info.innerHTML).not.toMatch(/ep_private_id|jobs\/run|slots/);
expect(elements.get("serverless-status-compact").textContent).toBe("Queued");
expect(elements.get("endpoint-id-compact").value).toBe("ep_private_id");
expect(Number(elements.get("jobs-per-run-compact").value)).toBe(1);
expect(Number(elements.get("concurrency-compact").value)).toBe(1);
const box = elements.get("controls-compact").children[0].children[0];
expect(box.textContent).toBe("Runpod Serverless endpoint");
expect(box.classList.values.has("dist-info-box")).toBe(true);
expect(box.style.cssText).toBe(ui.styles.infoBox);
expect(elements.get("status-compact").style.cssText).toContain("flex: 0 0 10px");
});
it("uses provider status instead of a direct ComfyUI URL probe", async () => {
const statusElement = { textContent: "" };
elements.set("serverless-status-w4", statusElement);
const status = {
ok: true, queueRemaining: 0, message: "Ready to scale",
jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w4", execution_id: "exec-1" }],
};
const extension = {
config: { workers: [{ id: "w4", name: "Runpod", type: "serverless", enabled: true }] },
api: { getServerlessStatus: vi.fn().mockResolvedValue(status), probeWorker: vi.fn() },
statusCheckAbortController: null,
state: {
setWorkerStatus: vi.fn(),
updateWorker: vi.fn(),
getWorkerStatus: vi.fn().mockReturnValue({ online: true, processing: false, queueCount: 0 }),
getWorker: vi.fn().mockReturnValue({}),
isWorkerLaunching: vi.fn().mockReturnValue(false),
},
updateWorkerCard: vi.fn().mockReturnValue(true),
log: vi.fn(),
};
await checkWorkerStatus(extension, extension.config.workers[0]);
expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w4");
expect(extension.api.probeWorker).not.toHaveBeenCalled();
expect(extension.state.setWorkerStatus).toHaveBeenCalledWith("w4", {
online: true, processing: true, queueCount: 0,
});
expect(statusElement.textContent).toBe("Running");
expect(statusElement.title).toBe("Ready to scale");
});
it("considers a healthy zero-worker Runpod pool ready and never builds a worker URL", async () => {
const worker = { id: "w-ready", name: "Runpod", type: "serverless", enabled: true };
const extension = {
api: {
getServerlessStatus: vi.fn().mockResolvedValue({ ok: true, queueRemaining: 0, message: "Ready to scale" }),
probeWorker: vi.fn(),
},
log: vi.fn(),
};
await expect(performPreflightCheck(extension, [worker])).resolves.toEqual([worker]);
expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w-ready");
expect(extension.api.probeWorker).not.toHaveBeenCalled();
expect(getWorkerUrl(extension, worker)).toBeNull();
});
it("does not render or invoke direct ComfyUI launch, stop, or log controls", async () => {
const ui = new DistributedUI();
const worker = { id: "w-controls", name: "Runpod", type: "serverless", enabled: true };
const controls = ui.createControlsSection({ dynamic: true }, worker, {}, true);
const extension = {
config: { workers: [worker] },
api: {
launchWorker: vi.fn(), stopWorker: vi.fn(), getRemoteWorkerLog: vi.fn(),
},
};
const box = controls.children[0].children[0];
expect(box.textContent).toBe("Runpod Serverless endpoint");
expect(box.tagName).toBe("div");
expect(controls.children[0].children).toHaveLength(1);
await launchWorker(extension, worker.id);
await stopWorker(extension, worker.id);
await viewWorkerLog(extension, worker.id, true);
expect(extension.api.launchWorker).not.toHaveBeenCalled();
expect(extension.api.stopWorker).not.toHaveBeenCalled();
expect(extension.api.getRemoteWorkerLog).not.toHaveBeenCalled();
});
it("uses the serverless status, test, and cancel routes with the frozen request shapes", async () => {
const originalFetch = globalThis.fetch;
const statusPayload = {
ok: true,
queueRemaining: 0,
message: "Ready to scale",
jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w5", execution_id: "exec-1" }],
};
globalThis.fetch = vi.fn()
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint test passed" }) })
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue(statusPayload) })
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true }) });
try {
const client = createApiClient("http://localhost:8188");
await client.testServerlessWorker("w5");
await expect(client.getServerlessStatus("w5")).resolves.toEqual(statusPayload);
await client.cancelServerlessWorker("w5");
expect(globalThis.fetch).toHaveBeenNthCalledWith(1, "http://localhost:8188/distributed/serverless/test", expect.objectContaining({
method: "POST", body: JSON.stringify({ worker_id: "w5" }),
}));
expect(globalThis.fetch).toHaveBeenNthCalledWith(2, "http://localhost:8188/distributed/serverless/status/w5", expect.any(Object));
expect(globalThis.fetch).toHaveBeenNthCalledWith(3, "http://localhost:8188/distributed/serverless/cancel", expect.objectContaining({
method: "POST", body: JSON.stringify({ worker_id: "w5" }),
}));
} finally {
globalThis.fetch = originalFetch;
}
});
it("tests only the saved Runpod endpoint and reports the backend message", async () => {
const button = document.createElement("button");
button.id = "test-serverless-w6";
button.textContent = "Test Saved Endpoint";
const extension = {
api: {
testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
queueDistributed: vi.fn(),
},
app: { extensionManager: { toast: { add: vi.fn() } } },
};
await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
expect(extension.api.queueDistributed).not.toHaveBeenCalled();
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
severity: "success", detail: "Endpoint reachable",
}));
expect(button.textContent).toBe("Test Saved Endpoint");
expect(button.disabled).toBe(false);
});
it("redacts API-key-bearing worker update errors from logs and user-visible exceptions", async () => {
const originalFetch = globalThis.fetch;
const logSpy = vi.spyOn(console, "log").mockImplementation(() => {});
const secret = "never-log-this-runpod-key";
globalThis.fetch = vi.fn().mockResolvedValue({
ok: false,
status: 400,
json: vi.fn().mockResolvedValue({ message: `Rejected key ${secret}` }),
});
try {
const client = createApiClient("http://localhost:8188");
await expect(client.updateWorker("w-secret", { api_key: secret })).rejects.toThrow("[redacted]");
expect(logSpy.mock.calls.flat().join(" ")).not.toContain(secret);
} finally {
globalThis.fetch = originalFetch;
}
});
it("never sends VRAM operations to serverless workers", async () => {
const originalFetch = globalThis.fetch;
globalThis.fetch = vi.fn();
vi.useFakeTimers();
const button = {
textContent: "Clear Worker VRAM", style: { cssText: "" }, disabled: false,
classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
};
const extension = { enabledWorkers: [{ id: "s-vram", name: "Runpod", type: "serverless", enabled: true }] };
try {
await handleClearMemory(extension, button);
expect(globalThis.fetch).not.toHaveBeenCalled();
} finally {
vi.runOnlyPendingTimers();
vi.useRealTimers();
globalThis.fetch = originalFetch;
}
});
it("cancels each enabled serverless worker once without changing normal worker interrupts", async () => {
const originalFetch = globalThis.fetch;
globalThis.fetch = vi.fn().mockResolvedValue({ ok: true });
vi.useFakeTimers();
const button = {
textContent: "Interrupt Workers", style: { cssText: "" }, disabled: false,
classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
};
const workers = [
{ id: "s1", name: "Serverless 1", type: "serverless", enabled: true },
{ id: "s2", name: "Serverless 2", type: "serverless", enabled: true },
{ id: "s3", name: "Disabled", type: "serverless", enabled: false },
{ id: "local", name: "Local", type: "local", enabled: true, port: 8189 },
];
const extension = {
enabledWorkers: workers.filter((worker) => worker.enabled),
api: { cancelServerlessWorker: vi.fn().mockResolvedValue({ ok: true }) },
log: vi.fn(),
};
try {
await handleInterruptWorkers(extension, button);
expect(extension.api.cancelServerlessWorker).toHaveBeenCalledTimes(2);
expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(1, "s1");
expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(2, "s2");
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
expect(globalThis.fetch.mock.calls[0][0]).toBe("http://localhost:8189/interrupt");
} finally {
vi.runOnlyPendingTimers();
vi.useRealTimers();
globalThis.fetch = originalFetch;
}
});
});
-5
View File
@@ -106,11 +106,6 @@ describe("buildWorkerWebSocketUrl", () => {
"wss://worker.example.com/distributed/worker_ws"
);
});
it("always appends /distributed/worker_ws", () => {
const url = buildWorkerWebSocketUrl("http://worker.local:8188");
expect(url.endsWith("/distributed/worker_ws")).toBe(true);
});
});
+3 -25
View File
@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { getWorkerUrl } from "../workerLifecycle.js";
describe("workerLifecycle URL construction", () => {
describe("workerLifecycle URL wiring", () => {
let originalWindow;
beforeEach(() => {
@@ -22,31 +22,9 @@ describe("workerLifecycle URL construction", () => {
globalThis.window = originalWindow;
});
it("builds local worker URL with explicit local port", () => {
// URL variants are covered in urlUtils.test.js; keep the wrapper wiring here.
it("forwards the worker and endpoint using window.location", () => {
const worker = { id: "w1", port: 8189, type: "local" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
});
it("builds remote worker URL with host:port", () => {
const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
});
it("builds cloud worker URL as https", () => {
const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
});
it("rewrites runpod proxy hostname for local worker ports", () => {
globalThis.window = {
location: {
hostname: "podabc.proxy.runpod.net",
protocol: "https:",
port: "",
origin: "https://podabc.proxy.runpod.net",
},
};
const worker = { id: "w4", port: 8189, type: "local" };
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
});
});
+16
View File
@@ -0,0 +1,16 @@
import { describe, expect, it } from "vitest";
import { PULSE_ANIMATION_CSS } from "../constants.js";
describe("worker settings layout", () => {
it("lets expanded forms grow instead of clipping the final controls", () => {
const expanded = PULSE_ANIMATION_CSS.match(/\.worker-settings\.expanded\s*\{([^}]+)\}/)?.[1];
expect(expanded).toBeDefined();
expect(expanded).toMatch(/max-height:\s*none\s*;/);
});
it("still hides collapsed settings", () => {
const collapsed = PULSE_ANIMATION_CSS.match(/\.worker-settings\s*\{([^}]+)\}/)?.[1];
expect(collapsed).toMatch(/max-height:\s*0\s*;/);
expect(collapsed).toMatch(/overflow:\s*hidden\s*;/);
});
});
+18 -2
View File
@@ -7,6 +7,14 @@ import { createLogModal } from './ui/logModal.js';
import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js';
import { isRemoteWorker } from './workerSettings.js';
const escapeWorkerInfo = (value) => String(value ?? "").replace(/[&<>"']/g, (character) => ({
"&": "&amp;",
"<": "&lt;",
">": "&gt;",
'"': "&quot;",
"'": "&#39;",
}[character]));
const cardConfigs = {
master: {
checkbox: {
@@ -58,7 +66,12 @@ const cardConfigs = {
infoText: (data, extension) => {
const isRemote = isRemoteWorker(extension, data);
const isCloud = data.type === 'cloud';
const isServerless = data.type === 'serverless';
if (isServerless) {
return `<span class="dist-worker-info__title">${escapeWorkerInfo(data.name)}</span>`;
}
if (isCloud) {
// For cloud workers, don't show port (it's always 443)
return `<span class="dist-worker-info__title">${data.name}</span><br><span class="dist-worker-info__meta">${data.host}</span>`;
@@ -590,7 +603,10 @@ export class DistributedUI {
}
controlsWrapper.appendChild(badge);
} else if (config.dynamic && data) {
if (isRemote) {
if (data.type === 'serverless') {
// Match the master/remote footer without exposing direct-worker actions.
controlsWrapper.appendChild(this.createInfoBox("Runpod Serverless endpoint"));
} else if (isRemote) {
const isCloud = data.type === 'cloud';
const workerTypeText = isCloud ? "Cloud worker" : "Remote worker";
const workerTypeBadge = this.createInfoBox(workerTypeText);
+12 -1
View File
@@ -1,4 +1,4 @@
import { updateWorkerControls, toggleWorkerExpanded } from "../workerLifecycle.js";
import { updateWorkerControls, toggleWorkerExpanded, formatServerlessStatus } from "../workerLifecycle.js";
import { isRemoteWorker } from "../workerSettings.js";
export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
@@ -74,6 +74,17 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
infoSpan.innerHTML = config.infoText(data, extension);
workerContent.appendChild(infoSpan);
if (isWorker && data.type === "serverless") {
const statusMessage = document.createElement("small");
statusMessage.id = `serverless-status-${data.id}`;
statusMessage.className = "dist-worker-info__meta dist-serverless-status";
statusMessage.setAttribute("aria-live", "polite");
const status = extension.state.getWorker(data.id)?.serverlessStatus;
statusMessage.textContent = data.enabled ? formatServerlessStatus(status) : "Disabled";
statusMessage.title = status?.message || "";
infoSpan.appendChild(statusMessage);
}
infoRow.appendChild(workerContent);
let settingsArrow;
+128 -45
View File
@@ -1,5 +1,5 @@
import { BUTTON_STYLES } from '../constants.js';
import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings } from '../workerSettings.js';
import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings, testServerlessWorker } from '../workerSettings.js';
export function createWorkerSettingsForm(ui, extension, worker) {
const form = document.createElement("div");
@@ -33,9 +33,14 @@ export function createWorkerSettingsForm(ui, extension, worker) {
cloudOption.value = "cloud";
cloudOption.textContent = "Cloud";
const serverlessOption = document.createElement("option");
serverlessOption.value = "serverless";
serverlessOption.textContent = "Serverless (Runpod)";
typeSelect.appendChild(localOption);
typeSelect.appendChild(remoteOption);
typeSelect.appendChild(cloudOption);
typeSelect.appendChild(serverlessOption);
const runpodText = document.createElement("a");
runpodText.id = `runpod-text-${worker.id}`;
@@ -44,50 +49,61 @@ export function createWorkerSettingsForm(ui, extension, worker) {
runpodText.textContent = "Deploy Cloud Worker with Runpod";
runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;";
const createOnChangeHandler = () => {
return (e) => {
const workerType = e.target.value;
const hostGroup = document.getElementById(`host-group-${worker.id}`);
const hostInput = document.getElementById(`host-${worker.id}`);
const portGroup = document.getElementById(`port-group-${worker.id}`);
const portInput = document.getElementById(`port-${worker.id}`);
const cudaGroup = document.getElementById(`cuda-group-${worker.id}`);
const argsGroup = document.getElementById(`args-group-${worker.id}`);
const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`);
const updateWorkerType = (workerType, applyDefaults = false) => {
// The form is initialized before its card is attached to document.
// Use this form's own elements, not document-wide ID lookups.
const hostInput = hostGroup.input;
const portInput = portGroup.input;
const serverlessGroups = [
endpointGroup.group,
apiKeyGroup.group,
clearApiKeyGroup,
jobsPerRunGroup.group,
concurrencyGroup.group,
timeoutGroup.group,
testServerlessGroup,
];
if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) {
return;
}
const hideServerless = () => serverlessGroups.forEach((group) => { group.style.display = "none"; });
if (workerType === "local") {
hostGroup.style.display = "none";
portGroup.style.display = "flex";
cudaGroup.style.display = "flex";
argsGroup.style.display = "flex";
runpodTextElem.style.display = "none";
} else if (workerType === "remote") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex";
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "none";
hostInput.placeholder = "e.g., 192.168.1.100";
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
}
} else if (workerType === "cloud") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex";
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "block";
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
portInput.value = "443";
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
}
if (workerType === "local") {
hideServerless();
hostGroup.group.style.display = "none";
portGroup.group.style.display = "flex";
cudaGroup.group.style.display = "flex";
argsGroup.group.style.display = "flex";
runpodText.style.display = "none";
} else if (workerType === "remote") {
hideServerless();
hostGroup.group.style.display = "flex";
portGroup.group.style.display = "flex";
cudaGroup.group.style.display = "none";
argsGroup.group.style.display = "none";
runpodText.style.display = "none";
hostInput.placeholder = "e.g., 192.168.1.100";
if (applyDefaults && (hostInput.value === "localhost" || hostInput.value === "127.0.0.1")) {
hostInput.value = "";
}
};
} else if (workerType === "cloud") {
hideServerless();
hostGroup.group.style.display = "flex";
portGroup.group.style.display = "flex";
cudaGroup.group.style.display = "none";
argsGroup.group.style.display = "none";
runpodText.style.display = "block";
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
if (applyDefaults) portInput.value = "443";
if (applyDefaults && (hostInput.value === "localhost" || hostInput.value === "127.0.0.1")) {
hostInput.value = "";
}
} else if (workerType === "serverless") {
hostGroup.group.style.display = "none";
portGroup.group.style.display = "none";
cudaGroup.group.style.display = "none";
argsGroup.group.style.display = "none";
runpodText.style.display = "none";
serverlessGroups.forEach((group) => { group.style.display = "flex"; });
}
};
typeGroup.appendChild(typeLabel);
@@ -114,6 +130,71 @@ export function createWorkerSettingsForm(ui, extension, worker) {
argsGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex";
form.appendChild(argsGroup.group);
const endpointGroup = ui.createFormGroup("Endpoint ID:", worker.endpoint_id || "", `endpoint-id-${worker.id}`, "text", "Runpod endpoint ID");
endpointGroup.group.id = `endpoint-id-group-${worker.id}`;
endpointGroup.group.style.display = "none";
form.appendChild(endpointGroup.group);
const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key");
apiKeyGroup.group.id = `api-key-group-${worker.id}`;
apiKeyGroup.input.autocomplete = "new-password";
apiKeyGroup.input.spellcheck = false;
apiKeyGroup.group.style.display = "none";
form.appendChild(apiKeyGroup.group);
const clearApiKeyGroup = document.createElement("div");
clearApiKeyGroup.id = `clear-api-key-group-${worker.id}`;
clearApiKeyGroup.style.cssText = "display: none; align-items: center; gap: 8px;";
const clearApiKeyInput = document.createElement("input");
clearApiKeyInput.type = "checkbox";
clearApiKeyInput.id = `clear-api-key-${worker.id}`;
const clearApiKeyLabel = document.createElement("label");
clearApiKeyLabel.htmlFor = clearApiKeyInput.id;
clearApiKeyLabel.textContent = "Clear saved API key (disable pool first)";
clearApiKeyGroup.appendChild(clearApiKeyInput);
clearApiKeyGroup.appendChild(clearApiKeyLabel);
clearApiKeyInput.onchange = () => {
if (clearApiKeyInput.checked) apiKeyGroup.input.value = "";
apiKeyGroup.input.disabled = clearApiKeyInput.checked;
};
form.appendChild(clearApiKeyGroup);
const jobsPerRunGroup = ui.createFormGroup("Jobs per run:", worker.jobs_per_run ?? 1, `jobs-per-run-${worker.id}`, "number");
jobsPerRunGroup.group.id = `jobs-per-run-group-${worker.id}`;
jobsPerRunGroup.input.min = "1";
jobsPerRunGroup.input.max = "16";
jobsPerRunGroup.input.step = "1";
jobsPerRunGroup.input.title = "Maximum jobs submitted in one run";
jobsPerRunGroup.group.style.display = "none";
form.appendChild(jobsPerRunGroup.group);
const concurrencyGroup = ui.createFormGroup("Concurrent jobs (slots):", worker.concurrency ?? 1, `concurrency-${worker.id}`, "number");
concurrencyGroup.group.id = `concurrency-group-${worker.id}`;
concurrencyGroup.input.min = "1";
concurrencyGroup.input.max = "16";
concurrencyGroup.input.step = "1";
concurrencyGroup.input.title = "Maximum concurrent jobs per serverless worker";
concurrencyGroup.group.style.display = "none";
form.appendChild(concurrencyGroup.group);
const timeoutGroup = ui.createFormGroup("Timeout (seconds):", worker.timeout_seconds ?? 1800, `timeout-seconds-${worker.id}`, "number");
timeoutGroup.group.id = `timeout-seconds-group-${worker.id}`;
timeoutGroup.input.min = "30";
timeoutGroup.input.max = "86400";
timeoutGroup.input.step = "1";
timeoutGroup.group.style.display = "none";
form.appendChild(timeoutGroup.group);
const testServerlessButton = ui.createButton("Test Saved Endpoint", () => testServerlessWorker(extension, worker.id));
testServerlessButton.id = `test-serverless-${worker.id}`;
testServerlessButton.title = "Test the saved Runpod endpoint without submitting inference";
testServerlessButton.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl;
const testServerlessGroup = document.createElement("div");
testServerlessGroup.id = `serverless-test-group-${worker.id}`;
testServerlessGroup.style.cssText = "display: none;";
testServerlessGroup.appendChild(testServerlessButton);
form.appendChild(testServerlessGroup);
const saveBtn = ui.createButton("Save", () => saveWorkerSettings(extension, worker.id), "background-color: #4a7c4a;");
saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success;
@@ -126,9 +207,11 @@ export function createWorkerSettingsForm(ui, extension, worker) {
const buttonGroup = ui.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;");
form.appendChild(buttonGroup);
typeSelect.onchange = createOnChangeHandler();
typeSelect.onchange = (event) => updateWorkerType(event.target.value, true);
if (worker.type === "cloud") {
if (worker.type === "serverless") {
typeSelect.value = "serverless";
} else if (worker.type === "cloud") {
typeSelect.value = "cloud";
runpodText.style.display = "block";
} else if (isRemoteWorker(extension, worker)) {
@@ -137,7 +220,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
typeSelect.value = "local";
}
typeSelect.dispatchEvent(new Event('change'));
updateWorkerType(typeSelect.value);
return form;
}
+106 -1
View File
@@ -41,6 +41,41 @@ function setButtonVisibility(button, visible) {
button.style.display = visible ? "" : "none";
}
export function formatServerlessStatus(status) {
if (!status) return "Checking status…";
if (!status.ok) return "Unavailable";
const jobs = (Array.isArray(status.jobs) ? status.jobs : [])
.map((job) => String(job?.status || "").trim().toUpperCase());
// Prefer current work over old ledger entries, and avoid raw provider enums.
const phases = [
["Downloading result", ["DOWNLOADING"]],
["Cancelling", ["CANCEL_REQUESTED"]],
["Running", ["RUNNING", "PROCESSING", "IN_PROGRESS"]],
["Starting", ["STARTING"]],
["Submitting", ["SUBMITTING"]],
["Queued", ["WAITING", "IN_QUEUE", "QUEUED", "PENDING"]],
];
for (const [label, states] of phases) {
const count = jobs.filter((state) => states.includes(state)).length;
if (count) return count > 1 ? `${label} · ${count} jobs` : label;
}
// Health counts queued AND running jobs, not just the waiting queue.
const remaining = Number(status.queueRemaining);
if (Number.isFinite(remaining) && remaining > 0) {
return `Active · ${remaining} ${remaining === 1 ? "job" : "jobs"}`;
}
if (jobs[0] === "FAILED") return "Last job failed";
if (jobs[0] === "TIMED_OUT") return "Last job timed out";
return "Ready";
}
function updateServerlessStatusMessage(workerId, status) {
const messageElement = document.getElementById(`serverless-status-${workerId}`);
if (!messageElement) return;
messageElement.textContent = formatServerlessStatus(status);
messageElement.title = status?.message || "";
}
export async function checkAllWorkerStatuses(extension) {
if (_statusCheckRunning || !extension.panelElement) {
return;
@@ -136,11 +171,60 @@ export async function checkMasterStatus(extension) {
// Helper to build worker URL
export function getWorkerUrl(extension, worker, endpoint = "") {
if (worker?.type === "serverless") {
return null;
}
return buildWorkerUrl(worker, endpoint, window.location);
}
export async function checkWorkerStatus(extension, worker) {
// Assume caller ensured enabled; proceed with check
if (worker.type === "serverless") {
try {
const result = await extension.api.getServerlessStatus(worker.id);
const queueRemaining = Number(result?.queueRemaining);
const normalizedQueue = Number.isFinite(queueRemaining) ? Math.max(0, queueRemaining) : 0;
const normalizedStatus = {
ok: Boolean(result?.ok),
queueRemaining: normalizedQueue,
message: typeof result?.message === "string" ? result.message : "",
jobs: (Array.isArray(result?.jobs) ? result.jobs : []).map((job) => ({
job_id: String(job?.job_id || ""),
status: String(job?.status || ""),
worker_id: String(job?.worker_id || ""),
execution_id: String(job?.execution_id || ""),
})),
};
const hasActiveJobs = normalizedStatus.jobs.some((job) =>
/^(waiting|submitting|in_queue|queued|pending|running|processing|starting|in_progress|downloading|cancel_requested)$/i.test(job.status.trim())
);
const status = {
online: normalizedStatus.ok,
processing: normalizedQueue > 0 || hasActiveJobs,
queueCount: normalizedQueue,
};
extension.state.updateWorker(worker.id, { serverlessStatus: normalizedStatus });
extension.state.setWorkerStatus(worker.id, status);
updateServerlessStatusMessage(worker.id, normalizedStatus);
const updatedInPlace = extension.updateWorkerCard?.(worker.id, status);
if (!updatedInPlace) {
updateWorkerControls(extension, worker.id);
}
} catch (error) {
if (error?.name === "AbortError") return;
const failedStatus = { ok: false, queueRemaining: 0, message: error?.message || "Runpod status unavailable", jobs: [] };
extension.state.updateWorker(worker.id, { serverlessStatus: failedStatus });
extension.state.setWorkerStatus(worker.id, { online: false, processing: false, queueCount: 0 });
updateServerlessStatusMessage(worker.id, failedStatus);
const updatedInPlace = extension.updateWorkerCard?.(worker.id, extension.state.getWorkerStatus(worker.id));
if (!updatedInPlace) {
updateWorkerControls(extension, worker.id);
}
extension.log(`Serverless worker ${worker.id} status check failed: ${error.message}`, "debug");
}
return;
}
const workerUrl = getWorkerUrl(extension, worker);
try {
@@ -206,6 +290,9 @@ export async function checkWorkerStatus(extension, worker) {
export async function launchWorker(extension, workerId) {
const worker = extension.config.workers.find((w) => w.id === workerId);
if (!worker || worker.type === "serverless") {
return;
}
// If worker is disabled, enable it first
if (!worker.enabled) {
@@ -273,6 +360,9 @@ export async function launchWorker(extension, workerId) {
export async function stopWorker(extension, workerId) {
const worker = extension.config.workers.find((w) => w.id === workerId);
if (!worker || worker.type === "serverless") {
return;
}
const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1];
// Provide immediate feedback
@@ -349,8 +439,12 @@ export async function loadManagedWorkers(extension) {
try {
const result = await extension.api.getManagedWorkers();
// Check for launching workers
const serverlessWorkerIds = new Set(
(extension.config?.workers || []).filter((worker) => worker.type === "serverless").map((worker) => String(worker.id))
);
// Managed-worker lifecycle state is ComfyUI-process-only; serverless is provider-managed.
for (const [workerId, info] of Object.entries(result.managed_workers)) {
if (serverlessWorkerIds.has(String(workerId))) continue;
extension.state.setWorkerManaged(workerId, info);
// If worker is marked as launching, add to launchingWorkers set
@@ -386,6 +480,13 @@ export function updateWorkerControls(extension, workerId) {
const stopBtn = document.getElementById(`stop-${workerId}`);
const logBtn = document.getElementById(`log-${workerId}`);
if (worker.type === "serverless") {
setButtonVisibility(launchBtn, false);
setButtonVisibility(stopBtn, false);
setButtonVisibility(logBtn, false);
return;
}
if (isRemoteWorker(extension, worker)) {
setButtonVisibility(launchBtn, false);
setButtonVisibility(stopBtn, false);
@@ -438,6 +539,9 @@ export function updateWorkerControls(extension, workerId) {
export async function viewWorkerLog(extension, workerId, isRemote = false) {
const worker = extension.config.workers.find((w) => w.id === workerId);
if (worker?.type === "serverless") {
return;
}
const isRemoteLog = isRemote || (worker ? isRemoteWorker(extension, worker) : false);
const managedInfo = extension.state.getWorker(workerId).managed;
if (!isRemoteLog && !managedInfo?.log_file) {
@@ -498,6 +602,7 @@ export async function refreshLog(extension, workerId, silent = false) {
try {
const worker = extension.config.workers.find((w) => w.id === workerId);
if (worker?.type === "serverless") return;
const isRemoteLog = worker ? isRemoteWorker(extension, worker) : false;
const data = isRemoteLog
? await extension.api.getRemoteWorkerLog(workerId, 300)
+149 -2
View File
@@ -18,7 +18,7 @@ export function isRemoteWorker(extension, worker) {
const workerType = String(worker?.type || "").toLowerCase();
// Explicit type always wins over host heuristics.
if (workerType === "cloud" || workerType === "remote") {
if (workerType === "cloud" || workerType === "remote" || workerType === "serverless") {
return true;
}
if (workerType === "local") {
@@ -37,6 +37,134 @@ export function isCloudWorker(extension, worker) {
return worker.type === "cloud";
}
function showServerlessValidationError(extension, detail) {
extension.app.extensionManager.toast.add({
severity: "error",
summary: "Validation Error",
detail,
life: 3000,
});
}
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
const endpointId = String(document.getElementById(`endpoint-id-${workerId}`).value || "").trim();
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
const concurrency = Number(document.getElementById(`concurrency-${workerId}`).value);
const timeoutSeconds = Number(document.getElementById(`timeout-seconds-${workerId}`).value);
if (!name.trim()) {
showServerlessValidationError(extension, "Worker name is required");
return;
}
if (!endpointId) {
showServerlessValidationError(extension, "Endpoint ID is required for Runpod serverless workers");
return;
}
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
showServerlessValidationError(extension, "Jobs per run must be an integer between 1 and 16");
return;
}
if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 16) {
showServerlessValidationError(extension, "Concurrency must be an integer between 1 and 16");
return;
}
if (!Number.isInteger(timeoutSeconds) || timeoutSeconds < 30 || timeoutSeconds > 86400) {
showServerlessValidationError(extension, "Timeout must be an integer between 30 and 86400 seconds");
return;
}
const updateData = {
name: name.trim(),
type: "serverless",
provider: "runpod",
endpoint_id: endpointId,
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
...(clearApiKey ? { clear_api_key: true } : {}),
jobs_per_run: jobsPerRun,
concurrency,
timeout_seconds: timeoutSeconds,
// Never enable a serverless worker as a side effect of saving credentials.
enabled: worker.enabled,
};
try {
await extension.api.updateWorker(workerId, updateData);
document.getElementById(`api-key-${workerId}`).value = "";
worker.name = name.trim();
worker.type = "serverless";
worker.provider = "runpod";
worker.endpoint_id = endpointId;
worker.jobs_per_run = jobsPerRun;
worker.concurrency = concurrency;
worker.timeout_seconds = timeoutSeconds;
worker.api_key_configured = clearApiKey
? false
: apiKey.length > 0 || Boolean(worker.api_key_configured);
// The credential is write-only: never copy it into extension config/state.
delete worker.api_key;
delete worker.host;
delete worker.port;
delete worker.cuda_device;
delete worker.extra_args;
extension.state.updateWorker(workerId, { enabled: worker.enabled });
emitWorkersChanged(extension);
extension.app.extensionManager.toast.add({
severity: "success",
summary: "Settings Saved",
detail: `Worker ${worker.name} settings updated`,
life: 3000,
});
extension.state.setWorkerExpanded(workerId, false);
if (extension.panelElement) {
renderSidebarContent(extension, extension.panelElement);
}
} catch (error) {
extension.app.extensionManager.toast.add({
severity: "error",
summary: "Save Failed",
detail: error.message,
life: 5000,
});
}
}
export async function testServerlessWorker(extension, workerId) {
const button = document.getElementById(`test-serverless-${workerId}`);
const originalText = button?.textContent;
if (button) {
button.disabled = true;
button.textContent = "Testing...";
}
try {
// This endpoint only checks the saved endpoint configuration; it does not enqueue inference.
const result = await extension.api.testServerlessWorker(workerId);
extension.app.extensionManager.toast.add({
severity: result?.ok ? "success" : "error",
summary: result?.ok ? "Runpod Endpoint Ready" : "Runpod Endpoint Test Failed",
detail: result?.message || (result?.ok ? "Saved endpoint is reachable" : "Endpoint test failed"),
life: 4000,
});
return result;
} catch (error) {
extension.app.extensionManager.toast.add({
severity: "error",
summary: "Runpod Endpoint Test Failed",
detail: error.message,
life: 5000,
});
return null;
} finally {
if (button) {
button.disabled = false;
button.textContent = originalText;
}
}
}
export async function saveWorkerSettings(extension, workerId) {
const worker = extension.config.workers.find((w) => w.id === workerId);
if (!worker) {
@@ -44,8 +172,12 @@ export async function saveWorkerSettings(extension, workerId) {
}
// Get form values
const name = document.getElementById(`name-${workerId}`).value;
const workerType = document.getElementById(`worker-type-${workerId}`).value;
if (workerType === "serverless") {
await saveServerlessWorkerSettings(extension, worker, workerId, document.getElementById(`name-${workerId}`).value);
return;
}
const name = document.getElementById(`name-${workerId}`).value;
const isRemote = workerType === "remote" || workerType === "cloud";
const isCloud = workerType === "cloud";
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
@@ -209,6 +341,21 @@ export function cancelWorkerSettings(extension, workerId) {
document.getElementById(`port-${workerId}`).value = worker.port;
document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0;
document.getElementById(`args-${workerId}`).value = worker.extra_args || "";
const endpointInput = document.getElementById(`endpoint-id-${workerId}`);
if (endpointInput) endpointInput.value = worker.endpoint_id || "";
const apiKeyInput = document.getElementById(`api-key-${workerId}`);
if (apiKeyInput) {
apiKeyInput.value = "";
apiKeyInput.disabled = false;
}
const clearApiKeyInput = document.getElementById(`clear-api-key-${workerId}`);
if (clearApiKeyInput) clearApiKeyInput.checked = false;
const jobsPerRunInput = document.getElementById(`jobs-per-run-${workerId}`);
if (jobsPerRunInput) jobsPerRunInput.value = worker.jobs_per_run ?? 1;
const concurrencyInput = document.getElementById(`concurrency-${workerId}`);
if (concurrencyInput) concurrencyInput.value = worker.concurrency ?? 1;
const timeoutInput = document.getElementById(`timeout-seconds-${workerId}`);
if (timeoutInput) timeoutInput.value = worker.timeout_seconds ?? 1800;
// Reset remote checkbox
const remoteCheckbox = document.getElementById(`remote-${workerId}`);
+20 -10
View File
@@ -19,12 +19,17 @@ export async function handleWorkerOperation(extension, button, operation, succes
setButtonStateClass("btn--working");
try {
const urlsToProcess = extension.enabledWorkers.map(w => ({
name: w.name,
url: getWorkerUrl(extension, w)
const workersToProcess = extension.enabledWorkers.filter((worker) =>
operation.endpoint !== ENDPOINTS.CLEAR_MEMORY || worker.type !== "serverless"
);
const targets = workersToProcess.map((worker) => ({
name: worker.name,
worker,
serverless: worker.type === "serverless",
url: worker.type === "serverless" ? null : getWorkerUrl(extension, worker),
}));
if (urlsToProcess.length === 0) {
if (targets.length === 0) {
button.textContent = "No Workers";
setButtonStateClass("btn--error");
setTimeout(() => {
@@ -37,14 +42,19 @@ export async function handleWorkerOperation(extension, button, operation, succes
return;
}
const promises = urlsToProcess.map(target =>
fetch(`${target.url}${operation.endpoint}`, {
method: 'POST',
const promises = targets.map((target) => {
if (target.serverless && operation.name === "Interrupt") {
return extension.api.cancelServerlessWorker(target.worker.id)
.then((result) => ({ ok: result?.ok !== false, name: target.name }))
.catch(() => ({ ok: false, name: target.name }));
}
return fetch(`${target.url}${operation.endpoint}`, {
method: 'POST',
mode: 'cors'
})
.then(response => ({ ok: response.ok, name: target.name }))
.catch(() => ({ ok: false, name: target.name }))
);
.catch(() => ({ ok: false, name: target.name }));
});
const results = await Promise.all(promises);
const failures = results.filter(r => !r.ok);
+2
View File
@@ -10,6 +10,8 @@ from ..utils.logging import debug_log
async def is_local_worker(worker_config):
"""Check if a worker is running on the same machine as the master."""
if worker_config.get('type') == 'serverless':
return False
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':
return True
+97
View File
@@ -0,0 +1,97 @@
"""Host-local worker port selection and launch preflight (no ComfyUI imports)."""
import errno
import os
import socket
from urllib.parse import urlsplit
def _port_number(value):
try:
port = int(value)
except (TypeError, ValueError):
raise ValueError("Worker port must be an integer between 1 and 65535") from None
if isinstance(value, (bool, float)) or not 1 <= port <= 65535:
raise ValueError("Worker port must be an integer between 1 and 65535")
return port
def _local_config(worker):
local_hosts = ("localhost", "127.0.0.1", "0.0.0.0", "::1", "::")
host = (worker.get("host") or "localhost").strip().lower()
if worker.get("type") == "local" or host in local_hosts:
return True
try:
host = urlsplit(host if "://" in host else "//" + host).hostname
except ValueError:
return False
return host in local_hosts
def _reserved_ports(workers, exclude_id=None):
reserved = {}
for worker in workers:
if exclude_id is not None and str(worker.get("id")) == str(exclude_id):
continue
if not _local_config(worker):
continue
try:
port = _port_number(worker.get("port"))
except ValueError:
continue
reserved[port] = worker.get("name") or worker.get("id") or "another local worker"
return reserved
def is_port_available(port):
"""Match asyncio's address reuse while rejecting active IPv4/IPv6 listeners."""
addresses = [(socket.AF_INET, "0.0.0.0")]
if socket.has_ipv6:
addresses.append((socket.AF_INET6, "::"))
for family, address in addresses:
try:
with socket.socket(family, socket.SOCK_STREAM) as probe:
if os.name != "nt":
# Like asyncio, allow a restart while old connections are
# in TIME_WAIT. Do not enable SO_REUSEPORT.
probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
elif hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
# Windows SO_REUSEADDR can steal an active listener's port.
probe.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
probe.bind((address, port))
probe.listen(1)
except OSError as exc:
if family == socket.AF_INET6 and exc.errno in (errno.EAFNOSUPPORT, errno.EADDRNOTAVAIL):
continue
return False
return True
def allocate_worker_ports(master_port, workers, count):
"""Select ports above the actual master, skipping reservations/listeners.
This is a snapshot, not a reservation. Launch must validate again because
another process can claim a suggested port before the user starts a worker.
"""
master_port = _port_number(master_port)
reserved = _reserved_ports(workers)
selected = []
if count == 0:
return selected
for port in range(master_port + 1, 65536):
if port not in reserved and is_port_available(port):
selected.append(port)
if len(selected) == count:
return selected
raise ValueError(f"Not enough available worker ports above master port {master_port}; need {count}")
def validate_worker_port(worker, master_port, workers):
"""Reject conflicts without silently changing a configured worker port."""
port = _port_number(worker.get("port"))
if port == _port_number(master_port):
raise ValueError(f"Worker port {port} conflicts with the master port {master_port}; choose another worker port")
reserved = _reserved_ports(workers, exclude_id=worker.get("id"))
if port in reserved:
raise ValueError(f"Worker port {port} is assigned to local worker {reserved[port]}; choose another worker port")
if not is_port_available(port):
raise ValueError(f"Worker port {port} is already in use on this host; choose another worker port")
+49
View File
@@ -1,4 +1,5 @@
import glob
import hashlib
import os
import shlex
import shutil
@@ -87,6 +88,51 @@ class LaunchCommandBuilder:
return wt_path
return None
@staticmethod
def _arg_value(cmd, flag):
"""Return the last value, matching argparse's override semantics."""
value = None
for index, arg in enumerate(cmd):
if arg == flag and index + 1 < len(cmd):
value = cmd[index + 1]
elif arg.startswith(flag + "="):
value = arg.split("=", 1)[1]
return value
def _add_worker_database(self, cmd, worker_config, comfy_root):
args = self._get_runtime_args()
# The parsed namespace advertises support without launching ComfyUI
# (or importing another checkout's argument parser).
if args is None or not hasattr(args, "database_url"):
return
if any(arg.split("=", 1)[0] in ("--database-url", "--disable-assets") for arg in cmd):
return
user_directory = self._arg_value(cmd, "--user-directory")
base_directory = self._arg_value(cmd, "--base-directory")
if not user_directory:
user_directory = os.path.join(base_directory or comfy_root, "user")
if not os.path.isabs(user_directory):
# Popen uses comfy_root as cwd, which can differ from the master's.
user_directory = os.path.join(comfy_root, user_directory)
database_directory = os.path.abspath(os.path.join(user_directory, "distributed", "workers"))
# Hash the immutable ID: names/GPU/ports may change, and IDs must not
# become filesystem paths or collide after filename sanitization.
worker_id = str(worker_config["id"])
filename = hashlib.sha256(worker_id.encode("utf-8")).hexdigest() + ".db"
database_path = os.path.join(database_directory, filename).replace(os.sep, "/")
# Use the same URL codec as ComfyUI; raw '?' and '%xx' can otherwise
# truncate/decode the path and make distinct workers share a database.
from sqlalchemy.engine import URL, make_url
database_url = URL.create("sqlite", database=database_path).render_as_string()
if make_url(database_url).database != database_path:
raise ValueError(
"Installed SQLAlchemy cannot safely encode the worker database path; "
"use a user directory without '?' or an explicit --database-url with a safe path"
)
os.makedirs(database_directory, exist_ok=True)
cmd.extend(["--database-url", database_url])
def build_launch_command(self, worker_config, comfy_root):
"""Build the command to launch a worker."""
main_py = os.path.join(comfy_root, "main.py")
@@ -141,4 +187,7 @@ class LaunchCommandBuilder:
raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {forbidden}")
cmd.extend(extra_args_list)
if int(self._arg_value(cmd, "--port")) != int(worker_config["port"]):
raise ValueError("extra_args must not override the configured worker port; set the port in the UI/config instead")
self._add_worker_database(cmd, worker_config, comfy_root)
return cmd
+13
View File
@@ -7,7 +7,9 @@ import time
from ...utils.config import load_config, save_config
from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL
from ...utils.logging import debug_log, log
from ...utils.network import get_server_port
from ...utils.process import get_python_executable, is_process_alive, terminate_process
from ..ports import validate_worker_port
try:
import psutil
@@ -28,6 +30,17 @@ class ProcessLifecycle:
"""Launch a worker process with logging."""
_ = show_window # Kept for API compatibility.
comfy_root = self._manager.find_comfy_root()
config = load_config()
try:
master_port = get_server_port()
except AttributeError:
master_port = None
if master_port is None:
# The auto-launch timer can fire during custom-node imports, before
# PromptServer assigns its port. Still reserve the configured port.
from comfy.cli_args import args
master_port = args.port
validate_worker_port(worker_config, master_port, config.get("workers", []))
env = os.environ.copy()
env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))
+2
View File
@@ -22,6 +22,8 @@ class WorkerProcessManager:
return self._launch_builder.build_launch_command(worker_config, comfy_root)
def launch_worker(self, worker_config, show_window=False):
if worker_config.get("type") == "serverless":
raise ValueError("Serverless pools cannot be launched as local processes")
return self._lifecycle.launch_worker(worker_config, show_window=show_window)
def stop_worker(self, worker_id):
+2
View File
@@ -35,6 +35,8 @@ def auto_launch_workers():
launched_count = 0
for worker in workers:
if worker.get('type') == 'serverless':
continue
if worker.get('enabled', False):
worker_id = worker.get('id')
worker_name = worker.get('name', f'Worker {worker_id}')

Some files were not shown because too many files have changed in this diff Show More