Compare commits
43
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
593047badf | ||
|
|
a9ef36f951 | ||
|
|
6db5692d78 | ||
|
|
591ffc7396 | ||
|
|
5f9a2f0fc6 | ||
|
|
b636fd5202 | ||
|
|
f0927f1166 | ||
|
|
c30036ae31 | ||
|
|
41caf02481 | ||
|
|
494d5506a0 | ||
|
|
7ed60b93ca | ||
|
|
6b94242110 | ||
|
|
40678a51c4 | ||
|
|
9345a8e71b | ||
|
|
a630f55bd9 | ||
|
|
6edd3fda38 | ||
|
|
6f85129e45 | ||
|
|
a53d88f7fb | ||
|
|
9be342f5fb | ||
|
|
a1faf7e916 | ||
|
|
32ac027e0e | ||
|
|
fec959e1ac | ||
|
|
cedc8d4591 | ||
|
|
db01824d6e | ||
|
|
6e7818bee6 | ||
|
|
c246d965e7 | ||
|
|
d6ce26f143 | ||
|
|
e2125c7529 | ||
|
|
a91f9fb081 | ||
|
|
6874d2735f | ||
|
|
0e9ee7fbac | ||
|
|
5386de10e3 | ||
|
|
41f9e44945 | ||
|
|
7792d613bb | ||
|
|
6e698512b7 | ||
|
|
79cc9f5ad5 | ||
|
|
4eba87cba3 | ||
|
|
27e08de94c | ||
|
|
c8453a139b | ||
|
|
aae831e1e5 | ||
|
|
e7ab67733b | ||
|
|
dd55ff740e | ||
|
|
a6d0b82d35 |
@@ -1,4 +1,3 @@
|
||||
# These are supported funding model platforms
|
||||
|
||||
github: robertvoy
|
||||
buy_me_a_coffee: robertvoy
|
||||
|
||||
@@ -19,6 +19,6 @@ jobs:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
@@ -7,8 +7,3 @@ __pycache__/
|
||||
node_modules/
|
||||
npm-debug.log*
|
||||
AGENTS.md
|
||||
.claude/
|
||||
|
||||
# Desloppify artifacts
|
||||
.desloppify/
|
||||
scorecard.png
|
||||
|
||||
@@ -15,7 +15,7 @@
|
||||
|
||||
---
|
||||
|
||||
## Key Features
|
||||
## Key Features
|
||||
|
||||
#### Parallel Workflow Processing
|
||||
- Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
|
||||
@@ -27,32 +27,28 @@
|
||||
- Intelligent distribution
|
||||
- Handles single images and videos
|
||||
|
||||
#### Ease of Use
|
||||
- Auto-setup local workers; easily add remote/cloud ones
|
||||
- Convert any workflow to distributed with 2 nodes
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Current Architecture
|
||||
|
||||
- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
|
||||
- There is **no Distributed Queue node** anymore.
|
||||
- With `load_balance=true`, orchestration selects one least-busy execution participant:
|
||||
- If master participation is enabled, master is included as a candidate.
|
||||
- If master is in orchestrator-only mode, only workers are considered.
|
||||
|
||||
---
|
||||
#### Dynamic Image-List Processing
|
||||
- Process a whole image list in parallel: the master and every worker take the next image when they free up, and results come back 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
|
||||
- JSON configuration with UI controls
|
||||
|
||||
---
|
||||
|
||||
## Worker Types
|
||||
|
||||
<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 endpoints for image Collector jobs; no Distributed install or inbound tunnel on the endpoint. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
@@ -61,7 +57,6 @@ ComfyUI Distributed supports three types of workers:
|
||||
## Requirements
|
||||
|
||||
- ComfyUI
|
||||
> Note: Desktop app not currently supported
|
||||
- Multiple NVIDIA GPUs
|
||||
> No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
|
||||
- That's it
|
||||
@@ -92,19 +87,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
|
||||
|
||||
## Workflow Examples
|
||||
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
### Basic Parallel Generation
|
||||
Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
|
||||
|
||||

|
||||
|
||||
> [Download workflow](/workflows/distributed-txt2img.json)
|
||||
|
||||
1. Open your ComfyUI workflow
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
|
||||
### Parallel WAN Generation
|
||||
Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
|
||||
@@ -122,12 +117,12 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
|
||||
7. Enable workers in the UI
|
||||
8. Run the workflow!
|
||||
|
||||
### Distributed Image Upscaling
|
||||
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
|
||||
### Distributed Image Upscaling
|
||||
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
|
||||
|
||||

|
||||
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
|
||||
1. Load your image
|
||||
2. Upscale with ESRGAN or similar
|
||||
@@ -153,42 +148,44 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
||||
|
||||
---
|
||||
|
||||
## Developer API
|
||||
## Developer API
|
||||
|
||||
Control your distributed cluster programmatically without opening the browser.
|
||||
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **Functionality:** Accepts a distributed queue payload (`prompt`, `client_id`, `enabled_worker_ids`, optional `workflow`/`delegate_master`/`trace_execution_id`), dispatches to healthy workers, and returns `{prompt_id, worker_count}`.
|
||||
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **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).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
|
||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
| **Distributed Image List Start** | Turns an image list or batch into one job per image, scheduled to the next available participant; mask and value inputs are shared across images or paired one per image |
|
||||
| **Distributed Image List End** | Closes the paired region and returns images, masks, and texts as lists in the original input order |
|
||||
|
||||
---
|
||||
|
||||
@@ -206,7 +203,7 @@ No, it does not speed up the generation of a single image or video. Instead, it
|
||||
|
||||
<details>
|
||||
<summary>Does it work with the ComfyUI desktop app?</summary>
|
||||
Currently, it is not compatible with the ComfyUI desktop app.
|
||||
Yes, it does now.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -231,6 +228,26 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do
|
||||
|
||||
---
|
||||
|
||||
## Community projects
|
||||
|
||||
Community-maintained add-ons and forks that extend ComfyUI-Distributed.
|
||||
|
||||
### [ComfyUI-H3MC-Distributed](https://github.com/ymiyad/ComfyUI-H3MC-Distributed) by ymiyad
|
||||
|
||||
Connects H3 Motion Context with ComfyUI-Distributed to generate multiple candidate takes of the next video clip across your workers. Compare the results, pick your preferred take, and use it as context for the next clip.
|
||||
|
||||
Each saved candidate video includes its continuation latent, so you don't need a shared folder or separate latent files. This add-on speeds up iteration through parallel candidates; it does not split a single clip across GPUs.
|
||||
|
||||
### [ComfyUI-Deadline-Distributed](https://github.com/doubletwisted/ComfyUI-Deadline-Distributed) by doubletwisted
|
||||
|
||||
A fork of ComfyUI-Distributed with Thinkbox/AWS Deadline 10 integration for render farms. Claim Deadline workers from the Distributed panel, register them automatically, and collect their results back into your master workflow.
|
||||
|
||||
Requires Deadline 10 and the author's ComfyUI-Deadline-Plugin.
|
||||
|
||||
See each repository for installation instructions, requirements, and support.
|
||||
|
||||
---
|
||||
|
||||
## Disclaimer
|
||||
|
||||
This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats.
|
||||
@@ -249,3 +266,4 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
+18
-7
@@ -1,11 +1,22 @@
|
||||
"""ComfyUI-Distributed package entrypoint."""
|
||||
from __future__ import annotations
|
||||
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||
|
||||
from .bootstrap.entrypoint import build_node_mappings, initialize_runtime
|
||||
from .nodes.v3 import NODES
|
||||
from .runtime.bootstrap import initialize
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
WEB_DIRECTORY = './web'
|
||||
|
||||
NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
|
||||
initialize_runtime()
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
class DistributedExtension(ComfyExtension):
|
||||
async def on_load(self) -> None:
|
||||
initialize()
|
||||
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return list(NODES)
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> DistributedExtension:
|
||||
return DistributedExtension()
|
||||
|
||||
|
||||
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||
|
||||
+7
-21
@@ -1,21 +1,7 @@
|
||||
"""Distributed API package with explicit route bootstrap."""
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
_ROUTE_MODULES = (
|
||||
"config_routes",
|
||||
"tunnel_routes",
|
||||
"worker_routes",
|
||||
"job_routes",
|
||||
"usdu_routes",
|
||||
)
|
||||
|
||||
|
||||
def bootstrap_routes() -> None:
|
||||
"""Import route modules once so aiohttp decorators register endpoints."""
|
||||
for module_name in _ROUTE_MODULES:
|
||||
import_module(f"{__name__}.{module_name}")
|
||||
|
||||
|
||||
__all__ = ["bootstrap_routes"]
|
||||
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
|
||||
|
||||
+438
-143
@@ -1,15 +1,39 @@
|
||||
from __future__ import annotations
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import config_transaction, load_config
|
||||
from ..utils.logging import debug_log
|
||||
try:
|
||||
from ..utils.config import config_transaction, load_config, save_config
|
||||
except ImportError:
|
||||
from ..utils.config import load_config
|
||||
|
||||
try:
|
||||
from ..utils.config import save_config
|
||||
except ImportError:
|
||||
def save_config(_config):
|
||||
return True
|
||||
|
||||
@asynccontextmanager
|
||||
async def config_transaction():
|
||||
config = load_config()
|
||||
original_snapshot = json.dumps(config, sort_keys=True)
|
||||
yield config
|
||||
if json.dumps(config, sort_keys=True) != original_snapshot:
|
||||
save_config(config)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host
|
||||
from .endpoint_policy import run_authorized_endpoint
|
||||
from ..utils.secret_fields import is_secret_field_name
|
||||
from ..utils.serverless_config import (
|
||||
ServerlessSecretChanges,
|
||||
public_config,
|
||||
validate_serverless_worker,
|
||||
)
|
||||
|
||||
|
||||
def _positive_int(value: int) -> bool:
|
||||
def _positive_int(value):
|
||||
return value > 0
|
||||
|
||||
|
||||
@@ -69,103 +93,247 @@ 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
|
||||
|
||||
|
||||
def _validate_workers_payload(workers_value: object) -> list[str]:
|
||||
"""Validate top-level workers payload shape before persisting config."""
|
||||
if not isinstance(workers_value, list):
|
||||
return ["workers: expected list"]
|
||||
_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",
|
||||
"media_profile",
|
||||
}
|
||||
_SERVERLESS_CONFIG_FIELDS = {
|
||||
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
|
||||
}
|
||||
|
||||
errors: list[str] = []
|
||||
seen_ids: set[str] = set()
|
||||
for index, worker in enumerate(workers_value):
|
||||
if not isinstance(worker, dict):
|
||||
errors.append(f"workers[{index}]: expected object")
|
||||
continue
|
||||
|
||||
worker_id = str(worker.get("id", "")).strip()
|
||||
if not worker_id:
|
||||
errors.append(f"workers[{index}].id: expected non-empty string")
|
||||
elif worker_id in seen_ids:
|
||||
errors.append(f"workers[{index}].id: duplicate id '{worker_id}'")
|
||||
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:
|
||||
seen_ids.add(worker_id)
|
||||
validated.append(worker)
|
||||
return validated
|
||||
|
||||
if "port" in worker and worker.get("port") is not None:
|
||||
try:
|
||||
port = int(worker.get("port"))
|
||||
except (TypeError, ValueError):
|
||||
errors.append(f"workers[{index}].port: expected integer")
|
||||
else:
|
||||
if port <= 0:
|
||||
errors.append(f"workers[{index}].port: expected positive integer")
|
||||
|
||||
return errors
|
||||
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":
|
||||
continue
|
||||
if worker.get("type") == "serverless":
|
||||
# Older clients do not know this field. Only an explicit selection
|
||||
# may change the endpoint contract, including on bulk saves.
|
||||
if "media_profile" not in worker and "media_profile" in old:
|
||||
worker["media_profile"] = old["media_profile"]
|
||||
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("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
|
||||
):
|
||||
_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: web.Request) -> web.StreamResponse:
|
||||
async def _operation() -> web.StreamResponse:
|
||||
return web.json_response(load_config())
|
||||
|
||||
return await run_authorized_endpoint(request, _operation)
|
||||
async def get_config_endpoint(request):
|
||||
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: web.Request) -> web.StreamResponse:
|
||||
"""Bulk config update with schema validation."""
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
async def update_config_endpoint(request):
|
||||
"""Bulk config update with schema validation and credential redaction."""
|
||||
try:
|
||||
data = await request.json()
|
||||
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)
|
||||
if not isinstance(data, dict):
|
||||
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 = {}
|
||||
errors = []
|
||||
validated_settings = {}
|
||||
validated_root = {}
|
||||
errors = []
|
||||
|
||||
for key, value in data.items():
|
||||
if key not in CONFIG_SCHEMA:
|
||||
errors.append(f"Unknown field: {key}")
|
||||
continue
|
||||
for key, value in data.items():
|
||||
if key not in CONFIG_SCHEMA:
|
||||
errors.append(f"Unknown field: {key}")
|
||||
continue
|
||||
|
||||
expected_type, validator = CONFIG_SCHEMA[key]
|
||||
if not isinstance(value, expected_type):
|
||||
errors.append(f"{key}: expected {expected_type.__name__}")
|
||||
continue
|
||||
expected_type, validator = CONFIG_SCHEMA[key]
|
||||
if not isinstance(value, expected_type):
|
||||
errors.append(f"{key}: expected {expected_type.__name__}")
|
||||
continue
|
||||
|
||||
if validator and not validator(value):
|
||||
errors.append(f"{key}: value {value!r} failed validation")
|
||||
continue
|
||||
if validator and not validator(value):
|
||||
errors.append(f"{key}: failed validation")
|
||||
continue
|
||||
|
||||
if key == "workers":
|
||||
errors.extend(_validate_workers_payload(value))
|
||||
if errors:
|
||||
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:
|
||||
validated_root[key] = value
|
||||
if key in _SETTINGS_FIELDS:
|
||||
validated_settings[key] = value
|
||||
else:
|
||||
validated_root[key] = value
|
||||
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
if errors:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"error": errors,
|
||||
"message": "; ".join(errors),
|
||||
}, status=400)
|
||||
|
||||
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})
|
||||
store = None
|
||||
response_config = None
|
||||
try:
|
||||
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 worker", 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)
|
||||
|
||||
return await run_authorized_endpoint(request, _operation)
|
||||
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 worker", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to update config", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
||||
async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def queue_status_endpoint(request):
|
||||
"""Check if a job queue is initialized."""
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
job_id = request.match_info['job_id']
|
||||
|
||||
# Import to ensure initialization
|
||||
@@ -177,89 +345,216 @@ async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
|
||||
return web.json_response({"exists": exists, "job_id": job_id})
|
||||
|
||||
return await run_authorized_endpoint(request, _operation, unexpected_status=500)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
||||
async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def _operation() -> web.StreamResponse:
|
||||
data = await request.json()
|
||||
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:
|
||||
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)
|
||||
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
|
||||
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")
|
||||
)
|
||||
media_profile_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
|
||||
)
|
||||
if current_type == "serverless" and (
|
||||
endpoint_changed or media_profile_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 worker; turn the worker off before clearing its 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}
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
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 worker", 409)
|
||||
except ValueError:
|
||||
return _error_response("Invalid worker configuration", 400)
|
||||
except Exception:
|
||||
return _error_response("Unable to update worker configuration", 500)
|
||||
|
||||
return await run_authorized_endpoint(request, _operation, unexpected_status=500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
||||
async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def _operation() -> web.StreamResponse:
|
||||
data = await request.json()
|
||||
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:
|
||||
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"
|
||||
})
|
||||
"message": "Worker deleted" if _worker_type(removed) == "serverless"
|
||||
else f"Worker {removed.get('name', worker_id)} deleted",
|
||||
}
|
||||
|
||||
return await run_authorized_endpoint(request, _operation, unexpected_status=500)
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 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: web.Request) -> web.StreamResponse:
|
||||
async def update_setting_endpoint(request):
|
||||
"""Updates a specific key in the settings object."""
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
data = await request.json()
|
||||
key = data.get("key")
|
||||
value = data.get("value")
|
||||
@@ -276,13 +571,13 @@ async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
config['settings'][key] = value
|
||||
|
||||
return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
|
||||
|
||||
return await run_authorized_endpoint(request, _operation, unexpected_status=500)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_master")
|
||||
async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def update_master_endpoint(request):
|
||||
"""Updates master configuration."""
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
data = await request.json()
|
||||
|
||||
async with config_transaction() as config:
|
||||
@@ -291,5 +586,5 @@ async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
_apply_field_patch(config['master'], data, _MASTER_FIELDS)
|
||||
|
||||
return web.json_response({"status": "success", "message": "Master configuration updated."})
|
||||
|
||||
return await run_authorized_endpoint(request, _operation, unexpected_status=500)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
"""Shared endpoint policy helpers for distributed API routes."""
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from ..utils.network import handle_api_error
|
||||
from .request_guards import authorization_error_or_none
|
||||
|
||||
|
||||
async def run_authorized_endpoint(
|
||||
request: web.Request,
|
||||
operation: Callable[[], Awaitable[web.StreamResponse]],
|
||||
*,
|
||||
unexpected_status: int = 500,
|
||||
) -> web.StreamResponse:
|
||||
"""Run endpoint operation with shared auth guard and fallback error mapping."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
return await operation()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, unexpected_status)
|
||||
@@ -0,0 +1,92 @@
|
||||
"""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,
|
||||
PROTOCOL_VERSION,
|
||||
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,
|
||||
"execution_targets": True,
|
||||
"value_results": True,
|
||||
"protocol_version": PROTOCOL_VERSION,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd", "DistributedImageListGetValue"],
|
||||
"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")
|
||||
identity = store_worker_item(data.get("token"), item)
|
||||
return web.json_response({"status": "stored", "item_index": item["index"],
|
||||
"item_identity": identity})
|
||||
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)
|
||||
+233
-109
@@ -1,5 +1,3 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import io
|
||||
@@ -7,7 +5,6 @@ import os
|
||||
import base64
|
||||
import binascii
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
@@ -18,21 +15,33 @@ from ..utils.logging import debug_log
|
||||
from ..utils.image import pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import handle_api_error
|
||||
from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
|
||||
from ..utils.runtime_state import ensure_distributed_runtime_state
|
||||
from .request_guards import authorization_error_or_none
|
||||
from .schemas import require_bool_literal
|
||||
from .queue_orchestration import orchestrate_distributed_execution
|
||||
from ..utils import transfer_codec
|
||||
try:
|
||||
from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
|
||||
except ImportError:
|
||||
from .queue_orchestration import orchestrate_distributed_execution
|
||||
|
||||
def ensure_distributed_state():
|
||||
return None
|
||||
from .queue_request import parse_queue_request_payload
|
||||
|
||||
# Canonical worker result envelope accepted by POST /distributed/job_complete:
|
||||
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
# Worker result envelopes accepted by POST /distributed/job_complete:
|
||||
#
|
||||
# application/json (legacy, one frame per request):
|
||||
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool,
|
||||
# "audio"?: {...} }
|
||||
#
|
||||
# multipart/form-data (compressed, many frames per request):
|
||||
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
|
||||
# "is_last", "audio"?: {...} }
|
||||
# media: one part per frame for still formats, or one video file for video formats
|
||||
# format "exact" adds envelope "descriptors" (one tensor-le-v1 descriptor per
|
||||
# frame) and sends one raw tensor part per frame; values are not quantized.
|
||||
|
||||
|
||||
def _runtime_state():
|
||||
return ensure_distributed_runtime_state()
|
||||
|
||||
|
||||
def _decode_image_sync(image_path: str) -> dict[str, Any]:
|
||||
def _decode_image_sync(image_path):
|
||||
"""Decode image/video file and compute hash in a threadpool worker."""
|
||||
import base64
|
||||
import hashlib
|
||||
@@ -78,7 +87,7 @@ def _decode_image_sync(image_path: str) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
|
||||
def _check_file_sync(filename, expected_hash):
|
||||
"""Check file presence and hash in a threadpool worker."""
|
||||
import hashlib
|
||||
import folder_paths
|
||||
@@ -103,7 +112,7 @@ def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
|
||||
def _decode_canonical_png_tensor(image_payload):
|
||||
"""Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
|
||||
if not isinstance(image_payload, str) or not image_payload.strip():
|
||||
raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
|
||||
@@ -134,7 +143,7 @@ def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
|
||||
raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
|
||||
|
||||
|
||||
def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
|
||||
def _decode_audio_payload(audio_payload):
|
||||
"""Decode canonical audio payload into an AUDIO dict."""
|
||||
from ..utils.audio_payload import decode_audio_payload
|
||||
|
||||
@@ -142,20 +151,17 @@ def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/prepare_job")
|
||||
async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
async def prepare_job_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id", 400)
|
||||
|
||||
runtime_state = _runtime_state()
|
||||
async with runtime_state.distributed_jobs_lock:
|
||||
if multi_job_id not in runtime_state.distributed_pending_jobs:
|
||||
runtime_state.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
|
||||
debug_log(f"Prepared queue for job {multi_job_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
@@ -163,13 +169,9 @@ async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/clear_memory")
|
||||
async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def clear_memory_endpoint(request):
|
||||
debug_log("Received request to clear VRAM.")
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
warnings: list[str] = []
|
||||
# Use ComfyUI's prompt server queue system like the /free endpoint does
|
||||
if hasattr(server.PromptServer.instance, 'prompt_queue'):
|
||||
server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
|
||||
@@ -186,16 +188,12 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
try:
|
||||
mm.unload_all_models()
|
||||
except AttributeError as e:
|
||||
warning = f"Model unload warning: {e}"
|
||||
warnings.append(warning)
|
||||
debug_log(warning)
|
||||
debug_log(f"Warning during model unload: {e}")
|
||||
|
||||
try:
|
||||
mm.soft_empty_cache()
|
||||
except Exception as e:
|
||||
warning = f"Cache clear warning: {e}"
|
||||
warnings.append(warning)
|
||||
debug_log(warning)
|
||||
debug_log(f"Warning during cache clear: {e}")
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
@@ -204,16 +202,6 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
if warnings:
|
||||
debug_log("VRAM cleared with warnings.")
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "partial",
|
||||
"message": "GPU memory cleared with warnings.",
|
||||
"warnings": warnings,
|
||||
}
|
||||
)
|
||||
|
||||
debug_log("VRAM cleared successfully.")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared."})
|
||||
except Exception as e:
|
||||
@@ -222,16 +210,13 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
debug_log(f"VRAM clear failed: {e}")
|
||||
return await handle_api_error(request, f"GPU memory clear failed: {e}", 500)
|
||||
debug_log(f"Partial VRAM clear completed with warning: {e}")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/queue")
|
||||
async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def distributed_queue_endpoint(request):
|
||||
"""Queue a distributed workflow, mirroring the UI orchestration pipeline."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
raw_payload = await request.json()
|
||||
except Exception as exc:
|
||||
@@ -243,7 +228,7 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
try:
|
||||
prompt_id, worker_count = await orchestrate_distributed_execution(
|
||||
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
|
||||
payload.prompt,
|
||||
payload.workflow_meta,
|
||||
payload.client_id,
|
||||
@@ -253,17 +238,17 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"number": prompt_number,
|
||||
"node_errors": node_errors,
|
||||
"worker_count": worker_count,
|
||||
"auto_prepare_supported": True,
|
||||
})
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/load_image")
|
||||
async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def load_image_endpoint(request):
|
||||
"""Load an image or video file and return it as base64 data with hash."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
data = await request.json()
|
||||
image_path = data.get("image_path")
|
||||
@@ -279,11 +264,8 @@ async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/check_file")
|
||||
async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def check_file_endpoint(request):
|
||||
"""Check if a file exists and matches the given hash."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
data = await request.json()
|
||||
filename = data.get("filename")
|
||||
@@ -299,11 +281,168 @@ async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
async def _enqueue_worker_items(multi_job_id, items):
|
||||
"""Put decoded items on the job queue, waiting briefly for the master to create it.
|
||||
|
||||
Returns the queue size after enqueueing, or None if the job never appeared.
|
||||
"""
|
||||
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
|
||||
while True:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
if pending is not None:
|
||||
for item in items:
|
||||
await pending.put(item)
|
||||
return pending.qsize()
|
||||
if time.monotonic() > deadline:
|
||||
return None
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
def _decode_compressed_chunk_sync(envelope, parts):
|
||||
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
|
||||
if envelope.get("format") == transfer_codec.EXACT_FORMAT:
|
||||
return _decode_exact_chunk_sync(envelope, parts)
|
||||
frames = transfer_codec.decode_chunk(envelope, parts)
|
||||
return [
|
||||
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
|
||||
for frame in frames
|
||||
]
|
||||
|
||||
|
||||
def _decode_exact_chunk_sync(envelope, parts):
|
||||
"""Decode raw exact IMAGE frames, checking each against the envelope geometry."""
|
||||
from ..serverless_bridge import codec as bridge_codec
|
||||
count, width, height = envelope.get("count"), envelope.get("width"), envelope.get("height")
|
||||
for name, value in (("count", count), ("width", width), ("height", height)):
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
|
||||
raise transfer_codec.TransferCodecError(f"envelope.{name}: expected positive integer")
|
||||
descriptors = envelope.get("descriptors")
|
||||
if not isinstance(descriptors, list) or len(descriptors) != count or len(parts) != count:
|
||||
raise transfer_codec.TransferCodecError(
|
||||
f"exact chunk expected {count} descriptors and frames, got "
|
||||
f"{len(descriptors) if isinstance(descriptors, list) else 'no'} descriptors and {len(parts)} frames"
|
||||
)
|
||||
tensors = []
|
||||
for descriptor, payload in zip(descriptors, parts):
|
||||
if not isinstance(descriptor, dict) or descriptor.get("encoding") != "tensor-le-v1":
|
||||
raise transfer_codec.TransferCodecError("exact frames require tensor-le-v1 encoding")
|
||||
shape = descriptor.get("shape") if isinstance(descriptor, dict) else None
|
||||
if not isinstance(shape, list) or len(shape) != 4 or shape[:3] != [1, height, width]:
|
||||
raise transfer_codec.TransferCodecError(
|
||||
f"exact frame shape {shape} does not match the envelope's 1x{height}x{width} frames"
|
||||
)
|
||||
try:
|
||||
tensors.append(bridge_codec.decode_payload("IMAGE", descriptor, payload))
|
||||
except ValueError as exc:
|
||||
raise transfer_codec.TransferCodecError(f"exact frame: {exc}") from None
|
||||
return tensors
|
||||
|
||||
|
||||
def _validate_compressed_envelope(envelope):
|
||||
errors = []
|
||||
if not isinstance(envelope, dict):
|
||||
return ["envelope: expected JSON object"]
|
||||
for key in ("job_id", "worker_id"):
|
||||
value = envelope.get(key)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
errors.append(f"{key}: expected non-empty string")
|
||||
start_index = envelope.get("start_index")
|
||||
if not isinstance(start_index, int) or isinstance(start_index, bool) or start_index < 0:
|
||||
errors.append("start_index: expected non-negative integer")
|
||||
if not isinstance(envelope.get("is_last"), bool):
|
||||
errors.append("is_last: expected boolean")
|
||||
try:
|
||||
if transfer_codec.normalize_format(envelope.get("format")) == transfer_codec.LEGACY_FORMAT:
|
||||
errors.append("format: legacy_png must be sent as JSON")
|
||||
except transfer_codec.TransferCodecError as exc:
|
||||
errors.append(f"format: {exc}")
|
||||
audio_payload = envelope.get("audio")
|
||||
if audio_payload is not None and not isinstance(audio_payload, dict):
|
||||
errors.append("audio: expected object when provided")
|
||||
return errors
|
||||
|
||||
|
||||
async def _job_complete_compressed(request):
|
||||
try:
|
||||
form = await request.post()
|
||||
except Exception as exc:
|
||||
too_large = getattr(web, "HTTPRequestEntityTooLarge", None)
|
||||
if too_large is not None and isinstance(exc, too_large):
|
||||
return await handle_api_error(
|
||||
request, f"Request body exceeds the master's upload limit (--max-upload-size): {exc}", 413)
|
||||
return await handle_api_error(request, f"Invalid multipart payload: {exc}", 400)
|
||||
|
||||
raw_envelope = form.get("envelope")
|
||||
if isinstance(raw_envelope, (bytes, bytearray)):
|
||||
raw_envelope = raw_envelope.decode("utf-8")
|
||||
elif hasattr(raw_envelope, "file"):
|
||||
raw_envelope = raw_envelope.file.read().decode("utf-8")
|
||||
try:
|
||||
envelope = json.loads(raw_envelope) if isinstance(raw_envelope, str) else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, f"envelope: invalid JSON ({exc})", 400)
|
||||
if envelope is None:
|
||||
return await handle_api_error(request, "envelope: missing field", 400)
|
||||
|
||||
errors = _validate_compressed_envelope(envelope)
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
|
||||
parts = []
|
||||
for field in form.getall("media", []):
|
||||
if hasattr(field, "file"):
|
||||
parts.append(field.file.read())
|
||||
elif isinstance(field, (bytes, bytearray)):
|
||||
parts.append(bytes(field))
|
||||
else:
|
||||
return await handle_api_error(request, "media: expected binary file parts", 400)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
try:
|
||||
tensors = await loop.run_in_executor(None, _decode_compressed_chunk_sync, envelope, parts)
|
||||
decoded_audio = _decode_audio_payload(envelope.get("audio")) if envelope.get("audio") is not None else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
multi_job_id = envelope["job_id"].strip()
|
||||
worker_id = envelope["worker_id"].strip()
|
||||
start_index = envelope["start_index"]
|
||||
is_last = envelope["is_last"]
|
||||
last_pos = len(tensors) - 1
|
||||
items = [
|
||||
{
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": start_index + pos,
|
||||
"is_last": bool(is_last and pos == last_pos),
|
||||
"audio": decoded_audio if pos == last_pos else None,
|
||||
"format": envelope["format"],
|
||||
}
|
||||
for pos, tensor in enumerate(tensors)
|
||||
]
|
||||
|
||||
queue_size = await _enqueue_worker_items(multi_job_id, items)
|
||||
if queue_size is None:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
|
||||
debug_log(
|
||||
f"job_complete received {envelope['format']} chunk - job_id: {multi_job_id}, worker: {worker_id}, "
|
||||
f"frames: {start_index}..{start_index + len(items) - 1}, "
|
||||
f"bytes: {sum(len(p) for p in parts)}, is_last: {is_last}, queue_size: {queue_size}"
|
||||
)
|
||||
return web.json_response({"status": "success", "frames": len(items)})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/job_complete")
|
||||
async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
async def job_complete_endpoint(request):
|
||||
content_type = str(getattr(request, "content_type", "") or "").lower()
|
||||
if content_type.startswith("multipart/"):
|
||||
try:
|
||||
return await _job_complete_compressed(request)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
@@ -318,7 +457,7 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
batch_idx = data.get("batch_idx")
|
||||
image_payload = data.get("image")
|
||||
audio_payload = data.get("audio")
|
||||
is_last_raw = data.get("is_last")
|
||||
is_last = data.get("is_last")
|
||||
|
||||
errors = []
|
||||
if not isinstance(job_id, str) or not job_id.strip():
|
||||
@@ -327,57 +466,42 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
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")
|
||||
|
||||
try:
|
||||
is_last = require_bool_literal(is_last_raw, field_name="is_last")
|
||||
except ValueError as exc:
|
||||
errors.append(str(exc))
|
||||
is_last = False
|
||||
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()
|
||||
|
||||
runtime_state = _runtime_state()
|
||||
allowed_workers = runtime_state.distributed_job_allowed_workers.get(multi_job_id)
|
||||
if allowed_workers is not None and worker_id not in allowed_workers:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Unauthorized worker_id for job {multi_job_id}",
|
||||
403,
|
||||
)
|
||||
|
||||
pending = None
|
||||
queue_size = 0
|
||||
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
|
||||
while pending is None:
|
||||
async with runtime_state.distributed_jobs_lock:
|
||||
pending = runtime_state.distributed_pending_jobs.get(multi_job_id)
|
||||
if pending is not None:
|
||||
queue_item = {
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
}
|
||||
if decoded_audio is not None:
|
||||
queue_item["audio"] = decoded_audio
|
||||
await pending.put(
|
||||
queue_item
|
||||
)
|
||||
queue_size = pending.qsize()
|
||||
break
|
||||
|
||||
if time.monotonic() > deadline:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
await asyncio.sleep(0.05)
|
||||
queue_size = await _enqueue_worker_items(
|
||||
multi_job_id,
|
||||
[
|
||||
{
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
"audio": decoded_audio,
|
||||
"format": transfer_codec.LEGACY_FORMAT,
|
||||
}
|
||||
],
|
||||
)
|
||||
if queue_size is None:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
|
||||
debug_log(
|
||||
f"job_complete received canonical envelope - job_id: {multi_job_id}, "
|
||||
|
||||
@@ -1,7 +1,4 @@
|
||||
"""Orchestration helpers used by distributed queue execution."""
|
||||
|
||||
__all__ = [
|
||||
"dispatch",
|
||||
"media_sync",
|
||||
"prompt_transform",
|
||||
]
|
||||
# Orchestration helpers split out from queue_orchestration.py:
|
||||
# - prompt_transform.py: graph pruning + hidden input overrides
|
||||
# - media_sync.py: remote media/path normalization
|
||||
# - dispatch.py: worker probe + prompt dispatch
|
||||
|
||||
+207
-112
@@ -1,33 +1,48 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.auth import distributed_auth_headers
|
||||
from ...utils.config import load_config
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session, probe_worker
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
from ..schemas import coerce_positive_int
|
||||
# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module
|
||||
# stays importable without the image-list runtime; a test pins them equal).
|
||||
IMAGE_LIST_PROTOCOL_VERSION = 2
|
||||
try:
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
except ImportError:
|
||||
def trace_debug(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def trace_info(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
try:
|
||||
from ..schemas import parse_positive_int
|
||||
except ImportError:
|
||||
def parse_positive_int(value, default):
|
||||
try:
|
||||
parsed = int(value)
|
||||
return parsed if parsed > 0 else default
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
_least_busy_rr_index = 0
|
||||
|
||||
|
||||
async def worker_is_active(worker: dict[str, Any]) -> bool:
|
||||
async def worker_is_active(worker):
|
||||
"""Ping worker's /prompt endpoint to confirm it's reachable."""
|
||||
url = build_worker_url(worker)
|
||||
return await probe_worker(url, timeout=3.0) is not None
|
||||
|
||||
|
||||
async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
|
||||
async def worker_ws_is_active(worker):
|
||||
"""Ping worker's websocket endpoint to confirm it's reachable."""
|
||||
session = await get_client_session()
|
||||
url = build_worker_url(worker, "/distributed/worker_ws")
|
||||
headers = distributed_auth_headers(load_config())
|
||||
try:
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3, headers=headers)
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
|
||||
await ws.close()
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
@@ -43,17 +58,16 @@ async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def _dispatch_via_websocket(
|
||||
worker_url,
|
||||
payload,
|
||||
client_id,
|
||||
headers: dict[str, str],
|
||||
timeout=60.0,
|
||||
):
|
||||
async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
|
||||
"""Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
|
||||
request_id = uuid.uuid4().hex
|
||||
ws_payload = {
|
||||
@@ -63,49 +77,70 @@ async def _dispatch_via_websocket(
|
||||
"workflow": payload.get("workflow"),
|
||||
"client_id": client_id,
|
||||
}
|
||||
if "partial_execution_targets" in payload:
|
||||
ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
|
||||
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,
|
||||
headers=headers,
|
||||
) 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(
|
||||
worker: dict[str, Any],
|
||||
prompt_obj: dict[str, Any],
|
||||
workflow_meta: dict[str, Any] | None,
|
||||
client_id: str | None = None,
|
||||
use_websocket: bool = False,
|
||||
trace_execution_id: str | None = None,
|
||||
) -> None:
|
||||
worker,
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id=None,
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
return_prompt_id=False,
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
url = build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
target_payload = {}
|
||||
if partial_execution_targets is not None:
|
||||
target_payload["partial_execution_targets"] = list(partial_execution_targets)
|
||||
payload.update(target_payload)
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
@@ -113,18 +148,19 @@ async def dispatch_worker_prompt(
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
if use_websocket:
|
||||
ws_headers = distributed_auth_headers(load_config())
|
||||
try:
|
||||
await _dispatch_via_websocket(
|
||||
prompt_id = await _dispatch_via_websocket(
|
||||
worker_url,
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
**target_payload,
|
||||
},
|
||||
client_id,
|
||||
ws_headers,
|
||||
)
|
||||
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:
|
||||
@@ -140,17 +176,119 @@ 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_value_results=False):
|
||||
"""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 != IMAGE_LIST_PROTOCOL_VERSION
|
||||
or not required_nodes.issubset(nodes)
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} is incompatible: requires protocol "
|
||||
f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes "
|
||||
f"{sorted(required_nodes)}; no work was submitted"
|
||||
)
|
||||
|
||||
if capability.get("execution_targets") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit execution_targets support; "
|
||||
"upgrade that worker before running this feature; no work was submitted"
|
||||
)
|
||||
|
||||
if require_value_results and capability.get("value_results") is not True:
|
||||
raise RuntimeError(
|
||||
f"Image-list worker {worker_id} lacks explicit value_results support; "
|
||||
"upgrade that worker before using End overrides; 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(
|
||||
workers: list[dict[str, Any]],
|
||||
use_websocket: bool,
|
||||
delegate_master: bool,
|
||||
trace_execution_id: str | None = None,
|
||||
probe_concurrency: int = 8,
|
||||
) -> tuple[list[dict[str, Any]], bool]:
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
):
|
||||
"""Probe workers and return (active_workers, updated_delegate_master)."""
|
||||
probe_limit = coerce_positive_int(probe_concurrency, 8)
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
|
||||
if trace_execution_id and workers:
|
||||
@@ -204,6 +342,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:
|
||||
@@ -224,16 +366,16 @@ def _select_idle_round_robin(statuses):
|
||||
|
||||
|
||||
async def select_least_busy_worker(
|
||||
workers: list[dict[str, Any]],
|
||||
trace_execution_id: str | None = None,
|
||||
probe_concurrency: int = 8,
|
||||
probe_timeout: float = 3.0,
|
||||
) -> dict[str, Any] | None:
|
||||
workers,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
probe_timeout=3.0,
|
||||
):
|
||||
"""Select one worker by queue depth, round-robin among idle workers."""
|
||||
if not workers:
|
||||
return None
|
||||
|
||||
probe_limit = coerce_positive_int(probe_concurrency, 8)
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
statuses = await asyncio.gather(
|
||||
*[
|
||||
@@ -267,50 +409,3 @@ async def select_least_busy_worker(
|
||||
f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
|
||||
)
|
||||
return worker
|
||||
|
||||
|
||||
async def rank_workers_by_load(
|
||||
workers: list[dict[str, Any]],
|
||||
trace_execution_id: str | None = None,
|
||||
probe_concurrency: int = 8,
|
||||
probe_timeout: float = 3.0,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return all workers sorted by queue depth (ascending), preserving unreachable workers at the end."""
|
||||
if not workers:
|
||||
return []
|
||||
|
||||
probe_limit = coerce_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
statuses = await asyncio.gather(
|
||||
*[
|
||||
_probe_worker_queue(worker, probe_semaphore, probe_timeout)
|
||||
for worker in workers
|
||||
]
|
||||
)
|
||||
|
||||
ranked_statuses = []
|
||||
unreachable_workers = []
|
||||
for worker, status in zip(workers, statuses):
|
||||
if status is None:
|
||||
unreachable_workers.append(worker)
|
||||
continue
|
||||
ranked_statuses.append(status)
|
||||
|
||||
ranked_statuses.sort(key=lambda status: status["queue_remaining"])
|
||||
ranked_workers = [status["worker"] for status in ranked_statuses] + unreachable_workers
|
||||
|
||||
if trace_execution_id:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
"Ranked workers by load: "
|
||||
+ ", ".join(
|
||||
f"{status['worker'].get('id')}={status['queue_remaining']}"
|
||||
for status in ranked_statuses
|
||||
)
|
||||
+ (
|
||||
f"; unreachable={','.join(worker.get('id', '') for worker in unreachable_workers)}"
|
||||
if unreachable_workers
|
||||
else ""
|
||||
),
|
||||
)
|
||||
return ranked_workers
|
||||
|
||||
@@ -3,12 +3,9 @@ import hashlib
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.auth import distributed_auth_headers
|
||||
from ...utils.config import load_config
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
@@ -32,11 +29,16 @@ 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
|
||||
|
||||
|
||||
def convert_paths_for_platform(obj: Any, target_separator: str) -> Any:
|
||||
def convert_paths_for_platform(obj, target_separator):
|
||||
"""Recursively normalize likely file paths for the worker platform separator."""
|
||||
if target_separator not in ("/", "\\"):
|
||||
return obj
|
||||
@@ -127,16 +129,12 @@ def _load_media_file_sync(filename):
|
||||
return file_bytes, file_hash, mime_type
|
||||
|
||||
|
||||
async def fetch_worker_path_separator(
|
||||
worker: dict[str, Any],
|
||||
trace_execution_id: str | None = None,
|
||||
) -> str | None:
|
||||
async def fetch_worker_path_separator(worker, trace_execution_id=None):
|
||||
"""Best-effort fetch of a worker's path separator from /distributed/system_info."""
|
||||
url = build_worker_url(worker, "/distributed/system_info")
|
||||
session = await get_client_session()
|
||||
headers = distributed_auth_headers(load_config())
|
||||
try:
|
||||
async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
|
||||
if resp.status != 200:
|
||||
return None
|
||||
payload = await resp.json()
|
||||
@@ -150,18 +148,46 @@ async def fetch_worker_path_separator(
|
||||
return None
|
||||
|
||||
|
||||
async def fetch_worker_has_node(worker, node_class, trace_execution_id=None):
|
||||
"""Best-effort check that a worker has ``node_class`` registered.
|
||||
|
||||
Uses ComfyUI's ``GET /object_info/{node_class}``, which returns ``{}`` for
|
||||
unknown classes. Any failure counts as "not available" so callers fall
|
||||
back to nodes every ComfyUI install has.
|
||||
"""
|
||||
url = build_worker_url(worker, f"/object_info/{node_class}")
|
||||
session = await get_client_session()
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
|
||||
if resp.status != 200:
|
||||
return False
|
||||
payload = await resp.json()
|
||||
return isinstance(payload, dict) and node_class in payload
|
||||
except Exception as exc:
|
||||
message = f"Failed to probe node '{node_class}' on worker {worker.get('id')}: {exc}"
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, message)
|
||||
else:
|
||||
debug_log(f"[Distributed] {message}")
|
||||
return False
|
||||
|
||||
|
||||
async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
|
||||
"""Upload one media file to worker iff missing or hash-mismatched."""
|
||||
session = await get_client_session()
|
||||
headers = distributed_auth_headers(load_config())
|
||||
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:
|
||||
async with session.post(
|
||||
check_url,
|
||||
json={"filename": normalized, "hash": file_hash},
|
||||
headers=headers,
|
||||
timeout=aiohttp.ClientTimeout(total=6),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
@@ -202,11 +228,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
|
||||
return True, worker_path
|
||||
|
||||
|
||||
async def sync_worker_media(
|
||||
worker: dict[str, Any],
|
||||
prompt_obj: dict[str, Any],
|
||||
trace_execution_id: str | None = None,
|
||||
) -> None:
|
||||
async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
|
||||
"""Sync referenced media files from master to a remote worker before dispatch."""
|
||||
media_refs = _find_media_references(prompt_obj)
|
||||
if not media_refs:
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import json
|
||||
from collections import deque
|
||||
from typing import Any
|
||||
|
||||
from ...utils.logging import debug_log
|
||||
|
||||
@@ -28,7 +27,7 @@ class PromptIndex:
|
||||
def nodes_for_class(self, class_name):
|
||||
return self.nodes_by_class.get(class_name, [])
|
||||
|
||||
def has_upstream(self, start_node_id: str, target_class: str) -> bool:
|
||||
def has_upstream(self, start_node_id, target_class):
|
||||
cache_key = (str(start_node_id), target_class)
|
||||
if cache_key in self._upstream_cache:
|
||||
return self._upstream_cache[cache_key]
|
||||
@@ -89,38 +88,6 @@ def _find_downstream_nodes(prompt_obj, start_ids):
|
||||
return connected
|
||||
|
||||
|
||||
def _find_downstream_of_output_slot(prompt_obj, node_id, slot_index):
|
||||
"""Return all nodes reachable from a specific output slot on one node."""
|
||||
source_id = str(node_id)
|
||||
try:
|
||||
slot = int(slot_index)
|
||||
except (TypeError, ValueError):
|
||||
return set()
|
||||
|
||||
direct_consumers = set()
|
||||
for candidate_id, candidate_node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = candidate_node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if not (isinstance(value, list) and len(value) == 2):
|
||||
continue
|
||||
if str(value[0]) != source_id:
|
||||
continue
|
||||
try:
|
||||
input_slot = int(value[1])
|
||||
except (TypeError, ValueError):
|
||||
debug_log(
|
||||
f"Prompt transform: skipping malformed slot reference for node {candidate_id}: {value}"
|
||||
)
|
||||
continue
|
||||
if input_slot == slot:
|
||||
direct_consumers.add(str(candidate_id))
|
||||
break
|
||||
|
||||
if not direct_consumers:
|
||||
return set()
|
||||
return _find_downstream_nodes(prompt_obj, list(direct_consumers))
|
||||
|
||||
|
||||
def _create_numeric_id_generator(prompt_obj):
|
||||
"""Return a closure that yields new numeric string IDs."""
|
||||
max_id = 0
|
||||
@@ -128,7 +95,6 @@ def _create_numeric_id_generator(prompt_obj):
|
||||
try:
|
||||
numeric = int(node_id)
|
||||
except (TypeError, ValueError):
|
||||
debug_log(f"Prompt transform: ignoring non-numeric node id while generating ids: {node_id}")
|
||||
continue
|
||||
max_id = max(max_id, numeric)
|
||||
|
||||
@@ -159,193 +125,229 @@ def _find_upstream_nodes(prompt_obj, start_ids):
|
||||
return connected
|
||||
|
||||
|
||||
def _resolve_participants(enabled_worker_ids, delegate_master):
|
||||
worker_ids = [str(worker_id) for worker_id in (enabled_worker_ids or [])]
|
||||
if delegate_master:
|
||||
return worker_ids
|
||||
return ["master"] + worker_ids
|
||||
_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 _remove_dangling_input_refs(prompt_obj):
|
||||
"""Drop input links that point to nodes no longer present in prompt_obj."""
|
||||
existing_ids = set(prompt_obj.keys())
|
||||
for _node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = node.get("inputs", {})
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
if isinstance(input_value, list) and len(input_value) == 2:
|
||||
source_id = str(input_value[0])
|
||||
if source_id not in existing_ids:
|
||||
inputs.pop(input_name, 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 _has_terminal_output_nodes(prompt_obj):
|
||||
"""Return True when the prompt already has at least one terminal output node."""
|
||||
for _node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") in {"PreviewImage", "SaveImage"}:
|
||||
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 _add_preview_node(prompt_obj, next_id_fn, source_node_id, slot_index, title_suffix):
|
||||
"""Attach a PreviewImage output node to the given source output slot."""
|
||||
preview_id = next_id_fn()
|
||||
prompt_obj[preview_id] = {
|
||||
"inputs": {
|
||||
"images": [str(source_node_id), int(slot_index)],
|
||||
},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {
|
||||
"title": f"Preview Image ({title_suffix})",
|
||||
},
|
||||
}
|
||||
return preview_id
|
||||
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 _ensure_worker_output_node(prompt_obj):
|
||||
"""Guarantee worker prompts retain at least one output node after pruning.
|
||||
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)
|
||||
)
|
||||
|
||||
Branch pruning can remove all explicit output nodes for worker prompts, which
|
||||
causes ComfyUI validation to reject the prompt (`prompt_no_outputs`).
|
||||
"""
|
||||
if _has_terminal_output_nodes(prompt_obj):
|
||||
return prompt_obj
|
||||
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
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
|
||||
|
||||
# Prefer a DistributedBranchCollector output tied to the assigned branch for this worker.
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") != "DistributedBranchCollector":
|
||||
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", {})
|
||||
try:
|
||||
assigned_branch = int(inputs.get("assigned_branch", -1))
|
||||
except (TypeError, ValueError):
|
||||
assigned_branch = -1
|
||||
if assigned_branch >= 0:
|
||||
_add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
|
||||
return prompt_obj
|
||||
|
||||
# Otherwise, anchor to the assigned DistributedBranch slot if available.
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") != "DistributedBranch":
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
try:
|
||||
assigned_branch = int(inputs.get("assigned_branch", -1))
|
||||
except (TypeError, ValueError):
|
||||
assigned_branch = -1
|
||||
if assigned_branch >= 0:
|
||||
_add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
|
||||
return prompt_obj
|
||||
|
||||
# Idle participants (assigned_branch=-1) still need a terminal node to pass
|
||||
# validation; keep this cheap by emitting a tiny synthetic image preview.
|
||||
empty_id = next_id()
|
||||
prompt_obj[empty_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
"inputs": {
|
||||
"height": 64,
|
||||
"width": 64,
|
||||
"channels": 3,
|
||||
},
|
||||
"_meta": {
|
||||
"title": "Distributed Empty Image (auto-added worker output)",
|
||||
},
|
||||
}
|
||||
_add_preview_node(prompt_obj, next_id, empty_id, 0, "auto-added worker output")
|
||||
return prompt_obj
|
||||
|
||||
|
||||
def _prune_worker_downstream_of_branch_collectors(prompt_obj):
|
||||
"""Drop worker-side nodes downstream of branch collectors."""
|
||||
collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
|
||||
if not collector_ids:
|
||||
return prompt_obj
|
||||
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
for node_id in downstream:
|
||||
if node_id in collector_ids:
|
||||
continue
|
||||
prompt_obj.pop(node_id, None)
|
||||
|
||||
_remove_dangling_input_refs(prompt_obj)
|
||||
return prompt_obj
|
||||
|
||||
|
||||
def prune_prompt_for_branch_worker(
|
||||
prompt_obj: dict[str, Any],
|
||||
branch_node_id: str,
|
||||
assigned_branch: int | list[int] | tuple[int, ...] | set[int],
|
||||
num_branches: int,
|
||||
) -> dict[str, Any]:
|
||||
"""Prune non-assigned branch paths while keeping shared downstream nodes."""
|
||||
branch_id = str(branch_node_id)
|
||||
|
||||
assigned_slots = set()
|
||||
if isinstance(assigned_branch, (list, tuple, set)):
|
||||
for value in assigned_branch:
|
||||
try:
|
||||
idx = int(value)
|
||||
except (TypeError, ValueError):
|
||||
debug_log(f"Prompt transform: invalid assigned branch slot ignored: {value}")
|
||||
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 idx >= 0:
|
||||
assigned_slots.add(idx)
|
||||
else:
|
||||
try:
|
||||
idx = int(assigned_branch)
|
||||
except (TypeError, ValueError):
|
||||
idx = -1
|
||||
if idx >= 0:
|
||||
assigned_slots.add(idx)
|
||||
|
||||
try:
|
||||
total_slots = int(num_branches)
|
||||
except (TypeError, ValueError):
|
||||
total_slots = 2
|
||||
total_slots = max(2, min(total_slots, 10))
|
||||
|
||||
downstream_by_slot = {}
|
||||
for slot_idx in range(total_slots):
|
||||
downstream = _find_downstream_of_output_slot(prompt_obj, branch_id, slot_idx)
|
||||
downstream.discard(branch_id)
|
||||
downstream_by_slot[slot_idx] = downstream
|
||||
|
||||
keep = {branch_id}
|
||||
for slot_idx in assigned_slots:
|
||||
keep.update(downstream_by_slot.get(slot_idx, set()))
|
||||
|
||||
remove = set()
|
||||
for slot_idx in range(total_slots):
|
||||
if slot_idx in assigned_slots:
|
||||
continue
|
||||
remove.update(downstream_by_slot.get(slot_idx, set()))
|
||||
|
||||
remove -= keep
|
||||
for node_id in remove:
|
||||
prompt_obj.pop(node_id, None)
|
||||
|
||||
_remove_dangling_input_refs(prompt_obj)
|
||||
return prompt_obj
|
||||
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: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Prune worker prompt to distributed nodes and their upstream dependencies."""
|
||||
WORKER_SINK_CLASS = "DistributedWorkerSink"
|
||||
|
||||
|
||||
def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
|
||||
"""Prune worker prompt to distributed nodes and their upstream dependencies.
|
||||
|
||||
Removing downstream nodes can leave the worker prompt without an output
|
||||
node. One is re-added after each such distributed node: the hidden
|
||||
``DistributedWorkerSink`` when ``use_worker_sink`` is set (the worker must
|
||||
have it registered), otherwise PreviewImage/PreviewAudio for workers
|
||||
running older versions. The preview nodes write every frame to the
|
||||
worker's temp directory, which the sink avoids.
|
||||
"""
|
||||
collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
|
||||
branch_collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
|
||||
upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
|
||||
branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch")
|
||||
distributed_ids = collector_ids + branch_collector_ids + upscale_ids + branch_ids
|
||||
distributed_ids = collector_ids + upscale_ids
|
||||
if not distributed_ids:
|
||||
return prompt_obj
|
||||
|
||||
connected = _find_upstream_nodes(prompt_obj, distributed_ids)
|
||||
if branch_ids:
|
||||
connected.update(_find_downstream_nodes(prompt_obj, branch_ids))
|
||||
if branch_collector_ids:
|
||||
downstream_of_collectors = _find_downstream_nodes(prompt_obj, branch_collector_ids)
|
||||
connected -= (downstream_of_collectors - set(branch_collector_ids))
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in connected:
|
||||
node = prompt_obj.get(node_id)
|
||||
@@ -358,30 +360,53 @@ def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
|
||||
if dist_id not in pruned_prompt:
|
||||
continue
|
||||
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
||||
has_removed_downstream = any(node_id != dist_id and node_id not in connected for node_id in downstream)
|
||||
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 use_worker_sink and (image_connected or audio_connected):
|
||||
pruned_prompt[next_id()] = {
|
||||
"inputs": {"value": [dist_id, 0 if image_connected else 1]},
|
||||
"class_type": WORKER_SINK_CLASS,
|
||||
"_meta": {"title": "Distributed Worker Sink (auto-added)"},
|
||||
}
|
||||
elif image_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"images": [dist_id, 0]},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {"title": "Preview Image (auto-added)"},
|
||||
}
|
||||
elif audio_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"audio": [dist_id, 1]},
|
||||
"class_type": "PreviewAudio",
|
||||
"_meta": {"title": "Preview Audio (auto-added)"},
|
||||
}
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def prepare_delegate_master_prompt(
|
||||
prompt_obj: dict[str, Any],
|
||||
collector_ids: list[str],
|
||||
) -> dict[str, Any]:
|
||||
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
nodes_to_keep.update(
|
||||
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
|
||||
)
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in nodes_to_keep:
|
||||
@@ -409,16 +434,9 @@ def prepare_delegate_master_prompt(
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
collector_class = collector_entry.get("class_type")
|
||||
if collector_class == "DistributedBranchCollector":
|
||||
# Branch collector does not require an image input in delegate-only mode.
|
||||
continue
|
||||
if collector_class == "UltimateSDUpscaleDistributed":
|
||||
# Swap to lightweight delegate collector — no model inputs needed.
|
||||
collector_entry["class_type"] = "USDUDelegateCollector"
|
||||
debug_log(
|
||||
f"Swapped USDU node {collector_id} to USDUDelegateCollector for delegate-only master prompt."
|
||||
)
|
||||
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] = {
|
||||
@@ -440,145 +458,76 @@ def prepare_delegate_master_prompt(
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str]:
|
||||
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("DistributedBranchCollector")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
+ prompt_index.nodes_for_class("DistributedBranch")
|
||||
+ prompt_index.nodes_for_class("DistributedImageListStart")
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
return job_map
|
||||
|
||||
|
||||
def _override_simple_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json, class_names):
|
||||
"""Configure simple distributed nodes (DistributedSeed, DistributedValue)."""
|
||||
for class_name in class_names:
|
||||
for node_id in prompt_index.nodes_for_class(class_name):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["worker_id"] = "" if is_master else str(participant_id)
|
||||
|
||||
|
||||
def _override_job_nodes(
|
||||
prompt_copy, prompt_index, is_master, participant_id,
|
||||
job_id_map, master_url, enabled_json, delegate_master,
|
||||
class_names, skip_if_upstream=None,
|
||||
):
|
||||
"""Configure job-aware distributed nodes (collector, upscale)."""
|
||||
for class_name in class_names:
|
||||
for node_id in prompt_index.nodes_for_class(class_name):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
if skip_if_upstream and prompt_index.has_upstream(node_id, skip_if_upstream):
|
||||
node.setdefault("inputs", {})["pass_through"] = True
|
||||
continue
|
||||
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
|
||||
def _override_branch_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
delegate_master,
|
||||
):
|
||||
"""Assign branch slots per participant and prune non-assigned branch paths."""
|
||||
participants = _resolve_participants(enabled_worker_ids, delegate_master)
|
||||
participant_id = str(participant_id)
|
||||
participant_count = len(participants)
|
||||
participant_pos = participants.index(participant_id) if participant_id in participants else None
|
||||
|
||||
for node_id in prompt_index.nodes_for_class("DistributedBranch"):
|
||||
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedSeed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
inputs = node.setdefault("inputs", {})
|
||||
try:
|
||||
num_branches = int(inputs.get("num_branches", 2))
|
||||
except (TypeError, ValueError):
|
||||
num_branches = 2
|
||||
num_branches = max(2, min(num_branches, 10))
|
||||
|
||||
assigned_slots = []
|
||||
if participant_pos is not None and participant_count > 0:
|
||||
assigned_slots = [slot for slot in range(num_branches) if slot % participant_count == participant_pos]
|
||||
|
||||
if len(assigned_slots) == 1:
|
||||
assigned_branch = assigned_slots[0]
|
||||
else:
|
||||
assigned_branch = -1
|
||||
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["worker_id"] = "" if is_master else participant_id
|
||||
inputs["assigned_branch"] = assigned_branch
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
|
||||
if is_master and delegate_master:
|
||||
continue
|
||||
|
||||
prompt_copy = prune_prompt_for_branch_worker(
|
||||
prompt_copy,
|
||||
node_id,
|
||||
assigned_slots,
|
||||
num_branches,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _find_upstream_branch_node(prompt_copy, start_node_id):
|
||||
"""Locate the nearest upstream DistributedBranch node for the provided node."""
|
||||
visited = set()
|
||||
stack = [str(start_node_id)]
|
||||
while stack:
|
||||
node_id = stack.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
inputs = node.get("inputs", {})
|
||||
for input_value in inputs.values():
|
||||
if not (isinstance(input_value, list) and len(input_value) == 2):
|
||||
continue
|
||||
source_id = str(input_value[0])
|
||||
source_node = prompt_copy.get(source_id)
|
||||
if not isinstance(source_node, dict):
|
||||
continue
|
||||
if source_node.get("class_type") == "DistributedBranch":
|
||||
return source_id
|
||||
stack.append(source_id)
|
||||
return None
|
||||
|
||||
|
||||
def _override_branch_collector_nodes(
|
||||
def _override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
@@ -588,77 +537,104 @@ def _override_branch_collector_nodes(
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
):
|
||||
"""Configure DistributedBranchCollector nodes for branch convergence."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedBranchCollector"):
|
||||
"""Configure DistributedCollector nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
upstream_branch_id = _find_upstream_branch_node(prompt_copy, node_id)
|
||||
assigned_branch = -1
|
||||
if upstream_branch_id is not None:
|
||||
branch_node = prompt_copy.get(upstream_branch_id, {})
|
||||
branch_inputs = branch_node.get("inputs", {}) if isinstance(branch_node, dict) else {}
|
||||
try:
|
||||
assigned_branch = int(branch_inputs.get("assigned_branch", -1))
|
||||
except (TypeError, ValueError):
|
||||
assigned_branch = -1
|
||||
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
||||
node.setdefault("inputs", {})["pass_through"] = True
|
||||
continue
|
||||
|
||||
inputs = node.setdefault("inputs", {})
|
||||
# Group all branch collectors under the same DistributedBranch into one queue.
|
||||
job_source_id = upstream_branch_id or node_id
|
||||
inputs["multi_job_id"] = job_id_map.get(job_source_id, job_source_id)
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["assigned_branch"] = assigned_branch
|
||||
if is_master:
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = str(participant_id)
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
|
||||
def _override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
):
|
||||
"""Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
|
||||
|
||||
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedValue nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedValue"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
|
||||
job_id_map, master_url, delegate_master, 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_participates"] = participant_id == "master" and not delegate_master
|
||||
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: dict[str, Any],
|
||||
participant_id: str,
|
||||
enabled_worker_ids: list[str],
|
||||
job_id_map: dict[str, str],
|
||||
master_url: str,
|
||||
delegate_master: bool,
|
||||
prompt_index: PromptIndex,
|
||||
) -> dict[str, Any]:
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
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"
|
||||
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
||||
enabled_json = json.dumps(enabled_worker_ids)
|
||||
|
||||
_override_simple_nodes(
|
||||
prompt_copy, prompt_index, is_master, participant_id, enabled_json,
|
||||
["DistributedSeed", "DistributedValue"],
|
||||
)
|
||||
_override_job_nodes(
|
||||
prompt_copy, prompt_index, is_master, participant_id,
|
||||
job_id_map, master_url, enabled_json, delegate_master,
|
||||
["DistributedCollector"],
|
||||
skip_if_upstream="UltimateSDUpscaleDistributed",
|
||||
)
|
||||
_override_job_nodes(
|
||||
prompt_copy, prompt_index, is_master, participant_id,
|
||||
job_id_map, master_url, enabled_json, delegate_master,
|
||||
["UltimateSDUpscaleDistributed"],
|
||||
)
|
||||
prompt_copy = _override_branch_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
delegate_master,
|
||||
)
|
||||
_override_branch_collector_nodes(
|
||||
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
@@ -668,9 +644,24 @@ def apply_participant_overrides(
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
)
|
||||
|
||||
if not is_master:
|
||||
prompt_copy = _prune_worker_downstream_of_branch_collectors(prompt_copy)
|
||||
prompt_copy = _ensure_worker_output_node(prompt_copy)
|
||||
_override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
)
|
||||
_override_image_list_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""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, *, media=None):
|
||||
"""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,
|
||||
},
|
||||
}
|
||||
if media is not None:
|
||||
from ...serverless_bridge.handler import _validate_request
|
||||
payload["input"]["media"] = media
|
||||
_validate_request(payload["input"])
|
||||
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, *, media=None):
|
||||
body = serialize_request(workflow, images, timeout_seconds, media=media)
|
||||
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)}")
|
||||
@@ -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])
|
||||
@@ -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)
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Pure compilation of prepared Distributed Collectors for comfy-media-v1 pools.
|
||||
|
||||
The legacy image compiler remains separate. This compiler emits one MediaExportV1
|
||||
execution target per connected Collector socket (``images`` and/or ``audio``),
|
||||
retaining only its upstream computation. It never imports ComfyUI or
|
||||
inspects/executes runtime node classes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
from graphlib import CycleError, TopologicalSorter
|
||||
|
||||
from ...serverless_bridge.contract import SOCKET_KINDS, identifier, validate_bindings
|
||||
from ...serverless_bridge.image_transfer import DEFAULT_FORMAT, validate_options
|
||||
from .serverless_prompt import (
|
||||
_PRIMITIVE_NODES,
|
||||
_REMOVED_SCALAR_NODES,
|
||||
_is_link,
|
||||
_is_unsupported_distributed,
|
||||
_make_scalar_resolver,
|
||||
_prune_graph,
|
||||
)
|
||||
|
||||
|
||||
_COLLECTOR = "DistributedCollector"
|
||||
# Collector inputs, in output-slot order, that a typed pool returns.
|
||||
COLLECTOR_SOCKETS = ("images", "audio")
|
||||
# Collector widget defaults (the saved-workflow contract is legacy_png / 90).
|
||||
COLLECTOR_DEFAULT_FORMAT = "legacy_png"
|
||||
COLLECTOR_DEFAULT_QUALITY = 90
|
||||
|
||||
|
||||
def collector_image_policy(inputs):
|
||||
"""Return the endpoint IMAGE (format, quality) for one Collector.
|
||||
|
||||
legacy_png is the Collector's saved default and names a JSON/base64 framing
|
||||
that only applies to direct workers; serverless returns the same 8-bit
|
||||
pixels as compressed png. exact stays opt-in. No other format is altered.
|
||||
"""
|
||||
fmt, quality = validate_options(inputs.get("transfer_format", COLLECTOR_DEFAULT_FORMAT),
|
||||
inputs.get("transfer_quality", COLLECTOR_DEFAULT_QUALITY))
|
||||
return ("png" if fmt == "legacy_png" else fmt), quality
|
||||
|
||||
|
||||
def _validate_identifier(value, context):
|
||||
# Validate the original token before any lookup. In particular, neither
|
||||
# integer node IDs nor whitespace in prepared identities is normalized.
|
||||
try:
|
||||
identifier(value)
|
||||
except ValueError:
|
||||
raise ValueError(f"{context} must be a valid nonempty media identifier") from None
|
||||
|
||||
|
||||
def _validate_prompt(prepared_prompt):
|
||||
"""Validate the entire source graph, including branches later discarded."""
|
||||
if not isinstance(prepared_prompt, dict) or not prepared_prompt:
|
||||
raise ValueError("prepared prompt must be a nonempty API graph mapping")
|
||||
|
||||
collectors = []
|
||||
for node_id, node in prepared_prompt.items():
|
||||
_validate_identifier(node_id, "prompt node ID")
|
||||
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) or not class_type.strip():
|
||||
raise ValueError(f"prompt node {node_id!r} must have a nonempty class_type")
|
||||
if class_type == _COLLECTOR:
|
||||
collectors.append(node_id)
|
||||
elif _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")
|
||||
if any(not isinstance(name, str) or not name for name in inputs):
|
||||
raise ValueError(f"prompt node {node_id!r} input names must be nonempty strings")
|
||||
|
||||
if not collectors:
|
||||
raise ValueError("prepared prompt contains no DistributedCollector nodes")
|
||||
|
||||
dependencies = {}
|
||||
for node_id, node in prepared_prompt.items():
|
||||
parents = set()
|
||||
for input_name, value in node.get("inputs", {}).items():
|
||||
# Ordinary inputs may contain literal lists. As in the legacy
|
||||
# compiler, exactly two elements denotes an API link, not media.
|
||||
if not _is_link(value):
|
||||
continue
|
||||
source_id, slot = value
|
||||
context = f"link at node {node_id!r} input {input_name!r}"
|
||||
_validate_identifier(source_id, context + " source ID")
|
||||
if source_id not in prepared_prompt:
|
||||
raise ValueError(f"invalid {context}: source {source_id!r} does not exist")
|
||||
if type(slot) is not int or slot < 0:
|
||||
raise ValueError(f"invalid {context}: output slot must be a nonnegative integer")
|
||||
source_class = prepared_prompt[source_id]["class_type"]
|
||||
if source_class in _REMOVED_SCALAR_NODES or source_class in _PRIMITIVE_NODES:
|
||||
if slot != 0:
|
||||
raise ValueError(f"invalid output slot {slot} for scalar node {source_id!r}")
|
||||
elif source_class == _COLLECTOR and slot >= len(COLLECTOR_SOCKETS):
|
||||
raise ValueError(f"invalid output slot {slot} for collector {source_id!r}")
|
||||
parents.add(source_id)
|
||||
dependencies[node_id] = parents
|
||||
|
||||
for collector_id in collectors:
|
||||
inputs = prepared_prompt[collector_id].get("inputs", {})
|
||||
collector_image_policy(inputs)
|
||||
for field in ("multi_job_id", "worker_id"):
|
||||
_validate_identifier(inputs.get(field), f"collector {collector_id!r} {field}")
|
||||
connected = [socket for socket in COLLECTOR_SOCKETS if socket in inputs]
|
||||
if not connected:
|
||||
raise ValueError(f"collector {collector_id!r} requires at least one connected media socket")
|
||||
for socket in connected:
|
||||
value = inputs[socket]
|
||||
if not _is_link(value):
|
||||
raise ValueError(f"collector {collector_id!r} {socket} must be linked; use PrimitiveString for literal text")
|
||||
if prepared_prompt[value[0]]["class_type"] in _REMOVED_SCALAR_NODES:
|
||||
# Export sockets must remain actual links after scalar folding.
|
||||
# A PrimitiveString wrapper supports distributed text values.
|
||||
raise ValueError(f"collector {collector_id!r} {socket} cannot link directly to a distributed scalar; use a typed node or PrimitiveString")
|
||||
|
||||
# Iterative topological sorting checks ordinary and unused cycles too, not
|
||||
# only the primitive chains that the shared scalar resolver traverses.
|
||||
try:
|
||||
ordered = list(TopologicalSorter(dependencies).static_order())
|
||||
except CycleError:
|
||||
raise ValueError("cycle in prepared media prompt") from None
|
||||
collector_set = set(collectors)
|
||||
has_collector_ancestor = {}
|
||||
for node_id in ordered:
|
||||
has_collector_ancestor[node_id] = any(
|
||||
parent in collector_set or has_collector_ancestor[parent]
|
||||
for parent in dependencies[node_id]
|
||||
)
|
||||
if node_id in collector_set and has_collector_ancestor[node_id]:
|
||||
raise ValueError("collector-to-collector dependence is unsupported")
|
||||
|
||||
return sorted(collectors)
|
||||
|
||||
|
||||
def _build_exports(source_nodes, collector_ids, request_id):
|
||||
bindings = []
|
||||
exports = {}
|
||||
reserved = set(source_nodes)
|
||||
for collector_id in collector_ids:
|
||||
inputs = source_nodes[collector_id].get("inputs", {})
|
||||
for socket in COLLECTOR_SOCKETS:
|
||||
kind = SOCKET_KINDS[socket]
|
||||
if socket not in inputs:
|
||||
continue
|
||||
# Hashing avoids separator ambiguity and keeps long valid source
|
||||
# identities within the bridge's identifier length bound. Identity
|
||||
# is stable across request IDs and dictionary insertion orders.
|
||||
identity = json.dumps([collector_id, socket], separators=(",", ":"))
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
base_id = f"serverless:media:{digest}"
|
||||
node_id = base_id
|
||||
suffix = 0
|
||||
while node_id in reserved:
|
||||
suffix += 1
|
||||
node_id = f"{base_id}:{suffix}"
|
||||
reserved.add(node_id)
|
||||
binding = {
|
||||
"binding_id": f"b:{digest}",
|
||||
"collector_id": collector_id,
|
||||
"job_id": inputs["multi_job_id"],
|
||||
"participant_id": inputs["worker_id"],
|
||||
"node_id": node_id,
|
||||
"socket": socket,
|
||||
"kind": kind,
|
||||
}
|
||||
bindings.append(binding)
|
||||
exports[node_id] = {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": request_id,
|
||||
"binding_json": json.dumps(binding),
|
||||
socket: copy.deepcopy(inputs[socket]),
|
||||
},
|
||||
}
|
||||
if kind == "IMAGE":
|
||||
fmt, quality = collector_image_policy(inputs)
|
||||
if fmt != DEFAULT_FORMAT:
|
||||
exports[node_id]["inputs"].update(transfer_format=fmt, transfer_quality=quality)
|
||||
validate_bindings(bindings)
|
||||
return bindings, exports
|
||||
|
||||
|
||||
def compile_media_prompt(prepared_prompt: dict, request_id: str) -> tuple[dict, list[dict]]:
|
||||
"""Return a worker graph and canonical typed-media export bindings.
|
||||
|
||||
Each DistributedCollector must have prepared literal ``multi_job_id`` and
|
||||
``worker_id`` identities and at least one linked ``images``/``audio`` input.
|
||||
Present media inputs must be API links. Collector transport controls become
|
||||
the IMAGE export policy (``legacy_png`` is sent as ``png``); other controls
|
||||
are not exported.
|
||||
|
||||
Only MediaExportV1 targets and their ancestors survive. Required ordinary
|
||||
output/utility nodes are preserved regardless of their runtime OUTPUT_NODE
|
||||
flag; callers must dispatch exactly the returned binding ``node_id`` values
|
||||
as partial_execution_targets. Distributed scalar dependencies use the same
|
||||
resolver as the legacy compiler. Direct scalar-to-media links need a typed
|
||||
wrapper so that each export input remains a link.
|
||||
|
||||
Validation and compilation never modify the caller's source graph. This
|
||||
structural compiler does not infer socket types or unknown nodes' output
|
||||
counts; ComfyUI performs runtime type validation at the endpoint.
|
||||
"""
|
||||
_validate_identifier(request_id, "request_id")
|
||||
collector_ids = _validate_prompt(prepared_prompt)
|
||||
bindings, exports = _build_exports(prepared_prompt, collector_ids, request_id)
|
||||
compiled = _prune_graph(
|
||||
prepared_prompt, bindings, exports, _make_scalar_resolver(prepared_prompt)
|
||||
)
|
||||
return {node_id: compiled[node_id] for node_id in sorted(compiled)}, bindings
|
||||
@@ -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
|
||||
@@ -0,0 +1,411 @@
|
||||
"""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,
|
||||
)
|
||||
|
||||
|
||||
# Fixed per-worker cap on this master's outstanding Runpod jobs; equals the
|
||||
# Images per run maximum (MAX_JOBS_PER_RUN in utils/serverless_config.py), so
|
||||
# every image of a run is submitted at once. Deliberately not configurable or
|
||||
# derived from jobs_per_run: unresolved jobs from earlier runs (ambiguous
|
||||
# submissions, unconfirmed cancels) still count against it, and a constant
|
||||
# never has to change while such jobs exist.
|
||||
PER_WORKER_JOB_LIMIT = 16
|
||||
|
||||
|
||||
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 < PER_WORKER_JOB_LIMIT:
|
||||
# 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):
|
||||
# The limit is a constant, so it never has to be swapped while this
|
||||
# worker still has unresolved jobs.
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
semaphore = self.pool_limits.get(pool_id)
|
||||
if semaphore is None:
|
||||
semaphore = asyncio.Semaphore(PER_WORKER_JOB_LIMIT)
|
||||
self.pool_limits[pool_id] = semaphore
|
||||
return semaphore
|
||||
|
||||
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))
|
||||
from .serverless_typed_results import validate_prepared_job
|
||||
# Validate the entire batch before creating even the first local slot.
|
||||
for prepared in prepared_jobs:
|
||||
validate_prepared_job(prepared)
|
||||
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, *, media=None):
|
||||
if media is not None and not error:
|
||||
# Validate all destinations before exposing any partial success.
|
||||
for binding in run["manifest"]:
|
||||
if binding["participant_id"] != run["worker"]["id"]:
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
values = media[binding["job_id"]][binding["socket"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
async with self.server.distributed_jobs_lock:
|
||||
delivered = set()
|
||||
for binding in run["manifest"]:
|
||||
job_id = binding["job_id"]
|
||||
if error or media is not None:
|
||||
if job_id in delivered:
|
||||
continue
|
||||
delivered.add(job_id)
|
||||
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})
|
||||
elif media is not None:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "media": media[job_id], "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:
|
||||
from .serverless_typed_results import validate_prepared_job, collect_typed_results
|
||||
typed = validate_prepared_job(run)
|
||||
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
|
||||
if typed:
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"], media=run["media"])
|
||||
else:
|
||||
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())
|
||||
if typed:
|
||||
media = await collect_typed_results(output, run["media"], run["manifest"])
|
||||
await self._deliver(run, media=media)
|
||||
else:
|
||||
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.pop("media", 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 worker")
|
||||
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)
|
||||
run.pop("media", 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)]]
|
||||
@@ -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)
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Bounded local native inputs for comfy-media-v1 requests.
|
||||
|
||||
Only stock LoadAudio/LoadVideo schemas are adapted. Original nodes and links
|
||||
remain in the private request graph for endpoint validation; only the existing
|
||||
image preparer rewrites filenames. No upload, transcode, resample or provider
|
||||
operation belongs here. The caller still enforces the complete request-body cap.
|
||||
"""
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from ...serverless_bridge import codec, contract, video_codec
|
||||
from . import serverless_media
|
||||
|
||||
|
||||
_LOADERS = {"LoadAudio": ("audio", "AUDIO"), "LoadVideo": ("file", "VIDEO")}
|
||||
Error = serverless_media.ServerlessMediaError
|
||||
|
||||
|
||||
def _loader_references(workflow):
|
||||
if not isinstance(workflow, dict):
|
||||
raise Error("Serverless workflow must be an API graph")
|
||||
loaders = {}
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str) or class_type not in _LOADERS:
|
||||
continue
|
||||
contract.identifier(node_id)
|
||||
key, kind = _LOADERS[class_type]
|
||||
inputs = node.get("inputs")
|
||||
if not isinstance(inputs, dict) or set(inputs) != {key}:
|
||||
raise Error(f"{class_type} inputs must contain only its native filename; extra controls are unsupported")
|
||||
if not isinstance(inputs[key], str):
|
||||
raise Error(f"{class_type} requires a literal filename; linked inputs are unsupported")
|
||||
normalized, _suffix = serverless_media._normalize_media_reference(inputs[key])
|
||||
loaders[node_id] = (kind, normalized)
|
||||
|
||||
references = []
|
||||
used = set()
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
|
||||
continue
|
||||
for key, value in node["inputs"].items():
|
||||
if not isinstance(value, (list, tuple)) or not value or not isinstance(value[0], str):
|
||||
continue
|
||||
if value[0] not in loaders:
|
||||
continue
|
||||
if type(value) is not list or len(value) != 2 or type(value[1]) is not int or value[1] != 0:
|
||||
raise Error("Native media loader links require output slot 0")
|
||||
contract.identifier(node_id)
|
||||
contract.identifier(key)
|
||||
references.append((node_id, key, value[0]))
|
||||
used.add(value[0])
|
||||
if used != set(loaders):
|
||||
raise Error("Native media loader has no consumer")
|
||||
return loaders, references
|
||||
|
||||
|
||||
def _source_path(value, folder_paths):
|
||||
"""Resolve native annotations; authorization/opening is owned by _read_path."""
|
||||
annotation = serverless_media._MEDIA_ANNOTATION.search(value)
|
||||
unannotated = value[:annotation.start()] if annotation else value
|
||||
if Path(unannotated).is_absolute():
|
||||
return unannotated
|
||||
try:
|
||||
return folder_paths.get_annotated_filepath(value)
|
||||
except Exception:
|
||||
raise Error("Native media input path is missing or invalid") from None
|
||||
|
||||
|
||||
def _decode_audio(data, limit):
|
||||
"""Match core nodes_audio.load, bounding each frame *before* accumulation.
|
||||
|
||||
Float32 output bytes are charged even for narrow integer PCM. Native float64
|
||||
remains float64, as in f32_pcm (there is no silent dtype conversion). Raw
|
||||
retained arrays and each live AV plane allocation have a separate <= limit
|
||||
guard, so wider/padded raw frames cannot bypass the output budget. PyAV owns
|
||||
the one frame being decoded; it is checked before to_ndarray or torch.cat.
|
||||
Concatenation/conversion use bounded temporary copies, not an RSS guarantee.
|
||||
"""
|
||||
if limit <= 0:
|
||||
raise Error("AUDIO decoded input exceeds byte limit")
|
||||
try:
|
||||
import av
|
||||
import torch
|
||||
from comfy_extras.nodes_audio import f32_pcm
|
||||
except ImportError:
|
||||
raise Error("Native AUDIO decoding is unavailable") from None
|
||||
|
||||
try:
|
||||
with av.open(io.BytesIO(data)) as container:
|
||||
if not container.streams.audio:
|
||||
raise Error("Native AUDIO input has no audio stream")
|
||||
stream = container.streams.audio[0]
|
||||
rate = stream.codec_context.sample_rate
|
||||
channels = stream.channels
|
||||
if type(rate) is not int or rate <= 0 or type(channels) is not int or channels <= 0:
|
||||
raise Error("Native AUDIO stream shape or sample rate is invalid")
|
||||
frames = []
|
||||
output_bytes = raw_bytes = 0
|
||||
dtype = None
|
||||
for frame in container.decode(streams=stream.index):
|
||||
samples = frame.samples
|
||||
if (type(samples) is not int or samples <= 0 or frame.sample_rate != rate
|
||||
or len(frame.layout.channels) != channels):
|
||||
raise Error("Native AUDIO frame shape or sample rate changed")
|
||||
if frame.format.name not in {"s16", "s16p", "s32", "s32p", "flt", "fltp", "dbl", "dblp"}:
|
||||
raise Error("Unsupported native AUDIO PCM format")
|
||||
itemsize = frame.format.bytes
|
||||
count = samples * channels
|
||||
# Native f32_pcm promotes integer PCM to float32, not float64.
|
||||
decoded_bytes = count * max(4, itemsize)
|
||||
frame_bytes = count * itemsize
|
||||
plane_bytes = sum(plane.buffer_size for plane in frame.planes)
|
||||
if (decoded_bytes > limit - output_bytes or frame_bytes > limit - raw_bytes
|
||||
or plane_bytes > limit - raw_bytes):
|
||||
raise Error("AUDIO decoded input or raw frame exceeds byte limit")
|
||||
output_bytes += decoded_bytes
|
||||
raw_bytes += frame_bytes
|
||||
array = frame.to_ndarray()
|
||||
if array.ndim != 2 or array.size != count or array.nbytes != frame_bytes:
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
buf = torch.from_numpy(array)
|
||||
if buf.dtype not in (torch.int16, torch.int32, torch.float32, torch.float64):
|
||||
raise Error("Unsupported native AUDIO PCM dtype")
|
||||
if dtype is not None and buf.dtype != dtype:
|
||||
raise Error("Native AUDIO frame dtype changed")
|
||||
dtype = buf.dtype
|
||||
if buf.shape[0] != channels:
|
||||
buf = buf.view(-1, channels).t()
|
||||
if tuple(buf.shape) != (channels, samples):
|
||||
raise Error("Native AUDIO frame shape is invalid")
|
||||
frames.append(buf)
|
||||
if not frames:
|
||||
raise Error("Native AUDIO input has no decoded frames")
|
||||
waveform = f32_pcm(torch.cat(frames, dim=1)).unsqueeze(0)
|
||||
return {"waveform": waveform, "sample_rate": rate}
|
||||
except Error:
|
||||
raise
|
||||
except Exception:
|
||||
raise Error("Native AUDIO input could not be decoded safely") from None
|
||||
|
||||
|
||||
def _artifact_id(kind, descriptor, payload):
|
||||
digest = hashlib.sha256()
|
||||
digest.update(kind.encode("ascii") + b"\0")
|
||||
# Include interpretation (not just waveform bytes), notably sample_rate.
|
||||
digest.update(json.dumps(descriptor, sort_keys=True, separators=(",", ":"), allow_nan=False).encode())
|
||||
digest.update(b"\0")
|
||||
digest.update(payload)
|
||||
return "input-" + kind.lower() + "-" + digest.hexdigest()
|
||||
|
||||
|
||||
async def prepare_media_inputs(workflow: dict) -> tuple[list, list, list]:
|
||||
"""Return legacy images, inline typed artifacts, and every consumer reference.
|
||||
|
||||
Paths are read once into bounded bytes through the video adapter's no-follow
|
||||
directory-fd reader. Decoders receive BytesIO, never a second path open.
|
||||
Source files and unique decoded artifacts have independent aggregate caps.
|
||||
Repeated content of the same kind is decoded once and charged once.
|
||||
"""
|
||||
loaders, references = _loader_references(workflow)
|
||||
image_view = {node_id: node for node_id, node in workflow.items() if node_id not in loaders}
|
||||
# Retain the strict scan for every other loader before expensive decoding.
|
||||
serverless_media._iter_image_references(image_view)
|
||||
if not loaders:
|
||||
return await serverless_media.prepare_images(image_view), [], []
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
raise Error("ComfyUI media directories are unavailable") from None
|
||||
roots = list(serverless_media._folder_roots(folder_paths).values())
|
||||
artifacts = {}
|
||||
source_artifacts = {}
|
||||
source_digests = set()
|
||||
filenames = {}
|
||||
loader_artifacts = {}
|
||||
file_total = decoded_total = 0
|
||||
for node_id, (kind, filename) in loaders.items():
|
||||
# Avoid retaining file buffers merely to deduplicate repeat references.
|
||||
if (kind, filename) in filenames:
|
||||
loader_artifacts[node_id] = filenames[kind, filename]
|
||||
continue
|
||||
path = _source_path(filename, folder_paths)
|
||||
data = video_codec._read_path(path, contract.MAX_ARTIFACT_BYTES, roots)
|
||||
source_digest = hashlib.sha256(data).hexdigest()
|
||||
if source_digest not in source_digests:
|
||||
if len(data) > contract.MAX_TOTAL_BYTES - file_total:
|
||||
raise Error("Native media source files exceed total byte limit")
|
||||
file_total += len(data)
|
||||
source_digests.add(source_digest)
|
||||
source_key = (kind, source_digest)
|
||||
if source_key not in source_artifacts:
|
||||
if len(artifacts) >= contract.MAX_ARTIFACTS:
|
||||
raise Error("Native media input artifact count exceeds limit")
|
||||
limit = min(contract.MAX_ARTIFACT_BYTES, contract.MAX_TOTAL_BYTES - decoded_total)
|
||||
if kind == "AUDIO":
|
||||
value = _decode_audio(data, limit)
|
||||
descriptor, payload = codec.encode_payload(kind, value, max_bytes=limit)
|
||||
del value
|
||||
else:
|
||||
from comfy_api.latest import InputImpl
|
||||
value = InputImpl.VideoFromFile(io.BytesIO(data))
|
||||
descriptor, payload = video_codec.encode_video(value, max_bytes=limit, allowed_roots=roots)
|
||||
del value
|
||||
aid = _artifact_id(kind, descriptor, payload)
|
||||
if aid not in artifacts:
|
||||
if len(payload) > limit:
|
||||
raise Error("Native media decoded inputs exceed total byte limit")
|
||||
artifacts[aid] = contract.artifact_from_payload(
|
||||
aid, kind, descriptor, payload, max_bytes=contract.MAX_ARTIFACT_BYTES)
|
||||
decoded_total += len(payload)
|
||||
source_artifacts[source_key] = aid
|
||||
del payload
|
||||
aid = source_artifacts[source_key]
|
||||
loader_artifacts[node_id] = filenames[kind, filename] = aid
|
||||
del data
|
||||
refs = [{"node_id": node_id, "input_name": key, "artifact_id": loader_artifacts[source]}
|
||||
for node_id, key, source in references]
|
||||
# Shared node dictionaries deliberately propagate only legacy image rewrites.
|
||||
images = await serverless_media.prepare_images(image_view)
|
||||
return images, list(artifacts.values()), refs
|
||||
@@ -0,0 +1,79 @@
|
||||
"""Explicit media-v1 preflight and atomic, lossless result materialization."""
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...serverless_bridge import contract
|
||||
from ...serverless_bridge.contract import MAX_ARTIFACT_BYTES, PROFILE, validate_bindings
|
||||
from .serverless_media import (
|
||||
MAX_REDIRECTS, _CHUNK_SIZE, _DOWNLOAD_TIMEOUT_SECONDS,
|
||||
_PublicResolver, _validate_artifact_url,
|
||||
)
|
||||
|
||||
|
||||
def validate_prepared_job(prepared):
|
||||
"""Fail before journalling/submission; legacy jobs keep their original path."""
|
||||
profile = prepared["worker"].get("media_profile", "legacy-image-v1")
|
||||
media = prepared.get("media")
|
||||
if profile == "legacy-image-v1" and media is None:
|
||||
return False
|
||||
if profile != PROFILE or media is None:
|
||||
raise ValueError("Serverless media profile/request mismatch")
|
||||
from .runpod_client import serialize_request
|
||||
# The shared endpoint validator also checks exporter controls and input refs.
|
||||
serialize_request(prepared["workflow"], prepared["images"],
|
||||
prepared["worker"]["timeout_seconds"], media=media)
|
||||
expected = validate_bindings(prepared["manifest"])
|
||||
if expected != validate_bindings(media["outputs"]):
|
||||
raise ValueError("Serverless media manifest identity mismatch")
|
||||
if any(binding["participant_id"] != prepared["worker"]["id"] for binding in expected.values()):
|
||||
raise ValueError("Serverless media participant identity mismatch")
|
||||
return True
|
||||
|
||||
|
||||
async def download_public_artifact(url, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""Contract-sized public HTTPS download, never using a provider session/key."""
|
||||
if type(max_bytes) is not int or not 0 <= max_bytes <= MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
current_url = _validate_artifact_url(url)
|
||||
connector = aiohttp.TCPConnector(resolver=_PublicResolver(), 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):
|
||||
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 ValueError("Media artifact redirect is invalid")
|
||||
current_url = _validate_artifact_url(urljoin(current_url, location))
|
||||
continue
|
||||
if response.status != 200:
|
||||
raise ValueError("Media artifact download failed")
|
||||
payload = bytearray()
|
||||
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
|
||||
if not isinstance(chunk, bytes) or len(payload) + len(chunk) > max_bytes:
|
||||
raise ValueError("Media artifact exceeds the permitted size")
|
||||
payload.extend(chunk)
|
||||
return bytes(payload)
|
||||
except Exception:
|
||||
# Signed URLs and transport exception details must not reach the ledger.
|
||||
raise ValueError("Media artifact download failed or exceeded its byte limit") from None
|
||||
raise ValueError("Media artifact download failed")
|
||||
|
||||
|
||||
async def collect_typed_results(output, media, manifest, *, download=None):
|
||||
"""Decode every binding before returning any collector's native value lists."""
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise ValueError("Invalid serverless media output")
|
||||
decoded = await contract.decode_response(output.get("media"), request_id=media["request_id"],
|
||||
bindings=manifest,
|
||||
download=download if download is not None else download_public_artifact)
|
||||
grouped = {}
|
||||
for binding in manifest:
|
||||
values = decoded[binding["binding_id"]]
|
||||
if not isinstance(values, list) or not values:
|
||||
raise ValueError("Missing serverless media values")
|
||||
grouped.setdefault(binding["job_id"], {}).setdefault(binding["socket"], []).extend(values)
|
||||
return grouped
|
||||
+378
-256
@@ -1,7 +1,9 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
from ..utils.config import load_config
|
||||
@@ -12,17 +14,20 @@ from ..utils.constants import (
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_master_url
|
||||
from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance
|
||||
from ..utils.network import build_master_url, build_master_callback_url
|
||||
from ..utils.trace_logger import trace_debug
|
||||
from .schemas import coerce_positive_float, coerce_positive_int
|
||||
from .schemas import parse_positive_float, parse_positive_int
|
||||
from .orchestration.dispatch import (
|
||||
dispatch_worker_prompt,
|
||||
rank_workers_by_load,
|
||||
select_active_workers,
|
||||
select_least_busy_worker,
|
||||
)
|
||||
from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
|
||||
from .orchestration.media_sync import (
|
||||
convert_paths_for_platform,
|
||||
fetch_worker_has_node,
|
||||
fetch_worker_path_separator,
|
||||
sync_worker_media,
|
||||
)
|
||||
from .orchestration.prompt_transform import (
|
||||
PromptIndex,
|
||||
apply_participant_overrides,
|
||||
@@ -30,7 +35,15 @@ from .orchestration.prompt_transform import (
|
||||
generate_job_id_map,
|
||||
prepare_delegate_master_prompt,
|
||||
prune_prompt_for_worker,
|
||||
WORKER_SINK_CLASS,
|
||||
_resolve_direct_workers_for_image_list,
|
||||
)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
from ..image_list.execution_plan import prepare_image_list_execution
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def _generate_execution_trace_id():
|
||||
return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
@@ -38,15 +51,23 @@ def _generate_execution_trace_id():
|
||||
|
||||
def ensure_distributed_state(server_instance=None):
|
||||
"""Ensure prompt_server has the state used by distributed queue orchestration."""
|
||||
ensure_distributed_runtime_state(server_instance)
|
||||
ps = server_instance or prompt_server
|
||||
if not hasattr(ps, "distributed_pending_jobs"):
|
||||
ps.distributed_pending_jobs = {}
|
||||
if not hasattr(ps, "distributed_jobs_lock"):
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
|
||||
|
||||
# Initialize top-level distributed queue state at module import time.
|
||||
ensure_distributed_state()
|
||||
|
||||
|
||||
async def _ensure_distributed_queue(job_id):
|
||||
"""Ensure a queue exists for the given distributed job ID."""
|
||||
runtime_state = ensure_distributed_runtime_state()
|
||||
async with runtime_state.distributed_jobs_lock:
|
||||
if job_id not in runtime_state.distributed_pending_jobs:
|
||||
runtime_state.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
|
||||
|
||||
def _resolve_enabled_workers(config, requested_ids=None):
|
||||
@@ -63,6 +84,18 @@ 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)
|
||||
if pool.get("media_profile") == "comfy-media-v1":
|
||||
from ..serverless_bridge.contract import identifier
|
||||
# Do not repair malformed media identities by stripping them.
|
||||
identifier(pool["id"])
|
||||
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)
|
||||
@@ -85,19 +118,19 @@ def _resolve_enabled_workers(config, requested_ids=None):
|
||||
def _resolve_orchestration_limits(config):
|
||||
"""Resolve bounded concurrency/timeouts for worker preparation pipeline."""
|
||||
settings = (config or {}).get("settings", {}) or {}
|
||||
worker_probe_concurrency = coerce_positive_int(
|
||||
worker_probe_concurrency = parse_positive_int(
|
||||
settings.get("worker_probe_concurrency"),
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
|
||||
)
|
||||
worker_prep_concurrency = coerce_positive_int(
|
||||
worker_prep_concurrency = parse_positive_int(
|
||||
settings.get("worker_prep_concurrency"),
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
media_sync_concurrency = coerce_positive_int(
|
||||
media_sync_concurrency = parse_positive_int(
|
||||
settings.get("media_sync_concurrency"),
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
|
||||
)
|
||||
media_sync_timeout_seconds = coerce_positive_float(
|
||||
media_sync_timeout_seconds = parse_positive_float(
|
||||
settings.get("media_sync_timeout_seconds"),
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
|
||||
)
|
||||
@@ -127,12 +160,39 @@ def _prompt_requests_load_balance(prompt_index):
|
||||
return False
|
||||
|
||||
|
||||
def _validate_media_workers(workers, prompt_index=None):
|
||||
"""Reject unknown serverless media profiles before probing or dispatching.
|
||||
|
||||
A Distributed Collector accepts direct workers, legacy-image-v1 pools and
|
||||
comfy-media-v1 pools in the same run; each pool's profile selects its
|
||||
endpoint contract. Unknown profiles fail rather than guessing a contract.
|
||||
legacy-image-v1 endpoints return 8-bit files, so a Collector set to exact
|
||||
cannot use them.
|
||||
"""
|
||||
exact = prompt_index is not None and any(
|
||||
str(prompt_index.inputs_by_node.get(node_id, {}).get("transfer_format", "")).strip().lower() == "exact"
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector")
|
||||
)
|
||||
for worker in workers:
|
||||
if worker.get("type") != "serverless":
|
||||
continue
|
||||
profile = worker.get("media_profile", "legacy-image-v1")
|
||||
if profile not in ("legacy-image-v1", "comfy-media-v1"):
|
||||
raise ValueError("Unsupported serverless media_profile")
|
||||
if exact and profile == "legacy-image-v1":
|
||||
raise ValueError(
|
||||
"transfer_format exact requires comfy-media-v1 serverless pools; "
|
||||
"legacy-image-v1 endpoints return 8-bit images"
|
||||
)
|
||||
|
||||
|
||||
async def _prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
trace_execution_id,
|
||||
worker_prep_semaphore,
|
||||
@@ -141,22 +201,73 @@ async def _prepare_worker_payload(
|
||||
):
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
_validate_media_workers([worker], prompt_index)
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
if worker.get("type") == "serverless":
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
if worker.get("media_profile", "legacy-image-v1") == "comfy-media-v1":
|
||||
from ..serverless_bridge.contract import PROFILE, VERSION
|
||||
from .orchestration.serverless_media_prompt import compile_media_prompt
|
||||
from .orchestration.serverless_typed_inputs import prepare_media_inputs
|
||||
# The typed compiler validates the complete private graph and
|
||||
# prunes it itself, retaining every connected Collector branch.
|
||||
worker_prompt = apply_participant_overrides(
|
||||
worker_prompt, worker["id"], enabled_ids, job_id_map,
|
||||
"", delegate_master, prompt_index,
|
||||
)
|
||||
request_id = str(uuid.uuid4())
|
||||
workflow, manifest = compile_media_prompt(worker_prompt, request_id)
|
||||
images, artifacts, refs = await prepare_media_inputs(workflow)
|
||||
media = {
|
||||
"version": VERSION, "profile": PROFILE, "request_id": request_id,
|
||||
"artifacts": artifacts, "inputs": refs, "outputs": copy.deepcopy(manifest),
|
||||
}
|
||||
serialize_request(workflow, images, worker["timeout_seconds"], media=media)
|
||||
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest,
|
||||
"images": images, "media": media}
|
||||
|
||||
from .orchestration.serverless_prompt import compile_serverless_prompt
|
||||
from .orchestration.serverless_media import prepare_images
|
||||
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,
|
||||
prompt_server_instance=prompt_server,
|
||||
)
|
||||
|
||||
worker_type = str(worker.get("type") or "local").strip().lower()
|
||||
is_remote_like = bool(worker.get("host")) and worker_type != "local"
|
||||
if is_remote_like:
|
||||
path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
|
||||
if path_separator:
|
||||
worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
|
||||
# Older workers lack the sink node; they keep the PreviewImage terminator.
|
||||
has_sink_task = asyncio.ensure_future(
|
||||
fetch_worker_has_node(worker, WORKER_SINK_CLASS, trace_execution_id=trace_execution_id)
|
||||
)
|
||||
try:
|
||||
if is_remote_like:
|
||||
path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
|
||||
if path_separator:
|
||||
worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
|
||||
use_worker_sink = await has_sink_task
|
||||
finally:
|
||||
if not has_sink_task.done():
|
||||
has_sink_task.cancel()
|
||||
|
||||
worker_prompt = prune_prompt_for_worker(worker_prompt)
|
||||
worker_prompt = prune_prompt_for_worker(worker_prompt, use_worker_sink=use_worker_sink)
|
||||
worker_prompt = apply_participant_overrides(
|
||||
worker_prompt,
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
worker_master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
@@ -180,27 +291,135 @@ async def _prepare_worker_payload(
|
||||
return worker, worker_prompt
|
||||
|
||||
|
||||
async def _select_execution_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
load_balance_requested,
|
||||
has_branch_nodes,
|
||||
master_url,
|
||||
execution_trace_id,
|
||||
worker_probe_concurrency,
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=None,
|
||||
delegate_master=None,
|
||||
trace_execution_id=None,
|
||||
):
|
||||
active_workers, delegate_master = await select_active_workers(
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
|
||||
(
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
image_list_mode = any(isinstance(node, dict) and node.get("class_type") == "DistributedImageListStart"
|
||||
for node in prompt_obj.values())
|
||||
target_kwargs = {}
|
||||
if image_list_mode:
|
||||
# Check the source graph before execution-root pruning can remove an
|
||||
# incompatible mode (or mask it with an unrelated boundary error).
|
||||
if any(isinstance(node, dict) and node.get("class_type") in (
|
||||
"DistributedCollector", "UltimateSDUpscaleDistributed") for node in prompt_obj.values()):
|
||||
raise ValueError(
|
||||
"Distributed Image List v1 cannot be mixed with Distributed Collector or "
|
||||
"Ultimate SD Upscale nodes in one prompt"
|
||||
)
|
||||
prompt_obj, outer_targets = prepare_image_list_execution(prompt_obj)
|
||||
target_kwargs["partial_execution_targets"] = outer_targets
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
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, explicit_targets=True)
|
||||
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 and delegate_master:
|
||||
raise RuntimeError(
|
||||
"Distributed Image List in orchestrator-only mode requires a selected direct worker; "
|
||||
"serverless pools are unsupported"
|
||||
)
|
||||
else:
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
_validate_media_workers(workers, prompt_index)
|
||||
if any(worker.get("media_profile") == "comfy-media-v1" for worker in workers):
|
||||
from ..serverless_bridge.contract import identifier
|
||||
# Validate the source, not PromptIndex's JSON copy (which coerces
|
||||
# integer keys to strings), before probing or dispatching workers.
|
||||
for node_id in prompt_obj:
|
||||
identifier(node_id)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
(
|
||||
f"Orchestration start: requested_workers={len(workers)}, "
|
||||
f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
|
||||
f"websocket={use_websocket}, "
|
||||
f"probe_concurrency={worker_probe_concurrency}, "
|
||||
f"prep_concurrency={worker_prep_concurrency}, "
|
||||
f"media_sync_concurrency={media_sync_concurrency}, "
|
||||
f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
|
||||
f"load_balance={load_balance_requested}"
|
||||
),
|
||||
)
|
||||
|
||||
# Keep the legacy Collector fallback, but never override image-list
|
||||
# orchestrator-only mode merely because its direct workers are unavailable.
|
||||
if not image_list_mode and not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, selected_delegate_master = await select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
False if image_list_mode else delegate_master,
|
||||
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 and delegate_master:
|
||||
raise RuntimeError(
|
||||
"No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
|
||||
)
|
||||
from ..image_list.runtime import END_CLASS, end_override_names
|
||||
requires_values = any(end_override_names(node.get("inputs", {}))
|
||||
for node in prompt_obj.values() if node.get("class_type") == END_CLASS)
|
||||
value_preflight = {"require_value_results": True} if requires_values else {}
|
||||
# Persist candidates, not this submission-time idle snapshot: a preceding
|
||||
# master prompt may still be using them. The image-list runtime rechecks
|
||||
# capability and availability before assigning the first item.
|
||||
if active_workers:
|
||||
await preflight_image_list_workers(active_workers, require_idle=False,
|
||||
**value_preflight)
|
||||
else:
|
||||
delegate_master = selected_delegate_master
|
||||
|
||||
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.
|
||||
candidate_workers.append(
|
||||
{
|
||||
"id": "master",
|
||||
@@ -225,6 +444,7 @@ async def _select_execution_workers(
|
||||
selected_worker = candidate_workers[0]
|
||||
|
||||
if selected_worker is not None and str(selected_worker.get("id")) == "master":
|
||||
# Master selected as least busy; run master workload only.
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
trace_debug(
|
||||
@@ -233,6 +453,7 @@ async def _select_execution_workers(
|
||||
)
|
||||
elif selected_worker is not None:
|
||||
active_workers = [selected_worker]
|
||||
# Worker selected as least busy; keep master orchestrator-only for this run.
|
||||
delegate_master = True
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
@@ -246,47 +467,27 @@ async def _select_execution_workers(
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
|
||||
if has_branch_nodes and len(active_workers) > 1:
|
||||
active_workers = await rank_workers_by_load(
|
||||
active_workers,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
queue_result = await queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
queue_result["number"],
|
||||
0,
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
return active_workers, delegate_master
|
||||
|
||||
|
||||
def _distributed_queue_nodes(prompt_index):
|
||||
return (
|
||||
prompt_index.nodes_for_class("DistributedCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedBranchCollector")
|
||||
+ prompt_index.nodes_for_class("DistributedBranch")
|
||||
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
|
||||
)
|
||||
|
||||
|
||||
def _register_job_allowed_workers(job_id_map, enabled_ids):
|
||||
runtime_state = ensure_distributed_runtime_state()
|
||||
allowed = {str(worker_id) for worker_id in (enabled_ids or []) if str(worker_id).strip()}
|
||||
for job_id in job_id_map.values():
|
||||
if job_id:
|
||||
runtime_state.distributed_job_allowed_workers[str(job_id)] = set(allowed)
|
||||
|
||||
|
||||
async def _ensure_job_queues(prompt_index, job_id_map):
|
||||
for node_id in _distributed_queue_nodes(prompt_index):
|
||||
job_id = job_id_map.get(node_id)
|
||||
if job_id:
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
|
||||
def _build_master_prompt(
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
):
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = apply_participant_overrides(
|
||||
master_prompt,
|
||||
@@ -296,173 +497,22 @@ def _build_master_prompt(
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
workflow_meta=workflow_meta,
|
||||
)
|
||||
|
||||
if not delegate_master:
|
||||
return master_prompt
|
||||
|
||||
collector_ids = (
|
||||
find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
+ find_nodes_by_class(master_prompt, "DistributedBranchCollector")
|
||||
)
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
# Include USDU nodes as collector-like for delegate pruning
|
||||
collector_ids.extend(upscale_nodes)
|
||||
if not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collector/branch-collector nodes found in master prompt. Running full prompt on master."
|
||||
)
|
||||
return master_prompt
|
||||
return prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
|
||||
async def _prepare_worker_payloads_for_dispatch(
|
||||
active_workers,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
):
|
||||
if not active_workers:
|
||||
return []
|
||||
|
||||
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
|
||||
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
|
||||
return await asyncio.gather(
|
||||
*[
|
||||
_prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
if delegate_master and not image_list_mode:
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
|
||||
)
|
||||
for worker in active_workers
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def _dispatch_worker_payloads(
|
||||
worker_payloads,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
use_websocket,
|
||||
execution_trace_id,
|
||||
):
|
||||
if not worker_payloads:
|
||||
return
|
||||
|
||||
await asyncio.gather(
|
||||
*[
|
||||
dispatch_worker_prompt(
|
||||
worker,
|
||||
wprompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
use_websocket=use_websocket,
|
||||
trace_execution_id=execution_trace_id,
|
||||
elif not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
|
||||
)
|
||||
for worker, wprompt in worker_payloads
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj: dict[str, Any],
|
||||
workflow_meta: dict[str, Any] | None,
|
||||
client_id: str | None,
|
||||
enabled_worker_ids: list[str] | set[str] | None = None,
|
||||
delegate_master: bool | None = None,
|
||||
trace_execution_id: str | None = None,
|
||||
) -> tuple[str, int]:
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (prompt_id, worker_count)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=get_prompt_server_instance())
|
||||
(
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
has_branch_nodes = bool(prompt_index.nodes_for_class("DistributedBranch"))
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
(
|
||||
f"Orchestration start: requested_workers={len(workers)}, "
|
||||
f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
|
||||
f"websocket={use_websocket}, "
|
||||
f"probe_concurrency={worker_probe_concurrency}, "
|
||||
f"prep_concurrency={worker_prep_concurrency}, "
|
||||
f"media_sync_concurrency={media_sync_concurrency}, "
|
||||
f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
|
||||
f"load_balance={load_balance_requested}, has_branch_nodes={has_branch_nodes}"
|
||||
),
|
||||
)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, delegate_master = await _select_execution_workers(
|
||||
workers=workers,
|
||||
use_websocket=use_websocket,
|
||||
delegate_master=delegate_master,
|
||||
load_balance_requested=load_balance_requested,
|
||||
has_branch_nodes=has_branch_nodes,
|
||||
master_url=master_url,
|
||||
execution_trace_id=execution_trace_id,
|
||||
worker_probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
|
||||
return prompt_id, 0
|
||||
|
||||
_register_job_allowed_workers(job_id_map, enabled_ids)
|
||||
await _ensure_job_queues(prompt_index, job_id_map)
|
||||
|
||||
master_prompt = _build_master_prompt(
|
||||
prompt_index=prompt_index,
|
||||
enabled_ids=enabled_ids,
|
||||
job_id_map=job_id_map,
|
||||
master_url=master_url,
|
||||
delegate_master=delegate_master,
|
||||
)
|
||||
else:
|
||||
master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
if active_workers:
|
||||
trace_debug(
|
||||
@@ -470,29 +520,101 @@ async def orchestrate_distributed_execution(
|
||||
"Active distributed workers: "
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
|
||||
)
|
||||
worker_payloads = await _prepare_worker_payloads_for_dispatch(
|
||||
active_workers=active_workers,
|
||||
prompt_index=prompt_index,
|
||||
enabled_ids=enabled_ids,
|
||||
job_id_map=job_id_map,
|
||||
master_url=master_url,
|
||||
delegate_master=delegate_master,
|
||||
execution_trace_id=execution_trace_id,
|
||||
worker_prep_concurrency=worker_prep_concurrency,
|
||||
media_sync_concurrency=media_sync_concurrency,
|
||||
media_sync_timeout_seconds=media_sync_timeout_seconds,
|
||||
)
|
||||
await _dispatch_worker_payloads(
|
||||
worker_payloads=worker_payloads,
|
||||
workflow_meta=workflow_meta,
|
||||
client_id=client_id,
|
||||
use_websocket=use_websocket,
|
||||
execution_trace_id=execution_trace_id,
|
||||
)
|
||||
worker_payloads = []
|
||||
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(
|
||||
*[
|
||||
_prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
)
|
||||
for worker in active_workers
|
||||
]
|
||||
)
|
||||
|
||||
prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
|
||||
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
|
||||
# Repeat the manager's whole-batch identity/profile preflight *before*
|
||||
# any direct dispatch, master queueing or collector queue allocation.
|
||||
# Pure legacy jobs are a no-op in that validator; keep their path free
|
||||
# of typed imports (and compatible with existing legacy media adapters).
|
||||
for prepared in serverless_jobs:
|
||||
if (prepared["worker"].get("media_profile", "legacy-image-v1") != "legacy-image-v1"
|
||||
or prepared.get("media") is not None):
|
||||
from .orchestration.serverless_typed_results import validate_prepared_job
|
||||
validate_prepared_job(prepared)
|
||||
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(
|
||||
worker,
|
||||
wprompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
use_websocket=use_websocket,
|
||||
trace_execution_id=execution_trace_id,
|
||||
)
|
||||
for worker, wprompt in regular_payloads
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
**target_kwargs,
|
||||
)
|
||||
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:
|
||||
# Typed manifests have several sockets per collector, but
|
||||
# each participant gets only one terminal collector event.
|
||||
for job_id in {binding["job_id"] for binding in prepared["manifest"]}:
|
||||
q = prompt_server.distributed_pending_jobs.get(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, len(worker_payloads)
|
||||
return prompt_id, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
|
||||
|
||||
+15
-20
@@ -1,20 +1,16 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from ..utils.worker_ids import require_enabled_worker_ids
|
||||
except ImportError: # pragma: no cover - supports direct module loading in isolated tests.
|
||||
from utils.worker_ids import require_enabled_worker_ids
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueRequestPayload:
|
||||
prompt: dict[str, Any]
|
||||
workflow_meta: dict[str, Any] | None
|
||||
prompt: Dict[str, Any]
|
||||
workflow_meta: Any
|
||||
client_id: str
|
||||
delegate_master: bool | None
|
||||
enabled_worker_ids: list[str]
|
||||
trace_execution_id: str | None
|
||||
delegate_master: Optional[bool]
|
||||
enabled_worker_ids: List[str]
|
||||
auto_prepare: bool
|
||||
trace_execution_id: Optional[str]
|
||||
|
||||
|
||||
def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
@@ -22,6 +18,11 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
|
||||
auto_prepare_raw = data.get("auto_prepare", True)
|
||||
if not isinstance(auto_prepare_raw, bool):
|
||||
raise ValueError("auto_prepare must be a boolean when provided")
|
||||
auto_prepare = auto_prepare_raw
|
||||
|
||||
prompt = data.get("prompt")
|
||||
# Auto-prepare is always on server-side; keep the field for wire compatibility.
|
||||
if prompt is None:
|
||||
@@ -34,10 +35,6 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
|
||||
workflow_meta = data.get("workflow")
|
||||
if workflow_meta is not None and not isinstance(workflow_meta, dict):
|
||||
raise ValueError("Field 'workflow' must be an object when provided")
|
||||
|
||||
enabled_ids_raw = data.get("enabled_worker_ids")
|
||||
workers_field = data.get("workers")
|
||||
if enabled_ids_raw is None and workers_field is not None:
|
||||
@@ -54,10 +51,7 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
else:
|
||||
if not isinstance(enabled_ids_raw, list):
|
||||
raise ValueError("enabled_worker_ids must be a list of worker IDs")
|
||||
try:
|
||||
enabled_ids = require_enabled_worker_ids(enabled_ids_raw)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"enabled_worker_ids invalid: {exc}") from exc
|
||||
enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
|
||||
|
||||
delegate_master = data.get("delegate_master")
|
||||
if delegate_master is not None and not isinstance(delegate_master, bool):
|
||||
@@ -76,9 +70,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
|
||||
return QueueRequestPayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=workflow_meta,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=delegate_master,
|
||||
enabled_worker_ids=enabled_ids,
|
||||
auto_prepare=auto_prepare,
|
||||
trace_execution_id=trace_execution_id,
|
||||
)
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
"""Shared request guards for distributed API endpoints."""
|
||||
from __future__ import annotations
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.network import handle_api_error
|
||||
from .schemas import is_authorized_request
|
||||
|
||||
|
||||
async def authorization_error_or_none(request: web.Request) -> web.StreamResponse | None:
|
||||
"""Return a 403 response when distributed API auth fails, else None."""
|
||||
if is_authorized_request(request, load_config()):
|
||||
return None
|
||||
return await handle_api_error(request, "Unauthorized", 403)
|
||||
+15
-61
@@ -1,29 +1,3 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ..utils.auth import (
|
||||
AUTH_HEADER_NAME,
|
||||
distributed_auth_headers,
|
||||
is_authorized_request,
|
||||
)
|
||||
from ..utils.parsing import coerce_bool, require_bool_literal
|
||||
|
||||
__all__ = [
|
||||
"AUTH_HEADER_NAME",
|
||||
"coerce_bool",
|
||||
"coerce_positive_float",
|
||||
"coerce_positive_int",
|
||||
"distributed_auth_headers",
|
||||
"is_authorized_request",
|
||||
"require_bool_literal",
|
||||
"require_fields",
|
||||
"require_worker_id",
|
||||
"require_positive_int",
|
||||
"validate_worker_id",
|
||||
]
|
||||
|
||||
|
||||
def require_fields(data: dict, *fields) -> list[str]:
|
||||
"""Return field names that are missing or empty in a JSON object."""
|
||||
if not isinstance(data, dict):
|
||||
@@ -44,43 +18,26 @@ def require_fields(data: dict, *fields) -> list[str]:
|
||||
return missing
|
||||
|
||||
|
||||
def validate_worker_id(worker_id: str, config: dict[str, Any]) -> bool:
|
||||
def validate_worker_id(worker_id: str, config: dict) -> bool:
|
||||
"""Return True when worker_id exists in config['workers']."""
|
||||
try:
|
||||
require_worker_id(worker_id, config)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def require_worker_id(worker_id: str, config: dict[str, Any], field_name: str = "worker_id") -> str:
|
||||
"""Strict worker-id validator for request boundaries."""
|
||||
worker_id_str = str(worker_id).strip()
|
||||
if not worker_id_str:
|
||||
raise ValueError(f"Field '{field_name}' must be a non-empty worker id.")
|
||||
|
||||
worker_id_str = str(worker_id)
|
||||
workers = (config or {}).get("workers", [])
|
||||
for worker in workers:
|
||||
if not isinstance(worker, dict):
|
||||
continue
|
||||
if str(worker.get("id")).strip() == worker_id_str:
|
||||
return worker_id_str
|
||||
raise ValueError(f"Worker {worker_id_str} not found.")
|
||||
return any(str(worker.get("id")) == worker_id_str for worker in workers)
|
||||
|
||||
|
||||
def require_positive_int(value, field_name: str = "value") -> int:
|
||||
"""Strict positive-int parser for request boundary validation."""
|
||||
def validate_positive_int(value, field_name: str) -> str | None:
|
||||
"""Validate positive integers and return an error string when invalid."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"Field '{field_name}' must be a positive integer.") from exc
|
||||
except (TypeError, ValueError):
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
if parsed <= 0:
|
||||
raise ValueError(f"Field '{field_name}' must be a positive integer.")
|
||||
return parsed
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
return None
|
||||
|
||||
|
||||
def coerce_positive_int(value, default: int) -> int:
|
||||
"""Permissive positive-int coercion used for non-boundary defaults."""
|
||||
def parse_positive_int(value, default: int) -> int:
|
||||
"""Parse value as positive int, returning default on failure."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
@@ -88,13 +45,10 @@ def coerce_positive_int(value, default: int) -> int:
|
||||
return max(1, parsed)
|
||||
|
||||
|
||||
def coerce_positive_float(value, default: float) -> float:
|
||||
"""Permissive positive-float coercion used for non-boundary defaults."""
|
||||
fallback = float(default)
|
||||
if fallback <= 0.0:
|
||||
fallback = 0.1
|
||||
def parse_positive_float(value, default: float) -> float:
|
||||
"""Parse value as positive float, returning default on failure."""
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return fallback
|
||||
return parsed if parsed > 0.0 else fallback
|
||||
return max(0.0, float(default))
|
||||
return max(0.0, parsed)
|
||||
|
||||
@@ -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 worker"}, 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)
|
||||
+29
-41
@@ -1,63 +1,51 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from .endpoint_policy import run_authorized_endpoint
|
||||
from ..utils.cloudflare import cloudflare_tunnel_manager
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
|
||||
async def tunnel_status_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def tunnel_status_endpoint(request):
|
||||
"""Return Cloudflare tunnel status and last known details."""
|
||||
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
status = cloudflare_tunnel_manager.get_status()
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host")
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"tunnel": status,
|
||||
"master_host": master_host,
|
||||
}
|
||||
)
|
||||
|
||||
return await run_authorized_endpoint(request, _operation)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": status,
|
||||
"master_host": master_host
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
|
||||
async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def tunnel_start_endpoint(request):
|
||||
"""Start a Cloudflare tunnel pointing at the current ComfyUI server."""
|
||||
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.start_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host"),
|
||||
}
|
||||
)
|
||||
|
||||
return await run_authorized_endpoint(request, _operation)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
|
||||
async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def tunnel_stop_endpoint(request):
|
||||
"""Stop the managed Cloudflare tunnel if running."""
|
||||
|
||||
async def _operation() -> web.StreamResponse:
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.stop_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host"),
|
||||
}
|
||||
)
|
||||
|
||||
return await run_authorized_endpoint(request, _operation)
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
+21
-137
@@ -1,5 +1,3 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import time
|
||||
@@ -8,78 +6,15 @@ from aiohttp import web
|
||||
from PIL import Image
|
||||
import server
|
||||
|
||||
from .request_guards import authorization_error_or_none
|
||||
from .schemas import require_bool_literal
|
||||
from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job
|
||||
from ..upscale.payload_parsers import parse_tiles_from_form
|
||||
from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
|
||||
from ..upscale.payload_parsers import _parse_tiles_from_form
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.network import handle_api_error
|
||||
from ..utils.usdu_management import MAX_PAYLOAD_SIZE
|
||||
|
||||
|
||||
def _parse_int_field(value, field_name: str, *, minimum: int | None = None) -> int:
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"{field_name} must be an integer") from exc
|
||||
if minimum is not None and parsed < minimum:
|
||||
raise ValueError(f"{field_name} must be >= {minimum}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_bool_field(
|
||||
value,
|
||||
field_name: str,
|
||||
*,
|
||||
default: bool = False,
|
||||
) -> bool:
|
||||
if value is None:
|
||||
return default
|
||||
return require_bool_literal(value, field_name=field_name)
|
||||
|
||||
|
||||
def _is_worker_allowed(job_data: BaseJobState, worker_id: str) -> bool:
|
||||
known_workers = {str(wid) for wid in getattr(job_data, "worker_status", {}).keys()}
|
||||
known_workers.update(str(wid) for wid in getattr(job_data, "assigned_to_workers", {}).keys())
|
||||
if not known_workers:
|
||||
return True
|
||||
return str(worker_id) in known_workers
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/init_dynamic_job")
|
||||
async def init_dynamic_job_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
"""Allow workers to initialize a dynamic job queue on the master.
|
||||
|
||||
Called by the first worker to reach the USDU node in delegate-only mode.
|
||||
Idempotent — subsequent calls are no-ops if the queue already exists.
|
||||
"""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
data = await request.json()
|
||||
multi_job_id = data.get("multi_job_id")
|
||||
if not multi_job_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id", 400)
|
||||
|
||||
batch_size = _parse_int_field(data.get("batch_size"), "batch_size", minimum=1)
|
||||
enabled_workers = data.get("enabled_workers") or []
|
||||
if isinstance(enabled_workers, str):
|
||||
import json as _json
|
||||
enabled_workers = _json.loads(enabled_workers)
|
||||
|
||||
await init_dynamic_job(multi_job_id, batch_size, enabled_workers)
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
@@ -93,8 +28,6 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if isinstance(job_data, BaseJobState):
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
debug_log(f"Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
@@ -105,38 +38,24 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
||||
async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def submit_tiles_endpoint(request):
|
||||
"""Endpoint for workers to submit processed tiles in static mode."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length:
|
||||
try:
|
||||
payload_size = int(content_length)
|
||||
except ValueError:
|
||||
return await handle_api_error(request, "Invalid content-length header", 400)
|
||||
if payload_size > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
try:
|
||||
is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
try:
|
||||
batch_size = _parse_int_field(data.get('batch_size', 0), "batch_size", minimum=0)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
|
||||
# Handle completion signal
|
||||
if batch_size == 0 and is_last:
|
||||
@@ -145,8 +64,6 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
@@ -156,7 +73,7 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
try:
|
||||
tiles = parse_tiles_from_form(data)
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
@@ -166,8 +83,6 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
|
||||
q = job_data.queue
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
@@ -191,28 +106,17 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
||||
async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def submit_image_endpoint(request):
|
||||
"""Endpoint for workers to submit processed images in dynamic mode."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length:
|
||||
try:
|
||||
payload_size = int(content_length)
|
||||
except ValueError:
|
||||
return await handle_api_error(request, "Invalid content-length header", 400)
|
||||
if payload_size > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
try:
|
||||
is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
@@ -221,10 +125,7 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
# Handle image submission
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
try:
|
||||
image_idx = _parse_int_field(data.get('image_idx'), "image_idx", minimum=0)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
@@ -235,8 +136,6 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
@@ -252,8 +151,6 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
@@ -269,11 +166,8 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
@@ -296,8 +190,6 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
pending_queue = job_data.pending_tasks
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid job configuration", 400)
|
||||
if not _is_worker_allowed(job_data, worker_id):
|
||||
return await handle_api_error(request, "Unauthorized worker_id", 403)
|
||||
|
||||
try:
|
||||
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
@@ -307,33 +199,25 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
if mode == 'dynamic':
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
||||
return web.json_response(
|
||||
{
|
||||
"kind": "image",
|
||||
"task_idx": task_idx,
|
||||
"estimated_remaining": remaining,
|
||||
}
|
||||
)
|
||||
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
||||
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
||||
return web.json_response({
|
||||
"kind": "tile",
|
||||
"task_idx": task_idx,
|
||||
"tile_idx": task_idx,
|
||||
"estimated_remaining": remaining,
|
||||
"batched_static": job_data.batched_static,
|
||||
})
|
||||
except asyncio.TimeoutError:
|
||||
return web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0})
|
||||
if mode == 'dynamic':
|
||||
return web.json_response({"image_idx": None})
|
||||
return web.json_response({"tile_idx": None})
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
|
||||
+80
-109
@@ -1,13 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import platform
|
||||
import subprocess # nosec B404 - commands are fixed and never shell-expanded
|
||||
import subprocess
|
||||
import socket
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import aiohttp
|
||||
@@ -19,33 +16,35 @@ 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 .request_guards import authorization_error_or_none
|
||||
from .schemas import (
|
||||
distributed_auth_headers,
|
||||
require_fields,
|
||||
require_worker_id,
|
||||
)
|
||||
from ..workers.ports import allocate_worker_ports
|
||||
from .schemas import require_fields, validate_worker_id
|
||||
from ..workers.detection import (
|
||||
get_machine_id,
|
||||
is_docker_environment,
|
||||
is_runpod_environment,
|
||||
)
|
||||
from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
|
||||
try:
|
||||
from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
|
||||
except ImportError:
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
def __init__(self, message, validation_error=None, node_errors=None):
|
||||
super().__init__(str(message))
|
||||
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_ws")
|
||||
async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def worker_ws_endpoint(request):
|
||||
"""WebSocket endpoint for worker prompt dispatch."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
|
||||
ws = web.WebSocketResponse(heartbeat=30)
|
||||
await ws.prepare(request)
|
||||
|
||||
@@ -62,6 +61,13 @@ async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
})
|
||||
continue
|
||||
|
||||
if not isinstance(data, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack", "request_id": None, "ok": False,
|
||||
"error": "Expected a websocket message object.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
@@ -82,10 +88,16 @@ async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
continue
|
||||
|
||||
try:
|
||||
target_kwargs = {}
|
||||
if "partial_execution_targets" in data:
|
||||
from ..utils.async_helpers import validate_execution_targets
|
||||
target_kwargs["partial_execution_targets"] = validate_execution_targets(
|
||||
prompt, data["partial_execution_targets"])
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
**target_kwargs,
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
@@ -116,11 +128,8 @@ async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
|
||||
async def clear_launching_state(request: web.Request) -> web.StreamResponse:
|
||||
async def clear_launching_state(request):
|
||||
"""Clear the launching flag when worker is confirmed running."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
@@ -130,10 +139,8 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
try:
|
||||
worker_id = require_worker_id(worker_id, config)
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 404)
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Clear launching flag in managed processes
|
||||
if worker_id in wm.processes:
|
||||
@@ -147,9 +154,8 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
def get_network_ips() -> list[str]:
|
||||
def get_network_ips():
|
||||
"""Get all network IPs, trying multiple methods."""
|
||||
command_timeout = 5.0
|
||||
ips = []
|
||||
hostname = socket.gethostname()
|
||||
|
||||
@@ -160,8 +166,8 @@ def get_network_ips() -> list[str]:
|
||||
ip = info[4][0]
|
||||
if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now
|
||||
ips.append(ip)
|
||||
except (socket.gaierror, OSError) as exc:
|
||||
debug_log(f"get_network_ips: getaddrinfo failed for hostname {hostname}: {exc}")
|
||||
except (socket.gaierror, OSError):
|
||||
pass
|
||||
|
||||
# Method 2: Try to connect to external server and get local IP
|
||||
try:
|
||||
@@ -171,19 +177,14 @@ def get_network_ips() -> list[str]:
|
||||
s.close()
|
||||
if local_ip not in ips:
|
||||
ips.append(local_ip)
|
||||
except (OSError, socket.error) as exc:
|
||||
debug_log(f"get_network_ips: UDP local IP probe failed: {exc}")
|
||||
except (OSError, socket.error):
|
||||
pass
|
||||
|
||||
# Method 3: Platform-specific commands
|
||||
try:
|
||||
if platform.system() == "Windows":
|
||||
# Windows ipconfig
|
||||
result = subprocess.run( # nosec B603 - static command, no user input
|
||||
["ipconfig"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=command_timeout,
|
||||
)
|
||||
result = subprocess.run(["ipconfig"], capture_output=True, text=True)
|
||||
lines = result.stdout.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
if 'IPv4' in line and i + 1 < len(lines):
|
||||
@@ -193,20 +194,10 @@ def get_network_ips() -> list[str]:
|
||||
else:
|
||||
# Unix/Linux/Mac ifconfig or ip addr
|
||||
try:
|
||||
result = subprocess.run( # nosec B603 - static command, no user input
|
||||
["ip", "addr"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=command_timeout,
|
||||
)
|
||||
result = subprocess.run(["ip", "addr"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
try:
|
||||
result = subprocess.run( # nosec B603 - static command, no user input
|
||||
["ifconfig"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=command_timeout,
|
||||
)
|
||||
result = subprocess.run(["ifconfig"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
result = None
|
||||
|
||||
@@ -217,13 +208,13 @@ def get_network_ips() -> list[str]:
|
||||
ip = match.group(1)
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
except (OSError, subprocess.SubprocessError) as exc:
|
||||
debug_log(f"get_network_ips: platform command probe failed: {exc}")
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
pass
|
||||
|
||||
return ips
|
||||
|
||||
|
||||
def get_recommended_ip(ips: list[str]) -> str | None:
|
||||
def get_recommended_ip(ips):
|
||||
"""Choose the best IP for master-worker communication."""
|
||||
# Priority order:
|
||||
# 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
|
||||
@@ -258,7 +249,7 @@ def get_recommended_ip(ips: list[str]) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _get_cuda_info() -> tuple[int | None, int, int]:
|
||||
def _get_cuda_info():
|
||||
"""Detect CUDA device index and total physical GPU count.
|
||||
|
||||
Returns (cuda_device, cuda_device_count, physical_device_count).
|
||||
@@ -274,7 +265,7 @@ def _get_cuda_info() -> tuple[int | None, int, int]:
|
||||
if visible_devices:
|
||||
cuda_device = visible_devices[0]
|
||||
try:
|
||||
result = subprocess.run( # nosec B603 - static command, no user input
|
||||
result = subprocess.run(
|
||||
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
@@ -298,9 +289,16 @@ def _get_cuda_info() -> tuple[int | None, int, int]:
|
||||
return None, 0, 0
|
||||
|
||||
|
||||
def _collect_network_info_sync() -> dict[str, Any]:
|
||||
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)
|
||||
@@ -309,11 +307,13 @@ def _collect_network_info_sync() -> dict[str, Any]:
|
||||
"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,
|
||||
}
|
||||
|
||||
|
||||
def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
|
||||
def _read_worker_log_sync(log_file, lines_to_read):
|
||||
"""Read worker log content from disk in a threadpool worker."""
|
||||
file_size = os.path.getsize(log_file)
|
||||
|
||||
@@ -349,12 +349,7 @@ def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _parse_positive_int_query(
|
||||
value: Any,
|
||||
default: int,
|
||||
minimum: int = 1,
|
||||
maximum: int | None = 10000,
|
||||
) -> int:
|
||||
def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
|
||||
"""Parse bounded positive integer query params with sane fallback."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
@@ -366,7 +361,7 @@ def _parse_positive_int_query(
|
||||
return parsed
|
||||
|
||||
|
||||
def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any] | None:
|
||||
def _find_worker_by_id(config, worker_id):
|
||||
worker_id_str = str(worker_id).strip()
|
||||
for worker in config.get("workers", []):
|
||||
if str(worker.get("id")).strip() == worker_id_str:
|
||||
@@ -375,11 +370,8 @@ def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any]
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local_log")
|
||||
async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_local_log_endpoint(request):
|
||||
"""Return this instance's in-memory ComfyUI log buffer."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
from app.logger import get_logs
|
||||
except Exception as e:
|
||||
@@ -423,11 +415,8 @@ async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/network_info")
|
||||
async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_network_info_endpoint(request):
|
||||
"""Get network interfaces and recommend best IP for master."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
info = await loop.run_in_executor(None, _collect_network_info_sync)
|
||||
@@ -441,11 +430,8 @@ async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/system_info")
|
||||
async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_system_info_endpoint(request):
|
||||
"""Get system information including machine ID for local worker detection."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
import socket
|
||||
|
||||
@@ -468,11 +454,8 @@ async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/launch_worker")
|
||||
async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def launch_worker_endpoint(request):
|
||||
"""Launch a worker process from the UI."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
@@ -484,14 +467,15 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
# Find worker config
|
||||
config = load_config()
|
||||
try:
|
||||
worker_id = require_worker_id(worker_id, config)
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 404)
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
return await handle_api_error(request, "Serverless workers are managed by Runpod, not local processes", 400)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = worker_id
|
||||
|
||||
@@ -506,7 +490,7 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
is_running = process.poll() is None
|
||||
else:
|
||||
# Restored process without subprocess object
|
||||
is_running = wm.is_process_running(proc_info['pid'])
|
||||
is_running = wm._is_process_running(proc_info['pid'])
|
||||
|
||||
if is_running:
|
||||
return await handle_api_error(request, "Worker already running (managed by UI)", 409)
|
||||
@@ -526,6 +510,8 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
"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)
|
||||
|
||||
@@ -534,11 +520,8 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/stop_worker")
|
||||
async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def stop_worker_endpoint(request):
|
||||
"""Stop a worker process that was launched from the UI."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
@@ -548,10 +531,8 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
try:
|
||||
worker_id = require_worker_id(worker_id, config)
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 404)
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
success, message = wm.stop_worker(worker_id)
|
||||
|
||||
@@ -569,11 +550,8 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/managed_workers")
|
||||
async def get_managed_workers_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_managed_workers_endpoint(request):
|
||||
"""Get list of workers managed by this UI instance."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
managed = get_worker_manager().get_managed_workers()
|
||||
return web.json_response({
|
||||
@@ -585,16 +563,16 @@ async def get_managed_workers_endpoint(request: web.Request) -> web.StreamRespon
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
|
||||
async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_local_worker_status_endpoint(request):
|
||||
"""Check status of all local workers (localhost/no host specified)."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
config = load_config()
|
||||
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"]:
|
||||
@@ -658,11 +636,8 @@ async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamRe
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
|
||||
async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_worker_log_endpoint(request):
|
||||
"""Get log content for a specific worker."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
worker_id = request.match_info['worker_id']
|
||||
@@ -704,11 +679,8 @@ async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
|
||||
async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
|
||||
async def get_remote_worker_log_endpoint(request):
|
||||
"""Proxy a remote worker log request to the worker's local in-memory log endpoint."""
|
||||
auth_error = await authorization_error_or_none(request)
|
||||
if auth_error is not None:
|
||||
return auth_error
|
||||
try:
|
||||
worker_id = str(request.match_info["worker_id"]).strip()
|
||||
config = load_config()
|
||||
@@ -731,7 +703,6 @@ async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResp
|
||||
async with session.get(
|
||||
worker_url,
|
||||
params={"lines": str(lines_to_read)},
|
||||
headers=distributed_auth_headers(config),
|
||||
timeout=aiohttp.ClientTimeout(total=5),
|
||||
) as resp:
|
||||
if resp.status >= 400:
|
||||
|
||||
@@ -1,5 +0,0 @@
|
||||
"""Bootstrap entrypoints for ComfyUI-Distributed."""
|
||||
|
||||
from .entrypoint import build_node_mappings, initialize_runtime
|
||||
|
||||
__all__ = ["build_node_mappings", "initialize_runtime"]
|
||||
@@ -1,66 +0,0 @@
|
||||
"""Single startup owner for node surface and runtime initialization."""
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def build_node_mappings() -> tuple[dict[str, Any], dict[str, str]]:
|
||||
"""Build and cache the full exported node mapping surface."""
|
||||
from ..nodes import (
|
||||
NODE_CLASS_MAPPINGS as distributed_class_mappings,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as distributed_display_name_mappings,
|
||||
)
|
||||
from ..nodes.distributed_upscale import (
|
||||
NODE_CLASS_MAPPINGS as upscale_class_mappings,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as upscale_display_name_mappings,
|
||||
)
|
||||
|
||||
return (
|
||||
{**distributed_class_mappings, **upscale_class_mappings},
|
||||
{**distributed_display_name_mappings, **upscale_display_name_mappings},
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _runtime_initialized() -> dict[str, bool]:
|
||||
return {"done": False}
|
||||
|
||||
|
||||
def initialize_runtime(prompt_server: Any | None = None) -> None:
|
||||
"""Initialize distributed runtime side effects once."""
|
||||
state = _runtime_initialized()
|
||||
if state["done"]:
|
||||
return
|
||||
|
||||
import server
|
||||
|
||||
from ..api import bootstrap_routes
|
||||
from ..api.queue_orchestration import ensure_distributed_state
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
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
|
||||
|
||||
ensure_config_exists()
|
||||
bootstrap_routes()
|
||||
|
||||
server_instance = prompt_server if prompt_server is not None else server.PromptServer.instance
|
||||
ensure_distributed_state(server_instance)
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
|
||||
if not os.environ.get("COMFYUI_IS_WORKER"):
|
||||
atexit.register(sync_cleanup)
|
||||
delayed_auto_launch()
|
||||
register_async_signals()
|
||||
|
||||
state["done"] = True
|
||||
|
||||
node_mappings, _ = build_node_mappings()
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(node_mappings.keys())}")
|
||||
+25
-1
@@ -1,4 +1,28 @@
|
||||
"""Project-level pytest collection rules for plugin-style package layout."""
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
# All actual package context is provided by each test module via
|
||||
# importlib.util.spec_from_file_location with synthetic stub packages.
|
||||
|
||||
from _pytest.python import Package
|
||||
|
||||
_orig_pkg_setup = Package.setup
|
||||
|
||||
|
||||
def _patched_pkg_setup(self) -> None:
|
||||
# Skip the root package setup — its __init__.py uses relative imports
|
||||
# that require a parent package (ComfyUI's plugin loader) which is not
|
||||
# available in the test environment.
|
||||
if self.path == self.config.rootpath:
|
||||
return
|
||||
_orig_pkg_setup(self)
|
||||
|
||||
|
||||
Package.setup = _patched_pkg_setup
|
||||
|
||||
collect_ignore = [
|
||||
"__init__.py",
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
"""Compatibility facade for legacy imports from `distributed`."""
|
||||
from __future__ import annotations
|
||||
|
||||
from .bootstrap.entrypoint import (
|
||||
build_node_mappings,
|
||||
initialize_runtime,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
|
||||
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
"initialize_runtime",
|
||||
]
|
||||
@@ -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.
|
||||
|
||||
@@ -53,6 +53,7 @@ Queue a workflow for distributed execution.
|
||||
"delegate_master": false,
|
||||
"enabled_worker_ids": ["1", "2"],
|
||||
"workers": ["1", "2"],
|
||||
"auto_prepare": true,
|
||||
"trace_execution_id": "exec_1700000000_ab12cd"
|
||||
}
|
||||
```
|
||||
@@ -60,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.
|
||||
@@ -73,6 +75,10 @@ Queue a workflow for distributed execution.
|
||||
- The explicit worker IDs to consider for this run.
|
||||
- `workers` (optional, array of strings or objects with `id`)
|
||||
- Transitional alias for `enabled_worker_ids` used by older clients.
|
||||
- `auto_prepare` (optional, boolean)
|
||||
- Kept for wire compatibility.
|
||||
- Backend orchestration always runs with auto-prepare semantics.
|
||||
- If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
|
||||
- `trace_execution_id` (optional, string)
|
||||
- Passed through to orchestration logs.
|
||||
- Server log lines include the marker as `[exec:<trace_execution_id>]`.
|
||||
@@ -100,7 +106,8 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
```json
|
||||
{
|
||||
"prompt_id": "<uuid>",
|
||||
"worker_count": 2
|
||||
"worker_count": 2,
|
||||
"auto_prepare_supported": true
|
||||
}
|
||||
```
|
||||
|
||||
@@ -117,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:
|
||||
|
||||
@@ -130,13 +141,16 @@ Also, for collector-based flows:
|
||||
|
||||
## Endpoint: `POST /distributed/job_complete`
|
||||
|
||||
Submit one completed worker image back to the master collector queue.
|
||||
Submit direct-worker image/audio results to the master Collector queue. The route
|
||||
accepts canonical JSON completions or multipart image chunks, including opt-in
|
||||
exact floating-point IMAGE transport. Serverless typed results use their separate
|
||||
bridge/provider contract.
|
||||
|
||||
### URL
|
||||
|
||||
- `http://<master-host>:<master-port>/distributed/job_complete`
|
||||
|
||||
### Request Body
|
||||
### JSON request body
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -153,13 +167,35 @@ Submit one completed worker image back to the master collector queue.
|
||||
- `job_id` (string, required)
|
||||
- `worker_id` (string, required)
|
||||
- `batch_idx` (integer >= 0, required)
|
||||
- `image` (string, required)
|
||||
- `image` (string, optional if `audio` is present)
|
||||
- PNG payload as either:
|
||||
- data URL: `data:image/png;base64,...`
|
||||
- raw base64 PNG bytes
|
||||
- `is_last` (boolean, required)
|
||||
|
||||
Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
|
||||
At least one of `image` or `audio` must be present. Optional `audio` is an object
|
||||
with `sample_rate` (positive integer), `shape` (`[batch, channels, samples]`),
|
||||
`dtype: "float32"`, and `data` (base64 waveform bytes). Audio-only completions
|
||||
still include `batch_idx`; the Collector uses `0` and `is_last: true`.
|
||||
|
||||
### Multipart image chunks
|
||||
|
||||
Send `multipart/form-data` with a JSON `envelope` field and repeated binary
|
||||
`media` file fields. The envelope contains `job_id`, `worker_id`, `format`,
|
||||
`start_index` (nonnegative integer), `count`, `width`, `height`, `is_last`, and
|
||||
optional canonical `audio`.
|
||||
|
||||
- `png`, `jpeg`, `webp`: one `media` part per frame.
|
||||
- `ffv1`, `h264`, `h265` (experimental), `av1`: one video part per chunk.
|
||||
- `exact`: one raw tensor part per frame and an envelope `descriptors` list with
|
||||
one `tensor-le-v1` IMAGE descriptor per part. Each descriptor shape must be
|
||||
`[1, height, width, channels]`; decoding preserves supported float dtype and
|
||||
values instead of quantizing to 8-bit.
|
||||
|
||||
`legacy_png` uses JSON, not multipart. Exact transport requires support on both
|
||||
sides and never retries as an 8-bit format when an older master rejects it.
|
||||
See [result transfer](compressed-result-transfer.md) for format controls,
|
||||
non-exact compatibility fallbacks and request limits.
|
||||
|
||||
### CORS note
|
||||
|
||||
@@ -195,7 +231,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" \
|
||||
@@ -203,18 +239,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
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
# Compressed and Exact Result Transfer
|
||||
|
||||
**Distributed Collector** (`DistributedCollector`) collects images and/or audio. Its `transfer_format` and `transfer_quality` widgets control **remote image returns only**; they do not re-encode master-local images or change audio transport. The saved defaults remain `legacy_png` and quality `90`.
|
||||
|
||||
| `transfer_format` | Fidelity | Direct-worker transport |
|
||||
|---|---|---|
|
||||
| `legacy_png` (default) | 8-bit pixels | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Compatible with legacy masters. |
|
||||
| `png` | Lossless 8-bit pixels | Compressed PNG, binary, several frames per request. |
|
||||
| `ffv1` | Lossless 8-bit pixels | FFV1 video chunks; useful for frame sequences. |
|
||||
| `jpeg`, `webp` | Lossy 8-bit images | Per-frame stills. |
|
||||
| `h264`, `h265` (**experimental**), `av1` | Lossy 8-bit images | Video chunks; useful for similar consecutive frames. |
|
||||
| `exact` (opt-in) | Original floating-point tensors | Binary multipart chunks with one raw tensor payload and `tensor-le-v1` descriptor per frame. No 8-bit quantization. |
|
||||
|
||||
“Lossless” PNG/FFV1 means lossless **after 8-bit quantization**, not preservation of the source float tensor. Direct non-exact returns decode to RGB float32. `exact` preserves supported floating-point dtype, shape, channels and values: float16, bfloat16, float32 or float64, finite dense BHWC tensors with 1, 3 or 4 channels. It generally costs substantially more bandwidth.
|
||||
|
||||
`transfer_quality` (1–100, default 90) applies to lossy formats, not `exact`, PNG or FFV1. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0). A quality of 100 is not a substitute for exact float transport.
|
||||
|
||||
## Serverless profiles
|
||||
|
||||
- **`comfy-media-v1`:** the same Collector supplies the endpoint's image-return policy. The saved `legacy_png` default is mapped to compressed **`png`** with the same 8-bit pixel precision; normal serverless returns are not exact by default. All other formats retain their selection, including opt-in `exact` and experimental `h265`. Missing codecs or rejected formats fail instead of selecting another codec. Typed image returns preserve channel count and carry alpha separately where needed. See [serverless images and audio](serverless-media.md).
|
||||
- **`legacy-image-v1`:** retains its image-file/`SaveImage` contract, not the typed return-format controls. Selecting a legacy pool when any Collector requests `exact` is rejected before dispatch; there is no 8-bit fallback. Audio requires direct workers or `comfy-media-v1`.
|
||||
|
||||
## Direct-worker compatibility and limits
|
||||
|
||||
- Choose the format on the master; the value is sent to workers with the workflow.
|
||||
- **`exact` requires support on both master and worker.** If an older master rejects an exact chunk, the worker raises an error and never resends it as `legacy_png` or another 8-bit format. An exact Collector also rejects direct-worker image returns labelled with a non-exact format. Update both sides or explicitly select a non-exact format.
|
||||
- For **non-exact compressed** formats only, the specific first-request older-master rejection (`400` with `Invalid JSON payload`) causes a retry using `legacy_png`. Other rejections, including HTTP 413, remain errors.
|
||||
- For direct video returns only, an unavailable encoder falls back to `h264` for lossy video if available, otherwise `png`; unavailable `ffv1` falls back to `png`. The worker logs this. This compatibility fallback does **not** apply to `exact` or typed serverless returns.
|
||||
- Multipart requests use a 32 MiB chunk target (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`) and a 96 MiB request budget (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`), reserving space for the envelope and accounting separately for audio. A single large frame may exceed the chunk target but must fit the request budget. Keep the request budget below the master's `--max-upload-size` (default 100 MB).
|
||||
- Non-exact compressed chunks are all encoded before sending. Exact transfer instead plans frame sizes from shape/dtype before sending and encodes one chunk at a time. An oversized individual frame is rejected before any chunks are sent; choose another format or raise the worker and master limits together.
|
||||
- Audio is attached to the final image chunk when it fits, otherwise sent separately. Audio-only completions use the canonical JSON envelope. Direct audio uses float32 waveform bytes regardless of the image format; `exact` is an IMAGE option, not a new audio codec.
|
||||
- Each worker logs the format and bytes sent. These request limits are for direct multipart transfers; serverless request/artifact limits are documented separately.
|
||||
|
||||
## Previously documented size example
|
||||
|
||||
The following measurements are retained from a 121-frame 1280×720 slow-pan batch (one worker's result); they are not new validation of the current implementation or exact transport.
|
||||
|
||||
| Format | Sent | vs legacy | PSNR |
|
||||
|---|---|---|---|
|
||||
| `legacy_png` | 446 MB | 1× | lossless at 8-bit precision |
|
||||
| `png` | 163 MB | 2.7× | lossless at 8-bit precision |
|
||||
| `ffv1` | 128 MB | 3.5× | lossless at 8-bit precision |
|
||||
| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
|
||||
| `webp` q90 | 13 MB | 33× | 40.6 dB |
|
||||
| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
|
||||
| `h265` q90 (experimental) | 1.0 MB | 434× | 40.1 dB |
|
||||
| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
|
||||
|
||||
Video codecs compress this slow pan especially well; gains depend on content and motion.
|
||||
@@ -0,0 +1,137 @@
|
||||
# Dynamic image-list queue
|
||||
|
||||
Process a list of images across the master and direct local or remote ComfyUI workers without assigning fixed chunks to each participant.
|
||||
|
||||
The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
|
||||
|
||||
The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
|
||||
|
||||
## 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. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation.
|
||||
2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. 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; for each assigned item, 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; each frame becomes one job.
|
||||
- Optional `mask`.
|
||||
- Optional `value_1`, `value_2`, … inputs. Only `value_1` is shown at first; connecting it adds `value_2`, and so on, up to `value_10`. Each accepts an IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN.
|
||||
|
||||
`mask` and every value follow one rule:
|
||||
|
||||
- **1 value** is shared by every image.
|
||||
- **One value per image** pairs them in order.
|
||||
- **Any other count** is an error. Values are never padded or repeated silently.
|
||||
|
||||
Only `images` is split into frames. A batch connected to a value input stays a single value, so a 2-image batch can travel with each item (for example, two reference images per job). To pass a different batch to each image, connect a list with one batch per image.
|
||||
|
||||
Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`, then one `value_N` output per connected `value_N` input, with the same type. Unused value outputs are hidden. Seeds and strengths are ordinary values: connect a single number to share it, a list for one per image, or derive it from `item_index`.
|
||||
|
||||
Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
|
||||
|
||||
End accepts `images`, optional `masks` / `texts`, and `context`. Its first three outputs remain `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 mask/text outputs contain `None` entries.
|
||||
|
||||
### Keep values through End without extra wires
|
||||
|
||||
Start's existing **context** connection now also lets End recover each item's original `value_1` … `value_10`. You do not need to wire those values individually into End. Context remains a lightweight execution reference, not a general-purpose data pipe.
|
||||
|
||||
End appends a **values** output (`DISTRIBUTED_IMAGE_LIST_VALUES`): one value bundle per image, in the same source order. Connect it to **Distributed Image List Get Value**, then choose `value_name` to extract the field you need. The getter preserves its type and ComfyUI maps it across the bundles, producing a list aligned with End's images. Missing fields fail clearly instead of producing empty placeholders.
|
||||
|
||||
For example, to preserve filenames:
|
||||
|
||||
```text
|
||||
filenames → Start.value_1
|
||||
Start.context ─────────────→ End.context
|
||||
Start.image → processing → End.images
|
||||
├─ image_list → Save.images
|
||||
└─ values → Get Value [value_1] → Save.filename_prefix
|
||||
```
|
||||
|
||||
**New or changed per-item values:** connect the calculated result to End's optional `value_N` input. It replaces the original field with that number, or adds it if Start did not supply it. Every connected override must produce exactly one supported value for each item. A tensor batch remains one value, just as it does at Start. Image, mask, string, integer, float and boolean values are supported; arbitrary objects/models/latents and non-finite numbers are not.
|
||||
|
||||
End shows override inputs progressively. Use **Add value override** to expose a later slot without connecting the earlier ones. Disconnecting an override never renumbers the other End fields: `value_3` always means `value_3`. Unlike Start's native Autogrow inputs, End overrides must keep these identities stable. Existing image/mask/text/context sockets and output indices are unchanged. If the frontend extension is unavailable, Python still declares all optional override inputs and the values output.
|
||||
|
||||
Original values are retained on the master; workers only return explicitly connected overrides, avoiding a redundant return trip for unchanged reference images or other values. Overrides use the existing exact value transport and size limits. Selected workers must advertise `value_results: true` when overrides are requested, checked before execution and again at dispatch. Original-value passthrough alone does not require that additional capability.
|
||||
|
||||
Values leaving the region must come from **End**, not directly from Start. To chain regions, extract an End value and connect it to the next Start's corresponding input. Values remain available downstream after End cleans up the execution context; they do not depend on a later lookup of expired job state.
|
||||
|
||||
Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. 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 socket became connection-only, 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 and are 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.
|
||||
|
||||
## Native cache reuse
|
||||
|
||||
Repeated items can reuse **ComfyUI's native processing cache** on the same participant. Start uses a deterministic identity of the exact image, optional mask, all supplied value slots and source index; tensor dtype, shape and bytes are preserved. Fresh job IDs, assignment tokens and callback addresses live at End, so they no longer invalidate otherwise unchanged upstream processing. End still validates the current assignment and delivers results in source order on every run.
|
||||
|
||||
- Processing settings, seeds, links and node fingerprints still participate in ComfyUI's normal signatures. Changed content invalidates reuse; nodes that declare themselves noncacheable still run.
|
||||
- The master's private executor retains one native LRU output cache between successful slots, capped at **128 entries when retained**. An active graph may temporarily exceed that target; oversized generations are not retained. Failed or cancelled local slots do not seed reuse, and `--cache-none` disables the local output cache. Direct workers use their own configured native caches.
|
||||
- Cache hits are best-effort, not guaranteed: eviction, restarts, changed ordering/index, changed settings and assignment to a different worker can cause misses. There is no worker affinity, disk cache or separate completed-result store.
|
||||
- Reuse retains tensor/model references and can increase RAM/VRAM use; the entry cap is **not a byte limit**. Exact identity calculation also adds CPU hashing and, for GPU-resident inputs, a CPU copy.
|
||||
- Remote items are still uploaded. Each upload returns a worker-computed identity; if a compatible older worker does not return a valid identity, that assignment uses the legacy graph without stable-input reuse. A worker downgrade between items also falls back safely. Assignment expiry and current-item validation apply even when processing hits the cache.
|
||||
- Keep side-effect nodes outside the region when they must run on every submission. A cached processing node follows ordinary ComfyUI semantics and does not repeat its save/preview side effects; End's delivery and downstream saving still run.
|
||||
|
||||
## Scope and failure behavior
|
||||
|
||||
- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
|
||||
- Queue submission checks selected direct workers for reachability and feature compatibility, but does not freeze the worker list to that moment's idle snapshot. A selected worker busy with an earlier prompt remains a candidate for the queued run. Immediately before a region starts assigning images, its saved candidates are checked again: workers still busy or transiently unreachable are excluded; healthy idle candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers excluded by the initial reachability probe are not retained, and newly selected workers are not added to an already-queued run. Workers becoming free after the region starts are not added mid-region.
|
||||
- **Orchestrator Only** requires a selected, reachable, feature-compatible direct worker at submission, but can queue while all selected workers are busy. At region start it requires at least one idle compatible direct worker and fails if none are available; it never falls back to master processing. The queue response's `worker_count` counts saved remote candidates, not guaranteed execution-time participants.
|
||||
- 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 in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
|
||||
- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
|
||||
- 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.
|
||||
- Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed.
|
||||
- Every selected worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation.
|
||||
- **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End.
|
||||
- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits.
|
||||
- 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; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. 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 cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/mask/value 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 workers using real ComfyUI HTTP handlers, `PromptQueue` records and CPU `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once required output utilities, discarded previews, unchanged source graphs, End-only worker roots, post-End capture, source-once materialization and cached local model-loader dependencies. WebSocket target propagation is covered by focused route/dispatch tests. 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.
|
||||
@@ -0,0 +1,187 @@
|
||||
# Serverless images and audio
|
||||
|
||||
**Distributed Collector** (`DistributedCollector`) collects `IMAGE` and/or `AUDIO`
|
||||
from direct workers and serverless pools. There is no separate
|
||||
`DistributedMediaCollector`; replace that removed node with `DistributedCollector`
|
||||
for image/audio workflows. The Collector has no MASK, STRING, or native VIDEO
|
||||
output sockets.
|
||||
|
||||
`comfy-media-v1` is an **explicit endpoint profile**, not an automatically detected
|
||||
capability. It returns typed image/audio items to the Collector. The standalone
|
||||
bridge's generic contract also supports `MASK`, `STRING`, and native ComfyUI
|
||||
`VIDEO`, but that does not add those outputs to the Collector. Omitting
|
||||
`media_profile` still selects the image-only `legacy-image-v1` contract.
|
||||
See [legacy serverless setup](serverless-workers.md) for credentials and recovery.
|
||||
|
||||
## Use the Collector with a typed endpoint
|
||||
|
||||
1. Prepare an endpoint with the [standalone media bridge](../serverless_bridge/README.md)
|
||||
and the workflow's models/custom nodes. The endpoint needs only the bridge,
|
||||
**not the full ComfyUI-Distributed extension**.
|
||||
2. In the serverless worker settings, set **Endpoint media profile** to
|
||||
**Typed media v1 (requires endpoint bridge)** (`media_profile: "comfy-media-v1"`).
|
||||
Start with **Images per run = 1**. This existing setting counts remote jobs,
|
||||
not the number of typed items each job returns.
|
||||
3. Use **Distributed Collector** (`DistributedCollector`), connect `images`,
|
||||
`audio`, or both, and connect the corresponding outputs downstream. Image-only
|
||||
and audio-only workflows use the same node.
|
||||
4. Queue through Distributed. **Delegate Only** excludes local upstream
|
||||
generation; normal collection includes the master's local values.
|
||||
|
||||
**Test Saved Endpoint** checks provider authentication/health, not bridge
|
||||
installation, model availability, or typed execution. Selecting a profile does
|
||||
not install or upgrade anything on the endpoint.
|
||||
|
||||
The Collector accepts optional `images` and `audio` inputs, including ComfyUI list
|
||||
inputs, and returns two ordinary outputs—not list outputs:
|
||||
|
||||
| Socket / output index | Native type | Aggregation |
|
||||
| --- | --- | --- |
|
||||
| `images` / 0 | `IMAGE` | Concatenate BHWC tensors along the batch dimension |
|
||||
| `audio` / 1 | `AUDIO` | Concatenate waveforms along the sample dimension |
|
||||
|
||||
For each connected socket, order is **master first, then prepared participant
|
||||
order, then that participant's item order**—never completion order. Delegate-only
|
||||
mode omits the master. Image dimensions/channels and audio batch/channel dimensions
|
||||
must be compatible for concatenation. Use a common audio sample rate; the Collector
|
||||
does not resample or align streams. Audio-only collection returns `None` for images;
|
||||
when no audio is supplied it returns a silent placeholder, not an empty list.
|
||||
|
||||
The typed wire contract carries 1–64 items per connected socket. The entire remote
|
||||
response is validated and decoded before any of its results are delivered;
|
||||
missing, duplicate, mismatched, or corrupt typed results fail collection. Validated
|
||||
wire lists are then combined into the Collector's ordinary outputs.
|
||||
|
||||
Direct workers, `legacy-image-v1` pools and `comfy-media-v1` pools can share an
|
||||
image-only Collector run. Audio requires direct workers or `comfy-media-v1` pools.
|
||||
If any Collector requests `exact`, a selected `legacy-image-v1` pool is rejected
|
||||
before dispatch instead of silently returning 8-bit images.
|
||||
|
||||
## Image return format and quality
|
||||
|
||||
The Collector's saved `transfer_format` default remains **`legacy_png`**, with
|
||||
`transfer_quality` **90**. For `comfy-media-v1`, the compiler maps `legacy_png` to
|
||||
**`png`**, so normal serverless returns use compressed 8-bit PNG. `exact` is opt-in.
|
||||
These controls affect **remote IMAGE returns only**, not master-local images,
|
||||
uploaded inputs or AUDIO. The bridge's generic MASK, STRING and native VIDEO
|
||||
codecs are also unaffected.
|
||||
|
||||
- `exact`: preserve supported floating-point dtype and values, with no image
|
||||
quantization. Available for direct workers and `comfy-media-v1` pools; it is
|
||||
never downgraded to an 8-bit format. An older master that rejects exact transfer
|
||||
causes an error, not a legacy retry. Update both sides before selecting it.
|
||||
- `legacy_png`, `png`, `ffv1`: lossless **8-bit image pixels**, not exact float
|
||||
tensors. For typed serverless returns, image values are clipped to [0,1] and
|
||||
quantized; `legacy_png` is sent as `png`. Direct workers retain the original
|
||||
JSON/base64 framing for `legacy_png`.
|
||||
- `jpeg`, `webp`: lossy still-image compression.
|
||||
- `h264`, `h265` (**experimental**), `av1`: lossy video encoding of each IMAGE
|
||||
tensor's frame batch. This is an image-return encoding, not a native VIDEO output.
|
||||
- Quality 1–100 affects lossy formats only, using the existing Collector's
|
||||
quality-to-CRF mapping for video. Dimensions, frame order and channel count
|
||||
are preserved for typed serverless returns; alpha is carried separately where
|
||||
needed, not silently dropped.
|
||||
|
||||
Both master and endpoint must have this bridge revision for compressed image
|
||||
returns. Pillow handles still formats and CPU/software PyAV encoders handle video;
|
||||
unavailable codecs fail clearly rather than silently changing the selection.
|
||||
Compression can reduce result egress and inline response size, but does not relax
|
||||
encoded/decoded byte budgets or the input request limit. Compressed IMAGE batches
|
||||
are capped at 240 frames and 8192 pixels per dimension, with a conservative
|
||||
256 MiB aggregate decode budget (including codec scratch and alpha). Direct
|
||||
workers have different limits and non-exact compatibility fallbacks; see
|
||||
[result transfer](compressed-result-transfer.md). The bridge's internal
|
||||
`MediaExportV1` default is still `exact` for generic callers that omit controls;
|
||||
the Collector compiler explicitly supplies `png` for the normal saved default.
|
||||
|
||||
## Inputs and fidelity
|
||||
|
||||
Automatic local-file transfer supports stock `LoadImage`, plus explicit stock
|
||||
`LoadAudio` (`audio` filename) and `LoadVideo` (`file` filename) adapters. Audio and
|
||||
video loaders require a literal filename, no extra controls, and output slot 0.
|
||||
Files must be under ComfyUI's input/output/temp roots. Custom file loaders are not
|
||||
supported; detected unsupported file references fail rather than silently
|
||||
becoming endpoint-local paths. Models and custom nodes are never transferred.
|
||||
|
||||
- `LoadImage` uses the existing bounded file/base64 image staging path.
|
||||
- `LoadAudio` is decoded on the master using native PCM semantics with bounded
|
||||
accumulation, then transferred as a waveform and sample rate. It does not
|
||||
resample; integer PCM conversion follows ComfyUI's `f32_pcm` behavior.
|
||||
- `LoadVideo` transfers opaque file bytes as native `VideoFromFile`, without
|
||||
decoding frames or transcoding the container.
|
||||
- Typed tensor transport is `tensor-le-v1`: contiguous little-endian float16,
|
||||
bfloat16, float32, or float64 bytes with exact dtype/shape/layout metadata.
|
||||
Only dense, finite tensors with positive dimensions are supported. IMAGE has
|
||||
1, 3, or 4 channels. Bridge export explicitly detaches/stages tensors to CPU;
|
||||
the pure codec itself rejects non-CPU storage. Remote decoding returns
|
||||
independently owned CPU values. STRING transport is exact UTF-8 (`utf8-v1`).
|
||||
|
||||
### Bridge-level native VIDEO support
|
||||
|
||||
The generic bridge can import/export native VIDEO, and the typed input adapter
|
||||
can stage `LoadVideo` for upstream processing. This is not a Collector VIDEO
|
||||
output or a generic video converter.
|
||||
|
||||
The adapter supports the pinned native classes **exactly**, not arbitrary
|
||||
subclasses or objects exposing `save_to()` / `get_components()`:
|
||||
|
||||
- `VideoFromComponents`: images, optional matching BHW or BHW1 alpha, optional
|
||||
audio, positive rational frame rate, bounded JSON metadata, positive integer
|
||||
bit depth, and `sRGB`, `HDR`, or `HDR PQ` color space are preserved.
|
||||
- `VideoFromFile`: regular local files within allowed media roots or an exact
|
||||
`BytesIO`; opaque bytes, finite signed start/duration, and optional four-integer
|
||||
crop are preserved without normalization.
|
||||
- `VideoFromList`: 1–64 flat exact `VideoFromFile` segments, plus optional
|
||||
`complete_audio`, in order.
|
||||
|
||||
Unknown classes, incompatible native state, nested/component segment lists,
|
||||
unsupported tensors, invalid metadata, and oversized payloads are rejected.
|
||||
There is no fallback transcode, frame extraction, color conversion, or container
|
||||
repair/validation. A downstream native consumer still needs to read the original
|
||||
container. Decoded file-backed VIDEO values own independent `BytesIO` storage;
|
||||
component tensors also own their storage, so staging-file cleanup cannot
|
||||
invalidate a returned VIDEO.
|
||||
|
||||
## Size and operational limits
|
||||
|
||||
- The typed contract caps each artifact at **256 MiB**, aggregate artifact bytes
|
||||
at **512 MiB**, and each artifact list at **128** entries. These are validation
|
||||
ceilings, not a promise that inline requests can carry that much data.
|
||||
- The built-in master packs inputs inline and enforces **10,000,000 bytes for the
|
||||
complete serialized Runpod `/run` request**, including workflow, policy and
|
||||
base64. Legacy image reads also retain their 10,000,000-byte aggregate guard.
|
||||
There is **no built-in master-side large-input uploader** in PR2.
|
||||
- The endpoint defaults to inline results and caps its result JSON at
|
||||
**10,000,000 bytes**. Larger results require an operator-supplied publisher
|
||||
returning expiring public HTTPS artifact URLs. See the bridge guide for
|
||||
bounded download/upload responsibilities and retention.
|
||||
- The generic five-kind typed envelope is a serverless bridge contract. Direct
|
||||
workers support Collector images/audio, including opt-in exact IMAGE transport,
|
||||
not generic MASK/STRING/VIDEO Collector outputs. The typed compiler rejects
|
||||
unsupported distributed/batch nodes and collector-to-collector dependencies.
|
||||
Downstream processing stays on the master. [Image-list scheduling](image-list-queue.md)
|
||||
is a separate execution mode.
|
||||
|
||||
## Timeout is not proof that Comfy stopped
|
||||
|
||||
Provider cancellation, a handler timeout, or cancellation of its Python task
|
||||
**does not establish that the local Comfy prompt stopped**. The bridge never
|
||||
posts `/interrupt`: the pinned core route checks a queue snapshot and then
|
||||
interrupts globally, so even a prompt ID does not make it race-free.
|
||||
|
||||
If submission/termination or callback teardown cannot be confirmed, the bridge
|
||||
raises `MediaExecutionUncertain` and retains owned input/exporter files—even
|
||||
after terminal Comfy history if a downloader/publisher ignores cancellation.
|
||||
Callback teardown has a separate 10-second grace; late callback exit does not
|
||||
clean retained files. Callbacks must be cooperative async code; synchronous
|
||||
event-loop blocking cannot be hard-timed-out.
|
||||
|
||||
The endpoint operator must **quarantine the owning worker/environment, reconcile
|
||||
the local prompt and readers or terminate that environment, and establish that
|
||||
they have stopped before cleanup or worker reuse**. Do not automatically retry
|
||||
the request or turn every handler exception into a safely reusable terminal
|
||||
provider job. Master-side provider status/reconciliation alone is not proof
|
||||
that a surviving local Comfy process or callback has stopped.
|
||||
|
||||
These docs describe the implementation, not a test result or live Runpod/provider
|
||||
or GPU end-to-end validation.
|
||||
@@ -0,0 +1,184 @@
|
||||
# Runpod serverless workers (preview)
|
||||
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. This page describes the
|
||||
default **`legacy-image-v1`** profile: Distributed submits normal ComfyUI API
|
||||
workflows and collects returned image files. For image/audio collection and
|
||||
opt-in exact IMAGE returns, use [the `comfy-media-v1` profile](serverless-media.md)
|
||||
with the same **Distributed Collector** (`DistributedCollector`). The endpoint
|
||||
workers do not need the full ComfyUI-Distributed extension installed.
|
||||
|
||||
## Setup
|
||||
|
||||
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
|
||||
2. Enter its name, Runpod **endpoint ID** (not a URL; a pasted `https://api.runpod.ai/v2/<id>/…`
|
||||
URL is trimmed to the ID), and API key.
|
||||
3. Start with **Images per run = 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 worker when you want queued Distributed workflows to use it.
|
||||
6. Connect the **Distributed Collector**'s `images` input and queue through
|
||||
Distributed as usual. Leave `audio` unconnected for this legacy profile.
|
||||
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. The checkbox is unavailable while the
|
||||
worker is enabled, so turn the worker off first. A saved key is required before
|
||||
enabling a worker. 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
|
||||
|
||||
- **Images per run** (1–16) is how many separate Runpod jobs this worker submits
|
||||
for each queued run. Each returns one image and receives its own Distributed
|
||||
Seed/Value overrides. All of them are submitted at once.
|
||||
- How many of those jobs actually run in parallel is decided by the Runpod
|
||||
endpoint's **max workers**, not by Distributed. Jobs beyond that wait in
|
||||
Runpod's queue. To run every image at the same time, keep Images per run at or
|
||||
below the endpoint's max workers.
|
||||
- Distributed also keeps a fixed limit of 16 outstanding jobs per serverless
|
||||
worker (and 64 across all of them). Retained jobs after a restart, ambiguous
|
||||
submissions, and unconfirmed cancellations keep counting against it until
|
||||
reconciled. Other applications using the same endpoint are unaffected.
|
||||
- Older configs may contain a `concurrency` field from the previous separate
|
||||
setting. It is still accepted when loading or saving, but ignored and dropped.
|
||||
- **Load balance** selects one participant, so it sends one job instead of
|
||||
Images per run.
|
||||
- **Timeout** includes 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 worker'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 worker's endpoint, API key, worker type,
|
||||
and deletion even when no local task remains. Disabling the worker 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_WORKER_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_WORKER_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.
|
||||
|
||||
Worker setting changes become available once all its local tasks and unresolved records
|
||||
are finished. No credentials or ledger files need to be edited by hand.
|
||||
|
||||
## Legacy profile limits
|
||||
|
||||
- Only the Collector's image branch is supported by `legacy-image-v1`; connected
|
||||
audio is rejected. `comfy-media-v1` adds audio through the same Collector and
|
||||
supports the configured image-return formats.
|
||||
- `exact` is opt-in and requires direct workers or `comfy-media-v1` pools. If any
|
||||
Collector requests it, a selected legacy pool is rejected before dispatch,
|
||||
not used as an 8-bit fallback. Other Collector transfer-format controls do not
|
||||
change this legacy profile's `SaveImage`/image-file contract. With
|
||||
`comfy-media-v1`, the saved `legacy_png` default is mapped to compressed `png`.
|
||||
- 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.
|
||||
|
||||
These docs describe the implementation, not a test result. A successful health
|
||||
check is not a paid end-to-end generation test; validate your deployed handler
|
||||
and model inventory with one image per run before increasing it.
|
||||
@@ -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.
|
||||
@@ -74,12 +74,14 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
|
||||
> ComfyUI instances running on a cloud service like Runpod.
|
||||
|
||||
For faster result transfers and lower bandwidth usage, see [Compressed Result Transfer](compressed-result-transfer.md).
|
||||
|
||||
### Deploy Cloud Worker on Runpod
|
||||
|
||||
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||
|
||||
**On Runpod:**
|
||||
> If using your own template, 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 +94,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
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Dynamic per-image image-list execution internals."""
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Prepare image-list execution without changing the saved/API prompt or registry."""
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from .runtime import (
|
||||
END_CLASS,
|
||||
END_RESULT_INPUTS,
|
||||
_ancestors_for_inputs,
|
||||
_downstream_until,
|
||||
_image_list_pairs,
|
||||
_is_output_node,
|
||||
_prompt_links,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
def _reachable(roots: set[str], edges: dict[str, set[str]]) -> set[str]:
|
||||
visited = set()
|
||||
pending = list(roots)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
pending.extend(edges[node_id] - visited)
|
||||
return visited
|
||||
|
||||
|
||||
def prepare_image_list_execution(prompt: dict) -> tuple[dict, list[str]]:
|
||||
"""Return an execution copy and its outer-master output targets.
|
||||
|
||||
Disposable per-item side branches are removed before boundary validation.
|
||||
Required output-node utilities remain normal dependencies, but must not be
|
||||
independent outer roots. Callers must use these targets (and End-only targets
|
||||
for item execution); this does not suppress side effects inside required nodes.
|
||||
"""
|
||||
pairs = _image_list_pairs(prompt)
|
||||
prepared = copy.deepcopy(prompt)
|
||||
upstream = {node_id: set() for node_id in prepared}
|
||||
downstream = {node_id: set() for node_id in prepared}
|
||||
for node_id, node in prepared.items():
|
||||
for _, source, _ in _prompt_links(node, prepared):
|
||||
upstream[node_id].add(source)
|
||||
downstream[source].add(node_id)
|
||||
|
||||
# Never erase another region or its dependencies to make a bad graph pass.
|
||||
boundaries = {node_id for pair in pairs for node_id in pair}
|
||||
required = _reachable(boundaries, upstream)
|
||||
discarded = set()
|
||||
for start_id, end_id in pairs:
|
||||
region_nodes = _ancestors_for_inputs(prepared, end_id, start_id, END_RESULT_INPUTS)
|
||||
outside = _downstream_until(prepared, start_id, end_id) - region_nodes
|
||||
# Preserve bypass joins and their whole chain into post-End computation.
|
||||
# They must fail boundary validation, not disappear with debug previews.
|
||||
post_end = _reachable({end_id}, downstream)
|
||||
protected = required | _reachable(post_end, upstream)
|
||||
discarded.update(outside - protected)
|
||||
for node_id in discarded:
|
||||
del prepared[node_id]
|
||||
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
per_item = set().union(*(region["branch_ids"] for region in regions))
|
||||
registry = getattr(sys.modules.get("nodes"), "NODE_CLASS_MAPPINGS", {})
|
||||
targets = [
|
||||
node_id for node_id, node in prepared.items()
|
||||
if node_id not in per_item and isinstance(node, dict)
|
||||
and (node.get("class_type") == END_CLASS or _is_output_node(node, registry))
|
||||
]
|
||||
return prepared, targets
|
||||
@@ -0,0 +1,174 @@
|
||||
"""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 has_safe_execution_targets(prompt_server: Any, prompt_id: str | None, prompt: dict) -> bool:
|
||||
"""Prove outer roots from the running queue, not a forgeable prompt marker.
|
||||
|
||||
Plain /prompt remains strict when ComfyUI selected all outputs, including a
|
||||
flagged per-item dependency. No saved workflow field grants this exemption.
|
||||
"""
|
||||
if not prompt_id:
|
||||
return False
|
||||
from .execution_plan import prepare_image_list_execution
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
if prepared != prompt or not targets:
|
||||
return False
|
||||
queue = getattr(prompt_server, "prompt_queue", None)
|
||||
for name in ("get_current_queue_volatile", "get_current_queue"):
|
||||
reader = getattr(queue, name, None)
|
||||
if not callable(reader):
|
||||
continue
|
||||
try:
|
||||
running, _pending = reader()
|
||||
for entry in running:
|
||||
if (isinstance(entry, (tuple, list)) and len(entry) > 4
|
||||
and entry[1] == prompt_id and entry[2] == prompt):
|
||||
roots = entry[4]
|
||||
return (isinstance(roots, (list, tuple))
|
||||
and all(isinstance(root, str) for root in roots)
|
||||
and len(roots) == len(targets) and set(roots) == set(targets))
|
||||
return False
|
||||
except Exception:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,223 @@
|
||||
"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.
|
||||
|
||||
Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
|
||||
process interrupt, resets progress handlers, changes RAM-cache callbacks and
|
||||
starts/ends a model tracker. Those belong to the suspended parent executor.
|
||||
Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
|
||||
with private provider-free caches and a private server facade. Model outputs are
|
||||
registered with the existing parent tracker, whose lifetime includes this work.
|
||||
|
||||
Core progress and GraphBuilder defaults are process-global. Borrow them only
|
||||
inside the owned execution thread, restore them in finally, and join that thread
|
||||
before End returns. Cancellation never clears or raises the global interrupt:
|
||||
it stops locally at node/progress boundaries. A non-cooperative synchronous node
|
||||
must finish before cleanup can return; abandoning its thread is not safe.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
from .runtime import ImageListCancelledError
|
||||
|
||||
|
||||
_EXECUTION_LOCK = threading.Lock()
|
||||
_CACHE_LOCK = threading.Lock()
|
||||
# One process-local native output cache, not a second completed-result cache.
|
||||
# Core's LRU can temporarily exceed its target for the active graph; oversized
|
||||
# generations are discarded rather than retained after an item/run finishes.
|
||||
NATIVE_CACHE_MAX_ENTRIES = 128
|
||||
_RETAINED_OUTPUTS = None
|
||||
|
||||
|
||||
def _take_outputs(execution):
|
||||
global _RETAINED_OUTPUTS
|
||||
with _CACHE_LOCK:
|
||||
outputs, _RETAINED_OUTPUTS = _RETAINED_OUTPUTS, None
|
||||
if getattr(execution.args, "cache_none", False):
|
||||
return execution.NullCache()
|
||||
return outputs if outputs is not None else execution.LRUCache(
|
||||
execution.CacheKeySetInputSignature, max_size=NATIVE_CACHE_MAX_ENTRIES)
|
||||
|
||||
|
||||
def _retain_outputs(outputs):
|
||||
global _RETAINED_OUTPUTS
|
||||
# Only the bounded native LRU is eligible. NullCache and test doubles are not.
|
||||
if getattr(outputs, "max_size", None) != NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
outputs.clean_unused()
|
||||
if len(outputs.cache) > NATIVE_CACHE_MAX_ENTRIES:
|
||||
return
|
||||
with _CACHE_LOCK:
|
||||
_RETAINED_OUTPUTS = outputs
|
||||
|
||||
|
||||
class LocalExecutionSlot:
|
||||
def __init__(self, *, model_callback=None):
|
||||
self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
|
||||
self._stop = threading.Event()
|
||||
self._future = None
|
||||
self._caches = None
|
||||
self._failed = False
|
||||
self._model_callback = model_callback
|
||||
self._closed = False
|
||||
self._loop = None
|
||||
self._task = None
|
||||
self._state_lock = threading.Lock()
|
||||
|
||||
def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
|
||||
if self._closed or self._stop.is_set():
|
||||
raise ImageListCancelledError("Local image-list slot has stopped")
|
||||
if self._future is not None and not self._future.done():
|
||||
raise RuntimeError("Local image-list slot already has an in-flight item")
|
||||
self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
|
||||
return self._future
|
||||
|
||||
def request_stop(self):
|
||||
self._stop.set()
|
||||
# Wakes a graph waiting on an async node as well as stopping synchronous
|
||||
# nodes at their next progress/between-node checkpoint.
|
||||
with self._state_lock:
|
||||
if self._loop is not None and self._task is not None:
|
||||
self._loop.call_soon_threadsafe(self._task.cancel)
|
||||
|
||||
async def close(self):
|
||||
if self._closed:
|
||||
return
|
||||
# A stop requested before close indicates cancellation, not normal drain.
|
||||
cancelled = self._stop.is_set()
|
||||
self.request_stop()
|
||||
try:
|
||||
if self._future is not None:
|
||||
# Do not cancel this Future: cancelling an asyncio wrapper cannot
|
||||
# kill its running thread, and falsely makes cleanup look done.
|
||||
waiter = asyncio.wrap_future(self._future)
|
||||
while not waiter.done():
|
||||
try:
|
||||
await asyncio.shield(waiter)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break # The scheduler propagates execution errors.
|
||||
if waiter.done() and not waiter.cancelled():
|
||||
waiter.exception()
|
||||
finally:
|
||||
self._pool.shutdown(wait=True, cancel_futures=True)
|
||||
if not cancelled and not self._failed and self._caches is not None:
|
||||
_retain_outputs(self._caches.outputs)
|
||||
self._caches = None
|
||||
self._closed = True
|
||||
|
||||
def _run(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
# The normal Comfy prompt worker is single-threaded and blocked in End.
|
||||
# Fail closed if another local slot is nevertheless executing, rather
|
||||
# than interleave writes to core's process-global execution state.
|
||||
if not _EXECUTION_LOCK.acquire(blocking=False):
|
||||
raise RuntimeError("Another local image-list executor is already active")
|
||||
try:
|
||||
asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
|
||||
except asyncio.CancelledError as exc:
|
||||
self._failed = True
|
||||
raise ImageListCancelledError("Local image-list execution cancelled") from exc
|
||||
except BaseException:
|
||||
self._failed = True
|
||||
raise
|
||||
finally:
|
||||
with self._state_lock:
|
||||
self._loop = self._task = None
|
||||
_EXECUTION_LOCK.release()
|
||||
|
||||
def _check_stopped(self):
|
||||
import comfy.model_management
|
||||
if self._stop.is_set() or comfy.model_management.processing_interrupted():
|
||||
raise ImageListCancelledError("Local image-list execution interrupted")
|
||||
|
||||
async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
|
||||
import execution
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
from comfy_execution import progress
|
||||
from comfy_execution.graph_utils import GraphBuilder
|
||||
|
||||
self._check_stopped()
|
||||
with self._state_lock:
|
||||
self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
|
||||
dynamic_prompt = execution.DynamicPrompt(prompt)
|
||||
private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
|
||||
parent_progress = progress.global_progress_registry
|
||||
parent_hook = comfy.utils.PROGRESS_BAR_HOOK
|
||||
prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
|
||||
parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
|
||||
progress.global_progress_registry = private_progress
|
||||
|
||||
def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
|
||||
self._check_stopped()
|
||||
# Never route a child progress/preview event through the parent's
|
||||
# client or reset its handlers. Keep the private progress accurate.
|
||||
if node_id is not None:
|
||||
private_progress.update_progress(node_id, value, total, preview)
|
||||
|
||||
comfy.utils.PROGRESS_BAR_HOOK = progress_hook
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
if self._caches is None:
|
||||
outputs = _take_outputs(execution)
|
||||
objects = execution.HierarchicalCache(execution.CacheKeySetID)
|
||||
self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
|
||||
caches = self._caches
|
||||
changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
|
||||
for cache in caches.all:
|
||||
await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
|
||||
cache.clean_unused()
|
||||
execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
|
||||
execution_list.add_node(end_id)
|
||||
server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
assets = execution.NoAssets(execution.args)
|
||||
extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
|
||||
executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
|
||||
while not execution_list.is_empty():
|
||||
self._check_stopped()
|
||||
node_id, error, exc = await execution_list.stage_node_execution()
|
||||
if error is not None:
|
||||
raise RuntimeError(f"Local image-list graph failed: {error}") from exc
|
||||
self._check_stopped()
|
||||
result, error, exc = await execution.execute(
|
||||
server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
|
||||
execution_list, pending_subgraphs, pending_async, ui, assets)
|
||||
if result == execution.ExecutionResult.FAILURE:
|
||||
if isinstance(exc, (comfy.model_management.InterruptProcessingException,
|
||||
ImageListCancelledError)):
|
||||
raise ImageListCancelledError("Local image-list execution interrupted") from exc
|
||||
raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
|
||||
if result == execution.ExecutionResult.PENDING:
|
||||
execution_list.unstage_node_execution()
|
||||
else:
|
||||
execution_list.complete_node_execution()
|
||||
self._check_stopped()
|
||||
finally:
|
||||
# Core execute creates async-node and unblock tasks. This is a private
|
||||
# event loop, so every other task is owned here. Drain/cancel them
|
||||
# before restoring globals, including when the graph failed early.
|
||||
with self._state_lock:
|
||||
self._task = None # Stop requests must not interrupt this drain.
|
||||
try:
|
||||
while True:
|
||||
pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
|
||||
if not pending:
|
||||
break
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
drain = asyncio.gather(*pending, return_exceptions=True)
|
||||
while not drain.done():
|
||||
try:
|
||||
await asyncio.shield(drain)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
finally:
|
||||
progress.global_progress_registry = parent_progress
|
||||
comfy.utils.PROGRESS_BAR_HOOK = parent_hook
|
||||
for key, value in parent_prefix.items():
|
||||
setattr(GraphBuilder, key, value)
|
||||
File diff suppressed because it is too large
Load Diff
+1
-37
@@ -1,37 +1 @@
|
||||
from .utilities import (
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
ImageBatchDivider,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from .collector import DistributedCollectorNode
|
||||
from .branch import DistributedBranch
|
||||
from .branch_collector import DistributedBranchCollector
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DistributedCollector": DistributedCollectorNode,
|
||||
"DistributedBranch": DistributedBranch,
|
||||
"DistributedBranchCollector": DistributedBranchCollector,
|
||||
"DistributedSeed": DistributedSeed,
|
||||
"DistributedModelName": DistributedModelName,
|
||||
"DistributedValue": DistributedValue,
|
||||
"ImageBatchDivider": ImageBatchDivider,
|
||||
"AudioBatchDivider": AudioBatchDivider,
|
||||
"DistributedEmptyImage": DistributedEmptyImage,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DistributedCollector": "Distributed Collector",
|
||||
"DistributedBranch": "Distributed Branch",
|
||||
"DistributedBranchCollector": "Distributed Branch 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."""
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
from .utilities import any_type
|
||||
|
||||
|
||||
class DistributedBranch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedBranch"]) -> dict[str, Any]:
|
||||
return {
|
||||
"required": {
|
||||
"input": (any_type,),
|
||||
"num_branches": ("INT", {"default": 2, "min": 2, "max": 10, "step": 1}),
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"assigned_branch": ("INT", {"default": -1, "min": -1, "max": 9}),
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = tuple([any_type] * 10)
|
||||
RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(10)])
|
||||
FUNCTION = "branch"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def branch(
|
||||
self,
|
||||
input: Any,
|
||||
num_branches: int = 2,
|
||||
**_kwargs: Any,
|
||||
) -> tuple[Any, ...]:
|
||||
return tuple([input] * 10)
|
||||
@@ -1,298 +0,0 @@
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
import torch
|
||||
|
||||
from ..utils.auth import distributed_auth_headers
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..utils.config import is_master_delegate_only, load_config
|
||||
from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
|
||||
from ..utils.logging import log
|
||||
from ..utils.network import get_client_session
|
||||
from ..utils.worker_ids import coerce_enabled_worker_ids
|
||||
from .context_kwargs import parse_distributed_hidden_context
|
||||
from .hidden_inputs import build_distributed_hidden_inputs
|
||||
from .queue_wait import collect_worker_queue_results
|
||||
from .runtime_helpers import (
|
||||
get_prompt_server_instance as _get_prompt_server_instance,
|
||||
throw_if_processing_interrupted as _throw_if_processing_interrupted,
|
||||
)
|
||||
from .utilities import any_type
|
||||
|
||||
|
||||
MAX_BRANCH_OUTPUTS = 10
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BranchRunContext:
|
||||
multi_job_id: str = ""
|
||||
is_worker: bool = False
|
||||
master_url: str = ""
|
||||
enabled_worker_ids: list[str] = field(default_factory=list)
|
||||
worker_id: str = ""
|
||||
assigned_branch: int = -1
|
||||
delegate_only: bool = False
|
||||
|
||||
class DistributedBranchCollector:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedBranchCollector"]) -> dict[str, Any]:
|
||||
optional_inputs = {
|
||||
f"branch_{idx + 1}": (any_type,)
|
||||
for idx in range(MAX_BRANCH_OUTPUTS)
|
||||
}
|
||||
return {
|
||||
"required": {
|
||||
"num_branches": ("INT", {"default": 2, "min": 2, "max": MAX_BRANCH_OUTPUTS, "step": 1}),
|
||||
},
|
||||
"optional": optional_inputs,
|
||||
"hidden": build_distributed_hidden_inputs(
|
||||
include_assigned_branch=True,
|
||||
assigned_branch_max=MAX_BRANCH_OUTPUTS - 1,
|
||||
),
|
||||
}
|
||||
|
||||
RETURN_TYPES = tuple([any_type] * MAX_BRANCH_OUTPUTS)
|
||||
RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(MAX_BRANCH_OUTPUTS)])
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def _branch_values(self, **kwargs):
|
||||
values = []
|
||||
for idx in range(MAX_BRANCH_OUTPUTS):
|
||||
values.append(kwargs.get(f"branch_{idx + 1}", None))
|
||||
return values
|
||||
|
||||
def _build_run_context(self, **kwargs: Any) -> BranchRunContext:
|
||||
try:
|
||||
assigned_branch = int(kwargs.get("assigned_branch", -1))
|
||||
except (TypeError, ValueError):
|
||||
assigned_branch = -1
|
||||
|
||||
common_context = parse_distributed_hidden_context(kwargs)
|
||||
return BranchRunContext(
|
||||
assigned_branch=assigned_branch,
|
||||
**common_context,
|
||||
)
|
||||
|
||||
def _resolve_local_branch_value(self, branch_values, assigned_branch):
|
||||
try:
|
||||
assigned_idx = int(assigned_branch)
|
||||
except (TypeError, ValueError):
|
||||
assigned_idx = -1
|
||||
|
||||
if 0 <= assigned_idx < MAX_BRANCH_OUTPUTS:
|
||||
assigned_value = branch_values[assigned_idx]
|
||||
if assigned_value is not None:
|
||||
return assigned_value, assigned_idx
|
||||
|
||||
# In participant-pruned prompts, each participant should typically have
|
||||
# exactly one local branch input. If socket names drifted (e.g. branch_2
|
||||
# connected for a participant assigned to slot 0), remap that sole value
|
||||
# to the assigned slot so the participant still contributes correctly.
|
||||
non_none_indices = [idx for idx, value in enumerate(branch_values) if value is not None]
|
||||
if len(non_none_indices) == 1:
|
||||
return branch_values[non_none_indices[0]], assigned_idx
|
||||
|
||||
for idx, value in enumerate(branch_values):
|
||||
if value is not None:
|
||||
return value, idx
|
||||
return None, assigned_idx
|
||||
|
||||
def _build_outputs(self, branch_values, num_branches):
|
||||
outputs = [None] * MAX_BRANCH_OUTPUTS
|
||||
try:
|
||||
branch_count = int(num_branches)
|
||||
except (TypeError, ValueError):
|
||||
branch_count = 2
|
||||
branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
|
||||
|
||||
for idx in range(branch_count):
|
||||
value = branch_values[idx]
|
||||
if value is not None:
|
||||
outputs[idx] = value
|
||||
return tuple(outputs)
|
||||
|
||||
def run(
|
||||
self,
|
||||
num_branches: int = 2,
|
||||
**kwargs: Any,
|
||||
) -> tuple[Any, ...]:
|
||||
branch_values = self._branch_values(**kwargs)
|
||||
context = self._build_run_context(**kwargs)
|
||||
|
||||
if not context.multi_job_id:
|
||||
return self._build_outputs(branch_values, num_branches)
|
||||
|
||||
return run_async_in_server_loop(
|
||||
self.execute(
|
||||
branch_values,
|
||||
num_branches=num_branches,
|
||||
context=context,
|
||||
)
|
||||
)
|
||||
|
||||
def _participant_branch_slots(self, enabled_workers, delegate_mode, num_branches):
|
||||
participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers)
|
||||
participant_count = len(participants)
|
||||
if participant_count <= 0:
|
||||
return {}
|
||||
|
||||
slots_by_participant = {}
|
||||
for branch_slot in range(int(num_branches)):
|
||||
participant_id = str(participants[branch_slot % participant_count])
|
||||
slots_by_participant.setdefault(participant_id, []).append(branch_slot)
|
||||
return slots_by_participant
|
||||
|
||||
def _expected_worker_ids_for_branches(self, enabled_workers, delegate_mode, num_branches):
|
||||
slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, num_branches)
|
||||
expected = []
|
||||
for participant_id, slots in slots_by_participant.items():
|
||||
if participant_id == "master" or not slots:
|
||||
continue
|
||||
expected.append(str(participant_id))
|
||||
return expected
|
||||
|
||||
def _fallback_for_missing_branch(self, source_value):
|
||||
if isinstance(source_value, torch.Tensor):
|
||||
tensor = source_value
|
||||
if tensor.ndim == 3:
|
||||
tensor = tensor.unsqueeze(0)
|
||||
if tensor.ndim == 4:
|
||||
fallback = torch.zeros_like(tensor)
|
||||
if fallback.is_cuda:
|
||||
fallback = fallback.cpu()
|
||||
return ensure_contiguous(fallback)
|
||||
return torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
|
||||
async def _send_branch_result_to_master(self, value, branch_idx, multi_job_id, master_url, worker_id):
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("DistributedBranchCollector currently supports torch.Tensor results for worker transfer.")
|
||||
|
||||
image_batch = value
|
||||
if image_batch.ndim == 3:
|
||||
image_batch = image_batch.unsqueeze(0)
|
||||
if image_batch.ndim != 4 or image_batch.shape[0] <= 0:
|
||||
raise ValueError(
|
||||
"DistributedBranchCollector tensor result must be IMAGE-like with shape [B,H,W,C] or [H,W,C]."
|
||||
)
|
||||
|
||||
payload = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": int(branch_idx),
|
||||
"image": encode_tensor_png_data_url(image_batch, 0),
|
||||
"is_last": True,
|
||||
}
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
headers=distributed_auth_headers(load_config()),
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
async def execute(self, *args: Any, **kwargs: Any) -> tuple[Any, ...]:
|
||||
"""Compatibility wrapper with a uniform execute() signature across collectors."""
|
||||
return await self._execute_branch(*args, **kwargs)
|
||||
|
||||
async def _execute_branch(
|
||||
self,
|
||||
branch_values: list[Any],
|
||||
num_branches: int = 2,
|
||||
context: BranchRunContext | None = None,
|
||||
) -> tuple[Any, ...]:
|
||||
run_context = context or BranchRunContext()
|
||||
try:
|
||||
branch_count = int(num_branches)
|
||||
except (TypeError, ValueError):
|
||||
branch_count = 2
|
||||
branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
|
||||
|
||||
outputs = [None] * MAX_BRANCH_OUTPUTS
|
||||
for idx in range(branch_count):
|
||||
if branch_values[idx] is not None:
|
||||
outputs[idx] = branch_values[idx]
|
||||
|
||||
local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, run_context.assigned_branch)
|
||||
if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
|
||||
outputs[local_branch_idx] = local_value
|
||||
|
||||
if run_context.is_worker:
|
||||
if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
|
||||
try:
|
||||
await self._send_branch_result_to_master(
|
||||
local_value,
|
||||
local_branch_idx,
|
||||
run_context.multi_job_id,
|
||||
run_context.master_url,
|
||||
run_context.worker_id,
|
||||
)
|
||||
outputs[local_branch_idx] = local_value
|
||||
except Exception as exc:
|
||||
log(f"Worker - DistributedBranchCollector failed to send branch result: {exc}")
|
||||
return tuple(outputs)
|
||||
|
||||
delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
|
||||
enabled_workers = coerce_enabled_worker_ids(run_context.enabled_worker_ids)
|
||||
slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count)
|
||||
expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count)
|
||||
expected_workers = set(expected_worker_ids)
|
||||
|
||||
if not expected_workers:
|
||||
return tuple(outputs)
|
||||
|
||||
prompt_server = _get_prompt_server_instance()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if run_context.multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue()
|
||||
|
||||
workers_done: set[str] = set()
|
||||
|
||||
def _handle_queue_result(result: dict[str, Any]) -> None:
|
||||
image_index = result.get("image_index")
|
||||
tensor = result.get("tensor")
|
||||
try:
|
||||
slot_idx = int(image_index)
|
||||
except (TypeError, ValueError):
|
||||
slot_idx = -1
|
||||
|
||||
if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None:
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
if tensor.is_cuda:
|
||||
tensor = tensor.cpu()
|
||||
tensor = ensure_contiguous(tensor)
|
||||
outputs[slot_idx] = tensor
|
||||
|
||||
try:
|
||||
workers_done = await collect_worker_queue_results(
|
||||
prompt_server=prompt_server,
|
||||
multi_job_id=run_context.multi_job_id,
|
||||
expected_workers=expected_workers,
|
||||
on_result=_handle_queue_result,
|
||||
timeout_log_prefix=(
|
||||
"Master - DistributedBranchCollector heartbeat timeout. "
|
||||
"Still waiting for workers: "
|
||||
),
|
||||
throw_if_interrupted=_throw_if_processing_interrupted,
|
||||
)
|
||||
|
||||
finally:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
|
||||
if hasattr(prompt_server, "distributed_job_allowed_workers"):
|
||||
prompt_server.distributed_job_allowed_workers.pop(run_context.multi_job_id, None)
|
||||
|
||||
missing_workers = sorted(expected_workers - workers_done)
|
||||
if missing_workers:
|
||||
fallback = self._fallback_for_missing_branch(local_value)
|
||||
for missing_worker_id in missing_workers:
|
||||
for slot_idx in slots_by_participant.get(str(missing_worker_id), []):
|
||||
if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and outputs[slot_idx] is None:
|
||||
outputs[slot_idx] = fallback
|
||||
|
||||
return tuple(outputs)
|
||||
+758
-330
File diff suppressed because it is too large
Load Diff
@@ -1,26 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
from ..utils.parsing import coerce_bool
|
||||
from ..utils.worker_ids import coerce_enabled_worker_ids
|
||||
|
||||
|
||||
def _parse_bool(value: Any, default: bool = False) -> bool:
|
||||
return coerce_bool(value, default=default)
|
||||
|
||||
|
||||
def _parse_enabled_workers(value: Any) -> list[str]:
|
||||
if value is None:
|
||||
return []
|
||||
return coerce_enabled_worker_ids(value)
|
||||
|
||||
|
||||
def parse_distributed_hidden_context(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Parse common distributed hidden inputs from node kwargs."""
|
||||
return {
|
||||
"multi_job_id": str(kwargs.get("multi_job_id", "") or ""),
|
||||
"is_worker": _parse_bool(kwargs.get("is_worker", False)),
|
||||
"master_url": str(kwargs.get("master_url", "") or ""),
|
||||
"enabled_worker_ids": _parse_enabled_workers(kwargs.get("enabled_worker_ids", "[]")),
|
||||
"worker_id": str(kwargs.get("worker_id", "") or ""),
|
||||
"delegate_only": _parse_bool(kwargs.get("delegate_only", False)),
|
||||
}
|
||||
+88
-335
@@ -1,28 +1,25 @@
|
||||
import json
|
||||
import math
|
||||
from functools import wraps
|
||||
from typing import Any, Callable
|
||||
|
||||
import comfy.samplers
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
from .hidden_inputs import build_distributed_hidden_inputs
|
||||
|
||||
from ..upscale.tile_ops import TileOpsMixin
|
||||
from ..upscale.result_collector import ResultCollectorMixin
|
||||
from ..upscale.worker_comms import WorkerCommsMixin
|
||||
from ..upscale.job_state import JobStateMixin
|
||||
from ..upscale.processing_args import UpscaleCoreArgs
|
||||
from ..upscale.modes.single_gpu import SingleGpuModeMixin
|
||||
from ..upscale.modes.static import StaticModeMixin
|
||||
from ..upscale.modes.dynamic import DynamicModeMixin
|
||||
from ..utils.worker_ids import coerce_enabled_worker_ids
|
||||
|
||||
def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
def sync_wrapper(async_func):
|
||||
"""Decorator to wrap async methods for synchronous execution."""
|
||||
@wraps(async_func)
|
||||
def sync_func(self, *args: Any, **kwargs: Any) -> Any:
|
||||
def sync_func(self, *args, **kwargs):
|
||||
# Use run_async_in_server_loop for ComfyUI compatibility
|
||||
return run_async_in_server_loop(
|
||||
async_func(self, *args, **kwargs),
|
||||
@@ -30,10 +27,21 @@ def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
)
|
||||
return sync_func
|
||||
|
||||
|
||||
def _parse_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]:
|
||||
"""Backward-compatible alias for enabled-worker normalization."""
|
||||
return coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
def _parse_enabled_worker_ids(enabled_worker_ids):
|
||||
"""Parse enabled worker IDs from either JSON or list input."""
|
||||
if isinstance(enabled_worker_ids, list):
|
||||
return [str(worker_id) for worker_id in enabled_worker_ids]
|
||||
if not enabled_worker_ids:
|
||||
return []
|
||||
if isinstance(enabled_worker_ids, str):
|
||||
try:
|
||||
parsed = json.loads(enabled_worker_ids)
|
||||
except json.JSONDecodeError:
|
||||
log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
|
||||
return []
|
||||
if isinstance(parsed, list):
|
||||
return [str(wid) for wid in parsed]
|
||||
return []
|
||||
|
||||
class UltimateSDUpscaleDistributed(
|
||||
DynamicModeMixin,
|
||||
@@ -48,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):
|
||||
@@ -73,14 +81,7 @@ class UltimateSDUpscaleDistributed(
|
||||
debug_log("UltimateSDUpscaleDistributed - Node initialized")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["UltimateSDUpscaleDistributed"]) -> dict[str, Any]:
|
||||
hidden_inputs = build_distributed_hidden_inputs()
|
||||
hidden_inputs.update(
|
||||
{
|
||||
"tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility
|
||||
"dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
|
||||
}
|
||||
)
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"upscaled_image": ("IMAGE",),
|
||||
@@ -101,7 +102,15 @@ class UltimateSDUpscaleDistributed(
|
||||
"force_uniform_tiles": ("BOOLEAN", {"default": True}),
|
||||
"tiled_decode": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"hidden": hidden_inputs,
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility
|
||||
"dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
@@ -113,47 +122,12 @@ class UltimateSDUpscaleDistributed(
|
||||
"""Force re-execution."""
|
||||
return float("nan") # Always re-execute
|
||||
|
||||
def run(
|
||||
self,
|
||||
upscaled_image: Any,
|
||||
model: Any,
|
||||
positive: Any,
|
||||
negative: Any,
|
||||
vae: Any,
|
||||
seed: int,
|
||||
steps: int,
|
||||
cfg: float,
|
||||
sampler_name: str,
|
||||
scheduler: str,
|
||||
denoise: float,
|
||||
tile_width: int,
|
||||
tile_height: int,
|
||||
padding: int,
|
||||
mask_blur: int,
|
||||
force_uniform_tiles: bool,
|
||||
tiled_decode: bool,
|
||||
multi_job_id: str = "",
|
||||
is_worker: bool = False,
|
||||
master_url: str = "",
|
||||
enabled_worker_ids: str = "[]",
|
||||
worker_id: str = "",
|
||||
tile_indices: str = "",
|
||||
dynamic_threshold: int = 8,
|
||||
) -> tuple[Any, ...]:
|
||||
def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode,
|
||||
multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]",
|
||||
worker_id="", tile_indices="", dynamic_threshold=8):
|
||||
"""Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
|
||||
core_args = UpscaleCoreArgs(
|
||||
model=model,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
seed=seed,
|
||||
steps=steps,
|
||||
cfg=cfg,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler,
|
||||
denoise=denoise,
|
||||
tiled_decode=tiled_decode,
|
||||
)
|
||||
# Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
|
||||
try:
|
||||
batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
|
||||
@@ -168,15 +142,9 @@ class UltimateSDUpscaleDistributed(
|
||||
)
|
||||
if not multi_job_id:
|
||||
# No distributed processing, run single GPU version
|
||||
return self.process_single_gpu(
|
||||
upscaled_image=upscaled_image,
|
||||
core_args=core_args,
|
||||
tile_width=tile_width,
|
||||
tile_height=tile_height,
|
||||
padding=padding,
|
||||
mask_blur=mask_blur,
|
||||
force_uniform_tiles=force_uniform_tiles,
|
||||
)
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
if is_worker:
|
||||
# Worker mode: process tiles synchronously
|
||||
@@ -187,64 +155,24 @@ class UltimateSDUpscaleDistributed(
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
else:
|
||||
# Master mode: distribute and collect synchronously
|
||||
return self.process_master(
|
||||
upscaled_image=upscaled_image,
|
||||
core_args=core_args,
|
||||
tile_width=tile_width,
|
||||
tile_height=tile_height,
|
||||
padding=padding,
|
||||
mask_blur=mask_blur,
|
||||
force_uniform_tiles=force_uniform_tiles,
|
||||
multi_job_id=multi_job_id,
|
||||
enabled_worker_ids=enabled_worker_ids,
|
||||
dynamic_threshold=dynamic_threshold,
|
||||
)
|
||||
return self.process_master(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold)
|
||||
|
||||
def process_worker(
|
||||
self,
|
||||
upscaled_image: Any,
|
||||
model: Any,
|
||||
positive: Any,
|
||||
negative: Any,
|
||||
vae: Any,
|
||||
seed: int,
|
||||
steps: int,
|
||||
cfg: float,
|
||||
sampler_name: str,
|
||||
scheduler: str,
|
||||
denoise: float,
|
||||
tile_width: int,
|
||||
tile_height: int,
|
||||
padding: int,
|
||||
mask_blur: int,
|
||||
force_uniform_tiles: bool,
|
||||
tiled_decode: bool,
|
||||
multi_job_id: str,
|
||||
master_url: str,
|
||||
worker_id: str,
|
||||
enabled_worker_ids: str,
|
||||
dynamic_threshold: int,
|
||||
) -> tuple[Any, ...]:
|
||||
def process_worker(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold):
|
||||
"""Unified worker processing - handles both static and dynamic modes."""
|
||||
core_args = UpscaleCoreArgs(
|
||||
model=model,
|
||||
positive=positive,
|
||||
negative=negative,
|
||||
vae=vae,
|
||||
seed=seed,
|
||||
steps=steps,
|
||||
cfg=cfg,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler,
|
||||
denoise=denoise,
|
||||
tiled_decode=tiled_decode,
|
||||
)
|
||||
# Get batch size to determine mode
|
||||
batch_size = upscaled_image.shape[0]
|
||||
|
||||
# Ensure mode consistency across master/workers via shared threshold
|
||||
# Determine mode (must match master's logic)
|
||||
enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
# Compute number of tiles for this image to decide if tile distribution makes sense
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
@@ -253,47 +181,31 @@ class UltimateSDUpscaleDistributed(
|
||||
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# For USDU-style processing, we want tile distribution whenever workers are available
|
||||
# and there is more than one tile to process for single-image runs.
|
||||
if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
|
||||
# and there is more than one tile to process, even if batch == 1.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
debug_log(f"USDU Dist Worker - Batch size {batch_size}")
|
||||
|
||||
if mode == "dynamic":
|
||||
return self.process_worker_dynamic(
|
||||
upscaled_image=upscaled_image,
|
||||
core_args=core_args,
|
||||
tile_width=tile_width,
|
||||
tile_height=tile_height,
|
||||
padding=padding,
|
||||
mask_blur=mask_blur,
|
||||
force_uniform_tiles=force_uniform_tiles,
|
||||
multi_job_id=multi_job_id,
|
||||
master_url=master_url,
|
||||
worker_id=worker_id,
|
||||
enabled_worker_ids=enabled_worker_ids,
|
||||
dynamic_threshold=dynamic_threshold,
|
||||
)
|
||||
return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
|
||||
# Static mode - enhanced with health monitoring and retry logic
|
||||
return self._process_worker_static_sync(upscaled_image, core_args,
|
||||
return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, multi_job_id, master_url,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_workers)
|
||||
|
||||
def process_master(
|
||||
self,
|
||||
upscaled_image: Any,
|
||||
core_args: UpscaleCoreArgs,
|
||||
tile_width: int,
|
||||
tile_height: int,
|
||||
padding: int,
|
||||
mask_blur: int,
|
||||
force_uniform_tiles: bool,
|
||||
multi_job_id: str,
|
||||
enabled_worker_ids: str,
|
||||
dynamic_threshold: int,
|
||||
) -> tuple[Any, ...]:
|
||||
def process_master(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold):
|
||||
"""Unified master processing with enhanced monitoring and failure handling."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
@@ -312,206 +224,47 @@ class UltimateSDUpscaleDistributed(
|
||||
)
|
||||
|
||||
# Parse enabled workers
|
||||
enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
# Determine processing mode
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# Prefer tile-based static distribution when workers are available and there are multiple tiles,
|
||||
# for single-image jobs to spread tiles across GPUs like the legacy tile queue.
|
||||
if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
|
||||
# even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
|
||||
|
||||
if mode == "single_gpu":
|
||||
# No workers, process all tiles locally
|
||||
return self.process_single_gpu(
|
||||
upscaled_image=upscaled_image,
|
||||
core_args=core_args,
|
||||
tile_width=tile_width,
|
||||
tile_height=tile_height,
|
||||
padding=padding,
|
||||
mask_blur=mask_blur,
|
||||
force_uniform_tiles=force_uniform_tiles,
|
||||
)
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
elif mode == "dynamic":
|
||||
# Dynamic mode for large batches
|
||||
return self.process_master_dynamic(
|
||||
upscaled_image=upscaled_image,
|
||||
core_args=core_args,
|
||||
tile_width=tile_width,
|
||||
tile_height=tile_height,
|
||||
padding=padding,
|
||||
mask_blur=mask_blur,
|
||||
force_uniform_tiles=force_uniform_tiles,
|
||||
multi_job_id=multi_job_id,
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
return self.process_master_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers)
|
||||
|
||||
# Static mode - enhanced with unified job management
|
||||
return self._process_master_static_sync(upscaled_image, core_args,
|
||||
return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, multi_job_id, enabled_workers,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
|
||||
all_tiles, num_tiles_per_image)
|
||||
|
||||
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
||||
"""Determine mode from worker availability and configured threshold."""
|
||||
"""Determines processing mode per requested policy:
|
||||
- any workers => prefer static (tile-based) for USDU
|
||||
- no workers => single_gpu
|
||||
"""
|
||||
if num_workers == 0:
|
||||
return "single_gpu"
|
||||
threshold = max(1, int(dynamic_threshold))
|
||||
if int(batch_size) >= threshold:
|
||||
return "dynamic"
|
||||
# Default to static when distributed; master/worker may still override if special cases arise
|
||||
return "static"
|
||||
|
||||
class USDUDelegateCollector:
|
||||
"""Lightweight stand-in used automatically in delegate-only master prompts.
|
||||
|
||||
When the master is in delegate-only mode, the orchestration code swaps the
|
||||
full ``UltimateSDUpscaleDistributed`` class for this one so that no upstream
|
||||
model/image nodes execute on the master. Workers initialise the dynamic job
|
||||
queue via ``/distributed/init_dynamic_job`` and this node simply waits for
|
||||
all images to arrive, then assembles the output tensor.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"delegate_only": ("BOOLEAN", {"default": True}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image/upscaling"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
def run(self, multi_job_id="", enabled_worker_ids="[]", **kwargs):
|
||||
import time
|
||||
import torch
|
||||
from ..utils.image import pil_to_tensor
|
||||
from ..upscale.job_models import ImageJobState
|
||||
|
||||
if not multi_job_id:
|
||||
log("USDUDelegateCollector: no multi_job_id, returning empty image")
|
||||
return (torch.zeros(1, 64, 64, 3),)
|
||||
|
||||
enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
log(f"USDU delegate-only: waiting for workers to init job {multi_job_id}")
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Wait for workers to create the job via /distributed/init_dynamic_job
|
||||
max_wait = 120.0
|
||||
poll_interval = 1.0
|
||||
start = time.time()
|
||||
job_data = None
|
||||
while time.time() - start < max_wait:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState) and job_data.batch_size > 0:
|
||||
break
|
||||
job_data = None
|
||||
time.sleep(poll_interval)
|
||||
|
||||
if job_data is None:
|
||||
log(f"USDUDelegateCollector: job {multi_job_id} was not initialized by workers within {max_wait}s")
|
||||
return (torch.zeros(1, 64, 64, 3),)
|
||||
|
||||
batch_size = job_data.batch_size
|
||||
log(f"USDU delegate-only: job ready, collecting {batch_size} images from {num_workers} workers")
|
||||
|
||||
# Collect all images using the existing result collector
|
||||
collected = run_async_in_server_loop(
|
||||
self._collect_all_images(multi_job_id, batch_size, num_workers, prompt_server),
|
||||
timeout=600.0,
|
||||
)
|
||||
|
||||
# Assemble tensor from collected PIL images
|
||||
result_images = []
|
||||
for idx in range(batch_size):
|
||||
img = collected.get(idx)
|
||||
if img is not None:
|
||||
result_images.append(pil_to_tensor(img))
|
||||
else:
|
||||
log(f"USDUDelegateCollector: missing image {idx}, using blank")
|
||||
result_images.append(torch.zeros(1, 64, 64, 3))
|
||||
|
||||
result_tensor = torch.cat(result_images, dim=0)
|
||||
log(f"USDU delegate-only: collected all {batch_size} images")
|
||||
return (result_tensor,)
|
||||
|
||||
@staticmethod
|
||||
async def _collect_all_images(multi_job_id, batch_size, num_workers, prompt_server):
|
||||
"""Wait for all images to arrive from workers."""
|
||||
import asyncio
|
||||
from ..upscale.job_models import ImageJobState
|
||||
from ..upscale.job_timeout import check_and_requeue_timed_out_workers
|
||||
|
||||
timeout_seconds = 300.0
|
||||
poll_interval = 2.0
|
||||
start = asyncio.get_event_loop().time()
|
||||
|
||||
while True:
|
||||
# Drain results from the queue
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
# Drain queue items into completed_images
|
||||
while True:
|
||||
try:
|
||||
result = job_data.queue.get_nowait()
|
||||
worker_id = result.get("worker_id")
|
||||
if "image_idx" in result and "image" in result:
|
||||
idx = result["image_idx"]
|
||||
if idx not in job_data.completed_images:
|
||||
job_data.completed_images[idx] = result["image"]
|
||||
debug_log(f"Delegate collected image {idx} from worker {worker_id}")
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
completed_count = len(job_data.completed_images)
|
||||
if completed_count >= batch_size:
|
||||
debug_log(f"Delegate: all {batch_size} images collected")
|
||||
completed = dict(job_data.completed_images)
|
||||
# Cleanup
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
|
||||
return completed
|
||||
|
||||
# Check for timeouts periodically
|
||||
await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
|
||||
|
||||
elapsed = asyncio.get_event_loop().time() - start
|
||||
if elapsed > timeout_seconds:
|
||||
log(f"USDUDelegateCollector: timed out after {elapsed:.0f}s with {len(getattr(job_data, 'completed_images', {}))} of {batch_size} images")
|
||||
if isinstance(job_data, ImageJobState):
|
||||
completed = dict(job_data.completed_images)
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
|
||||
return completed
|
||||
return {}
|
||||
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
|
||||
# Node registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
||||
"USDUDelegateCollector": USDUDelegateCollector,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
||||
# No display name for USDUDelegateCollector — it's an internal node
|
||||
}
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
|
||||
def build_worker_identity_hidden_inputs() -> dict[str, tuple[str, dict[str, Any]]]:
|
||||
"""Hidden-input contract for worker identity context."""
|
||||
return {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
}
|
||||
|
||||
|
||||
def build_distributed_hidden_inputs(
|
||||
*,
|
||||
include_assigned_branch: bool = False,
|
||||
assigned_branch_max: int = 9,
|
||||
include_worker_batch_size: bool = False,
|
||||
include_pass_through: bool = False,
|
||||
) -> dict[str, tuple[str, dict[str, Any]]]:
|
||||
hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"delegate_only": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
hidden_inputs.update(build_worker_identity_hidden_inputs())
|
||||
if include_assigned_branch:
|
||||
hidden_inputs["assigned_branch"] = (
|
||||
"INT",
|
||||
{"default": -1, "min": -1, "max": assigned_branch_max},
|
||||
)
|
||||
if include_worker_batch_size:
|
||||
hidden_inputs["worker_batch_size"] = (
|
||||
"INT",
|
||||
{"default": 1, "min": 1, "max": 1024},
|
||||
)
|
||||
if include_pass_through:
|
||||
hidden_inputs["pass_through"] = ("BOOLEAN", {"default": False})
|
||||
return hidden_inputs
|
||||
@@ -0,0 +1,376 @@
|
||||
"""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 (
|
||||
MAX_VALUE_SLOTS,
|
||||
VALUE_NAMES,
|
||||
VALUES_INPUT,
|
||||
end_override_names,
|
||||
create_master_job,
|
||||
normalize_image_list_inputs,
|
||||
normalize_worker_result,
|
||||
run_image_list_job,
|
||||
encode_result,
|
||||
encode_request_payload,
|
||||
unwrap_single,
|
||||
is_valid_item_identity,
|
||||
)
|
||||
|
||||
|
||||
def _stable_worker_binding(prompt, start_id, identity):
|
||||
"""Resolve transient state only from the exact Start -> End context wire."""
|
||||
if not is_valid_item_identity(identity):
|
||||
raise ValueError("Invalid worker item identity")
|
||||
start = prompt.get(start_id) if isinstance(prompt, dict) else None
|
||||
if (not isinstance(start, dict) or start.get("class_type") != "DistributedImageListStart"
|
||||
or start.get("inputs", {}).get("worker_item_identity") != identity):
|
||||
raise RuntimeError("Worker item identity is missing its paired Start")
|
||||
candidates = []
|
||||
for end_id, node in prompt.items():
|
||||
if not isinstance(node, dict) or node.get("class_type") != "DistributedImageListEnd":
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
link = inputs.get("context")
|
||||
if isinstance(link, (list, tuple)) and list(link) == [start_id, 3]:
|
||||
candidates.append((end_id, inputs))
|
||||
if len(candidates) != 1:
|
||||
raise RuntimeError("Worker Start must have exactly one paired End context wire")
|
||||
end_id, inputs = candidates[0]
|
||||
if inputs.get("worker_item_identity") != identity:
|
||||
raise ValueError("Paired End worker item identity does not match Start")
|
||||
for name in ("worker_result_token", "worker_context_id"):
|
||||
if not isinstance(inputs.get(name), str) or not inputs[name]:
|
||||
raise RuntimeError(f"Worker paired End is missing {name}")
|
||||
index = inputs.get("worker_item_index")
|
||||
if isinstance(index, bool) or not isinstance(index, int) or index < 0:
|
||||
raise ValueError("Worker paired End item index is invalid")
|
||||
return end_id, inputs
|
||||
|
||||
|
||||
# Types the transport can carry exactly. Anything else (models, latents, ...)
|
||||
# must be loaded inside the region or shared as a model dependency instead.
|
||||
_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
|
||||
assert len(VALUE_NAMES) == MAX_VALUE_SLOTS
|
||||
|
||||
|
||||
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. Each image becomes one job. "
|
||||
"mask and value_N inputs take 1 value (shared by every image) or exactly "
|
||||
"one value per image, paired in order."),
|
||||
inputs=[
|
||||
io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one job."),
|
||||
io.Mask.Input("mask", optional=True,
|
||||
tooltip="1 mask shared by every image, or one mask per image."),
|
||||
io.Autogrow.Input(VALUES_INPUT, optional=True, template=io.Autogrow.TemplateNames(
|
||||
io.MultiType.Input("value", types=_VALUE_TYPES,
|
||||
tooltip=("Any image, mask, text, number or boolean. 1 value is shared by "
|
||||
"every image; one value per image pairs them in order. "
|
||||
"A batch stays one value.")),
|
||||
names=list(VALUE_NAMES), min=0)),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="image", is_output_list=True),
|
||||
io.Mask.Output(display_name="mask", 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),
|
||||
# value_N outputs mirror value_N inputs. The web extension hides
|
||||
# unused ones and shows the connected input's type.
|
||||
*[io.AnyType.Output(id=name, display_name=name, is_output_list=True) for name in VALUE_NAMES],
|
||||
],
|
||||
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, mask=None, values=None,
|
||||
worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
|
||||
image_list_master_callback="", image_list_execution_id="",
|
||||
image_list_workflow_meta=None, image_list_master_participates=None,
|
||||
image_list_local_execution=False, worker_item_identity=None):
|
||||
worker_item_token = unwrap_single(worker_item_token)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
|
||||
if worker_item_identity is not None:
|
||||
prompt = getattr(cls.hidden, "prompt", None)
|
||||
start_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
_, delivery = _stable_worker_binding(prompt, start_id, worker_item_identity)
|
||||
# Only cache misses execute Start. Never put this token/context in
|
||||
# Start's literal inputs or returned context: both would poison reuse.
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(delivery["worker_result_token"],
|
||||
item_identity=worker_item_identity,
|
||||
item_index=delivery["worker_item_index"])
|
||||
item_values = item.get("values") or {}
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("mask")], [item["index"]], [worker_item_identity],
|
||||
*[[item_values.get(name)] for name in VALUE_NAMES],
|
||||
)
|
||||
if worker_item_token:
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import consume_worker_item
|
||||
item = consume_worker_item(worker_item_token)
|
||||
else:
|
||||
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")
|
||||
item_values = item.get("values") or {}
|
||||
return io.NodeOutput(
|
||||
[item["image"]], [item.get("mask")], [item["index"]], [context_id],
|
||||
*[[item_values.get(name)] for name in VALUE_NAMES],
|
||||
)
|
||||
|
||||
items = normalize_image_list_inputs(images, mask, values)
|
||||
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 comfy_execution.utils import get_executing_context
|
||||
from ..image_list.lifecycle import has_safe_execution_targets
|
||||
import server
|
||||
executing = get_executing_context()
|
||||
explicit_targets = has_safe_execution_targets(
|
||||
getattr(server.PromptServer, "instance", None),
|
||||
executing.prompt_id if executing is not None else None, prompt)
|
||||
from ..image_list.runtime import validate_image_list_regions
|
||||
regions = validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
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}")
|
||||
master_participates = unwrap_single(image_list_master_participates)
|
||||
if master_participates is None:
|
||||
from ..utils.config import load_config
|
||||
master_participates = not bool((load_config().get("settings", {}) or {}).get("master_delegate_only", False))
|
||||
if not isinstance(master_participates, bool):
|
||||
raise ValueError("image_list_master_participates must be a boolean")
|
||||
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,
|
||||
master_participates=master_participates,
|
||||
explicit_targets=explicit_targets,
|
||||
)
|
||||
return io.NodeOutput(
|
||||
[item["image"] for item in items],
|
||||
[item.get("mask") for item in items],
|
||||
[item["index"] for item in items],
|
||||
[context_id for _ in items],
|
||||
*[[item["values"].get(name) for item in items] for name in VALUE_NAMES],
|
||||
)
|
||||
|
||||
|
||||
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, text and values lists. "
|
||||
"Start values are retained; connected value_N inputs replace or add that key per item."),
|
||||
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"),
|
||||
*[io.MultiType.Input(name, types=_VALUE_TYPES, optional=True, lazy=True,
|
||||
tooltip="Replace or add this value for each item; exactly one value per item.")
|
||||
for name in VALUE_NAMES],
|
||||
],
|
||||
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),
|
||||
io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Output(
|
||||
id="values", display_name="values", 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, io.Hidden.execution_list],
|
||||
)
|
||||
|
||||
@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),
|
||||
*((name, kwargs.get(name)) for name in end_override_names(end_inputs))):
|
||||
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="",
|
||||
image_list_local_execution=False, worker_context_id="",
|
||||
worker_item_identity=None, **values):
|
||||
worker_result_token = unwrap_single(worker_result_token)
|
||||
master_callback_url = unwrap_single(master_callback_url)
|
||||
worker_context_id = unwrap_single(worker_context_id)
|
||||
worker_item_identity = unwrap_single(worker_item_identity)
|
||||
if worker_item_identity is not None and not worker_result_token:
|
||||
raise RuntimeError("Worker item identity is missing its delivery token")
|
||||
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")
|
||||
prompt = getattr(cls.hidden, "prompt", {}) or {}
|
||||
context_id = context_values[0]
|
||||
if worker_item_identity is not None:
|
||||
end_id = str(getattr(cls.hidden, "unique_id", ""))
|
||||
node = prompt.get(end_id, {})
|
||||
link = node.get("inputs", {}).get("context")
|
||||
if (node.get("class_type") != "DistributedImageListEnd"
|
||||
or not isinstance(link, (list, tuple)) or len(link) != 2):
|
||||
raise RuntimeError("Worker End is missing its paired Start context wire")
|
||||
paired_id, end_inputs = _stable_worker_binding(prompt, link[0], worker_item_identity)
|
||||
if (paired_id != end_id or context_values != [worker_item_identity]
|
||||
or end_inputs["worker_context_id"] != worker_context_id
|
||||
or end_inputs["worker_result_token"] != worker_result_token
|
||||
or end_inputs["worker_item_index"] != unwrap_single(worker_item_index)
|
||||
or end_inputs.get("master_callback_url", "") != master_callback_url):
|
||||
raise RuntimeError("Worker End identity/context does not match its current delivery")
|
||||
from ..image_list.runtime import finish_worker_item
|
||||
# On a miss this consumes Start's receipt; on a hit it validates
|
||||
# and drops the current staged tensors. Both paths enforce TTL.
|
||||
finish_worker_item(worker_result_token, worker_item_identity,
|
||||
unwrap_single(worker_item_index))
|
||||
context_id = worker_context_id
|
||||
else:
|
||||
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), {})
|
||||
expected = end_override_names(end_inputs)
|
||||
# Missing connected lazy values must fail instead of becoming an
|
||||
# implicit fallback to the original Start value.
|
||||
supplied = {name: values.get(name) for name in expected}
|
||||
supplied.update(values)
|
||||
result = normalize_worker_result(images, masks, texts, supplied)
|
||||
if unwrap_single(image_list_local_execution):
|
||||
from ..image_list.runtime import set_master_result
|
||||
set_master_result(context_id, int(unwrap_single(worker_item_index)),
|
||||
str(worker_result_token), result)
|
||||
else:
|
||||
payload = encode_result(result)
|
||||
run_async_in_server_loop(
|
||||
submit_worker_result(str(master_callback_url), str(worker_result_token),
|
||||
context_id, int(unwrap_single(worker_item_index)), payload), timeout=60,
|
||||
)
|
||||
return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")],
|
||||
[result.get("values", {})])
|
||||
|
||||
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.
|
||||
execution_list = getattr(cls.hidden, "execution_list", None)
|
||||
# Local model outputs belong to this still-running outer prompt.
|
||||
# Add them to its existing tracker; never start/end another tracker
|
||||
# that could unpin a model still needed by the parent.
|
||||
model_callback = getattr(execution_list, "output_link_callback", None)
|
||||
result_lists = run_async_in_server_loop(run_image_list_job(
|
||||
context_id, model_callback=model_callback, include_values=True))
|
||||
except ImageListCancelledError:
|
||||
import comfy.model_management
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
images_out, masks_out, texts_out, values_out = result_lists
|
||||
if len({len(output) for output in result_lists}) != 1:
|
||||
raise RuntimeError("Image-list collector received mismatched image, mask, text or values result counts")
|
||||
return io.NodeOutput(images_out, masks_out, texts_out, values_out)
|
||||
|
||||
|
||||
class DistributedImageListGetValue(io.ComfyNode):
|
||||
"""Ordinary scalar execution lets core map over End's ordered values list."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="DistributedImageListGetValue",
|
||||
display_name="Distributed Image List Get Value",
|
||||
category="image/distributed",
|
||||
inputs=[io.Custom("DISTRIBUTED_IMAGE_LIST_VALUES").Input("values"),
|
||||
io.Combo.Input("value_name", options=list(VALUE_NAMES), default="value_1")],
|
||||
outputs=[io.AnyType.Output(id="value", display_name="value")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, values, value_name="value_1"):
|
||||
if value_name not in VALUE_NAMES:
|
||||
raise ValueError(f"Unknown image-list value slot {value_name!r}")
|
||||
if not isinstance(values, dict):
|
||||
raise TypeError("Image List Get Value expects one values mapping from Image List End")
|
||||
if value_name not in values:
|
||||
raise ValueError(f"Image List Get Value: {value_name} is absent from this item's values")
|
||||
return io.NodeOutput(values[value_name])
|
||||
|
||||
|
||||
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}")
|
||||
@@ -1,46 +0,0 @@
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from ..utils.config import get_worker_timeout_seconds
|
||||
from ..utils.constants import HEARTBEAT_INTERVAL
|
||||
from ..utils.logging import log
|
||||
|
||||
|
||||
async def collect_worker_queue_results(
|
||||
*,
|
||||
prompt_server: Any,
|
||||
multi_job_id: str,
|
||||
expected_workers: set[str],
|
||||
on_result: Callable[[dict[str, Any]], None],
|
||||
timeout_log_prefix: str,
|
||||
throw_if_interrupted: Callable[[], None],
|
||||
) -> set[str]:
|
||||
workers_done: set[str] = set()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
|
||||
last_activity = time.time()
|
||||
|
||||
while len(workers_done) < len(expected_workers):
|
||||
throw_if_interrupted()
|
||||
try:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
queue = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
if (time.time() - last_activity) < base_timeout:
|
||||
continue
|
||||
missing_workers = sorted(expected_workers - workers_done)
|
||||
log(f"{timeout_log_prefix}{missing_workers}")
|
||||
break
|
||||
|
||||
on_result(result)
|
||||
worker_id_value = str(result.get("worker_id", ""))
|
||||
is_last = bool(result.get("is_last", False))
|
||||
last_activity = time.time()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
if is_last and worker_id_value in expected_workers:
|
||||
workers_done.add(worker_id_value)
|
||||
|
||||
return workers_done
|
||||
@@ -1,16 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
|
||||
def get_prompt_server_instance() -> Any:
|
||||
import server as _server
|
||||
|
||||
return _server.PromptServer.instance
|
||||
|
||||
|
||||
def throw_if_processing_interrupted() -> None:
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
model_management.throw_exception_if_processing_interrupted()
|
||||
except Exception:
|
||||
return
|
||||
+50
-75
@@ -1,10 +1,7 @@
|
||||
import torch
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.worker_ids import coerce_enabled_worker_ids, parse_worker_index, worker_value_key
|
||||
from .hidden_inputs import build_worker_identity_hidden_inputs
|
||||
|
||||
|
||||
def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]:
|
||||
@@ -31,7 +28,7 @@ class DistributedSeed:
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedSeed"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"seed": ("INT", {
|
||||
@@ -42,7 +39,8 @@ class DistributedSeed:
|
||||
}),
|
||||
},
|
||||
"hidden": {
|
||||
**build_worker_identity_hidden_inputs(),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -51,29 +49,30 @@ class DistributedSeed:
|
||||
FUNCTION = "distribute"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def distribute(
|
||||
self,
|
||||
seed: int,
|
||||
is_worker: bool = False,
|
||||
worker_id: str = "",
|
||||
enabled_worker_ids: str = "[]",
|
||||
) -> tuple[int]:
|
||||
def distribute(self, seed, is_worker=False, worker_id=""):
|
||||
if not is_worker:
|
||||
# Master node: pass through original values
|
||||
debug_log(f"Distributor - Master: seed={seed}")
|
||||
return (seed,)
|
||||
else:
|
||||
enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
worker_index = parse_worker_index(worker_id, enabled_workers)
|
||||
if worker_index is not None:
|
||||
# Worker node: apply offset based on worker index
|
||||
# Find worker index from enabled_worker_ids
|
||||
try:
|
||||
# Worker IDs are passed as "worker_0", "worker_1", etc.
|
||||
if worker_id.startswith("worker_"):
|
||||
worker_index = int(worker_id.split("_")[1])
|
||||
else:
|
||||
# Fallback: try to parse as direct index
|
||||
worker_index = int(worker_id)
|
||||
|
||||
offset = worker_index + 1
|
||||
new_seed = seed + offset
|
||||
debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
|
||||
return (new_seed,)
|
||||
|
||||
debug_log(f"Distributor - Error parsing worker_id '{worker_id}': no worker index resolved")
|
||||
# Fallback: return original seed
|
||||
return (seed,)
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
|
||||
# Fallback: return original seed
|
||||
return (seed,)
|
||||
|
||||
|
||||
# Define ByPassTypeTuple for flexible return types
|
||||
@@ -93,14 +92,15 @@ class DistributedValue:
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedValue"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"default_value": ("STRING", {"default": ""}),
|
||||
"worker_values": ("STRING", {"default": "{}"}),
|
||||
},
|
||||
"hidden": {
|
||||
**build_worker_identity_hidden_inputs(),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -110,7 +110,7 @@ class DistributedValue:
|
||||
CATEGORY = "utils"
|
||||
|
||||
@staticmethod
|
||||
def _coerce(value: Any, value_type: str) -> Any:
|
||||
def _coerce(value, value_type):
|
||||
"""Convert a string value to the requested type."""
|
||||
if value_type == "INT":
|
||||
return int(float(value))
|
||||
@@ -119,71 +119,51 @@ class DistributedValue:
|
||||
return value # STRING and COMBO stay as strings
|
||||
|
||||
@staticmethod
|
||||
def _coerce_safe(value: Any, value_type: str) -> Any:
|
||||
def _coerce_safe(value, value_type):
|
||||
"""Best-effort coercion with graceful fallback to original value."""
|
||||
try:
|
||||
return DistributedValue._coerce(value, value_type)
|
||||
except (TypeError, ValueError):
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _infer_value_type(value: Any) -> str:
|
||||
"""Infer coercion type from the provided default value."""
|
||||
if isinstance(value, bool):
|
||||
return "STRING"
|
||||
if isinstance(value, int):
|
||||
return "INT"
|
||||
if isinstance(value, float):
|
||||
return "FLOAT"
|
||||
return "STRING"
|
||||
|
||||
def distribute(
|
||||
self,
|
||||
default_value: Any,
|
||||
worker_values: str | dict[str, Any] = "{}",
|
||||
is_worker: bool = False,
|
||||
worker_id: str = "",
|
||||
enabled_worker_ids: str = "[]",
|
||||
) -> tuple[Any]:
|
||||
def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
|
||||
values = {}
|
||||
value_type = "STRING"
|
||||
|
||||
try:
|
||||
raw_values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
|
||||
values = dict(raw_values) if isinstance(raw_values, dict) else {}
|
||||
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
|
||||
if not isinstance(values, dict):
|
||||
values = {}
|
||||
except json.JSONDecodeError as e:
|
||||
debug_log(f"DistributedValue - Error parsing worker_values: {e}")
|
||||
values = {}
|
||||
|
||||
inferred_type = self._infer_value_type(default_value)
|
||||
value_type = values.get("_type", inferred_type)
|
||||
if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}:
|
||||
value_type = inferred_type
|
||||
value_type = values.get("_type", "STRING")
|
||||
coerced_default = self._coerce_safe(default_value, value_type)
|
||||
|
||||
if not is_worker:
|
||||
debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
|
||||
direct_key = str(worker_id).strip()
|
||||
lookup_key = direct_key if direct_key in values else worker_value_key(worker_id, enabled_workers)
|
||||
raw = values.get(lookup_key)
|
||||
if raw is not None and raw != "":
|
||||
try:
|
||||
try:
|
||||
if worker_id.startswith("worker_"):
|
||||
idx = int(worker_id.split("_")[1])
|
||||
else:
|
||||
idx = int(worker_id)
|
||||
key = str(idx + 1) # worker_0 → key "1" (1-indexed)
|
||||
raw = values.get(key, "")
|
||||
if raw:
|
||||
coerced = self._coerce(raw, value_type)
|
||||
debug_log(f"DistributedValue - Worker key {lookup_key}: returning '{coerced}'")
|
||||
debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
|
||||
return (coerced,)
|
||||
except (TypeError, ValueError) as e:
|
||||
debug_log(f"DistributedValue - Error coercing worker value for key {lookup_key}: {e}")
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"DistributedValue - Error: {e}")
|
||||
debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
class DistributedModelName:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedModelName"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"default": ""}),
|
||||
@@ -200,7 +180,7 @@ class DistributedModelName:
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "utils"
|
||||
|
||||
def _stringify(self, value: Any) -> str:
|
||||
def _stringify(self, value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (int, float, bool)):
|
||||
@@ -210,7 +190,7 @@ class DistributedModelName:
|
||||
except Exception:
|
||||
return str(value)
|
||||
|
||||
def _update_workflow(self, extra_pnginfo: Any, unique_id: Any, values: list[str]) -> None:
|
||||
def _update_workflow(self, extra_pnginfo, unique_id, values):
|
||||
if not extra_pnginfo:
|
||||
return
|
||||
info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
|
||||
@@ -228,12 +208,7 @@ class DistributedModelName:
|
||||
if node:
|
||||
node["widgets_values"] = [values]
|
||||
|
||||
def log_input(
|
||||
self,
|
||||
text: Any,
|
||||
unique_id: Any = None,
|
||||
extra_pnginfo: Any = None,
|
||||
) -> dict[str, Any]:
|
||||
def log_input(self, text, unique_id=None, extra_pnginfo=None):
|
||||
values = []
|
||||
if isinstance(text, list):
|
||||
for val in text:
|
||||
@@ -249,7 +224,7 @@ class DistributedModelName:
|
||||
return {"ui": {"text": values}, "result": (values,)}
|
||||
|
||||
class ByPassTypeTuple(tuple):
|
||||
def __getitem__(self, index: int) -> Any:
|
||||
def __getitem__(self, index):
|
||||
if index > 0:
|
||||
index = 0
|
||||
item = super().__getitem__(index)
|
||||
@@ -259,7 +234,7 @@ class ByPassTypeTuple(tuple):
|
||||
|
||||
class ImageBatchDivider:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["ImageBatchDivider"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
@@ -280,7 +255,7 @@ class ImageBatchDivider:
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "image"
|
||||
|
||||
def divide_batch(self, images: torch.Tensor, divide_by: int) -> tuple[torch.Tensor, ...]:
|
||||
def divide_batch(self, images, divide_by):
|
||||
total_splits = max(1, min(int(divide_by), 10))
|
||||
total_frames = images.shape[0]
|
||||
empty_tensor = images[:0]
|
||||
@@ -294,10 +269,10 @@ 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(cls: type["AudioBatchDivider"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
@@ -307,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"
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -318,7 +293,7 @@ class AudioBatchDivider:
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "audio"
|
||||
|
||||
def divide_audio(self, audio: dict[str, Any], divide_by: int) -> tuple[dict[str, Any], ...]:
|
||||
def divide_audio(self, audio, divide_by):
|
||||
import torch
|
||||
|
||||
waveform = audio.get("waveform")
|
||||
@@ -358,7 +333,7 @@ class DistributedEmptyImage:
|
||||
"""Produces an empty IMAGE batch used when the master delegates all work."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls: type["DistributedEmptyImage"]) -> dict[str, Any]:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
|
||||
@@ -371,7 +346,7 @@ class DistributedEmptyImage:
|
||||
FUNCTION = "create"
|
||||
CATEGORY = "image"
|
||||
|
||||
def create(self, height: int, width: int, channels: int) -> tuple[torch.Tensor]:
|
||||
def create(self, height, width, channels):
|
||||
import torch
|
||||
|
||||
shape = (0, height, width, channels)
|
||||
|
||||
+215
@@ -0,0 +1,215 @@
|
||||
"""Native V3 schemas with the existing execution algorithms kept intact.
|
||||
|
||||
Runtime objects are private, per-execution helpers, not registered V1 nodes.
|
||||
This avoids sharing mutable instance state through V3's sanitized class clones.
|
||||
"""
|
||||
import comfy.samplers
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from ..utils import transfer_codec as _transfer_codec
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
from .image_list import DistributedImageListEnd, DistributedImageListStart, DistributedImageListGetValue
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
|
||||
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
|
||||
max=1125899906842624, force_input=False)],
|
||||
outputs=[io.Int.Output(display_name='seed')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed, is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedValue(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedValue', display_name='Distributed Value', category='utils',
|
||||
inputs=[io.String.Input('default_value', default=''),
|
||||
io.String.Input('worker_values', default='{}')],
|
||||
outputs=[io.AnyType.Output(display_name='value')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedValue().distribute(
|
||||
default_value, worker_values, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedModelName(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
|
||||
inputs=[io.String.Input('text', default='')],
|
||||
outputs=[io.AnyType.Output(display_name='output')],
|
||||
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text):
|
||||
result = _utilities.DistributedModelName().log_input(
|
||||
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
|
||||
return io.NodeOutput(*result['result'], ui=result['ui'])
|
||||
|
||||
|
||||
class ImageBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
|
||||
inputs=[io.Image.Input('images'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of parts to divide the batch into')],
|
||||
# The existing frontend still displays only divide_by sockets.
|
||||
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, divide_by):
|
||||
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
|
||||
|
||||
|
||||
class AudioBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
|
||||
inputs=[io.Audio.Input('audio'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of sequential time segments to create')],
|
||||
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, divide_by):
|
||||
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
|
||||
|
||||
|
||||
class DistributedEmptyImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
|
||||
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('width', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
|
||||
outputs=[io.Image.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, height, width, channels):
|
||||
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
|
||||
|
||||
|
||||
class DistributedCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False,
|
||||
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
|
||||
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
|
||||
io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
|
||||
default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,
|
||||
tooltip=_transfer_codec.TRANSFER_FORMAT_TOOLTIP),
|
||||
io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
|
||||
min=1, max=100, step=1, optional=True,
|
||||
tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
|
||||
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
|
||||
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
|
||||
worker_id='', pass_through=False, delegate_only=False,
|
||||
transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
|
||||
transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
|
||||
return io.NodeOutput(*_CollectorRuntime().run(
|
||||
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
|
||||
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
|
||||
worker_batch_size=worker_batch_size, worker_id=worker_id,
|
||||
pass_through=pass_through, delegate_only=delegate_only,
|
||||
transfer_format=transfer_format, transfer_quality=transfer_quality))
|
||||
|
||||
|
||||
class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='UltimateSDUpscaleDistributed',
|
||||
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
|
||||
inputs=[
|
||||
io.Image.Input('upscaled_image'), io.Model.Input('model'),
|
||||
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
|
||||
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Int.Input('steps', default=20, min=1, max=10000),
|
||||
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
|
||||
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
|
||||
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
|
||||
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
|
||||
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('padding', default=32, min=0, max=256, step=8),
|
||||
io.Int.Input('mask_blur', default=8, min=0, max=256),
|
||||
io.Boolean.Input('force_uniform_tiles', default=True),
|
||||
io.Boolean.Input('tiled_decode', default=False),
|
||||
], outputs=[io.Image.Output()], accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
return _UpscaleRuntime.IS_CHANGED(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
|
||||
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
|
||||
return io.NodeOutput(*_UpscaleRuntime().run(
|
||||
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
|
||||
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
|
||||
|
||||
|
||||
class DistributedWorkerSink(io.ComfyNode):
|
||||
"""Internal output node that terminates pruned worker prompts.
|
||||
|
||||
Worker prompts drop everything downstream of the distributed node, which
|
||||
leaves no output node for ComfyUI to execute towards. This sink provides
|
||||
one without writing preview files (PreviewImage wrote every frame to the
|
||||
worker's temp directory). Hidden from search/menus; inserted by the master.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedWorkerSink', display_name='Distributed Worker Sink (internal)', category='utils',
|
||||
inputs=[io.AnyType.Input('value')],
|
||||
is_output_node=True, is_dev_only=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, value):
|
||||
return io.NodeOutput()
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
|
||||
DistributedWorkerSink, DistributedImageListStart, DistributedImageListEnd, DistributedImageListGetValue]
|
||||
+2
-5
@@ -1,12 +1,9 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.4.1"
|
||||
version = "1.5.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = [
|
||||
"aiohttp>=3.9,<4",
|
||||
"Pillow>=10,<12",
|
||||
]
|
||||
dependencies = []
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Internal extension lifecycle support (not a node provider)."""
|
||||
@@ -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}')
|
||||
@@ -0,0 +1,266 @@
|
||||
# Standalone `comfy-media-v1` endpoint bridge
|
||||
|
||||
This directory is independently installable. It registers two internal Comfy V3
|
||||
nodes, `MediaImportV1` and `MediaExportV1`, and provides the asynchronous HTTP
|
||||
executor `execute_media_request`. It does **not** start ComfyUI, provision an
|
||||
endpoint, install models, or provide a provider worker/supervisor. The endpoint
|
||||
needs this bridge, not the full ComfyUI-Distributed extension.
|
||||
|
||||
The bridge's generic contract supports `IMAGE`, `AUDIO`, `MASK`, `STRING`, and
|
||||
native ComfyUI `VIDEO`. **Distributed Collector** (`DistributedCollector`) uses
|
||||
only the image/audio outputs and combines their returned lists into ordinary
|
||||
IMAGE/AUDIO outputs. It has no MASK/STRING/VIDEO outputs, and the separate
|
||||
`DistributedMediaCollector` node has been removed. For Collector setup, ordering,
|
||||
and limits, see [serverless images and audio](../docs/serverless-media.md).
|
||||
|
||||
## Install in the endpoint image
|
||||
|
||||
1. Copy this complete directory to `<ComfyUI>/custom_nodes/serverless_bridge`.
|
||||
Use a compatible ComfyUI with `comfy_api.v0_0_2` and the native VIDEO classes
|
||||
used by [video_codec.py](video_codec.py). These exact-class/private-state
|
||||
adapters are tied to the supported core implementation, not arbitrary VIDEO
|
||||
plugins or future core layouts.
|
||||
2. Restart ComfyUI so `comfy_entrypoint()` registers the two nodes. Keep the
|
||||
workflow's own models and custom nodes installed separately.
|
||||
3. Make `<ComfyUI>/custom_nodes` importable by the handler process (for example,
|
||||
via its `PYTHONPATH`). The handler uses `aiohttp`; the nodes/codecs use
|
||||
ComfyUI's PyTorch/native media dependencies. The handler module itself does
|
||||
not import ComfyUI, PyTorch, Runpod, or Distributed bootstrap code.
|
||||
4. Give the handler filesystem access to **the same input directory** used by
|
||||
that ComfyUI instance, plus its trusted HTTP URL. These are operator settings,
|
||||
never request-controlled paths/URLs. Staging/cleanup requires POSIX directory
|
||||
descriptor and no-follow file operations.
|
||||
|
||||
Installing the nodes alone is insufficient: the endpoint must route media-v1
|
||||
inputs through this executor and implement the uncertainty policy below. A
|
||||
legacy image handler cannot interpret the typed envelope.
|
||||
|
||||
## Callable integration
|
||||
|
||||
```python
|
||||
from serverless_bridge.handler import execute_media_request
|
||||
|
||||
|
||||
async def execute_endpoint_input(
|
||||
input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
download=None, publish=None,
|
||||
):
|
||||
# Invoke only under a supervisor implementing the uncertainty policy below.
|
||||
# Exceptions intentionally propagate; this is NOT a provider worker loop.
|
||||
return await execute_media_request(
|
||||
input_data,
|
||||
comfy_url=comfy_url,
|
||||
input_dir=input_dir,
|
||||
timeout_seconds=1800,
|
||||
session=None,
|
||||
download=download,
|
||||
publish=publish,
|
||||
temp_dir=temp_dir,
|
||||
)
|
||||
```
|
||||
|
||||
Pass the provider job's **`input` object**, not the whole job. Return the resulting
|
||||
`{"media": ...}` object as provider output; Runpod exposes it as `output.media`.
|
||||
Do not add another `output` wrapper inside it. `timeout_seconds` must be finite
|
||||
and positive. If supplied, `session` is a caller-owned `aiohttp.ClientSession`
|
||||
for local Comfy HTTP and remains open; otherwise the executor creates/closes its
|
||||
own session. Input dictionaries are copied before graph rewriting.
|
||||
|
||||
**This callable is not a turnkey Runpod handler.** An outer SDK that catches every
|
||||
exception and immediately reuses the worker is unsafe. The supervisor must
|
||||
recognize `serverless_bridge.handler.MediaExecutionUncertain`, stop intake, and
|
||||
quarantine/reconcile or terminate the owning environment before reuse. Merely
|
||||
returning `{"error": ...}` or marking a provider job `FAILED` is not that policy.
|
||||
|
||||
## Request and response
|
||||
|
||||
A minimal generic bridge text request builder (no files or model required by this
|
||||
graph). This illustrates the bridge contract, not a Collector text output:
|
||||
|
||||
```python
|
||||
import json
|
||||
import uuid
|
||||
|
||||
|
||||
def make_text_request(text="hello"):
|
||||
request_id = uuid.uuid4().hex
|
||||
binding = {
|
||||
"binding_id": "text-binding",
|
||||
"collector_id": "collector-1",
|
||||
"job_id": "collector-job-1",
|
||||
"participant_id": "serverless:example:0",
|
||||
"node_id": "export-1",
|
||||
"socket": "texts",
|
||||
"kind": "STRING",
|
||||
}
|
||||
return {
|
||||
"workflow": {
|
||||
"text-1": {"class_type": "PrimitiveString", "inputs": {"value": text}},
|
||||
"export-1": {
|
||||
"class_type": "MediaExportV1",
|
||||
"inputs": {
|
||||
"request_id": request_id,
|
||||
"binding_json": json.dumps(binding),
|
||||
"texts": ["text-1", 0],
|
||||
},
|
||||
},
|
||||
},
|
||||
"images": [],
|
||||
"media": {
|
||||
"version": 1,
|
||||
"profile": "comfy-media-v1",
|
||||
"request_id": request_id,
|
||||
"artifacts": [],
|
||||
"inputs": [],
|
||||
"outputs": [binding],
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
`workflow` is an API graph, not an editor workflow. Top-level input fields are
|
||||
exactly `workflow`, `media`, and optional `images`. `media` has exactly the six
|
||||
fields shown above. `media.outputs` declares one exporter per connected socket;
|
||||
its binding must match that node's `binding_json` and request ID exactly. Socket
|
||||
names/types are `images/IMAGE`, `audio/AUDIO`, `masks/MASK`, `texts/STRING`, and
|
||||
`video/VIDEO`. Use unique request IDs. Identifiers are 1–200 characters from
|
||||
`[A-Za-z0-9_.:/-]`; request scopes additionally reject empty, `.` or `..` slash
|
||||
components. The executor validates the graph and submits **once**, with exactly
|
||||
the declared exporter IDs as `partial_execution_targets`. Unrelated graph roots,
|
||||
cycles, undeclared exporters, or caller-supplied `MediaImportV1` nodes fail.
|
||||
|
||||
For typed inputs, `media.artifacts` contains records with exactly:
|
||||
|
||||
| Field | Value |
|
||||
| --- | --- |
|
||||
| `id`, `kind` | Unique artifact ID and one of the five native types |
|
||||
| `descriptor` | Kind-specific codec metadata; use the shared codec, not pickle |
|
||||
| `byte_length`, `sha256` | Exact payload byte count and lowercase SHA-256 hex |
|
||||
| `transport` | `{"type": "inline", "data_base64": "..."}` or the URL transport below |
|
||||
|
||||
Each `media.inputs` entry is
|
||||
`{"node_id": "consumer", "input_name": "socket", "artifact_id": "artifact-id"}`.
|
||||
It replaces an existing consumer input with a generated import link; every
|
||||
artifact must be used, and destinations cannot repeat. The executor verifies
|
||||
bytes/checksums, stages owned files, inserts imports, then prunes disconnected
|
||||
loader ancestors. Optional legacy `images` entries are `{ "name": "input.png",
|
||||
"image": "<base64>" }`, must match stock `LoadImage` references, and are rewritten
|
||||
to owned staging paths. They are not an arbitrary filename-upload interface.
|
||||
|
||||
The returned `media` object has exactly `version`, `profile`, `request_id`,
|
||||
`bindings`, `artifacts`, and `items`:
|
||||
|
||||
- `bindings`: every declared binding, unchanged, plus `count` (1–64).
|
||||
- `artifacts`: the same artifact schema above with inline or URL transport.
|
||||
- `items`: `{ "binding_id": "...", "item_index": 0, "artifact_id": "..." }`
|
||||
records covering every index exactly once, ordered by `item_index` when decoded.
|
||||
|
||||
All identity, completeness, type, size, and checksum checks must pass; do not
|
||||
consume a partial result. The internal history transport `comfy-file` is local
|
||||
only and must not be returned to the master. The executor reads validated
|
||||
`media_exports` history records and bounded `/view` payloads before constructing
|
||||
the external envelope. [contract.py](contract.py) is the schema authority.
|
||||
|
||||
## Optional image return encoding
|
||||
|
||||
IMAGE `MediaExportV1` nodes may also contain literal `transfer_format` and
|
||||
`transfer_quality` controls. The **bridge-level defaults** are `exact` and `90`;
|
||||
generic callers that omit the controls retain that representation. Other media
|
||||
exports must not carry image controls. Callers should not alter binding identity
|
||||
metadata.
|
||||
|
||||
The **Collector default is different**: its saved widget value remains
|
||||
`legacy_png` with quality `90`. For a `comfy-media-v1` pool, the master compiler
|
||||
maps that value to an explicit `png` export, so normal serverless image returns
|
||||
use compressed 8-bit PNG. `exact` is opt-in at the Collector. Other selected
|
||||
formats are unchanged. Direct workers also support opt-in exact IMAGE transport;
|
||||
an older master rejection is an error, not an 8-bit retry. A Collector requesting
|
||||
`exact` cannot dispatch to a `legacy-image-v1` pool.
|
||||
|
||||
Formats: `exact`, `legacy_png`, `png`, `ffv1`, `jpeg`, `webp`, `h264`,
|
||||
`h265` (**experimental**), `av1`.
|
||||
`exact` preserves float tensor bits. The image/video codecs quantize RGB/grayscale
|
||||
to 8-bit; PNG/FFV1 are lossless at that precision, while the other codecs are
|
||||
lossy. Quality is an integer 1–100 and affects only lossy formats. IMAGE batch
|
||||
order, shape and channel count are retained, and alpha is not discarded.
|
||||
Native VIDEO/MASK/AUDIO/STRING bridge values still use their original exact codecs;
|
||||
the image-format controls do not add those sockets to the Collector.
|
||||
|
||||
Copy/update the complete bridge directory on both sides, including
|
||||
`image_transfer.py` and `image_codec.py`. Pillow and the selected CPU/software
|
||||
PyAV codec must be available; missing encoders fail without automatic fallback.
|
||||
Compressed IMAGE artifacts use `image-transfer-v1`; decode through the shared
|
||||
`codec.decode_payload`, which enforces payload and decoded-shape limits.
|
||||
See [image transfer choices](../docs/serverless-media.md#image-return-format-and-quality).
|
||||
|
||||
## Transport budgets and ownership
|
||||
|
||||
- Contract maxima: 256 MiB per artifact, 512 MiB aggregate artifact bytes,
|
||||
128 artifacts, 64 items per binding. Request input totals include decoded
|
||||
legacy-image bytes. Decoders validate tensor/VIDEO sizes and layouts before
|
||||
allocating native storage; these are payload limits, **not peak-RSS guarantees**.
|
||||
- The built-in master currently emits inline inputs only. Its complete Runpod
|
||||
request JSON is capped at **10,000,000 bytes**, including base64, graph and
|
||||
policy. The larger contract limits do not bypass this cap; there is no
|
||||
built-in large-input/object-storage uploader.
|
||||
- The executor's result JSON is capped at **10,000,000 bytes**. Inline results
|
||||
are the default. When needed, `publish` is called for selected larger
|
||||
artifacts to fit that budget; supplying it does not force all results to URLs.
|
||||
Without a publisher, an oversized response fails instead of being truncated.
|
||||
- Optional **async** `download(url, *, max_bytes)` returns `bytes`. The executor
|
||||
supplies the declared artifact length and checks exact length and SHA-256.
|
||||
The callback must enforce public HTTPS, DNS/IP restrictions, redirect checks,
|
||||
and bounded streaming **before allocation**; syntactic contract validation
|
||||
does not provide SSRF protection. No downloader is installed by default.
|
||||
- Optional **async** `publish(metadata, payload)` uploads the exact `bytes` and
|
||||
returns exactly `{"type": "url", "url": "https://...", "expires_at": "..."}`.
|
||||
`metadata` is the artifact record without `transport`. Expiry must be a future,
|
||||
timezone-aware ISO timestamp; URLs use HTTPS port 443, no userinfo/fragment,
|
||||
and at most 8192 characters. Keep objects available long enough for collection.
|
||||
The callback/operator owns upload credentials, object expiry and cleanup,
|
||||
including orphan uploads after failure. There is no bridge storage backend.
|
||||
- The master uses a separate bounded public-HTTPS downloader without provider
|
||||
credentials. Never send provider credentials to artifact hosts or log signed
|
||||
URLs. Endpoint callbacks must enforce equivalent protections themselves.
|
||||
|
||||
Callbacks must cooperate with async cancellation, finish/abort their own I/O,
|
||||
and not detach readers/uploads that outlive the callback. Synchronous code that
|
||||
blocks the event loop cannot be hard-timed-out by this executor.
|
||||
|
||||
## Cancellation, quarantine, and file retention
|
||||
|
||||
`timeout_seconds` covers staging, submission, polling, downloads and publication.
|
||||
On failure/cancellation, there is a separate **10-second callback-drain grace**
|
||||
and, when needed, a **10-second prompt-reconciliation grace**. These are recovery
|
||||
allowances, not evidence that execution stopped. Ambiguous `/prompt` submission
|
||||
is never retried. An identified pending prompt may be deleted from `/queue`,
|
||||
with exact read-back; running prompts are never sent `/interrupt`. The pinned
|
||||
core route has a snapshot/global-interrupt race even when given a prompt ID.
|
||||
|
||||
`MediaExecutionUncertain` means prompt acceptance/termination **or callback
|
||||
teardown** is unconfirmed. Even successful terminal Comfy history is insufficient
|
||||
if an async downloader/publisher ignores cancellation. The bridge retains input
|
||||
and exporter files in that case; late callback exit does not automatically
|
||||
remove retained files. Before accepting another job on that environment:
|
||||
|
||||
1. Quarantine it and stop new intake; do not automatically resubmit the request.
|
||||
2. Reconcile the local prompt and all readers, or terminate the owning execution
|
||||
environment when their termination cannot be established.
|
||||
3. Only then clean retained files and allow reuse/replacement. A provider terminal
|
||||
job status alone cannot certify a surviving local Comfy process or callback.
|
||||
|
||||
Staged inputs use random exclusive files under
|
||||
`input_dir/distributed-media-v1/<sha256(request_id)>/`. Cleanup removes only files
|
||||
created by that invocation, after confirmed terminal state (or before submission)
|
||||
and reader teardown; it never recursively deletes a shared namespace.
|
||||
|
||||
`temp_dir=None` leaves exporter retention to the Comfy temp-directory/environment
|
||||
lifecycle. Optionally pass the **trusted local Comfy temp root** to enable
|
||||
best-effort cleanup of validated, invocation-owned exporter UUID `.bin` files
|
||||
under the corresponding hashed namespace, only after terminal history and
|
||||
collection/publisher teardown. Symlinked, hard-linked, replaced, missing, or
|
||||
unidentifiable files are not claimed; directories and all unidentified leftovers
|
||||
remain lifecycle-owned. This is not a temp-directory sweep. Independently decoded
|
||||
VIDEO objects do not depend on either staging or exporter files.
|
||||
|
||||
This guide describes the source contract, not a test result. No live provider/GPU
|
||||
validation is claimed, and this directory is not a turnkey provider deployment.
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Standalone typed-media bridge; copy this directory into endpoint custom_nodes.
|
||||
|
||||
Importing codecs does not load ComfyUI, providers or Distributed bootstrap.
|
||||
"""
|
||||
|
||||
async def comfy_entrypoint():
|
||||
from comfy_api.v0_0_2 import ComfyExtension
|
||||
from .nodes import MediaImportV1, MediaExportV1
|
||||
|
||||
class MediaBridgeExtension(ComfyExtension):
|
||||
async def get_node_list(self):
|
||||
return [MediaImportV1, MediaExportV1]
|
||||
|
||||
return MediaBridgeExtension()
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Lossless, bounded CPU payloads shared with the standalone endpoint bridge.
|
||||
|
||||
Descriptors contain only kind-specific metadata. Tensor bytes are contiguous,
|
||||
little-endian floating-point values, never a pickle or an image/audio container.
|
||||
The optional byte budget can lower, but never raise, MAX_PAYLOAD_BYTES.
|
||||
"""
|
||||
|
||||
import ctypes
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
MAX_PAYLOAD_BYTES = 256 * 1024 * 1024
|
||||
|
||||
_DTYPES = {
|
||||
"float16": (torch.float16, 2),
|
||||
"bfloat16": (torch.bfloat16, 2),
|
||||
"float32": (torch.float32, 4),
|
||||
"float64": (torch.float64, 8),
|
||||
}
|
||||
_DTYPE_NAMES = {dtype: name for name, (dtype, _) in _DTYPES.items()}
|
||||
_LAYOUTS = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}
|
||||
|
||||
|
||||
def _validate_kind(kind):
|
||||
if type(kind) is not str or kind not in ("IMAGE", "MASK", "STRING", "AUDIO"):
|
||||
raise ValueError("Unsupported payload kind")
|
||||
|
||||
|
||||
def _byte_limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _sample_rate(value):
|
||||
if type(value) is not int or value <= 0:
|
||||
raise ValueError("AUDIO sample_rate must be a positive integer")
|
||||
return value
|
||||
|
||||
|
||||
def _tensor_size(kind, shape, itemsize, limit):
|
||||
if type(shape) is not list or len(shape) != len(_LAYOUTS[kind]):
|
||||
raise ValueError("Tensor shape has an invalid rank")
|
||||
if any(type(dimension) is not int or dimension <= 0 for dimension in shape):
|
||||
raise ValueError("Tensor shape dimensions must be positive integers")
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("IMAGE shape must have 1, 3 or 4 channels")
|
||||
size = itemsize
|
||||
for dimension in shape:
|
||||
# Divide before multiplying, including for arbitrary-size Python ints.
|
||||
if dimension > limit // size:
|
||||
raise ValueError("Tensor payload exceeds byte limit")
|
||||
size *= dimension
|
||||
return size
|
||||
|
||||
|
||||
def _native_little_endian(payload, itemsize):
|
||||
"""Convert native to wire byte order, or vice versa (an involution)."""
|
||||
if sys.byteorder == "little":
|
||||
return payload
|
||||
swapped = bytearray(len(payload))
|
||||
for offset in range(itemsize):
|
||||
swapped[offset::itemsize] = payload[itemsize - 1 - offset::itemsize]
|
||||
return bytes(swapped)
|
||||
|
||||
|
||||
def _encode_string(value, limit):
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("STRING value must be text")
|
||||
# Check the exact UTF-8 size before allocating the encoded payload.
|
||||
if len(value) > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
size = 0
|
||||
for character in value:
|
||||
codepoint = ord(character)
|
||||
if 0xD800 <= codepoint <= 0xDFFF:
|
||||
raise ValueError("STRING value must be valid UTF-8 text")
|
||||
size += 1 + (codepoint > 0x7F) + (codepoint > 0x7FF) + (codepoint > 0xFFFF)
|
||||
if size > limit:
|
||||
raise ValueError("STRING payload exceeds byte limit")
|
||||
return {"encoding": "utf8-v1"}, value.encode("utf-8")
|
||||
|
||||
|
||||
def encode_payload(kind: str, value, *, max_bytes: int = MAX_PAYLOAD_BYTES) -> tuple[dict, bytes]:
|
||||
"""Encode a typed value without dtype conversion, normalization or GPU use.
|
||||
|
||||
IMAGE/MASK accept dense CPU tensors; AUDIO accepts a waveform CPU tensor
|
||||
plus a positive integer sample_rate. STRING accepts exact Unicode text.
|
||||
Invalid values and exceeded budgets raise nonsensitive ValueError messages.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if kind == "STRING":
|
||||
return _encode_string(value, limit)
|
||||
|
||||
sample_rate = None
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(value, dict) or set(value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
sample_rate = _sample_rate(value["sample_rate"])
|
||||
value = value["waveform"]
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
if value.device.type != "cpu":
|
||||
raise ValueError("Tensor value must be on CPU")
|
||||
if value.layout != torch.strided or value.is_nested:
|
||||
raise ValueError("Tensor value must have a dense strided layout")
|
||||
if value.dtype not in _DTYPE_NAMES:
|
||||
raise ValueError("Unsupported tensor dtype")
|
||||
|
||||
dtype_name = _DTYPE_NAMES[value.dtype]
|
||||
_, itemsize = _DTYPES[dtype_name]
|
||||
shape = list(value.shape)
|
||||
size = _tensor_size(kind, shape, itemsize, limit)
|
||||
tensor = value.detach().resolve_neg().contiguous()
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor values must be finite")
|
||||
# A uint16 view preserves BF16 bits without conversion through float32.
|
||||
if tensor.dtype == torch.bfloat16:
|
||||
tensor = tensor.view(torch.uint16)
|
||||
# The live, contiguous CPU tensor owns this pointer for the entire copy.
|
||||
payload = _native_little_endian(ctypes.string_at(tensor.data_ptr(), size), itemsize)
|
||||
descriptor = {
|
||||
"encoding": "tensor-le-v1", "dtype": dtype_name,
|
||||
"shape": shape, "layout": _LAYOUTS[kind],
|
||||
}
|
||||
if kind == "AUDIO":
|
||||
descriptor["sample_rate"] = sample_rate
|
||||
return descriptor, payload
|
||||
|
||||
|
||||
def decode_payload(kind: str, descriptor: dict, payload: bytes, *, max_bytes: int = MAX_PAYLOAD_BYTES):
|
||||
"""Decode validated wire bytes into independent contiguous CPU tensors/text.
|
||||
|
||||
All metadata, shape products and byte lengths are checked before tensor
|
||||
allocation. No GPU device discovery, transfers or object deserialization
|
||||
occur. AUDIO returns {'waveform': tensor, 'sample_rate': int}.
|
||||
"""
|
||||
_validate_kind(kind)
|
||||
limit = _byte_limit(max_bytes)
|
||||
if not isinstance(descriptor, dict):
|
||||
raise ValueError("Payload descriptor must be a dictionary")
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("Payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("Payload exceeds byte limit")
|
||||
if kind == "IMAGE" and descriptor.get("encoding") == "image-transfer-v1":
|
||||
# Lazy and IMAGE-only: exact tensors and other kinds stay unchanged.
|
||||
from .image_codec import decode_image
|
||||
return decode_image(descriptor, payload, max_bytes=limit)
|
||||
if kind == "STRING":
|
||||
if set(descriptor) != {"encoding"} or descriptor["encoding"] != "utf8-v1":
|
||||
raise ValueError("Invalid STRING descriptor encoding or fields")
|
||||
try:
|
||||
return payload.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
raise ValueError("STRING payload must be valid UTF-8 bytes") from None
|
||||
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
if set(descriptor) != fields:
|
||||
raise ValueError("Invalid tensor descriptor fields")
|
||||
if descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Unsupported tensor descriptor encoding")
|
||||
if descriptor["layout"] != _LAYOUTS[kind]:
|
||||
raise ValueError("Invalid tensor descriptor layout")
|
||||
dtype_name = descriptor["dtype"]
|
||||
if type(dtype_name) is not str or dtype_name not in _DTYPES:
|
||||
raise ValueError("Unsupported tensor descriptor dtype")
|
||||
dtype, itemsize = _DTYPES[dtype_name]
|
||||
size = _tensor_size(kind, descriptor["shape"], itemsize, limit)
|
||||
if len(payload) != size:
|
||||
raise ValueError("Tensor payload length does not match shape and dtype")
|
||||
sample_rate = _sample_rate(descriptor["sample_rate"]) if kind == "AUDIO" else None
|
||||
|
||||
# frombuffer is always CPU. Copy immutable wire bytes into a writable private
|
||||
# buffer, then clone so the returned tensor owns its storage independently.
|
||||
buffer = bytearray(_native_little_endian(payload, itemsize))
|
||||
storage_dtype = torch.uint16 if dtype == torch.bfloat16 else dtype
|
||||
tensor = torch.frombuffer(buffer, dtype=storage_dtype).view(dtype)
|
||||
if not torch.isfinite(tensor).all().item():
|
||||
raise ValueError("Tensor payload values must be finite")
|
||||
tensor = tensor.clone().reshape(descriptor["shape"])
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": tensor, "sample_rate": sample_rate}
|
||||
return tensor
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Versioned, bounded media envelopes shared by master and endpoint bridge.
|
||||
|
||||
No provider or ComfyUI bootstrap imports. Network I/O is supplied by the caller,
|
||||
which must enforce public-HTTPS resolution, redirect and streaming byte limits.
|
||||
"""
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timezone
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
PROFILE = "comfy-media-v1"
|
||||
VERSION = 1
|
||||
MAX_ARTIFACT_BYTES = 256 * 1024 * 1024
|
||||
MAX_TOTAL_BYTES = 512 * 1024 * 1024
|
||||
MAX_ARTIFACTS = 128
|
||||
MAX_ITEMS = 64
|
||||
MAX_DESCRIPTOR_BYTES = 128 * 1024
|
||||
SOCKET_KINDS = {"images": "IMAGE", "audio": "AUDIO", "masks": "MASK", "texts": "STRING", "video": "VIDEO"}
|
||||
BINDING_FIELDS = {"binding_id", "collector_id", "job_id", "participant_id", "node_id", "socket", "kind"}
|
||||
_ID = re.compile(r"[A-Za-z0-9_.:/-]{1,200}\Z")
|
||||
_HASH = re.compile(r"[a-f0-9]{64}\Z")
|
||||
|
||||
|
||||
def identifier(value):
|
||||
if not isinstance(value, str) or not _ID.fullmatch(value):
|
||||
raise ValueError("Invalid media identifier")
|
||||
return value
|
||||
|
||||
|
||||
def bounded_int(value, maximum, *, minimum=0):
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
raise ValueError("Invalid media count or size")
|
||||
return value
|
||||
|
||||
|
||||
def validate_json(value, *, max_bytes=MAX_DESCRIPTOR_BYTES):
|
||||
"""Bound metadata depth and size without accepting non-finite JSON numbers."""
|
||||
def walk(item, depth=0):
|
||||
if depth > 12:
|
||||
raise ValueError("Media metadata nesting exceeds limit")
|
||||
if isinstance(item, dict):
|
||||
if len(item) > 4096 or any(not isinstance(k, str) for k in item):
|
||||
raise ValueError("Invalid media metadata object")
|
||||
for child in item.values():
|
||||
walk(child, depth + 1)
|
||||
elif isinstance(item, list):
|
||||
if len(item) > 4096:
|
||||
raise ValueError("Media metadata list exceeds limit")
|
||||
for child in item:
|
||||
walk(child, depth + 1)
|
||||
elif item is not None and type(item) not in (str, int, float, bool):
|
||||
raise ValueError("Invalid media metadata value")
|
||||
walk(value)
|
||||
try:
|
||||
raw = json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")
|
||||
except (ValueError, TypeError, UnicodeError, RecursionError):
|
||||
raise ValueError("Invalid media JSON") from None
|
||||
if len(raw) > max_bytes:
|
||||
raise ValueError("Media JSON exceeds byte limit")
|
||||
return raw
|
||||
|
||||
|
||||
def validate_binding(binding):
|
||||
if not isinstance(binding, dict) or set(binding) != BINDING_FIELDS:
|
||||
raise ValueError("Invalid media binding")
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
identifier(binding[field])
|
||||
if SOCKET_KINDS.get(binding["socket"]) != binding["kind"]:
|
||||
raise ValueError("Media socket/type mismatch")
|
||||
return binding
|
||||
|
||||
|
||||
def validate_bindings(bindings):
|
||||
if not isinstance(bindings, list) or not 1 <= len(bindings) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media bindings count")
|
||||
result = {}
|
||||
destinations = set()
|
||||
for binding in bindings:
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
destination = (binding["job_id"], binding["participant_id"], binding["collector_id"], binding["socket"])
|
||||
if bid in result or destination in destinations:
|
||||
raise ValueError("Duplicate media binding")
|
||||
result[bid] = binding
|
||||
destinations.add(destination)
|
||||
return result
|
||||
|
||||
|
||||
def _validate_url_transport(transport):
|
||||
if set(transport) != {"type", "url", "expires_at"}:
|
||||
raise ValueError("Invalid media URL transport")
|
||||
url = transport["url"]
|
||||
if not isinstance(url, str) or len(url) > 8192:
|
||||
raise ValueError("Invalid media URL")
|
||||
try:
|
||||
parsed = urlsplit(url)
|
||||
if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
|
||||
raise ValueError()
|
||||
if parsed.port not in (None, 443):
|
||||
raise ValueError()
|
||||
expiry = datetime.fromisoformat(transport["expires_at"].replace("Z", "+00:00"))
|
||||
if expiry.tzinfo is None or expiry <= datetime.now(timezone.utc):
|
||||
raise ValueError()
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
raise ValueError("Invalid or expired media URL") from None
|
||||
|
||||
|
||||
def validate_artifact(artifact, *, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256", "transport"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid media artifact")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Unsupported media kind")
|
||||
if not isinstance(artifact["descriptor"], dict):
|
||||
raise ValueError("Invalid media descriptor")
|
||||
validate_json(artifact["descriptor"])
|
||||
size = bounded_int(artifact["byte_length"], max_bytes)
|
||||
digest = artifact["sha256"]
|
||||
if not isinstance(digest, str) or not _HASH.fullmatch(digest):
|
||||
raise ValueError("Invalid media checksum")
|
||||
transport = artifact["transport"]
|
||||
if not isinstance(transport, dict):
|
||||
raise ValueError("Invalid media transport")
|
||||
if transport.get("type") == "inline":
|
||||
if set(transport) != {"type", "data_base64"} or not isinstance(transport["data_base64"], str):
|
||||
raise ValueError("Invalid inline media transport")
|
||||
if len(transport["data_base64"]) != 4 * ((size + 2) // 3):
|
||||
raise ValueError("Inline media length mismatch")
|
||||
elif transport.get("type") == "url":
|
||||
_validate_url_transport(transport)
|
||||
else:
|
||||
raise ValueError("Unsupported media transport")
|
||||
return artifact
|
||||
|
||||
|
||||
def artifact_from_payload(artifact_id, kind, descriptor, payload, *, transport=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("Media payload must be bytes")
|
||||
bounded_int(len(payload), max_bytes)
|
||||
result = {"id": artifact_id, "kind": kind, "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": transport if transport is not None else {
|
||||
"type": "inline", "data_base64": base64.b64encode(payload).decode("ascii")}}
|
||||
return validate_artifact(result, max_bytes=max_bytes)
|
||||
|
||||
|
||||
async def resolve_artifact(artifact, *, download=None, max_bytes=MAX_ARTIFACT_BYTES):
|
||||
"""download(url, *, max_bytes) must return bounded bytes; never log its URL."""
|
||||
validate_artifact(artifact, max_bytes=max_bytes)
|
||||
transport = artifact["transport"]
|
||||
if transport["type"] == "inline":
|
||||
try:
|
||||
payload = base64.b64decode(transport["data_base64"], validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid media base64") from None
|
||||
else:
|
||||
if download is None:
|
||||
raise ValueError("Media URL download is not configured")
|
||||
payload = await download(transport["url"], max_bytes=artifact["byte_length"])
|
||||
if not isinstance(payload, bytes) or len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def validate_response(response, *, request_id, bindings, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
"""Validate ALL identity/completeness before downloading or delivering anything."""
|
||||
expected = validate_bindings(bindings)
|
||||
identifier(request_id)
|
||||
if not isinstance(response, dict) or set(response) != {"version", "profile", "request_id", "bindings", "artifacts", "items"}:
|
||||
raise ValueError("Invalid media response envelope")
|
||||
if type(response["version"]) is not int or response["version"] != VERSION or response["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media response version/profile")
|
||||
if response["request_id"] != request_id:
|
||||
raise ValueError("Media request identity mismatch")
|
||||
returned = response["bindings"]
|
||||
if not isinstance(returned, list) or len(returned) != len(expected):
|
||||
raise ValueError("Missing or extra media bindings")
|
||||
counts = {}
|
||||
for row in returned:
|
||||
if not isinstance(row, dict) or set(row) != BINDING_FIELDS | {"count"}:
|
||||
raise ValueError("Invalid result binding")
|
||||
binding = {key: row[key] for key in BINDING_FIELDS}
|
||||
validate_binding(binding)
|
||||
bid = binding["binding_id"]
|
||||
if bid in counts or expected.get(bid) != binding:
|
||||
raise ValueError("Media binding identity mismatch")
|
||||
# Connected outputs must produce at least one value. Empty STRING is a value.
|
||||
counts[bid] = bounded_int(row["count"], MAX_ITEMS, minimum=1)
|
||||
artifacts = response["artifacts"]
|
||||
if not isinstance(artifacts, list) or not 1 <= len(artifacts) <= MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid media artifacts count")
|
||||
indexed = {}
|
||||
total = 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = artifact["id"]
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate media artifact")
|
||||
total += artifact["byte_length"]
|
||||
if total > max_total_bytes:
|
||||
raise ValueError("Media response exceeds total byte limit")
|
||||
indexed[aid] = artifact
|
||||
items = response["items"]
|
||||
if not isinstance(items, list) or len(items) != sum(counts.values()):
|
||||
raise ValueError("Missing or extra media items")
|
||||
ordered = {bid: [None] * count for bid, count in counts.items()}
|
||||
used = set()
|
||||
for item in items:
|
||||
if not isinstance(item, dict) or set(item) != {"binding_id", "item_index", "artifact_id"}:
|
||||
raise ValueError("Invalid media item")
|
||||
bid, aid = identifier(item["binding_id"]), identifier(item["artifact_id"])
|
||||
if bid not in counts or aid not in indexed:
|
||||
raise ValueError("Unknown media item identity")
|
||||
index = bounded_int(item["item_index"], counts[bid] - 1)
|
||||
if ordered[bid][index] is not None or indexed[aid]["kind"] != expected[bid]["kind"]:
|
||||
raise ValueError("Duplicate or mistyped media item")
|
||||
ordered[bid][index] = indexed[aid]
|
||||
used.add(aid)
|
||||
if used != set(indexed) or any(value is None for values in ordered.values() for value in values):
|
||||
raise ValueError("Incomplete media result manifest")
|
||||
return ordered
|
||||
|
||||
|
||||
async def decode_response(response, *, request_id, bindings, download=None, max_total_bytes=MAX_TOTAL_BYTES):
|
||||
ordered = validate_response(response, request_id=request_id, bindings=bindings, max_total_bytes=max_total_bytes)
|
||||
# Resolve and decode before exposing a single result to the collector queue.
|
||||
from .codec import decode_payload
|
||||
cache = {}
|
||||
decoded = {}
|
||||
for bid, artifacts in ordered.items():
|
||||
values = []
|
||||
for artifact in artifacts:
|
||||
aid = artifact["id"]
|
||||
if aid not in cache:
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
if artifact["kind"] == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
cache[aid] = decode_video(artifact["descriptor"], payload)
|
||||
else:
|
||||
cache[aid] = decode_payload(artifact["kind"], artifact["descriptor"], payload)
|
||||
values.append(cache[aid])
|
||||
decoded[bid] = values
|
||||
return decoded
|
||||
@@ -0,0 +1,794 @@
|
||||
"""Pure HTTP executor for media-v1 requests against a local ComfyUI instance.
|
||||
|
||||
No ComfyUI, tensor, or provider imports. ``download(url, *, max_bytes)`` must
|
||||
provide public-HTTPS/redirect/DNS protection and bounded streaming. Optional
|
||||
``publish(metadata, payload)`` returns a contract-valid, expiring HTTPS transport.
|
||||
|
||||
Only identified invocation-owned files are deleted. Optional trusted local
|
||||
``temp_dir`` enables exporter cleanup after terminal history and collection.
|
||||
Without it, exporter retention belongs to the ComfyUI temp-directory lifecycle.
|
||||
The pinned native /interrupt route is not atomic: running-prompt cancellation
|
||||
fails closed and retains files instead of risking an unrelated successor prompt.
|
||||
"""
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import aiohttp
|
||||
|
||||
from .contract import (
|
||||
BINDING_FIELDS, MAX_ARTIFACT_BYTES, MAX_ARTIFACTS, MAX_TOTAL_BYTES,
|
||||
PROFILE, VERSION, artifact_from_payload, identifier, resolve_artifact,
|
||||
validate_artifact, validate_bindings, validate_json, validate_response,
|
||||
)
|
||||
|
||||
MAX_RESPONSE_JSON_BYTES = 10_000_000
|
||||
MAX_HTTP_JSON_BYTES = 32 * 1024 * 1024
|
||||
POLL_INTERVAL_SECONDS = 0.25
|
||||
CANCEL_GRACE_SECONDS = 10.0
|
||||
DRAIN_GRACE_SECONDS = 10.0
|
||||
HTTP_TIMEOUT_SECONDS = 15.0
|
||||
_UNCERTAIN_WORK = set()
|
||||
GET_ATTEMPTS = 3
|
||||
_IMPORT_SLOTS = {"IMAGE": 0, "AUDIO": 1, "MASK": 2, "STRING": 3, "VIDEO": 4}
|
||||
# Used ONLY to run the shared manifest validator before /view. Never fetched.
|
||||
_VALIDATION_TRANSPORT = {"type": "url", "url": "https://media-validation.invalid/artifact",
|
||||
"expires_at": "9999-01-01T00:00:00Z"}
|
||||
|
||||
|
||||
class MediaExecutionUncertain(RuntimeError):
|
||||
"""The local prompt may still execute; staged files were deliberately retained."""
|
||||
|
||||
|
||||
class _HTTPError(RuntimeError):
|
||||
def __init__(self, status):
|
||||
super().__init__(f"Local Comfy HTTP request failed (status {status})")
|
||||
self.status = status
|
||||
|
||||
|
||||
def _scope_id(value):
|
||||
identifier(value)
|
||||
if any(part in ("", ".", "..") for part in value.split("/")):
|
||||
raise ValueError("Unsafe media scope identifier")
|
||||
return value
|
||||
|
||||
|
||||
def _json_bytes(value):
|
||||
return json.dumps(value, allow_nan=False).encode("utf-8")
|
||||
|
||||
|
||||
def _link(value, workflow):
|
||||
if (not isinstance(value, list) or len(value) != 2 or
|
||||
not isinstance(value[0], str) or value[0] not in workflow or
|
||||
type(value[1]) is not int or value[1] < 0):
|
||||
raise ValueError("Invalid workflow socket link")
|
||||
return value[0]
|
||||
|
||||
|
||||
def _validate_request(input_data):
|
||||
if (not isinstance(input_data, dict) or not {"workflow", "media"} <= set(input_data)
|
||||
or set(input_data) - {"workflow", "media", "images"}):
|
||||
raise ValueError("Invalid media request envelope")
|
||||
media = input_data["media"]
|
||||
if (not isinstance(media, dict) or set(media) !=
|
||||
{"version", "profile", "request_id", "artifacts", "inputs", "outputs"}):
|
||||
raise ValueError("Invalid media request controls")
|
||||
if type(media["version"]) is not int or media["version"] != VERSION or media["profile"] != PROFILE:
|
||||
raise ValueError("Unsupported media request version/profile")
|
||||
_scope_id(media["request_id"])
|
||||
bindings = validate_bindings(media["outputs"])
|
||||
targets = []
|
||||
for binding in bindings.values():
|
||||
for field in BINDING_FIELDS - {"socket", "kind"}:
|
||||
_scope_id(binding[field])
|
||||
targets.append(binding["node_id"])
|
||||
if len(set(targets)) != len(targets):
|
||||
raise ValueError("Duplicate media exporter node")
|
||||
|
||||
workflow = input_data["workflow"]
|
||||
if not isinstance(workflow, dict) or not workflow:
|
||||
raise ValueError("Invalid API workflow")
|
||||
validate_json(workflow, max_bytes=16 * 1024 * 1024)
|
||||
dependencies = {}
|
||||
for node_id, node in workflow.items():
|
||||
_scope_id(node_id)
|
||||
if (not isinstance(node, dict) or not isinstance(node.get("class_type"), str)
|
||||
or not node["class_type"] or not isinstance(node.get("inputs"), dict)):
|
||||
raise ValueError("Invalid API workflow node")
|
||||
if node["class_type"] == "MediaImportV1":
|
||||
raise ValueError("Request may not supply its own media import filenames")
|
||||
if node["class_type"] == "MediaExportV1" and node_id not in targets:
|
||||
raise ValueError("Undeclared media exporter")
|
||||
dependencies[node_id] = set()
|
||||
for name, value in node["inputs"].items():
|
||||
_scope_id(name)
|
||||
if isinstance(value, list):
|
||||
dependencies[node_id].add(_link(value, workflow))
|
||||
for binding in bindings.values():
|
||||
node = workflow.get(binding["node_id"])
|
||||
if node is None or node["class_type"] != "MediaExportV1":
|
||||
raise ValueError("Missing declared MediaExportV1 node")
|
||||
controls = node["inputs"]
|
||||
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
required = {"request_id", "binding_json", binding["socket"]}
|
||||
optional = {"transfer_format", "transfer_quality"} if binding["kind"] == "IMAGE" else set()
|
||||
if not required <= set(controls) or set(controls) - required - optional:
|
||||
raise ValueError("Media exporter must have exactly its declared socket and controls")
|
||||
validate_options(controls.get("transfer_format", DEFAULT_FORMAT),
|
||||
controls.get("transfer_quality", DEFAULT_QUALITY))
|
||||
if controls["request_id"] != media["request_id"]:
|
||||
raise ValueError("Media exporter request identity mismatch")
|
||||
text = controls["binding_json"]
|
||||
if not isinstance(text, str) or len(text) > 128 * 1024:
|
||||
raise ValueError("Invalid media exporter binding JSON")
|
||||
try:
|
||||
declared = json.loads(text)
|
||||
except (ValueError, RecursionError):
|
||||
raise ValueError("Invalid media exporter binding JSON") from None
|
||||
if declared != binding:
|
||||
raise ValueError("Media exporter binding identity mismatch")
|
||||
_link(controls[binding["socket"]], workflow)
|
||||
# The execution roots are exactly the declared exporters, not every class
|
||||
# with OUTPUT_NODE=True. Required upstream output utilities may remain.
|
||||
visited = set()
|
||||
pending = list(targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id not in visited:
|
||||
visited.add(node_id)
|
||||
pending.extend(dependencies[node_id])
|
||||
if visited != set(workflow):
|
||||
raise ValueError("Workflow contains nodes outside declared output roots")
|
||||
if any(set(targets) & sources for sources in dependencies.values()):
|
||||
raise ValueError("A media exporter cannot be an upstream dependency")
|
||||
# Kahn traversal bounds cycle checking without relying on recursion depth.
|
||||
counts = {node: len(sources) for node, sources in dependencies.items()}
|
||||
consumers = {node: [] for node in workflow}
|
||||
for node, sources in dependencies.items():
|
||||
for source in sources:
|
||||
consumers[source].append(node)
|
||||
pending = [node for node, count in counts.items() if not count]
|
||||
processed = 0
|
||||
while pending:
|
||||
processed += 1
|
||||
for consumer in consumers[pending.pop()]:
|
||||
counts[consumer] -= 1
|
||||
if not counts[consumer]:
|
||||
pending.append(consumer)
|
||||
if processed != len(workflow):
|
||||
raise ValueError("Workflow contains a dependency cycle")
|
||||
|
||||
artifacts = media["artifacts"]
|
||||
if not isinstance(artifacts, list) or len(artifacts) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid input artifacts count")
|
||||
indexed, total = {}, 0
|
||||
for artifact in artifacts:
|
||||
validate_artifact(artifact)
|
||||
aid = _scope_id(artifact["id"])
|
||||
if aid in indexed:
|
||||
raise ValueError("Duplicate input artifact")
|
||||
indexed[aid] = artifact
|
||||
total += artifact["byte_length"]
|
||||
inputs = media["inputs"]
|
||||
if not isinstance(inputs, list) or len(inputs) > 4096:
|
||||
raise ValueError("Invalid input references")
|
||||
destinations, used = set(), set()
|
||||
exporters = {binding["node_id"]: binding for binding in bindings.values()}
|
||||
for ref in inputs:
|
||||
if not isinstance(ref, dict) or set(ref) != {"node_id", "input_name", "artifact_id"}:
|
||||
raise ValueError("Invalid media input reference")
|
||||
for value in ref.values():
|
||||
_scope_id(value)
|
||||
node_id, name, aid = ref["node_id"], ref["input_name"], ref["artifact_id"]
|
||||
if (node_id not in workflow or name not in workflow[node_id]["inputs"]
|
||||
or aid not in indexed or (node_id, name) in destinations):
|
||||
raise ValueError("Unknown or duplicate media input reference")
|
||||
binding = exporters.get(node_id)
|
||||
if binding is not None and (name != binding["socket"] or indexed[aid]["kind"] != binding["kind"]):
|
||||
raise ValueError("Exporter input must match its declared media socket and kind")
|
||||
destinations.add((node_id, name))
|
||||
used.add(aid)
|
||||
if used != set(indexed):
|
||||
raise ValueError("Every input artifact must be used")
|
||||
|
||||
images = input_data.get("images", [])
|
||||
if not isinstance(images, list) or len(images) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid legacy images count")
|
||||
names = set()
|
||||
legacy = []
|
||||
for image in images:
|
||||
if not isinstance(image, dict) or set(image) != {"name", "image"}:
|
||||
raise ValueError("Invalid legacy image")
|
||||
name, encoded = image["name"], image["image"]
|
||||
if (not isinstance(name, str) or not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_. -]{0,254}", name)
|
||||
or name in names or not isinstance(encoded, str)):
|
||||
raise ValueError("Invalid or duplicate legacy image basename")
|
||||
names.add(name)
|
||||
if len(encoded) > 4 * ((MAX_ARTIFACT_BYTES + 2) // 3):
|
||||
raise ValueError("Legacy image exceeds byte limit")
|
||||
# Exact decoded length, including padding, before allocating bytes.
|
||||
size = (len(encoded) // 4) * 3 - (len(encoded) - len(encoded.rstrip("=")))
|
||||
if len(encoded) % 4 or size < 0 or size > MAX_ARTIFACT_BYTES:
|
||||
raise ValueError("Invalid legacy image base64 length")
|
||||
total += size
|
||||
matches = [node_id for node_id, node in workflow.items()
|
||||
if node["class_type"] == "LoadImage" and node["inputs"].get("image") == name]
|
||||
if not matches or any((node_id, "image") in destinations for node_id in matches):
|
||||
raise ValueError("Unused or conflicting legacy image binding")
|
||||
legacy.append((encoded, size, matches))
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media inputs exceed total byte limit")
|
||||
return copy.deepcopy(workflow), copy.deepcopy(media), legacy, targets
|
||||
|
||||
|
||||
class _Staging:
|
||||
"""Exclusive random files under a hashed namespace; never delete siblings."""
|
||||
def __init__(self, input_dir, request_id):
|
||||
self.root = Path(input_dir).resolve()
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.relative = f"distributed-media-v1/{self.key}"
|
||||
self.names = []
|
||||
self.fd = None
|
||||
|
||||
def _open(self):
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
directory = self.root
|
||||
for component in ("distributed-media-v1", self.key):
|
||||
directory = directory / component
|
||||
if directory.is_symlink():
|
||||
raise ValueError("Unsafe media staging directory symlink")
|
||||
directory.mkdir(exist_ok=True, mode=0o700)
|
||||
if not directory.resolve().is_relative_to(self.root):
|
||||
raise ValueError("Unsafe media staging directory")
|
||||
self.fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
||||
|
||||
def write(self, payload):
|
||||
if self.fd is None:
|
||||
self._open()
|
||||
name = uuid.uuid4().hex + ".bin"
|
||||
fd = os.open(name, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600, dir_fd=self.fd)
|
||||
self.names.append(name)
|
||||
with os.fdopen(fd, "wb") as stream:
|
||||
stream.write(payload)
|
||||
return f"{self.relative}/{name}"
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is not None:
|
||||
try:
|
||||
if cleanup:
|
||||
for name in self.names:
|
||||
try:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
if cleanup:
|
||||
# Never recursively remove a shared namespace or another run.
|
||||
for directory in (self.root / self.relative, self.root / "distributed-media-v1"):
|
||||
try:
|
||||
directory.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class _OwnedExports:
|
||||
"""Best-effort retirement of validated terminal-history files, never a sweep.
|
||||
|
||||
Walk every directory component with openat/O_NOFOLLOW, then pin the directory
|
||||
fd and each regular file's identity. Missing, linked, replaced, or unidentifiable
|
||||
files stay with the endpoint lifecycle. Never use history-supplied path strings.
|
||||
"""
|
||||
def __init__(self, temp_dir, request_id):
|
||||
self.root = Path(temp_dir).absolute() if temp_dir is not None else None
|
||||
self.key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
self.fd = None
|
||||
self.files = {}
|
||||
|
||||
def _open(self):
|
||||
if self.fd is not None:
|
||||
return True
|
||||
if self.root is None or ".." in self.root.parts:
|
||||
return False
|
||||
flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
fd = None
|
||||
try:
|
||||
fd = os.open(self.root.anchor, flags)
|
||||
for component in (*self.root.parts[1:], "distributed-media-v1", self.key):
|
||||
child = os.open(component, flags, dir_fd=fd)
|
||||
os.close(fd)
|
||||
fd = child
|
||||
self.fd, fd = fd, None
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
finally:
|
||||
if fd is not None:
|
||||
os.close(fd)
|
||||
|
||||
def _identity(self, name):
|
||||
info = os.stat(name, dir_fd=self.fd, follow_symlinks=False)
|
||||
if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1:
|
||||
return None
|
||||
return (info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns)
|
||||
|
||||
def identify(self, entry, media):
|
||||
if self.root is None or not isinstance(entry.get("outputs"), dict):
|
||||
return
|
||||
# A failed prompt may contain valid completed exporters alongside missing
|
||||
# or malformed ones. Claim only independently validated declared records.
|
||||
for binding in media["outputs"]:
|
||||
node_id = binding["node_id"]
|
||||
try:
|
||||
_, files = _manifest({"outputs": {node_id: entry["outputs"].get(node_id)}},
|
||||
dict(media, outputs=[binding]))
|
||||
except ValueError:
|
||||
continue
|
||||
if not files or not self._open():
|
||||
continue
|
||||
for artifact in files:
|
||||
name = artifact["id"] + ".bin" # _manifest enforces UUID hex.
|
||||
try:
|
||||
identity = self._identity(name)
|
||||
except OSError:
|
||||
continue
|
||||
if identity is not None:
|
||||
self.files.setdefault(name, identity)
|
||||
|
||||
def close(self, *, cleanup):
|
||||
if self.fd is None:
|
||||
return
|
||||
try:
|
||||
if cleanup:
|
||||
for name, identity in self.files.items():
|
||||
try:
|
||||
if self._identity(name) == identity:
|
||||
os.unlink(name, dir_fd=self.fd)
|
||||
except OSError:
|
||||
pass # Endpoint lifecycle owns inaccessible leftovers.
|
||||
finally:
|
||||
os.close(self.fd)
|
||||
self.fd = None
|
||||
# Deliberately retain directories: they may belong to concurrent runs.
|
||||
|
||||
|
||||
async def _body(response, maximum):
|
||||
length = response.content_length
|
||||
if length is not None and length > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body = bytearray()
|
||||
async for chunk in response.content.iter_chunked(64 * 1024):
|
||||
if len(body) + len(chunk) > maximum:
|
||||
raise ValueError("Local Comfy response exceeds declared byte length")
|
||||
body.extend(chunk)
|
||||
return bytes(body)
|
||||
|
||||
|
||||
class _Execution:
|
||||
def __init__(self, session, url, stage, media, targets, deadline, exports):
|
||||
self.session, self.url, self.stage = session, url, stage
|
||||
self.media, self.targets = media, targets
|
||||
self.deadline = deadline
|
||||
self.run_id = uuid.uuid4().hex
|
||||
self.prompt_id = None
|
||||
self.attempted = False
|
||||
self.terminal = False
|
||||
self.exports = exports
|
||||
self.interruption_blocked = False
|
||||
|
||||
def remaining(self):
|
||||
remaining = self.deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
raise asyncio.TimeoutError("Media execution deadline expired")
|
||||
return remaining
|
||||
|
||||
async def json_request(self, method, path, *, body=None, permit_rejection=False):
|
||||
async with self.session.request(method, self.url + path, json=body,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining())),
|
||||
allow_redirects=False) as response:
|
||||
raw = await _body(response, MAX_HTTP_JSON_BYTES)
|
||||
# Native queue-delete and interrupt routes return an empty 200.
|
||||
if response.status == 200 and not raw and method == "POST" and path in {"/queue", "/interrupt"}:
|
||||
return {}
|
||||
try:
|
||||
value = json.loads(raw)
|
||||
except (ValueError, UnicodeError, RecursionError):
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status) from None
|
||||
raise ValueError("Invalid local Comfy JSON response") from None
|
||||
if response.status != 200:
|
||||
# Native /prompt rejects the entire graph with a 400 error.
|
||||
# Other failures may hide acceptance and are never retried.
|
||||
if (permit_rejection and response.status == 400 and isinstance(value, dict)
|
||||
and value.get("error") and "prompt_id" not in value):
|
||||
self.terminal = True
|
||||
raise _HTTPError(response.status)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("Invalid local Comfy JSON object")
|
||||
return value
|
||||
|
||||
async def get_json(self, path):
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
return await self.json_request("GET", path)
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
|
||||
async def history(self):
|
||||
response = await self.get_json("/history/" + quote(self.prompt_id, safe=""))
|
||||
entry = response.get(self.prompt_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if not isinstance(entry, dict) or not isinstance(entry.get("status"), dict):
|
||||
raise ValueError("Malformed Comfy prompt history")
|
||||
status = entry["status"]
|
||||
messages = status.get("messages", [])
|
||||
if not isinstance(messages, list):
|
||||
raise ValueError("Malformed Comfy history messages")
|
||||
failed = status.get("status_str") in {"error", "interrupted"} or any(
|
||||
isinstance(message, list) and message and message[0] in {"execution_error", "execution_interrupted"}
|
||||
for message in messages)
|
||||
if failed or status.get("completed") is True:
|
||||
self.terminal = True
|
||||
self.exports.identify(entry, self.media)
|
||||
if failed or status.get("status_str") != "success":
|
||||
raise RuntimeError("Local Comfy prompt execution failed")
|
||||
return entry
|
||||
return None
|
||||
|
||||
async def submit_and_wait(self, workflow):
|
||||
body = {"prompt": workflow, "client_id": self.run_id,
|
||||
"extra_data": {"distributed_media_run_id": self.run_id},
|
||||
"partial_execution_targets": list(self.targets)}
|
||||
self.remaining() # CPU staging must not spend past the absolute deadline.
|
||||
self.attempted = True
|
||||
response = await self.json_request("POST", "/prompt", body=body, permit_rejection=True)
|
||||
self.prompt_id = _scope_id(response.get("prompt_id"))
|
||||
if response.get("node_errors"):
|
||||
raise RuntimeError("Local Comfy accepted an incomplete set of execution targets")
|
||||
while True:
|
||||
entry = await self.history()
|
||||
if entry is not None:
|
||||
return entry
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def settle(self):
|
||||
"""Bounded recovery; caller shields this from task cancellation."""
|
||||
deleted = False
|
||||
while True:
|
||||
if self.prompt_id:
|
||||
try:
|
||||
await self.history()
|
||||
except RuntimeError:
|
||||
if not self.terminal:
|
||||
raise
|
||||
if self.terminal:
|
||||
return
|
||||
queue = await self.get_json("/queue")
|
||||
groups = []
|
||||
for key in ("queue_running", "queue_pending"):
|
||||
rows = queue.get(key)
|
||||
if not isinstance(rows, list):
|
||||
raise ValueError("Malformed Comfy queue")
|
||||
group = {}
|
||||
for row in rows:
|
||||
if not isinstance(row, list) or len(row) < 4 or not isinstance(row[3], dict):
|
||||
raise ValueError("Malformed Comfy queue entry")
|
||||
pid = _scope_id(row[1])
|
||||
if pid in group:
|
||||
raise ValueError("Duplicate Comfy queue identity")
|
||||
group[pid] = row[3]
|
||||
groups.append(group)
|
||||
running, pending = groups
|
||||
if self.prompt_id is None:
|
||||
own = [pid for group in groups for pid, extra in group.items()
|
||||
if extra.get("distributed_media_run_id") == self.run_id
|
||||
and extra.get("client_id") == self.run_id]
|
||||
if len(own) > 1:
|
||||
raise ValueError("Ambiguous Comfy queue ownership")
|
||||
if own:
|
||||
self.prompt_id = own[0]
|
||||
if self.prompt_id in pending and not deleted:
|
||||
await self.json_request("POST", "/queue", body={"delete": [self.prompt_id]})
|
||||
deleted = True
|
||||
elif self.prompt_id in running:
|
||||
# Pinned ComfyUI server.py checks a snapshot then globally
|
||||
# interrupts, outside the queue mutex. Supplying prompt_id does
|
||||
# NOT close that race. No safe capability negotiation exists on
|
||||
# this route; do not invoke it, even after our own queue read.
|
||||
self.interruption_blocked = True
|
||||
raise RuntimeError("Native prompt-scoped interruption unavailable")
|
||||
elif deleted and self.prompt_id not in running and self.prompt_id not in pending:
|
||||
self.terminal = True # queue deletion confirmed by exact read-back
|
||||
return
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
|
||||
async def fetch(self, artifact):
|
||||
transport = artifact["transport"]
|
||||
params = {"filename": transport["filename"], "subfolder": transport["subfolder"], "type": "temp"}
|
||||
for attempt in range(GET_ATTEMPTS):
|
||||
try:
|
||||
async with self.session.get(self.url + "/view", params=params, allow_redirects=False,
|
||||
timeout=aiohttp.ClientTimeout(total=min(HTTP_TIMEOUT_SECONDS, self.remaining()))) as response:
|
||||
if response.status != 200:
|
||||
raise _HTTPError(response.status)
|
||||
payload = await _body(response, artifact["byte_length"])
|
||||
break
|
||||
except (_HTTPError, aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
if isinstance(exc, _HTTPError) and exc.status != 429 and exc.status < 500:
|
||||
raise
|
||||
if attempt == GET_ATTEMPTS - 1:
|
||||
raise
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS * (2 ** attempt))
|
||||
if len(payload) != artifact["byte_length"]:
|
||||
raise ValueError("Media output byte length mismatch")
|
||||
if hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Media output checksum mismatch")
|
||||
return payload
|
||||
|
||||
|
||||
def _manifest(entry, media):
|
||||
outputs = entry.get("outputs")
|
||||
if not isinstance(outputs, dict):
|
||||
raise ValueError("Missing media history outputs")
|
||||
expected_ids = {binding["node_id"] for binding in media["outputs"]}
|
||||
if any(node_id not in expected_ids and isinstance(value, dict) and "media_exports" in value
|
||||
for node_id, value in outputs.items()):
|
||||
raise ValueError("Unexpected media export history")
|
||||
response = {"version": VERSION, "profile": PROFILE, "request_id": media["request_id"],
|
||||
"bindings": [], "artifacts": [], "items": []}
|
||||
files = []
|
||||
subfolder = "distributed-media-v1/" + hashlib.sha256(media["request_id"].encode()).hexdigest()
|
||||
for binding in media["outputs"]:
|
||||
node = outputs.get(binding["node_id"], {})
|
||||
records = node.get("media_exports") if isinstance(node, dict) else None
|
||||
if not isinstance(records, list) or len(records) != 1:
|
||||
raise ValueError("Missing or duplicate media export result")
|
||||
record = records[0]
|
||||
if (not isinstance(record, dict) or set(record) != {"request_id", "binding", "artifacts", "items"}
|
||||
or record["request_id"] != media["request_id"]):
|
||||
raise ValueError("Media export request identity mismatch")
|
||||
if not isinstance(record["artifacts"], list) or len(record["artifacts"]) > MAX_ARTIFACTS:
|
||||
raise ValueError("Invalid exported artifacts count")
|
||||
converted = []
|
||||
for artifact in record["artifacts"]:
|
||||
if not isinstance(artifact, dict) or not isinstance(artifact.get("transport"), dict):
|
||||
raise ValueError("Invalid exported artifact")
|
||||
transport = artifact["transport"]
|
||||
aid = artifact.get("id")
|
||||
if (not isinstance(aid, str) or not re.fullmatch(r"[a-f0-9]{32}", aid)
|
||||
or transport != {"type": "comfy-file", "filename": aid + ".bin",
|
||||
"subfolder": subfolder, "folder_type": "temp"}):
|
||||
raise ValueError("Unsafe or unowned exported media file")
|
||||
converted.append(dict(artifact, transport=dict(_VALIDATION_TRANSPORT)))
|
||||
partial = dict(response, bindings=[record["binding"]], artifacts=converted, items=record["items"])
|
||||
validate_response(partial, request_id=media["request_id"], bindings=[binding])
|
||||
response["bindings"].append(copy.deepcopy(record["binding"]))
|
||||
response["artifacts"].extend(converted)
|
||||
response["items"].extend(copy.deepcopy(record["items"]))
|
||||
files.extend(copy.deepcopy(record["artifacts"]))
|
||||
# Validate the whole manifest BEFORE the first /view or publisher callback.
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
return response, files
|
||||
|
||||
|
||||
def _publication_plan(response, files, publish):
|
||||
skeleton = dict(response, artifacts=[dict(a, transport={"type": "inline", "data_base64": ""}) for a in files])
|
||||
base_size = len(_json_bytes({"media": skeleton}))
|
||||
sizes = {a["id"]: 4 * ((a["byte_length"] + 2) // 3) for a in files}
|
||||
predicted = base_size + sum(sizes.values())
|
||||
published = set()
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES and publish is None:
|
||||
raise ValueError("Media response exceeds inline JSON limit; a publisher is required")
|
||||
# Reserve enough space for a maximum-size escaped URL, including metadata.
|
||||
# Small remaining inline values preserve their exact bytes and ordering.
|
||||
for aid in sorted(sizes, key=sizes.get, reverse=True):
|
||||
if predicted <= MAX_RESPONSE_JSON_BYTES:
|
||||
break
|
||||
if sizes[aid] > 65_536:
|
||||
published.add(aid)
|
||||
predicted -= sizes[aid] - 65_536
|
||||
if predicted > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response metadata exceeds JSON limit")
|
||||
return published
|
||||
|
||||
|
||||
async def _run(execution, workflow, legacy, download, publish):
|
||||
media = execution.media
|
||||
imports = {}
|
||||
for artifact in media["artifacts"]:
|
||||
execution.remaining()
|
||||
payload = await resolve_artifact(artifact, download=download)
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
while node_id in workflow:
|
||||
node_id = "distributed-media-import:" + uuid.uuid4().hex
|
||||
metadata = {key: value for key, value in artifact.items() if key != "transport"}
|
||||
workflow[node_id] = {"class_type": "MediaImportV1", "inputs": {
|
||||
"artifact_json": json.dumps(metadata, allow_nan=False), "filename": filename}}
|
||||
imports[artifact["id"]] = [node_id, _IMPORT_SLOTS[artifact["kind"]]]
|
||||
for ref in media["inputs"]:
|
||||
workflow[ref["node_id"]]["inputs"][ref["input_name"]] = list(imports[ref["artifact_id"]])
|
||||
for encoded, size, matches in legacy:
|
||||
execution.remaining()
|
||||
try:
|
||||
payload = base64.b64decode(encoded, validate=True)
|
||||
except (ValueError, binascii.Error):
|
||||
raise ValueError("Invalid legacy image base64") from None
|
||||
if len(payload) != size:
|
||||
raise ValueError("Legacy image byte length mismatch")
|
||||
execution.remaining()
|
||||
filename = execution.stage.write(payload)
|
||||
for node_id in matches:
|
||||
workflow[node_id]["inputs"]["image"] = filename
|
||||
# Replacements can disconnect file loaders whose original endpoint-local
|
||||
# paths are unavailable. Keep exactly the post-rewrite exporter closure;
|
||||
# shared ancestors still needed by another branch must survive.
|
||||
reachable, pending = set(), list(execution.targets)
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in reachable:
|
||||
continue
|
||||
reachable.add(node_id)
|
||||
pending.extend(_link(value, workflow) for value in workflow[node_id]["inputs"].values()
|
||||
if isinstance(value, list))
|
||||
workflow = {node_id: node for node_id, node in workflow.items() if node_id in reachable}
|
||||
entry = await execution.submit_and_wait(workflow)
|
||||
response, files = _manifest(entry, media)
|
||||
published = _publication_plan(response, files, publish)
|
||||
artifacts = []
|
||||
for artifact in files:
|
||||
payload = await execution.fetch(artifact)
|
||||
transport = None
|
||||
if artifact["id"] in published:
|
||||
metadata = {key: copy.deepcopy(value) for key, value in artifact.items() if key != "transport"}
|
||||
transport = await publish(metadata, payload)
|
||||
if not isinstance(transport, dict) or transport.get("type") != "url":
|
||||
raise ValueError("Publisher must return an expiring HTTPS URL transport")
|
||||
artifacts.append(artifact_from_payload(artifact["id"], artifact["kind"], artifact["descriptor"],
|
||||
payload, transport=transport))
|
||||
response["artifacts"] = artifacts
|
||||
validate_response(response, request_id=media["request_id"], bindings=media["outputs"])
|
||||
result = {"media": response}
|
||||
if len(_json_bytes(result)) > MAX_RESPONSE_JSON_BYTES:
|
||||
raise ValueError("Media response exceeds JSON byte limit")
|
||||
execution.remaining()
|
||||
return result
|
||||
|
||||
|
||||
def _release_uncertain_work(task):
|
||||
_UNCERTAIN_WORK.discard(task)
|
||||
try:
|
||||
task.result()
|
||||
except BaseException:
|
||||
pass
|
||||
|
||||
|
||||
async def _drain_work(work):
|
||||
"""Cancel once; a misbehaving callback cannot make teardown unbounded."""
|
||||
if not work.done():
|
||||
work.cancel()
|
||||
deadline = asyncio.get_running_loop().time() + DRAIN_GRACE_SECONDS
|
||||
while not work.done():
|
||||
remaining = deadline - asyncio.get_running_loop().time()
|
||||
if remaining <= 0:
|
||||
# Keep a strong reference until it exits, but never delete its files
|
||||
# asynchronously: uncertain ownership belongs to endpoint recovery.
|
||||
_UNCERTAIN_WORK.add(work)
|
||||
work.add_done_callback(_release_uncertain_work)
|
||||
return False
|
||||
try:
|
||||
# Unlike wait_for, wait's timeout does not wait for cancellation.
|
||||
await asyncio.wait({work}, timeout=remaining)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
_release_uncertain_work(work)
|
||||
return True
|
||||
|
||||
|
||||
async def execute_media_request(input_data, *, comfy_url, input_dir, temp_dir=None,
|
||||
timeout_seconds=1800, session=None, download=None, publish=None) -> dict:
|
||||
"""Validate, stage, submit once, collect exact media, and retire owned inputs.
|
||||
|
||||
``timeout_seconds`` covers staging, HTTP, polling, and publishing. A separate
|
||||
bounded cancellation grace is allowed ONLY to establish the local prompt is
|
||||
terminal. On uncertain acceptance/termination, raise MediaExecutionUncertain
|
||||
and retain input files rather than racing an executing prompt. Never retry an
|
||||
ambiguous submission, even when the exception is a timeout.
|
||||
Async callback teardown has its own bounded grace; non-cooperative callbacks
|
||||
also raise MediaExecutionUncertain and retain files, even after terminal history.
|
||||
|
||||
``temp_dir`` is a trusted endpoint setting for this ComfyUI's local temp root,
|
||||
never a request-supplied path. Only validated exporter UUID.bin files beneath
|
||||
distributed-media-v1/sha256(request_id) are retired, after all collection and
|
||||
publisher work has unwound. Unidentified files remain lifecycle-owned. With
|
||||
``temp_dir=None`` all exporter files remain lifecycle-owned as before.
|
||||
|
||||
The pinned native /interrupt route has a check/global-interrupt race. Running
|
||||
prompts therefore fail closed with MediaExecutionUncertain, without posting
|
||||
/interrupt. A prompt-scoped request body alone is not a safety guarantee.
|
||||
"""
|
||||
if (type(timeout_seconds) not in (int, float) or not math.isfinite(timeout_seconds)
|
||||
or timeout_seconds <= 0):
|
||||
raise ValueError("Invalid media execution timeout")
|
||||
if not isinstance(comfy_url, str):
|
||||
raise ValueError("Invalid local Comfy URL")
|
||||
try:
|
||||
parsed = urlsplit(comfy_url)
|
||||
if (parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username
|
||||
or parsed.password or parsed.query or parsed.fragment):
|
||||
raise ValueError()
|
||||
parsed.port
|
||||
except ValueError:
|
||||
raise ValueError("Invalid local Comfy URL") from None
|
||||
deadline = asyncio.get_running_loop().time() + timeout_seconds
|
||||
workflow, media, legacy, targets = _validate_request(input_data)
|
||||
stage = _Staging(input_dir, media["request_id"])
|
||||
exports = _OwnedExports(temp_dir, media["request_id"])
|
||||
owned_session = session is None
|
||||
if owned_session:
|
||||
session = aiohttp.ClientSession()
|
||||
execution = _Execution(session, comfy_url.rstrip("/"), stage, media, targets, deadline, exports)
|
||||
work = None
|
||||
try:
|
||||
remaining = execution.remaining()
|
||||
work = asyncio.create_task(_run(execution, workflow, legacy, download, publish))
|
||||
return await asyncio.wait_for(asyncio.shield(work), remaining)
|
||||
except BaseException as exc:
|
||||
# Retire files only after readers unwind. Cooperative async callbacks
|
||||
# normally drain immediately; an uncooperative one requires quarantine.
|
||||
if work is not None and not await _drain_work(work):
|
||||
raise MediaExecutionUncertain(
|
||||
"Media callback teardown unconfirmed; owned files retained; "
|
||||
"quarantine this endpoint before reuse; do not resubmit automatically"
|
||||
) from exc
|
||||
if execution.attempted and not execution.terminal:
|
||||
execution.deadline = asyncio.get_running_loop().time() + CANCEL_GRACE_SECONDS
|
||||
recovery = asyncio.create_task(asyncio.wait_for(execution.settle(), CANCEL_GRACE_SECONDS))
|
||||
# Shield the bounded recovery even from repeated task.cancel() calls.
|
||||
while not recovery.done():
|
||||
try:
|
||||
await asyncio.shield(recovery)
|
||||
except asyncio.CancelledError:
|
||||
continue
|
||||
except Exception:
|
||||
break
|
||||
try:
|
||||
recovery.result()
|
||||
except BaseException:
|
||||
pass
|
||||
if not execution.terminal:
|
||||
raise MediaExecutionUncertain(
|
||||
"Local Comfy prompt terminal state unconfirmed; owned staged files retained "
|
||||
f"under {stage.relative}; prompt_id={execution.prompt_id or 'unknown'}; "
|
||||
"do not resubmit this request automatically"
|
||||
+ ("; native prompt-scoped interruption unavailable (non-atomic /interrupt route)"
|
||||
if execution.interruption_blocked else "")
|
||||
) from exc
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
try:
|
||||
exports.close(cleanup=(work is None or work.done()) and execution.terminal)
|
||||
finally:
|
||||
stage.close(cleanup=(work is None or work.done()) and
|
||||
(not execution.attempted or execution.terminal))
|
||||
finally:
|
||||
if owned_session:
|
||||
await session.close()
|
||||
@@ -0,0 +1,634 @@
|
||||
"""Bounded optional IMAGE returns, independent of ComfyUI and Distributed.
|
||||
|
||||
Exact mode is the existing tensor wire format, unchanged. All other modes clip
|
||||
and truncate RGB/luma to uint8 (as the Collector does), normalize in float32,
|
||||
then cast to the source dtype. RGBA always carries an exact tensor alpha
|
||||
sidecar, including PNG/WebP. ``legacy_png`` means uncompressed binary PNG here,
|
||||
not the old Collector's JSON/base64 transport.
|
||||
|
||||
The compressed wire format concatenates length-delimited stills or one Matroska
|
||||
video, followed by optional alpha bytes. No archive, path, URL or pickle is
|
||||
accepted. Limits cover wire bytes, descriptor structure, output tensor bytes,
|
||||
and aggregate padded pixel buffers, not process RSS/native codec overhead.
|
||||
Video preflight reads EBML track dimensions and block counts before creating a
|
||||
decoder. We extract packets ourselves, avoiding native container stream probing
|
||||
(which can decode before caller checks). FFmpeg max_pixels also bounds
|
||||
bitstream-driven resolution changes, including native buffer alignment.
|
||||
Software codecs only, one thread, no fallback to another encoder/format.
|
||||
"""
|
||||
|
||||
import io
|
||||
import struct
|
||||
import zlib
|
||||
from fractions import Fraction
|
||||
|
||||
import torch
|
||||
|
||||
from . import codec
|
||||
from .codec import MAX_PAYLOAD_BYTES
|
||||
from .image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options
|
||||
|
||||
MAX_FRAMES = 240
|
||||
MAX_DIMENSION = 8192
|
||||
MAX_DESCRIPTOR_BYTES = 64 * 1024
|
||||
MAX_CONTAINER_ELEMENTS = 16384
|
||||
_VIDEO_CODECS = {
|
||||
"ffv1": ("ffv1", None, "bgr0"),
|
||||
"h264": ("libx264", 51, "yuv420p"),
|
||||
"h265": ("libx265", 51, "yuv420p"),
|
||||
"av1": ("libsvtav1", 63, "yuv420p"),
|
||||
}
|
||||
_CODEC_IDS = {"ffv1": b"V_FFV1", "h264": b"V_MPEG4/ISO/AVC",
|
||||
"h265": b"V_MPEGH/ISO/HEVC", "av1": b"V_AV1"}
|
||||
|
||||
|
||||
class _BoundedIO(io.BytesIO):
|
||||
"""Enforce the final file size on every write, including muxer flush/seek."""
|
||||
def __init__(self, limit):
|
||||
super().__init__()
|
||||
self.limit = limit
|
||||
|
||||
def write(self, data):
|
||||
if self.tell() + len(data) > self.limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return super().write(data)
|
||||
|
||||
def seek(self, offset, whence=0):
|
||||
if whence == 0:
|
||||
position = offset
|
||||
elif whence == 1:
|
||||
position = self.tell() + offset
|
||||
elif whence == 2:
|
||||
position = self.getbuffer().nbytes + offset
|
||||
else:
|
||||
raise ValueError("Invalid image buffer seek")
|
||||
if not 0 <= position <= self.limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return super().seek(offset, whence)
|
||||
|
||||
|
||||
def _quality_to_crf(quality, maximum):
|
||||
return int(round(maximum * (1.0 - quality / 100.0) ** 0.6))
|
||||
|
||||
|
||||
def _padded_size(fmt, height, width):
|
||||
if fmt in ("h264", "h265", "av1"):
|
||||
height, width = height + height % 2, width + width % 2
|
||||
# libx265 / SVT-AV1 cannot safely encode arbitrarily small pictures.
|
||||
if fmt in ("h265", "av1"):
|
||||
height, width = max(64, height), max(64, width)
|
||||
return height, width
|
||||
|
||||
|
||||
def _shape_budget(shape, fmt, limit, itemsize=4, alpha_size=0, coded_frames=None):
|
||||
size = codec._tensor_size("IMAGE", shape, itemsize, limit)
|
||||
count, height, width, _ = shape
|
||||
if count > MAX_FRAMES or max(height, width) > MAX_DIMENSION:
|
||||
raise ValueError("IMAGE frame count or dimensions exceed limit")
|
||||
ph, pw = _padded_size(fmt, height, width)
|
||||
# Reserve output and padded RGBA-sized decoded buffers BEFORE allocation.
|
||||
bh, bw = ((ph + 63) // 64 * 64, (pw + 63) // 64 * 64) if fmt in _VIDEO_CODECS else (ph, pw)
|
||||
# Slice/OBU counts conservatively bound native frames, including AV1 hidden
|
||||
# frames and multiple outputs from one packet. Reserve 16 bytes per pixel
|
||||
# for hostile high-bit-depth streams, not only our usual 8-bit RGB frames.
|
||||
native_count = max(count, coded_frames or count)
|
||||
if native_count > MAX_FRAMES * 4:
|
||||
raise ValueError("IMAGE coded frame count exceeds limit")
|
||||
pixels = codec._tensor_size("IMAGE", [native_count, bh, bw, 4], 4, limit)
|
||||
if size + pixels + alpha_size > limit:
|
||||
raise ValueError("IMAGE decoded aggregate exceeds byte limit")
|
||||
return ph, pw
|
||||
|
||||
|
||||
def _descriptor_budget(value):
|
||||
"""Bound the JSON-shaped tree without serializing attacker-sized values."""
|
||||
remaining = MAX_DESCRIPTOR_BYTES
|
||||
def walk(obj, depth):
|
||||
nonlocal remaining
|
||||
if depth > 5:
|
||||
raise ValueError("IMAGE descriptor nesting exceeds limit")
|
||||
remaining -= 8
|
||||
if remaining < 0:
|
||||
raise ValueError("IMAGE descriptor exceeds byte limit")
|
||||
if obj is None:
|
||||
return
|
||||
if type(obj) is str:
|
||||
remaining -= len(obj) * 6
|
||||
elif type(obj) is int:
|
||||
if obj.bit_length() > 63:
|
||||
raise ValueError("IMAGE descriptor integer exceeds limit")
|
||||
remaining -= 20
|
||||
elif type(obj) is list:
|
||||
if len(obj) > MAX_FRAMES:
|
||||
raise ValueError("IMAGE descriptor list exceeds limit")
|
||||
for item in obj:
|
||||
walk(item, depth + 1)
|
||||
elif type(obj) is dict:
|
||||
if len(obj) > 10:
|
||||
raise ValueError("IMAGE descriptor fields exceed limit")
|
||||
for key, item in obj.items():
|
||||
if type(key) is not str:
|
||||
raise ValueError("Invalid IMAGE descriptor key")
|
||||
walk(key, depth + 1)
|
||||
walk(item, depth + 1)
|
||||
else:
|
||||
raise ValueError("Invalid IMAGE descriptor value")
|
||||
if remaining < 0:
|
||||
raise ValueError("IMAGE descriptor exceeds byte limit")
|
||||
walk(value, 0)
|
||||
|
||||
|
||||
def _alpha_metadata(alpha, shape, limit):
|
||||
if shape[-1] != 4:
|
||||
if alpha is not None:
|
||||
raise ValueError("Unexpected IMAGE alpha sidecar")
|
||||
return 0, torch.float32
|
||||
if type(alpha) is not dict or set(alpha) != {"descriptor", "size"}:
|
||||
raise ValueError("Missing or invalid IMAGE alpha sidecar")
|
||||
d = alpha["descriptor"]
|
||||
if type(d) is not dict or set(d) != {"encoding", "dtype", "shape", "layout"}:
|
||||
raise ValueError("Invalid IMAGE alpha descriptor")
|
||||
if d["encoding"] != "tensor-le-v1" or d["layout"] != "BHWC" or d["shape"] != shape[:3] + [1]:
|
||||
raise ValueError("Invalid IMAGE alpha shape or encoding")
|
||||
if type(d["dtype"]) is not str or d["dtype"] not in codec._DTYPES:
|
||||
raise ValueError("Invalid IMAGE alpha dtype")
|
||||
dtype, itemsize = codec._DTYPES[d["dtype"]]
|
||||
size = codec._tensor_size("IMAGE", d["shape"], itemsize, limit)
|
||||
if type(alpha["size"]) is not int or alpha["size"] != size:
|
||||
raise ValueError("Invalid IMAGE alpha size")
|
||||
return size, torch.float64 if dtype == torch.float64 else torch.float32
|
||||
|
||||
|
||||
def _import_av():
|
||||
try:
|
||||
import av
|
||||
except ImportError:
|
||||
raise ValueError("PyAV is required for IMAGE video transfer") from None
|
||||
return av
|
||||
|
||||
|
||||
def _encode_still(frame, fmt, quality, limit):
|
||||
from PIL import Image
|
||||
sink = _BoundedIO(limit)
|
||||
if frame.shape[-1] == 1:
|
||||
frame = frame[..., 0]
|
||||
image = Image.fromarray(frame)
|
||||
if fmt in ("legacy_png", "png"):
|
||||
image.save(sink, format="PNG", compress_level=0 if fmt == "legacy_png" else 4)
|
||||
elif fmt == "jpeg":
|
||||
image.save(sink, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
|
||||
else:
|
||||
image.save(sink, format="WEBP", quality=quality, method=4)
|
||||
return sink.getvalue()
|
||||
|
||||
|
||||
def _encode_video(frames, fmt, quality, height, width, limit):
|
||||
import numpy as np
|
||||
av = _import_av()
|
||||
encoder, maximum, pix_fmt = _VIDEO_CODECS[fmt]
|
||||
try:
|
||||
av.codec.Codec(encoder, "w")
|
||||
except Exception:
|
||||
raise ValueError("IMAGE software encoder unavailable: " + encoder) from None
|
||||
sink = _BoundedIO(limit)
|
||||
options = {}
|
||||
if maximum is not None:
|
||||
options["crf"] = str(_quality_to_crf(quality, maximum))
|
||||
if fmt == "h264":
|
||||
options["preset"] = "medium"
|
||||
elif fmt == "h265":
|
||||
options.update(preset="fast", **{"x265-params": "log-level=error:pools=none:frame-threads=1"})
|
||||
elif fmt == "av1":
|
||||
options.update(preset="8", **{"svtav1-params": "lp=1"})
|
||||
else:
|
||||
options.update(level="3", slicecrc="1")
|
||||
try:
|
||||
with av.open(sink, mode="w", format="matroska") as container:
|
||||
stream = container.add_stream(encoder, rate=24)
|
||||
stream.width, stream.height = width, height
|
||||
stream.pix_fmt = pix_fmt
|
||||
stream.codec_context.thread_count = 1
|
||||
stream.options = options
|
||||
for index, frame in enumerate(frames):
|
||||
if frame.shape[-1] == 1:
|
||||
frame = np.repeat(frame, 3, axis=-1)
|
||||
h, w = frame.shape[:2]
|
||||
frame = np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
|
||||
vf = av.VideoFrame.from_ndarray(frame, format="rgb24")
|
||||
vf.pts = index
|
||||
vf.time_base = Fraction(1, 24)
|
||||
for packet in stream.encode(vf):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
data = sink.getvalue()
|
||||
if len(data) > limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return data
|
||||
except av.error.FFmpegError:
|
||||
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ValueError("IMAGE software video encode failed (codec, dimensions or byte limit)") from None
|
||||
|
||||
|
||||
def encode_image(value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY,
|
||||
max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Return (descriptor, bytes); reject unsupported inputs, never fall back."""
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if fmt == "exact":
|
||||
return codec.encode_payload("IMAGE", value, max_bytes=max_bytes)
|
||||
limit = codec._byte_limit(max_bytes)
|
||||
if not isinstance(value, torch.Tensor) or value.device.type != "cpu":
|
||||
raise ValueError("IMAGE must be a CPU tensor")
|
||||
if value.layout != torch.strided or value.is_nested:
|
||||
raise ValueError("IMAGE must have dense strided layout")
|
||||
if value.dtype not in codec._DTYPE_NAMES:
|
||||
raise ValueError("Unsupported IMAGE dtype")
|
||||
shape = list(value.shape)
|
||||
codec._tensor_size("IMAGE", shape, value.element_size(), limit)
|
||||
alpha_size = shape[0] * shape[1] * shape[2] * value.element_size() if shape[-1] == 4 else 0
|
||||
itemsize = value.element_size()
|
||||
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
|
||||
if not torch.isfinite(value).all().item():
|
||||
raise ValueError("IMAGE values must be finite")
|
||||
value = value.detach().resolve_neg()
|
||||
alpha, alpha_payload = None, b""
|
||||
if shape[-1] == 4:
|
||||
ad, alpha_payload = codec.encode_payload("IMAGE", value[..., 3:4], max_bytes=limit)
|
||||
alpha = {"descriptor": ad, "size": len(alpha_payload)}
|
||||
import numpy as np
|
||||
def frames():
|
||||
for frame in value:
|
||||
arr = frame[..., :3].float().numpy()
|
||||
with np.errstate(over="ignore", invalid="ignore"):
|
||||
yield np.ascontiguousarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))
|
||||
budget = limit - len(alpha_payload)
|
||||
if fmt in _VIDEO_CODECS:
|
||||
parts = [_encode_video(frames(), fmt, quality, ph, pw, budget)]
|
||||
header = _video_header(parts[0], fmt, ph, pw, shape[0])
|
||||
coded_frames = _coded_frame_bound(parts[0], fmt, header)
|
||||
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
|
||||
else:
|
||||
parts = []
|
||||
for frame in frames():
|
||||
part = _encode_still(frame, fmt, quality, budget)
|
||||
budget -= len(part)
|
||||
parts.append(part)
|
||||
descriptor = {"encoding": "image-transfer-v1", "format": fmt, "quality": quality,
|
||||
"shape": shape, "layout": "BHWC", "dtype": codec._DTYPE_NAMES[value.dtype],
|
||||
"parts": [len(p) for p in parts], "alpha": alpha}
|
||||
_descriptor_budget(descriptor)
|
||||
payload = b"".join(parts) + alpha_payload
|
||||
if len(payload) > limit:
|
||||
raise ValueError("Encoded IMAGE exceeds byte limit")
|
||||
return descriptor, payload
|
||||
|
||||
|
||||
def _still_header(data, fmt):
|
||||
"""Read dimensions before Pillow (WebP may allocate pixels during open)."""
|
||||
if fmt in ("png", "legacy_png"):
|
||||
if data[:8] != b"\x89PNG\r\n\x1a\n":
|
||||
raise ValueError("Invalid IMAGE PNG header")
|
||||
offset, dimensions, saw_data, chunks = 8, None, False, 0
|
||||
while offset < len(data):
|
||||
chunks += 1
|
||||
if chunks > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE PNG chunk count exceeds limit")
|
||||
if offset + 12 > len(data):
|
||||
raise ValueError("Truncated IMAGE PNG chunk")
|
||||
length = int.from_bytes(data[offset:offset + 4], "big")
|
||||
end = offset + 12 + length
|
||||
if end > len(data):
|
||||
raise ValueError("Truncated IMAGE PNG chunk")
|
||||
tag = data[offset + 4:offset + 8]
|
||||
body = data[offset + 8:end - 4]
|
||||
if zlib.crc32(data[offset + 4:end - 4]) != int.from_bytes(data[end - 4:end], "big"):
|
||||
raise ValueError("Invalid IMAGE PNG checksum")
|
||||
if tag == b"IHDR" and offset == 8 and length == 13:
|
||||
width, height, depth, color, compression, filtering, interlace = struct.unpack(">IIBBBBB", body)
|
||||
if depth != 8 or color not in (0, 2) or (compression, filtering, interlace) != (0, 0, 0):
|
||||
raise ValueError("Unsupported IMAGE PNG layout")
|
||||
dimensions = (height, width)
|
||||
elif tag == b"IDAT" and dimensions:
|
||||
saw_data = True
|
||||
elif tag == b"IEND" and not length and end == len(data) and saw_data:
|
||||
return dimensions
|
||||
else:
|
||||
raise ValueError("Unsupported IMAGE PNG chunk")
|
||||
offset = end
|
||||
raise ValueError("Truncated IMAGE PNG")
|
||||
if fmt == "webp":
|
||||
# Our encoder writes a single lossy VP8 still. Reject animation/alpha/
|
||||
# extended WebP before WebPAnimDecoder can allocate its canvas.
|
||||
if len(data) < 30 or data[:4] != b"RIFF" or data[8:16] != b"WEBPVP8 ":
|
||||
raise ValueError("Unsupported IMAGE WebP header")
|
||||
size = int.from_bytes(data[16:20], "little")
|
||||
if int.from_bytes(data[4:8], "little") + 8 != len(data) or 20 + size + size % 2 != len(data):
|
||||
raise ValueError("Invalid IMAGE WebP length")
|
||||
if data[23:26] != b"\x9d\x01\x2a":
|
||||
raise ValueError("Invalid IMAGE WebP frame")
|
||||
width, height = struct.unpack("<HH", data[26:30])
|
||||
return height & 0x3fff, width & 0x3fff
|
||||
if data[:2] != b"\xff\xd8" or data[-2:] != b"\xff\xd9":
|
||||
raise ValueError("Invalid IMAGE JPEG header or trailer")
|
||||
offset, dimensions = 2, None
|
||||
while offset + 4 <= len(data):
|
||||
if data[offset] != 0xff:
|
||||
raise ValueError("Invalid IMAGE JPEG marker")
|
||||
marker = data[offset + 1]
|
||||
size = int.from_bytes(data[offset + 2:offset + 4], "big")
|
||||
end = offset + 2 + size
|
||||
if size < 2 or end > len(data):
|
||||
raise ValueError("Truncated IMAGE JPEG marker")
|
||||
if marker in (0xc0, 0xc2):
|
||||
if dimensions or size < 8 or data[offset + 4] != 8 or data[offset + 9] not in (1, 3):
|
||||
raise ValueError("Unsupported IMAGE JPEG layout")
|
||||
dimensions = struct.unpack(">HH", data[offset + 5:offset + 9])
|
||||
elif marker == 0xda:
|
||||
if dimensions is None:
|
||||
raise ValueError("Missing IMAGE JPEG dimensions")
|
||||
return dimensions
|
||||
offset = end
|
||||
raise ValueError("Missing IMAGE JPEG pixels")
|
||||
|
||||
|
||||
def _video_header(data, fmt, height, width, count):
|
||||
"""Strict bounded EBML preflight: one track, no lacing/attachments/encodings.
|
||||
|
||||
Native stream probing is not the first dimension check: it may instantiate a
|
||||
decoder. Validate the Matroska header and count blocks without decompressing.
|
||||
"""
|
||||
nodes = 0
|
||||
def vint(pos, end, identifier=False):
|
||||
if pos >= end or data[pos] == 0:
|
||||
raise ValueError("Invalid IMAGE EBML integer")
|
||||
length = 9 - data[pos].bit_length()
|
||||
if length > (4 if identifier else 8) or pos + length > end:
|
||||
raise ValueError("Truncated IMAGE EBML integer")
|
||||
value = int.from_bytes(data[pos:pos + length], "big")
|
||||
if not identifier:
|
||||
value &= (1 << (7 * length)) - 1
|
||||
if value == (1 << (7 * length)) - 1:
|
||||
raise ValueError("Unknown-sized IMAGE EBML elements are forbidden")
|
||||
return value, pos + length
|
||||
def elements(start, end):
|
||||
nonlocal nodes
|
||||
while start < end:
|
||||
nodes += 1
|
||||
if nodes > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE container element count exceeds limit")
|
||||
tag, pos = vint(start, end, True)
|
||||
length, body = vint(pos, end)
|
||||
finish = body + length
|
||||
if finish > end:
|
||||
raise ValueError("Truncated IMAGE EBML element")
|
||||
yield tag, body, finish
|
||||
start = finish
|
||||
def integer(start, end):
|
||||
if not 1 <= end - start <= 8:
|
||||
raise ValueError("Invalid IMAGE EBML numeric field")
|
||||
return int.from_bytes(data[start:end], "big")
|
||||
roots = list(elements(0, len(data)))
|
||||
if len(roots) != 2 or [r[0] for r in roots] != [0x1a45dfa3, 0x18538067]:
|
||||
raise ValueError("Invalid IMAGE Matroska structure")
|
||||
tracks, blocks, packets = [], 0, []
|
||||
for tag, start, end in elements(roots[1][1], roots[1][2]):
|
||||
if tag == 0x1654ae6b:
|
||||
for entry, es, ee in elements(start, end):
|
||||
if entry in (0xbf, 0xec):
|
||||
continue
|
||||
if entry != 0xae:
|
||||
raise ValueError("Unexpected IMAGE track element")
|
||||
fields, dims = {}, {}
|
||||
for field, fs, fe in elements(es, ee):
|
||||
if field in fields:
|
||||
raise ValueError("Duplicate IMAGE track field")
|
||||
fields[field] = (fs, fe)
|
||||
if field == 0x6d80:
|
||||
raise ValueError("IMAGE track encodings are forbidden")
|
||||
if field == 0xe0:
|
||||
for dim, ds, de in elements(fs, fe):
|
||||
if dim in dims:
|
||||
raise ValueError("Duplicate IMAGE video field")
|
||||
dims[dim] = (ds, de)
|
||||
if any(key not in fields for key in (0xd7, 0x83, 0x86)) or any(key not in dims for key in (0xb0, 0xba)):
|
||||
raise ValueError("Incomplete IMAGE video track")
|
||||
if integer(*fields[0xd7]) != 1 or integer(*fields[0x83]) != 1:
|
||||
raise ValueError("Unsupported IMAGE video track")
|
||||
if fields[0x86][1] - fields[0x86][0] > 32:
|
||||
raise ValueError("IMAGE codec identifier exceeds limit")
|
||||
codec_id = bytes(data[slice(*fields[0x86])])
|
||||
if codec_id != _CODEC_IDS[fmt] and not (fmt == "ffv1" and codec_id == b"V_MS/VFW/FOURCC"):
|
||||
raise ValueError("IMAGE video codec mismatch")
|
||||
if integer(*dims[0xb0]) != width or integer(*dims[0xba]) != height:
|
||||
raise ValueError("IMAGE video header dimensions mismatch")
|
||||
if codec_id == b"V_MS/VFW/FOURCC":
|
||||
cs, ce = fields.get(0x63a2, (0, 0))
|
||||
if ce - cs < 40 or data[cs + 16:cs + 20] != b"FFV1":
|
||||
raise ValueError("Invalid IMAGE FFV1 codec header")
|
||||
if struct.unpack("<ii", data[cs + 4:cs + 12]) != (width, height):
|
||||
raise ValueError("IMAGE FFV1 header dimensions mismatch")
|
||||
tracks.append(fields)
|
||||
elif tag == 0x1f43b675:
|
||||
def block(bs, be):
|
||||
nonlocal blocks
|
||||
track, pos = vint(bs, be)
|
||||
if track != 1 or pos + 3 >= be or data[pos + 2] & 0x06:
|
||||
raise ValueError("Invalid or laced IMAGE video block")
|
||||
blocks += 1
|
||||
if blocks > count:
|
||||
raise ValueError("IMAGE video frame count exceeds descriptor")
|
||||
packets.append((pos + 3, be))
|
||||
for field, fs, fe in elements(start, end):
|
||||
if field == 0xa3:
|
||||
block(fs, fe)
|
||||
elif field == 0xa0:
|
||||
group = [e for e in elements(fs, fe) if e[0] == 0xa1]
|
||||
if len(group) != 1:
|
||||
raise ValueError("Invalid IMAGE video block group")
|
||||
block(group[0][1], group[0][2])
|
||||
elif field not in (0xe7, 0xec, 0xbf):
|
||||
raise ValueError("Unsupported IMAGE cluster field")
|
||||
elif tag not in (0x114d9b74, 0x1549a966, 0x1c53bb6b, 0x1254c367, 0xec, 0xbf):
|
||||
raise ValueError("Unsupported IMAGE Matroska element")
|
||||
if len(tracks) != 1 or blocks != count:
|
||||
raise ValueError("IMAGE video track or frame count mismatch")
|
||||
cs, ce = tracks[0].get(0x63a2, (0, 0))
|
||||
if not 0 < ce - cs <= MAX_DESCRIPTOR_BYTES:
|
||||
raise ValueError("IMAGE codec configuration exceeds limit or is missing")
|
||||
if bytes(data[slice(*tracks[0][0x86])]) == b"V_MS/VFW/FOURCC":
|
||||
cs += 40 # Skip BITMAPINFOHEADER; FFV1 receives only its codec config.
|
||||
return bytes(data[cs:ce]), packets
|
||||
|
||||
|
||||
def _coded_frame_bound(data, fmt, header):
|
||||
"""Conservative per-packet allocation bound before native decode.
|
||||
|
||||
FFV1 has one picture per packet. Every AVC/HEVC picture needs at least one
|
||||
VCL NAL; every AV1 output (including show-existing) needs a frame/header OBU.
|
||||
Counting slices overestimates pictures, which is safe. Do not trust a
|
||||
Matroska block count alone: a malicious packet may contain several pictures.
|
||||
"""
|
||||
config, packets = header
|
||||
if fmt == "ffv1":
|
||||
return len(packets)
|
||||
total, units = 0, 0
|
||||
if fmt in ("h264", "h265"):
|
||||
index = 4 if fmt == "h264" else 21
|
||||
if len(config) <= index or config[0] != 1:
|
||||
raise ValueError("Invalid IMAGE video configuration")
|
||||
length_size = (config[index] & 3) + 1
|
||||
for start, end in packets:
|
||||
while start < end:
|
||||
units += 1
|
||||
if units > MAX_CONTAINER_ELEMENTS:
|
||||
raise ValueError("IMAGE coded unit count exceeds limit")
|
||||
if fmt in ("h264", "h265"):
|
||||
if start + length_size > end:
|
||||
raise ValueError("Truncated IMAGE NAL length")
|
||||
length = int.from_bytes(data[start:start + length_size], "big")
|
||||
start += length_size
|
||||
if length < (1 if fmt == "h264" else 2) or length > end - start:
|
||||
raise ValueError("Truncated IMAGE NAL unit")
|
||||
kind = (data[start] & 31) if fmt == "h264" else (data[start] >> 1) & 63
|
||||
total += (1 <= kind <= 5) if fmt == "h264" else (kind <= 31)
|
||||
else:
|
||||
tag = data[start]
|
||||
start += 1
|
||||
if tag & 0x81 or not tag & 2:
|
||||
raise ValueError("Unsupported IMAGE AV1 OBU header")
|
||||
kind = (tag >> 3) & 15
|
||||
if tag & 4:
|
||||
start += 1
|
||||
length = 0
|
||||
for shift in range(0, 56, 7):
|
||||
if start >= end:
|
||||
raise ValueError("Truncated IMAGE AV1 OBU length")
|
||||
value = data[start]
|
||||
start += 1
|
||||
length |= (value & 127) << shift
|
||||
if not value & 128:
|
||||
break
|
||||
else:
|
||||
raise ValueError("IMAGE AV1 OBU length exceeds limit")
|
||||
if length > end - start:
|
||||
raise ValueError("Truncated IMAGE AV1 OBU")
|
||||
total += kind in (3, 6)
|
||||
start += length
|
||||
if total > MAX_FRAMES * 4:
|
||||
raise ValueError("IMAGE coded frame count exceeds limit")
|
||||
if total < len(packets):
|
||||
raise ValueError("IMAGE video has missing coded frames")
|
||||
return total
|
||||
|
||||
|
||||
def _decode_video(data, fmt, height, width, count, header):
|
||||
av = _import_av()
|
||||
maximum = ((height + 63) // 64 * 64) * ((width + 63) // 64 * 64)
|
||||
ctx = av.CodecContext.create("libdav1d" if fmt == "av1" else "hevc" if fmt == "h265" else fmt, "r")
|
||||
ctx.width, ctx.height = width, height
|
||||
ctx.extradata, packets = header
|
||||
ctx.thread_count = 1
|
||||
ctx.options = {"max_pixels": str(maximum), "threads": "1", "err_detect": "explode"}
|
||||
ctx.open()
|
||||
if "max_pixels" in ctx.options:
|
||||
raise ValueError("IMAGE decoder does not support max_pixels safety limit")
|
||||
if (ctx.width, ctx.height) != (width, height):
|
||||
raise ValueError("IMAGE bitstream dimensions mismatch")
|
||||
decoded = 0
|
||||
for bounds in packets + [None]:
|
||||
packet = av.Packet(data[slice(*bounds)]) if bounds is not None else None
|
||||
for frame in ctx.decode(packet):
|
||||
if decoded >= count or (frame.width, frame.height) != (width, height):
|
||||
raise ValueError("IMAGE decoded dimensions or frame count mismatch")
|
||||
decoded += 1
|
||||
yield frame.to_ndarray(format="rgb24")
|
||||
if decoded != count:
|
||||
raise ValueError("IMAGE decoded frame count mismatch")
|
||||
|
||||
|
||||
def decode_image(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Decode to an independent contiguous CPU tensor after bounded preflight."""
|
||||
limit = codec._byte_limit(max_bytes)
|
||||
if not isinstance(descriptor, dict):
|
||||
raise ValueError("IMAGE descriptor must be a dictionary")
|
||||
if descriptor.get("encoding") == "tensor-le-v1":
|
||||
return codec.decode_payload("IMAGE", descriptor, payload, max_bytes=limit)
|
||||
_descriptor_budget(descriptor)
|
||||
if set(descriptor) != {"encoding", "format", "quality", "shape", "layout", "dtype", "parts", "alpha"}:
|
||||
raise ValueError("Invalid IMAGE descriptor fields")
|
||||
if descriptor["encoding"] != "image-transfer-v1" or descriptor["layout"] != "BHWC":
|
||||
raise ValueError("Invalid IMAGE descriptor encoding or layout")
|
||||
fmt, _ = validate_options(descriptor["format"], descriptor["quality"])
|
||||
if fmt == "exact":
|
||||
raise ValueError("Exact IMAGE requires the tensor encoding")
|
||||
if not isinstance(payload, bytes):
|
||||
raise ValueError("IMAGE payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("IMAGE payload exceeds byte limit")
|
||||
shape = descriptor["shape"]
|
||||
if type(descriptor["dtype"]) is not str or descriptor["dtype"] not in codec._DTYPES:
|
||||
raise ValueError("Unsupported IMAGE descriptor dtype")
|
||||
dtype, itemsize = codec._DTYPES[descriptor["dtype"]]
|
||||
codec._tensor_size("IMAGE", shape, itemsize, limit)
|
||||
alpha_size, _ = _alpha_metadata(descriptor["alpha"], shape, limit)
|
||||
if alpha_size and descriptor["alpha"]["descriptor"]["dtype"] != descriptor["dtype"]:
|
||||
raise ValueError("IMAGE alpha dtype mismatch")
|
||||
ph, pw = _shape_budget(shape, fmt, limit, itemsize, alpha_size)
|
||||
count, height, width, channels = shape
|
||||
lengths = descriptor["parts"]
|
||||
if type(lengths) is not list or len(lengths) != (1 if fmt in _VIDEO_CODECS else count):
|
||||
raise ValueError("IMAGE part count mismatch")
|
||||
total = alpha_size
|
||||
for length in lengths:
|
||||
if type(length) is not int or not 0 < length <= limit - total:
|
||||
raise ValueError("IMAGE part length exceeds byte limit")
|
||||
total += length
|
||||
if total != len(payload):
|
||||
raise ValueError("IMAGE payload length does not match parts")
|
||||
parts, offset = [], 0
|
||||
for length in lengths:
|
||||
parts.append(memoryview(payload)[offset:offset + length])
|
||||
offset += length
|
||||
try:
|
||||
if fmt in _VIDEO_CODECS:
|
||||
video_header = _video_header(parts[0], fmt, ph, pw, count)
|
||||
coded_frames = _coded_frame_bound(parts[0], fmt, video_header)
|
||||
_shape_budget(shape, fmt, limit, itemsize, alpha_size, coded_frames)
|
||||
else:
|
||||
for part in parts:
|
||||
if _still_header(part, fmt) != (height, width):
|
||||
raise ValueError("IMAGE header dimensions mismatch")
|
||||
# All headers and aggregate budgets are checked before output allocation.
|
||||
alpha = None
|
||||
if alpha_size:
|
||||
alpha = codec.decode_payload("IMAGE", descriptor["alpha"]["descriptor"], payload[offset:], max_bytes=limit)
|
||||
result = torch.empty(shape, dtype=dtype, device="cpu")
|
||||
if fmt in _VIDEO_CODECS:
|
||||
frames = _decode_video(parts[0], fmt, ph, pw, count, video_header)
|
||||
else:
|
||||
def stills():
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
for part in parts:
|
||||
with Image.open(io.BytesIO(part)) as image:
|
||||
if image.size != (width, height) or getattr(image, "n_frames", 1) != 1:
|
||||
raise ValueError("IMAGE still dimensions or frame count mismatch")
|
||||
if image.mode not in ("L", "RGB"):
|
||||
raise ValueError("Unsupported IMAGE still pixel mode")
|
||||
yield np.array(image.convert("RGB"), dtype=np.uint8)
|
||||
frames = stills()
|
||||
color = 3 if channels == 4 else channels
|
||||
for index, frame in enumerate(frames):
|
||||
# Float32 division followed by the source dtype cast reproduces
|
||||
# Collector normalized pixels without narrowing exact alpha.
|
||||
result[index, ..., :color] = torch.from_numpy(frame[:height, :width, :color].copy()).float().div_(255)
|
||||
if alpha is not None:
|
||||
result[..., 3:4] = alpha
|
||||
return result
|
||||
except ValueError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ValueError("Malformed or unsupported IMAGE compressed payload") from None
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Image-return policy shared by the master and standalone endpoint bridge.
|
||||
|
||||
This module is deliberately dependency-free so graph compilation does not import
|
||||
Torch, Pillow, PyAV or ComfyUI. Existing Collector settings stay unchanged.
|
||||
"""
|
||||
|
||||
FORMATS = ("exact", "legacy_png", "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
|
||||
DEFAULT_FORMAT = "exact"
|
||||
DEFAULT_QUALITY = 90
|
||||
FORMAT_TOOLTIP = (
|
||||
"Image returns only; master-local images and other media stay unchanged. "
|
||||
"exact: original floating-point tensor, no quantization. "
|
||||
"legacy_png / png / ffv1: lossless 8-bit pixels, not exact float tensors. "
|
||||
"jpeg / webp: lossy stills; h264 / h265 / av1: lossy frame-batch video. "
|
||||
"Alpha is preserved separately when the codec cannot carry it. "
|
||||
"Requires matching transfer support on the worker/endpoint."
|
||||
)
|
||||
QUALITY_TOOLTIP = (
|
||||
"Quality 1–100 for lossy image returns; ignored by exact and lossless modes. "
|
||||
"JPEG/WebP use it directly; video formats use the existing Collector CRF mapping."
|
||||
)
|
||||
|
||||
|
||||
def validate_options(transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
"""Reject malformed controls rather than repairing or silently substituting."""
|
||||
if type(transfer_format) is not str or transfer_format not in FORMATS:
|
||||
raise ValueError("Unsupported image transfer_format")
|
||||
if type(transfer_quality) is not int or not 1 <= transfer_quality <= 100:
|
||||
raise ValueError("Image transfer_quality must be an integer from 1 to 100")
|
||||
return transfer_format, transfer_quality
|
||||
@@ -0,0 +1,175 @@
|
||||
"""Internal ComfyUI boundary nodes for the standalone media-v1 handler."""
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from .codec import decode_payload, encode_payload
|
||||
from .image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY,
|
||||
FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options)
|
||||
from .contract import (
|
||||
MAX_ARTIFACT_BYTES, MAX_ITEMS, MAX_TOTAL_BYTES, SOCKET_KINDS,
|
||||
bounded_int, identifier, validate_binding, validate_json,
|
||||
)
|
||||
|
||||
|
||||
def _scalar(value):
|
||||
if isinstance(value, list):
|
||||
if len(value) != 1:
|
||||
raise ValueError("Media bridge control must have exactly one value")
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
|
||||
def _json(text):
|
||||
if not isinstance(text, str) or len(text) > 256 * 1024:
|
||||
raise ValueError("Invalid media bridge JSON")
|
||||
try:
|
||||
value = json.loads(text)
|
||||
except (ValueError, TypeError, RecursionError):
|
||||
raise ValueError("Invalid media bridge JSON") from None
|
||||
validate_json(value)
|
||||
return value
|
||||
|
||||
|
||||
def _cpu(kind, value):
|
||||
# Runtime transfer only, never sampling or device discovery. Pure codecs
|
||||
# deliberately reject non-CPU storage rather than making hidden transfers.
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
return value.detach().cpu()
|
||||
if kind == "AUDIO":
|
||||
return {"waveform": value["waveform"].detach().cpu(), "sample_rate": value["sample_rate"]}
|
||||
return value
|
||||
|
||||
|
||||
def encode_value(kind, value, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if kind == "IMAGE" and fmt != DEFAULT_FORMAT:
|
||||
from .image_codec import encode_image
|
||||
return encode_image(_cpu(kind, value), transfer_format=fmt, transfer_quality=quality)
|
||||
if kind == "VIDEO":
|
||||
from .video_codec import encode_video_cpu
|
||||
return encode_video_cpu(value)
|
||||
return encode_payload(kind, _cpu(kind, value))
|
||||
|
||||
|
||||
def decode_value(kind, descriptor, payload):
|
||||
if kind == "VIDEO":
|
||||
from .video_codec import decode_video
|
||||
return decode_video(descriptor, payload)
|
||||
return decode_payload(kind, descriptor, payload)
|
||||
|
||||
|
||||
def export_values(request_id, binding, values, *, transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
"""Write bounded owned payloads; history contains metadata, never base64."""
|
||||
identifier(request_id)
|
||||
validate_binding(binding)
|
||||
fmt, quality = validate_options(transfer_format, transfer_quality)
|
||||
if binding["kind"] != "IMAGE" and fmt != DEFAULT_FORMAT:
|
||||
raise ValueError("Image transfer controls apply only to IMAGE exports")
|
||||
if not isinstance(values, list) or not 1 <= len(values) <= MAX_ITEMS:
|
||||
raise ValueError("Connected media output has no values or exceeds item limit")
|
||||
request_key = hashlib.sha256(request_id.encode()).hexdigest()
|
||||
subfolder = f"distributed-media-v1/{request_key}"
|
||||
root = Path(folder_paths.get_temp_directory()).resolve()
|
||||
directory = (root / subfolder).resolve()
|
||||
if not directory.is_relative_to(root):
|
||||
raise ValueError("Unsafe media output directory")
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
written, artifacts, items = [], [], []
|
||||
total = 0
|
||||
try:
|
||||
for index, value in enumerate(values):
|
||||
descriptor, payload = encode_value(binding["kind"], value,
|
||||
transfer_format=fmt, transfer_quality=quality)
|
||||
validate_json(descriptor)
|
||||
total += len(payload)
|
||||
if total > MAX_TOTAL_BYTES:
|
||||
raise ValueError("Media output exceeds total byte limit")
|
||||
aid = uuid.uuid4().hex
|
||||
filename = aid + ".bin"
|
||||
path = directory / filename
|
||||
with path.open("xb") as stream:
|
||||
written.append(path)
|
||||
stream.write(payload)
|
||||
artifacts.append({"id": aid, "kind": binding["kind"], "descriptor": descriptor,
|
||||
"byte_length": len(payload), "sha256": hashlib.sha256(payload).hexdigest(),
|
||||
"transport": {"type": "comfy-file", "filename": filename,
|
||||
"subfolder": subfolder, "folder_type": "temp"}})
|
||||
items.append({"binding_id": binding["binding_id"], "item_index": index, "artifact_id": aid})
|
||||
except BaseException:
|
||||
for path in written:
|
||||
path.unlink(missing_ok=True)
|
||||
raise
|
||||
return {"request_id": request_id, "binding": dict(binding, count=len(values)),
|
||||
"artifacts": artifacts, "items": items}
|
||||
|
||||
|
||||
def import_value(artifact, filename):
|
||||
"""Read a handler-staged, checksummed file within native media roots only."""
|
||||
fields = {"id", "kind", "descriptor", "byte_length", "sha256"}
|
||||
if not isinstance(artifact, dict) or set(artifact) != fields:
|
||||
raise ValueError("Invalid staged media metadata")
|
||||
identifier(artifact["id"])
|
||||
if artifact["kind"] not in SOCKET_KINDS.values():
|
||||
raise ValueError("Invalid staged media kind")
|
||||
size = bounded_int(artifact["byte_length"], MAX_ARTIFACT_BYTES)
|
||||
if not isinstance(filename, str):
|
||||
raise ValueError("Invalid staged media filename")
|
||||
roots = [Path(getter()).resolve() for getter in (
|
||||
folder_paths.get_input_directory, folder_paths.get_output_directory, folder_paths.get_temp_directory)]
|
||||
path = Path(folder_paths.get_annotated_filepath(filename)).resolve()
|
||||
if not any(path.is_relative_to(root) for root in roots) or not path.is_file():
|
||||
raise ValueError("Staged media file is outside allowed roots or unavailable")
|
||||
if path.stat().st_size != size:
|
||||
raise ValueError("Staged media file length mismatch")
|
||||
with path.open("rb") as stream:
|
||||
payload = stream.read(size + 1)
|
||||
if len(payload) != size or hashlib.sha256(payload).hexdigest() != artifact["sha256"]:
|
||||
raise ValueError("Staged media file failed integrity validation")
|
||||
return decode_value(artifact["kind"], artifact["descriptor"], payload)
|
||||
|
||||
|
||||
class MediaImportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaImportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("artifact_json"), io.String.Input("filename")],
|
||||
outputs=[io.Image.Output(), io.Audio.Output(), io.Mask.Output(), io.String.Output(), io.Video.Output()])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, artifact_json, filename):
|
||||
artifact = _json(artifact_json)
|
||||
value = import_value(artifact, filename)
|
||||
return io.NodeOutput(*(value if kind == artifact["kind"] else None for kind in SOCKET_KINDS.values()))
|
||||
|
||||
|
||||
class MediaExportV1(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="MediaExportV1", category="distributed/internal", is_deprecated=True,
|
||||
inputs=[io.String.Input("request_id"), io.String.Input("binding_json"),
|
||||
io.Image.Input("images", optional=True), io.Audio.Input("audio", optional=True),
|
||||
io.Mask.Input("masks", optional=True), io.String.Input("texts", optional=True, force_input=True),
|
||||
io.Video.Input("video", optional=True),
|
||||
io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT,
|
||||
optional=True, tooltip=FORMAT_TOOLTIP),
|
||||
io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100,
|
||||
optional=True, tooltip=QUALITY_TOOLTIP)],
|
||||
outputs=[], is_output_node=True, is_input_list=True)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, request_id, binding_json, images=None, audio=None, masks=None, texts=None, video=None,
|
||||
transfer_format=DEFAULT_FORMAT, transfer_quality=DEFAULT_QUALITY):
|
||||
binding = _json(_scalar(binding_json))
|
||||
validate_binding(binding)
|
||||
provided = dict(images=images, audio=audio, masks=masks, texts=texts, video=video)
|
||||
if {name for name, value in provided.items() if value is not None} != {binding["socket"]}:
|
||||
raise ValueError("Media export must receive exactly its declared typed socket")
|
||||
exported = export_values(_scalar(request_id), binding, provided[binding["socket"]],
|
||||
transfer_format=_scalar(transfer_format),
|
||||
transfer_quality=_scalar(transfer_quality))
|
||||
return io.NodeOutput(ui={"media_exports": [exported]})
|
||||
@@ -0,0 +1,482 @@
|
||||
"""Lossless transport of pinned native ComfyUI VIDEO representations.
|
||||
|
||||
No save_to/get_components fallback is allowed: an unknown class must not silently
|
||||
transcode. Tensor components use the shared CPU codec. File bytes are opaque and
|
||||
retain signed trim/crop parameters; decoded files own independent BytesIOs.
|
||||
VideoFromList's native constructor flattens its inputs, so wire segments are a
|
||||
bounded, flat list of exact VideoFromFile objects, never recursive playlists.
|
||||
|
||||
Offsets are relative to the containing VIDEO payload. Canonical part order is
|
||||
images/alpha/audio, or segments/complete_audio. Every byte is covered exactly
|
||||
once. All descriptor structure and tensor sizes are preflighted before decoding
|
||||
any media storage. File containers are not parsed until a native consumer reads
|
||||
them; this adapter does not attempt to validate or repair container codecs.
|
||||
"""
|
||||
from fractions import Fraction
|
||||
import io
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import stat
|
||||
|
||||
from .codec import MAX_PAYLOAD_BYTES, decode_payload, encode_payload
|
||||
|
||||
|
||||
MAX_SEGMENTS = 64
|
||||
MAX_METADATA_BYTES = 64 * 1024
|
||||
MAX_METADATA_DEPTH = 16
|
||||
_ENCODING = "native-video-v1"
|
||||
_BASE_FIELDS = {"encoding", "representation", "byte_length"}
|
||||
_COMPONENT_FIELDS = {"images", "alpha", "audio", "frame_rate", "metadata", "bit_depth", "color_space"}
|
||||
_FILE_FIELDS = {"data", "start_time", "duration", "crop"}
|
||||
_SEGMENT_FIELDS = {"segments", "complete_audio"}
|
||||
_TENSOR_BYTES = {"float16": 2, "bfloat16": 2, "float32": 4, "float64": 8}
|
||||
|
||||
|
||||
def _limit(max_bytes):
|
||||
if type(max_bytes) is not int or max_bytes < 0:
|
||||
raise ValueError("VIDEO max_bytes must be a nonnegative integer")
|
||||
return min(max_bytes, MAX_PAYLOAD_BYTES)
|
||||
|
||||
|
||||
def _native_types():
|
||||
# Lazy imports keep non-VIDEO users independent of ComfyUI and PyAV.
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
return VideoFromComponents, VideoFromFile, VideoFromList, VideoComponents
|
||||
|
||||
|
||||
def _fields(value, fields, label):
|
||||
if type(value) is not dict or set(value) != fields:
|
||||
raise ValueError(f"Invalid VIDEO {label} fields")
|
||||
|
||||
|
||||
def _integer(value, *, positive=False):
|
||||
if type(value) is not int or value < (1 if positive else 0):
|
||||
raise ValueError("VIDEO length, offset and dimensions must be integers in range")
|
||||
return value
|
||||
|
||||
|
||||
def _number(value):
|
||||
if type(value) not in (int, float) or (type(value) is float and not math.isfinite(value)):
|
||||
raise ValueError("VIDEO trim values must be finite numbers")
|
||||
return value
|
||||
|
||||
|
||||
def _crop(value, *, wire=False):
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not (list if wire else tuple) or len(value) != 4:
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
if any(type(part) is not int for part in value):
|
||||
raise ValueError("VIDEO crop must contain four integers")
|
||||
# Native crop normalization is a consumer concern, not a transport operation.
|
||||
return list(value) if not wire else tuple(value)
|
||||
|
||||
|
||||
def _frame_rate(value):
|
||||
_fields(value, {"numerator", "denominator"}, "frame_rate")
|
||||
_integer(value["numerator"], positive=True)
|
||||
_integer(value["denominator"], positive=True)
|
||||
|
||||
|
||||
def _encoding_options(bit_depth, color_space):
|
||||
_integer(bit_depth, positive=True)
|
||||
if type(color_space) is not str or color_space not in ("sRGB", "HDR", "HDR PQ"):
|
||||
raise ValueError("Unsupported native VIDEO color_space")
|
||||
|
||||
|
||||
def _metadata(value):
|
||||
"""Validate bounded JSON before allocating its serialized form; return a copy."""
|
||||
if value is None:
|
||||
return None
|
||||
if type(value) is not dict:
|
||||
raise ValueError("VIDEO metadata must be a JSON object or null")
|
||||
remaining = MAX_METADATA_BYTES
|
||||
|
||||
def charge(size):
|
||||
nonlocal remaining
|
||||
remaining -= size
|
||||
if remaining < 0:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
|
||||
def visit(item, depth):
|
||||
if depth > MAX_METADATA_DEPTH:
|
||||
raise ValueError("VIDEO metadata exceeds depth limit")
|
||||
kind = type(item)
|
||||
if kind in (dict, list):
|
||||
# Each member costs at least one JSON byte. Also bounds iteration
|
||||
# before walking potentially huge containers or cyclic structures.
|
||||
if len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
charge(2 + max(0, len(item) - 1))
|
||||
if kind is dict:
|
||||
for key, child in item.items():
|
||||
if type(key) is not str:
|
||||
raise ValueError("VIDEO metadata keys must be strings")
|
||||
visit(key, depth + 1)
|
||||
charge(1)
|
||||
visit(child, depth + 1)
|
||||
else:
|
||||
for child in item:
|
||||
visit(child, depth + 1)
|
||||
elif kind in (str, int, float, bool) or item is None:
|
||||
if kind is str and len(item) > remaining:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
# Avoid converting an unbounded Python integer to decimal first.
|
||||
if kind is int and item.bit_length() > MAX_METADATA_BYTES * 4:
|
||||
raise ValueError("VIDEO metadata exceeds byte limit")
|
||||
if kind is float and not math.isfinite(item):
|
||||
raise ValueError("VIDEO metadata numbers must be finite")
|
||||
try:
|
||||
charge(len(json.dumps(item, ensure_ascii=True, allow_nan=False)))
|
||||
except (ValueError, OverflowError):
|
||||
raise ValueError("Invalid or oversized VIDEO metadata scalar") from None
|
||||
else:
|
||||
raise ValueError("VIDEO metadata must contain only JSON values")
|
||||
|
||||
visit(value, 0)
|
||||
return json.loads(json.dumps(value, ensure_ascii=True, allow_nan=False, separators=(",", ":")))
|
||||
|
||||
|
||||
def _attrs(value, names):
|
||||
attributes = vars(value)
|
||||
if not all(name in attributes for name in names):
|
||||
raise ValueError("Native VIDEO attributes do not match the pinned adapter")
|
||||
return [attributes[name] for name in names]
|
||||
|
||||
|
||||
def _roots(allowed_roots):
|
||||
if allowed_roots is None:
|
||||
import folder_paths
|
||||
allowed_roots = [folder_paths.get_input_directory(), folder_paths.get_output_directory(),
|
||||
folder_paths.get_temp_directory()]
|
||||
if type(allowed_roots) not in (list, tuple):
|
||||
raise ValueError("VIDEO allowed_roots must be a list or tuple")
|
||||
try:
|
||||
roots = [Path(root).resolve() for root in allowed_roots]
|
||||
except (OSError, TypeError, ValueError, RuntimeError):
|
||||
raise ValueError("Invalid VIDEO allowed roots") from None
|
||||
return roots
|
||||
|
||||
|
||||
def _read_path(value, limit, allowed_roots):
|
||||
"""Resolve authorization, reject special files, and bound all reads.
|
||||
|
||||
Open each resolved path component with NOFOLLOW using directory fds, so an
|
||||
intermediate symlink swapped in after resolve() cannot bypass authorization.
|
||||
NONBLOCK prevents an attacker swapping in a FIFO from hanging the worker.
|
||||
"""
|
||||
roots = _roots(allowed_roots)
|
||||
directory_fd = file_fd = None
|
||||
try:
|
||||
path = Path(value).resolve(strict=True)
|
||||
if not any(path.is_relative_to(root) for root in roots):
|
||||
raise ValueError("VIDEO file is outside allowed roots")
|
||||
directory_flags = os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW
|
||||
directory_fd = os.open(path.anchor, directory_flags)
|
||||
for part in path.parts[1:-1]:
|
||||
next_fd = os.open(part, directory_flags, dir_fd=directory_fd)
|
||||
os.close(directory_fd)
|
||||
directory_fd = next_fd
|
||||
file_fd = os.open(path.name, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=directory_fd)
|
||||
info = os.fstat(file_fd)
|
||||
if not stat.S_ISREG(info.st_mode):
|
||||
raise ValueError("VIDEO source must be a regular file")
|
||||
if info.st_size > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
# Read at most the stat size plus a sentinel; reject growth or truncation.
|
||||
chunks = []
|
||||
count = 0
|
||||
while count <= info.st_size:
|
||||
chunk = os.read(file_fd, min(1024 * 1024, info.st_size + 1 - count))
|
||||
if not chunk:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
count += len(chunk)
|
||||
if count != info.st_size:
|
||||
raise ValueError("VIDEO source changed size during bounded read")
|
||||
return b"".join(chunks)
|
||||
except (OSError, RuntimeError):
|
||||
raise ValueError("VIDEO file cannot be read safely") from None
|
||||
finally:
|
||||
if file_fd is not None:
|
||||
os.close(file_fd)
|
||||
if directory_fd is not None:
|
||||
os.close(directory_fd)
|
||||
|
||||
|
||||
def _file_bytes(value, limit, allowed_roots):
|
||||
if type(value) is io.BytesIO:
|
||||
try:
|
||||
with value.getbuffer() as view:
|
||||
if view.nbytes > limit:
|
||||
raise ValueError("VIDEO file exceeds byte limit")
|
||||
return bytes(view)
|
||||
except ValueError:
|
||||
raise ValueError("VIDEO BytesIO is closed or exceeds byte limit") from None
|
||||
if isinstance(value, (str, os.PathLike)):
|
||||
return _read_path(value, limit, allowed_roots)
|
||||
raise ValueError("VIDEO file must be a local path or exact BytesIO")
|
||||
|
||||
|
||||
def encode_video(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Encode an exact native VIDEO without opening/transcoding its container.
|
||||
|
||||
Paths are authorized only under allowed_roots (when supplied), otherwise
|
||||
under ComfyUI input/output/temp roots. Budgets may lower, never raise, the
|
||||
shared maximum. Unsupported native state raises ValueError, not a fallback.
|
||||
Non-CPU tensors are rejected; runtime staging belongs to the bridge boundary.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=False)
|
||||
|
||||
|
||||
def encode_video_cpu(value, *, max_bytes=MAX_PAYLOAD_BYTES, allowed_roots=None):
|
||||
"""Bridge-only boundary: detach/stage native tensors, then losslessly encode.
|
||||
|
||||
Share the exact-type adapter with encode_video rather than reconstructing
|
||||
native videos, calling get_components, or duplicating private attributes.
|
||||
Images, alpha and both audio representations move without dtype conversion;
|
||||
metadata/options and opaque file bytes/trim/crop take the unchanged path.
|
||||
"""
|
||||
return _encode_video(value, max_bytes=max_bytes, allowed_roots=allowed_roots, stage_cpu=True)
|
||||
|
||||
|
||||
def _stage_tensor_cpu(value):
|
||||
import torch
|
||||
if not isinstance(value, torch.Tensor):
|
||||
raise ValueError("Tensor value must be a torch.Tensor")
|
||||
return value.detach().cpu()
|
||||
|
||||
|
||||
def _encode_video(value, *, max_bytes, allowed_roots, stage_cpu):
|
||||
limit = _limit(max_bytes)
|
||||
components_class, file_class, list_class, components_type = _native_types()
|
||||
descriptor = {"encoding": _ENCODING}
|
||||
chunks = []
|
||||
size = 0
|
||||
|
||||
def add(payload, nested=None):
|
||||
nonlocal size
|
||||
if len(payload) > limit - size:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
slot = {"offset": size, "length": len(payload)}
|
||||
if nested is not None:
|
||||
slot["descriptor"] = nested
|
||||
chunks.append(payload)
|
||||
size += len(payload)
|
||||
return slot
|
||||
|
||||
def tensor(kind, tensor_value):
|
||||
if stage_cpu:
|
||||
if kind == "AUDIO":
|
||||
if not isinstance(tensor_value, dict) or set(tensor_value) != {"waveform", "sample_rate"}:
|
||||
raise ValueError("AUDIO value must contain waveform and sample_rate")
|
||||
tensor_value = {**tensor_value, "waveform": _stage_tensor_cpu(tensor_value["waveform"])}
|
||||
else:
|
||||
tensor_value = _stage_tensor_cpu(tensor_value)
|
||||
nested, payload = encode_payload(kind, tensor_value, max_bytes=limit - size)
|
||||
return add(payload, nested)
|
||||
|
||||
if type(value) is components_class:
|
||||
component, bit_depth, color_space = _attrs(value, [
|
||||
"_VideoFromComponents__components", "_VideoFromComponents__bit_depth",
|
||||
"_VideoFromComponents__color_space",
|
||||
])
|
||||
if type(component) is not components_type:
|
||||
raise ValueError("Unsupported native VIDEO components class")
|
||||
images, fps, audio, metadata, alpha = _attrs(component, [
|
||||
"images", "frame_rate", "audio", "metadata", "alpha",
|
||||
])
|
||||
if type(fps) is not Fraction or fps <= 0:
|
||||
raise ValueError("VIDEO frame_rate must be a positive Fraction")
|
||||
_encoding_options(bit_depth, color_space)
|
||||
descriptor.update(representation="components", bit_depth=bit_depth, color_space=color_space,
|
||||
frame_rate={"numerator": fps.numerator, "denominator": fps.denominator},
|
||||
metadata=_metadata(metadata))
|
||||
descriptor["images"] = tensor("IMAGE", images)
|
||||
# Native get_components() returns BHW1 alpha; MASK nodes use BHW.
|
||||
# Select an existing lossless tensor layout without squeezing/converting.
|
||||
alpha_kind = "IMAGE" if getattr(alpha, "ndim", None) == 4 else "MASK"
|
||||
descriptor["alpha"] = tensor(alpha_kind, alpha) if alpha is not None else None
|
||||
if alpha is not None:
|
||||
_alpha_header(descriptor["alpha"], descriptor["images"])
|
||||
descriptor["audio"] = tensor("AUDIO", audio) if audio is not None else None
|
||||
elif type(value) is file_class:
|
||||
source, start, duration, crop = _attrs(value, [
|
||||
"_VideoFromFile__file", "_VideoFromFile__start_time", "_VideoFromFile__duration",
|
||||
"_VideoFromFile__crop",
|
||||
])
|
||||
descriptor.update(representation="file", start_time=_number(start), duration=_number(duration),
|
||||
crop=_crop(crop))
|
||||
descriptor["data"] = add(_file_bytes(source, limit, allowed_roots))
|
||||
elif type(value) is list_class:
|
||||
videos, complete_audio = _attrs(value, ["videos", "complete_audio"])
|
||||
if type(videos) is not list or not 1 <= len(videos) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
if any(type(segment) is not file_class for segment in videos):
|
||||
raise ValueError("VIDEO segments must be exact native VideoFromFile objects")
|
||||
descriptor.update(representation="segments", segments=[])
|
||||
for segment in videos:
|
||||
nested, payload = encode_video(segment, max_bytes=limit - size, allowed_roots=allowed_roots)
|
||||
descriptor["segments"].append(add(payload, nested))
|
||||
descriptor["complete_audio"] = tensor("AUDIO", complete_audio) if complete_audio is not None else None
|
||||
else:
|
||||
raise ValueError("Unsupported native VIDEO class")
|
||||
descriptor["byte_length"] = size
|
||||
return descriptor, b"".join(chunks)
|
||||
|
||||
|
||||
def _tensor_header(kind, descriptor, length):
|
||||
"""Allocation-free preflight of the shared tensor-le-v1 descriptor.
|
||||
|
||||
The shared decoder still validates the bytes/finite values. Doing its small
|
||||
structural check here first prevents allocation for an early part when a
|
||||
later part's descriptor is already known to be invalid.
|
||||
"""
|
||||
fields = {"encoding", "dtype", "shape", "layout"}
|
||||
if kind == "AUDIO":
|
||||
fields.add("sample_rate")
|
||||
_fields(descriptor, fields, "tensor")
|
||||
layout = {"IMAGE": "BHWC", "MASK": "BHW", "AUDIO": "BCT"}[kind]
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != "tensor-le-v1":
|
||||
raise ValueError("Invalid VIDEO tensor encoding")
|
||||
if type(descriptor["layout"]) is not str or descriptor["layout"] != layout:
|
||||
raise ValueError("Invalid VIDEO tensor layout")
|
||||
dtype = descriptor["dtype"]
|
||||
if type(dtype) is not str or dtype not in _TENSOR_BYTES:
|
||||
raise ValueError("Invalid VIDEO tensor dtype")
|
||||
shape = descriptor["shape"]
|
||||
if type(shape) is not list or len(shape) != len(layout):
|
||||
raise ValueError("Invalid VIDEO tensor rank")
|
||||
size = _TENSOR_BYTES[dtype]
|
||||
for dimension in shape:
|
||||
_integer(dimension, positive=True)
|
||||
if dimension > length // size:
|
||||
raise ValueError("VIDEO tensor shape exceeds byte limit")
|
||||
size *= dimension
|
||||
if kind == "IMAGE" and shape[-1] not in (1, 3, 4):
|
||||
raise ValueError("Invalid VIDEO image channels")
|
||||
if size != length:
|
||||
raise ValueError("VIDEO tensor length does not match shape")
|
||||
if kind == "AUDIO":
|
||||
_integer(descriptor["sample_rate"], positive=True)
|
||||
|
||||
|
||||
def _alpha_header(alpha, images):
|
||||
_fields(alpha, {"offset", "length", "descriptor"}, "alpha part")
|
||||
descriptor = alpha["descriptor"]
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("Invalid VIDEO alpha descriptor")
|
||||
layout = descriptor.get("layout")
|
||||
if type(layout) is not str or layout not in ("BHW", "BHWC"):
|
||||
raise ValueError("Invalid VIDEO alpha layout")
|
||||
kind = "MASK" if layout == "BHW" else "IMAGE"
|
||||
_tensor_header(kind, descriptor, _integer(alpha["length"]))
|
||||
expected = images["descriptor"]["shape"][:3] + ([1] if kind == "IMAGE" else [])
|
||||
if descriptor["shape"] != expected:
|
||||
raise ValueError("VIDEO alpha shape must match images")
|
||||
return kind
|
||||
|
||||
|
||||
def _validate_video(descriptor, size, *, file_only=False):
|
||||
if type(descriptor) is not dict:
|
||||
raise ValueError("VIDEO descriptor must be a dictionary")
|
||||
representation = descriptor.get("representation")
|
||||
if type(representation) is not str or representation not in ("components", "file", "segments"):
|
||||
raise ValueError("Unsupported VIDEO representation")
|
||||
if file_only and representation != "file":
|
||||
# Depth is at most one: native lists already flatten their segments.
|
||||
raise ValueError("VIDEO segments cannot contain nested lists or components")
|
||||
fields = {"components": _COMPONENT_FIELDS, "file": _FILE_FIELDS, "segments": _SEGMENT_FIELDS}[representation]
|
||||
_fields(descriptor, _BASE_FIELDS | fields, "descriptor")
|
||||
if type(descriptor["encoding"]) is not str or descriptor["encoding"] != _ENCODING:
|
||||
raise ValueError("Unsupported VIDEO encoding")
|
||||
if _integer(descriptor["byte_length"]) != size:
|
||||
raise ValueError("VIDEO payload length does not match descriptor")
|
||||
result = dict(descriptor)
|
||||
cursor = 0
|
||||
|
||||
def part(slot, kind=None, video=False):
|
||||
nonlocal cursor
|
||||
fields = {"offset", "length", "descriptor"} if kind or video else {"offset", "length"}
|
||||
_fields(slot, fields, "part")
|
||||
offset = _integer(slot["offset"])
|
||||
length = _integer(slot["length"])
|
||||
if offset != cursor or length > size - cursor:
|
||||
raise ValueError("VIDEO parts must not overlap, leave gaps or exceed bounds")
|
||||
cursor += length
|
||||
if kind:
|
||||
_tensor_header(kind, slot["descriptor"], length)
|
||||
if video:
|
||||
return {**slot, "descriptor": _validate_video(slot["descriptor"], length, file_only=True)}
|
||||
return slot
|
||||
|
||||
if representation == "components":
|
||||
_frame_rate(descriptor["frame_rate"])
|
||||
_encoding_options(descriptor["bit_depth"], descriptor["color_space"])
|
||||
result["metadata"] = _metadata(descriptor["metadata"])
|
||||
part(descriptor["images"], "IMAGE")
|
||||
if descriptor["alpha"] is not None:
|
||||
part(descriptor["alpha"], _alpha_header(descriptor["alpha"], descriptor["images"]))
|
||||
if descriptor["audio"] is not None:
|
||||
part(descriptor["audio"], "AUDIO")
|
||||
elif representation == "file":
|
||||
_number(descriptor["start_time"])
|
||||
_number(descriptor["duration"])
|
||||
_crop(descriptor["crop"], wire=True)
|
||||
part(descriptor["data"])
|
||||
else:
|
||||
segments = descriptor["segments"]
|
||||
if type(segments) is not list or not 1 <= len(segments) <= MAX_SEGMENTS:
|
||||
raise ValueError("VIDEO must have between 1 and 64 segments")
|
||||
result["segments"] = [part(segment, video=True) for segment in segments]
|
||||
if descriptor["complete_audio"] is not None:
|
||||
part(descriptor["complete_audio"], "AUDIO")
|
||||
if cursor != size:
|
||||
raise ValueError("VIDEO payload contains trailing bytes")
|
||||
return result
|
||||
|
||||
|
||||
def _decode_validated(descriptor, payload, types):
|
||||
components_class, file_class, list_class, components_type = types
|
||||
|
||||
def data(slot):
|
||||
return payload[slot["offset"]:slot["offset"] + slot["length"]]
|
||||
|
||||
def tensor(kind, slot):
|
||||
if slot is None:
|
||||
return None
|
||||
return decode_payload(kind, slot["descriptor"], data(slot), max_bytes=slot["length"])
|
||||
|
||||
if descriptor["representation"] == "file":
|
||||
# VideoFromFile retains the BytesIO, unlike a temporary filesystem path.
|
||||
return file_class(io.BytesIO(data(descriptor["data"])), start_time=descriptor["start_time"],
|
||||
duration=descriptor["duration"], crop=_crop(descriptor["crop"], wire=True))
|
||||
if descriptor["representation"] == "components":
|
||||
fps = descriptor["frame_rate"]
|
||||
return components_class(components_type(
|
||||
images=tensor("IMAGE", descriptor["images"]),
|
||||
frame_rate=Fraction(fps["numerator"], fps["denominator"]),
|
||||
audio=tensor("AUDIO", descriptor["audio"]),
|
||||
metadata=descriptor["metadata"],
|
||||
alpha=tensor("IMAGE" if descriptor["alpha"] is not None and
|
||||
descriptor["alpha"]["descriptor"]["layout"] == "BHWC" else "MASK",
|
||||
descriptor["alpha"]),
|
||||
), bit_depth=descriptor["bit_depth"], color_space=descriptor["color_space"])
|
||||
videos = [_decode_validated(slot["descriptor"], data(slot), types) for slot in descriptor["segments"]]
|
||||
return list_class(videos, complete_audio=tensor("AUDIO", descriptor["complete_audio"]))
|
||||
|
||||
|
||||
def decode_video(descriptor, payload, *, max_bytes=MAX_PAYLOAD_BYTES):
|
||||
"""Decode a bounded VIDEO into independently owned native CPU media objects."""
|
||||
limit = _limit(max_bytes)
|
||||
if type(payload) is not bytes:
|
||||
raise ValueError("VIDEO payload must be bytes")
|
||||
if len(payload) > limit:
|
||||
raise ValueError("VIDEO payload exceeds byte limit")
|
||||
validated = _validate_video(descriptor, len(payload))
|
||||
return _decode_validated(validated, payload, _native_types())
|
||||
@@ -1,163 +0,0 @@
|
||||
"""Shared stubs/utilities for API route module unit tests."""
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
class RoutesStub:
|
||||
"""Simple decorator-style route registrar used in unit tests."""
|
||||
|
||||
def get(self, _path: str):
|
||||
return lambda fn: fn
|
||||
|
||||
def post(self, _path: str):
|
||||
return lambda fn: fn
|
||||
|
||||
|
||||
class ClientTimeoutStub:
|
||||
def __init__(self, total: float | None = None):
|
||||
self.total = total
|
||||
|
||||
|
||||
class WSMsgTypeStub:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
|
||||
class WebSocketResponseStub:
|
||||
def __init__(self, *args: Any, **kwargs: Any):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def prepare(self, _request: Any) -> None:
|
||||
return None
|
||||
|
||||
async def send_json(self, _payload: Any) -> None:
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
async def _empty():
|
||||
if False:
|
||||
yield None
|
||||
|
||||
return _empty()
|
||||
|
||||
|
||||
class FormDataStub:
|
||||
def add_field(self, *_args: Any, **_kwargs: Any) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def reset_package_namespace(package_name: str) -> None:
|
||||
"""Clear test package modules from sys.modules."""
|
||||
prefix = f"{package_name}."
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(prefix):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
|
||||
def ensure_namespace_package(module_name: str) -> types.ModuleType:
|
||||
"""Create a namespace-like package module in sys.modules."""
|
||||
module = types.ModuleType(module_name)
|
||||
module.__path__ = []
|
||||
sys.modules[module_name] = module
|
||||
return module
|
||||
|
||||
|
||||
def bootstrap_test_package(
|
||||
package_name: str,
|
||||
*,
|
||||
with_api: bool = True,
|
||||
with_utils: bool = True,
|
||||
with_workers: bool = False,
|
||||
with_upscale: bool = False,
|
||||
with_orchestration: bool = False,
|
||||
) -> None:
|
||||
"""Create a clean package scaffold for module-level import tests."""
|
||||
reset_package_namespace(package_name)
|
||||
ensure_namespace_package(package_name)
|
||||
|
||||
if with_api:
|
||||
ensure_namespace_package(f"{package_name}.api")
|
||||
if with_utils:
|
||||
ensure_namespace_package(f"{package_name}.utils")
|
||||
if with_workers:
|
||||
ensure_namespace_package(f"{package_name}.workers")
|
||||
if with_upscale:
|
||||
ensure_namespace_package(f"{package_name}.upscale")
|
||||
if with_orchestration:
|
||||
ensure_namespace_package(f"{package_name}.api.orchestration")
|
||||
|
||||
|
||||
def install_server_stub(prompt_server_instance: Any | None = None) -> Any:
|
||||
"""Install a minimal `server` module exposing PromptServer.instance.routes."""
|
||||
if prompt_server_instance is None:
|
||||
prompt_server_instance = types.SimpleNamespace()
|
||||
if not hasattr(prompt_server_instance, "routes"):
|
||||
prompt_server_instance.routes = RoutesStub()
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
||||
sys.modules["server"] = server_module
|
||||
return prompt_server_instance
|
||||
|
||||
|
||||
def install_aiohttp_stub(
|
||||
json_response_factory: Callable[..., Any],
|
||||
) -> bool:
|
||||
"""Install/augment `aiohttp` with common members used by API tests."""
|
||||
created = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
aiohttp_module = sys.modules["aiohttp"]
|
||||
if not hasattr(aiohttp_module, "ClientTimeout"):
|
||||
aiohttp_module.ClientTimeout = ClientTimeoutStub
|
||||
if not hasattr(aiohttp_module, "WSMsgType"):
|
||||
aiohttp_module.WSMsgType = WSMsgTypeStub
|
||||
if not hasattr(aiohttp_module, "FormData"):
|
||||
aiohttp_module.FormData = FormDataStub
|
||||
|
||||
web_obj = getattr(aiohttp_module, "web", None)
|
||||
if web_obj is None:
|
||||
web_obj = types.SimpleNamespace()
|
||||
aiohttp_module.web = web_obj
|
||||
web_obj.json_response = json_response_factory
|
||||
if not hasattr(web_obj, "WebSocketResponse"):
|
||||
web_obj.WebSocketResponse = WebSocketResponseStub
|
||||
|
||||
return created
|
||||
|
||||
|
||||
def cleanup_optional_module(module_name: str, created: bool) -> None:
|
||||
"""Remove stubbed module only when this test created it."""
|
||||
if created:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
|
||||
def install_request_guards_stub(package_name: str) -> None:
|
||||
"""Install permissive request guard for isolated route tests."""
|
||||
request_guards_module = types.ModuleType(f"{package_name}.api.request_guards")
|
||||
|
||||
async def _authorization_error_or_none(_request):
|
||||
return None
|
||||
|
||||
request_guards_module.authorization_error_or_none = _authorization_error_or_none
|
||||
sys.modules[f"{package_name}.api.request_guards"] = request_guards_module
|
||||
|
||||
|
||||
def install_endpoint_policy_passthrough_stub(package_name: str) -> None:
|
||||
"""Install endpoint policy stub that executes operation directly."""
|
||||
endpoint_policy_module = types.ModuleType(f"{package_name}.api.endpoint_policy")
|
||||
|
||||
async def _run_authorized_endpoint(_request, operation, unexpected_status=500):
|
||||
_ = unexpected_status
|
||||
return await operation()
|
||||
|
||||
endpoint_policy_module.run_authorized_endpoint = _run_authorized_endpoint
|
||||
sys.modules[f"{package_name}.api.endpoint_policy"] = endpoint_policy_module
|
||||
@@ -1,423 +0,0 @@
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
AUTH_TOKEN = "integration-secret"
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, *, headers=None, json_payload=None, post_payload=None, query=None):
|
||||
self.headers = headers or {}
|
||||
self._json_payload = json_payload or {}
|
||||
self._post_payload = post_payload or {}
|
||||
self.query = query or {}
|
||||
|
||||
async def json(self):
|
||||
return self._json_payload
|
||||
|
||||
async def post(self):
|
||||
return self._post_payload
|
||||
|
||||
|
||||
def _load_module(package_name: str, rel_path: str, module_name: str):
|
||||
module_path = Path(__file__).resolve().parents[2] / rel_path
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.{module_name}", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _install_common_utils(package_name: str, config_data: dict):
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: copy.deepcopy(config_data)
|
||||
config_module.save_config = lambda _cfg: True
|
||||
|
||||
@asynccontextmanager
|
||||
async def _config_transaction():
|
||||
yield config_module.load_config()
|
||||
|
||||
config_module.config_transaction = _config_transaction
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
network_module.build_worker_url = lambda worker, path="": f"http://{worker.get('host', '127.0.0.1')}{path}"
|
||||
network_module.get_client_session = lambda: None
|
||||
network_module.probe_worker = lambda *_args, **_kwargs: {"ok": True}
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
|
||||
def _load_real_guard_stack(package_name: str):
|
||||
_load_module(package_name, "utils/auth.py", "utils.auth")
|
||||
_load_module(package_name, "utils/parsing.py", "utils.parsing")
|
||||
_load_module(package_name, "api/schemas.py", "api.schemas")
|
||||
_load_module(package_name, "api/request_guards.py", "api.request_guards")
|
||||
_load_module(package_name, "api/endpoint_policy.py", "api.endpoint_policy")
|
||||
|
||||
|
||||
def _load_request_guards_only_module():
|
||||
package_name = "dist_auth_guard_testpkg"
|
||||
config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
request_guards = sys.modules[f"{package_name}.api.request_guards"]
|
||||
schemas = sys.modules[f"{package_name}.api.schemas"]
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
return request_guards, schemas
|
||||
|
||||
|
||||
def _load_config_routes_real_guard_module():
|
||||
package_name = "dist_config_auth_integration_testpkg"
|
||||
config_data = {
|
||||
"workers": [],
|
||||
"master": {"host": "127.0.0.1"},
|
||||
"settings": {"distributed_api_token": AUTH_TOKEN},
|
||||
"tunnel": {},
|
||||
}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
module = _load_module(package_name, "api/config_routes.py", "api.config_routes")
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
return module
|
||||
|
||||
|
||||
def _load_tunnel_routes_real_guard_module():
|
||||
package_name = "dist_tunnel_auth_integration_testpkg"
|
||||
config_data = {
|
||||
"workers": [],
|
||||
"master": {"host": "127.0.0.1"},
|
||||
"settings": {"distributed_api_token": AUTH_TOKEN},
|
||||
"tunnel": {},
|
||||
}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
|
||||
cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
|
||||
cloudflare_module.cloudflare_tunnel_manager = types.SimpleNamespace(
|
||||
get_status=lambda: {"active": False},
|
||||
start_tunnel=lambda: {"active": True},
|
||||
stop_tunnel=lambda: {"active": False},
|
||||
)
|
||||
sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
|
||||
|
||||
module = _load_module(package_name, "api/tunnel_routes.py", "api.tunnel_routes")
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
return module
|
||||
|
||||
|
||||
def _load_usdu_routes_real_guard_module():
|
||||
package_name = "dist_usdu_auth_integration_testpkg"
|
||||
config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
|
||||
usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
|
||||
usdu_management_module.MAX_PAYLOAD_SIZE = 1024 * 1024
|
||||
sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
|
||||
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
job_store_module.init_dynamic_job = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class _BaseJobState:
|
||||
pass
|
||||
|
||||
class _ImageJobState(_BaseJobState):
|
||||
def __init__(self):
|
||||
self.mode = "dynamic"
|
||||
self.queue = asyncio.Queue()
|
||||
self.pending_images = asyncio.Queue()
|
||||
self.completed_images = {}
|
||||
self.worker_status = {}
|
||||
self.assigned_to_workers = {}
|
||||
|
||||
class _TileJobState(_BaseJobState):
|
||||
def __init__(self):
|
||||
self.mode = "static"
|
||||
self.queue = asyncio.Queue()
|
||||
self.pending_tasks = asyncio.Queue()
|
||||
self.completed_tasks = {}
|
||||
self.worker_status = {}
|
||||
self.assigned_to_workers = {}
|
||||
self.batched_static = False
|
||||
|
||||
job_models_module.BaseJobState = _BaseJobState
|
||||
job_models_module.ImageJobState = _ImageJobState
|
||||
job_models_module.TileJobState = _TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
|
||||
parsers_module.parse_tiles_from_form = lambda _data: []
|
||||
sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
|
||||
|
||||
module = _load_module(package_name, "api/usdu_routes.py", "api.usdu_routes")
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._ImageJobState = _ImageJobState
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
return module
|
||||
|
||||
|
||||
def _install_torch_stub_if_missing():
|
||||
if "torch" in sys.modules:
|
||||
return False
|
||||
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
return True
|
||||
|
||||
|
||||
def _load_worker_routes_real_guard_module():
|
||||
package_name = "dist_worker_auth_integration_testpkg"
|
||||
config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}, "workers": []}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
created_torch_stub = _install_torch_stub_if_missing()
|
||||
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.CHUNK_SIZE = 65536
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
workers_module = types.ModuleType(f"{package_name}.workers")
|
||||
workers_module.get_worker_manager = lambda: types.SimpleNamespace(
|
||||
processes={},
|
||||
save_processes=lambda: None,
|
||||
is_process_running=lambda _pid: False,
|
||||
)
|
||||
sys.modules[f"{package_name}.workers"] = workers_module
|
||||
|
||||
detection_module = types.ModuleType(f"{package_name}.workers.detection")
|
||||
detection_module.get_machine_id = lambda: "machine-id"
|
||||
detection_module.is_docker_environment = lambda: False
|
||||
detection_module.is_runpod_environment = lambda: False
|
||||
sys.modules[f"{package_name}.workers.detection"] = detection_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
class _PromptValidationError(Exception):
|
||||
def __init__(self, message="validation error"):
|
||||
super().__init__(message)
|
||||
self.validation_error = {}
|
||||
self.node_errors = {}
|
||||
|
||||
async def _queue_prompt_payload(*_args, **_kwargs):
|
||||
return "prompt-id"
|
||||
|
||||
async_helpers_module.PromptValidationError = _PromptValidationError
|
||||
async_helpers_module.queue_prompt_payload = _queue_prompt_payload
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
module = _load_module(package_name, "api/worker_routes.py", "api.worker_routes")
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
return module
|
||||
|
||||
|
||||
def _load_job_routes_real_guard_module():
|
||||
package_name = "dist_job_auth_integration_testpkg"
|
||||
config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
created_torch_stub = _install_torch_stub_if_missing()
|
||||
|
||||
_install_common_utils(package_name, config_data)
|
||||
_load_real_guard_stack(package_name)
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
image_module.pil_to_tensor = lambda _img: None
|
||||
image_module.ensure_contiguous = lambda tensor: tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.JOB_INIT_GRACE_PERIOD = 0.1
|
||||
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
|
||||
runtime_state = types.SimpleNamespace(
|
||||
distributed_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
runtime_state_module.ensure_distributed_runtime_state = lambda: runtime_state
|
||||
sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
|
||||
|
||||
orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
orchestration_module.orchestrate_distributed_execution = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = orchestration_module
|
||||
|
||||
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
||||
queue_request_module.parse_queue_request_payload = lambda data: data
|
||||
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
||||
|
||||
module = _load_module(package_name, "api/job_routes.py", "api.job_routes")
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
return module
|
||||
|
||||
|
||||
request_guards_module, schemas_module = _load_request_guards_only_module()
|
||||
config_routes_module = _load_config_routes_real_guard_module()
|
||||
tunnel_routes_module = _load_tunnel_routes_real_guard_module()
|
||||
usdu_routes_module = _load_usdu_routes_real_guard_module()
|
||||
worker_routes_module = _load_worker_routes_real_guard_module()
|
||||
job_routes_module = _load_job_routes_real_guard_module()
|
||||
|
||||
|
||||
class AuthGuardContractTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_is_authorized_request_missing_token_fails(self):
|
||||
request = _FakeRequest(headers={})
|
||||
config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
|
||||
self.assertFalse(schemas_module.is_authorized_request(request, config))
|
||||
|
||||
async def test_is_authorized_request_matching_token_succeeds(self):
|
||||
request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
|
||||
config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
|
||||
self.assertTrue(schemas_module.is_authorized_request(request, config))
|
||||
|
||||
async def test_authorization_error_or_none_returns_403_for_invalid_token(self):
|
||||
request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: "wrong-token"})
|
||||
response = await request_guards_module.authorization_error_or_none(request)
|
||||
self.assertIsNotNone(response)
|
||||
self.assertEqual(response.status, 403)
|
||||
|
||||
async def test_authorization_error_or_none_returns_none_for_valid_token(self):
|
||||
request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
|
||||
response = await request_guards_module.authorization_error_or_none(request)
|
||||
self.assertIsNone(response)
|
||||
|
||||
|
||||
class RouteAuthIntegrationTests(unittest.IsolatedAsyncioTestCase):
|
||||
def _auth_headers(self):
|
||||
return {schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN}
|
||||
|
||||
async def test_config_endpoint_enforces_real_guard(self):
|
||||
unauthorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers={}))
|
||||
self.assertEqual(unauthorized.status, 403)
|
||||
|
||||
authorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers=self._auth_headers()))
|
||||
self.assertEqual(authorized.status, 200)
|
||||
self.assertIn("settings", authorized.payload)
|
||||
|
||||
async def test_tunnel_endpoint_enforces_real_guard(self):
|
||||
unauthorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers={}))
|
||||
self.assertEqual(unauthorized.status, 403)
|
||||
|
||||
authorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers=self._auth_headers()))
|
||||
self.assertEqual(authorized.status, 200)
|
||||
self.assertEqual(authorized.payload.get("status"), "success")
|
||||
|
||||
async def test_usdu_endpoint_enforces_real_guard(self):
|
||||
unauthorized = await usdu_routes_module.job_status_endpoint(_FakeRequest(headers={}))
|
||||
self.assertEqual(unauthorized.status, 403)
|
||||
|
||||
authorized = await usdu_routes_module.job_status_endpoint(
|
||||
_FakeRequest(headers=self._auth_headers(), query={"multi_job_id": "job-1"})
|
||||
)
|
||||
self.assertEqual(authorized.status, 200)
|
||||
self.assertIn("ready", authorized.payload)
|
||||
|
||||
async def test_worker_endpoint_enforces_real_guard(self):
|
||||
worker_routes_module._collect_network_info_sync = lambda: {
|
||||
"interfaces": [],
|
||||
"recommended_host": "127.0.0.1",
|
||||
}
|
||||
unauthorized = await worker_routes_module.get_network_info_endpoint(_FakeRequest(headers={}))
|
||||
self.assertEqual(unauthorized.status, 403)
|
||||
|
||||
authorized = await worker_routes_module.get_network_info_endpoint(
|
||||
_FakeRequest(headers=self._auth_headers())
|
||||
)
|
||||
self.assertEqual(authorized.status, 200)
|
||||
self.assertEqual(authorized.payload.get("status"), "success")
|
||||
|
||||
async def test_job_endpoint_enforces_real_guard(self):
|
||||
unauthorized = await job_routes_module.prepare_job_endpoint(
|
||||
_FakeRequest(headers={}, json_payload={"multi_job_id": "job-1"})
|
||||
)
|
||||
self.assertEqual(unauthorized.status, 403)
|
||||
|
||||
authorized = await job_routes_module.prepare_job_endpoint(
|
||||
_FakeRequest(headers=self._auth_headers(), json_payload={"multi_job_id": "job-1"})
|
||||
)
|
||||
self.assertEqual(authorized.status, 200)
|
||||
self.assertEqual(authorized.payload.get("status"), "success")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -3,19 +3,9 @@ import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
install_endpoint_policy_passthrough_stub,
|
||||
install_request_guards_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
@@ -35,18 +25,38 @@ def _load_config_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
|
||||
package_name = "dist_api_config_testpkg"
|
||||
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
schemas_module.is_authorized_request = lambda _request, _config: True
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
install_request_guards_stub(package_name)
|
||||
install_endpoint_policy_passthrough_stub(package_name)
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
install_server_stub()
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = [str(module_path.parents[1] / "utils")]
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
import aiohttp # Load the real dependency; isolate the fake response at the tested module.
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
@@ -56,16 +66,7 @@ def _load_config_routes_module():
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
if isinstance(error, list):
|
||||
message = "; ".join(str(item) for item in error)
|
||||
error_payload = [str(item) for item in error]
|
||||
else:
|
||||
message = str(error)
|
||||
error_payload = str(error)
|
||||
return _FakeResponse(
|
||||
{"status": "error", "error": error_payload, "message": message},
|
||||
status=status,
|
||||
)
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
@@ -81,12 +82,6 @@ def _load_config_routes_module():
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: copy.deepcopy(default_config)
|
||||
config_module.save_config = lambda _cfg: True
|
||||
|
||||
@asynccontextmanager
|
||||
async def _config_transaction():
|
||||
yield config_module.load_config()
|
||||
|
||||
config_module.config_transaction = _config_transaction
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path)
|
||||
@@ -94,7 +89,14 @@ def _load_config_routes_module():
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
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
|
||||
|
||||
@@ -115,7 +117,9 @@ class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
async def test_update_config_valid_field_persists(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
|
||||
config_routes, "save_config", return_value=True
|
||||
):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
|
||||
@@ -0,0 +1,555 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload):
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
||||
codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
||||
codec_module = importlib.util.module_from_spec(codec_spec)
|
||||
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
||||
codec_spec.loader.exec_module(codec_module)
|
||||
utils_pkg.transfer_codec = codec_module
|
||||
return codec_module
|
||||
|
||||
|
||||
def _load_job_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
||||
package_name = "dist_api_queue_testpkg"
|
||||
|
||||
# Reset package namespace to avoid stale module state across test runs.
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
_register_real_transfer_codec(package_name, utils_pkg)
|
||||
|
||||
# aiohttp.web stub
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
# server module stub with route decorators
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
prompt_server_instance = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
# torch stub (only needed to satisfy import)
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
# PIL stub (only needed to satisfy import)
|
||||
created_pil_stub = False
|
||||
if "PIL" not in sys.modules:
|
||||
created_pil_stub = True
|
||||
pil_module = types.ModuleType("PIL")
|
||||
image_module = types.ModuleType("PIL.Image")
|
||||
pil_module.Image = image_module
|
||||
sys.modules["PIL"] = pil_module
|
||||
sys.modules["PIL.Image"] = image_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
||||
image_module.ensure_contiguous = lambda tensor: tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
|
||||
def _decode_audio_payload(payload):
|
||||
if payload is None:
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Field 'audio' must be an object when provided.")
|
||||
|
||||
encoded = payload.get("data")
|
||||
shape = payload.get("shape")
|
||||
dtype = payload.get("dtype", "float32")
|
||||
sample_rate = payload.get("sample_rate", 44100)
|
||||
if not isinstance(encoded, str) or not encoded.strip():
|
||||
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
||||
if not isinstance(shape, list) or len(shape) != 3:
|
||||
raise ValueError("Field 'audio.shape' must be a 3-item list.")
|
||||
if dtype != "float32":
|
||||
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
||||
try:
|
||||
shape_tuple = tuple(int(dim) for dim in shape)
|
||||
except Exception as exc:
|
||||
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
||||
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError("Field 'audio.data' byte size mismatch.")
|
||||
|
||||
waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
|
||||
return {"waveform": waveform, "sample_rate": int(sample_rate)}
|
||||
|
||||
audio_payload_module.decode_audio_payload = _decode_audio_payload
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
||||
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _QueuePayload:
|
||||
prompt: dict
|
||||
workflow_meta: object
|
||||
client_id: str
|
||||
delegate_master: object
|
||||
enabled_worker_ids: list
|
||||
auto_prepare: bool
|
||||
trace_execution_id: object
|
||||
|
||||
def _parse_queue_request_payload(data):
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
enabled = data.get("enabled_worker_ids")
|
||||
if not isinstance(enabled, list):
|
||||
raise ValueError("enabled_worker_ids required")
|
||||
client_id = data.get("client_id")
|
||||
if not isinstance(client_id, str) or not client_id.strip():
|
||||
raise ValueError("client_id required")
|
||||
return _QueuePayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=data.get("delegate_master"),
|
||||
enabled_worker_ids=enabled,
|
||||
auto_prepare=bool(data.get("auto_prepare", True)),
|
||||
trace_execution_id=data.get("trace_execution_id"),
|
||||
)
|
||||
|
||||
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
||||
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
||||
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
# 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)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
if created_pil_stub:
|
||||
sys.modules.pop("PIL.Image", None)
|
||||
sys.modules.pop("PIL", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
job_routes = _load_job_routes_module()
|
||||
|
||||
|
||||
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
"auto_prepare": True,
|
||||
}
|
||||
)
|
||||
with patch.object(
|
||||
job_routes,
|
||||
"orchestrate_distributed_execution",
|
||||
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
||||
):
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
||||
self.assertEqual(response.payload.get("number"), 42)
|
||||
self.assertEqual(response.payload.get("node_errors"), {})
|
||||
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
||||
|
||||
async def test_distributed_queue_missing_prompt_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("prompt", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
||||
|
||||
|
||||
class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
def _encoded_audio_payload(self):
|
||||
waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
|
||||
return {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(waveform.tobytes()).decode("ascii"),
|
||||
}
|
||||
|
||||
async def test_job_complete_accepts_audio_payload(self):
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-1")
|
||||
self.assertTrue(queued["is_last"])
|
||||
self.assertIsNotNone(queued["audio"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_job_complete_accepts_audio_without_image(self):
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "audio-only-job",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
decode_image.assert_not_called()
|
||||
queued = await queue.get()
|
||||
self.assertIsNone(queued["tensor"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
|
||||
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": 123,
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("image", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_job_complete_returns_400_for_malformed_audio(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_shape(self):
|
||||
bad = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(bad)
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_dtype(self):
|
||||
payload = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float16",
|
||||
"data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(payload)
|
||||
|
||||
|
||||
class _FakeFileField:
|
||||
def __init__(self, data):
|
||||
self.file = io.BytesIO(data)
|
||||
|
||||
|
||||
class _FakeMultipartRequest:
|
||||
content_type = "multipart/form-data"
|
||||
|
||||
def __init__(self, envelope, parts):
|
||||
from multidict import MultiDict
|
||||
|
||||
fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
|
||||
fields += [("media", _FakeFileField(part)) for part in parts]
|
||||
self._form = MultiDict(fields)
|
||||
|
||||
async def post(self):
|
||||
return self._form
|
||||
|
||||
async def json(self):
|
||||
raise ValueError("not JSON")
|
||||
|
||||
|
||||
class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
codec = job_routes.transfer_codec
|
||||
|
||||
def setUp(self):
|
||||
self.queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
|
||||
|
||||
def _frames(self, count=4, height=20, width=30):
|
||||
rng = np.random.default_rng(1)
|
||||
return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
|
||||
|
||||
def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
|
||||
chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
|
||||
parts = chunk.pop("parts")
|
||||
envelope = {
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"is_last": is_last,
|
||||
**chunk,
|
||||
"start_index": chunk["start_index"] + start_offset,
|
||||
}
|
||||
if audio is not None:
|
||||
envelope["audio"] = audio
|
||||
envelope.update(overrides)
|
||||
return _FakeMultipartRequest(envelope, parts)
|
||||
|
||||
async def _drain(self):
|
||||
items = []
|
||||
while not self.queue.empty():
|
||||
items.append(await self.queue.get())
|
||||
return items
|
||||
|
||||
async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
|
||||
frames = self._frames()
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
|
||||
self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
|
||||
self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
|
||||
for item, frame in zip(items, frames):
|
||||
self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
|
||||
self.assertEqual(item["tensor"].dtype, torch.float32)
|
||||
expected = frame.astype(np.float32) / 255.0
|
||||
self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
|
||||
|
||||
async def test_non_final_chunk_never_marks_worker_done(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "png", is_last=False)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertFalse(any(item["is_last"] for item in await self._drain()))
|
||||
|
||||
async def test_audio_is_attached_only_to_final_frame(self):
|
||||
audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=3), "jpeg", audio=audio)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
|
||||
self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_rejects_frame_count_mismatch(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=2), "ffv1", count=5)
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("frames", response.payload["message"])
|
||||
self.assertTrue(self.queue.empty())
|
||||
|
||||
async def test_rejects_bad_envelope(self):
|
||||
response = await job_routes.job_complete_endpoint(
|
||||
self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
|
||||
)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("format", response.payload["message"])
|
||||
self.assertIn("is_last", response.payload["message"])
|
||||
|
||||
async def test_rejects_invalid_envelope_json(self):
|
||||
response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("envelope", response.payload["message"])
|
||||
|
||||
async def test_unknown_job_returns_404(self):
|
||||
job_routes.prompt_server.distributed_pending_jobs = {}
|
||||
with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
|
||||
response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
|
||||
|
||||
self.assertEqual(response.status, 404)
|
||||
|
||||
async def test_legacy_json_payload_still_accepted(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-z",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 3,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
items = await self._drain()
|
||||
self.assertEqual(len(items), 1)
|
||||
self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,360 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
import asyncio
|
||||
import base64
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
install_request_guards_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload, headers=None):
|
||||
self._payload = payload
|
||||
self.headers = headers or {}
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def _load_job_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
||||
package_name = "dist_api_queue_testpkg"
|
||||
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
|
||||
# aiohttp.web stub
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
|
||||
# server module stub with route decorators
|
||||
prompt_server_instance = types.SimpleNamespace(
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
install_server_stub(prompt_server_instance)
|
||||
|
||||
# torch stub (only needed to satisfy import)
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
# PIL stub (only needed to satisfy import)
|
||||
created_pil_stub = False
|
||||
if "PIL" not in sys.modules:
|
||||
created_pil_stub = True
|
||||
pil_module = types.ModuleType("PIL")
|
||||
image_module = types.ModuleType("PIL.Image")
|
||||
pil_module.Image = image_module
|
||||
sys.modules["PIL"] = pil_module
|
||||
sys.modules["PIL.Image"] = image_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
||||
image_module.ensure_contiguous = lambda tensor: tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
|
||||
def _decode_audio_payload(payload):
|
||||
if payload is None:
|
||||
return None
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Field 'audio' must be an object when provided.")
|
||||
|
||||
encoded = payload.get("data")
|
||||
shape = payload.get("shape")
|
||||
dtype = payload.get("dtype", "float32")
|
||||
sample_rate = payload.get("sample_rate", 44100)
|
||||
if not isinstance(encoded, str) or not encoded.strip():
|
||||
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
||||
if not isinstance(shape, list) or len(shape) != 3:
|
||||
raise ValueError("Field 'audio.shape' must be a 3-item list.")
|
||||
if dtype != "float32":
|
||||
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
||||
try:
|
||||
shape_tuple = tuple(int(dim) for dim in shape)
|
||||
except Exception as exc:
|
||||
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
||||
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError("Field 'audio.data' byte size mismatch.")
|
||||
|
||||
waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
|
||||
return {"waveform": waveform, "sample_rate": int(sample_rate)}
|
||||
|
||||
audio_payload_module.decode_audio_payload = _decode_audio_payload
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
||||
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"settings": {}}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
|
||||
|
||||
def _ensure_distributed_runtime_state(server_instance=None):
|
||||
ps = server_instance or prompt_server_instance
|
||||
if not hasattr(ps, "distributed_pending_jobs"):
|
||||
ps.distributed_pending_jobs = {}
|
||||
if not hasattr(ps, "distributed_jobs_lock") or ps.distributed_jobs_lock is None:
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
if not hasattr(ps, "distributed_job_allowed_workers"):
|
||||
ps.distributed_job_allowed_workers = {}
|
||||
if not hasattr(ps, "distributed_pending_tile_jobs"):
|
||||
ps.distributed_pending_tile_jobs = {}
|
||||
if not hasattr(ps, "distributed_tile_jobs_lock"):
|
||||
ps.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
return types.SimpleNamespace(
|
||||
distributed_pending_jobs=ps.distributed_pending_jobs,
|
||||
distributed_jobs_lock=ps.distributed_jobs_lock,
|
||||
distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
|
||||
distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
|
||||
distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
|
||||
)
|
||||
|
||||
runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
|
||||
runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
|
||||
sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
schemas_module.is_authorized_request = lambda _request, _config: True
|
||||
schemas_module.require_bool_literal = (
|
||||
lambda value, field_name="value": value
|
||||
if isinstance(value, bool)
|
||||
else (_ for _ in ()).throw(ValueError(f"Field '{field_name}' must be a boolean literal."))
|
||||
)
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
|
||||
install_request_guards_stub(package_name)
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
||||
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
|
||||
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _QueuePayload:
|
||||
prompt: dict
|
||||
workflow_meta: object
|
||||
client_id: str
|
||||
delegate_master: object
|
||||
enabled_worker_ids: list
|
||||
trace_execution_id: object
|
||||
|
||||
def _parse_queue_request_payload(data):
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
enabled = data.get("enabled_worker_ids")
|
||||
if not isinstance(enabled, list):
|
||||
raise ValueError("enabled_worker_ids required")
|
||||
client_id = data.get("client_id")
|
||||
if not isinstance(client_id, str) or not client_id.strip():
|
||||
raise ValueError("client_id required")
|
||||
return _QueuePayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=data.get("delegate_master"),
|
||||
enabled_worker_ids=enabled,
|
||||
trace_execution_id=data.get("trace_execution_id"),
|
||||
)
|
||||
|
||||
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
||||
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
||||
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
cleanup_optional_module("torch", created_torch_stub)
|
||||
if created_pil_stub:
|
||||
cleanup_optional_module("PIL.Image", True)
|
||||
cleanup_optional_module("PIL", True)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
job_routes = _load_job_routes_module()
|
||||
|
||||
|
||||
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_distributed_queue_happy_path_returns_prompt_id(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
"auto_prepare": True,
|
||||
}
|
||||
)
|
||||
with patch.object(
|
||||
job_routes,
|
||||
"orchestrate_distributed_execution",
|
||||
new=AsyncMock(return_value=("prompt_123", 2)),
|
||||
):
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
||||
|
||||
async def test_distributed_queue_missing_prompt_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"enabled_worker_ids": ["w1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("prompt", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"prompt": {"1": {"class_type": "Node"}},
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
response = await job_routes.distributed_queue_endpoint(request)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
||||
|
||||
|
||||
class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
def _encoded_audio_payload(self):
|
||||
waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
|
||||
return {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(waveform.tobytes()).decode("ascii"),
|
||||
}
|
||||
|
||||
async def test_job_complete_accepts_audio_payload(self):
|
||||
queue = asyncio.Queue()
|
||||
prompt_server = job_routes.server.PromptServer.instance
|
||||
prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
prompt_server.distributed_pending_jobs = {"job-1": queue}
|
||||
prompt_server.distributed_job_allowed_workers = {"job-1": {"worker-1"}}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-1",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"audio": self._encoded_audio_payload(),
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-1")
|
||||
self.assertTrue(queued["is_last"])
|
||||
self.assertIsNotNone(queued["audio"])
|
||||
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
||||
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
||||
|
||||
async def test_job_complete_rejects_worker_not_in_job_allowlist(self):
|
||||
queue = asyncio.Queue()
|
||||
prompt_server = job_routes.server.PromptServer.instance
|
||||
prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
prompt_server.distributed_pending_jobs = {"job-allow": queue}
|
||||
prompt_server.distributed_job_allowed_workers = {"job-allow": {"worker-expected"}}
|
||||
request = _FakeRequest(
|
||||
{
|
||||
"job_id": "job-allow",
|
||||
"worker_id": "worker-unexpected",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,AAAA",
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
|
||||
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
||||
response = await job_routes.job_complete_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 403)
|
||||
self.assertIn("unauthorized", response.payload.get("message", "").lower())
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_shape(self):
|
||||
bad = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(bad)
|
||||
|
||||
def test_decode_audio_payload_rejects_bad_dtype(self):
|
||||
payload = {
|
||||
"sample_rate": 44100,
|
||||
"shape": [1, 2, 4],
|
||||
"dtype": "float16",
|
||||
"data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
job_routes._decode_audio_payload(payload)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+151
-47
@@ -4,29 +4,30 @@ import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
)
|
||||
|
||||
|
||||
class _AiohttpResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
def _load_media_sync_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
|
||||
package_name = "dist_ms_testpkg"
|
||||
|
||||
bootstrap_test_package(
|
||||
package_name,
|
||||
with_api=True,
|
||||
with_utils=True,
|
||||
with_orchestration=True,
|
||||
)
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
@@ -47,17 +48,22 @@ def _load_media_sync_module():
|
||||
trace_module.trace_info = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
|
||||
|
||||
auth_module = types.ModuleType(f"{package_name}.utils.auth")
|
||||
auth_module.distributed_auth_headers = lambda _config=None: {}
|
||||
sys.modules[f"{package_name}.utils.auth"] = auth_module
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"settings": {}}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _AiohttpResponse(payload, status=status)
|
||||
)
|
||||
class _FormData:
|
||||
def add_field(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.FormData = _FormData
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.media_sync",
|
||||
@@ -67,7 +73,8 @@ def _load_media_sync_module():
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
@@ -157,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)
|
||||
@@ -251,9 +243,121 @@ 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()
|
||||
class _AiohttpResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fetch_worker_has_node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status, payload):
|
||||
self.status = status
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
if isinstance(self._payload, Exception):
|
||||
raise self._payload
|
||||
return self._payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, response=None, error=None):
|
||||
self.response = response
|
||||
self.error = error
|
||||
self.urls = []
|
||||
|
||||
def get(self, url, timeout=None):
|
||||
self.urls.append(url)
|
||||
if self.error:
|
||||
raise self.error
|
||||
return self.response
|
||||
|
||||
|
||||
class FetchWorkerHasNodeTests(unittest.TestCase):
|
||||
def _probe(self, session):
|
||||
import asyncio
|
||||
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
original = ms.get_client_session
|
||||
ms.get_client_session = _get_session
|
||||
try:
|
||||
return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
|
||||
finally:
|
||||
ms.get_client_session = original
|
||||
|
||||
def test_registered_node_is_reported(self):
|
||||
session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
|
||||
self.assertTrue(self._probe(session))
|
||||
self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
|
||||
|
||||
def test_unknown_node_returns_empty_object(self):
|
||||
# ComfyUI answers 200 {} for classes that are not registered.
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
|
||||
|
||||
def test_http_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
|
||||
|
||||
def test_bad_json_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
|
||||
|
||||
def test_connection_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))
|
||||
|
||||
@@ -1,195 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
def _bootstrap_package(package_name):
|
||||
bootstrap_test_package(
|
||||
package_name,
|
||||
with_api=True,
|
||||
with_utils=True,
|
||||
with_orchestration=True,
|
||||
)
|
||||
|
||||
prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
install_server_stub(prompt_server_instance)
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
async def _queue_prompt_payload(prompt_obj, workflow_meta, client_id):
|
||||
_ = (prompt_obj, workflow_meta, client_id)
|
||||
return "prompt-id"
|
||||
|
||||
async_helpers_module.queue_prompt_payload = _queue_prompt_payload
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"settings": {}, "workers": []}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
|
||||
constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120.0
|
||||
constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
|
||||
constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_master_url = lambda **_kwargs: "http://127.0.0.1:8188"
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
trace_logger_module = types.ModuleType(f"{package_name}.utils.trace_logger")
|
||||
trace_logger_module.trace_debug = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module
|
||||
|
||||
runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
|
||||
|
||||
def _ensure_distributed_runtime_state(server_instance=None):
|
||||
ps = server_instance or prompt_server_instance
|
||||
if not hasattr(ps, "distributed_pending_jobs"):
|
||||
ps.distributed_pending_jobs = {}
|
||||
if not hasattr(ps, "distributed_jobs_lock"):
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
if not hasattr(ps, "distributed_job_allowed_workers"):
|
||||
ps.distributed_job_allowed_workers = {}
|
||||
if not hasattr(ps, "distributed_pending_tile_jobs"):
|
||||
ps.distributed_pending_tile_jobs = {}
|
||||
if not hasattr(ps, "distributed_tile_jobs_lock"):
|
||||
ps.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
return types.SimpleNamespace(
|
||||
distributed_pending_jobs=ps.distributed_pending_jobs,
|
||||
distributed_jobs_lock=ps.distributed_jobs_lock,
|
||||
distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
|
||||
distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
|
||||
distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
|
||||
)
|
||||
|
||||
runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
|
||||
runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
|
||||
sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
schemas_module.coerce_positive_int = (
|
||||
lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
|
||||
)
|
||||
schemas_module.coerce_positive_float = (
|
||||
lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 else default
|
||||
)
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
|
||||
dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch")
|
||||
dispatch_module.dispatch_worker_prompt = lambda *args, **kwargs: None
|
||||
dispatch_module.rank_workers_by_load = lambda workers, **kwargs: workers
|
||||
dispatch_module.select_active_workers = lambda workers, *_args, **_kwargs: (workers, False)
|
||||
dispatch_module.select_least_busy_worker = lambda workers, **_kwargs: workers[0] if workers else None
|
||||
sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module
|
||||
|
||||
media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync")
|
||||
media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt
|
||||
media_sync_module.fetch_worker_path_separator = lambda *_args, **_kwargs: "/"
|
||||
media_sync_module.sync_worker_media = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
|
||||
|
||||
prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
|
||||
prompt_transform_module.PromptIndex = object
|
||||
prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
|
||||
prompt_transform_module.find_nodes_by_class = lambda _prompt, _class_name: []
|
||||
prompt_transform_module.generate_job_id_map = lambda _index, _prefix: {}
|
||||
prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _ids: prompt
|
||||
prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
|
||||
sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
|
||||
|
||||
|
||||
def _load_module(package_name, module_rel_path, module_name):
|
||||
module_path = Path(__file__).resolve().parents[2] / module_rel_path
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.{module_name}",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _load_queue_orchestration_module():
|
||||
package_name = "dist_queue_orch_testpkg"
|
||||
_bootstrap_package(package_name)
|
||||
return _load_module(package_name, "api/queue_orchestration.py", "api.queue_orchestration")
|
||||
|
||||
|
||||
queue_orchestration = _load_queue_orchestration_module()
|
||||
|
||||
|
||||
class QueueOrchestrationHelpersTests(unittest.TestCase):
|
||||
def test_resolve_enabled_workers_filters_disabled_and_requested(self):
|
||||
config = {
|
||||
"workers": [
|
||||
{"id": "1", "enabled": True, "host": "a", "port": 8188},
|
||||
{"id": "2", "enabled": False, "host": "b", "port": 8189},
|
||||
{"id": "3", "enabled": True, "host": "c", "port": "bad-port"},
|
||||
]
|
||||
}
|
||||
|
||||
enabled_only = queue_orchestration._resolve_enabled_workers(config)
|
||||
requested = queue_orchestration._resolve_enabled_workers(config, requested_ids={"2", "3"})
|
||||
|
||||
self.assertEqual([worker["id"] for worker in enabled_only], ["1", "3"])
|
||||
self.assertEqual([worker["id"] for worker in requested], ["2", "3"])
|
||||
self.assertEqual(requested[1]["port"], 8188)
|
||||
|
||||
def test_resolve_orchestration_limits_uses_defaults_for_invalid_values(self):
|
||||
config = {
|
||||
"settings": {
|
||||
"worker_probe_concurrency": "x",
|
||||
"worker_prep_concurrency": "3",
|
||||
"media_sync_concurrency": "-1",
|
||||
"media_sync_timeout_seconds": "42.5",
|
||||
}
|
||||
}
|
||||
|
||||
limits = queue_orchestration._resolve_orchestration_limits(config)
|
||||
|
||||
self.assertEqual(limits[0], 8) # default
|
||||
self.assertEqual(limits[1], 3) # parsed
|
||||
self.assertEqual(limits[2], 2) # default
|
||||
self.assertEqual(limits[3], 42.5)
|
||||
|
||||
def test_is_load_balance_enabled_accepts_common_forms(self):
|
||||
self.assertTrue(queue_orchestration._is_load_balance_enabled(True))
|
||||
self.assertTrue(queue_orchestration._is_load_balance_enabled(1))
|
||||
self.assertTrue(queue_orchestration._is_load_balance_enabled(" yes "))
|
||||
self.assertFalse(queue_orchestration._is_load_balance_enabled(0))
|
||||
self.assertFalse(queue_orchestration._is_load_balance_enabled("off"))
|
||||
|
||||
def test_prompt_requests_load_balance_detects_collector_flag(self):
|
||||
prompt_index = types.SimpleNamespace(
|
||||
nodes_for_class=lambda class_name: ["10"] if class_name == "DistributedCollector" else [],
|
||||
inputs_by_node={"10": {"load_balance": "true"}},
|
||||
)
|
||||
self.assertTrue(queue_orchestration._prompt_requests_load_balance(prompt_index))
|
||||
|
||||
def test_ensure_distributed_state_initializes_attrs(self):
|
||||
prompt_server = types.SimpleNamespace()
|
||||
queue_orchestration.ensure_distributed_state(prompt_server)
|
||||
self.assertTrue(hasattr(prompt_server, "distributed_pending_jobs"))
|
||||
self.assertTrue(hasattr(prompt_server, "distributed_jobs_lock"))
|
||||
self.assertIsInstance(prompt_server.distributed_jobs_lock, asyncio.Lock)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,219 @@
|
||||
"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
_TESTS_DIR = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _load_from(path, name):
|
||||
spec = importlib.util.spec_from_file_location(name, path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _load_modules():
|
||||
real_aiohttp = sys.modules["aiohttp"] # imported inside the isolated test
|
||||
queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
|
||||
collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
|
||||
job_routes = queue_tests._load_job_routes_module()
|
||||
collector = collector_tests._load_collector_module()
|
||||
|
||||
# Error responses must be real HTTP responses here (the unit-test harness
|
||||
# returns plain objects); mirror utils.network.handle_api_error.
|
||||
from aiohttp import web
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
if isinstance(error, list):
|
||||
return web.json_response({"errors": [str(item) for item in error]}, status=status)
|
||||
return web.json_response({"error": str(error)}, status=status)
|
||||
|
||||
job_routes.handle_api_error = handle_api_error
|
||||
# The queue-test harness also fakes aiohttp.web for route modules; the
|
||||
# endpoint must build real responses for the real server used here.
|
||||
job_routes.web = web
|
||||
# The collector harness stubs aiohttp; this test needs the real client.
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
collector.aiohttp = real_aiohttp
|
||||
return job_routes, collector
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_sys_modules():
|
||||
"""The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
|
||||
saved = dict(sys.modules)
|
||||
# Earlier test modules may have left a stub `aiohttp` behind; drop it so the
|
||||
# real package is imported for this test (restored afterwards).
|
||||
for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
|
||||
if getattr(sys.modules[name], "__file__", None) is None:
|
||||
del sys.modules[name]
|
||||
yield
|
||||
sys.modules.clear()
|
||||
sys.modules.update(saved)
|
||||
|
||||
|
||||
async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
|
||||
request_limit, chunk_target, audio=None):
|
||||
"""Send `images` from a real worker node to the real completion route.
|
||||
|
||||
Returns (queued items, sizes of every request body the server received,
|
||||
exception raised by the worker or None).
|
||||
"""
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
body_sizes = []
|
||||
|
||||
@web.middleware
|
||||
async def record_size(request, handler):
|
||||
body_sizes.append(request.content_length)
|
||||
return await handler(request)
|
||||
|
||||
app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
error = None
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
try:
|
||||
await node.send_batch_to_master(
|
||||
images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
except Exception as exc: # surfaced to the caller for assertions
|
||||
error = exc
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items, body_sizes, error
|
||||
|
||||
|
||||
def _noise_images(count, height=360, width=640, seed=97):
|
||||
return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
|
||||
def test_high_entropy_video_stays_under_server_upload_limit(fmt):
|
||||
"""Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(24)
|
||||
limit = 2 * 1024 * 1024
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, fmt,
|
||||
client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
|
||||
|
||||
assert error is None, error
|
||||
assert len(body_sizes) > 1 and max(body_sizes) <= limit
|
||||
assert [item["image_index"] for item in items] == list(range(24))
|
||||
assert [item["is_last"] for item in items] == [False] * 23 + [True]
|
||||
|
||||
|
||||
def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
|
||||
"""A worker configured above the master's limit fails loudly instead of
|
||||
silently resending the whole batch as uncompressed PNG."""
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(8)
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "ffv1",
|
||||
client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
|
||||
|
||||
assert isinstance(error, RuntimeError)
|
||||
assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
|
||||
assert len(body_sizes) == 1
|
||||
assert items == []
|
||||
|
||||
|
||||
def test_frame_that_cannot_fit_sends_nothing():
|
||||
pytest.importorskip("aiohttp")
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = _noise_images(3)
|
||||
reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
items, body_sizes, error = asyncio.run(_run_transfer(
|
||||
job_routes, collector_module, images, "png",
|
||||
client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
|
||||
|
||||
assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
|
||||
assert body_sizes == []
|
||||
assert items == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
|
||||
def test_worker_to_master_over_real_http(fmt):
|
||||
# Import here, not at module level: other test harnesses only install their
|
||||
# aiohttp stubs when the real package has not been imported yet.
|
||||
aiohttp = pytest.importorskip("aiohttp")
|
||||
from aiohttp import web
|
||||
|
||||
job_routes, collector_module = _load_modules()
|
||||
images = torch.rand(5, 24, 36, 3)
|
||||
|
||||
async def scenario():
|
||||
queue = asyncio.Queue()
|
||||
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
collector_module.get_client_session = _get_session
|
||||
collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000 # several requests for still formats
|
||||
node = collector_module.DistributedCollectorNode()
|
||||
await node.send_batch_to_master(
|
||||
images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
|
||||
)
|
||||
await runner.cleanup()
|
||||
|
||||
items = []
|
||||
while not queue.empty():
|
||||
items.append(queue.get_nowait())
|
||||
return items
|
||||
|
||||
items = asyncio.run(scenario())
|
||||
|
||||
assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
|
||||
assert [item["is_last"] for item in items] == [False, False, False, False, True]
|
||||
expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
|
||||
got = torch.cat([item["tensor"] for item in items]).numpy()
|
||||
assert got.shape == expected.shape
|
||||
if fmt in ("png", "ffv1"):
|
||||
assert np.array_equal(got, expected)
|
||||
else:
|
||||
assert np.abs(got - expected).mean() < 0.25 # random noise: only sanity-check lossy codecs
|
||||
@@ -1,95 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
install_aiohttp_stub,
|
||||
install_endpoint_policy_passthrough_stub,
|
||||
install_request_guards_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
def _load_tunnel_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py"
|
||||
package_name = "dist_tunnel_routes_testpkg"
|
||||
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True)
|
||||
install_server_stub()
|
||||
install_aiohttp_stub(lambda payload, status=200: _FakeResponse(payload, status=status))
|
||||
|
||||
cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
|
||||
|
||||
class _TunnelManager:
|
||||
def get_status(self):
|
||||
return {"status": "running"}
|
||||
|
||||
async def start_tunnel(self):
|
||||
return {"status": "running", "public_url": "https://example.trycloudflare.com"}
|
||||
|
||||
async def stop_tunnel(self):
|
||||
return {"status": "stopped"}
|
||||
|
||||
cloudflare_module.cloudflare_tunnel_manager = _TunnelManager()
|
||||
sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"master": {"host": "https://master.example.com"}}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
install_request_guards_stub(package_name)
|
||||
install_endpoint_policy_passthrough_stub(package_name)
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.tunnel_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
tunnel_routes = _load_tunnel_routes_module()
|
||||
|
||||
|
||||
class TunnelRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_status_endpoint_returns_tunnel_and_master_host(self):
|
||||
response = await tunnel_routes.tunnel_status_endpoint(request=None)
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload["status"], "success")
|
||||
self.assertEqual(response.payload["tunnel"]["status"], "running")
|
||||
self.assertEqual(response.payload["master_host"], "https://master.example.com")
|
||||
|
||||
async def test_start_and_stop_endpoints_return_success_payload(self):
|
||||
start_response = await tunnel_routes.tunnel_start_endpoint(request=None)
|
||||
stop_response = await tunnel_routes.tunnel_stop_endpoint(request=None)
|
||||
|
||||
self.assertEqual(start_response.status, 200)
|
||||
self.assertEqual(start_response.payload["tunnel"]["status"], "running")
|
||||
self.assertEqual(stop_response.status, 200)
|
||||
self.assertEqual(stop_response.payload["tunnel"]["status"], "stopped")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -9,14 +9,6 @@ from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
install_request_guards_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
@@ -38,30 +30,43 @@ class _FakeRequest:
|
||||
return self._post_payload
|
||||
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
def _load_usdu_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
|
||||
package_name = "dist_api_usdu_testpkg"
|
||||
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
def _require_bool_literal(value, field_name: str = "value"):
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
if normalized == "true":
|
||||
return True
|
||||
if normalized == "false":
|
||||
return False
|
||||
raise ValueError(f"Field '{field_name}' must be a boolean literal.")
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
schemas_module.require_bool_literal = _require_bool_literal
|
||||
schemas_module.is_authorized_request = lambda _request, _config: True
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
install_request_guards_stub(package_name)
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
@@ -70,23 +75,23 @@ def _load_usdu_routes_module():
|
||||
)
|
||||
}
|
||||
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"settings": {}}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
|
||||
usdu_management_module.MAX_PAYLOAD_SIZE = 1024
|
||||
sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
@@ -98,21 +103,6 @@ def _load_usdu_routes_module():
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.MAX_PAYLOAD_SIZE = 1024
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
|
||||
async def _init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None):
|
||||
prompt_server = prompt_server_holder["value"]
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
return
|
||||
job_data = ImageJobState(multi_job_id=multi_job_id)
|
||||
job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])}
|
||||
job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])}
|
||||
indices = all_indices if all_indices is not None else list(range(int(batch_size or 0)))
|
||||
for idx in indices:
|
||||
await job_data.pending_images.put(int(idx))
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
job_store_module.init_dynamic_job = _init_dynamic_job
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
@@ -160,7 +150,6 @@ def _load_usdu_routes_module():
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
|
||||
parsers_module.parse_tiles_from_form = lambda _data: []
|
||||
parsers_module._parse_tiles_from_form = lambda _data: []
|
||||
sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
|
||||
|
||||
@@ -169,7 +158,8 @@ def _load_usdu_routes_module():
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
@@ -232,8 +222,7 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("kind"), "image")
|
||||
self.assertEqual(response.payload.get("task_idx"), 7)
|
||||
self.assertEqual(response.payload.get("image_idx"), 7)
|
||||
self.assertEqual(response.payload.get("estimated_remaining"), 0)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
|
||||
self.assertIn("worker-a", job_data.worker_status)
|
||||
@@ -249,8 +238,7 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("kind"), "tile")
|
||||
self.assertEqual(response.payload.get("task_idx"), 4)
|
||||
self.assertEqual(response.payload.get("tile_idx"), 4)
|
||||
self.assertTrue(response.payload.get("batched_static"))
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
|
||||
|
||||
|
||||
+181
-34
@@ -8,14 +8,6 @@ from collections import deque
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from tests.api.harness import (
|
||||
bootstrap_test_package,
|
||||
cleanup_optional_module,
|
||||
install_aiohttp_stub,
|
||||
install_request_guards_stub,
|
||||
install_server_stub,
|
||||
)
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
@@ -24,11 +16,10 @@ class _FakeResponse:
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None, match_info=None, query=None, headers=None):
|
||||
def __init__(self, payload=None, match_info=None, query=None):
|
||||
self._payload = payload
|
||||
self.match_info = match_info or {}
|
||||
self.query = query or {}
|
||||
self.headers = headers or {}
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
@@ -58,8 +49,8 @@ class _FakeHTTPClientSession:
|
||||
self._status = status
|
||||
self.calls = []
|
||||
|
||||
def get(self, url, params=None, headers=None, timeout=None):
|
||||
self.calls.append({"url": url, "params": params, "headers": headers, "timeout": timeout})
|
||||
def get(self, url, params=None, timeout=None):
|
||||
self.calls.append({"url": url, "params": params, "timeout": timeout})
|
||||
return _FakeHTTPClientResponse(self._payload, status=self._status)
|
||||
|
||||
|
||||
@@ -71,12 +62,12 @@ class _DummyWorkerManager:
|
||||
worker_id = str(worker["id"])
|
||||
self.processes[worker_id] = {
|
||||
"pid": 12345,
|
||||
"log_file": f"/tmp/distributed_worker_{worker_id}.log", # nosec B108 - deterministic fake test path
|
||||
"log_file": f"/tmp/distributed_worker_{worker_id}.log",
|
||||
"process": None,
|
||||
}
|
||||
return 12345
|
||||
|
||||
def is_process_running(self, _pid):
|
||||
def _is_process_running(self, _pid):
|
||||
return False
|
||||
|
||||
def save_processes(self):
|
||||
@@ -98,10 +89,24 @@ def _load_worker_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
|
||||
package_name = "dist_api_worker_testpkg"
|
||||
|
||||
bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
|
||||
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
@@ -114,10 +119,59 @@ def _load_worker_routes_module():
|
||||
detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
|
||||
sys.modules[f"{package_name}.workers.detection"] = detection_module
|
||||
|
||||
created_aiohttp_stub = install_aiohttp_stub(
|
||||
lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
install_server_stub()
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _WebSocketResponse:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def prepare(self, _request):
|
||||
return None
|
||||
|
||||
async def send_json(self, _payload):
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
async def _empty():
|
||||
if False:
|
||||
yield None
|
||||
return _empty()
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
|
||||
WebSocketResponse=_WebSocketResponse,
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
@@ -159,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")
|
||||
@@ -194,29 +249,21 @@ def _load_worker_routes_module():
|
||||
def _validate_worker_id(worker_id, config):
|
||||
return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
|
||||
|
||||
def _require_worker_id(worker_id, config, field_name="worker_id"):
|
||||
_ = field_name
|
||||
worker_id_str = str(worker_id).strip()
|
||||
if any(str(worker.get("id")) == worker_id_str for worker in config.get("workers", [])):
|
||||
return worker_id_str
|
||||
raise ValueError(f"Worker {worker_id_str} not found")
|
||||
|
||||
schemas_module.require_fields = _require_fields
|
||||
schemas_module.validate_worker_id = _validate_worker_id
|
||||
schemas_module.require_worker_id = _require_worker_id
|
||||
schemas_module.is_authorized_request = lambda _request, _config: True
|
||||
schemas_module.distributed_auth_headers = lambda _config: {}
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
|
||||
install_request_guards_stub(package_name)
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
# Keep this module's response double independent of aiohttp import order.
|
||||
module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
|
||||
|
||||
cleanup_optional_module("aiohttp", created_aiohttp_stub)
|
||||
cleanup_optional_module("torch", created_torch_stub)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
|
||||
return module
|
||||
|
||||
@@ -224,7 +271,107 @@ def _load_worker_routes_module():
|
||||
worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WebsocketExecutionTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
from test_async_helpers import _load_async_helpers_module
|
||||
self.modules = patch.dict(sys.modules)
|
||||
self.modules.start()
|
||||
self.addCleanup(self.modules.stop)
|
||||
self.helper, self.server = _load_async_helpers_module()
|
||||
sys.modules['dist_api_worker_testpkg.utils.async_helpers'] = self.helper
|
||||
self.prompt = {'9': {'class_type': 'Output'}, '10': {'class_type': 'Output'}}
|
||||
|
||||
async def receive(self, payload):
|
||||
import json
|
||||
sent = []
|
||||
class Socket:
|
||||
async def prepare(self, request):
|
||||
pass
|
||||
async def send_json(self, value):
|
||||
sent.append(value)
|
||||
def __aiter__(self):
|
||||
async def messages():
|
||||
yield types.SimpleNamespace(type=worker_routes.aiohttp.WSMsgType.TEXT,
|
||||
data=json.dumps(payload))
|
||||
return messages()
|
||||
with patch.object(worker_routes, 'web', types.SimpleNamespace(WebSocketResponse=lambda **kw: Socket())), \
|
||||
patch.object(worker_routes, 'queue_prompt_payload', self.helper.queue_prompt_payload), \
|
||||
patch.object(worker_routes, 'PromptValidationError', self.helper.PromptValidationError):
|
||||
await worker_routes.worker_ws_endpoint(_FakeRequest())
|
||||
return sent[0]
|
||||
|
||||
def payload(self, targets):
|
||||
return {'type': 'dispatch_prompt', 'request_id': 'req', 'prompt': self.prompt,
|
||||
'partial_execution_targets': targets}
|
||||
|
||||
async def test_explicit_targets_reach_queue_and_acknowledgement(self):
|
||||
ack = await self.receive(self.payload(['9']))
|
||||
self.assertTrue(ack['ok'], ack)
|
||||
self.assertEqual(ack['request_id'], 'req')
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ['9'])
|
||||
|
||||
async def test_malformed_targets_and_messages_are_rejected(self):
|
||||
for targets in (None, [], '9', [True], [9], ['absent'], ['9', '9']):
|
||||
with self.subTest(targets=targets):
|
||||
ack = await self.receive(self.payload(targets))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertIn('partial_execution_targets', ack['error'])
|
||||
self.assertFalse((await self.receive([]))['ok'])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_validation_cannot_silently_queue_a_partial_subset(self):
|
||||
ack = await self.receive(self.payload(['9', '10']))
|
||||
self.assertFalse(ack['ok'])
|
||||
self.assertEqual(ack['validation_error']['type'], 'partial_execution_targets_failed_validation')
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_missing_target_field_keeps_legacy_queue_call(self):
|
||||
from unittest.mock import AsyncMock
|
||||
queued = AsyncMock(return_value='legacy')
|
||||
original = self.helper.queue_prompt_payload
|
||||
self.helper.queue_prompt_payload = queued
|
||||
try:
|
||||
payload = self.payload(['9'])
|
||||
payload.pop('partial_execution_targets')
|
||||
self.assertTrue((await self.receive(payload))['ok'])
|
||||
self.assertNotIn('partial_execution_targets', queued.call_args.kwargs)
|
||||
finally:
|
||||
self.helper.queue_prompt_payload = original
|
||||
|
||||
|
||||
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}]}
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# This conftest.py marks the tests/ directory as the pytest collection root,
|
||||
# preventing pytest from traversing into the parent package's __init__.py.
|
||||
+691
@@ -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)"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Offline CPU proof against ComfyUI's real signatures and PromptExecutor.
|
||||
|
||||
Exercises one native worker executor and separate private master-slot lifetimes,
|
||||
including their shared stable Start / fresh End contract.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def main(core):
|
||||
sys.path.insert(0, str(Path(core).resolve()))
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True # Must precede all model-management imports.
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
import torch
|
||||
from comfy_execution.caching import CacheKeySetInputSignature
|
||||
|
||||
name = '_native_cache_identity_probe'
|
||||
for suffix in ('', '.image_list', '.nodes', '.utils'):
|
||||
package = types.ModuleType(name + suffix)
|
||||
package.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
|
||||
sys.modules[name + suffix] = package
|
||||
|
||||
def forbidden(*args, **kwargs):
|
||||
raise AssertionError('offline native proof attempted a server-loop bridge')
|
||||
|
||||
sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
|
||||
rt = importlib.import_module(name + '.image_list.runtime')
|
||||
boundary = importlib.import_module(name + '.nodes.image_list')
|
||||
calls = []
|
||||
|
||||
class Operation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'mask': ('MASK',),
|
||||
'text': ('STRING', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
|
||||
def run(self, image, mask, text):
|
||||
calls.append(text)
|
||||
return image + 1, mask + 1, text + ':done'
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
|
||||
DistributedImageListEnd=boundary.DistributedImageListEnd,
|
||||
Operation=Operation)
|
||||
prompt = {
|
||||
'start': {'class_type': rt.START_CLASS, 'inputs': {'image_list_execution_id': 'outer-A'}},
|
||||
'op': {'class_type': 'Operation', 'inputs': {
|
||||
'image': ['start', 0], 'mask': ['start', 1], 'text': ['start', 4]}},
|
||||
'end': {'class_type': rt.END_CLASS, 'inputs': {
|
||||
'images': ['op', 0], 'masks': ['op', 1], 'texts': ['op', 2],
|
||||
'context': ['start', 3], 'value_3': ['start', 2]}},
|
||||
}
|
||||
data = rt.normalize_image_list_inputs(torch.zeros(1, 2, 3, 3, dtype=torch.float64),
|
||||
torch.ones(1, 2, 3, dtype=torch.bfloat16),
|
||||
{'value_1': 'original', 'value_2': 3})[0]
|
||||
region = rt.validate_image_list_regions(prompt)[0]
|
||||
|
||||
def prepare(data):
|
||||
token = rt.create_worker_task_token()
|
||||
context = rt.create_master_job('start', 'end', prompt, [data], ['worker'], '')
|
||||
rt.register_assignment(context, data['index'], 'worker', token)
|
||||
identity = rt.store_worker_item(token, rt.encode_item(data))
|
||||
result = rt.build_worker_prompt(prompt, region, token, context, data['index'], '', item_identity=identity)
|
||||
result['start']['inputs']['image_list_local_execution'] = True
|
||||
result['end']['inputs']['image_list_local_execution'] = True
|
||||
return token, context, result
|
||||
|
||||
async def signature(graph, executor):
|
||||
dynamic = execution.DynamicPrompt(graph)
|
||||
changed = execution.IsChangedCache('signature-only', dynamic, executor.caches.outputs)
|
||||
keys = CacheKeySetInputSignature(dynamic, graph.keys(), changed)
|
||||
await keys.add_keys(graph.keys())
|
||||
return {node: keys.get_data_key(node) for node in graph}
|
||||
|
||||
async def probe():
|
||||
server = types.SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
|
||||
executor = execution.PromptExecutor(server, cache_type=execution.CacheType.LRU,
|
||||
cache_args={'lru': 64, 'ram': 0, 'ram_inactive': 0})
|
||||
ta, ca, a = prepare(data)
|
||||
prompt['start']['inputs']['image_list_execution_id'] = 'outer-B'
|
||||
tb, cb, b = prepare(data)
|
||||
assert ta != tb and ca != cb
|
||||
ka, kb = await signature(a, executor), await signature(b, executor)
|
||||
assert ka['start'] == kb['start'] and ka['op'] == kb['op']
|
||||
assert ka['end'] != kb['end']
|
||||
assert a['start']['inputs'] == b['start']['inputs']
|
||||
# Exercise actual V3 list mapping, hidden.prompt/unique_id, lazy inputs,
|
||||
# native node signatures and cache hits, not a reimplementation of them.
|
||||
for number, (token, context, graph) in enumerate(((ta, ca, a), (tb, cb, b))):
|
||||
valid = await execution.validate_prompt(f'native-{number}', graph, ['end'])
|
||||
assert valid[0], valid
|
||||
await executor.execute_async(graph, f'native-{number}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert torch.equal(result['mask'], data['mask'] + 1)
|
||||
assert result['text'] == 'original:done'
|
||||
assert result['values'] == {'value_3': 0}
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
if number:
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for failure in ('expired', 'mismatched'):
|
||||
token, context, graph = prepare(data)
|
||||
expires, staged = rt._WORKER_ITEMS[token]
|
||||
if failure == 'expired':
|
||||
rt._WORKER_ITEMS[token] = (0, staged)
|
||||
else:
|
||||
staged['image'] = staged['image'] + 9
|
||||
await executor.execute_async(graph, f'rejected-{failure}', {}, ['end'])
|
||||
assert not executor.success, failure
|
||||
assert not rt.get_master_job(context).received_indexes
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
cached = next(entry['nodes'] for event, entry in executor.status_messages if event == 'execution_cached')
|
||||
assert 'start' in cached and 'op' in cached and 'end' not in cached, cached
|
||||
assert calls == ['original'], calls
|
||||
for part in ('image', 'mask', 'value', 'index'):
|
||||
changed = copy.deepcopy(data)
|
||||
if part in ('image', 'mask'):
|
||||
changed[part].reshape(-1)[0] += 1
|
||||
elif part == 'value':
|
||||
changed['values']['value_2'] = 4 # Even an unconnected slot is content.
|
||||
else:
|
||||
changed['index'] = 1
|
||||
token, context, graph = prepare(changed)
|
||||
keys = await signature(graph, executor)
|
||||
assert keys['start'] != ka['start'] and keys['op'] != ka['op'], part
|
||||
await executor.execute_async(graph, f'changed-{part}', {}, ['end'])
|
||||
assert executor.success, executor.status_messages
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
assert rt.get_master_job(context).received_indexes == {changed['index']}
|
||||
assert len(calls) == 5, calls
|
||||
# Combine the real Start/End identity path with cache transfer between
|
||||
# private master slots, not just independent signature/retention probes.
|
||||
local = importlib.import_module(name + '.image_list.local_executor')
|
||||
before = len(calls)
|
||||
for number in range(2):
|
||||
token, context, graph = prepare(data)
|
||||
slot = local.LocalExecutionSlot()
|
||||
try:
|
||||
await asyncio.wrap_future(slot.submit(graph, f'local-{number}', 'end'))
|
||||
result = rt.get_master_job(context).results[0]
|
||||
assert torch.equal(result['image'], data['image'] + 1)
|
||||
assert result['values'] == {'value_3': 0}
|
||||
assert token not in rt._WORKER_ITEMS
|
||||
finally:
|
||||
await slot.close()
|
||||
rt.remove_master_job(context)
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
assert len(calls) == before + 1, calls
|
||||
assert not rt._WORKER_ITEMS
|
||||
assert not torch.cuda.is_initialized()
|
||||
print('NATIVE_CACHE_IDENTITY_OK stable Start/op keys; fresh End key; '
|
||||
'1 operation for 2 same-data jobs; 4 changed-content misses; '
|
||||
'expired/mismatched cached inputs rejected; current delivery and cleanup')
|
||||
|
||||
asyncio.run(probe())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(sys.argv[1])
|
||||
@@ -0,0 +1,688 @@
|
||||
"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
|
||||
|
||||
Isolated listeners use real core /prompt, /queue, /history and /interrupt handlers.
|
||||
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 copy
|
||||
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, MODEL_CALLS = [], [], [], []
|
||||
UTILITY_CALLS, ITEM_EXECUTIONS, VALUE_CAPTURES = [], [], []
|
||||
ALLOW_MASTER = False
|
||||
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((2, 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 TestSharedFlag:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('BOOLEAN',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
assert not WORKER, 'shared value source leaked into worker prompt'
|
||||
return (True,)
|
||||
|
||||
|
||||
class TestCachedModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {}}
|
||||
RETURN_TYPES = ('MODEL',)
|
||||
FUNCTION = 'run'
|
||||
def run(self):
|
||||
MODEL_CALLS.append(True)
|
||||
return (object(),)
|
||||
|
||||
|
||||
def assert_item_execution(prompt, index):
|
||||
"""Prove the operation belongs to an item, not merely a participating host."""
|
||||
from comfy_execution.utils import get_executing_context
|
||||
context = get_executing_context()
|
||||
assert context is not None
|
||||
assert prompt['4']['inputs'].get('worker_result_token'), 'per-item node ran in outer graph'
|
||||
assert (prompt['2']['inputs'].get('worker_item_identity')
|
||||
or prompt['2']['inputs'].get('worker_item_token')), 'per-item source binding missing'
|
||||
running, _ = server.PromptServer.instance.prompt_queue.get_current_queue_volatile()
|
||||
assert len(running) == 1, running
|
||||
entry = running[0]
|
||||
if WORKER:
|
||||
assert entry[1] == context.prompt_id and entry[2] == prompt
|
||||
assert entry[4] == ['4'], entry[4]
|
||||
else:
|
||||
assert ALLOW_MASTER, 'per-item node ran on delegate-only master'
|
||||
assert entry[1] != context.prompt_id, 'per-item node ran in outer executor'
|
||||
assert set(entry[4]) == {'4', '5'}, entry[4]
|
||||
return {'index': index, 'prompt_id': context.prompt_id,
|
||||
'role': 'worker' if WORKER else 'master-item', 'queue_prompt_id': entry[1]}
|
||||
|
||||
|
||||
class TestRequiredOutputUtility:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',), 'index': ('INT', {'forceInput': True})},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ('IMAGE',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, index, prompt):
|
||||
assert_item_execution(prompt, index)
|
||||
UTILITY_CALLS.append(index)
|
||||
return (image,)
|
||||
|
||||
|
||||
class TestDiscardedPreview:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'image': ('IMAGE',)}}
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, image):
|
||||
raise AssertionError('discarded per-item output preview executed')
|
||||
|
||||
|
||||
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}),
|
||||
'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)},
|
||||
'hidden': {'prompt': 'PROMPT'}}
|
||||
RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
|
||||
FUNCTION = 'run'
|
||||
def run(self, image, reference_image, mask, text, strength, seed, index, flag, model, prompt):
|
||||
assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
|
||||
ITEM_EXECUTIONS.append(assert_item_execution(prompt, index))
|
||||
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 flag is True
|
||||
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 TestValueOverride:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'text': ('STRING', {'forceInput': True}),
|
||||
'index': ('INT', {'forceInput': True})}}
|
||||
RETURN_TYPES = ('STRING', 'INT')
|
||||
FUNCTION = 'run'
|
||||
def run(self, text, index):
|
||||
assert WORKER or ALLOW_MASTER, 'override branch executed on delegate-only master'
|
||||
return (text + ':override', 200 + index)
|
||||
|
||||
|
||||
class TestCapture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})},
|
||||
'optional': {'values': ('DISTRIBUTED_IMAGE_LIST_VALUES',),
|
||||
'picked': ('STRING', {'forceInput': True})}}
|
||||
INPUT_IS_LIST = True
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = 'run'
|
||||
def run(self, images, masks, texts, values=None, picked=None):
|
||||
assert not WORKER, 'post-End capture leaked into a worker prompt'
|
||||
CAPTURED.append((images, masks, texts))
|
||||
VALUE_CAPTURES.append((values, picked))
|
||||
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_')]
|
||||
boundary = importlib.import_module(module.__name__ + '.nodes.image_list')
|
||||
encode_result = boundary.encode_result
|
||||
def checked_encode(result):
|
||||
encoded = encode_result(result)
|
||||
if 'values' in encoded:
|
||||
assert set(encoded['values']) == {f'value_{n}' for n in (2, 3, 4, 5, 7, 8, 9, 10)}
|
||||
return encoded
|
||||
boundary.encode_result = checked_encode
|
||||
nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
|
||||
TestImageOperation=TestImageOperation, TestCapture=TestCapture,
|
||||
TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag,
|
||||
TestRequiredOutputUtility=TestRequiredOutputUtility,
|
||||
TestDiscardedPreview=TestDiscardedPreview, TestValueOverride=TestValueOverride)
|
||||
# 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]
|
||||
|
||||
|
||||
def finish_queue_item(ps, item_id, executor):
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result,
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if executor.success else 'error',
|
||||
completed=executor.success, messages=executor.status_messages),
|
||||
process_item=lambda item: item[:5] + item[6:])
|
||||
|
||||
|
||||
async def worker_main():
|
||||
ps, assets, module, app = await setup()
|
||||
wake = asyncio.Event()
|
||||
interrupted_ids, submissions = [], []
|
||||
core_handlers = {(route.method, route.path): route.handler for route in ps.routes}
|
||||
async def prompt(request):
|
||||
data = await request.json()
|
||||
# Use core's actual validation/enqueue/response schema, including
|
||||
# partial_execution_targets; only observe it and wake our CPU consumer.
|
||||
response = await core_handlers['POST', '/prompt'](request)
|
||||
if response.status == 200:
|
||||
accepted = json.loads(response.body)
|
||||
_, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
entry = next(item for item in pending if item[1] == accepted['prompt_id'])
|
||||
submissions.append({'prompt_id': entry[1], 'requested': data.get('partial_execution_targets'),
|
||||
'queued': entry[4], 'node_ids': sorted(entry[2]),
|
||||
'end_ids': [key for key, node in entry[2].items()
|
||||
if node['class_type'] == 'DistributedImageListEnd']})
|
||||
wake.set()
|
||||
return response
|
||||
async def interrupt(request):
|
||||
interrupted_ids.append((await request.json())['prompt_id'])
|
||||
return await core_handlers['POST', '/interrupt'](request)
|
||||
async def stats(request):
|
||||
running, pending = ps.prompt_queue.get_current_queue_volatile()
|
||||
return web.json_response({'calls': BRANCH_CALLS, 'utility_calls': UTILITY_CALLS,
|
||||
'item_executions': ITEM_EXECUTIONS, 'submissions': submissions,
|
||||
'running': [e[1] for e in running], 'pending': [e[1] for e in pending],
|
||||
'interrupted': interrupted_ids, 'history': ps.prompt_queue.get_history()})
|
||||
async def consume():
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
while True:
|
||||
await wake.wait()
|
||||
while (queued := ps.prompt_queue.get(timeout=0)) is not None:
|
||||
entry, item_id = queued
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
wake.clear()
|
||||
app.add_routes([route for route in ps.routes if
|
||||
(route.path in ('/queue', '/history/{prompt_id}') or
|
||||
(route.path == '/prompt' and route.method == 'GET'))])
|
||||
app.router.add_post('/prompt', prompt)
|
||||
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': {}},
|
||||
# Generic Autogrow value slots carry per-item reference images, text,
|
||||
# strength and seed (one value per image), plus a shared boolean.
|
||||
'2': {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': ['1', 0], 'mask': ['1', 2],
|
||||
'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
|
||||
'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
|
||||
'values.value_6': ['1', 2],
|
||||
'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), ('mask', 1), ('index', 2),
|
||||
('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
|
||||
| {'image': ['8', 0], 'model': ['6', 0]}},
|
||||
'4': {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3],
|
||||
'value_2': ['10', 0], 'value_3': ['2', 6], 'value_4': ['2', 7], 'value_5': ['2', 8],
|
||||
'value_7': ['3', 1], 'value_8': ['2', 4], 'value_9': ['2', 9], 'value_10': ['10', 1]}},
|
||||
'5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2],
|
||||
'values': ['4', 3], 'picked': ['get', 0]}},
|
||||
'6': {'class_type': 'TestCachedModel', 'inputs': {}},
|
||||
'7': {'class_type': 'TestSharedFlag', 'inputs': {}},
|
||||
'8': {'class_type': 'TestRequiredOutputUtility', 'inputs': {'image': ['2', 0], 'index': ['2', 2]}},
|
||||
'9': {'class_type': 'TestDiscardedPreview', 'inputs': {'image': ['8', 0]}},
|
||||
'10': {'class_type': 'TestValueOverride', 'inputs': {'text': ['2', 5], 'index': ['2', 2]}},
|
||||
'get': {'class_type': 'DistributedImageListGetValue', 'inputs': {'values': ['4', 3], 'value_name': 'value_2'}},
|
||||
}
|
||||
|
||||
|
||||
def check_carried_values():
|
||||
values, picked = VALUE_CAPTURES[-1]
|
||||
assert len(values) == len(picked) == COUNT
|
||||
for index, (actual, source) in enumerate(zip(values, source_rows())):
|
||||
assert set(actual) == {f'value_{n}' for n in range(1, 11)}, actual.keys()
|
||||
assert torch.equal(actual['value_1'], source['reference_image'])
|
||||
assert actual['value_2'] == picked[index] == f'item-{index}:override'
|
||||
assert actual['value_3'] == source['strength']
|
||||
assert actual['value_4'] == source['seed']
|
||||
assert actual['value_5'] is True
|
||||
assert torch.equal(actual['value_6'], source['mask'])
|
||||
assert torch.equal(actual['value_7'], source['mask'] * 0.5)
|
||||
assert torch.equal(actual['value_8'], source['reference_image'])
|
||||
assert actual['value_8'].shape[0] == 2, 'value tensor batch was split'
|
||||
assert torch.equal(actual['value_9'], source['mask'])
|
||||
assert actual['value_10'] == 200 + index
|
||||
|
||||
|
||||
async def coordinator():
|
||||
global ALLOW_MASTER
|
||||
ps, assets, module, app = await setup()
|
||||
runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
|
||||
execution_plan = importlib.import_module(module.__name__ + '.image_list.execution_plan')
|
||||
lifecycle = importlib.import_module(module.__name__ + '.image_list.lifecycle')
|
||||
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_delegate_only': True},
|
||||
'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
|
||||
assert {'4', '5', '8', '9'} == set(valid[2]), 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')
|
||||
saved_prompt = copy.deepcopy(prompt)
|
||||
outer_executions = []
|
||||
async def enqueue(prepared, workflow, client_id, **kwargs):
|
||||
targets = kwargs['partial_execution_targets']
|
||||
planned, planned_targets = execution_plan.prepare_image_list_execution(prepared)
|
||||
assert planned == prepared and targets == planned_targets
|
||||
if '9' in saved_prompt and 'TestImageOperation' in {
|
||||
node['class_type'] for node in prepared.values()}:
|
||||
assert set(prepared) == set(saved_prompt) - {'9'}, prepared
|
||||
assert prepared['8'] == saved_prompt['8'], 'required utility changed'
|
||||
assert prepared['5'] == saved_prompt['5'], 'post-End output changed'
|
||||
assert set(targets) == {'4', '5'}, targets
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestRequiredOutputUtility'].OUTPUT_NODE is True
|
||||
assert nodes.NODE_CLASS_MAPPINGS['TestDiscardedPreview'].OUTPUT_NODE is True
|
||||
pid = 'network-acceptance' if not outer_executions else str(uuid.uuid4())
|
||||
checked = await execution.validate_prompt(pid, prepared, targets)
|
||||
assert checked[0], checked
|
||||
assert not checked[3] and Counter(checked[2]) == Counter(targets), checked
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
# Exactly the core queue tuple/put/get/task_done lifecycle. Start
|
||||
# must prove its roots from this real running record, not a mock.
|
||||
number = len(outer_executions)
|
||||
ps.prompt_queue.put((number, pid, prepared, {'client_id': client_id}, checked[2], {}))
|
||||
entry, item_id = ps.prompt_queue.get(timeout=0)
|
||||
assert lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
before_captures = len(CAPTURED)
|
||||
await asyncio.to_thread(executor.execute, entry[2], entry[1], entry[3], entry[4])
|
||||
finish_queue_item(ps, item_id, executor)
|
||||
assert ps.prompt_queue.get_current_queue_volatile() == ([], [])
|
||||
history = ps.prompt_queue.get_history(pid)[pid]
|
||||
assert Counter(history['prompt'][4]) == Counter(targets)
|
||||
assert not lifecycle.has_safe_execution_targets(ps, pid, prepared)
|
||||
outer_executions.append({'prompt_id': pid, 'targets': targets,
|
||||
'captures': len(CAPTURED) - before_captures})
|
||||
if executor.success:
|
||||
assert len(CAPTURED) == before_captures + 1, 'post-End Capture must run exactly once'
|
||||
return {'prompt_id': pid, 'number': number, 'node_errors': checked[3]}
|
||||
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'
|
||||
assert prompt == saved_prompt, 'queue preparation changed the saved/source graph'
|
||||
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 and not UTILITY_CALLS
|
||||
assert len(CAPTURED) == 1
|
||||
check_carried_values()
|
||||
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))
|
||||
for row in stats:
|
||||
assert Counter(row['utility_calls']) == Counter(row['calls']), row
|
||||
assert len(row['submissions']) == len(row['calls'])
|
||||
assert len({event['prompt_id'] for event in row['item_executions']}) == len(row['calls'])
|
||||
for submission in row['submissions']:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
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}}}
|
||||
chained['constant'] = {'class_type': 'TestSharedFlag', 'inputs': {}}
|
||||
previous = '1'
|
||||
value_source = 'constant'
|
||||
for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
|
||||
chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
|
||||
'images': [previous, 0], 'values.value_1': [value_source, 0]}}
|
||||
chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
|
||||
'images': [start, 0], 'context': [start, 3]}}
|
||||
value_source = f'get-{end}'
|
||||
chained[value_source] = {'class_type': 'DistributedImageListGetValue', 'inputs': {
|
||||
'values': [end, 3], 'value_name': 'value_1'}}
|
||||
previous = end
|
||||
chained['8'] = {'class_type': 'TestCapture', 'inputs': {
|
||||
'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2], 'values': ['7', 3]}}
|
||||
saved_chained = copy.deepcopy(chained)
|
||||
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 chained == saved_chained
|
||||
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 VALUE_CAPTURES[1][0] == [{'value_1': True}, {'value_1': True}]
|
||||
assert not runtime._MASTER_JOBS
|
||||
print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
|
||||
# Same public queue path, now with master participation enabled. The
|
||||
# parent queue stays occupied while its local slot processes items.
|
||||
ALLOW_MASTER = True
|
||||
config['settings']['master_delegate_only'] = False
|
||||
participation_counts = {}
|
||||
for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
|
||||
# Keep cold-run exact-once assertions independent of other modes.
|
||||
importlib.import_module(module.__name__ + '.image_list.local_executor')._RETAINED_OUTPUTS = None
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
|
||||
before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
before_worker_stats = []
|
||||
for worker in workers:
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
before_worker_stats.append(await response.json())
|
||||
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', selected), 45)
|
||||
assert prompt == saved_prompt, (mode, 'saved/source graph changed')
|
||||
assert executor.success, (mode, executor.status_messages)
|
||||
assert queued[2] == len(selected)
|
||||
assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
|
||||
local_calls = BRANCH_CALLS[before_local:]
|
||||
assert Counter(UTILITY_CALLS[before_utilities:]) == Counter(local_calls), mode
|
||||
assert len({event['prompt_id'] for event in ITEM_EXECUTIONS[before_events:]}) == len(local_calls)
|
||||
assert local_calls, f'{mode}: master did not process any items'
|
||||
assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
|
||||
remote_calls = []
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for worker, before_stats in zip(workers, before_worker_stats):
|
||||
async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
|
||||
current = await response.json()
|
||||
calls = current['calls'][len(before_stats['calls']):]
|
||||
utilities = current['utility_calls'][len(before_stats['utility_calls']):]
|
||||
submissions = current['submissions'][len(before_stats['submissions']):]
|
||||
events = current['item_executions'][len(before_stats['item_executions']):]
|
||||
assert Counter(utilities) == Counter(calls), (mode, current)
|
||||
assert len(submissions) == len(calls) == len({event['prompt_id'] for event in events})
|
||||
for submission in submissions:
|
||||
assert submission['requested'] == submission['queued'] == ['4'], submission
|
||||
assert '8' in submission['node_ids'] and '9' not in submission['node_ids']
|
||||
assert '5' not in submission['node_ids']
|
||||
remote_calls.extend(calls)
|
||||
assert not current['running'] and not current['pending'], current
|
||||
if mode == 'mixed':
|
||||
assert remote_calls, 'mixed execution failed to use direct workers'
|
||||
else:
|
||||
assert not remote_calls and len(local_calls) == COUNT
|
||||
assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
|
||||
check_carried_values()
|
||||
images, masks, texts = CAPTURED[-1]
|
||||
assert len(images) == len(masks) == len(texts) == COUNT
|
||||
for i, row in enumerate(source_rows()):
|
||||
assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
|
||||
assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
|
||||
assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
|
||||
# Repeat the completed master-only queue with a fresh outer executor.
|
||||
# Every item must be delivered again without recomputing its branch.
|
||||
before_sources, before_local = len(SOURCE_CALLS), len(BRANCH_CALLS)
|
||||
before_models, before_utilities = len(MODEL_CALLS), len(UTILITY_CALLS)
|
||||
before_captures = len(CAPTURED)
|
||||
previous_outputs = CAPTURED[-1]
|
||||
executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
|
||||
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(
|
||||
prompt, None, 'test', []), 45)
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(SOURCE_CALLS) == before_sources + 1
|
||||
assert len(BRANCH_CALLS) == before_local
|
||||
assert len(MODEL_CALLS) == before_models
|
||||
assert len(UTILITY_CALLS) == before_utilities
|
||||
assert len(CAPTURED) == before_captures + 1
|
||||
for before, after in zip(previous_outputs[:2], CAPTURED[-1][:2]):
|
||||
assert len(before) == len(after) == COUNT
|
||||
assert all(torch.equal(a, b) for a, b in zip(before, after))
|
||||
assert previous_outputs[2] == CAPTURED[-1][2]
|
||||
check_carried_values()
|
||||
assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
|
||||
assert prompt == saved_prompt
|
||||
print('IMAGE_LIST_CACHE_REQUEUE_OK unchanged items reused; fresh ordered delivery and cleanup')
|
||||
print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
|
||||
print('IMAGE_LIST_VALUES_OK passthrough all six types; sparse override/add; Get; sequential Get-to-Start; direct/mixed/master-only')
|
||||
assert len(outer_executions) == 5 and all(row['captures'] == 1 for row in outer_executions)
|
||||
print('IMAGE_LIST_TARGETS_OK ' + json.dumps({
|
||||
'modes': ['delegate-only', 'mixed', 'master-only'],
|
||||
'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'],
|
||||
'required_utility_once_per_item': True, 'preview_discarded': True,
|
||||
'saved_prompt_unchanged': prompt == saved_prompt, 'capture_once_per_outer_prompt': True}))
|
||||
ALLOW_MASTER = False
|
||||
config['settings']['master_delegate_only'] = True
|
||||
# Preserve the original direct-runtime cancellation/failure fixture:
|
||||
# it has no outer queue and must not claim target-control evidence.
|
||||
direct_prompt = copy.deepcopy(prompt)
|
||||
del direct_prompt['8'], direct_prompt['9']
|
||||
direct_prompt['3']['inputs']['image'] = ['2', 0]
|
||||
# Real Comfy user-interrupt flag while a real worker executor is active.
|
||||
rows = source_rows()
|
||||
items = runtime.normalize_image_list_inputs(
|
||||
[row['image'] for row in rows], [row['mask'] for row in rows],
|
||||
{f'value_{n}': [row[key] for row in rows]
|
||||
for n, key in ((1, 'reference_image'), (2, 'text'), (3, 'strength'), (4, 'seed'), (6, 'mask'))}
|
||||
| {'value_5': True})
|
||||
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', direct_prompt, items, ['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, include_values=True))
|
||||
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 = copy.deepcopy(items)
|
||||
bad_rows[0]['values']['value_4'] = -1
|
||||
context = runtime.create_master_job('2', '4', direct_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, include_values=True), 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())
|
||||
@@ -0,0 +1,114 @@
|
||||
"""CPU-only native Comfy execution across private-slot lifetimes."""
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
sys.path.insert(0, sys.argv[1])
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
import execution
|
||||
import nodes
|
||||
|
||||
pkg = '_local_retention_probe'
|
||||
for name, path in [(pkg, ROOT), (pkg + '.image_list', ROOT / 'image_list')]:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = [str(path)]
|
||||
sys.modules[name] = module
|
||||
local = importlib.import_module(pkg + '.image_list.local_executor')
|
||||
|
||||
CALLS = []
|
||||
OUTPUTS = []
|
||||
|
||||
class Counted:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',)}}
|
||||
RETURN_TYPES = ('INT',)
|
||||
FUNCTION = 'run'
|
||||
def run(self, value):
|
||||
if value < 0:
|
||||
raise ValueError('intentional retention failure')
|
||||
CALLS.append(value)
|
||||
return (value,)
|
||||
|
||||
class Uncached(Counted):
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float('nan')
|
||||
|
||||
class Capture:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {'required': {'value': ('INT',), 'delivery': ('STRING',)}}
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = 'run'
|
||||
def run(self, value, delivery):
|
||||
OUTPUTS.append((value, delivery))
|
||||
return ()
|
||||
|
||||
nodes.NODE_CLASS_MAPPINGS.update(RetentionCounted=Counted, RetentionUncached=Uncached, RetentionCapture=Capture)
|
||||
|
||||
async def run(slot, value, delivery, uncached=False):
|
||||
prompt = {
|
||||
'source': {'class_type': 'RetentionUncached' if uncached else 'RetentionCounted', 'inputs': {'value': value}},
|
||||
'end': {'class_type': 'RetentionCapture', 'inputs': {'value': ['source', 0], 'delivery': delivery}},
|
||||
}
|
||||
await asyncio.wrap_future(slot.submit(prompt, delivery, 'end'))
|
||||
|
||||
async def main():
|
||||
first = local.LocalExecutionSlot()
|
||||
await run(first, 1, 'first-1')
|
||||
await run(first, 2, 'first-2')
|
||||
await first.close()
|
||||
assert local._RETAINED_OUTPUTS is not None
|
||||
second = local.LocalExecutionSlot()
|
||||
await run(second, 1, 'second-1')
|
||||
await run(second, 2, 'second-2')
|
||||
await run(second, 3, 'changed')
|
||||
await second.close()
|
||||
assert CALLS == [1, 2, 3], CALLS
|
||||
assert OUTPUTS == [(1, 'first-1'), (2, 'first-2'), (1, 'second-1'), (2, 'second-2'), (3, 'changed')], OUTPUTS
|
||||
third = local.LocalExecutionSlot()
|
||||
await run(third, 4, 'uncached-1', True)
|
||||
await run(third, 4, 'uncached-2', True)
|
||||
await third.close()
|
||||
assert CALLS[-2:] == [4, 4]
|
||||
comfy.cli_args.args.cache_none = True
|
||||
disabled = local.LocalExecutionSlot()
|
||||
await run(disabled, 1, 'disabled-1')
|
||||
await run(disabled, 1, 'disabled-2')
|
||||
await disabled.close()
|
||||
assert CALLS[-2:] == [1, 1]
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
comfy.cli_args.args.cache_none = False
|
||||
failed = local.LocalExecutionSlot()
|
||||
await run(failed, 7, 'before-failure')
|
||||
try:
|
||||
await run(failed, -1, 'failure')
|
||||
except RuntimeError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError('failure fixture unexpectedly succeeded')
|
||||
await failed.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
cancelled = local.LocalExecutionSlot()
|
||||
await run(cancelled, 8, 'cancelled')
|
||||
cancelled.request_stop()
|
||||
await cancelled.close()
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
oversized = execution.LRUCache(execution.CacheKeySetInputSignature, max_size=local.NATIVE_CACHE_MAX_ENTRIES)
|
||||
empty = execution.DynamicPrompt({})
|
||||
await oversized.set_prompt(empty, [], execution.IsChangedCache('bound', empty, oversized))
|
||||
oversized.min_generation = oversized.generation
|
||||
oversized.cache = {n: n for n in range(local.NATIVE_CACHE_MAX_ENTRIES + 1)}
|
||||
local._retain_outputs(oversized)
|
||||
assert local._RETAINED_OUTPUTS is None
|
||||
print('PASS: cross-run native cache reuse, fresh delivery, input change, noncacheable nodes, cache_none, cancellation, retention bound')
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,512 @@
|
||||
"""Real CPU Comfy V3 + HTTP proof; invoked only by the opt-in parent test.
|
||||
|
||||
The only harness pieces are deterministic source/assertion nodes and a bounded
|
||||
queue worker. /prompt, /history and /view are Comfy's real handlers; compilation,
|
||||
packing, staging, codecs, bridge execution, response decoding and Collector are
|
||||
unmodified production code. No normal instance or Runpod client is contacted.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import importlib
|
||||
import io as bytes_io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import sys
|
||||
import types
|
||||
import wave
|
||||
from contextlib import asynccontextmanager
|
||||
from fractions import Fraction
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def bootstrap():
|
||||
if len(sys.argv) != 3:
|
||||
raise SystemExit("Use test_media_end_to_end.py with explicit COMFYUI_SOURCE_ROOT")
|
||||
core, run_root = (Path(value).resolve() for value in sys.argv[1:])
|
||||
assert run_root.is_dir() and run_root != core and not run_root.is_relative_to(core)
|
||||
for name in ("CUDA_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"):
|
||||
assert os.environ.get(name) == "", f"CPU isolation missing: {name}"
|
||||
assert not any(name in sys.modules for name in ("torch", "execution", "folder_paths"))
|
||||
sys.path.insert(0, str(core))
|
||||
import comfy.cli_args
|
||||
args = comfy.cli_args.args
|
||||
# Set CPU and all writable roots BEFORE any framework/tensor import.
|
||||
args.cpu = True
|
||||
args.disable_assets = True
|
||||
args.disable_all_custom_nodes = True
|
||||
args.enable_manager = False
|
||||
args.offline = True
|
||||
args.base_directory = str(run_root)
|
||||
args.front_end_root = str(run_root)
|
||||
args.database_url = "sqlite:///:memory:"
|
||||
for name in ("input", "output", "temp", "user"):
|
||||
(run_root / name).mkdir()
|
||||
setattr(args, name + "_directory", str(run_root / name))
|
||||
return run_root
|
||||
|
||||
|
||||
async def check(base):
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
import execution
|
||||
import folder_paths
|
||||
import nodes
|
||||
import server
|
||||
from aiohttp import web
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
from comfy_api.latest._input_impl.video_types import (
|
||||
VideoFromComponents, VideoFromFile, VideoFromList,
|
||||
)
|
||||
from comfy_api.latest._util import VideoComponents
|
||||
from comfy_extras import nodes_audio
|
||||
|
||||
assert not torch.cuda.is_initialized()
|
||||
torch.set_num_threads(1)
|
||||
assert execution.comfy.model_management.get_torch_device().type == "cpu"
|
||||
for name in ("input", "output", "temp", "user"):
|
||||
assert Path(getattr(folder_paths, "get_" + name + "_directory")()).is_relative_to(base)
|
||||
|
||||
# Namespace scaffolding avoids the unrelated pack bootstrap (service/config
|
||||
# side effects). Every imported module below is the real source on disk.
|
||||
package_name = "_media_e2e_acceptance"
|
||||
for suffix in ("", ".api", ".api.orchestration", ".nodes"):
|
||||
package = types.ModuleType(package_name + suffix)
|
||||
package.__path__ = [str(ROOT.joinpath(*suffix.lstrip(".").split(".")))] if suffix else [str(ROOT)]
|
||||
sys.modules[package.__name__] = package
|
||||
def load(name):
|
||||
return importlib.import_module(package_name + "." + name)
|
||||
|
||||
asset_manager = default_asset_manager()
|
||||
ps = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
ps.distributed_pending_jobs = {}
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
compiler = load("api.orchestration.serverless_media_prompt")
|
||||
packer = load("api.orchestration.serverless_typed_inputs")
|
||||
decoder = load("api.orchestration.serverless_typed_results")
|
||||
serializer = load("api.orchestration.runpod_client").serialize_request
|
||||
handler = load("serverless_bridge.handler")
|
||||
contract = load("serverless_bridge.contract")
|
||||
video_codec = load("serverless_bridge.video_codec")
|
||||
bridge = load("serverless_bridge")
|
||||
collector = load("nodes.v3").DistributedCollector
|
||||
extension = await bridge.comfy_entrypoint()
|
||||
bridge_classes = await extension.get_node_list()
|
||||
assert {cls.GET_SCHEMA().node_id for cls in bridge_classes} == {"MediaImportV1", "MediaExportV1"}
|
||||
sockets = list(contract.SOCKET_KINDS)
|
||||
collector_sockets = ["images", "audio"]
|
||||
kinds = list(contract.SOCKET_KINDS.values())
|
||||
assert kinds == ["IMAGE", "AUDIO", "MASK", "STRING", "VIDEO"]
|
||||
io_types = [io.Image, io.Audio, io.Mask, io.String, io.Video]
|
||||
fixtures, expected_sinks, seen_sinks, source_calls = {}, {}, {}, []
|
||||
|
||||
class AcceptanceMediaSource(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="AcceptanceMediaSource", inputs=[io.String.Input("key")],
|
||||
outputs=[kind.Output(is_output_list=True) for kind in io_types])
|
||||
|
||||
@classmethod
|
||||
def execute(cls, key):
|
||||
source_calls.append(key)
|
||||
return io.NodeOutput(*(fixtures[key][socket] for socket in sockets))
|
||||
|
||||
def tensor_equal(actual, expected):
|
||||
assert actual.device.type == "cpu" and actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
# Byte comparison preserves signed zero and catches dtype narrowing.
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
def exact(kind, actual, expected):
|
||||
if kind in ("IMAGE", "MASK"):
|
||||
tensor_equal(actual, expected)
|
||||
elif kind == "AUDIO":
|
||||
assert actual["sample_rate"] == expected["sample_rate"]
|
||||
tensor_equal(actual["waveform"], expected["waveform"])
|
||||
elif kind == "STRING":
|
||||
assert type(actual) is str and actual.encode("utf-8") == expected.encode("utf-8")
|
||||
else:
|
||||
assert video_codec.encode_video(actual) == expected
|
||||
if type(actual) is VideoFromFile:
|
||||
buffer = vars(actual)["_VideoFromFile__file"]
|
||||
assert type(buffer) is bytes_io.BytesIO and not buffer.closed
|
||||
position = buffer.tell()
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
assert len(list(container.decode(video=0))) == 3
|
||||
buffer.seek(position)
|
||||
elif type(actual) is VideoFromList:
|
||||
for item in actual.videos:
|
||||
exact("VIDEO", item, video_codec.encode_video(item))
|
||||
else:
|
||||
assert type(actual) is VideoFromComponents
|
||||
components = actual.get_components()
|
||||
assert type(components.frame_rate) is Fraction
|
||||
assert components.images.dtype == torch.float64
|
||||
assert components.alpha.dtype == torch.bfloat16
|
||||
assert components.audio["waveform"].dtype == torch.float64
|
||||
assert actual.get_bit_depth() == 10 and actual.get_color_space() == "HDR PQ"
|
||||
|
||||
class AcceptanceMediaAssert(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(node_id="AcceptanceMediaAssert", inputs=[io.String.Input("tag"),
|
||||
*[kind.Input(socket, optional=True, **({"force_input": True} if socket == "texts" else {}))
|
||||
for socket, kind in zip(sockets, io_types)]], outputs=[],
|
||||
is_input_list=True, is_output_node=True)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, tag, **media):
|
||||
assert len(tag) == 1 and tag[0] not in seen_sinks
|
||||
expected = expected_sinks[tag[0]]
|
||||
assert set(media) == set(expected)
|
||||
for socket in expected:
|
||||
assert len(media[socket]) == len(expected[socket]), (tag, socket)
|
||||
# Existing Collector outputs are a batched IMAGE and one AUDIO,
|
||||
# not lists. Native V3 wraps each output once for this list sink.
|
||||
for actual, wanted in zip(media[socket], expected[socket]):
|
||||
exact(contract.SOCKET_KINDS[socket], actual, wanted)
|
||||
counts = {socket: len(media[socket]) for socket in expected}
|
||||
seen_sinks[tag[0]] = counts
|
||||
return io.NodeOutput(ui={"accepted": [counts]})
|
||||
|
||||
for cls in [*bridge_classes, collector, AcceptanceMediaSource, AcceptanceMediaAssert]:
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
cls.GET_NODE_INFO_V1()
|
||||
nodes.NODE_CLASS_MAPPINGS[cls.GET_SCHEMA().node_id] = cls
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
def add_bridge_exports(workflow, manifest, request_id, participant, job_id, links):
|
||||
"""Exercise bridge-only kinds without inventing Collector sockets."""
|
||||
for socket, link in links.items():
|
||||
export_id = "bridge-export:" + socket
|
||||
binding = {"binding_id": "bridge-binding:" + socket, "collector_id": "bridge-only",
|
||||
"job_id": job_id, "participant_id": participant, "node_id": export_id,
|
||||
"socket": socket, "kind": contract.SOCKET_KINDS[socket]}
|
||||
assert export_id not in workflow
|
||||
manifest.append(binding)
|
||||
workflow[export_id] = node("MediaExportV1", request_id=request_id,
|
||||
binding_json=json.dumps(binding), **{socket: link})
|
||||
contract.validate_bindings(manifest)
|
||||
|
||||
def tiny_video(variant):
|
||||
buffer = bytes_io.BytesIO()
|
||||
with av.open(buffer, "w", format="matroska") as container:
|
||||
stream = container.add_stream("ffv1", rate=Fraction(24, 1))
|
||||
stream.width = stream.height = 16
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for level in (16 + variant, 64 + variant, 128 + variant):
|
||||
frame = av.VideoFrame.from_ndarray(np.full((16, 16, 3), level, dtype=np.uint8), format="rgb24")
|
||||
for packet in stream.encode(frame):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode():
|
||||
container.mux(packet)
|
||||
return buffer.getvalue()
|
||||
|
||||
def native_fixtures(variant):
|
||||
image = torch.arange(48, dtype=torch.float64).reshape(2, 2, 3, 4) / 997 + variant / 100
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
mask = torch.tensor([[[-0.0, 0.125, 0.875]]], dtype=torch.bfloat16)
|
||||
audio = {"waveform": torch.arange(38, dtype=torch.float64).reshape(1, 2, 19) / 997,
|
||||
"sample_rate": 44100}
|
||||
audio["waveform"][0, 0, 0] = -0.0
|
||||
component = VideoFromComponents(VideoComponents(
|
||||
images=image[..., :3], frame_rate=Fraction(30000, 1001), alpha=mask.reshape(1, 1, 3, 1).repeat(2, 2, 1, 1),
|
||||
audio=audio, metadata={"title": "雪", "nested": [True, None, -4, 0.125]}),
|
||||
bit_depth=10, color_space="HDR PQ")
|
||||
payload = tiny_video(variant)
|
||||
file = VideoFromFile(bytes_io.BytesIO(payload), start_time=-0.0, duration=0, crop=(0, 0, 8, 8))
|
||||
sequence = VideoFromList([VideoFromFile(bytes_io.BytesIO(payload)), file], complete_audio=audio)
|
||||
return {"images": [image, image[:1].clone()], "audio": [audio], "masks": [mask],
|
||||
"texts": [f" café 🎥 {variant}\r\n雪 ", ""], "video": [component, sequence]}
|
||||
|
||||
for key, variant in (("local", 0), ("worker-a", 1), ("worker-b", 2)):
|
||||
fixtures[key] = native_fixtures(variant)
|
||||
native_root = base / "master-input"
|
||||
native_root.mkdir()
|
||||
wav = native_root / "stereo.wav"
|
||||
pcm = np.array([-32768, -12001, -1, 0, 1, 8193, 32767, 17] * 2, dtype="<i2")
|
||||
with wave.open(str(wav), "wb") as out:
|
||||
out.setnchannels(2)
|
||||
out.setsampwidth(2)
|
||||
out.setframerate(44100)
|
||||
out.writeframes(pcm.tobytes())
|
||||
waveform, sample_rate = nodes_audio.load(str(wav))
|
||||
clip = native_root / "native.mkv"
|
||||
clip.write_bytes(tiny_video(3))
|
||||
native_expected = copy.copy(fixtures["worker-a"])
|
||||
native_expected["audio"] = [{"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}]
|
||||
native_expected["video"] = [VideoFromFile(str(clip))]
|
||||
expectations = {}
|
||||
for key, values in (("worker-a", native_expected), ("worker-b", fixtures["worker-b"])):
|
||||
expectations[key] = {socket: [video_codec.encode_video(v, allowed_roots=[native_root]) if socket == "video" else v
|
||||
for v in values[socket]] for socket in sockets}
|
||||
|
||||
# Native routes only, bound to a new ephemeral loopback listener. No fake
|
||||
# history, permissive validator, bridge/codec mock, or existing service.
|
||||
allowed_paths = {"/prompt", "/history/{prompt_id}", "/view", "/queue"}
|
||||
ps.app.add_routes([route for route in ps.routes if route.path in allowed_paths])
|
||||
completed = []
|
||||
stop = asyncio.Event()
|
||||
|
||||
async def queue_worker():
|
||||
while not stop.is_set():
|
||||
item = await asyncio.to_thread(ps.prompt_queue.get, 0.1)
|
||||
if item is None:
|
||||
continue
|
||||
queued, item_id = item
|
||||
assert len(queued) in (5, 6), "Unsupported core PromptQueue tuple; adapt acceptance worker explicitly"
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, queued[2], queued[1], queued[3], queued[4])
|
||||
status = execution.PromptQueue.ExecutionStatus(
|
||||
status_str="success" if executor.success else "error", completed=executor.success,
|
||||
messages=executor.status_messages)
|
||||
ps.prompt_queue.task_done(item_id, executor.history_result, status=status)
|
||||
completed.append((queued, executor.success, executor.status_messages))
|
||||
|
||||
@asynccontextmanager
|
||||
async def serve():
|
||||
runner = web.AppRunner(ps.app)
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
ALLOWED_CONNECTIONS.add(("127.0.0.1", port))
|
||||
worker = asyncio.create_task(queue_worker())
|
||||
try:
|
||||
yield f"http://127.0.0.1:{port}"
|
||||
finally:
|
||||
stop.set()
|
||||
await asyncio.wait_for(worker, 10)
|
||||
await runner.cleanup()
|
||||
ALLOWED_CONNECTIONS.clear()
|
||||
|
||||
grouped_results = {}
|
||||
async with serve() as url:
|
||||
for key in ("worker-a", "worker-b"):
|
||||
participant = "serverless:" + key
|
||||
request_id = "acceptance-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedCollector", load_balance=False, transfer_format="exact",
|
||||
multi_job_id="job-main", worker_id=participant,
|
||||
**{socket: ["source", index] for index, socket in enumerate(collector_sockets)}),
|
||||
"audio-collector": node("DistributedCollector", load_balance=False,
|
||||
multi_job_id="job-audio", worker_id=participant, audio=["source", 1])}
|
||||
if key == "worker-a":
|
||||
graph.update({"wav": node("LoadAudio", audio=wav.name), "clip": node("LoadVideo", file=clip.name)})
|
||||
graph["collector"]["inputs"]["audio"] = ["wav", 0]
|
||||
before = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
assert graph == before
|
||||
bridge_links = {socket: ["source", sockets.index(socket)] for socket in sockets[2:]}
|
||||
if key == "worker-a":
|
||||
workflow["clip"] = copy.deepcopy(graph["clip"])
|
||||
bridge_links["video"] = ["clip", 0]
|
||||
add_bridge_exports(workflow, manifest, request_id, participant, "job-bridge", bridge_links)
|
||||
assert len(manifest) == 6 and {b["kind"] for b in manifest} == set(kinds)
|
||||
folder_paths.set_input_directory(str(native_root))
|
||||
workflow_before = copy.deepcopy(workflow)
|
||||
images, artifacts, refs = await packer.prepare_media_inputs(workflow)
|
||||
assert images == [] and workflow == workflow_before
|
||||
if key == "worker-a":
|
||||
assert {a["kind"] for a in artifacts} == {"AUDIO", "VIDEO"} and len(artifacts) == 2
|
||||
assert len(refs) == 2
|
||||
assert {ref["input_name"] for ref in refs} == {"audio", "video"}
|
||||
assert all(workflow[ref["node_id"]]["class_type"] == "MediaExportV1" for ref in refs)
|
||||
# Endpoint cannot fall back to the master's local files/loaders.
|
||||
wav.unlink()
|
||||
clip.unlink()
|
||||
print("NATIVE_INPUTS_AND_DIRECT_EXPORT_REFS_OK WAV + native file VIDEO")
|
||||
else:
|
||||
assert not artifacts and not refs
|
||||
folder_paths.set_input_directory(str(base / "input"))
|
||||
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
|
||||
"artifacts": artifacts, "inputs": refs, "outputs": manifest}
|
||||
body = json.loads(serializer(workflow, images, 30, media=media))
|
||||
assert decoder.validate_prepared_job({"worker": {"id": participant, "media_profile": contract.PROFILE,
|
||||
"timeout_seconds": 30}, "workflow": workflow, "images": images,
|
||||
"media": media, "manifest": manifest}) is True
|
||||
before = copy.deepcopy(body["input"])
|
||||
result = await handler.execute_media_request(body["input"], comfy_url=url,
|
||||
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
|
||||
assert body["input"] == before
|
||||
assert completed and completed[-1][1], completed
|
||||
queued = completed[-1][0]
|
||||
assert set(queued[4]) == {b["node_id"] for b in manifest}
|
||||
submitted = queued[2]
|
||||
if key == "worker-a":
|
||||
assert "wav" not in submitted and "clip" not in submitted
|
||||
imports = [n for n in submitted.values() if n["class_type"] == "MediaImportV1"]
|
||||
assert len(imports) == 2
|
||||
assert {json.loads(n["inputs"]["artifact_json"])["kind"] for n in imports} == {"AUDIO", "VIDEO"}
|
||||
assert not list((base / "input").rglob("*.bin"))
|
||||
assert not list((base / "temp").rglob("*.bin"))
|
||||
if key == "worker-b":
|
||||
# Retire the originals as well as handler-owned staging/export
|
||||
# files before constructing the result's lazy native videos.
|
||||
for original in fixtures[key]["video"][1].videos:
|
||||
vars(original)["_VideoFromFile__file"].close()
|
||||
# Exercise the wire format, not an in-memory handler-only object.
|
||||
result = json.loads(json.dumps(result, allow_nan=False))
|
||||
grouped = await decoder.collect_typed_results(result, media, manifest)
|
||||
assert set(grouped) == {"job-main", "job-audio", "job-bridge"}
|
||||
assert set(grouped["job-main"]) == set(collector_sockets)
|
||||
assert set(grouped["job-bridge"]) == set(sockets[2:])
|
||||
exact("AUDIO", grouped["job-audio"]["audio"][0], fixtures[key]["audio"][0])
|
||||
all_media = {**grouped["job-main"], **grouped["job-bridge"]}
|
||||
for socket, kind in zip(sockets, kinds):
|
||||
values, expected = all_media[socket], expectations[key][socket]
|
||||
assert len(values) == len(expected), (key, socket)
|
||||
for actual, wanted in zip(values, expected):
|
||||
exact(kind, actual, wanted)
|
||||
if key == "worker-b":
|
||||
original = fixtures[key]["video"][0].get_components()
|
||||
decoded = grouped["job-bridge"]["video"][0].get_components()
|
||||
for field in ("images", "alpha"):
|
||||
assert getattr(decoded, field).data_ptr() != getattr(original, field).data_ptr()
|
||||
assert decoded.audio["waveform"].data_ptr() != original.audio["waveform"].data_ptr()
|
||||
assert decoded.metadata == original.metadata and decoded.metadata is not original.metadata
|
||||
grouped_results[participant] = grouped
|
||||
# Late invalid artifact must reject the entire grouped response.
|
||||
corrupt = copy.deepcopy(result)
|
||||
corrupt["media"]["artifacts"][-1]["sha256"] = "0" * 64
|
||||
try:
|
||||
await decoder.collect_typed_results(corrupt, media, manifest)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
raise AssertionError("Result decoder accepted a corrupt final artifact")
|
||||
assert ps.distributed_pending_jobs == {}
|
||||
# Exercise user-selected image return encoding through the actual
|
||||
# compiler, HTTP handler, V3 exporter and atomic response decoder.
|
||||
for fmt in ("legacy_png", "png", "jpeg", "h264"):
|
||||
key = "encoded-" + fmt
|
||||
fixtures[key] = native_fixtures(3)
|
||||
# Keep full alpha precision and use nontrivial image dimensions.
|
||||
torch.manual_seed(41)
|
||||
rgba = torch.rand((2, 64, 64, 4), dtype=torch.float64)
|
||||
fixtures[key]["images"] = [rgba]
|
||||
request_id = "request-" + key
|
||||
graph = {"source": node("AcceptanceMediaSource", key=key),
|
||||
"collector": node("DistributedCollector", transfer_format=fmt,
|
||||
transfer_quality=73, multi_job_id="job-" + key,
|
||||
worker_id="serverless:" + key,
|
||||
**{socket: ["source", i] for i, socket in enumerate(collector_sockets)})}
|
||||
if fmt == "legacy_png":
|
||||
# A saved Collector with no newly selected controls still uses
|
||||
# legacy_png / 90, mapped to PNG by the serverless compiler.
|
||||
graph["collector"]["inputs"].pop("transfer_format")
|
||||
graph["collector"]["inputs"].pop("transfer_quality")
|
||||
original = copy.deepcopy(graph)
|
||||
workflow, manifest = compiler.compile_media_prompt(graph, request_id)
|
||||
if fmt == "legacy_png":
|
||||
binding = next(b for b in manifest if b["kind"] == "IMAGE")
|
||||
controls = workflow[binding["node_id"]]["inputs"]
|
||||
assert (controls["transfer_format"], controls["transfer_quality"]) == ("png", 90)
|
||||
add_bridge_exports(workflow, manifest, request_id, "serverless:" + key, "job-bridge-" + key,
|
||||
{socket: ["source", sockets.index(socket)] for socket in sockets[2:]})
|
||||
media = {"version": 1, "profile": contract.PROFILE, "request_id": request_id,
|
||||
"artifacts": [], "inputs": [], "outputs": manifest}
|
||||
body = json.loads(serializer(workflow, [], 30, media=media))
|
||||
result = await handler.execute_media_request(body["input"], comfy_url=url,
|
||||
input_dir=base / "input", temp_dir=base / "temp", timeout_seconds=30)
|
||||
assert graph == original
|
||||
grouped = await decoder.collect_typed_results(result, media, manifest)
|
||||
actual = {**grouped["job-" + key], **grouped["job-bridge-" + key]}
|
||||
image = actual["images"][0]
|
||||
assert image.shape == rgba.shape and image.dtype == rgba.dtype
|
||||
assert torch.equal(image[..., 3], rgba[..., 3]), "alpha must not be discarded/narrowed"
|
||||
assert not torch.equal(image[..., :3], rgba[..., :3]), "selection must actually encode"
|
||||
if fmt in ("legacy_png", "png"):
|
||||
pixels = (rgba[..., :3].float().clamp(0, 1) * 255).to(torch.uint8)
|
||||
assert torch.equal(image[..., :3], (pixels.float() / 255).to(rgba.dtype))
|
||||
for socket, kind in zip(sockets[1:], kinds[1:]):
|
||||
for got, wanted in zip(actual[socket], fixtures[key][socket]):
|
||||
exact(kind, got, video_codec.encode_video(wanted) if kind == "VIDEO" else wanted)
|
||||
assert not list((base / "temp").rglob("*.bin"))
|
||||
print("IMAGE_TRANSFER_HTTP_OK default legacy_png -> png/90; png jpeg h264; alpha and other media exact")
|
||||
print("REAL_HTTP_BRIDGE_RESULTS_OK Collector IMAGE/AUDIO + independent MASK/STRING/VIDEO bridge exports; native VIDEO survives cleanup")
|
||||
print("ATOMIC_DECODE_REJECTION_OK corrupt final artifact rejected before delivery")
|
||||
|
||||
participants = list(grouped_results)
|
||||
|
||||
def expected_batches(ordered, include_images=True):
|
||||
audio_items = [item for media in ordered for item in media["audio"]]
|
||||
expected = {"audio": [{"waveform": torch.cat([item["waveform"] for item in audio_items], dim=-1),
|
||||
"sample_rate": audio_items[0]["sample_rate"]}]}
|
||||
if include_images:
|
||||
expected["images"] = [torch.cat([image for media in ordered for image in media["images"]])]
|
||||
return expected
|
||||
|
||||
for delegate in (False, True):
|
||||
tag = "delegate" if delegate else "master"
|
||||
job = "collect-" + tag
|
||||
queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await queue.put({"worker_id": participant, "media": grouped_results[participant]["job-main"], "is_last": True})
|
||||
ps.distributed_pending_jobs[job] = queue
|
||||
ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-main"] for p in participants]
|
||||
expected_sinks[tag] = expected_batches(ordered)
|
||||
graph = {"collector": node("DistributedCollector", load_balance=False, multi_job_id=job,
|
||||
enabled_worker_ids=json.dumps(participants), delegate_only=delegate, transfer_format="exact"),
|
||||
"assert": node("AcceptanceMediaAssert", tag=tag,
|
||||
**{socket: ["collector", index] for index, socket in enumerate(collector_sockets)})}
|
||||
if not delegate:
|
||||
graph["source"] = node("AcceptanceMediaSource", key="local")
|
||||
graph["collector"]["inputs"].update({socket: ["source", index] for index, socket in enumerate(collector_sockets)})
|
||||
# A second audio-only Collector consumes its own job group from the same
|
||||
# atomic response. It must return None for images and batch audio in order.
|
||||
audio_job, audio_tag = job + "-audio", tag + "-audio"
|
||||
audio_queue = asyncio.Queue()
|
||||
for participant in reversed(participants):
|
||||
await audio_queue.put({"worker_id": participant,
|
||||
"media": grouped_results[participant]["job-audio"], "is_last": True})
|
||||
ps.distributed_pending_jobs[audio_job] = audio_queue
|
||||
audio_ordered = ([] if delegate else [fixtures["local"]]) + [grouped_results[p]["job-audio"] for p in participants]
|
||||
expected_sinks[audio_tag] = expected_batches(audio_ordered, include_images=False)
|
||||
graph["audio-collector"] = node("DistributedCollector", load_balance=False, multi_job_id=audio_job,
|
||||
enabled_worker_ids=json.dumps(participants), delegate_only=delegate)
|
||||
graph["audio-assert"] = node("AcceptanceMediaAssert", tag=audio_tag, audio=["audio-collector", 1])
|
||||
if not delegate:
|
||||
graph["audio-collector"]["inputs"]["audio"] = ["source", 1]
|
||||
valid = await execution.validate_prompt(job, graph, None)
|
||||
assert valid[0] and not valid[3], valid
|
||||
executor = execution.PromptExecutor(ps, cache_args={"ram": 0, "ram_inactive": 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, job, {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert executor.history_result["outputs"]["assert"]["accepted"] == [seen_sinks[tag]]
|
||||
assert executor.history_result["outputs"]["audio-assert"]["accepted"] == [seen_sinks[audio_tag]]
|
||||
assert job not in ps.distributed_pending_jobs and audio_job not in ps.distributed_pending_jobs
|
||||
assert source_calls == ["worker-a", "worker-b", "encoded-legacy_png", "encoded-png", "encoded-jpeg", "encoded-h264", "local"], source_calls
|
||||
assert set(seen_sinks) == {"master", "delegate", "master-audio", "delegate-audio"}
|
||||
|
||||
assert not list((base / "output").rglob("*"))
|
||||
assert not torch.cuda.is_initialized()
|
||||
print("MASTER_AND_DELEGATE_ORDER_OK", json.dumps(seen_sinks, sort_keys=True))
|
||||
print("MEDIA_END_TO_END_CPU_OK real PromptExecutor/HTTP; no queue-orchestration or provider acceptance claimed")
|
||||
|
||||
|
||||
ALLOWED_CONNECTIONS = set()
|
||||
|
||||
|
||||
def network_guard(event, args):
|
||||
"""Fail closed on provider/normal-instance traffic, including during imports."""
|
||||
if event == "socket.connect":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address in ALLOWED_CONNECTIONS, "Acceptance forbids non-owned network connections"
|
||||
elif event == "socket.bind":
|
||||
sock, address = args
|
||||
if sock.family in (socket.AF_INET, socket.AF_INET6):
|
||||
assert address == ("127.0.0.1", 0), "Acceptance binds only ephemeral loopback"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.addaudithook(network_guard)
|
||||
asyncio.run(asyncio.wait_for(check(bootstrap()), timeout=120))
|
||||
@@ -0,0 +1,518 @@
|
||||
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
|
||||
|
||||
No HTTP listener, workers, model downloads or live installation changes.
|
||||
The original contracts were captured from 32ac027 using the same core.
|
||||
"""
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
|
||||
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
|
||||
|
||||
# Inputs intentionally added after the V1 baseline was captured. They must be
|
||||
# optional, appended after the baseline inputs, and keep legacy behaviour as
|
||||
# their default so saved workflows are unaffected.
|
||||
ADDED_OPTIONAL_INPUTS = {
|
||||
'DistributedCollector': {
|
||||
'transfer_format': [
|
||||
['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1', 'exact'],
|
||||
{'default': 'legacy_png'},
|
||||
],
|
||||
'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
|
||||
},
|
||||
}
|
||||
|
||||
# Nodes added after the V1 baseline; not part of saved-workflow contracts.
|
||||
# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
|
||||
ADDED_NODES = {'DistributedWorkerSink', 'DistributedImageListStart', 'DistributedImageListEnd',
|
||||
'DistributedImageListGetValue'}
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
kind = value[0]
|
||||
opts = dict(value[1]) if len(value) > 1 else {}
|
||||
if kind == 'STRING':
|
||||
# V3 serializes the same default single-line widget explicitly.
|
||||
opts.setdefault('multiline', False)
|
||||
if kind == 'COMBO':
|
||||
kind = opts['options']
|
||||
opts = {key: val for key, val in opts.items() if key != 'options'}
|
||||
if isinstance(kind, list):
|
||||
# Single-selection is the V1 dropdown default as well.
|
||||
opts.setdefault('multiselect', False)
|
||||
return [kind, opts]
|
||||
|
||||
|
||||
async def map_node(cls, values, extra=None):
|
||||
"""Use the actual input parser and sanitized V3 executor clones."""
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
|
||||
assert not missing, missing
|
||||
# Use a worker thread for synchronous nodes: collector/upscale bridge back
|
||||
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
|
||||
def run():
|
||||
return asyncio.run(execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
|
||||
return await asyncio.to_thread(run)
|
||||
|
||||
|
||||
def check_saved_workflows(mapping):
|
||||
checked = 0
|
||||
occurrences = 0
|
||||
for path in sorted((ROOT / 'workflows').glob('*.json')):
|
||||
workflow = json.loads(path.read_text())
|
||||
local_nodes = {node['id']: node for node in workflow['nodes']}
|
||||
for saved in workflow['nodes']:
|
||||
if saved['type'] not in mapping:
|
||||
continue
|
||||
occurrences += 1
|
||||
info = mapping[saved['type']].GET_NODE_INFO_V1()
|
||||
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
|
||||
for socket in saved.get('inputs', []):
|
||||
assert socket['name'] in inputs, (path.name, saved['id'], socket)
|
||||
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
|
||||
for index, output in enumerate(saved.get('outputs', [])):
|
||||
assert output['type'] == info['output'][index], (path.name, index)
|
||||
for link in workflow['links']:
|
||||
if link[1] == saved['id']:
|
||||
assert link[2] < len(info['output']), (path.name, link)
|
||||
target = local_nodes[link[3]]['inputs'][link[4]]
|
||||
assert target['type'] == info['output'][link[2]], (path.name, link)
|
||||
checked += 1
|
||||
assert checked == 5 and occurrences == 10, (checked, occurrences)
|
||||
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
|
||||
|
||||
|
||||
async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
result = await map_node(mapping['DistributedSeed'],
|
||||
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (126,)
|
||||
values = json.dumps({'_type': 'INT', '3': '17'})
|
||||
result = await map_node(mapping['DistributedValue'],
|
||||
{'default_value': '4', 'worker_values': values,
|
||||
'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (17,)
|
||||
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
|
||||
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
|
||||
{'extra_pnginfo': metadata})
|
||||
assert result[0].result == ('model.ckpt',)
|
||||
assert result[0].ui == {'text': ['model.ckpt']}
|
||||
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
|
||||
result = await map_node(mapping['DistributedEmptyImage'],
|
||||
{'height': 8, 'width': 8, 'channels': 3})
|
||||
empty = result[0].result[0]
|
||||
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
|
||||
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
|
||||
for count in (1, 3, 10):
|
||||
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat(result[0].result[:count]), images)
|
||||
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
|
||||
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
|
||||
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
|
||||
collector = mapping['DistributedCollector']
|
||||
prepared, missing, hidden = execution.get_input_data(
|
||||
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
|
||||
assert not missing
|
||||
prepared['images'] = [images[:2], images[2:]]
|
||||
prepared['audio'] = [audio, audio]
|
||||
result = await execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
|
||||
assert torch.equal(result[0].result[0], images)
|
||||
assert result[0].result[1]['waveform'].shape[-1] == 66
|
||||
# Exercise real async aggregation without an HTTP endpoint or worker.
|
||||
queue = asyncio.Queue()
|
||||
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
|
||||
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
|
||||
result = await map_node(collector,
|
||||
{'images': images[:2], 'load_balance': False,
|
||||
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
|
||||
assert torch.equal(result[0].result[0], images[:3])
|
||||
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
|
||||
|
||||
# The registered node must forward the transfer options to the runtime
|
||||
# helper, and default to the legacy transfer when they are not set.
|
||||
collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
|
||||
forwarded = []
|
||||
def fake_run(self, **kwargs):
|
||||
forwarded.append(kwargs)
|
||||
return (kwargs['images'], kwargs['audio'])
|
||||
with patch.object(collector_runtime, 'run', fake_run):
|
||||
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
|
||||
'transfer_format': 'av1', 'transfer_quality': 70})
|
||||
await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
|
||||
# Set values arrive list-wrapped (is_input_list); unset ones use the
|
||||
# execute() defaults. The runtime helper unwraps both forms.
|
||||
assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
|
||||
(['av1'], [70]), ('legacy_png', 90)], forwarded
|
||||
|
||||
# Prove V3's per-execution class clones do not leak helper instance state;
|
||||
# replace only the GPU/model boundary, not the input parser or executor.
|
||||
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
|
||||
seen = []
|
||||
def fake_upscale(self, *args):
|
||||
assert not hasattr(self, 'acceptance_marker')
|
||||
self.acceptance_marker = True
|
||||
seen.append(args)
|
||||
return (args[0],)
|
||||
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
|
||||
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
|
||||
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
|
||||
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
|
||||
'force_uniform_tiles': True, 'tiled_decode': False,
|
||||
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
|
||||
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
|
||||
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
|
||||
with patch.object(upscale_runtime, 'run', fake_upscale):
|
||||
for _ in range(2):
|
||||
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
|
||||
assert result[0].result[0] is images
|
||||
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
|
||||
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
|
||||
'tile_indices', 'dynamic_threshold')), seen[0]
|
||||
import math
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
||||
'pass_through': True}},
|
||||
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
|
||||
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('v3-graph', graph, None)
|
||||
assert valid[0], valid
|
||||
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
|
||||
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
|
||||
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
|
||||
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
|
||||
assert not rejected[0]
|
||||
errors = [error for entry in rejected[3].values() for error in entry['errors']]
|
||||
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
|
||||
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
history = executor.history_result
|
||||
record = history['outputs']['4']['images'][0]
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
|
||||
|
||||
# A worker prompt terminated by the sink validates and runs through the
|
||||
# real executor without writing anything to the temp directory.
|
||||
sink_graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 4, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
||||
'pass_through': True}},
|
||||
'3': {'class_type': 'DistributedWorkerSink', 'inputs': {'value': ['2', 0]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('v3-sink', sink_graph, None)
|
||||
assert valid[0], valid
|
||||
assert valid[2] == ['3'], valid[2]
|
||||
collected = []
|
||||
def record_run(self, **kwargs):
|
||||
collected.append(kwargs['images'])
|
||||
return (kwargs['images'][0], None)
|
||||
with tempfile.TemporaryDirectory(prefix='v3-sink-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp), patch.object(collector_runtime, 'run', record_run):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, sink_graph, 'v3-sink', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
assert len(collected) == 1 and collected[0][0].shape == (4, 8, 8, 3), collected
|
||||
assert list(Path(temp).rglob('*')) == [], list(Path(temp).rglob('*'))
|
||||
assert executor.history_result['outputs'].get('3', {}) == {}, executor.history_result
|
||||
print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing')
|
||||
|
||||
|
||||
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, *, model_callback=None, include_values=False):
|
||||
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], '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', 3],
|
||||
}},
|
||||
'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', 4]},
|
||||
}
|
||||
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) + len(ADDED_NODES) == 12
|
||||
assert set(mapping) == set(BASELINE) | ADDED_NODES
|
||||
for node_id in ADDED_NODES:
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
|
||||
assert 'DistributedMediaCollector' not in mapping
|
||||
collector_info = mapping['DistributedCollector'].GET_NODE_INFO_V1()
|
||||
assert collector_info['output'] == ['IMAGE', 'AUDIO'], collector_info
|
||||
assert not any(collector_info['output_is_list']), collector_info
|
||||
assert list(collector_info['input']['optional']) == [
|
||||
'images', 'audio', 'transfer_format', 'transfer_quality']
|
||||
sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
|
||||
assert sink_info['output_node'] and sink_info['dev_only'], sink_info
|
||||
for node_id, cls in mapping.items():
|
||||
if node_id in ADDED_NODES:
|
||||
continue
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
|
||||
old = dict(BASELINE[node_id])
|
||||
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
|
||||
# V1's ByPassTypeTuple advertises '*' when indexed, while its
|
||||
# underlying tuple and existing frontend declare IMAGE/AUDIO.
|
||||
# Native V3 declares all ten existing typed sockets explicitly.
|
||||
assert old['output'] == ['*'] * 10
|
||||
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
|
||||
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
|
||||
for group in ('required', 'optional'):
|
||||
old_inputs = old['input'].get(group, {})
|
||||
new_inputs = dict(new['input'].get(group, {}))
|
||||
if group == 'optional':
|
||||
for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
|
||||
assert name in new_inputs, (node_id, name, 'added input missing')
|
||||
actual = normalized_input(new_inputs.pop(name))
|
||||
kind, opts = normalized_input(expected)
|
||||
assert actual[0] == kind, (node_id, name, actual)
|
||||
for key, value in opts.items():
|
||||
assert actual[1].get(key) == value, (node_id, name, key, actual)
|
||||
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
|
||||
for name, original in old_inputs.items():
|
||||
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
|
||||
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
|
||||
assert old[key] == new[key], (node_id, key, old[key], new[key])
|
||||
# Standard context lives in cls.hidden; orchestrator metadata remains
|
||||
# accepted by its original kwarg name, without creating new widgets.
|
||||
signature = inspect.signature(cls.execute)
|
||||
for name, field in old['input'].get('hidden', {}).items():
|
||||
if isinstance(field, list):
|
||||
assert cls.GET_SCHEMA().accept_all_inputs, node_id
|
||||
assert name in signature.parameters, (node_id, name)
|
||||
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
|
||||
else:
|
||||
assert name in new['input']['hidden'], (node_id, name)
|
||||
expected_hidden = []
|
||||
if node_id == 'DistributedModelName':
|
||||
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
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'
|
||||
runtime = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
assert [input.id for input in end_schema.inputs] == ['images', 'masks', 'texts', 'context'] + list(runtime.VALUE_NAMES)
|
||||
assert [output.io_type for output in end_schema.outputs] == [
|
||||
'IMAGE', 'MASK', 'STRING', 'DISTRIBUTED_IMAGE_LIST_VALUES']
|
||||
assert end_schema.outputs[3].id == end_schema.outputs[3].display_name == 'values'
|
||||
assert all(output.is_output_list for output in end_schema.outputs)
|
||||
for override in end_schema.inputs[4:]:
|
||||
assert override.optional and override.lazy
|
||||
assert override.get_io_type() == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN'
|
||||
getter = mapping['DistributedImageListGetValue']
|
||||
getter_schema = getter.GET_SCHEMA()
|
||||
assert getter_schema.display_name == 'Distributed Image List Get Value'
|
||||
assert not getter_schema.is_input_list
|
||||
assert getter_schema.outputs[0].id == getter_schema.outputs[0].display_name == 'value'
|
||||
assert getter_schema.outputs[0].io_type == '*'
|
||||
assert getter_schema.inputs[1].options == list(runtime.VALUE_NAMES)
|
||||
assert getter_schema.inputs[1].default == 'value_1'
|
||||
# Real core maps scalar mappings, never arbitrary tensor batches.
|
||||
tensor = torch.ones(3, 2, 2, 3)
|
||||
data, missing, hidden = execution.get_input_data({'value_name': 'value_1'}, getter, 'get-value')
|
||||
data['values'] = [{'value_1': tensor}, {'value_1': 17}, {'value_1': False}]
|
||||
mapped = await execution._async_map_node_over_list('get-value', 'get-value', getter, data,
|
||||
getter.FUNCTION, v3_data=hidden)
|
||||
mapped = await execution.resolve_map_node_over_list_results(mapped)
|
||||
assert len(mapped) == 3 and mapped[0].result[0] is tensor
|
||||
assert mapped[1].result == (17,) and mapped[2].result == (False,)
|
||||
try:
|
||||
getter.execute({}, 'value_1')
|
||||
raise AssertionError('missing Get Value key accepted')
|
||||
except ValueError as exc:
|
||||
assert 'absent' in str(exc)
|
||||
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', 'MASK', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + ['*'] * 10
|
||||
assert [out.display_name for out in start_schema.outputs][4:] == [f'value_{i}' for i in range(1, 11)]
|
||||
assert all(out.is_output_list for out in start_schema.outputs)
|
||||
assert start_schema.is_input_list and start_schema.accept_all_inputs
|
||||
start_info = json.loads(json.dumps(mapping['DistributedImageListStart'].GET_NODE_INFO_V1()))
|
||||
assert list(start_info['input']['optional']) == ['images', 'mask', 'values'], start_info['input']
|
||||
values_spec = start_info['input']['optional']['values']
|
||||
assert values_spec[0] == 'COMFY_AUTOGROW_V3', values_spec
|
||||
template = values_spec[1]['template']
|
||||
assert template['names'] == [f'value_{i}' for i in range(1, 11)] and template['min'] == 0, template
|
||||
assert template['input']['required']['value'][0] == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN', template
|
||||
# Autogrow + INPUT_IS_LIST through the real parser/list mapper, values
|
||||
# still split 1-or-N, unconnected slots ignored.
|
||||
# get_input_data must accept Autogrow names; constants are wrapped once.
|
||||
parsed, missing, hidden = execution.get_input_data(
|
||||
{'values.value_1': 'shared', 'values.value_3': 4},
|
||||
mapping['DistributedImageListStart'], unique_id='autogrow-probe', extra_data={})
|
||||
assert not missing and parsed['values.value_1'] == ['shared'], parsed
|
||||
# Linked inputs arrive as the upstream output list (one element per value).
|
||||
prepared = {'images': [torch.zeros((3, 2, 2, 3))], 'values.value_1': ['shared'],
|
||||
'values.value_3': [4, 5, 6]}
|
||||
seen = {}
|
||||
runtime_mod = sys.modules[module.__name__ + '.image_list.runtime']
|
||||
def capture(_images, _mask, values):
|
||||
seen['values'] = values
|
||||
raise RuntimeError('captured')
|
||||
nodes_mod = sys.modules[module.__name__ + '.nodes.image_list']
|
||||
with patch.object(nodes_mod, 'normalize_image_list_inputs', capture):
|
||||
try:
|
||||
await asyncio.to_thread(lambda: asyncio.run(execution._async_map_node_over_list(
|
||||
'autogrow-probe', 'autogrow-probe', mapping['DistributedImageListStart'], prepared,
|
||||
'execute', v3_data=hidden)))
|
||||
except RuntimeError as exc:
|
||||
assert 'captured' in str(exc), exc
|
||||
assert seen['values'] == {'value_1': ['shared'], 'value_3': [4, 5, 6]}, seen
|
||||
rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
|
||||
assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
|
||||
print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
|
||||
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())
|
||||
@@ -0,0 +1,146 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _load_async_helpers_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
|
||||
package_name = "dist_async_helpers_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
execution_module = types.ModuleType("execution")
|
||||
|
||||
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
|
||||
return (True, None, ["9"], {})
|
||||
|
||||
execution_module.validate_prompt = _validate_prompt
|
||||
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
|
||||
sys.modules["execution"] = execution_module
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
trigger_on_prompt=lambda payload: payload,
|
||||
number=12,
|
||||
prompt_queue=_PromptQueue(),
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_loop = lambda: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module, prompt_server
|
||||
|
||||
|
||||
async_helpers, prompt_server = _load_async_helpers_module()
|
||||
|
||||
|
||||
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
|
||||
result = await async_helpers.queue_prompt_payload(
|
||||
{"1": {"class_type": "Node"}},
|
||||
workflow_meta={"id": "workflow-1"},
|
||||
client_id="client-1",
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(result["prompt_id"], str)
|
||||
self.assertTrue(result["prompt_id"])
|
||||
self.assertEqual(result["number"], 12)
|
||||
self.assertEqual(result["node_errors"], {})
|
||||
|
||||
self.assertEqual(prompt_server.number, 13)
|
||||
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
|
||||
queued_item = prompt_server.prompt_queue.items[0]
|
||||
self.assertEqual(queued_item[0], 12)
|
||||
extra_data = queued_item[3]
|
||||
self.assertEqual(extra_data["client_id"], "client-1")
|
||||
self.assertIn("create_time", extra_data)
|
||||
self.assertIsInstance(extra_data["create_time"], int)
|
||||
self.assertGreater(extra_data["create_time"], 0)
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
class ExplicitTargetTests(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.module, self.server = _load_async_helpers_module()
|
||||
self.prompt = {"9": {"class_type": "Output"}, "10": {"class_type": "Output"}}
|
||||
self.validate = AsyncMock(return_value=(True, None, ["9"], {}))
|
||||
self.module.execution.validate_prompt = self.validate
|
||||
|
||||
async def test_targets_survive_callback_omission_and_are_queued_exactly(self):
|
||||
seen = []
|
||||
def callback(payload):
|
||||
seen.append(payload["partial_execution_targets"][:])
|
||||
payload["partial_execution_targets"].append("10")
|
||||
return {"prompt": payload["prompt"]}
|
||||
self.server.trigger_on_prompt = callback
|
||||
targets = ["9"]
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.assertEqual(seen, [["9"]])
|
||||
self.assertEqual(targets, ["9"])
|
||||
self.assertEqual(self.validate.call_args.args[2], ["9"])
|
||||
self.assertEqual(self.server.prompt_queue.items[0][4], ["9"])
|
||||
|
||||
async def test_invalid_requests_never_reach_validation_or_queue(self):
|
||||
for targets in ([], "9", [9], [True], ["9", "9"], ["missing"], {}):
|
||||
with self.subTest(targets=targets):
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=targets)
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_removed_target_after_callback_is_rejected(self):
|
||||
self.server.trigger_on_prompt = lambda payload: {"prompt": {"10": self.prompt["10"]}}
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9"])
|
||||
self.validate.assert_not_awaited()
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_partial_subset_extra_roots_and_target_errors_fail_closed(self):
|
||||
for roots, errors in ((["9"], {}), (["9", "10", "extra"], {}),
|
||||
(["9", "10"], {"9": {"errors": [{"message": "bad"}]}}),
|
||||
(["9", "10"], {"dep": {"dependent_outputs": ["10"], "errors": []}})):
|
||||
with self.subTest(roots=roots, errors=errors):
|
||||
self.validate.return_value = (True, None, roots, errors)
|
||||
with self.assertRaises(self.module.PromptValidationError):
|
||||
await self.module.queue_prompt_payload(self.prompt, partial_execution_targets=["9", "10"])
|
||||
self.assertEqual(self.server.prompt_queue.items, [])
|
||||
|
||||
async def test_none_preserves_legacy_payload_and_validation(self):
|
||||
seen = []
|
||||
self.server.trigger_on_prompt = lambda payload: seen.append(payload.copy()) or payload
|
||||
await self.module.queue_prompt_payload(self.prompt)
|
||||
self.assertNotIn("partial_execution_targets", seen[0])
|
||||
self.assertIsNone(self.validate.call_args.args[2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -8,7 +8,7 @@ import torch
|
||||
|
||||
|
||||
def _load_utilities_module():
|
||||
module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
package_name = "dist_divider_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
@@ -0,0 +1,524 @@
|
||||
"""Exact Collector transport through real serializers, HTTP, route and queue.
|
||||
|
||||
Only ComfyUI's process-wide server/model-management and unrelated orchestration
|
||||
are replaced. Tensor/audio codecs, multipart bodies and aiohttp are real; no
|
||||
running ComfyUI service, GPU, provider or external network access is required.
|
||||
"""
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
import importlib
|
||||
import importlib.util
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_collector_exact_testpkg"
|
||||
JOB_ID = "exact-job"
|
||||
WORKER_ID = "worker-1"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def transport(monkeypatch):
|
||||
# Other legacy test loaders may leave aiohttp stubs behind. Replace only
|
||||
# those entries, never unload/reload real extension-backed dependencies.
|
||||
for name, loaded_module in list(sys.modules.items()):
|
||||
if (name == "aiohttp" or name.startswith("aiohttp.")) and not getattr(loaded_module, "__file__", None):
|
||||
monkeypatch.delitem(sys.modules, name)
|
||||
aiohttp = importlib.import_module("aiohttp")
|
||||
web = importlib.import_module("aiohttp.web")
|
||||
|
||||
def module(name, **attrs):
|
||||
result = types.ModuleType(name)
|
||||
result.__dict__.update(attrs)
|
||||
monkeypatch.setitem(sys.modules, name, result)
|
||||
return result
|
||||
|
||||
def load(relative):
|
||||
name = f"{PKG}.{relative}"
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT.joinpath(*relative.split(".")).with_suffix(".py"))
|
||||
assert spec is not None and spec.loader is not None
|
||||
result = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, name, result)
|
||||
spec.loader.exec_module(result)
|
||||
parent, _, child = name.rpartition(".")
|
||||
setattr(sys.modules[parent], child, result)
|
||||
return result
|
||||
|
||||
module(PKG, __path__=[str(ROOT)])
|
||||
for package in ("nodes", "api", "utils", "serverless_bridge"):
|
||||
module(f"{PKG}.{package}", __path__=[str(ROOT / package)])
|
||||
server = types.SimpleNamespace(routes=web.RouteTableDef(), distributed_pending_jobs={})
|
||||
module("server", PromptServer=types.SimpleNamespace(instance=server))
|
||||
mm = module("comfy.model_management", throw_exception_if_processing_interrupted=lambda: None)
|
||||
progress = module("comfy.utils", ProgressBar=lambda _total: types.SimpleNamespace(update=lambda _n: None))
|
||||
module("comfy", __path__=[], model_management=mm, utils=progress)
|
||||
module(f"{PKG}.utils.logging", debug_log=lambda *_a, **_kw: None, log=lambda *_a, **_kw: None)
|
||||
module(f"{PKG}.utils.config", get_worker_timeout_seconds=lambda: 0.1,
|
||||
load_config=lambda: {"workers": []}, is_master_delegate_only=lambda: False)
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
return web.json_response({"error": str(error)}, status=status)
|
||||
|
||||
module(f"{PKG}.utils.network", handle_api_error=handle_api_error,
|
||||
get_client_session=AsyncMock(), build_worker_url=lambda _worker: "http://unused.invalid",
|
||||
probe_worker=AsyncMock(return_value=None))
|
||||
module(f"{PKG}.utils.async_helpers", run_async_in_server_loop=asyncio.run)
|
||||
module(f"{PKG}.api.queue_orchestration", ensure_distributed_state=lambda: None,
|
||||
orchestrate_distributed_execution=AsyncMock())
|
||||
module(f"{PKG}.api.queue_request", parse_queue_request_payload=lambda value: value)
|
||||
for name in ("utils.constants", "utils.image", "utils.audio_payload", "utils.transfer_codec"):
|
||||
load(name)
|
||||
bridge_codec = load("serverless_bridge.codec")
|
||||
load("serverless_bridge.image_transfer")
|
||||
image_codec = load("serverless_bridge.image_codec")
|
||||
collector = load("nodes.collector")
|
||||
routes = load("api.job_routes")
|
||||
return types.SimpleNamespace(aiohttp=aiohttp, web=web, server=server,
|
||||
collector=collector, routes=routes, bridge_codec=bridge_codec,
|
||||
image_codec=image_codec)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def endpoint(transport, monkeypatch, *, request_limit=2 * 1024 * 1024,
|
||||
chunk_target=1024 * 1024, client_max_size=None, handler=None):
|
||||
"""A disposable loopback endpoint; record actual HTTP lengths, not estimates."""
|
||||
queue = asyncio.Queue()
|
||||
transport.server.distributed_jobs_lock = asyncio.Lock()
|
||||
transport.server.distributed_pending_jobs = {JOB_ID: queue}
|
||||
requests = []
|
||||
web = transport.web
|
||||
|
||||
@web.middleware
|
||||
async def record_request(request, handler):
|
||||
record = {"bytes": request.content_length, "content_type": request.content_type}
|
||||
requests.append(record)
|
||||
response = await handler(request)
|
||||
record["status"] = response.status
|
||||
if response.status < 400 and request.content_type.startswith("multipart/"):
|
||||
form = await request.post()
|
||||
record["envelope"] = json.loads(form["envelope"])
|
||||
return response
|
||||
|
||||
app = web.Application(client_max_size=client_max_size or request_limit, middlewares=[record_request])
|
||||
app.router.add_post("/distributed/job_complete", handler or transport.routes.job_complete_endpoint)
|
||||
runner = web.AppRunner(app)
|
||||
await runner.setup()
|
||||
try:
|
||||
site = web.TCPSite(runner, "127.0.0.1", 0)
|
||||
await site.start()
|
||||
port = site._server.sockets[0].getsockname()[1]
|
||||
async with transport.aiohttp.ClientSession() as session:
|
||||
monkeypatch.setattr(transport.collector, "get_client_session", AsyncMock(return_value=session))
|
||||
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_REQUEST_BYTES", request_limit)
|
||||
monkeypatch.setattr(transport.collector, "TRANSFER_MAX_CHUNK_BYTES", chunk_target)
|
||||
yield types.SimpleNamespace(
|
||||
session=session, url=f"http://127.0.0.1:{port}", requests=requests, queue=queue,
|
||||
node=transport.collector.DistributedCollectorNode(),
|
||||
)
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
|
||||
|
||||
def frames(dtype=torch.float32, channels=3, count=5, height=7, width=9):
|
||||
# Non-contiguous, out-of-range and non-8-bit-representable values exercise
|
||||
# detach/contiguous without permitting clipping, alpha loss or quantization.
|
||||
values = torch.arange(count * height * width * channels, dtype=torch.float64)
|
||||
values = (values.remainder(41) - 11) / 17
|
||||
images = values.reshape(count, width, height, channels).transpose(1, 2).to(dtype)
|
||||
images[0, 0, 0, 0] = -0.0
|
||||
return images.requires_grad_(True)
|
||||
|
||||
|
||||
def assert_exact(actual, expected):
|
||||
assert actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert actual.device.type == "cpu"
|
||||
assert actual.is_contiguous()
|
||||
# Compare raw bits as well as values, including BF16 and signed zero.
|
||||
assert torch.equal(actual, expected)
|
||||
assert torch.equal(actual.view(torch.uint8), expected.detach().contiguous().view(torch.uint8))
|
||||
|
||||
|
||||
def queued_items(queue):
|
||||
return [queue.get_nowait() for _ in range(queue.qsize())]
|
||||
|
||||
|
||||
async def send(ctx, images, *, audio=None, fmt="exact"):
|
||||
await ctx.node.send_batch_to_master(images, audio, JOB_ID, ctx.url, WORKER_ID, fmt, 17)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate_only", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
|
||||
@pytest.mark.parametrize("channels", [1, 3, 4])
|
||||
def test_exact_sender_route_and_collector_preserve_dtype_values_and_channels(
|
||||
transport, monkeypatch, dtype, channels, delegate_only):
|
||||
images = frames(dtype, channels)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + 2 * frame_bytes,
|
||||
chunk_target=2 * frame_bytes) as ctx:
|
||||
await send(ctx, images)
|
||||
items = queued_items(ctx.queue)
|
||||
assert [item["image_index"] for item in items] == list(range(len(images)))
|
||||
assert [item["is_last"] for item in items] == [False] * 4 + [True]
|
||||
assert all(item["worker_id"] == WORKER_ID and item["format"] == "exact" for item in items)
|
||||
for item, expected in zip(items, images.split(1)):
|
||||
assert_exact(item["tensor"], expected)
|
||||
await ctx.queue.put(item)
|
||||
envelopes = [record["envelope"] for record in ctx.requests]
|
||||
assert [(e["start_index"], e["count"]) for e in envelopes] == [(0, 2), (2, 2), (4, 1)]
|
||||
assert all(record["bytes"] <= reserve + 2 * frame_bytes for record in ctx.requests)
|
||||
master_images = None if delegate_only else images[:1].detach().clone().add(0.03125)
|
||||
actual, _audio = await ctx.node.execute(
|
||||
master_images, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
|
||||
delegate_only=delegate_only, transfer_format="exact",
|
||||
)
|
||||
expected = images if delegate_only else torch.cat([master_images, images])
|
||||
assert_exact(actual, expected)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_single_frame_can_exceed_chunk_target_but_not_request_limit(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes,
|
||||
chunk_target=frame_bytes // 2) as ctx:
|
||||
await send(ctx, images)
|
||||
assert [record["envelope"]["count"] for record in ctx.requests] == [1, 1, 1]
|
||||
assert all(record["bytes"] <= reserve + frame_bytes for record in ctx.requests)
|
||||
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_frame_is_rejected_before_encoding_or_post(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
encode = AsyncMock(side_effect=AssertionError("oversized frame must not be encoded"))
|
||||
monkeypatch.setattr(transport.collector, "_encode_exact_frames", encode)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + frame_bytes - 1) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
|
||||
await send(ctx, images)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
encode.assert_not_called()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(0, 2, 3, 3), (1, 0, 3, 3), (1, 2, 0, 3), (1, 2, 3, 0),
|
||||
(1, 2, 3, 2), (1, 2, 3, 5), (1, 2, 3), (2, 3)])
|
||||
def test_exact_sender_rejects_invalid_geometry_before_post(transport, monkeypatch, shape):
|
||||
image = torch.empty(shape)
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError):
|
||||
transport.collector._exact_frame_bytes(image)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
# An empty batch is also a no-media completion error (ValueError).
|
||||
with pytest.raises(ValueError):
|
||||
await send(ctx, image)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.uint8, torch.int32, torch.bool, torch.complex64])
|
||||
def test_exact_sender_rejects_unsupported_dtype_before_post(transport, monkeypatch, dtype):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dtype"):
|
||||
await send(ctx, torch.zeros(1, 2, 3, 3, dtype=dtype))
|
||||
assert ctx.requests == []
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_frame_planning_enforces_bridge_byte_limit(transport, monkeypatch):
|
||||
monkeypatch.setattr(transport.bridge_codec, "MAX_PAYLOAD_BYTES", 32)
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="[Bb]yte|limit"):
|
||||
transport.collector._exact_frame_bytes(torch.zeros(1, 2, 3, 3))
|
||||
|
||||
|
||||
def test_exact_frame_planning_rejects_non_dense_layout(transport):
|
||||
image = torch.zeros(1, 2, 3, 3).to_sparse()
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="dense|layout|strided"):
|
||||
transport.collector._exact_frame_bytes(image)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")])
|
||||
def test_exact_sender_rejects_nonfinite_payload(transport, monkeypatch, value):
|
||||
images = torch.full((1, 2, 3, 3), value)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.TransferCodecError, match="finite"):
|
||||
await send(ctx, images)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def exact_chunk(transport, count=2):
|
||||
images = frames(count=count).detach()
|
||||
encoded = [transport.bridge_codec.encode_payload("IMAGE", frame) for frame in images.split(1)]
|
||||
return {
|
||||
"job_id": JOB_ID, "worker_id": WORKER_ID, "format": "exact",
|
||||
"start_index": 0, "count": count, "width": images.shape[2], "height": images.shape[1],
|
||||
"is_last": True, "descriptors": [descriptor for descriptor, _ in encoded],
|
||||
}, [payload for _, payload in encoded]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", [
|
||||
"zero_count", "bool_width", "negative_height", "negative_index", "bool_index",
|
||||
"missing_descriptor", "extra_descriptor", "missing_part", "extra_part", "not_a_descriptor",
|
||||
"wrong_height", "wrong_batch", "bad_channels", "bad_dtype", "bad_layout", "bad_encoding",
|
||||
"extra_descriptor_field", "truncated", "trailing_bytes", "nonfinite", "compressed_spoof",
|
||||
])
|
||||
def test_exact_route_rejects_bad_geometry_and_payload_atomically(transport, monkeypatch, case):
|
||||
envelope, parts = exact_chunk(transport)
|
||||
descriptor = envelope["descriptors"][-1]
|
||||
if case == "zero_count":
|
||||
envelope["count"] = 0
|
||||
elif case == "bool_width":
|
||||
envelope["width"] = True
|
||||
elif case == "negative_height":
|
||||
envelope["height"] = -1
|
||||
elif case == "negative_index":
|
||||
envelope["start_index"] = -1
|
||||
elif case == "bool_index":
|
||||
envelope["start_index"] = True
|
||||
elif case == "missing_descriptor":
|
||||
envelope["descriptors"].pop()
|
||||
elif case == "extra_descriptor":
|
||||
envelope["descriptors"].append(dict(descriptor))
|
||||
elif case == "missing_part":
|
||||
parts.pop()
|
||||
elif case == "extra_part":
|
||||
parts.append(parts[-1])
|
||||
elif case == "not_a_descriptor":
|
||||
envelope["descriptors"][-1] = None
|
||||
elif case == "wrong_height":
|
||||
descriptor["shape"][1] += 1
|
||||
elif case == "wrong_batch":
|
||||
descriptor["shape"][0] = 2
|
||||
elif case == "bad_channels":
|
||||
descriptor["shape"][-1] = 2
|
||||
elif case == "bad_dtype":
|
||||
descriptor["dtype"] = "uint8"
|
||||
elif case == "bad_layout":
|
||||
descriptor["layout"] = "BCHW"
|
||||
elif case == "bad_encoding":
|
||||
descriptor["encoding"] = "pickle"
|
||||
elif case == "extra_descriptor_field":
|
||||
descriptor["unknown"] = True
|
||||
elif case == "truncated":
|
||||
parts[-1] = parts[-1][:-1]
|
||||
elif case == "trailing_bytes":
|
||||
parts[-1] += b"\0"
|
||||
elif case == "nonfinite":
|
||||
bad = torch.full(tuple(descriptor["shape"]), float("nan"))
|
||||
parts[-1] = bad.numpy().tobytes()
|
||||
elif case == "compressed_spoof":
|
||||
# A valid serverless image-transfer descriptor must NOT be accepted as
|
||||
# exact by the direct route, even though the shared bridge can decode it.
|
||||
envelope["descriptors"][-1], parts[-1] = transport.image_codec.encode_image(
|
||||
frames(count=1).detach(), transfer_format="png", transfer_quality=90,
|
||||
)
|
||||
# Prove the spoof is a valid compressed bridge payload, not simply junk.
|
||||
decoded = transport.bridge_codec.decode_payload("IMAGE", envelope["descriptors"][-1], parts[-1])
|
||||
assert list(decoded.shape) == descriptor["shape"]
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
async with ctx.session.post(
|
||||
ctx.url + "/distributed/job_complete",
|
||||
data=ctx.node._build_multipart(envelope, parts),
|
||||
) as response:
|
||||
body = await response.text()
|
||||
assert response.status == 400, body
|
||||
# A corrupt second frame must not leave the first in the queue.
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status,body", [(400, "Invalid JSON payload: old master"),
|
||||
(413, "upload limit"), (500, "unexpected failure")])
|
||||
def test_exact_master_rejection_never_falls_back(transport, monkeypatch, status, body):
|
||||
async def reject(_request):
|
||||
return transport.web.Response(status=status, text=body)
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, handler=reject, chunk_target=1) as ctx:
|
||||
legacy = AsyncMock(side_effect=AssertionError("exact must never use legacy transport"))
|
||||
compressed = AsyncMock(side_effect=AssertionError("exact must never use compressed transport"))
|
||||
monkeypatch.setattr(ctx.node, "_send_legacy_batch", legacy)
|
||||
monkeypatch.setattr(ctx.node, "_send_compressed_batch", compressed)
|
||||
with pytest.raises(RuntimeError, match=f"HTTP {status}") as raised:
|
||||
await send(ctx, frames(count=3))
|
||||
assert "never" in str(raised.value) and "8-bit" in str(raised.value)
|
||||
assert len(ctx.requests) == 1
|
||||
assert ctx.requests[0]["content_type"].startswith("multipart/")
|
||||
assert ctx.queue.empty()
|
||||
legacy.assert_not_called()
|
||||
compressed.assert_not_called()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fmt", ["legacy_png", "png"])
|
||||
def test_exact_collector_rejects_old_worker_8bit_result(transport, monkeypatch, fmt):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, frames(count=1).detach(), fmt=fmt)
|
||||
assert ctx.queue.qsize() == 1
|
||||
with pytest.raises(RuntimeError, match="requested exact"):
|
||||
await ctx.node.execute(
|
||||
None, None, multi_job_id=JOB_ID, enabled_worker_ids=json.dumps([WORKER_ID]),
|
||||
delegate_only=True, transfer_format="exact",
|
||||
)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_collector_shape_mismatch_raises_instead_of_returning_local_images(transport, monkeypatch):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
await send(ctx, frames(count=1, channels=4))
|
||||
with pytest.raises(RuntimeError, match="shape|combin|mismatch"):
|
||||
await ctx.node.execute(
|
||||
frames(count=1, channels=3), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_incomplete_worker_raises_and_cleans_queue(transport, monkeypatch):
|
||||
"""A timed-out worker's first frame is not a successful partial result."""
|
||||
envelope, parts = exact_chunk(transport, count=1)
|
||||
envelope["is_last"] = False
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch) as ctx:
|
||||
async with ctx.session.post(
|
||||
ctx.url + "/distributed/job_complete", data=ctx.node._build_multipart(envelope, parts),
|
||||
) as response:
|
||||
assert response.status == 200, await response.text()
|
||||
assert ctx.queue.qsize() == 1
|
||||
with pytest.raises(RuntimeError, match="Exact transfer incomplete"):
|
||||
await asyncio.wait_for(ctx.node.execute(
|
||||
frames(count=1).detach(), None, multi_job_id=JOB_ID,
|
||||
enabled_worker_ids=json.dumps([WORKER_ID]), transfer_format="exact",
|
||||
), timeout=2)
|
||||
assert JOB_ID not in transport.server.distributed_pending_jobs
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_multipart_metadata_is_rejected_before_post(transport, monkeypatch):
|
||||
images = frames(count=1)
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
limit = reserve + images.numel() * images.element_size()
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="request limit"):
|
||||
await ctx.node.send_batch_to_master(
|
||||
images, None, JOB_ID, ctx.url, "w" * limit, "exact", 90,
|
||||
)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_actual_server_upload_rejection_is_413_without_retry(transport, monkeypatch):
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, client_max_size=1024) as ctx:
|
||||
with pytest.raises(RuntimeError, match="HTTP 413"):
|
||||
await send(ctx, frames(count=3, height=64, width=64))
|
||||
assert len(ctx.requests) == 1
|
||||
assert ctx.requests[0]["status"] == 413
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_oversized_audio_is_rejected_before_any_image_post(transport, monkeypatch):
|
||||
images = frames(count=3)
|
||||
audio = {"waveform": torch.zeros(1, 1, 100000), "sample_rate": 48000}
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=reserve + 1024) as ctx:
|
||||
with pytest.raises(transport.collector.transfer_codec.FrameTooLargeError, match="audio.*request limit"):
|
||||
await send(ctx, images, audio=audio)
|
||||
assert ctx.requests == []
|
||||
assert ctx.queue.empty()
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
def test_exact_tiny_frames_cannot_overflow_multipart_metadata_budget(transport, monkeypatch):
|
||||
images = torch.full((2048, 1, 1, 1), 0.125, dtype=torch.float16)
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
limit = reserve + images.numel() * images.element_size()
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=limit) as ctx:
|
||||
await send(ctx, images)
|
||||
assert len(ctx.requests) > 1
|
||||
assert all(record["bytes"] <= limit for record in ctx.requests)
|
||||
assert_exact(torch.cat([item["tensor"] for item in queued_items(ctx.queue)]), images)
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("separate", [False, True])
|
||||
def test_exact_audio_is_preserved_and_only_final_completion_is_terminal(transport, monkeypatch, separate):
|
||||
images = frames(count=3)
|
||||
audio = {"waveform": torch.linspace(-0.9, 0.9, 60).reshape(1, 1, -1), "sample_rate": 48000}
|
||||
frame_bytes = images[0].numel() * images.element_size()
|
||||
reserve = transport.collector.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
media_budget = frame_bytes if separate else frame_bytes + 8192
|
||||
limit = reserve + media_budget
|
||||
|
||||
async def scenario():
|
||||
async with endpoint(transport, monkeypatch, request_limit=limit, chunk_target=frame_bytes) as ctx:
|
||||
await send(ctx, images, audio=audio)
|
||||
items = queued_items(ctx.queue)
|
||||
assert len(ctx.requests) == 3 + int(separate)
|
||||
assert [item["is_last"] for item in items] == [False] * (len(items) - 1) + [True]
|
||||
assert all(item["audio"] is None for item in items[:-1])
|
||||
assert items[-1]["audio"]["sample_rate"] == audio["sample_rate"]
|
||||
assert_exact(items[-1]["audio"]["waveform"], audio["waveform"])
|
||||
assert_exact(torch.cat([item["tensor"] for item in items if item["tensor"] is not None]), images)
|
||||
assert all(record["bytes"] <= limit for record in ctx.requests)
|
||||
if separate:
|
||||
assert items[-1]["tensor"] is None
|
||||
assert ctx.requests[-1]["content_type"] == "application/json"
|
||||
|
||||
asyncio.run(scenario())
|
||||
@@ -0,0 +1,638 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
|
||||
def _register_real_transfer_codec(package_name, utils_pkg):
|
||||
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
||||
codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
|
||||
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
||||
codec_module = importlib.util.module_from_spec(codec_spec)
|
||||
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
||||
codec_spec.loader.exec_module(codec_module)
|
||||
utils_pkg.transfer_codec = codec_module
|
||||
return codec_module
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
# Restore shared dependencies even when a different test imports this loader.
|
||||
# Keep real modules imported during loading (notably extension modules) cached.
|
||||
with pytest.MonkeyPatch.context() as monkeypatch:
|
||||
return _load_collector_module_with_stubs(monkeypatch)
|
||||
|
||||
|
||||
def _load_collector_module_with_stubs(monkeypatch):
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
_register_real_transfer_codec(package_name, utils_pkg)
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
def get(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
monkeypatch.setitem(sys.modules, "server", server_module)
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
|
||||
class InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_management.InterruptProcessingException = InterruptProcessingException
|
||||
model_management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy_module.model_management = model_management
|
||||
|
||||
comfy_utils = types.ModuleType("comfy.utils")
|
||||
|
||||
class ProgressBar:
|
||||
def __init__(self, _total):
|
||||
self.total = _total
|
||||
self.updates = []
|
||||
|
||||
def update(self, value):
|
||||
self.updates.append(value)
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
monkeypatch.setitem(sys.modules, "comfy", comfy_module)
|
||||
monkeypatch.setitem(sys.modules, "comfy.model_management", model_management)
|
||||
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils)
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
monkeypatch.setitem(sys.modules, "aiohttp", aiohttp_module)
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 0.1
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
config_module.is_master_delegate_only = lambda: False
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 1.0
|
||||
constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
|
||||
constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
def _ensure_contiguous(tensor):
|
||||
return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor
|
||||
|
||||
image_module.ensure_contiguous = _ensure_contiguous
|
||||
image_module.tensor_to_pil = lambda *_args, **_kwargs: None
|
||||
image_module.pil_to_tensor = lambda value: value
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: "http://worker"
|
||||
network_module.get_client_session = lambda: None
|
||||
network_module.probe_worker = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
||||
audio_payload_module.encode_audio_payload = lambda audio: audio
|
||||
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_collector_opts_into_comfyui_list_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode
|
||||
|
||||
assert collector.INPUT_IS_LIST is True
|
||||
|
||||
|
||||
def test_collector_exposes_images_as_optional_input():
|
||||
input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
|
||||
|
||||
assert "images" not in input_types["required"]
|
||||
assert input_types["optional"]["images"] == ("IMAGE",)
|
||||
|
||||
|
||||
def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
|
||||
images, returned_audio = collector.run(images=None, audio=[audio])
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
|
||||
|
||||
def test_collector_rejects_missing_images_and_audio():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
try:
|
||||
collector.run(images=None, audio=None)
|
||||
except ValueError as exc:
|
||||
assert "image or audio" in str(exc).lower()
|
||||
else:
|
||||
raise AssertionError("Expected collector to reject a run with no media input")
|
||||
|
||||
|
||||
def test_delegate_only_master_allows_no_local_media_input():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
|
||||
images, audio = collector.run(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=["delegate-audio-job"],
|
||||
delegate_only=[True],
|
||||
enabled_worker_ids=["[]"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
|
||||
collector = _load_collector_module().DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
load_balance=[False],
|
||||
audio=[None],
|
||||
multi_job_id=[""],
|
||||
is_worker=[False],
|
||||
master_url=[""],
|
||||
enabled_worker_ids=["[]"],
|
||||
worker_batch_size=[1],
|
||||
worker_id=[""],
|
||||
pass_through=[False],
|
||||
delegate_only=[False],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert torch.equal(images[0:1], first)
|
||||
assert torch.equal(images[1:2], second)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
|
||||
|
||||
def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
first = torch.zeros(1, 2, 2, 3)
|
||||
second = torch.ones(1, 2, 2, 3)
|
||||
posted_payloads = []
|
||||
|
||||
class _FakeImage:
|
||||
def save(self, fp, format=None, compress_level=None):
|
||||
fp.write(b"png-bytes")
|
||||
|
||||
class _FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, timeout):
|
||||
posted_payloads.append(json)
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda _audio: None
|
||||
|
||||
images, audio = collector.run(
|
||||
images=[first, second],
|
||||
load_balance=[False],
|
||||
audio=[None],
|
||||
multi_job_id=["job-list-1"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
enabled_worker_ids=["[]"],
|
||||
worker_batch_size=[1],
|
||||
worker_id=["worker-a"],
|
||||
pass_through=[False],
|
||||
delegate_only=[False],
|
||||
)
|
||||
|
||||
assert tuple(images.shape) == (2, 2, 2, 3)
|
||||
assert tuple(audio["waveform"].shape) == (1, 2, 1)
|
||||
assert len(posted_payloads) == 2
|
||||
assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
|
||||
assert [payload["is_last"] for payload in posted_payloads] == [False, True]
|
||||
assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
|
||||
assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
|
||||
|
||||
|
||||
def test_audio_only_worker_sends_one_completion_without_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
|
||||
posted = []
|
||||
|
||||
class _FakeResponse:
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
class _FakeSession:
|
||||
def post(self, url, json, timeout):
|
||||
posted.append((url, json, timeout.total))
|
||||
return _FakeResponse()
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return _FakeSession()
|
||||
|
||||
module.get_client_session = _fake_get_client_session
|
||||
module.encode_audio_payload = lambda value: {"encoded": value is audio}
|
||||
|
||||
images, returned_audio = collector.run(
|
||||
images=None,
|
||||
audio=[audio],
|
||||
multi_job_id=["audio-job"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
worker_id=["worker-a"],
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert returned_audio is audio
|
||||
assert len(posted) == 1
|
||||
assert posted[0][0] == "http://master/distributed/job_complete"
|
||||
assert posted[0][2] == 600
|
||||
assert posted[0][1] == {
|
||||
"job_id": "audio-job",
|
||||
"worker_id": "worker-a",
|
||||
"batch_idx": 0,
|
||||
"audio": {"encoded": True},
|
||||
"is_last": True,
|
||||
}
|
||||
|
||||
|
||||
def test_audio_only_master_combines_local_and_worker_audio():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=master_audio,
|
||||
multi_job_id="audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
|
||||
assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
|
||||
assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
|
||||
|
||||
|
||||
def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait(
|
||||
{
|
||||
"worker_id": "worker-a",
|
||||
"image_index": 0,
|
||||
"tensor": None,
|
||||
"audio": worker_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
|
||||
|
||||
images, combined_audio = asyncio.run(
|
||||
collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id="delegate-audio-job",
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
delegate_only=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert images is None
|
||||
assert combined_audio["sample_rate"] == 48000
|
||||
assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Compressed worker -> master transfer
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
class _RecordingResponse:
|
||||
def __init__(self, status=200, text=""):
|
||||
self.status = status
|
||||
self._text = text
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def text(self):
|
||||
return self._text
|
||||
|
||||
def raise_for_status(self):
|
||||
if self.status >= 400:
|
||||
raise RuntimeError(f"HTTP {self.status}")
|
||||
|
||||
|
||||
class _RecordingSession:
|
||||
def __init__(self, statuses=None, text="Invalid JSON payload"):
|
||||
self.calls = []
|
||||
self._statuses = list(statuses or [])
|
||||
self._text = text
|
||||
|
||||
def post(self, url, timeout, json=None, data=None):
|
||||
self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
|
||||
status = self._statuses.pop(0) if self._statuses else 200
|
||||
return _RecordingResponse(status=status, text=self._text)
|
||||
|
||||
|
||||
def _install_fake_transport(module, session):
|
||||
async def _fake_get_client_session():
|
||||
return session
|
||||
|
||||
module.get_client_session = _fake_get_client_session
|
||||
# Record the multipart content instead of building a real aiohttp.FormData.
|
||||
module.DistributedCollectorNode._build_multipart = staticmethod(
|
||||
lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
|
||||
)
|
||||
|
||||
|
||||
def _run_worker(collector, images, audio=None, **kwargs):
|
||||
return collector.run(
|
||||
images=[images],
|
||||
audio=[audio],
|
||||
multi_job_id=["job-c"],
|
||||
is_worker=[True],
|
||||
master_url=["http://master"],
|
||||
worker_id=["worker-a"],
|
||||
**{key: [value] for key, value in kwargs.items()},
|
||||
)
|
||||
|
||||
|
||||
def test_collector_exposes_transfer_widgets_with_legacy_default():
|
||||
module = _load_collector_module()
|
||||
optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
|
||||
|
||||
formats, options = optional["transfer_format"]
|
||||
assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
|
||||
assert options["default"] == "legacy_png"
|
||||
assert optional["transfer_quality"][0] == "INT"
|
||||
|
||||
|
||||
def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
|
||||
images = torch.rand(6, 32, 48, 3)
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
images,
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="ffv1",
|
||||
transfer_quality=90,
|
||||
)
|
||||
|
||||
assert len(session.calls) == 1
|
||||
call = session.calls[0]
|
||||
assert call["url"] == "http://master/distributed/job_complete"
|
||||
assert call["json"] is None
|
||||
envelope = call["data"]["envelope"]
|
||||
assert envelope["format"] == "ffv1"
|
||||
assert (envelope["start_index"], envelope["count"]) == (0, 6)
|
||||
assert (envelope["width"], envelope["height"]) == (48, 32)
|
||||
assert envelope["is_last"] is True
|
||||
assert envelope["audio"] == {"encoded": True}
|
||||
assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
|
||||
assert len(call["data"]["parts"]) == 1
|
||||
|
||||
|
||||
def test_worker_marks_only_final_chunk_as_last():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 1 # force one frame per request
|
||||
|
||||
_run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
|
||||
|
||||
envelopes = [call["data"]["envelope"] for call in session.calls]
|
||||
assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
|
||||
(0, 1, False),
|
||||
(1, 1, False),
|
||||
(2, 1, True),
|
||||
]
|
||||
assert "audio" not in envelopes[-1]
|
||||
|
||||
|
||||
def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession(statuses=[400])
|
||||
_install_fake_transport(module, session)
|
||||
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
||||
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
||||
)
|
||||
|
||||
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
|
||||
|
||||
assert session.calls[0]["data"] is not None # compressed attempt
|
||||
legacy = [call["json"] for call in session.calls[1:]]
|
||||
assert [payload["batch_idx"] for payload in legacy] == [0, 1]
|
||||
assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
|
||||
assert [payload["is_last"] for payload in legacy] == [False, True]
|
||||
|
||||
|
||||
def test_worker_uses_legacy_path_for_unknown_transfer_format():
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
|
||||
(255 * tensor[idx].cpu().numpy()).astype(np.uint8)
|
||||
)
|
||||
|
||||
_run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
|
||||
|
||||
assert len(session.calls) == 1
|
||||
assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
def test_worker_falls_back_when_video_encoder_missing():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
codec = module.transfer_codec
|
||||
original = codec.video_codec_available
|
||||
codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
|
||||
try:
|
||||
_run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
|
||||
finally:
|
||||
codec.video_codec_available = original
|
||||
|
||||
assert session.calls[0]["data"]["envelope"]["format"] == "h264"
|
||||
|
||||
|
||||
def test_worker_does_not_fall_back_on_non_legacy_rejection():
|
||||
import pytest
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
|
||||
_install_fake_transport(module, session)
|
||||
|
||||
with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
|
||||
_run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
|
||||
|
||||
# One compressed attempt, no legacy resend of the batch.
|
||||
assert len(session.calls) == 1
|
||||
assert session.calls[0]["json"] is None
|
||||
|
||||
|
||||
def test_worker_sends_nothing_when_a_frame_cannot_fit():
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
|
||||
# Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
|
||||
images = torch.zeros(3, 128, 128, 3)
|
||||
images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
|
||||
|
||||
with pytest.raises(module.transfer_codec.FrameTooLargeError):
|
||||
_run_worker(collector, images, transfer_format="png")
|
||||
|
||||
assert session.calls == []
|
||||
|
||||
|
||||
def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
|
||||
module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
|
||||
module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
torch.zeros(2, 16, 16, 3),
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="png",
|
||||
)
|
||||
|
||||
multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
|
||||
assert [e["is_last"] for e in multipart] == [False]
|
||||
assert all("audio" not in e for e in multipart)
|
||||
final = session.calls[-1]["json"]
|
||||
assert final["is_last"] is True and final["audio"]["data"].startswith("x")
|
||||
assert "image" not in final
|
||||
|
||||
|
||||
def test_worker_attaches_audio_to_last_chunk_when_it_fits():
|
||||
module = _load_collector_module()
|
||||
collector = module.DistributedCollectorNode()
|
||||
session = _RecordingSession()
|
||||
_install_fake_transport(module, session)
|
||||
module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
|
||||
|
||||
_run_worker(
|
||||
collector,
|
||||
torch.zeros(2, 16, 16, 3),
|
||||
audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
|
||||
transfer_format="png",
|
||||
)
|
||||
|
||||
assert len(session.calls) == 1
|
||||
envelope = session.calls[0]["data"]["envelope"]
|
||||
assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Real Collector aggregation of typed bridge deliveries and direct results.
|
||||
|
||||
Only Comfy's host services are stubbed by the shared Collector loader. Tensor
|
||||
normalization, participant ordering, queue consumption and cleanup are real.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import types
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def collector_module():
|
||||
module = _load_collector_module()
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_serverless = types.SimpleNamespace(cancel=AsyncMock())
|
||||
return module
|
||||
|
||||
|
||||
def native_media(value=0):
|
||||
image = torch.arange(24, dtype=torch.float64).reshape(1, 2, 3, 4) / 997 + value
|
||||
image[0, 0, 0, 0] = -0.0
|
||||
audio = {"waveform": torch.arange(6, dtype=torch.float64).reshape(1, 2, 3) / 997 + value,
|
||||
"sample_rate": 24000}
|
||||
return {"images": image, "audio": audio}
|
||||
|
||||
|
||||
def assert_exact(actual, expected):
|
||||
assert actual.dtype == expected.dtype and actual.shape == expected.shape
|
||||
assert torch.equal(actual.contiguous().view(torch.uint8), expected.contiguous().view(torch.uint8))
|
||||
|
||||
|
||||
def collect_queue(module, messages, participants, *, local=None, delegate=True):
|
||||
queue = asyncio.Queue()
|
||||
for message in messages:
|
||||
queue.put_nowait(message)
|
||||
module.prompt_server.distributed_pending_jobs["typed-job"] = queue
|
||||
values = local or {"images": None, "audio": None}
|
||||
return asyncio.run(asyncio.wait_for(module.DistributedCollectorNode().execute(
|
||||
**values, multi_job_id="typed-job", enabled_worker_ids=json.dumps(participants),
|
||||
delegate_only=delegate), timeout=2))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delegate", [False, True])
|
||||
@pytest.mark.parametrize("socket_set", [("images",), ("audio",), ("images", "audio")])
|
||||
def test_typed_and_direct_results_preserve_batch_and_audio_order(collector_module, delegate, socket_set):
|
||||
module = collector_module
|
||||
local, typed_a, typed_b, direct = [native_media(n) for n in range(4)]
|
||||
typed = {socket: [typed_a[socket], typed_b[socket]] for socket in socket_set}
|
||||
direct_message = {"worker_id": "direct", "is_last": True}
|
||||
if "images" in socket_set:
|
||||
direct_message.update(tensor=direct["images"], image_index=0)
|
||||
if "audio" in socket_set:
|
||||
direct_message["audio"] = direct["audio"]
|
||||
messages = [direct_message, {"worker_id": "serverless:pool:1", "media": typed, "is_last": True}]
|
||||
# Delivery order differs from participant order. Duplicate configured IDs
|
||||
# must not duplicate a participant's data or leave collection waiting.
|
||||
participants = ["serverless:pool:1", "direct", "serverless:pool:1"]
|
||||
local_inputs = {socket: local[socket] if socket in socket_set else None for socket in ("images", "audio")}
|
||||
images, audio = collect_queue(module, messages, participants, local=local_inputs, delegate=delegate)
|
||||
expected = ([] if delegate else [local]) + [typed_a, typed_b, direct]
|
||||
if "images" in socket_set:
|
||||
assert_exact(images, torch.cat([row["images"] for row in expected]))
|
||||
else:
|
||||
assert images is None
|
||||
if "audio" in socket_set:
|
||||
assert audio["sample_rate"] == 24000
|
||||
assert_exact(audio["waveform"], torch.cat([row["audio"]["waveform"] for row in expected], dim=-1))
|
||||
else:
|
||||
assert torch.equal(audio["waveform"], module.DistributedCollectorNode.EMPTY_AUDIO["waveform"])
|
||||
assert module.prompt_server.distributed_pending_jobs == {}
|
||||
module.prompt_server.distributed_serverless.cancel.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("media", [{}, {"texts": ["not a Collector socket"]}, {"images": []},
|
||||
{"audio": [None]}, {"images": "not a list"},
|
||||
{"images": [torch.zeros(2, 2, 3)]}])
|
||||
def test_invalid_typed_delivery_cancels_and_removes_queue(collector_module, media):
|
||||
module = collector_module
|
||||
with pytest.raises(RuntimeError, match="media|items|IMAGE"):
|
||||
collect_queue(module, [{"worker_id": "serverless:pool:1", "media": media, "is_last": True}],
|
||||
["serverless:pool:1"])
|
||||
assert module.prompt_server.distributed_pending_jobs == {}
|
||||
module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["nonfinal", "unexpected", "duplicate", "provider-error", "shape"])
|
||||
def test_typed_delivery_protocol_failures_are_not_partial_success(collector_module, failure):
|
||||
media = {"images": [native_media()["images"]]}
|
||||
first = {"worker_id": "serverless:pool:1", "media": media, "is_last": True}
|
||||
messages = [first]
|
||||
participants = ["serverless:pool:1", "serverless:pool:2"]
|
||||
if failure == "nonfinal":
|
||||
first["is_last"] = False
|
||||
elif failure == "unexpected":
|
||||
first["worker_id"] = "other"
|
||||
elif failure == "duplicate":
|
||||
messages.append(dict(first))
|
||||
elif failure == "provider-error":
|
||||
messages.append({"worker_id": "serverless:pool:2", "serverless_error": "offline failure", "is_last": True})
|
||||
else:
|
||||
messages.append({"worker_id": "serverless:pool:2", "media": {"images": [torch.zeros(1, 4, 5, 3)]},
|
||||
"is_last": True})
|
||||
with pytest.raises(RuntimeError, match="atomic|Unexpected|Duplicate|failed|shape mismatch"):
|
||||
collect_queue(collector_module, messages, participants)
|
||||
assert collector_module.prompt_server.distributed_pending_jobs == {}
|
||||
collector_module.prompt_server.distributed_serverless.cancel.assert_awaited_once_with(collector_job_id="typed-job")
|
||||
@@ -10,7 +10,7 @@ from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_config_module():
|
||||
module_path = Path(__file__).resolve().parents[3] / "utils" / "config.py"
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
|
||||
package_name = "dist_cfg_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
@@ -18,16 +18,16 @@ 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")
|
||||
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.GPU_CONFIG_FILE = "gpu_config.json"
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
|
||||
constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
|
||||
constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
|
||||
constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120
|
||||
sys.modules[f"{package_name}.constants"] = constants_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)
|
||||
@@ -213,7 +213,7 @@ class SaveConfigTests(unittest.TestCase):
|
||||
data = config.get_default_config()
|
||||
config.save_config(data)
|
||||
# Cache is now None; load_config should re-read
|
||||
self.assertIsNone(config._config_state().cache)
|
||||
self.assertIsNone(config._config_cache)
|
||||
|
||||
def test_written_file_is_valid_json(self):
|
||||
data = config.get_default_config()
|
||||
@@ -7,7 +7,7 @@ from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_detection_module():
|
||||
module_path = Path(__file__).resolve().parents[3] / "workers" / "detection.py"
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
|
||||
package_name = "dist_det_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
@@ -32,7 +32,6 @@ def _load_detection_module():
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.normalize_host = lambda value: value
|
||||
network_module.build_worker_url = lambda worker, endpoint="": f"http://{worker.get('host', 'localhost')}{endpoint}"
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in these tests")
|
||||
@@ -157,7 +156,7 @@ class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_0_0_0_0(self):
|
||||
result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188}) # nosec B104 - explicit wildcard host test case
|
||||
result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_when_type_is_local(self):
|
||||
@@ -0,0 +1,594 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
|
||||
|
||||
package_name = "dist_dispatch_testpkg"
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
|
||||
|
||||
async def _probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("get_client_session should be mocked in these tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _ClientError(Exception):
|
||||
pass
|
||||
|
||||
class _ClientConnectorError(_ClientError):
|
||||
pass
|
||||
|
||||
class _ClientConnectionError(_ClientError):
|
||||
pass
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _TCPConnector:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _ClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
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
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: {"payload": payload, "status": status}
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.dispatch",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
return module
|
||||
|
||||
|
||||
dispatch = _load_dispatch_module()
|
||||
|
||||
|
||||
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,
|
||||
"execution_targets": True,
|
||||
"protocol_version": 2,
|
||||
"node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
|
||||
}
|
||||
|
||||
|
||||
def _valid_queue(running=None, pending=None):
|
||||
return {"queue_running": running or [], "queue_pending": pending or []}
|
||||
|
||||
|
||||
class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_http_preserves_targets_and_omits_default_field(self):
|
||||
response = types.SimpleNamespace(raise_for_status=lambda: None)
|
||||
class Response:
|
||||
async def __aenter__(self): return response
|
||||
async def __aexit__(self, *args): pass
|
||||
calls = []
|
||||
def post(url, **kwargs):
|
||||
calls.append(kwargs['json'])
|
||||
return Response()
|
||||
session = types.SimpleNamespace(post=post)
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None)
|
||||
self.assertEqual(calls[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', calls[1])
|
||||
|
||||
async def test_websocket_sends_targets_all_the_way_to_wire(self):
|
||||
import json
|
||||
sent = []
|
||||
async def receive():
|
||||
return types.SimpleNamespace(type=dispatch.aiohttp.WSMsgType.TEXT, data=json.dumps({
|
||||
'type': 'dispatch_ack', 'request_id': sent[-1]['request_id'],
|
||||
'ok': True, 'prompt_id': 'worker-prompt'}))
|
||||
ws = types.SimpleNamespace(send_json=AsyncMock(side_effect=sent.append),
|
||||
receive=receive, close=AsyncMock())
|
||||
session = types.SimpleNamespace(ws_connect=AsyncMock(return_value=ws))
|
||||
with patch.object(dispatch, 'get_client_session', AsyncMock(return_value=session)):
|
||||
result = await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None,
|
||||
use_websocket=True, return_prompt_id=True, partial_execution_targets=['9'])
|
||||
await dispatch.dispatch_worker_prompt({'id': 'worker'}, {'9': {}}, None, use_websocket=True)
|
||||
self.assertEqual(result, 'worker-prompt')
|
||||
self.assertEqual(sent[0]['partial_execution_targets'], ['9'])
|
||||
self.assertNotIn('partial_execution_targets', sent[1])
|
||||
|
||||
|
||||
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_select_active_workers_filters_offline(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
|
||||
async def fake_probe(worker):
|
||||
return worker["id"] != "w2"
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=3,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_disables_delegate_when_all_offline(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker):
|
||||
return False
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=True,
|
||||
probe_concurrency=1,
|
||||
)
|
||||
|
||||
self.assertEqual(active_workers, [])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
|
||||
|
||||
async def fake_http_probe(_worker):
|
||||
return False
|
||||
|
||||
async def fake_ws_probe(_worker):
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
|
||||
dispatch,
|
||||
"worker_ws_is_active",
|
||||
side_effect=fake_ws_probe,
|
||||
) as ws_probe:
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=True,
|
||||
delegate_master=False,
|
||||
probe_concurrency=4,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
|
||||
self.assertEqual(ws_probe.call_count, 2)
|
||||
self.assertEqual(http_probe.call_count, 0)
|
||||
|
||||
async def test_probe_concurrency_is_bounded(self):
|
||||
workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
|
||||
state = {"in_flight": 0, "max_in_flight": 0}
|
||||
|
||||
async def fake_probe(_worker):
|
||||
state["in_flight"] += 1
|
||||
state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
|
||||
await asyncio.sleep(0.01)
|
||||
state["in_flight"] -= 1
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=2,
|
||||
)
|
||||
|
||||
self.assertEqual(len(active_workers), len(workers))
|
||||
self.assertLessEqual(state["max_in_flight"], 2)
|
||||
self.assertGreaterEqual(state["max_in_flight"], 2)
|
||||
|
||||
async def test_select_least_busy_worker_round_robins_idle_workers(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 0, "w2": 0, "w3": 2}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
dispatch._least_busy_rr_index = 0
|
||||
selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
|
||||
self.assertEqual(selected1["id"], "w1")
|
||||
self.assertEqual(selected2["id"], "w2")
|
||||
self.assertEqual(selected3["id"], "w1")
|
||||
|
||||
async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 5, "w2": 2, "w3": 4}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
|
||||
|
||||
self.assertEqual(selected["id"], "w2")
|
||||
|
||||
async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker_url, timeout=3.0):
|
||||
return None
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
|
||||
|
||||
self.assertIsNone(selected)
|
||||
|
||||
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_requires_exact_execution_targets_support(self):
|
||||
for value in (None, False, 1, "true", []):
|
||||
with self.subTest(value=value):
|
||||
capability = _valid_capability()
|
||||
if value is None:
|
||||
capability.pop("execution_targets")
|
||||
else:
|
||||
capability["execution_targets"] = value
|
||||
session = _FakeSession({("worker", "capabilities"): (200, capability)})
|
||||
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, "execution_targets.*upgrade"):
|
||||
await dispatch.preflight_image_list_workers([{"id": "worker"}])
|
||||
|
||||
async def test_image_list_value_results_requires_exact_true_only_for_overrides(self):
|
||||
for value in (None, False, 1, "true", [], True):
|
||||
capability = _valid_capability()
|
||||
if value is not None:
|
||||
capability['value_results'] = value
|
||||
session = _FakeSession({('worker', 'capabilities'): (200, capability),
|
||||
('worker', 'queue'): (200, _valid_queue())})
|
||||
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([{'id': 'worker'}]), {'worker': True})
|
||||
if value is True:
|
||||
self.assertEqual(await dispatch.preflight_image_list_workers(
|
||||
[{'id': 'worker'}], require_value_results=True), {'worker': True})
|
||||
else:
|
||||
with self.assertRaisesRegex(RuntimeError, 'value_results.*upgrade'):
|
||||
await dispatch.preflight_image_list_workers([{'id': 'worker'}], require_value_results=True)
|
||||
|
||||
async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
|
||||
malformed = {"supported": True, "protocol_version": 2, "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"] = 1
|
||||
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()
|
||||
@@ -13,7 +13,7 @@ class DistributedValueTests(unittest.TestCase):
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
pkg_name = "dv_test_pkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
@@ -28,18 +28,6 @@ class DistributedValueTests(unittest.TestCase):
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[pkg_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{pkg_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{pkg_name}.nodes"] = nodes_pkg
|
||||
|
||||
hidden_inputs_mod = types.ModuleType(f"{pkg_name}.nodes.hidden_inputs")
|
||||
hidden_inputs_mod.build_worker_identity_hidden_inputs = lambda: {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
}
|
||||
sys.modules[f"{pkg_name}.nodes.hidden_inputs"] = hidden_inputs_mod
|
||||
|
||||
utils_pkg = types.ModuleType(f"{pkg_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{pkg_name}.utils"] = utils_pkg
|
||||
@@ -49,16 +37,6 @@ class DistributedValueTests(unittest.TestCase):
|
||||
logging_mod.log = lambda *_a, **_k: None
|
||||
sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
|
||||
|
||||
worker_ids_path = Path(__file__).resolve().parents[3] / "utils" / "worker_ids.py"
|
||||
worker_ids_spec = importlib.util.spec_from_file_location(
|
||||
f"{pkg_name}.utils.worker_ids",
|
||||
worker_ids_path,
|
||||
)
|
||||
worker_ids_mod = importlib.util.module_from_spec(worker_ids_spec)
|
||||
assert worker_ids_spec is not None and worker_ids_spec.loader is not None
|
||||
worker_ids_spec.loader.exec_module(worker_ids_mod)
|
||||
sys.modules[f"{pkg_name}.utils.worker_ids"] = worker_ids_mod
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{pkg_name}.nodes.utilities", module_path
|
||||
)
|
||||
@@ -0,0 +1,341 @@
|
||||
"""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,
|
||||
VALUE_OUT_BASE,
|
||||
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}}
|
||||
|
||||
|
||||
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", VALUE_OUT_BASE],
|
||||
}},
|
||||
"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, masks, {
|
||||
"value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"],
|
||||
"value_3": [0.2], "value_4": [37],
|
||||
})
|
||||
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]["values"]["value_1"], refs[:1]))
|
||||
self.assertEqual(rows[1]["mask"].dtype, torch.float32)
|
||||
self.assertEqual([row["index"] for row in rows], [0, 1])
|
||||
self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"])
|
||||
self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2])
|
||||
self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37])
|
||||
|
||||
def test_one_value_is_shared_and_n_values_pair_in_order(self):
|
||||
image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
|
||||
rows = normalize_image_list_inputs(image, values={
|
||||
"value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]})
|
||||
self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3)
|
||||
self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3)
|
||||
self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3)
|
||||
self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3)
|
||||
rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]})
|
||||
self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}])
|
||||
with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"):
|
||||
normalize_image_list_inputs(image, values={"value_1": ["one", "two"]})
|
||||
with self.assertRaisesRegex(ValueError, "mask has 2 values"):
|
||||
normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)))
|
||||
shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4)))
|
||||
self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask))
|
||||
|
||||
def test_value_batches_stay_whole_and_only_images_split(self):
|
||||
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
pair = torch.rand((2, 4, 4, 3)) # e.g. two references for one item
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": [pair]})
|
||||
self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows))
|
||||
per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))]
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": per_item})
|
||||
self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1]))
|
||||
|
||||
def test_rejects_unsupported_values_and_counts(self):
|
||||
images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
with self.assertRaisesRegex(ValueError, "value_1 has 3 values"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]})
|
||||
with self.assertRaisesRegex(TypeError, "unsupported type dict"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]})
|
||||
with self.assertRaisesRegex(TypeError, "floating-point"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]})
|
||||
with self.assertRaisesRegex(TypeError, "finite"):
|
||||
normalize_image_list_inputs(images, values={"value_1": [float("nan")]})
|
||||
with self.assertRaisesRegex(ValueError, "Unknown Start value slot"):
|
||||
normalize_image_list_inputs(images, values={"value_11": [1]})
|
||||
with self.assertRaisesRegex(ValueError, "at least one image"):
|
||||
normalize_image_list_inputs(None)
|
||||
# Unconnected Autogrow slots arrive as None and are ignored.
|
||||
rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None})
|
||||
self.assertEqual([row["values"] for row in rows], [{}, {}])
|
||||
|
||||
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,
|
||||
"mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
|
||||
"values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16),
|
||||
"value_2": "safe text", "value_3": 0.625, "value_4": 42,
|
||||
"value_5": False},
|
||||
"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["values"]["value_1"], item["values"]["value_1"]))
|
||||
self.assertEqual(restored["values"]["value_1"].dtype, torch.float16)
|
||||
self.assertTrue(torch.equal(restored["mask"], item["mask"]))
|
||||
self.assertEqual(restored["values"]["value_2"], "safe text")
|
||||
self.assertEqual(restored["values"]["value_3"], 0.625)
|
||||
self.assertIs(restored["values"]["value_4"], 42)
|
||||
self.assertIs(restored["values"]["value_5"], False)
|
||||
with self.assertRaisesRegex(ValueError, "unknown slot"):
|
||||
decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}})
|
||||
with self.assertRaisesRegex(ValueError, "unsupported type"):
|
||||
decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}})
|
||||
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"])
|
||||
|
||||
def test_worker_prompt_drops_per_item_value_links(self):
|
||||
prompt = graph_one_region()
|
||||
prompt["8"] = {"class_type": "StringList", "inputs": {}}
|
||||
prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]})
|
||||
import sys, types
|
||||
from unittest.mock import patch
|
||||
registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})}
|
||||
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_12345678901234567890123456789012", "ctx", 0, "http://m")
|
||||
self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")])
|
||||
self.assertNotIn("8", worker)
|
||||
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]
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
|
||||
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]
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "typed context output"):
|
||||
validate_image_list_regions(prompt, explicit_targets=explicit_targets)
|
||||
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]}}
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
|
||||
validate_image_list_regions(nested, explicit_targets=explicit_targets)
|
||||
outside = graph_one_region()
|
||||
outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
|
||||
for explicit_targets in (False, True):
|
||||
with self.subTest(explicit_targets=explicit_targets):
|
||||
with self.assertRaisesRegex(ValueError, "cross-boundary"):
|
||||
validate_image_list_regions(outside, explicit_targets=explicit_targets)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,245 @@
|
||||
"""CPU-only execution-copy planning; no ComfyUI imports or prompt execution."""
|
||||
import copy
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
from image_list.execution_plan import prepare_image_list_execution
|
||||
from image_list.runtime import (
|
||||
CONTEXT_OUT,
|
||||
END_CLASS,
|
||||
IMAGE_OUT,
|
||||
START_CLASS,
|
||||
build_worker_prompt,
|
||||
validate_image_list_regions,
|
||||
)
|
||||
|
||||
|
||||
class OutputNode:
|
||||
OUTPUT_NODE = True
|
||||
|
||||
|
||||
class SchemaOutputNode:
|
||||
OUTPUT_NODE = False # V3's schema, not this fallback, is authoritative.
|
||||
|
||||
@classmethod
|
||||
def GET_NODE_INFO_V1(cls):
|
||||
return {"output_node": True, "output": ("IMAGE",)}
|
||||
|
||||
|
||||
class ModelLoader:
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def registry(monkeypatch):
|
||||
classes = {
|
||||
END_CLASS: SchemaOutputNode,
|
||||
"SaveImage": OutputNode,
|
||||
"PreviewImage": OutputNode,
|
||||
"OutputUtility": OutputNode,
|
||||
"SchemaUtility": SchemaOutputNode,
|
||||
"ModelLoader": ModelLoader,
|
||||
}
|
||||
monkeypatch.setitem(sys.modules, "nodes", types.SimpleNamespace(NODE_CLASS_MAPPINGS=classes))
|
||||
return classes
|
||||
|
||||
|
||||
def node(class_type, **inputs):
|
||||
return {"class_type": class_type, "inputs": inputs}
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
"loader": node("ModelLoader"),
|
||||
"source": node("ImageSource"),
|
||||
"start": node(START_CLASS, images=["source", 0]),
|
||||
"process": node("FrameOperation", image=["start", IMAGE_OUT], model=["loader", 0]),
|
||||
"end": node(END_CLASS, images=["process", 0], context=["start", CONTEXT_OUT]),
|
||||
"save": node("SaveImage", images=["end", 0]),
|
||||
}
|
||||
|
||||
|
||||
def test_discards_multistep_previews_but_preserves_other_output_roots(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"debug-transform": node("DebugTransform", images=["start", IMAGE_OUT]),
|
||||
"debug-preview": node("PreviewImage", images=["debug-transform", 0]),
|
||||
"process-preview": node("PreviewImage", images=["process", 0]),
|
||||
"unrelated-source": node("ImageSource"),
|
||||
"unrelated-preview": node("PreviewImage", images=["unrelated-source", 0]),
|
||||
"upstream-preview": node("PreviewImage", images=["source", 0]),
|
||||
"post-transform": node("FrameOperation", images=["end", 0]),
|
||||
"post-preview": node("PreviewImage", images=["post-transform", 0]),
|
||||
})
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
validate_image_list_regions(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"debug-transform", "debug-preview", "process-preview"}
|
||||
assert targets == ["end", "save", "unrelated-preview", "upstream-preview", "post-preview"]
|
||||
validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("class_type", ["OutputUtility", "SchemaUtility"])
|
||||
def test_required_output_utility_keeps_complete_dependency_chain(registry, class_type):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"utility-config": node("Constant", value=7),
|
||||
"utility-first": node(class_type, images=["process", 0], config=["utility-config", 0]),
|
||||
"utility-second": node(class_type, images=["utility-first", 0]),
|
||||
"result": node("FrameOperation", image=["utility-second", 0]),
|
||||
"discarded-preview": node("PreviewImage", images=["utility-first", 0]),
|
||||
})
|
||||
prompt["end"]["inputs"]["images"] = ["result", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
registered = dict(registry)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {"discarded-preview"}
|
||||
assert targets == ["end", "save"]
|
||||
with pytest.raises(ValueError, match="Output node"):
|
||||
validate_image_list_regions(prepared)
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
worker = build_worker_prompt(prepared, region, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == set(prepared) - {"source", "save"}
|
||||
assert worker["utility-second"]["inputs"] == {"images": ["utility-first", 0]}
|
||||
assert prompt == before
|
||||
assert registry == registered
|
||||
assert OutputNode.OUTPUT_NODE is True
|
||||
assert SchemaOutputNode.OUTPUT_NODE is False
|
||||
assert SchemaOutputNode.GET_NODE_INFO_V1()["output_node"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["start", "process"])
|
||||
@pytest.mark.parametrize("direct", [False, True])
|
||||
def test_internal_bypass_into_post_end_computation_is_rejected(registry, source, direct):
|
||||
prompt = graph()
|
||||
prompt["post-process"] = node("Composite", images=["end", 0])
|
||||
if not direct:
|
||||
prompt["bypass"] = node("FrameOperation", image=[source, 0])
|
||||
source = "bypass"
|
||||
prompt["post-process"]["inputs"]["reference"] = [source, 0]
|
||||
prompt["save"]["inputs"]["images"] = ["post-process", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError, match="cross-boundary"):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_sequential_namespaced_regions_and_intermediate_output_are_preserved(registry):
|
||||
prompt = graph()
|
||||
prompt.update({
|
||||
"next-start": node(START_CLASS, images=["end", 0]),
|
||||
"next-process": node("SchemaUtility", images=["next-start", IMAGE_OUT], model=["loader", 0]),
|
||||
"next-end": node(END_CLASS, images=["next-process", 0], context=["next-start", CONTEXT_OUT]),
|
||||
"next-save": node("SaveImage", images=["next-end", 0]),
|
||||
"first-debug": node("PreviewImage", images=["process", 0]),
|
||||
"second-debug": node("PreviewImage", images=["next-process", 0]),
|
||||
})
|
||||
renamed = {key: f"group:7/{key}" for key in prompt}
|
||||
for spec in prompt.values():
|
||||
for name, value in spec["inputs"].items():
|
||||
if isinstance(value, list):
|
||||
spec["inputs"][name] = [renamed[value[0]], value[1]]
|
||||
prompt = {renamed[key]: spec for key, spec in prompt.items()}
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert set(prepared) == set(prompt) - {renamed["first-debug"], renamed["second-debug"]}
|
||||
assert targets == [renamed[name] for name in ("end", "save", "next-end", "next-save")]
|
||||
regions = validate_image_list_regions(prepared, explicit_targets=True)
|
||||
assert {(r["start_id"], r["end_id"]) for r in regions} == {
|
||||
(renamed["start"], renamed["end"]), (renamed["next-start"], renamed["next-end"]),
|
||||
}
|
||||
second = next(r for r in regions if r["start_id"] == renamed["next-start"])
|
||||
worker = build_worker_prompt(prepared, second, "token", "context", 0, "http://master.invalid")
|
||||
assert set(worker) == {renamed[name] for name in ("loader", "next-start", "next-process", "next-end")}
|
||||
|
||||
|
||||
def test_shared_loader_and_its_independent_output_root_are_preserved(registry):
|
||||
registry["ModelLoader"] = type("OutputLoader", (ModelLoader, OutputNode), {})
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["independent-save"] = node("SaveImage", model=["loader", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["loader", "end", "save", "independent-save"]
|
||||
region = validate_image_list_regions(prepared, explicit_targets=True)[0]
|
||||
assert "loader" in region["node_ids"]
|
||||
assert "loader" not in region["branch_ids"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("unsafe_class", [None, type("ListLoader", (), {
|
||||
"RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,),
|
||||
})])
|
||||
def test_unknown_or_list_producing_shared_loader_still_fails_closed(registry, unsafe_class):
|
||||
if unsafe_class is None:
|
||||
registry.pop("ModelLoader")
|
||||
else:
|
||||
registry["ModelLoader"] = unsafe_class
|
||||
prompt = graph()
|
||||
prompt["source"]["inputs"]["model"] = ["loader", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*loader"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
def test_upstream_per_item_bypass_is_not_hidden_by_preview_pruning(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["inputs"]["reference"] = ["source", 0]
|
||||
prompt["debug"] = node("PreviewImage", images=["process", 0])
|
||||
with pytest.raises(ValueError, match="bypasses Start.*source"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nested", [False, True])
|
||||
def test_overlapping_or_nested_regions_cannot_be_discarded_as_side_branches(registry, nested):
|
||||
prompt = graph()
|
||||
prompt["other-start"] = node(START_CLASS, images=["process" if nested else "source", 0])
|
||||
prompt["other-process"] = node("FrameOperation", image=["other-start", IMAGE_OUT], reference=["process", 0])
|
||||
prompt["other-end"] = node(END_CLASS, images=["other-process", 0], context=["other-start", CONTEXT_OUT])
|
||||
with pytest.raises(ValueError, match="Nested/ambiguous|cross-boundary|ambiguous processing"):
|
||||
prepare_image_list_execution(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid", ["missing-end", "wrong-context", "missing-images", "unrelated-images"])
|
||||
def test_pairing_and_result_validation_cannot_be_pruned_away(registry, invalid):
|
||||
prompt = graph()
|
||||
if invalid == "missing-end":
|
||||
del prompt["end"]
|
||||
elif invalid == "wrong-context":
|
||||
prompt["end"]["inputs"]["context"] = ["source", 0]
|
||||
elif invalid == "missing-images":
|
||||
del prompt["end"]["inputs"]["images"]
|
||||
else:
|
||||
prompt["end"]["inputs"]["images"] = ["source", 0]
|
||||
before = copy.deepcopy(prompt)
|
||||
with pytest.raises(ValueError):
|
||||
prepare_image_list_execution(prompt)
|
||||
assert prompt == before
|
||||
|
||||
|
||||
def test_execution_copy_is_deep_and_repeated_preparation_is_stable(registry):
|
||||
prompt = graph()
|
||||
prompt["process"]["_meta"] = {"title": "Original", "tags": ["saved"]}
|
||||
before = copy.deepcopy(prompt)
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
prepared_again, targets_again = prepare_image_list_execution(prepared)
|
||||
assert prepared_again == prepared
|
||||
assert targets_again == targets
|
||||
prepared["start"]["inputs"]["images"][0] = "changed"
|
||||
prepared["process"]["_meta"]["tags"].append("execution-only")
|
||||
assert prompt == before
|
||||
assert prepared_again == before
|
||||
|
||||
|
||||
def test_v3_false_output_flag_overrides_legacy_true(registry):
|
||||
registry["NotOutput"] = type("NotOutput", (), {
|
||||
"OUTPUT_NODE": True,
|
||||
"GET_NODE_INFO_V1": classmethod(lambda cls: {"output_node": False}),
|
||||
})
|
||||
prompt = graph()
|
||||
prompt["not-output"] = node("NotOutput", images=["end", 0])
|
||||
prepared, targets = prepare_image_list_execution(prompt)
|
||||
assert prepared == prompt
|
||||
assert targets == ["end", "save"]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user