Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 robertvoy
|
1d4dcb985e |
@@ -1,3 +1,4 @@
|
||||
# These are supported funding model platforms
|
||||
|
||||
github: robertvoy
|
||||
buy_me_a_coffee: robertvoy
|
||||
|
||||
@@ -19,6 +19,6 @@ jobs:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
+1
-6
@@ -1,9 +1,4 @@
|
||||
bin/
|
||||
logs/
|
||||
gpu_config.json
|
||||
__pycache__/
|
||||
**/__pycache__/
|
||||
*.py[cod]
|
||||
node_modules/
|
||||
npm-debug.log*
|
||||
AGENTS.md
|
||||
__pycache__/
|
||||
@@ -38,12 +38,11 @@
|
||||
|
||||
<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 four types of workers:
|
||||
ComfyUI Distributed supports three types of workers:
|
||||
|
||||
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
|
||||
- **Remote Workers** - GPUs on other computers in your network
|
||||
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
|
||||
- **Serverless (Runpod)** - Queue-based endpoint pools for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
@@ -52,6 +51,7 @@ ComfyUI Distributed supports four 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,9 +92,8 @@ Generate multiple images in the time it takes to generate one. Each worker uses
|
||||
1. Open your ComfyUI workflow
|
||||
2. Add **Distributed Seed** → connect to sampler's seed
|
||||
3. Add **Distributed Collector** → after VAE Decode
|
||||
4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
|
||||
5. Enable workers in the UI
|
||||
6. Run the workflow!
|
||||
4. Enable workers in the UI
|
||||
5. Run the workflow!
|
||||
|
||||
### Parallel WAN Generation
|
||||
Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
|
||||
@@ -117,7 +116,7 @@ Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, w
|
||||
|
||||
|
||||
|
||||
> [Download workflow](/workflows/distributed-upscale.json)
|
||||
> [Download workflow](/workflows/distributed-txt2img.json)
|
||||
|
||||
1. Load your image
|
||||
2. Upscale with ESRGAN or similar
|
||||
@@ -148,40 +147,13 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
|
||||
Control your distributed cluster programmatically without opening the browser.
|
||||
|
||||
* **Endpoint:** `POST /distributed/queue`
|
||||
* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
|
||||
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
|
||||
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
|
||||
|
||||
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
|
||||
|
||||
---
|
||||
|
||||
## Distributed Value
|
||||
|
||||
Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
|
||||
|
||||
- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
|
||||
- The node shows only currently enabled workers.
|
||||
- If worker enablement changes, worker fields update automatically.
|
||||
- When disconnected, it resets to default string mode and clears per-worker overrides.
|
||||
- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
|
||||
|
||||
---
|
||||
|
||||
## Nodes
|
||||
|
||||
| Node | Description |
|
||||
|------|-------------|
|
||||
| **Distributed Seed** | Generates unique seeds for each worker |
|
||||
| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
|
||||
| **Distributed Value** | Outputs per-worker override values with fallback to default |
|
||||
| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
|
||||
| **Image Batch Divider** | Splits image batches for multi-GPU output |
|
||||
| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
|
||||
| **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
|
||||
| **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
|
||||
|
||||
---
|
||||
|
||||
## FAQ
|
||||
|
||||
<details>
|
||||
@@ -196,12 +168,12 @@ No, it does not speed up the generation of a single image or video. Instead, it
|
||||
|
||||
<details>
|
||||
<summary>Does it work with the ComfyUI desktop app?</summary>
|
||||
Yes, it does now.
|
||||
Currently, it is not compatible with the ComfyUI desktop app.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I combine my RTX 5090 with a GTX 980 to get faster results?</summary>
|
||||
Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks.
|
||||
Yes, you can combine different GPUs, but performance is optimized when using similar GPUs. A significant performance imbalance between GPUs may cause bottlenecks. For upscaling, setting `static_distribution` to `false` allows the faster GPU to handle more processing, which can mitigate some bottlenecks. Note that this setting only applies to upscaling tasks.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -209,11 +181,6 @@ Yes, you can combine different GPUs, but performance is optimized when using sim
|
||||
Yes, it is compatible with cloud providers. Refer to the setup guides for detailed instructions.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I use my main machine just to coordinate workers without rendering?</summary>
|
||||
Yes. Open the Distributed panel and uncheck the master toggle to run in orchestrator-only mode. The master will distribute work to workers but won't render locally. If all workers become unavailable, the master automatically re-enables to ensure your workflow still runs.
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>Can I make this work with my Docker setup?</summary>
|
||||
Yes, it is compatible with Docker setups, but you will need to configure your Docker environment yourself. Unfortunately, assistance with Docker configuration is not provided.
|
||||
@@ -237,6 +204,3 @@ Your support helps keep this project thriving.
|
||||
|
||||
Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
+56
-13
@@ -1,22 +1,65 @@
|
||||
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||
import os
|
||||
import sys
|
||||
|
||||
from .nodes.v3 import NODES
|
||||
from .runtime.bootstrap import initialize
|
||||
# Add the directory to Python path
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
|
||||
WEB_DIRECTORY = './web'
|
||||
# Patch ComfyUI execution validation for DistributedVideoCollector
|
||||
try:
|
||||
import execution
|
||||
from .distributed import ImageBatchDivider
|
||||
|
||||
# Store original validate function if it exists
|
||||
if hasattr(execution, 'validate_outputs'):
|
||||
original_validate_outputs = execution.validate_outputs
|
||||
|
||||
def patched_validate_outputs(executor, node_id, result, node_class):
|
||||
if node_class == ImageBatchDivider:
|
||||
return # Skip validation for our dynamic output node
|
||||
return original_validate_outputs(executor, node_id, result, node_class)
|
||||
|
||||
execution.validate_outputs = patched_validate_outputs
|
||||
else:
|
||||
# Fallback: patch the main execute method if validate_outputs doesn't exist
|
||||
if hasattr(execution.PromptExecutor, 'execute'):
|
||||
original_execute = execution.PromptExecutor.execute
|
||||
|
||||
def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
|
||||
# This is a more complex patch - for now just call original
|
||||
# The ByPassTypeTuple should handle most validation issues
|
||||
return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
|
||||
|
||||
execution.PromptExecutor.execute = patched_execute
|
||||
|
||||
except ImportError:
|
||||
pass # ComfyUI execution module not available during import
|
||||
|
||||
# Import everything needed from the main module
|
||||
from .distributed import (
|
||||
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
|
||||
class DistributedExtension(ComfyExtension):
|
||||
async def on_load(self) -> None:
|
||||
initialize()
|
||||
# Import utilities
|
||||
from .utils.config import ensure_config_exists, CONFIG_FILE
|
||||
from .utils.logging import debug_log
|
||||
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return list(NODES)
|
||||
# Import distributed upscale nodes
|
||||
from .distributed_upscale import (
|
||||
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
async def comfy_entrypoint() -> DistributedExtension:
|
||||
return DistributedExtension()
|
||||
ensure_config_exists()
|
||||
|
||||
# Merge node mappings
|
||||
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
|
||||
|
||||
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
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
|
||||
@@ -1,580 +0,0 @@
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
try:
|
||||
from ..utils.config import config_transaction, load_config, save_config
|
||||
except ImportError:
|
||||
from ..utils.config import load_config
|
||||
|
||||
try:
|
||||
from ..utils.config import save_config
|
||||
except ImportError:
|
||||
def save_config(_config):
|
||||
return True
|
||||
|
||||
@asynccontextmanager
|
||||
async def config_transaction():
|
||||
config = load_config()
|
||||
original_snapshot = json.dumps(config, sort_keys=True)
|
||||
yield config
|
||||
if json.dumps(config, sort_keys=True) != original_snapshot:
|
||||
save_config(config)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host
|
||||
from ..utils.secret_fields import is_secret_field_name
|
||||
from ..utils.serverless_config import (
|
||||
ServerlessSecretChanges,
|
||||
public_config,
|
||||
validate_serverless_worker,
|
||||
)
|
||||
|
||||
|
||||
def _positive_int(value):
|
||||
return value > 0
|
||||
|
||||
|
||||
CONFIG_SCHEMA = {
|
||||
"workers": (list, None),
|
||||
"master": (dict, None),
|
||||
"settings": (dict, None),
|
||||
"tunnel": (dict, None),
|
||||
"managed_processes": (dict, None),
|
||||
"worker_timeout_seconds": (int, _positive_int),
|
||||
"debug": (bool, None),
|
||||
"auto_launch_workers": (bool, None),
|
||||
"stop_workers_on_master_exit": (bool, None),
|
||||
"master_delegate_only": (bool, None),
|
||||
"websocket_orchestration": (bool, None),
|
||||
"has_auto_populated_workers": (bool, None),
|
||||
}
|
||||
|
||||
_SETTINGS_FIELDS = {
|
||||
"worker_timeout_seconds",
|
||||
"debug",
|
||||
"auto_launch_workers",
|
||||
"stop_workers_on_master_exit",
|
||||
"master_delegate_only",
|
||||
"websocket_orchestration",
|
||||
"has_auto_populated_workers",
|
||||
}
|
||||
|
||||
_WORKER_FIELDS = [
|
||||
("enabled", None, False),
|
||||
("name", None, False),
|
||||
("port", None, False),
|
||||
("host", normalize_host, True),
|
||||
("cuda_device", None, True),
|
||||
("extra_args", None, True),
|
||||
("type", None, False),
|
||||
]
|
||||
|
||||
_MASTER_FIELDS = [
|
||||
("name", None, False),
|
||||
("host", normalize_host, True),
|
||||
("port", None, False),
|
||||
("cuda_device", None, False),
|
||||
("extra_args", None, False),
|
||||
]
|
||||
|
||||
|
||||
def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
|
||||
"""Apply a partial update to a target dict based on field rules."""
|
||||
for key, normalizer, remove_on_none in field_rules:
|
||||
if key not in data:
|
||||
continue
|
||||
value = data[key]
|
||||
if value is None and remove_on_none:
|
||||
target.pop(key, None)
|
||||
else:
|
||||
target[key] = normalizer(value) if (normalizer and value is not None) else value
|
||||
|
||||
|
||||
_SECRET_FIELDS = {"api_key", "clear_api_key"}
|
||||
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
||||
_SERVERLESS_ONLY_FIELDS = {
|
||||
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
|
||||
}
|
||||
_SERVERLESS_CONFIG_FIELDS = {
|
||||
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
|
||||
}
|
||||
|
||||
|
||||
class _ActivePoolCheckError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _error_response(message, status):
|
||||
return web.json_response({"status": "error", "message": message}, status=status)
|
||||
|
||||
|
||||
def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0):
|
||||
if isinstance(value, dict):
|
||||
for key, item in value.items():
|
||||
field_name = key.casefold() if isinstance(key, str) else ""
|
||||
if is_secret_field_name(field_name):
|
||||
if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS):
|
||||
return True
|
||||
if _contains_secret_fields(
|
||||
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
|
||||
):
|
||||
return True
|
||||
elif isinstance(value, (list, tuple)):
|
||||
return any(
|
||||
_contains_secret_fields(
|
||||
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
|
||||
)
|
||||
for item in value
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _serverless_secret_store():
|
||||
from .orchestration.serverless_runtime import get_secret_store
|
||||
return get_secret_store()
|
||||
|
||||
|
||||
def _serverless_pool_is_active(worker_id):
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
return get_manager().pool_is_active(worker_id)
|
||||
|
||||
|
||||
def _ensure_pool_idle(worker_id):
|
||||
try:
|
||||
active = _serverless_pool_is_active(worker_id)
|
||||
except Exception:
|
||||
raise _ActivePoolCheckError from None
|
||||
if active:
|
||||
raise _ActivePoolCheckError
|
||||
|
||||
|
||||
def _worker_by_id(workers, worker_id):
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("id") == worker_id:
|
||||
return worker
|
||||
return None
|
||||
|
||||
|
||||
def _worker_type(worker):
|
||||
return worker.get("type", "local") if isinstance(worker, dict) else None
|
||||
|
||||
|
||||
def _strip_sensitive_fields(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: _strip_sensitive_fields(item)
|
||||
for key, item in value.items()
|
||||
if not is_secret_field_name(key)
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [_strip_sensitive_fields(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def _clean_serverless_worker(worker):
|
||||
clean = _strip_sensitive_fields(dict(worker))
|
||||
for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}:
|
||||
clean.pop(field, None)
|
||||
return clean
|
||||
|
||||
|
||||
def _validate_bulk_workers(workers):
|
||||
validated = []
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
||||
clean = _clean_serverless_worker(worker)
|
||||
validated.append(validate_serverless_worker(clean))
|
||||
else:
|
||||
validated.append(worker)
|
||||
return validated
|
||||
|
||||
|
||||
def _clean_bulk_type_conversions(old_workers, new_workers):
|
||||
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
for index, worker in enumerate(new_workers):
|
||||
if not isinstance(worker, dict) or worker.get("id") is None:
|
||||
continue
|
||||
old = old_by_id.get(str(worker["id"]))
|
||||
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
|
||||
continue
|
||||
converted = dict(worker)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
converted.pop(field, None)
|
||||
new_workers[index] = converted
|
||||
return new_workers
|
||||
|
||||
|
||||
def _bulk_serverless_changes(old_workers, new_workers):
|
||||
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
secret_ids = set()
|
||||
for worker_id, old in old_by_id.items():
|
||||
new = new_by_id.get(worker_id)
|
||||
if old.get("type") == "serverless":
|
||||
if new is None or new.get("type") != "serverless":
|
||||
_ensure_pool_idle(worker_id)
|
||||
secret_ids.add(worker_id)
|
||||
elif (old.get("endpoint_id") != new.get("endpoint_id")
|
||||
or old.get("concurrency", 1) != new.get("concurrency", 1)):
|
||||
_ensure_pool_idle(worker_id)
|
||||
elif new is not None and new.get("type") == "serverless":
|
||||
secret_ids.add(worker_id)
|
||||
return secret_ids
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/config")
|
||||
async def get_config_endpoint(request):
|
||||
try:
|
||||
config = load_config()
|
||||
return web.json_response(public_config(config, _serverless_secret_store()))
|
||||
except Exception:
|
||||
return _error_response("Unable to read config", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config")
|
||||
async def update_config_endpoint(request):
|
||||
"""Bulk config update with schema validation and credential redaction."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
return _error_response("Invalid JSON payload", 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 = []
|
||||
|
||||
for key, value in data.items():
|
||||
if key not in CONFIG_SCHEMA:
|
||||
errors.append(f"Unknown field: {key}")
|
||||
continue
|
||||
|
||||
expected_type, validator = CONFIG_SCHEMA[key]
|
||||
if not isinstance(value, expected_type):
|
||||
errors.append(f"{key}: expected {expected_type.__name__}")
|
||||
continue
|
||||
|
||||
if validator and not validator(value):
|
||||
errors.append(f"{key}: failed validation")
|
||||
continue
|
||||
|
||||
if key == "workers":
|
||||
try:
|
||||
value = _validate_bulk_workers(value)
|
||||
except (TypeError, ValueError):
|
||||
return _error_response("Invalid serverless worker configuration", 400)
|
||||
|
||||
if key in _SETTINGS_FIELDS:
|
||||
validated_settings[key] = value
|
||||
else:
|
||||
validated_root[key] = value
|
||||
|
||||
if errors:
|
||||
return web.json_response({
|
||||
"status": "error",
|
||||
"error": errors,
|
||||
"message": "; ".join(errors),
|
||||
}, status=400)
|
||||
|
||||
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 pool", 400)
|
||||
if changed_ids:
|
||||
store = store or _serverless_secret_store()
|
||||
for worker_id in sorted(changed_ids):
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
if prior_secret is not None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
|
||||
settings = config.setdefault("settings", {})
|
||||
settings.update(validated_settings)
|
||||
for key, value in validated_root.items():
|
||||
config[key] = value
|
||||
response_config = public_config(config, store or _serverless_secret_store())
|
||||
response = web.json_response({"status": "success", "config": response_config})
|
||||
secret_changes.commit()
|
||||
return response
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to update config", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
||||
async def queue_status_endpoint(request):
|
||||
"""Check if a job queue is initialized."""
|
||||
try:
|
||||
job_id = request.match_info['job_id']
|
||||
|
||||
# Import to ensure initialization
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
exists = job_id in prompt_server.distributed_pending_tile_jobs
|
||||
|
||||
debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
|
||||
return web.json_response({"exists": exists, "job_id": job_id})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
||||
async def update_worker_endpoint(request):
|
||||
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 _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:
|
||||
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")
|
||||
|
||||
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)
|
||||
|
||||
candidate = dict(existing or {})
|
||||
candidate["id"] = worker_id
|
||||
candidate["type"] = "serverless"
|
||||
candidate.setdefault("provider", "runpod")
|
||||
candidate.setdefault("enabled", False)
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "enabled", "name"):
|
||||
if field in data:
|
||||
candidate[field] = data[field]
|
||||
candidate = _clean_serverless_worker(candidate)
|
||||
validated = validate_serverless_worker(candidate)
|
||||
if "enabled" in validated and type(validated["enabled"]) is not bool:
|
||||
return _error_response("enabled must be a boolean", 400)
|
||||
if "name" in validated and not isinstance(validated["name"], str):
|
||||
return _error_response("name must be a string", 400)
|
||||
|
||||
store = _serverless_secret_store()
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
converting_to_serverless = current_type != "serverless"
|
||||
next_secret = None if converting_to_serverless else prior_secret
|
||||
if supplied_key:
|
||||
next_secret = supplied_key
|
||||
elif clear_key:
|
||||
next_secret = None
|
||||
|
||||
endpoint_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("endpoint_id") != validated.get("endpoint_id")
|
||||
)
|
||||
concurrency_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("concurrency", 1) != validated.get("concurrency", 1)
|
||||
)
|
||||
if current_type == "serverless" and (
|
||||
endpoint_changed or concurrency_changed or next_secret != prior_secret
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
|
||||
if validated.get("enabled", False) and next_secret is None:
|
||||
return _error_response("Save an API key before enabling this pool; disable it before clearing the key", 400)
|
||||
|
||||
if next_secret != prior_secret:
|
||||
if next_secret is None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
else:
|
||||
secret_changes.set(store, worker_id, next_secret)
|
||||
|
||||
if existing is None:
|
||||
workers.append(validated)
|
||||
else:
|
||||
existing.clear()
|
||||
existing.update(validated)
|
||||
result = {"status": "success", "api_key_configured": next_secret is not None}
|
||||
|
||||
elif existing is not None:
|
||||
cleaned_existing = _strip_sensitive_fields(existing)
|
||||
existing.clear()
|
||||
existing.update(cleaned_existing)
|
||||
if current_type == "serverless":
|
||||
if _contains_secret_fields(data):
|
||||
return _error_response("Invalid worker credentials", 400)
|
||||
_ensure_pool_idle(worker_id)
|
||||
store = _serverless_secret_store()
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
if prior_secret is not None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
existing.pop(field, None)
|
||||
elif _contains_secret_fields(data):
|
||||
return _error_response("API keys are only accepted for serverless workers", 400)
|
||||
_apply_field_patch(existing, data, _WORKER_FIELDS)
|
||||
result = {"status": "success"}
|
||||
elif all(key in data for key in ["name", "port", "cuda_device"]):
|
||||
if _contains_secret_fields(data):
|
||||
return _error_response("API keys are only accepted for serverless workers", 400)
|
||||
new_worker = {
|
||||
"id": worker_id,
|
||||
"name": data["name"],
|
||||
"host": normalize_host(data.get("host", "localhost")),
|
||||
"port": data["port"],
|
||||
"cuda_device": data["cuda_device"],
|
||||
"enabled": data.get("enabled", False),
|
||||
"extra_args": data.get("extra_args", ""),
|
||||
"type": data.get("type", "local")
|
||||
}
|
||||
workers.append(new_worker)
|
||||
result = {"status": "success"}
|
||||
else:
|
||||
return _error_response(
|
||||
"Worker not found and missing required fields for creation", 404
|
||||
)
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except ValueError:
|
||||
return _error_response("Invalid worker configuration", 400)
|
||||
except Exception:
|
||||
return _error_response("Unable to update worker configuration", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
||||
async def delete_worker_endpoint(request):
|
||||
return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
|
||||
|
||||
|
||||
async def _delete_worker_endpoint(request, secret_changes):
|
||||
store = None
|
||||
worker_id = None
|
||||
try:
|
||||
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 _error_response("Missing worker_id", 400)
|
||||
|
||||
result = None
|
||||
async with config_transaction() as config:
|
||||
workers = config.get("workers", [])
|
||||
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": "Worker deleted" if _worker_type(removed) == "serverless"
|
||||
else f"Worker {removed.get('name', worker_id)} deleted",
|
||||
}
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to delete worker", 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
|
||||
async def update_setting_endpoint(request):
|
||||
"""Updates a specific key in the settings object."""
|
||||
try:
|
||||
data = await request.json()
|
||||
key = data.get("key")
|
||||
value = data.get("value")
|
||||
|
||||
if not key or value is None:
|
||||
return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
|
||||
if key not in _SETTINGS_FIELDS:
|
||||
return await handle_api_error(request, f"Unknown setting: {key}", 400)
|
||||
|
||||
async with config_transaction() as config:
|
||||
if 'settings' not in config:
|
||||
config['settings'] = {}
|
||||
|
||||
config['settings'][key] = value
|
||||
|
||||
return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_master")
|
||||
async def update_master_endpoint(request):
|
||||
"""Updates master configuration."""
|
||||
try:
|
||||
data = await request.json()
|
||||
|
||||
async with config_transaction() as config:
|
||||
if 'master' not in config:
|
||||
config['master'] = {}
|
||||
_apply_field_patch(config['master'], data, _MASTER_FIELDS)
|
||||
|
||||
return web.json_response({"status": "success", "message": "Master configuration updated."})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
@@ -1,350 +0,0 @@
|
||||
import json
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import base64
|
||||
import binascii
|
||||
import time
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.image import pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import handle_api_error
|
||||
from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
|
||||
try:
|
||||
from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
|
||||
except ImportError:
|
||||
from .queue_orchestration import orchestrate_distributed_execution
|
||||
|
||||
def ensure_distributed_state():
|
||||
return None
|
||||
from .queue_request import parse_queue_request_payload
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
# 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 }
|
||||
|
||||
|
||||
def _decode_image_sync(image_path):
|
||||
"""Decode image/video file and compute hash in a threadpool worker."""
|
||||
import base64
|
||||
import hashlib
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(image_path)
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(image_path)
|
||||
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.webm'}
|
||||
file_ext = os.path.splitext(full_path)[1].lower()
|
||||
|
||||
if file_ext in video_extensions:
|
||||
with open(full_path, 'rb') as f:
|
||||
file_data = f.read()
|
||||
mime_types = {
|
||||
'.mp4': 'video/mp4',
|
||||
'.avi': 'video/x-msvideo',
|
||||
'.mov': 'video/quicktime',
|
||||
'.mkv': 'video/x-matroska',
|
||||
'.webm': 'video/webm'
|
||||
}
|
||||
mime_type = mime_types.get(file_ext, 'video/mp4')
|
||||
image_data = f"data:{mime_type};base64,{base64.b64encode(file_data).decode('utf-8')}"
|
||||
else:
|
||||
with Image.open(full_path) as img:
|
||||
if img.mode not in ('RGB', 'RGBA'):
|
||||
img = img.convert('RGB')
|
||||
buffer = io.BytesIO()
|
||||
img.save(buffer, format='PNG', compress_level=1)
|
||||
image_data = f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"image_data": image_data,
|
||||
"hash": file_hash,
|
||||
}
|
||||
|
||||
|
||||
def _check_file_sync(filename, expected_hash):
|
||||
"""Check file presence and hash in a threadpool worker."""
|
||||
import hashlib
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(filename)
|
||||
if not os.path.exists(full_path):
|
||||
return {
|
||||
"status": "success",
|
||||
"exists": False,
|
||||
}
|
||||
|
||||
hash_md5 = hashlib.md5()
|
||||
with open(full_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(4096), b""):
|
||||
hash_md5.update(chunk)
|
||||
file_hash = hash_md5.hexdigest()
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"exists": True,
|
||||
"hash_matches": file_hash == expected_hash,
|
||||
}
|
||||
|
||||
|
||||
def _decode_canonical_png_tensor(image_payload):
|
||||
"""Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
|
||||
if not isinstance(image_payload, str) or not image_payload.strip():
|
||||
raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
|
||||
|
||||
encoded = image_payload.strip()
|
||||
if encoded.startswith("data:"):
|
||||
header, sep, data_part = encoded.partition(",")
|
||||
if not sep:
|
||||
raise ValueError("Field 'image' data URL is malformed.")
|
||||
if not header.lower().startswith("data:image/png;base64"):
|
||||
raise ValueError("Field 'image' must be a PNG data URL when using data:* format.")
|
||||
encoded = data_part
|
||||
|
||||
try:
|
||||
png_bytes = base64.b64decode(encoded, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError("Field 'image' is not valid base64 PNG data.") from exc
|
||||
|
||||
if not png_bytes:
|
||||
raise ValueError("Field 'image' decoded to empty PNG data.")
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(png_bytes)) as img:
|
||||
img = img.convert("RGB")
|
||||
tensor = pil_to_tensor(img)
|
||||
return ensure_contiguous(tensor)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
|
||||
|
||||
|
||||
def _decode_audio_payload(audio_payload):
|
||||
"""Decode canonical audio payload into an AUDIO dict."""
|
||||
from ..utils.audio_payload import decode_audio_payload
|
||||
|
||||
return decode_audio_payload(audio_payload)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/prepare_job")
|
||||
async def prepare_job_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return await handle_api_error(request, "Missing multi_job_id", 400)
|
||||
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
|
||||
debug_log(f"Prepared queue for job {multi_job_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/clear_memory")
|
||||
async def clear_memory_endpoint(request):
|
||||
debug_log("Received request to clear VRAM.")
|
||||
try:
|
||||
# Use ComfyUI's prompt server queue system like the /free endpoint does
|
||||
if hasattr(server.PromptServer.instance, 'prompt_queue'):
|
||||
server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
|
||||
server.PromptServer.instance.prompt_queue.set_flag("free_memory", True)
|
||||
debug_log("Set queue flags for memory clearing.")
|
||||
|
||||
# Wait a bit for the queue to process
|
||||
await asyncio.sleep(MEMORY_CLEAR_DELAY)
|
||||
|
||||
# Also do direct cleanup as backup, but with error handling
|
||||
import gc
|
||||
import comfy.model_management as mm
|
||||
|
||||
try:
|
||||
mm.unload_all_models()
|
||||
except AttributeError as e:
|
||||
debug_log(f"Warning during model unload: {e}")
|
||||
|
||||
try:
|
||||
mm.soft_empty_cache()
|
||||
except Exception as e:
|
||||
debug_log(f"Warning during cache clear: {e}")
|
||||
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.ipc_collect()
|
||||
|
||||
debug_log("VRAM cleared successfully.")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared."})
|
||||
except Exception as e:
|
||||
# Even if there's an error, try to do basic cleanup
|
||||
import gc
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
debug_log(f"Partial VRAM clear completed with warning: {e}")
|
||||
return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/queue")
|
||||
async def distributed_queue_endpoint(request):
|
||||
"""Queue a distributed workflow, mirroring the UI orchestration pipeline."""
|
||||
try:
|
||||
raw_payload = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
|
||||
try:
|
||||
payload = parse_queue_request_payload(raw_payload)
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
|
||||
try:
|
||||
prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
|
||||
payload.prompt,
|
||||
payload.workflow_meta,
|
||||
payload.client_id,
|
||||
enabled_worker_ids=payload.enabled_worker_ids,
|
||||
delegate_master=payload.delegate_master,
|
||||
trace_execution_id=payload.trace_execution_id,
|
||||
)
|
||||
return web.json_response({
|
||||
"prompt_id": prompt_id,
|
||||
"number": prompt_number,
|
||||
"node_errors": node_errors,
|
||||
"worker_count": worker_count,
|
||||
"auto_prepare_supported": True,
|
||||
})
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, exc, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/load_image")
|
||||
async def load_image_endpoint(request):
|
||||
"""Load an image or video file and return it as base64 data with hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
image_path = data.get("image_path")
|
||||
|
||||
if not image_path:
|
||||
return await handle_api_error(request, "Missing image_path", 400)
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _decode_image_sync, image_path)
|
||||
return web.json_response(payload)
|
||||
except FileNotFoundError:
|
||||
return await handle_api_error(request, f"File not found: {image_path}", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/check_file")
|
||||
async def check_file_endpoint(request):
|
||||
"""Check if a file exists and matches the given hash."""
|
||||
try:
|
||||
data = await request.json()
|
||||
filename = data.get("filename")
|
||||
expected_hash = data.get("hash")
|
||||
|
||||
if not filename or not expected_hash:
|
||||
return await handle_api_error(request, "Missing filename or hash", 400)
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _check_file_sync, filename, expected_hash)
|
||||
return web.json_response(payload)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/job_complete")
|
||||
async def job_complete_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return await handle_api_error(request, "Expected a JSON object body", 400)
|
||||
|
||||
try:
|
||||
job_id = data.get("job_id")
|
||||
worker_id = data.get("worker_id")
|
||||
batch_idx = data.get("batch_idx")
|
||||
image_payload = data.get("image")
|
||||
audio_payload = data.get("audio")
|
||||
is_last = data.get("is_last")
|
||||
|
||||
errors = []
|
||||
if not isinstance(job_id, str) or not job_id.strip():
|
||||
errors.append("job_id: expected non-empty string")
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
errors.append("worker_id: expected non-empty string")
|
||||
if not isinstance(batch_idx, int) or batch_idx < 0:
|
||||
errors.append("batch_idx: expected non-negative integer")
|
||||
image_provided = image_payload is not None
|
||||
audio_provided = audio_payload is not None
|
||||
if not image_provided and not audio_provided:
|
||||
errors.append("expected at least one of image or audio")
|
||||
if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
|
||||
errors.append("image: expected non-empty base64 PNG string")
|
||||
if audio_provided and not isinstance(audio_payload, dict):
|
||||
errors.append("audio: expected object when provided")
|
||||
if not isinstance(is_last, bool):
|
||||
errors.append("is_last: expected boolean")
|
||||
if errors:
|
||||
return await handle_api_error(request, errors, 400)
|
||||
|
||||
try:
|
||||
tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
|
||||
decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
|
||||
except ValueError as exc:
|
||||
return await handle_api_error(request, exc, 400)
|
||||
multi_job_id = job_id.strip()
|
||||
worker_id = worker_id.strip()
|
||||
|
||||
pending = None
|
||||
queue_size = 0
|
||||
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
|
||||
while pending is None:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
if pending is not None:
|
||||
await pending.put(
|
||||
{
|
||||
"tensor": tensor,
|
||||
"worker_id": worker_id,
|
||||
"image_index": int(batch_idx),
|
||||
"is_last": is_last,
|
||||
"audio": decoded_audio,
|
||||
}
|
||||
)
|
||||
queue_size = pending.qsize()
|
||||
break
|
||||
|
||||
if time.monotonic() > deadline:
|
||||
return await handle_api_error(request, "job not initialized", 404)
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
debug_log(
|
||||
f"job_complete received canonical envelope - job_id: {multi_job_id}, "
|
||||
f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, "
|
||||
f"queue_size: {queue_size}"
|
||||
)
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
@@ -1,4 +0,0 @@
|
||||
# Orchestration helpers split out from queue_orchestration.py:
|
||||
# - prompt_transform.py: graph pruning + hidden input overrides
|
||||
# - media_sync.py: remote media/path normalization
|
||||
# - dispatch.py: worker probe + prompt dispatch
|
||||
@@ -1,277 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session, probe_worker
|
||||
try:
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
except ImportError:
|
||||
def trace_debug(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
def trace_info(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
try:
|
||||
from ..schemas import parse_positive_int
|
||||
except ImportError:
|
||||
def parse_positive_int(value, default):
|
||||
try:
|
||||
parsed = int(value)
|
||||
return parsed if parsed > 0 else default
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
_least_busy_rr_index = 0
|
||||
|
||||
|
||||
async def worker_is_active(worker):
|
||||
"""Ping worker's /prompt endpoint to confirm it's reachable."""
|
||||
url = build_worker_url(worker)
|
||||
return await probe_worker(url, timeout=3.0) is not None
|
||||
|
||||
|
||||
async def worker_ws_is_active(worker):
|
||||
"""Ping worker's websocket endpoint to confirm it's reachable."""
|
||||
session = await get_client_session()
|
||||
url = build_worker_url(worker, "/distributed/worker_ws")
|
||||
try:
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
|
||||
await ws.close()
|
||||
return True
|
||||
except asyncio.TimeoutError:
|
||||
debug_log(f"[Distributed] Worker WS probe timed out: {url}")
|
||||
return False
|
||||
except aiohttp.ClientConnectorError:
|
||||
debug_log(f"[Distributed] Worker WS unreachable: {url}")
|
||||
return False
|
||||
except Exception as e:
|
||||
debug_log(f"[Distributed] Worker WS probe unexpected error: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def _probe_worker_active(worker, use_websocket, semaphore):
|
||||
async with semaphore:
|
||||
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, timeout=60.0):
|
||||
"""Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
|
||||
request_id = uuid.uuid4().hex
|
||||
ws_payload = {
|
||||
"type": "dispatch_prompt",
|
||||
"request_id": request_id,
|
||||
"prompt": payload.get("prompt"),
|
||||
"workflow": payload.get("workflow"),
|
||||
"client_id": client_id,
|
||||
}
|
||||
ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
|
||||
ws_url = f"{ws_url}/distributed/worker_ws"
|
||||
session = await get_client_session()
|
||||
|
||||
async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
|
||||
await ws.send_json(ws_payload)
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(msg.data or "{}")
|
||||
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
|
||||
if data.get("ok"):
|
||||
return
|
||||
error_text = data.get("error") or "Worker rejected websocket dispatch."
|
||||
validation_error = data.get("validation_error")
|
||||
node_errors = data.get("node_errors")
|
||||
if validation_error:
|
||||
error_text = f"{error_text} | validation_error={validation_error}"
|
||||
if node_errors:
|
||||
error_text = f"{error_text} | node_errors={node_errors}"
|
||||
raise RuntimeError(error_text)
|
||||
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
|
||||
raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
|
||||
|
||||
raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
|
||||
|
||||
|
||||
async def dispatch_worker_prompt(
|
||||
worker,
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id=None,
|
||||
use_websocket=False,
|
||||
trace_execution_id=None,
|
||||
):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
worker_url = build_worker_url(worker)
|
||||
url = build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if extra_data:
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
if use_websocket:
|
||||
try:
|
||||
await _dispatch_via_websocket(
|
||||
worker_url,
|
||||
{
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
},
|
||||
client_id,
|
||||
)
|
||||
return
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
raise
|
||||
|
||||
session = await get_client_session()
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
|
||||
|
||||
async def select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
):
|
||||
"""Probe workers and return (active_workers, updated_delegate_master)."""
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
|
||||
if trace_execution_id and workers:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
|
||||
)
|
||||
|
||||
probe_results = await asyncio.gather(
|
||||
*[
|
||||
_probe_worker_active(worker, use_websocket, probe_semaphore)
|
||||
for worker in workers
|
||||
]
|
||||
)
|
||||
|
||||
active_workers = []
|
||||
for worker, is_active in probe_results:
|
||||
if is_active:
|
||||
active_workers.append(worker)
|
||||
else:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
else:
|
||||
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
|
||||
if trace_execution_id and workers:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Worker probe complete: active={len(active_workers)}/{len(workers)}",
|
||||
)
|
||||
|
||||
if not active_workers and delegate_master:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.")
|
||||
else:
|
||||
debug_log("All workers offline while delegate-only requested; enabling master participation.")
|
||||
delegate_master = False
|
||||
|
||||
return active_workers, delegate_master
|
||||
|
||||
|
||||
def _extract_queue_remaining(payload):
|
||||
if not isinstance(payload, dict):
|
||||
return 0
|
||||
try:
|
||||
queue_remaining = int(payload.get("exec_info", {}).get("queue_remaining", 0))
|
||||
except (TypeError, ValueError):
|
||||
queue_remaining = 0
|
||||
return max(queue_remaining, 0)
|
||||
|
||||
|
||||
async def _probe_worker_queue(worker, semaphore, probe_timeout):
|
||||
async with semaphore:
|
||||
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:
|
||||
return None
|
||||
return {
|
||||
"worker": worker,
|
||||
"queue_remaining": _extract_queue_remaining(payload),
|
||||
}
|
||||
|
||||
|
||||
def _select_idle_round_robin(statuses):
|
||||
global _least_busy_rr_index
|
||||
if not statuses:
|
||||
return None
|
||||
index = _least_busy_rr_index % len(statuses)
|
||||
_least_busy_rr_index += 1
|
||||
return statuses[index]
|
||||
|
||||
|
||||
async def select_least_busy_worker(
|
||||
workers,
|
||||
trace_execution_id=None,
|
||||
probe_concurrency=8,
|
||||
probe_timeout=3.0,
|
||||
):
|
||||
"""Select one worker by queue depth, round-robin among idle workers."""
|
||||
if not workers:
|
||||
return None
|
||||
|
||||
probe_limit = parse_positive_int(probe_concurrency, 8)
|
||||
probe_semaphore = asyncio.Semaphore(probe_limit)
|
||||
statuses = await asyncio.gather(
|
||||
*[
|
||||
_probe_worker_queue(worker, probe_semaphore, probe_timeout)
|
||||
for worker in workers
|
||||
]
|
||||
)
|
||||
statuses = [status for status in statuses if status is not None]
|
||||
if not statuses:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
|
||||
else:
|
||||
log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
|
||||
return None
|
||||
|
||||
idle_statuses = [status for status in statuses if status["queue_remaining"] == 0]
|
||||
if idle_statuses:
|
||||
selected = _select_idle_round_robin(idle_statuses)
|
||||
else:
|
||||
selected = min(statuses, key=lambda status: status["queue_remaining"])
|
||||
|
||||
worker = selected["worker"]
|
||||
queue_remaining = selected["queue_remaining"]
|
||||
if trace_execution_id:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
|
||||
)
|
||||
else:
|
||||
debug_log(
|
||||
f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
|
||||
)
|
||||
return worker
|
||||
@@ -1,256 +0,0 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
|
||||
import aiohttp
|
||||
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import build_worker_url, get_client_session
|
||||
from ...utils.trace_logger import trace_debug, trace_info
|
||||
|
||||
|
||||
LIKELY_FILENAME_RE = re.compile(
|
||||
r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|"
|
||||
r"wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma|latent|txt|vae|lora|embedding)"
|
||||
r"(\s*\[\w+\])?$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
MEDIA_FILE_RE = re.compile(
|
||||
r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm|wav|mp3|flac|m4a|aac|ogg|opus|aiff|aif|wma)(\s*\[\w+\])?$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_media_reference(value):
|
||||
"""Normalize one media string value to a path-like reference or None."""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
|
||||
if MEDIA_FILE_RE.search(cleaned):
|
||||
return cleaned
|
||||
return None
|
||||
|
||||
|
||||
def convert_paths_for_platform(obj, target_separator):
|
||||
"""Recursively normalize likely file paths for the worker platform separator."""
|
||||
if target_separator not in ("/", "\\"):
|
||||
return obj
|
||||
|
||||
def _convert(value):
|
||||
if isinstance(value, str):
|
||||
if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value):
|
||||
trimmed = value.strip()
|
||||
has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed))
|
||||
is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\")
|
||||
has_protocol = bool(re.match(r"^\w+://", trimmed))
|
||||
|
||||
# URLs are not local paths and should never be separator-normalized.
|
||||
if has_protocol:
|
||||
return trimmed
|
||||
|
||||
# Keep relative media-style paths in forward-slash form (Comfy-style annotated paths).
|
||||
if not has_drive and not is_absolute and not has_protocol and MEDIA_FILE_RE.search(trimmed):
|
||||
return re.sub(r"[\\]+", "/", trimmed)
|
||||
|
||||
if target_separator == "\\":
|
||||
return re.sub(r"[\\/]+", r"\\", trimmed)
|
||||
return re.sub(r"[\\/]+", "/", trimmed)
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return [_convert(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: _convert(item) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
return _convert(obj)
|
||||
|
||||
|
||||
def _find_media_references(prompt_obj):
|
||||
"""Find media file references in image/video/audio/file inputs used by worker prompts."""
|
||||
media_refs = set()
|
||||
for node in prompt_obj.values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
for key in ("image", "video", "audio", "file"):
|
||||
cleaned = _normalize_media_reference(inputs.get(key))
|
||||
if cleaned:
|
||||
media_refs.add(cleaned)
|
||||
return sorted(media_refs)
|
||||
|
||||
|
||||
def _rewrite_prompt_media_inputs(prompt_obj, worker_media_paths):
|
||||
"""Rewrite media string inputs to worker-local uploaded paths."""
|
||||
if not isinstance(worker_media_paths, dict) or not worker_media_paths:
|
||||
return
|
||||
|
||||
for node in prompt_obj.values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.get("inputs", {})
|
||||
if not isinstance(inputs, dict):
|
||||
continue
|
||||
for key in ("image", "video", "audio", "file"):
|
||||
value = inputs.get(key)
|
||||
cleaned = _normalize_media_reference(value)
|
||||
if not cleaned:
|
||||
continue
|
||||
worker_path = worker_media_paths.get(cleaned)
|
||||
if worker_path:
|
||||
inputs[key] = worker_path
|
||||
|
||||
|
||||
def _load_media_file_sync(filename):
|
||||
"""Load local media bytes and hash for worker upload sync."""
|
||||
import folder_paths
|
||||
|
||||
full_path = folder_paths.get_annotated_filepath(filename)
|
||||
if not os.path.exists(full_path):
|
||||
raise FileNotFoundError(filename)
|
||||
|
||||
with open(full_path, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
file_hash = hashlib.md5(file_bytes).hexdigest()
|
||||
mime_type = mimetypes.guess_type(full_path)[0]
|
||||
if not mime_type:
|
||||
ext = os.path.splitext(full_path)[1].lower()
|
||||
if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}:
|
||||
mime_type = "video/mp4"
|
||||
else:
|
||||
mime_type = "image/png"
|
||||
return file_bytes, file_hash, mime_type
|
||||
|
||||
|
||||
async def fetch_worker_path_separator(worker, trace_execution_id=None):
|
||||
"""Best-effort fetch of a worker's path separator from /distributed/system_info."""
|
||||
url = build_worker_url(worker, "/distributed/system_info")
|
||||
session = await get_client_session()
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
|
||||
if resp.status != 200:
|
||||
return None
|
||||
payload = await resp.json()
|
||||
separator = ((payload or {}).get("platform") or {}).get("path_separator")
|
||||
return separator if separator in ("/", "\\") else None
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}")
|
||||
else:
|
||||
debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}")
|
||||
return None
|
||||
|
||||
|
||||
async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
|
||||
"""Upload one media file to worker iff missing or hash-mismatched."""
|
||||
session = await get_client_session()
|
||||
normalized = filename.replace("\\", "/")
|
||||
|
||||
check_url = build_worker_url(worker, "/distributed/check_file")
|
||||
try:
|
||||
async with session.post(
|
||||
check_url,
|
||||
json={"filename": normalized, "hash": file_hash},
|
||||
timeout=aiohttp.ClientTimeout(total=6),
|
||||
) as resp:
|
||||
if resp.status == 200:
|
||||
payload = await resp.json()
|
||||
if payload.get("exists") and payload.get("hash_matches"):
|
||||
return False, normalized
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
|
||||
else:
|
||||
debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
|
||||
|
||||
parts = normalized.split("/")
|
||||
clean_name = parts[-1]
|
||||
subfolder = "/".join(parts[:-1])
|
||||
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type)
|
||||
form.add_field("type", "input")
|
||||
form.add_field("subfolder", subfolder)
|
||||
form.add_field("overwrite", "true")
|
||||
|
||||
upload_url = build_worker_url(worker, "/upload/image")
|
||||
async with session.post(
|
||||
upload_url,
|
||||
data=form,
|
||||
timeout=aiohttp.ClientTimeout(total=30),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
try:
|
||||
payload = await resp.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
|
||||
name = str((payload or {}).get("name") or clean_name).strip()
|
||||
subfolder = str((payload or {}).get("subfolder") or "").strip().replace("\\", "/").strip("/")
|
||||
worker_path = f"{subfolder}/{name}" if subfolder else name
|
||||
return True, worker_path
|
||||
|
||||
|
||||
async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
|
||||
"""Sync referenced media files from master to a remote worker before dispatch."""
|
||||
media_refs = _find_media_references(prompt_obj)
|
||||
if not media_refs:
|
||||
return
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
uploaded = 0
|
||||
skipped = 0
|
||||
missing = 0
|
||||
worker_media_paths = {}
|
||||
for filename in media_refs:
|
||||
try:
|
||||
file_bytes, file_hash, mime_type = await loop.run_in_executor(
|
||||
None, _load_media_file_sync, filename
|
||||
)
|
||||
except FileNotFoundError:
|
||||
missing += 1
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.")
|
||||
else:
|
||||
log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.")
|
||||
continue
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}")
|
||||
continue
|
||||
|
||||
try:
|
||||
changed, worker_path = await _upload_media_to_worker(
|
||||
worker,
|
||||
filename,
|
||||
file_bytes,
|
||||
file_hash,
|
||||
mime_type,
|
||||
trace_execution_id=trace_execution_id,
|
||||
)
|
||||
if worker_path:
|
||||
worker_media_paths[filename] = worker_path
|
||||
if changed:
|
||||
uploaded += 1
|
||||
else:
|
||||
skipped += 1
|
||||
except Exception as exc:
|
||||
if trace_execution_id:
|
||||
trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
|
||||
else:
|
||||
log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
|
||||
|
||||
_rewrite_prompt_media_inputs(prompt_obj, worker_media_paths)
|
||||
|
||||
summary = (
|
||||
f"Media sync for worker {worker.get('id')}: "
|
||||
f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}"
|
||||
)
|
||||
if trace_execution_id:
|
||||
trace_debug(trace_execution_id, summary)
|
||||
else:
|
||||
debug_log(f"[Distributed] {summary}")
|
||||
@@ -1,579 +0,0 @@
|
||||
import json
|
||||
from collections import deque
|
||||
|
||||
from ...utils.logging import debug_log
|
||||
|
||||
|
||||
class PromptIndex:
|
||||
"""Cache prompt metadata for faster worker/master prompt preparation."""
|
||||
|
||||
def __init__(self, prompt_obj):
|
||||
self._prompt_json = json.dumps(prompt_obj)
|
||||
self.nodes_by_class = {}
|
||||
self.class_by_node = {}
|
||||
self.inputs_by_node = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
class_type = node.get("class_type")
|
||||
node_id_str = str(node_id)
|
||||
if class_type:
|
||||
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
|
||||
self.class_by_node[node_id_str] = class_type
|
||||
self.inputs_by_node[node_id_str] = node.get("inputs", {})
|
||||
self._upstream_cache = {}
|
||||
|
||||
def copy_prompt(self):
|
||||
return json.loads(self._prompt_json)
|
||||
|
||||
def nodes_for_class(self, class_name):
|
||||
return self.nodes_by_class.get(class_name, [])
|
||||
|
||||
def has_upstream(self, start_node_id, target_class):
|
||||
cache_key = (str(start_node_id), target_class)
|
||||
if cache_key in self._upstream_cache:
|
||||
return self._upstream_cache[cache_key]
|
||||
|
||||
visited = set()
|
||||
stack = [str(start_node_id)]
|
||||
while stack:
|
||||
node_id = stack.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
inputs = self.inputs_by_node.get(node_id, {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
upstream_id = str(value[0])
|
||||
if self.class_by_node.get(upstream_id) == target_class:
|
||||
self._upstream_cache[cache_key] = True
|
||||
return True
|
||||
if upstream_id in self.inputs_by_node:
|
||||
stack.append(upstream_id)
|
||||
|
||||
self._upstream_cache[cache_key] = False
|
||||
return False
|
||||
|
||||
|
||||
def _iter_prompt_nodes(prompt_obj):
|
||||
for node_id, node in prompt_obj.items():
|
||||
if isinstance(node, dict):
|
||||
yield str(node_id), node
|
||||
|
||||
|
||||
def find_nodes_by_class(prompt_obj, class_name):
|
||||
nodes = []
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") == class_name:
|
||||
nodes.append(node_id)
|
||||
return nodes
|
||||
|
||||
|
||||
def _find_downstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable downstream from the provided IDs."""
|
||||
adjacency = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
adjacency.setdefault(source_id, set()).add(str(node_id))
|
||||
|
||||
connected = set(start_ids)
|
||||
queue = deque(start_ids)
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
|
||||
if dependent not in connected:
|
||||
connected.add(dependent)
|
||||
queue.append(dependent)
|
||||
return connected
|
||||
|
||||
|
||||
def _create_numeric_id_generator(prompt_obj):
|
||||
"""Return a closure that yields new numeric string IDs."""
|
||||
max_id = 0
|
||||
for node_id in prompt_obj.keys():
|
||||
try:
|
||||
numeric = int(node_id)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
max_id = max(max_id, numeric)
|
||||
|
||||
counter = max_id
|
||||
|
||||
def _next_id():
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
return str(counter)
|
||||
|
||||
return _next_id
|
||||
|
||||
|
||||
def _find_upstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable upstream from start_ids, including start nodes."""
|
||||
connected = set(str(node_id) for node_id in start_ids)
|
||||
queue = deque(connected)
|
||||
while queue:
|
||||
node_id = queue.popleft()
|
||||
node = prompt_obj.get(node_id) or {}
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
if source_id in prompt_obj and source_id not in connected:
|
||||
connected.add(source_id)
|
||||
queue.append(source_id)
|
||||
return connected
|
||||
|
||||
|
||||
_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
|
||||
"PrimitiveBoolean",
|
||||
"PrimitiveFloat",
|
||||
"PrimitiveInt",
|
||||
"PrimitiveNode",
|
||||
"PrimitiveString",
|
||||
}
|
||||
|
||||
_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
|
||||
"LoadImage",
|
||||
}
|
||||
|
||||
_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
|
||||
_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
|
||||
|
||||
# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
|
||||
# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
|
||||
# only after its connected inputs recursively prove safe.
|
||||
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
|
||||
_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
|
||||
"CreateList": ("inputs.",),
|
||||
}
|
||||
|
||||
# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
|
||||
_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
|
||||
|
||||
|
||||
def _get_delegate_master_node_class_mappings():
|
||||
"""Return ComfyUI node-class mappings when available."""
|
||||
if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
|
||||
return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
|
||||
try:
|
||||
import nodes as comfy_nodes # type: ignore
|
||||
except Exception: # pragma: no cover - depends on ComfyUI runtime imports
|
||||
return {}
|
||||
return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
|
||||
|
||||
|
||||
def _get_delegate_master_node_class(class_type):
|
||||
mappings = _get_delegate_master_node_class_mappings()
|
||||
return mappings.get(class_type) if isinstance(mappings, dict) else None
|
||||
|
||||
|
||||
def _normalize_delegate_master_return_type(return_type):
|
||||
if return_type is None:
|
||||
return ""
|
||||
return str(return_type).strip().upper()
|
||||
|
||||
|
||||
def _delegate_master_type_is_safe_scalar(type_name):
|
||||
return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
|
||||
|
||||
|
||||
def _delegate_master_type_is_safe_config(type_name):
|
||||
return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
|
||||
|
||||
|
||||
def _delegate_master_output_is_safe_scalar(class_type, output_index):
|
||||
"""Return True when a registered node output is lightweight config data."""
|
||||
if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
|
||||
return True
|
||||
node_class = _get_delegate_master_node_class(class_type)
|
||||
return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
|
||||
try:
|
||||
output_type = return_types[int(output_index)]
|
||||
except (IndexError, TypeError, ValueError):
|
||||
return False
|
||||
return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
|
||||
|
||||
|
||||
def _get_delegate_master_input_types(class_type):
|
||||
node_class = _get_delegate_master_node_class(class_type)
|
||||
input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
|
||||
if callable(input_types):
|
||||
try:
|
||||
input_types = input_types()
|
||||
except TypeError:
|
||||
return {}
|
||||
return input_types if isinstance(input_types, dict) else {}
|
||||
|
||||
|
||||
def _normalize_delegate_master_input_type(input_spec):
|
||||
if isinstance(input_spec, (list, tuple)) and input_spec:
|
||||
return _normalize_delegate_master_return_type(input_spec[0])
|
||||
return _normalize_delegate_master_return_type(input_spec)
|
||||
|
||||
|
||||
def _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
"""Return True when a registered downstream input expects config data."""
|
||||
for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
|
||||
if input_name.startswith(prefix):
|
||||
return True
|
||||
input_types = _get_delegate_master_input_types(class_type)
|
||||
for section_name in ("required", "optional"):
|
||||
section = input_types.get(section_name, {})
|
||||
if isinstance(section, dict) and input_name in section:
|
||||
input_type = _normalize_delegate_master_input_type(section[input_name])
|
||||
return _delegate_master_type_is_safe_config(input_type)
|
||||
return False
|
||||
|
||||
|
||||
def _is_delegate_master_always_retained_upstream_node(node):
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
class_type = node.get("class_type")
|
||||
return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
|
||||
|
||||
|
||||
def _is_delegate_master_retained_upstream_node(node, output_index=0):
|
||||
"""Return True for lightweight upstream nodes safe to keep on the master."""
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str):
|
||||
return False
|
||||
return (
|
||||
class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
|
||||
or class_type.startswith("Primitive")
|
||||
or _delegate_master_output_is_safe_scalar(class_type, output_index)
|
||||
)
|
||||
|
||||
|
||||
def _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
node_id,
|
||||
output_index,
|
||||
memo,
|
||||
visiting,
|
||||
):
|
||||
"""Return safe retained branch nodes, or None when the branch is not safe."""
|
||||
node_id = str(node_id)
|
||||
cache_key = (node_id, output_index)
|
||||
if cache_key in memo:
|
||||
cached = memo[cache_key]
|
||||
return None if cached is None else set(cached)
|
||||
if cache_key in visiting:
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
|
||||
node = prompt_obj.get(node_id)
|
||||
if not _is_delegate_master_retained_upstream_node(node, output_index):
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
|
||||
visiting.add(cache_key)
|
||||
retained = {node_id}
|
||||
inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
|
||||
class_type = node.get("class_type") if isinstance(node, dict) else None
|
||||
for input_name, value in inputs.items():
|
||||
if not (isinstance(value, list) and len(value) == 2):
|
||||
continue
|
||||
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
source_id = str(value[0])
|
||||
branch = _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
source_id,
|
||||
value[1],
|
||||
memo,
|
||||
visiting,
|
||||
)
|
||||
if branch is None:
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = None
|
||||
return None
|
||||
retained.update(branch)
|
||||
|
||||
visiting.remove(cache_key)
|
||||
memo[cache_key] = frozenset(retained)
|
||||
return retained
|
||||
|
||||
|
||||
def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
|
||||
"""Return lightweight upstream nodes needed by kept delegate-master nodes."""
|
||||
connected = set()
|
||||
memo = {}
|
||||
for node_id in start_ids:
|
||||
node = prompt_obj.get(str(node_id)) or {}
|
||||
inputs = node.get("inputs", {})
|
||||
class_type = node.get("class_type") if isinstance(node, dict) else None
|
||||
for input_name, value in inputs.items():
|
||||
if not (isinstance(value, list) and len(value) == 2):
|
||||
continue
|
||||
source_node = prompt_obj.get(str(value[0]))
|
||||
if _is_delegate_master_always_retained_upstream_node(source_node):
|
||||
connected.add(str(value[0]))
|
||||
continue
|
||||
if not _delegate_master_input_is_safe_scalar(class_type, input_name):
|
||||
continue
|
||||
branch = _collect_delegate_master_retained_upstream_branch(
|
||||
prompt_obj,
|
||||
value[0],
|
||||
value[1],
|
||||
memo,
|
||||
set(),
|
||||
)
|
||||
if branch is not None:
|
||||
connected.update(branch)
|
||||
return connected
|
||||
|
||||
|
||||
def prune_prompt_for_worker(prompt_obj):
|
||||
"""Prune worker prompt to distributed nodes and their upstream dependencies."""
|
||||
collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
|
||||
upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
|
||||
distributed_ids = collector_ids + upscale_ids
|
||||
if not distributed_ids:
|
||||
return prompt_obj
|
||||
|
||||
connected = _find_upstream_nodes(prompt_obj, distributed_ids)
|
||||
pruned_prompt = {}
|
||||
for node_id in connected:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
# Generate IDs from the original prompt so we never reuse IDs from pruned downstream nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for dist_id in distributed_ids:
|
||||
if dist_id not in pruned_prompt:
|
||||
continue
|
||||
downstream = _find_downstream_nodes(prompt_obj, [dist_id])
|
||||
has_removed_downstream = any(node_id != dist_id for node_id in downstream)
|
||||
if has_removed_downstream:
|
||||
original_node = prompt_obj.get(str(dist_id), {})
|
||||
class_type = original_node.get("class_type")
|
||||
inputs = original_node.get("inputs", {})
|
||||
image_connected = class_type != "DistributedCollector" or (
|
||||
isinstance(inputs.get("images"), list)
|
||||
and len(inputs["images"]) == 2
|
||||
)
|
||||
audio_connected = (
|
||||
class_type == "DistributedCollector"
|
||||
and isinstance(inputs.get("audio"), list)
|
||||
and len(inputs["audio"]) == 2
|
||||
)
|
||||
|
||||
if image_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"images": [dist_id, 0]},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {"title": "Preview Image (auto-added)"},
|
||||
}
|
||||
elif audio_connected:
|
||||
preview_id = next_id()
|
||||
pruned_prompt[preview_id] = {
|
||||
"inputs": {"audio": [dist_id, 1]},
|
||||
"class_type": "PreviewAudio",
|
||||
"_meta": {"title": "Preview Audio (auto-added)"},
|
||||
}
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
nodes_to_keep.update(
|
||||
_find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
|
||||
)
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in nodes_to_keep:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
pruned_ids = set(pruned_prompt.keys())
|
||||
for node_id, node in pruned_prompt.items():
|
||||
inputs = node.get("inputs")
|
||||
if not inputs:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
if isinstance(input_value, list) and len(input_value) == 2:
|
||||
source_id = str(input_value[0])
|
||||
if source_id not in pruned_ids:
|
||||
inputs.pop(input_name, None)
|
||||
debug_log(
|
||||
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
||||
)
|
||||
|
||||
# Generate IDs from the original prompt to avoid ID collisions with pruned nodes.
|
||||
next_id = _create_numeric_id_generator(prompt_obj)
|
||||
for collector_id in collector_ids:
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
|
||||
original_images = original_inputs.get("images")
|
||||
if not (isinstance(original_images, list) and len(original_images) == 2):
|
||||
continue
|
||||
placeholder_id = next_id()
|
||||
pruned_prompt[placeholder_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
"inputs": {
|
||||
"height": 64,
|
||||
"width": 64,
|
||||
"channels": 3,
|
||||
},
|
||||
"_meta": {
|
||||
"title": "Distributed Empty Image (auto-added)",
|
||||
},
|
||||
}
|
||||
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
|
||||
debug_log(
|
||||
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
|
||||
)
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
def generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
return job_map
|
||||
|
||||
|
||||
def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedSeed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def _override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
):
|
||||
"""Configure DistributedCollector nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
|
||||
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
||||
node.setdefault("inputs", {})["pass_through"] = True
|
||||
continue
|
||||
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
|
||||
def _override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
):
|
||||
"""Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
|
||||
inputs["is_worker"] = not is_master
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
if is_master:
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
|
||||
|
||||
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
|
||||
"""Configure DistributedValue nodes for master or worker role."""
|
||||
for node_id in prompt_index.nodes_for_class("DistributedValue"):
|
||||
node = prompt_copy.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["worker_id"] = ""
|
||||
else:
|
||||
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
|
||||
|
||||
|
||||
def apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
):
|
||||
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
||||
is_master = participant_id == "master"
|
||||
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
||||
enabled_json = json.dumps(enabled_worker_ids)
|
||||
|
||||
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
|
||||
_override_collector_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
delegate_master,
|
||||
)
|
||||
_override_upscale_nodes(
|
||||
prompt_copy,
|
||||
prompt_index,
|
||||
is_master,
|
||||
participant_id,
|
||||
job_id_map,
|
||||
master_url,
|
||||
enabled_json,
|
||||
)
|
||||
|
||||
return prompt_copy
|
||||
@@ -1,123 +0,0 @@
|
||||
"""Small Runpod queue API client. Never retries a potentially billable POST."""
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
|
||||
import aiohttp
|
||||
|
||||
API_BASE = "https://api.runpod.ai/v2"
|
||||
MAX_REQUEST_BYTES = 10_000_000
|
||||
_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z")
|
||||
|
||||
|
||||
class RunpodError(RuntimeError):
|
||||
"""Sanitized provider failure safe to display in the panel."""
|
||||
|
||||
|
||||
class AmbiguousSubmission(RunpodError):
|
||||
"""Submission may have been accepted: never retry it automatically."""
|
||||
|
||||
|
||||
def serialize_request(workflow, images, timeout_seconds):
|
||||
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
|
||||
payload = {
|
||||
"input": {"workflow": workflow, "images": images},
|
||||
"policy": {
|
||||
"executionTimeout": timeout_seconds * 1000,
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
body = json.dumps(
|
||||
payload,
|
||||
allow_nan=False,
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
size = len(body)
|
||||
if size > MAX_REQUEST_BYTES:
|
||||
raise ValueError(
|
||||
f"Serialized serverless request body is {size:,} bytes; Runpod allows at most "
|
||||
f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw "
|
||||
"image bytes; reduce image or workflow size."
|
||||
)
|
||||
return body
|
||||
|
||||
|
||||
def checked_id(value):
|
||||
if not isinstance(value, str) or not _ID.fullmatch(value):
|
||||
raise ValueError("Invalid Runpod endpoint or job ID")
|
||||
return value
|
||||
|
||||
|
||||
class RunpodClient:
|
||||
def __init__(self, session, endpoint_id, api_key):
|
||||
self.session = session
|
||||
self.endpoint_id = checked_id(endpoint_id)
|
||||
if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"):
|
||||
raise ValueError("A valid Runpod API key is required")
|
||||
self._api_key = api_key
|
||||
|
||||
async def _request(self, method, path, payload=None, request_body=None):
|
||||
# Base origin is code-owned, never a user-supplied URL. Do not follow redirects
|
||||
# with credentials, or include response bodies/transport errors in exceptions.
|
||||
url = f"{API_BASE}/{self.endpoint_id}/{path}"
|
||||
attempts = 3 if method == "GET" else 1
|
||||
for attempt in range(attempts):
|
||||
try:
|
||||
request_kwargs = {
|
||||
"headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"},
|
||||
"timeout": aiohttp.ClientTimeout(total=30),
|
||||
"allow_redirects": False,
|
||||
}
|
||||
if request_body is None:
|
||||
request_kwargs["json"] = payload
|
||||
else:
|
||||
request_kwargs["data"] = request_body
|
||||
async with self.session.request(
|
||||
method, url,
|
||||
**request_kwargs,
|
||||
) as response:
|
||||
status = response.status
|
||||
if status != 200:
|
||||
if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts:
|
||||
await asyncio.sleep(0.5 * (2 ** attempt))
|
||||
continue
|
||||
if method == "POST" and path == "run" and status >= 500:
|
||||
raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.")
|
||||
raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})")
|
||||
try:
|
||||
result = await response.json()
|
||||
except (ValueError, aiohttp.ContentTypeError):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None
|
||||
raise RunpodError("Runpod returned an invalid JSON response") from None
|
||||
if not isinstance(result, dict):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.")
|
||||
raise RunpodError("Runpod returned an invalid response")
|
||||
return result
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError, OSError):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None
|
||||
if method == "GET" and attempt + 1 < attempts:
|
||||
await asyncio.sleep(0.5 * (2 ** attempt))
|
||||
continue
|
||||
raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None
|
||||
raise RunpodError("Runpod request failed")
|
||||
|
||||
async def health(self):
|
||||
return await self._request("GET", "health")
|
||||
|
||||
async def submit(self, workflow, images, timeout_seconds):
|
||||
body = serialize_request(workflow, images, timeout_seconds)
|
||||
result = await self._request("POST", "run", request_body=body)
|
||||
try:
|
||||
return checked_id(result.get("id"))
|
||||
except ValueError:
|
||||
raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None
|
||||
|
||||
async def status(self, job_id):
|
||||
return await self._request("GET", f"status/{checked_id(job_id)}")
|
||||
|
||||
async def cancel(self, job_id):
|
||||
return await self._request("POST", f"cancel/{checked_id(job_id)}")
|
||||
@@ -1,57 +0,0 @@
|
||||
"""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])
|
||||
@@ -1,601 +0,0 @@
|
||||
"""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)
|
||||
@@ -1,420 +0,0 @@
|
||||
"""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
|
||||
@@ -1,375 +0,0 @@
|
||||
"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
|
||||
from .serverless_ledger import JobLedger
|
||||
from .serverless_state import (
|
||||
PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
|
||||
occupies_capacity, terminal_provider_status,
|
||||
)
|
||||
|
||||
|
||||
async def decode_results(output, manifest):
|
||||
from .serverless_media import collect_images
|
||||
import torch
|
||||
images = await collect_images(output, manifest)
|
||||
return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
|
||||
|
||||
|
||||
def state_directory():
|
||||
import folder_paths
|
||||
return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
|
||||
|
||||
|
||||
def get_secret_store():
|
||||
from ...utils.serverless_config import SecretStore
|
||||
return SecretStore(state_directory() / "credentials.json")
|
||||
|
||||
|
||||
def get_manager():
|
||||
import server
|
||||
from ...utils.network import get_client_session
|
||||
ps = server.PromptServer.instance
|
||||
manager = getattr(ps, "distributed_serverless", None)
|
||||
if manager is None:
|
||||
manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
|
||||
ps.distributed_serverless = manager
|
||||
return manager
|
||||
|
||||
|
||||
class ServerlessManager:
|
||||
def __init__(self, prompt_server, secrets, ledger_path, session_factory):
|
||||
self.server = prompt_server
|
||||
self.secrets = secrets
|
||||
self.ledger = JobLedger(ledger_path)
|
||||
self.session_factory = session_factory
|
||||
self.tasks = {}
|
||||
self.runs = {}
|
||||
self.pool_limits = {}
|
||||
self.health_cache = {}
|
||||
self.health_locks = {}
|
||||
self.poll_seconds = 2.0
|
||||
self.capacity_changed = asyncio.Condition()
|
||||
self.recovery_lock = asyncio.Lock()
|
||||
self._recover_local_records()
|
||||
|
||||
def _recover_local_records(self):
|
||||
# WAITING is committed before admission; it cannot have reached Runpod.
|
||||
# SUBMITTING is an intent record, so a lost ID must remain reserved.
|
||||
for record in self.ledger.all():
|
||||
if record.get("status") == "WAITING" and not record.get("job_id"):
|
||||
self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time())
|
||||
elif record.get("status") == "SUBMITTING" and not record.get("job_id"):
|
||||
self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time())
|
||||
|
||||
def _task_is_active(self, record_id):
|
||||
task = self.tasks.get(record_id)
|
||||
return task is not None and not task.done()
|
||||
|
||||
def pool_is_active(self, worker_id):
|
||||
"""Protect both live local work and unresolved pre-restart provider work."""
|
||||
for rid, run in self.runs.items():
|
||||
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
|
||||
if str(pool_id) == str(worker_id) and self._task_is_active(rid):
|
||||
return True
|
||||
return any(belongs_to_pool(record, worker_id) and is_unresolved(record)
|
||||
for record in self.ledger.all())
|
||||
|
||||
async def _reserve_provider_capacity(self, run):
|
||||
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
|
||||
async with self.capacity_changed:
|
||||
while True:
|
||||
self._check_master(run)
|
||||
outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record)
|
||||
for record in self.ledger.all())
|
||||
if outstanding < run["worker"]["concurrency"]:
|
||||
# Claim atomically before any await that could submit a billed job.
|
||||
self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time())
|
||||
return
|
||||
await self.capacity_changed.wait()
|
||||
|
||||
async def _notify_capacity(self):
|
||||
async with self.capacity_changed:
|
||||
self.capacity_changed.notify_all()
|
||||
|
||||
async def _client(self, worker):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
key = self.secrets.get(pool_id)
|
||||
return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
|
||||
|
||||
async def health(self, worker, *, fresh=False):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
cache_key = (pool_id, worker["endpoint_id"])
|
||||
lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
|
||||
async with lock:
|
||||
cached = self.health_cache.get(cache_key)
|
||||
if cached and not fresh and time.monotonic() - cached[0] < 10:
|
||||
return cached[1]
|
||||
client = await self._client(worker)
|
||||
data = await client.health()
|
||||
jobs = data.get("jobs", {})
|
||||
if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
|
||||
raise RunpodError("Runpod returned an invalid endpoint health record")
|
||||
try:
|
||||
pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
|
||||
except (ValueError, TypeError):
|
||||
raise RunpodError("Runpod returned invalid queue counts") from None
|
||||
result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
|
||||
self.health_cache[cache_key] = (time.monotonic(), result)
|
||||
return result
|
||||
|
||||
def _semaphore(self, worker):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
limit = worker["concurrency"]
|
||||
current = self.pool_limits.get(pool_id)
|
||||
if current is None:
|
||||
current = (limit, asyncio.Semaphore(limit))
|
||||
self.pool_limits[pool_id] = current
|
||||
elif current[0] != limit:
|
||||
if self.pool_is_active(pool_id):
|
||||
raise ValueError("Wait for this serverless pool's current jobs before changing concurrency")
|
||||
current = (limit, asyncio.Semaphore(limit))
|
||||
self.pool_limits[pool_id] = current
|
||||
return current[1]
|
||||
|
||||
def ensure_capacity(self, count):
|
||||
if sum(not task.done() for task in self.tasks.values()) + count > 64:
|
||||
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
|
||||
|
||||
async def start(self, execution_id, prompt_id, prepared_jobs):
|
||||
self.ensure_capacity(len(prepared_jobs))
|
||||
new_runs = []
|
||||
# Journal every local slot before starting any billed submission.
|
||||
for prepared in prepared_jobs:
|
||||
worker = prepared["worker"]
|
||||
semaphore = self._semaphore(worker)
|
||||
record_id = uuid.uuid4().hex
|
||||
now = time.time()
|
||||
self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
|
||||
worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
|
||||
endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
|
||||
run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
|
||||
provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
|
||||
self.runs[record_id] = run
|
||||
new_runs.append((record_id, run, semaphore))
|
||||
for record_id, run, semaphore in new_runs:
|
||||
self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
|
||||
# Bound in-memory completed tracking as well as the on-disk journal.
|
||||
done = [rid for rid, task in self.tasks.items() if task.done()]
|
||||
for rid in done[:-200]:
|
||||
self.tasks.pop(rid, None)
|
||||
self.runs.pop(rid, None)
|
||||
|
||||
def is_pending(self, collector_job_id, participant_id):
|
||||
return any(not task.done() and run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"])
|
||||
for rid, task in self.tasks.items() for run in [self.runs[rid]])
|
||||
|
||||
def owns(self, collector_job_id, participant_id):
|
||||
return any(run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
|
||||
|
||||
async def _deliver(self, run, images=None, error=None):
|
||||
async with self.server.distributed_jobs_lock:
|
||||
for binding in run["manifest"]:
|
||||
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if queue is None:
|
||||
continue # A cancelled/finished collector must never be resurrected.
|
||||
if error:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
|
||||
else:
|
||||
batch = images[binding["job_id"]]
|
||||
for index, tensor in enumerate(batch):
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
|
||||
"tensor": tensor, "is_last": index == len(batch) - 1})
|
||||
|
||||
async def _cancel_provider(self, run):
|
||||
if run["provider_job_id"] and run["client"]:
|
||||
try:
|
||||
result = await run["client"].cancel(run["provider_job_id"])
|
||||
return terminal_provider_status(result)
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def _check_master(self, run):
|
||||
queue = getattr(self.server, "prompt_queue", None)
|
||||
if queue is None:
|
||||
return
|
||||
history = queue.get_history(run["prompt_id"])
|
||||
record = history.get(run["prompt_id"], {})
|
||||
if record.get("status", {}).get("status_str") == "error":
|
||||
raise RunpodError("Master workflow failed; serverless work stopped")
|
||||
if not record and hasattr(queue, "get_current_queue"):
|
||||
running, pending = queue.get_current_queue()
|
||||
if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
|
||||
raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
|
||||
|
||||
async def _run(self, run, semaphore):
|
||||
rid = run["record_id"]
|
||||
submitted = False
|
||||
provider_terminal = False
|
||||
try:
|
||||
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
|
||||
async with semaphore:
|
||||
await self._reserve_provider_capacity(run)
|
||||
client = await self._client(run["worker"])
|
||||
run["client"] = client
|
||||
submitted = True
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
run["provider_job_id"] = job_id
|
||||
# Persist immediately: all subsequent actions use this exact ID.
|
||||
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
|
||||
while True:
|
||||
self._check_master(run)
|
||||
response = await client.status(job_id)
|
||||
status = response.get("status")
|
||||
if not isinstance(status, str) or status not in PROVIDER_STATES:
|
||||
raise RunpodError("Runpod returned an unknown job state")
|
||||
provider_terminal = status in PROVIDER_TERMINAL
|
||||
if status == "COMPLETED":
|
||||
output = response.get("output")
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
|
||||
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
|
||||
return
|
||||
self.ledger.update(rid, status=status, updated_at=time.time())
|
||||
if provider_terminal:
|
||||
raise RunpodError(f"Serverless job ended with {status}")
|
||||
await asyncio.sleep(self.poll_seconds)
|
||||
except AmbiguousSubmission as exc:
|
||||
self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time())
|
||||
await self._deliver(run, error=str(exc))
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError) as exc:
|
||||
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
|
||||
if submitted and not run["provider_job_id"]:
|
||||
status = "UNKNOWN"
|
||||
error = "Submission interrupted before its job ID was received; check Runpod before retrying"
|
||||
elif isinstance(exc, asyncio.TimeoutError):
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
|
||||
error = "Serverless job exceeded its queue/execution deadline"
|
||||
else:
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED")
|
||||
error = "Serverless work stopped" if terminal or not submitted else "Serverless cancellation requested; provider confirmation is pending"
|
||||
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
|
||||
await self._deliver(run, error=error)
|
||||
except Exception as exc:
|
||||
# If polling/decoding failed, stop only our own provider job, never the pool.
|
||||
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
|
||||
error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract"
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED")
|
||||
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
|
||||
await self._deliver(run, error=error)
|
||||
finally:
|
||||
# Completed tracking retains IDs/bindings only, not potentially large
|
||||
# base64 inputs, workflow metadata, or a client holding credentials.
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run["client"] = None
|
||||
# A released local semaphore alone cannot release a durable reservation.
|
||||
await self._notify_capacity()
|
||||
|
||||
async def wait_for_execution(self, execution_id):
|
||||
tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id]
|
||||
await asyncio.gather(*tasks, return_exceptions=False)
|
||||
|
||||
def _detached_records(self, worker_id=None, execution_id=None):
|
||||
return [record for record in self.ledger.all()
|
||||
if is_unresolved(record) and not self._task_is_active(record["record_id"])
|
||||
and (worker_id is None or belongs_to_pool(record, worker_id))
|
||||
and (execution_id is None or record.get("execution_id") == execution_id)]
|
||||
|
||||
async def _recover_provider_record(self, record, *, cancel=False):
|
||||
"""Reconcile an exact saved ID, never submit or resume a workflow."""
|
||||
if not record.get("job_id"):
|
||||
return False
|
||||
status = None
|
||||
try:
|
||||
client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]})
|
||||
response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"]))
|
||||
status = terminal_provider_status(response)
|
||||
except Exception:
|
||||
pass # Unreachable/expired is not proof of provider termination.
|
||||
if status is not None or cancel:
|
||||
self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time())
|
||||
return status is not None
|
||||
|
||||
async def reconcile(self, worker_id=None):
|
||||
"""Read terminal states for detached jobs and release only proven capacity."""
|
||||
resolved = 0
|
||||
async with self.recovery_lock:
|
||||
for record in self._detached_records(worker_id):
|
||||
resolved += await self._recover_provider_record(record)
|
||||
await self._notify_capacity()
|
||||
return resolved
|
||||
|
||||
async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False):
|
||||
"""Explicit operator escape hatch for lost IDs/expired provider records.
|
||||
|
||||
This records a human's verification, not a cancellation or provider
|
||||
status. It is never called automatically or allowed for live tasks.
|
||||
"""
|
||||
if provider_checked is not True:
|
||||
raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it")
|
||||
async with self.recovery_lock:
|
||||
record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None)
|
||||
if record is None:
|
||||
raise ValueError("Unresolved inactive job not found for this pool")
|
||||
self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(),
|
||||
error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed")
|
||||
await self._notify_capacity()
|
||||
|
||||
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
|
||||
targets = {}
|
||||
for rid, task in self.tasks.items():
|
||||
run = self.runs[rid]
|
||||
if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
|
||||
continue
|
||||
if execution_id and run["execution_id"] != execution_id:
|
||||
continue
|
||||
if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
|
||||
continue
|
||||
if not task.done():
|
||||
# Repeated requests must not interrupt the first cancellation's cleanup.
|
||||
if not task.cancelling():
|
||||
task.cancel()
|
||||
targets[rid] = task
|
||||
if targets:
|
||||
await asyncio.gather(*targets.values(), return_exceptions=True)
|
||||
recovered = 0
|
||||
async with self.recovery_lock:
|
||||
for record in self._detached_records(worker_id, execution_id):
|
||||
rid = record["record_id"]
|
||||
# A task cancelled before its first step never enters _run's handlers.
|
||||
# A failed start() may also leave a journalled but unstarted local slot.
|
||||
if record.get("status") == "WAITING" and not record.get("job_id"):
|
||||
run = self.runs.get(rid)
|
||||
if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])):
|
||||
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
await self._deliver(run, error="Serverless job cancelled before submission")
|
||||
continue
|
||||
# Do not send two cancellation requests for the same live job in
|
||||
# one action. Its unresolved reservation survives for later recovery.
|
||||
if rid in targets or collector_job_id or not record.get("job_id"):
|
||||
continue
|
||||
await self._recover_provider_record(record, cancel=True)
|
||||
recovered += 1
|
||||
await self._notify_capacity()
|
||||
return len(targets) + recovered
|
||||
|
||||
def public_jobs(self, worker_id):
|
||||
records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)]
|
||||
# Never hide unresolved IDs behind the bounded completed history.
|
||||
visible = [record for record in records if is_unresolved(record)]
|
||||
visible.extend(record for record in records if not is_unresolved(record))
|
||||
unresolved_count = sum(is_unresolved(record) for record in records)
|
||||
return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")}
|
||||
for record in visible[:max(20, unresolved_count)]]
|
||||
@@ -1,29 +0,0 @@
|
||||
"""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)
|
||||
@@ -1,471 +0,0 @@
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
from ..utils.config import load_config
|
||||
from ..utils.constants import (
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_master_url, build_master_callback_url
|
||||
from ..utils.trace_logger import trace_debug
|
||||
from .schemas import parse_positive_float, parse_positive_int
|
||||
from .orchestration.dispatch import (
|
||||
dispatch_worker_prompt,
|
||||
select_active_workers,
|
||||
select_least_busy_worker,
|
||||
)
|
||||
from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
|
||||
from .orchestration.prompt_transform import (
|
||||
PromptIndex,
|
||||
apply_participant_overrides,
|
||||
find_nodes_by_class,
|
||||
generate_job_id_map,
|
||||
prepare_delegate_master_prompt,
|
||||
prune_prompt_for_worker,
|
||||
)
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def _generate_execution_trace_id():
|
||||
return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
|
||||
|
||||
def ensure_distributed_state(server_instance=None):
|
||||
"""Ensure prompt_server has the state used by distributed queue orchestration."""
|
||||
ps = server_instance or prompt_server
|
||||
if not hasattr(ps, "distributed_pending_jobs"):
|
||||
ps.distributed_pending_jobs = {}
|
||||
if not hasattr(ps, "distributed_jobs_lock"):
|
||||
ps.distributed_jobs_lock = asyncio.Lock()
|
||||
|
||||
|
||||
# Initialize top-level distributed queue state at module import time.
|
||||
ensure_distributed_state()
|
||||
|
||||
|
||||
async def _ensure_distributed_queue(job_id):
|
||||
"""Ensure a queue exists for the given distributed job ID."""
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
|
||||
|
||||
def _resolve_enabled_workers(config, requested_ids=None):
|
||||
"""Return a list of worker configs that should participate."""
|
||||
workers = []
|
||||
for worker in config.get("workers", []):
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id:
|
||||
continue
|
||||
|
||||
if requested_ids is not None:
|
||||
if worker_id not in requested_ids:
|
||||
continue
|
||||
elif not worker.get("enabled", False):
|
||||
continue
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
from ..utils.serverless_config import validate_serverless_worker
|
||||
pool = validate_serverless_worker(worker)
|
||||
for slot in range(pool["jobs_per_run"]):
|
||||
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
|
||||
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
|
||||
continue
|
||||
|
||||
raw_port = worker.get("port", worker.get("listen_port", 8188))
|
||||
try:
|
||||
port = int(raw_port or 8188)
|
||||
except (TypeError, ValueError):
|
||||
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
|
||||
port = 8188
|
||||
|
||||
workers.append(
|
||||
{
|
||||
"id": worker_id,
|
||||
"name": worker.get("name", worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
}
|
||||
)
|
||||
return workers
|
||||
|
||||
|
||||
def _resolve_orchestration_limits(config):
|
||||
"""Resolve bounded concurrency/timeouts for worker preparation pipeline."""
|
||||
settings = (config or {}).get("settings", {}) or {}
|
||||
worker_probe_concurrency = parse_positive_int(
|
||||
settings.get("worker_probe_concurrency"),
|
||||
ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
|
||||
)
|
||||
worker_prep_concurrency = parse_positive_int(
|
||||
settings.get("worker_prep_concurrency"),
|
||||
ORCHESTRATION_WORKER_PREP_CONCURRENCY,
|
||||
)
|
||||
media_sync_concurrency = parse_positive_int(
|
||||
settings.get("media_sync_concurrency"),
|
||||
ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
|
||||
)
|
||||
media_sync_timeout_seconds = parse_positive_float(
|
||||
settings.get("media_sync_timeout_seconds"),
|
||||
ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
|
||||
)
|
||||
return (
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
)
|
||||
|
||||
|
||||
def _is_load_balance_enabled(value):
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (int, float)):
|
||||
return bool(value)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return False
|
||||
|
||||
|
||||
def _prompt_requests_load_balance(prompt_index):
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
inputs = prompt_index.inputs_by_node.get(node_id, {})
|
||||
if _is_load_balance_enabled(inputs.get("load_balance", False)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def _prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
trace_execution_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
):
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
if worker.get("type") == "serverless":
|
||||
from .orchestration.serverless_prompt import compile_serverless_prompt
|
||||
from .orchestration.serverless_media import prepare_images
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
worker_prompt = apply_participant_overrides(
|
||||
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
|
||||
job_id_map, "", delegate_master, prompt_index,
|
||||
)
|
||||
workflow, manifest = compile_serverless_prompt(worker_prompt)
|
||||
images = await prepare_images(workflow)
|
||||
# Fail during preparation, before regular workers or the master are
|
||||
# dispatched. submit() uses this same serializer for the sent bytes.
|
||||
serialize_request(workflow, images, worker["timeout_seconds"])
|
||||
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images}
|
||||
|
||||
worker_master_url = build_master_callback_url(
|
||||
worker,
|
||||
config=config,
|
||||
prompt_server_instance=prompt_server,
|
||||
)
|
||||
|
||||
worker_type = str(worker.get("type") or "local").strip().lower()
|
||||
is_remote_like = bool(worker.get("host")) and worker_type != "local"
|
||||
if is_remote_like:
|
||||
path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
|
||||
if path_separator:
|
||||
worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
|
||||
|
||||
worker_prompt = prune_prompt_for_worker(worker_prompt)
|
||||
worker_prompt = apply_participant_overrides(
|
||||
worker_prompt,
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
worker_master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if is_remote_like:
|
||||
async with media_sync_semaphore:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
sync_worker_media(worker, worker_prompt, trace_execution_id=trace_execution_id),
|
||||
timeout=media_sync_timeout_seconds,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
trace_debug(
|
||||
trace_execution_id,
|
||||
(
|
||||
f"Media sync timed out after {media_sync_timeout_seconds:.1f}s "
|
||||
f"for worker {worker.get('name')} ({worker.get('id')}); continuing dispatch."
|
||||
),
|
||||
)
|
||||
|
||||
return worker, worker_prompt
|
||||
|
||||
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=None,
|
||||
delegate_master=None,
|
||||
trace_execution_id=None,
|
||||
):
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
execution_trace_id = trace_execution_id or _generate_execution_trace_id()
|
||||
|
||||
config = load_config()
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
|
||||
(
|
||||
worker_probe_concurrency,
|
||||
worker_prep_concurrency,
|
||||
media_sync_concurrency,
|
||||
media_sync_timeout_seconds,
|
||||
) = _resolve_orchestration_limits(config)
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
load_balance_requested = _prompt_requests_load_balance(prompt_index)
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
(
|
||||
f"Orchestration start: requested_workers={len(workers)}, "
|
||||
f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
|
||||
f"websocket={use_websocket}, "
|
||||
f"probe_concurrency={worker_probe_concurrency}, "
|
||||
f"prep_concurrency={worker_prep_concurrency}, "
|
||||
f"media_sync_concurrency={media_sync_concurrency}, "
|
||||
f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
|
||||
f"load_balance={load_balance_requested}"
|
||||
),
|
||||
)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Delegate-only requested but no workers are enabled. Falling back to master execution.",
|
||||
)
|
||||
delegate_master = False
|
||||
|
||||
active_workers, delegate_master = await select_active_workers(
|
||||
workers,
|
||||
use_websocket,
|
||||
delegate_master,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
|
||||
if load_balance_requested:
|
||||
candidate_workers = list(active_workers)
|
||||
if not delegate_master:
|
||||
# Include master in load balancing only when master participation is enabled.
|
||||
candidate_workers.append(
|
||||
{
|
||||
"id": "master",
|
||||
"name": "Master",
|
||||
"host": master_url,
|
||||
"type": "local",
|
||||
}
|
||||
)
|
||||
|
||||
selected_worker = None
|
||||
if candidate_workers:
|
||||
selected_worker = await select_least_busy_worker(
|
||||
candidate_workers,
|
||||
trace_execution_id=execution_trace_id,
|
||||
probe_concurrency=worker_probe_concurrency,
|
||||
)
|
||||
if selected_worker is None and candidate_workers:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance selection probe failed; using first available candidate.",
|
||||
)
|
||||
selected_worker = candidate_workers[0]
|
||||
|
||||
if selected_worker is not None and str(selected_worker.get("id")) == "master":
|
||||
# Master selected as least busy; run master workload only.
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance selected master for execution (workers skipped).",
|
||||
)
|
||||
elif selected_worker is not None:
|
||||
active_workers = [selected_worker]
|
||||
# Worker selected as least busy; keep master orchestrator-only for this run.
|
||||
delegate_master = True
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
|
||||
)
|
||||
else:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Load-balance requested but no execution candidates were available.",
|
||||
)
|
||||
active_workers = []
|
||||
delegate_master = False
|
||||
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
|
||||
queue_result = await queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
return (
|
||||
queue_result["prompt_id"],
|
||||
queue_result["number"],
|
||||
0,
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = apply_participant_overrides(
|
||||
master_prompt,
|
||||
"master",
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
|
||||
)
|
||||
elif not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
|
||||
)
|
||||
else:
|
||||
master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
if active_workers:
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
"Active distributed workers: "
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
|
||||
)
|
||||
worker_payloads = []
|
||||
if active_workers:
|
||||
worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
|
||||
media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
|
||||
worker_payloads = await asyncio.gather(
|
||||
*[
|
||||
_prepare_worker_payload(
|
||||
worker,
|
||||
prompt_index,
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
config,
|
||||
delegate_master,
|
||||
execution_trace_id,
|
||||
worker_prep_semaphore,
|
||||
media_sync_semaphore,
|
||||
media_sync_timeout_seconds,
|
||||
)
|
||||
for worker in active_workers
|
||||
]
|
||||
)
|
||||
|
||||
serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"]
|
||||
regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"]
|
||||
manager = None
|
||||
if serverless_jobs:
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
manager = get_manager()
|
||||
manager.ensure_capacity(len(serverless_jobs))
|
||||
|
||||
# Do not leave collector queues behind when offline compilation/media validation fails.
|
||||
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,
|
||||
)
|
||||
except Exception:
|
||||
if serverless_jobs:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
for job_id in job_id_map.values():
|
||||
prompt_server.distributed_pending_jobs.pop(job_id, None)
|
||||
raise
|
||||
prompt_id = queue_result["prompt_id"]
|
||||
if manager is not None:
|
||||
try:
|
||||
await manager.start(execution_trace_id, prompt_id, serverless_jobs)
|
||||
except Exception:
|
||||
# Wake queued collectors with an explicit error if journal/capacity setup failed.
|
||||
await manager.cancel(execution_id=execution_trace_id)
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
for prepared in serverless_jobs:
|
||||
for binding in prepared["manifest"]:
|
||||
q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if q is not None:
|
||||
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
|
||||
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
|
||||
raise
|
||||
prompt_number = queue_result["number"]
|
||||
node_errors = queue_result.get("node_errors", {})
|
||||
trace_debug(
|
||||
execution_trace_id,
|
||||
f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
|
||||
)
|
||||
return prompt_id, prompt_number, len(worker_payloads), node_errors
|
||||
@@ -1,79 +0,0 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueueRequestPayload:
|
||||
prompt: Dict[str, Any]
|
||||
workflow_meta: Any
|
||||
client_id: str
|
||||
delegate_master: Optional[bool]
|
||||
enabled_worker_ids: List[str]
|
||||
auto_prepare: bool
|
||||
trace_execution_id: Optional[str]
|
||||
|
||||
|
||||
def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
|
||||
"""Parse and validate /distributed/queue payload into a normalized shape."""
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object body")
|
||||
|
||||
auto_prepare_raw = data.get("auto_prepare", True)
|
||||
if not isinstance(auto_prepare_raw, bool):
|
||||
raise ValueError("auto_prepare must be a boolean when provided")
|
||||
auto_prepare = auto_prepare_raw
|
||||
|
||||
prompt = data.get("prompt")
|
||||
# Auto-prepare is always on server-side; keep the field for wire compatibility.
|
||||
if prompt is None:
|
||||
workflow_payload = data.get("workflow")
|
||||
if isinstance(workflow_payload, dict):
|
||||
candidate_prompt = workflow_payload.get("prompt")
|
||||
if isinstance(candidate_prompt, dict):
|
||||
prompt = candidate_prompt
|
||||
|
||||
if not isinstance(prompt, dict):
|
||||
raise ValueError("Field 'prompt' must be an object")
|
||||
|
||||
enabled_ids_raw = data.get("enabled_worker_ids")
|
||||
workers_field = data.get("workers")
|
||||
if enabled_ids_raw is None and workers_field is not None:
|
||||
if not isinstance(workers_field, list):
|
||||
raise ValueError("Field 'workers' must be a list when provided")
|
||||
enabled_ids_raw = []
|
||||
for entry in workers_field:
|
||||
worker_id = entry.get("id") if isinstance(entry, dict) else entry
|
||||
if worker_id is not None:
|
||||
enabled_ids_raw.append(str(worker_id))
|
||||
|
||||
if enabled_ids_raw is None:
|
||||
raise ValueError("enabled_worker_ids required")
|
||||
else:
|
||||
if not isinstance(enabled_ids_raw, list):
|
||||
raise ValueError("enabled_worker_ids must be a list of worker IDs")
|
||||
enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
|
||||
|
||||
delegate_master = data.get("delegate_master")
|
||||
if delegate_master is not None and not isinstance(delegate_master, bool):
|
||||
raise ValueError("delegate_master must be a boolean when provided")
|
||||
|
||||
client_id = data.get("client_id")
|
||||
if not isinstance(client_id, str) or not client_id.strip():
|
||||
raise ValueError("client_id required")
|
||||
client_id = client_id.strip()
|
||||
|
||||
trace_execution_id = data.get("trace_execution_id")
|
||||
if trace_execution_id is not None:
|
||||
if not isinstance(trace_execution_id, str):
|
||||
raise ValueError("trace_execution_id must be a string when provided")
|
||||
trace_execution_id = trace_execution_id.strip() or None
|
||||
|
||||
return QueueRequestPayload(
|
||||
prompt=prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=client_id,
|
||||
delegate_master=delegate_master,
|
||||
enabled_worker_ids=enabled_ids,
|
||||
auto_prepare=auto_prepare,
|
||||
trace_execution_id=trace_execution_id,
|
||||
)
|
||||
@@ -1,54 +0,0 @@
|
||||
def require_fields(data: dict, *fields) -> list[str]:
|
||||
"""Return field names that are missing or empty in a JSON object."""
|
||||
if not isinstance(data, dict):
|
||||
return list(fields)
|
||||
|
||||
missing = []
|
||||
for field in fields:
|
||||
if field not in data:
|
||||
missing.append(field)
|
||||
continue
|
||||
value = data.get(field)
|
||||
if value is None:
|
||||
missing.append(field)
|
||||
continue
|
||||
if isinstance(value, str) and not value.strip():
|
||||
missing.append(field)
|
||||
|
||||
return missing
|
||||
|
||||
|
||||
def validate_worker_id(worker_id: str, config: dict) -> bool:
|
||||
"""Return True when worker_id exists in config['workers']."""
|
||||
worker_id_str = str(worker_id)
|
||||
workers = (config or {}).get("workers", [])
|
||||
return any(str(worker.get("id")) == worker_id_str for worker in workers)
|
||||
|
||||
|
||||
def validate_positive_int(value, field_name: str) -> str | None:
|
||||
"""Validate positive integers and return an error string when invalid."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
if parsed <= 0:
|
||||
return f"Field '{field_name}' must be a positive integer."
|
||||
return None
|
||||
|
||||
|
||||
def parse_positive_int(value, default: int) -> int:
|
||||
"""Parse value as positive int, returning default on failure."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return max(1, int(default))
|
||||
return max(1, parsed)
|
||||
|
||||
|
||||
def parse_positive_float(value, default: float) -> float:
|
||||
"""Parse value as positive float, returning default on failure."""
|
||||
try:
|
||||
parsed = float(value)
|
||||
except (TypeError, ValueError):
|
||||
return max(0.0, float(default))
|
||||
return max(0.0, parsed)
|
||||
@@ -1,98 +0,0 @@
|
||||
"""Master-only Runpod pool monitoring. Health checks never submit inference."""
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
from .orchestration.runpod_client import RunpodError
|
||||
|
||||
|
||||
def _pool(worker_id):
|
||||
return next((w for w in load_config().get("workers", [])
|
||||
if w.get("id") == worker_id and w.get("type") == "serverless"), None)
|
||||
|
||||
|
||||
async def _status(worker_id, fresh=False):
|
||||
worker = _pool(worker_id)
|
||||
if worker is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
jobs = []
|
||||
try:
|
||||
manager = get_manager()
|
||||
jobs = manager.public_jobs(worker_id)
|
||||
result = await manager.health(worker, fresh=fresh)
|
||||
return web.json_response(dict(result, jobs=jobs))
|
||||
except Exception as exc:
|
||||
message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
|
||||
return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/test")
|
||||
async def test_serverless_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
||||
return await _status(worker_id, fresh=True)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
|
||||
async def serverless_status_endpoint(request):
|
||||
return await _status(request.match_info["worker_id"])
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
|
||||
async def cancel_serverless_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
||||
if _pool(worker_id) is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
try:
|
||||
manager = get_manager()
|
||||
count = await manager.cancel(worker_id=worker_id)
|
||||
return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
|
||||
"message": "Cancellation requested for this master's jobs only"})
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
|
||||
|
||||
|
||||
async def _recover(request, *, acknowledge=False):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
raise ValueError()
|
||||
record_id = data.get("record_id")
|
||||
if acknowledge:
|
||||
if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
|
||||
raise ValueError()
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
|
||||
if _pool(worker_id) is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
try:
|
||||
manager = get_manager()
|
||||
if acknowledge:
|
||||
await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
|
||||
result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
|
||||
else:
|
||||
result = {"resolved": await manager.reconcile(worker_id=worker_id)}
|
||||
return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
|
||||
except ValueError:
|
||||
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this pool"}, status=409)
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
|
||||
async def reconcile_serverless_endpoint(request):
|
||||
return await _recover(request)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
|
||||
async def resolve_serverless_endpoint(request):
|
||||
return await _recover(request, acknowledge=True)
|
||||
@@ -1,51 +0,0 @@
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.cloudflare import cloudflare_tunnel_manager
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/tunnel/status")
|
||||
async def tunnel_status_endpoint(request):
|
||||
"""Return Cloudflare tunnel status and last known details."""
|
||||
try:
|
||||
status = cloudflare_tunnel_manager.get_status()
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host")
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": status,
|
||||
"master_host": master_host
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/start")
|
||||
async def tunnel_start_endpoint(request):
|
||||
"""Start a Cloudflare tunnel pointing at the current ComfyUI server."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.start_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
|
||||
async def tunnel_stop_endpoint(request):
|
||||
"""Stop the managed Cloudflare tunnel if running."""
|
||||
try:
|
||||
result = await cloudflare_tunnel_manager.stop_tunnel()
|
||||
config = load_config()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"tunnel": result,
|
||||
"master_host": (config.get("master") or {}).get("host")
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
@@ -1,228 +0,0 @@
|
||||
import asyncio
|
||||
import io
|
||||
import time
|
||||
|
||||
from aiohttp import web
|
||||
from PIL import Image
|
||||
import server
|
||||
|
||||
from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
|
||||
from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
|
||||
from ..upscale.payload_parsers import _parse_tiles_from_form
|
||||
from ..utils.logging import debug_log
|
||||
from ..utils.network import handle_api_error
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if isinstance(job_data, BaseJobState):
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
debug_log(f"Heartbeat from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
return await handle_api_error(request, "Worker status tracking not available", 400)
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
|
||||
async def submit_tiles_endpoint(request):
|
||||
"""Endpoint for workers to submit processed tiles in static mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
|
||||
# Handle completion signal
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [],
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
try:
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, TileJobState):
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
|
||||
q = job_data.queue
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last,
|
||||
})
|
||||
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': [],
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/submit_image")
|
||||
async def submit_image_endpoint(request):
|
||||
"""Endpoint for workers to submit processed images in dynamic mode."""
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
# Handle image submission
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last,
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle completion signal (no image data)
|
||||
elif is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
await job_data.queue.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Missing image data or invalid request", 400)
|
||||
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_image")
|
||||
async def request_image_endpoint(request):
|
||||
"""Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return await handle_api_error(request, "Invalid job data structure", 500)
|
||||
|
||||
mode = job_data.mode
|
||||
if isinstance(job_data, ImageJobState):
|
||||
pending_queue = job_data.pending_images
|
||||
elif isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
else:
|
||||
return await handle_api_error(request, "Invalid job configuration", 400)
|
||||
|
||||
try:
|
||||
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
|
||||
job_data.worker_status[worker_id] = time.time()
|
||||
remaining = pending_queue.qsize()
|
||||
|
||||
if mode == 'dynamic':
|
||||
debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
|
||||
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
||||
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
||||
return web.json_response({
|
||||
"tile_idx": task_idx,
|
||||
"estimated_remaining": remaining,
|
||||
"batched_static": job_data.batched_static,
|
||||
})
|
||||
except asyncio.TimeoutError:
|
||||
if mode == 'dynamic':
|
||||
return web.json_response({"image_idx": None})
|
||||
return web.json_response({"tile_idx": None})
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/job_status")
|
||||
async def job_status_endpoint(request):
|
||||
"""Endpoint to check if a job is ready."""
|
||||
multi_job_id = request.query.get('multi_job_id')
|
||||
if not multi_job_id:
|
||||
return web.json_response({"ready": False})
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
|
||||
return web.json_response({"ready": ready})
|
||||
@@ -1,714 +0,0 @@
|
||||
import json
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
import platform
|
||||
import subprocess
|
||||
import socket
|
||||
|
||||
import torch
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import (
|
||||
build_worker_url,
|
||||
get_client_session,
|
||||
get_server_port,
|
||||
handle_api_error,
|
||||
normalize_host,
|
||||
probe_worker,
|
||||
)
|
||||
from ..utils.constants import CHUNK_SIZE
|
||||
from ..workers import get_worker_manager
|
||||
from ..workers.ports import allocate_worker_ports
|
||||
from .schemas import require_fields, validate_worker_id
|
||||
from ..workers.detection import (
|
||||
get_machine_id,
|
||||
is_docker_environment,
|
||||
is_runpod_environment,
|
||||
)
|
||||
try:
|
||||
from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
|
||||
except ImportError:
|
||||
from ..utils.async_helpers import queue_prompt_payload
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
def __init__(self, message, validation_error=None, node_errors=None):
|
||||
super().__init__(str(message))
|
||||
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_ws")
|
||||
async def worker_ws_endpoint(request):
|
||||
"""WebSocket endpoint for worker prompt dispatch."""
|
||||
ws = web.WebSocketResponse(heartbeat=30)
|
||||
await ws.prepare(request)
|
||||
|
||||
async for msg in ws:
|
||||
if msg.type == aiohttp.WSMsgType.TEXT:
|
||||
try:
|
||||
data = json.loads(msg.data or "{}")
|
||||
except json.JSONDecodeError:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": None,
|
||||
"ok": False,
|
||||
"error": "Invalid JSON payload.",
|
||||
})
|
||||
continue
|
||||
|
||||
if data.get("type") != "dispatch_prompt":
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Unsupported websocket message type.",
|
||||
})
|
||||
continue
|
||||
|
||||
prompt = data.get("prompt")
|
||||
if not isinstance(prompt, dict):
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": "Field 'prompt' must be an object.",
|
||||
})
|
||||
continue
|
||||
|
||||
try:
|
||||
prompt_id = await queue_prompt_payload(
|
||||
prompt,
|
||||
workflow_meta=data.get("workflow"),
|
||||
client_id=data.get("client_id"),
|
||||
)
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": True,
|
||||
"prompt_id": prompt_id,
|
||||
})
|
||||
except PromptValidationError as exc:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
"validation_error": exc.validation_error,
|
||||
"node_errors": exc.node_errors,
|
||||
})
|
||||
except Exception as exc:
|
||||
await ws.send_json({
|
||||
"type": "dispatch_ack",
|
||||
"request_id": data.get("request_id"),
|
||||
"ok": False,
|
||||
"error": str(exc),
|
||||
})
|
||||
elif msg.type == aiohttp.WSMsgType.ERROR:
|
||||
log(f"[Distributed] Worker websocket error: {ws.exception()}")
|
||||
|
||||
return ws
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
|
||||
async def clear_launching_state(request):
|
||||
"""Clear the launching flag when worker is confirmed running."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Clear launching flag in managed processes
|
||||
if worker_id in wm.processes:
|
||||
if 'launching' in wm.processes[worker_id]:
|
||||
del wm.processes[worker_id]['launching']
|
||||
wm.save_processes()
|
||||
debug_log(f"Cleared launching state for worker {worker_id}")
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
def get_network_ips():
|
||||
"""Get all network IPs, trying multiple methods."""
|
||||
ips = []
|
||||
hostname = socket.gethostname()
|
||||
|
||||
# Method 1: Try socket.getaddrinfo
|
||||
try:
|
||||
addr_info = socket.getaddrinfo(hostname, None)
|
||||
for info in addr_info:
|
||||
ip = info[4][0]
|
||||
if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now
|
||||
ips.append(ip)
|
||||
except (socket.gaierror, OSError):
|
||||
pass
|
||||
|
||||
# Method 2: Try to connect to external server and get local IP
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
s.connect(("8.8.8.8", 80)) # Google DNS
|
||||
local_ip = s.getsockname()[0]
|
||||
s.close()
|
||||
if local_ip not in ips:
|
||||
ips.append(local_ip)
|
||||
except (OSError, socket.error):
|
||||
pass
|
||||
|
||||
# Method 3: Platform-specific commands
|
||||
try:
|
||||
if platform.system() == "Windows":
|
||||
# Windows ipconfig
|
||||
result = subprocess.run(["ipconfig"], capture_output=True, text=True)
|
||||
lines = result.stdout.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
if 'IPv4' in line and i + 1 < len(lines):
|
||||
ip = lines[i].split(':')[-1].strip()
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
else:
|
||||
# Unix/Linux/Mac ifconfig or ip addr
|
||||
try:
|
||||
result = subprocess.run(["ip", "addr"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
try:
|
||||
result = subprocess.run(["ifconfig"], capture_output=True, text=True)
|
||||
except (FileNotFoundError, OSError):
|
||||
result = None
|
||||
|
||||
import re
|
||||
ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)')
|
||||
if result is not None:
|
||||
for match in ip_pattern.finditer(result.stdout):
|
||||
ip = match.group(1)
|
||||
if ip and ip not in ips:
|
||||
ips.append(ip)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
pass
|
||||
|
||||
return ips
|
||||
|
||||
|
||||
def get_recommended_ip(ips):
|
||||
"""Choose the best IP for master-worker communication."""
|
||||
# Priority order:
|
||||
# 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
|
||||
# 2. Other non-localhost IPs
|
||||
# 3. Localhost as last resort
|
||||
|
||||
private_ips = []
|
||||
public_ips = []
|
||||
|
||||
for ip in ips:
|
||||
if ip.startswith('127.') or ip == 'localhost':
|
||||
continue
|
||||
elif (ip.startswith('192.168.')
|
||||
or ip.startswith('10.')
|
||||
or (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)):
|
||||
private_ips.append(ip)
|
||||
else:
|
||||
public_ips.append(ip)
|
||||
|
||||
# Prefer private IPs
|
||||
if private_ips:
|
||||
# Prefer 192.168 range as it's most common
|
||||
for ip in private_ips:
|
||||
if ip.startswith('192.168.'):
|
||||
return ip
|
||||
return private_ips[0]
|
||||
elif public_ips:
|
||||
return public_ips[0]
|
||||
elif ips:
|
||||
return ips[0]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _get_cuda_info():
|
||||
"""Detect CUDA device index and total physical GPU count.
|
||||
|
||||
Returns (cuda_device, cuda_device_count, physical_device_count).
|
||||
All three are 0/None if CUDA is unavailable.
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
return None, 0, 0
|
||||
try:
|
||||
cuda_device_count = torch.cuda.device_count()
|
||||
cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '')
|
||||
if cuda_visible and cuda_visible.strip():
|
||||
visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()]
|
||||
if visible_devices:
|
||||
cuda_device = visible_devices[0]
|
||||
try:
|
||||
result = subprocess.run(
|
||||
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
)
|
||||
physical_device_count = (
|
||||
len(result.stdout.strip().split('\n'))
|
||||
if result.returncode == 0
|
||||
else max(visible_devices) + 1
|
||||
)
|
||||
except (FileNotFoundError, OSError, subprocess.SubprocessError):
|
||||
physical_device_count = max(visible_devices) + 1
|
||||
return cuda_device, cuda_device_count, physical_device_count
|
||||
else:
|
||||
return 0, cuda_device_count, cuda_device_count
|
||||
else:
|
||||
cuda_device = torch.cuda.current_device()
|
||||
return cuda_device, cuda_device_count, cuda_device_count
|
||||
except Exception as e:
|
||||
debug_log(f"CUDA detection error: {e}")
|
||||
return None, 0, 0
|
||||
|
||||
|
||||
def _collect_network_info_sync():
|
||||
"""Collect network/cuda info in a worker thread to avoid blocking route handlers."""
|
||||
cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
|
||||
device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
|
||||
master_port = get_server_port()
|
||||
config = load_config()
|
||||
worker_ports = []
|
||||
if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
|
||||
worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
|
||||
worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
|
||||
hostname = socket.gethostname()
|
||||
all_ips = get_network_ips()
|
||||
recommended_ip = get_recommended_ip(all_ips)
|
||||
return {
|
||||
"hostname": hostname,
|
||||
"all_ips": all_ips,
|
||||
"recommended_ip": recommended_ip,
|
||||
"cuda_device": cuda_device,
|
||||
"cuda_device_count": device_count,
|
||||
"master_port": master_port,
|
||||
"local_worker_ports": worker_ports,
|
||||
}
|
||||
|
||||
|
||||
def _read_worker_log_sync(log_file, lines_to_read):
|
||||
"""Read worker log content from disk in a threadpool worker."""
|
||||
file_size = os.path.getsize(log_file)
|
||||
|
||||
with open(log_file, 'r', encoding='utf-8', errors='replace') as f:
|
||||
if lines_to_read > 0 and file_size > 1024 * 1024:
|
||||
# Read last N lines efficiently from end of file.
|
||||
lines = []
|
||||
f.seek(0, 2)
|
||||
file_length = f.tell()
|
||||
chunk_size = min(CHUNK_SIZE, file_length)
|
||||
|
||||
while len(lines) < lines_to_read and f.tell() > 0:
|
||||
current_pos = max(0, f.tell() - chunk_size)
|
||||
f.seek(current_pos)
|
||||
chunk = f.read(chunk_size)
|
||||
chunk_lines = chunk.splitlines()
|
||||
if current_pos > 0:
|
||||
chunk_lines = chunk_lines[1:]
|
||||
lines = chunk_lines + lines
|
||||
f.seek(current_pos)
|
||||
|
||||
content = '\n'.join(lines[-lines_to_read:])
|
||||
truncated = len(lines) > lines_to_read
|
||||
else:
|
||||
content = f.read()
|
||||
truncated = False
|
||||
|
||||
return {
|
||||
"content": content,
|
||||
"file_size": file_size,
|
||||
"truncated": truncated,
|
||||
"lines_shown": lines_to_read if truncated else content.count('\n') + 1,
|
||||
}
|
||||
|
||||
|
||||
def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
|
||||
"""Parse bounded positive integer query params with sane fallback."""
|
||||
try:
|
||||
parsed = int(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
parsed = max(minimum, parsed)
|
||||
if maximum is not None:
|
||||
parsed = min(maximum, parsed)
|
||||
return parsed
|
||||
|
||||
|
||||
def _find_worker_by_id(config, worker_id):
|
||||
worker_id_str = str(worker_id).strip()
|
||||
for worker in config.get("workers", []):
|
||||
if str(worker.get("id")).strip() == worker_id_str:
|
||||
return worker
|
||||
return None
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local_log")
|
||||
async def get_local_log_endpoint(request):
|
||||
"""Return this instance's in-memory ComfyUI log buffer."""
|
||||
try:
|
||||
from app.logger import get_logs
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to import app.logger: {e}", 500)
|
||||
|
||||
try:
|
||||
lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
|
||||
logs = get_logs()
|
||||
if logs is None:
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"content": "",
|
||||
"entries": 0,
|
||||
"source": "memory",
|
||||
"truncated": False,
|
||||
"lines_shown": 0,
|
||||
}
|
||||
)
|
||||
|
||||
entries = list(logs)
|
||||
selected_entries = entries[-lines_to_read:]
|
||||
content = "".join(
|
||||
entry.get("m", "") if isinstance(entry, dict) else str(entry)
|
||||
for entry in selected_entries
|
||||
)
|
||||
lines_shown = content.count("\n") + (1 if content else 0)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"status": "success",
|
||||
"content": content,
|
||||
"entries": len(selected_entries),
|
||||
"source": "memory",
|
||||
"truncated": len(entries) > len(selected_entries),
|
||||
"lines_shown": lines_shown,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/network_info")
|
||||
async def get_network_info_endpoint(request):
|
||||
"""Get network interfaces and recommend best IP for master."""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
info = await loop.run_in_executor(None, _collect_network_info_sync)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
**info,
|
||||
"message": "Auto-detected network configuration"
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/system_info")
|
||||
async def get_system_info_endpoint(request):
|
||||
"""Get system information including machine ID for local worker detection."""
|
||||
try:
|
||||
import socket
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"hostname": socket.gethostname(),
|
||||
"machine_id": get_machine_id(),
|
||||
"platform": {
|
||||
"system": platform.system(),
|
||||
"machine": platform.machine(),
|
||||
"node": platform.node(),
|
||||
"path_separator": os.sep, # Add path separator
|
||||
"os_name": os.name # Add OS name (posix, nt, etc.)
|
||||
},
|
||||
"is_docker": is_docker_environment(),
|
||||
"is_runpod": is_runpod_environment(),
|
||||
"runpod_pod_id": os.environ.get('RUNPOD_POD_ID')
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/launch_worker")
|
||||
async def launch_worker_endpoint(request):
|
||||
"""Launch a worker process from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
|
||||
# Find worker config
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
return await handle_api_error(request, "Serverless pools are managed by Runpod, not local processes", 400)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = worker_id
|
||||
|
||||
# Check if already running (managed by this instance)
|
||||
if worker_id_str in wm.processes:
|
||||
proc_info = wm.processes[worker_id_str]
|
||||
process = proc_info.get('process')
|
||||
|
||||
# Check if still running
|
||||
is_running = False
|
||||
if process:
|
||||
is_running = process.poll() is None
|
||||
else:
|
||||
# Restored process without subprocess object
|
||||
is_running = wm._is_process_running(proc_info['pid'])
|
||||
|
||||
if is_running:
|
||||
return await handle_api_error(request, "Worker already running (managed by UI)", 409)
|
||||
else:
|
||||
# Process is dead, remove it
|
||||
del wm.processes[worker_id_str]
|
||||
wm.save_processes()
|
||||
|
||||
# Launch the worker
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
pid = await loop.run_in_executor(None, wm.launch_worker, worker)
|
||||
log_file = wm.processes[worker_id_str].get('log_file')
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"pid": pid,
|
||||
"message": f"Worker {worker['name']} launched",
|
||||
"log_file": log_file
|
||||
})
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/stop_worker")
|
||||
async def stop_worker_endpoint(request):
|
||||
"""Stop a worker process that was launched from the UI."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
data = await request.json()
|
||||
missing = require_fields(data, "worker_id")
|
||||
if missing:
|
||||
return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
|
||||
|
||||
worker_id = str(data.get("worker_id")).strip()
|
||||
config = load_config()
|
||||
if not validate_worker_id(worker_id, config):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
success, message = wm.stop_worker(worker_id)
|
||||
|
||||
if success:
|
||||
return web.json_response({"status": "success", "message": message})
|
||||
else:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
message,
|
||||
404 if "not managed" in message else 409,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/managed_workers")
|
||||
async def get_managed_workers_endpoint(request):
|
||||
"""Get list of workers managed by this UI instance."""
|
||||
try:
|
||||
managed = get_worker_manager().get_managed_workers()
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"managed_workers": managed
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/local-worker-status")
|
||||
async def get_local_worker_status_endpoint(request):
|
||||
"""Check status of all local workers (localhost/no host specified)."""
|
||||
try:
|
||||
config = load_config()
|
||||
worker_statuses = {}
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
# 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"]:
|
||||
worker_id = worker["id"]
|
||||
port = worker["port"]
|
||||
|
||||
# Check if worker is enabled
|
||||
if not worker.get("enabled", False):
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": False,
|
||||
"processing": False,
|
||||
"queue_count": 0
|
||||
}
|
||||
continue
|
||||
|
||||
# Try to connect to worker
|
||||
try:
|
||||
worker_url = build_worker_url(worker)
|
||||
data = await probe_worker(worker_url, timeout=2.0)
|
||||
if data is None:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Unavailable",
|
||||
}
|
||||
continue
|
||||
queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
|
||||
worker_statuses[worker_id] = {
|
||||
"online": True,
|
||||
"enabled": True,
|
||||
"processing": queue_remaining > 0,
|
||||
"queue_count": queue_remaining
|
||||
}
|
||||
except asyncio.TimeoutError:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": "Timeout"
|
||||
}
|
||||
except Exception as e:
|
||||
worker_statuses[worker_id] = {
|
||||
"online": False,
|
||||
"enabled": True,
|
||||
"processing": False,
|
||||
"queue_count": 0,
|
||||
"error": str(e)
|
||||
}
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"worker_statuses": worker_statuses
|
||||
})
|
||||
except Exception as e:
|
||||
debug_log(f"Error checking local worker status: {e}")
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
|
||||
async def get_worker_log_endpoint(request):
|
||||
"""Get log content for a specific worker."""
|
||||
try:
|
||||
wm = get_worker_manager()
|
||||
worker_id = request.match_info['worker_id']
|
||||
|
||||
# Ensure worker_id is string
|
||||
worker_id = str(worker_id)
|
||||
|
||||
# Check if we manage this worker
|
||||
if worker_id not in wm.processes:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404)
|
||||
|
||||
proc_info = wm.processes[worker_id]
|
||||
log_file = proc_info.get('log_file')
|
||||
|
||||
if not log_file or not os.path.exists(log_file):
|
||||
return await handle_api_error(request, "Log file not found", 404)
|
||||
|
||||
# Read last N lines (or full file if small)
|
||||
lines_to_read = _parse_positive_int_query(request.query.get('lines'), default=1000)
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
payload = await loop.run_in_executor(None, _read_worker_log_sync, log_file, lines_to_read)
|
||||
|
||||
return web.json_response({
|
||||
"status": "success",
|
||||
"content": payload["content"],
|
||||
"log_file": log_file,
|
||||
"file_size": payload["file_size"],
|
||||
"truncated": payload["truncated"],
|
||||
"lines_shown": payload["lines_shown"],
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Error reading log file: {str(e)}", 500)
|
||||
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
|
||||
async def get_remote_worker_log_endpoint(request):
|
||||
"""Proxy a remote worker log request to the worker's local in-memory log endpoint."""
|
||||
try:
|
||||
worker_id = str(request.match_info["worker_id"]).strip()
|
||||
config = load_config()
|
||||
worker = _find_worker_by_id(config, worker_id)
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Remote log proxy is only meaningful for remote/cloud workers.
|
||||
host = normalize_host(worker.get("host")) or ""
|
||||
if not host:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Worker {worker_id} is local; use /distributed/worker_log/{worker_id} instead.",
|
||||
400,
|
||||
)
|
||||
|
||||
lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000)
|
||||
worker_url = build_worker_url(worker, "/distributed/local_log")
|
||||
session = await get_client_session()
|
||||
async with session.get(
|
||||
worker_url,
|
||||
params={"lines": str(lines_to_read)},
|
||||
timeout=aiohttp.ClientTimeout(total=5),
|
||||
) as resp:
|
||||
if resp.status >= 400:
|
||||
body = await resp.text()
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Remote worker {worker_id} returned HTTP {resp.status}: {body[:400]}",
|
||||
resp.status,
|
||||
)
|
||||
|
||||
try:
|
||||
data = await resp.json()
|
||||
except Exception as e:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Remote worker {worker_id} returned invalid JSON: {e}",
|
||||
502,
|
||||
)
|
||||
|
||||
return web.json_response(data)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
-31
@@ -1,31 +0,0 @@
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
# All actual package context is provided by each test module via
|
||||
# importlib.util.spec_from_file_location with synthetic stub packages.
|
||||
|
||||
from _pytest.python import Package
|
||||
|
||||
_orig_pkg_setup = Package.setup
|
||||
|
||||
|
||||
def _patched_pkg_setup(self) -> None:
|
||||
# Skip the root package setup — its __init__.py uses relative imports
|
||||
# that require a parent package (ComfyUI's plugin loader) which is not
|
||||
# available in the test environment.
|
||||
if self.path == self.config.rootpath:
|
||||
return
|
||||
_orig_pkg_setup(self)
|
||||
|
||||
|
||||
Package.setup = _patched_pkg_setup
|
||||
|
||||
collect_ignore = [
|
||||
"__init__.py",
|
||||
"distributed.py",
|
||||
"distributed_upscale.py",
|
||||
]
|
||||
+2900
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,605 @@
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
|
||||
import aiohttp
|
||||
|
||||
import execution
|
||||
import server
|
||||
|
||||
from .utils.config import load_config
|
||||
from .utils.logging import debug_log, log
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def ensure_distributed_state():
|
||||
"""Ensure prompt_server has the state used by distributed queue orchestration."""
|
||||
if not hasattr(prompt_server, "distributed_pending_jobs"):
|
||||
prompt_server.distributed_pending_jobs = {}
|
||||
if not hasattr(prompt_server, "distributed_jobs_lock"):
|
||||
prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
if not hasattr(prompt_server, "distributed_worker_ws"):
|
||||
prompt_server.distributed_worker_ws = {}
|
||||
|
||||
|
||||
async def _get_client_session():
|
||||
"""Get or create aiohttp client session (shared with distributed.py)."""
|
||||
if not hasattr(prompt_server, "_distributed_session"):
|
||||
prompt_server._distributed_session = aiohttp.ClientSession()
|
||||
return prompt_server._distributed_session
|
||||
|
||||
|
||||
class PromptIndex:
|
||||
"""Cache prompt metadata for faster worker/master prompt preparation."""
|
||||
|
||||
def __init__(self, prompt_obj):
|
||||
self._prompt_json = json.dumps(prompt_obj)
|
||||
self.nodes_by_class = {}
|
||||
self.class_by_node = {}
|
||||
self.inputs_by_node = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
class_type = node.get("class_type")
|
||||
node_id_str = str(node_id)
|
||||
if class_type:
|
||||
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
|
||||
self.class_by_node[node_id_str] = class_type
|
||||
self.inputs_by_node[node_id_str] = node.get("inputs", {})
|
||||
self._upstream_cache = {}
|
||||
|
||||
def copy_prompt(self):
|
||||
return json.loads(self._prompt_json)
|
||||
|
||||
def nodes_for_class(self, class_name):
|
||||
return self.nodes_by_class.get(class_name, [])
|
||||
|
||||
def has_upstream(self, start_node_id, target_class):
|
||||
cache_key = (str(start_node_id), target_class)
|
||||
if cache_key in self._upstream_cache:
|
||||
return self._upstream_cache[cache_key]
|
||||
|
||||
visited = set()
|
||||
stack = [str(start_node_id)]
|
||||
while stack:
|
||||
node_id = stack.pop()
|
||||
if node_id in visited:
|
||||
continue
|
||||
visited.add(node_id)
|
||||
inputs = self.inputs_by_node.get(node_id, {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
upstream_id = str(value[0])
|
||||
if self.class_by_node.get(upstream_id) == target_class:
|
||||
self._upstream_cache[cache_key] = True
|
||||
return True
|
||||
if upstream_id in self.inputs_by_node:
|
||||
stack.append(upstream_id)
|
||||
|
||||
self._upstream_cache[cache_key] = False
|
||||
return False
|
||||
|
||||
|
||||
def _iter_prompt_nodes(prompt_obj):
|
||||
for node_id, node in prompt_obj.items():
|
||||
if isinstance(node, dict):
|
||||
yield str(node_id), node
|
||||
|
||||
|
||||
def _find_nodes_by_class(prompt_obj, class_name):
|
||||
nodes = []
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
if node.get("class_type") == class_name:
|
||||
nodes.append(node_id)
|
||||
return nodes
|
||||
|
||||
|
||||
def _find_downstream_nodes(prompt_obj, start_ids):
|
||||
"""Return all nodes reachable downstream from the provided IDs."""
|
||||
adjacency = {}
|
||||
for node_id, node in _iter_prompt_nodes(prompt_obj):
|
||||
inputs = node.get("inputs", {})
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and len(value) == 2:
|
||||
source_id = str(value[0])
|
||||
adjacency.setdefault(source_id, set()).add(str(node_id))
|
||||
|
||||
connected = set(start_ids)
|
||||
queue = deque(start_ids)
|
||||
while queue:
|
||||
current = queue.popleft()
|
||||
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
|
||||
if dependent not in connected:
|
||||
connected.add(dependent)
|
||||
queue.append(dependent)
|
||||
return connected
|
||||
|
||||
|
||||
def _create_numeric_id_generator(prompt_obj):
|
||||
"""Return a closure that yields new numeric string IDs."""
|
||||
max_id = 0
|
||||
for node_id in prompt_obj.keys():
|
||||
try:
|
||||
numeric = int(node_id)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
max_id = max(max_id, numeric)
|
||||
|
||||
counter = max_id
|
||||
|
||||
def _next_id():
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
return str(counter)
|
||||
|
||||
return _next_id
|
||||
|
||||
|
||||
def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
|
||||
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
|
||||
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
|
||||
nodes_to_keep = set(collector_ids)
|
||||
nodes_to_keep.update(downstream)
|
||||
|
||||
pruned_prompt = {}
|
||||
for node_id in nodes_to_keep:
|
||||
node = prompt_obj.get(node_id)
|
||||
if node is not None:
|
||||
pruned_prompt[node_id] = json.loads(json.dumps(node))
|
||||
|
||||
pruned_ids = set(pruned_prompt.keys())
|
||||
for node_id, node in pruned_prompt.items():
|
||||
inputs = node.get("inputs")
|
||||
if not inputs:
|
||||
continue
|
||||
for input_name, input_value in list(inputs.items()):
|
||||
if isinstance(input_value, list) and len(input_value) == 2:
|
||||
source_id = str(input_value[0])
|
||||
if source_id not in pruned_ids:
|
||||
inputs.pop(input_name, None)
|
||||
debug_log(
|
||||
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
|
||||
)
|
||||
|
||||
next_id = _create_numeric_id_generator(pruned_prompt)
|
||||
for collector_id in collector_ids:
|
||||
collector_entry = pruned_prompt.get(collector_id)
|
||||
if not collector_entry:
|
||||
continue
|
||||
placeholder_id = next_id()
|
||||
pruned_prompt[placeholder_id] = {
|
||||
"class_type": "DistributedEmptyImage",
|
||||
"inputs": {
|
||||
"height": 64,
|
||||
"width": 64,
|
||||
"channels": 3,
|
||||
},
|
||||
"_meta": {
|
||||
"title": "Distributed Empty Image (auto-added)",
|
||||
},
|
||||
}
|
||||
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
|
||||
debug_log(
|
||||
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
|
||||
)
|
||||
|
||||
return pruned_prompt
|
||||
|
||||
|
||||
async def _ensure_distributed_queue(job_id):
|
||||
"""Ensure a queue exists for the given distributed job ID."""
|
||||
ensure_distributed_state()
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
|
||||
|
||||
|
||||
def _generate_job_id_map(prompt_index, prefix):
|
||||
"""Create stable per-node job IDs for distributed nodes."""
|
||||
job_map = {}
|
||||
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
|
||||
"UltimateSDUpscaleDistributed"
|
||||
)
|
||||
for node_id in distributed_nodes:
|
||||
job_map[node_id] = f"{prefix}_{node_id}"
|
||||
return job_map
|
||||
|
||||
|
||||
def _resolve_enabled_workers(config, requested_ids=None):
|
||||
"""Return a list of worker configs that should participate."""
|
||||
workers = []
|
||||
for worker in config.get("workers", []):
|
||||
worker_id = str(worker.get("id") or "").strip()
|
||||
if not worker_id:
|
||||
continue
|
||||
|
||||
if requested_ids is not None:
|
||||
if worker_id not in requested_ids:
|
||||
continue
|
||||
elif not worker.get("enabled", False):
|
||||
continue
|
||||
|
||||
raw_port = worker.get("port", worker.get("listen_port", 8188))
|
||||
try:
|
||||
port = int(raw_port or 8188)
|
||||
except (TypeError, ValueError):
|
||||
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
|
||||
port = 8188
|
||||
|
||||
workers.append(
|
||||
{
|
||||
"id": worker_id,
|
||||
"name": worker.get("name", worker_id),
|
||||
"host": worker.get("host"),
|
||||
"port": port,
|
||||
"type": worker.get("type", "local"),
|
||||
}
|
||||
)
|
||||
return workers
|
||||
|
||||
|
||||
def _resolve_master_url():
|
||||
"""Best-effort reconstruction of the master's public URL."""
|
||||
cfg = load_config()
|
||||
master_cfg = cfg.get("master", {}) or {}
|
||||
configured_host = (master_cfg.get("host") or "").strip()
|
||||
configured_port = master_cfg.get("port")
|
||||
default_port = getattr(prompt_server, "port", 8188) or 8188
|
||||
port = int(configured_port or default_port)
|
||||
|
||||
def _needs_https(hostname):
|
||||
hostname = hostname.lower()
|
||||
https_domains = (
|
||||
".proxy.runpod.net",
|
||||
".ngrok-free.app",
|
||||
".ngrok-free.dev",
|
||||
".ngrok.io",
|
||||
".trycloudflare.com",
|
||||
".cloudflare.dev",
|
||||
)
|
||||
return any(hostname.endswith(suffix) for suffix in https_domains)
|
||||
|
||||
if configured_host:
|
||||
if configured_host.startswith(("http://", "https://")):
|
||||
return configured_host.rstrip("/")
|
||||
|
||||
host = configured_host
|
||||
scheme = "https" if _needs_https(host) or port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
# For ngrok/cloud domains without explicit port, default to the scheme's default
|
||||
if configured_port is None and scheme == "https" and _needs_https(host):
|
||||
port = default_port_for_scheme
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{host}{port_part}"
|
||||
|
||||
address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
|
||||
if address in ("0.0.0.0", "::"):
|
||||
address = "127.0.0.1"
|
||||
scheme = "https" if port == 443 else "http"
|
||||
default_port_for_scheme = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port_for_scheme else f":{port}"
|
||||
return f"{scheme}://{address}{port_part}"
|
||||
|
||||
|
||||
def _build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
host = (worker.get("host") or "").strip()
|
||||
port = int(worker.get("port", 8188) or 8188)
|
||||
|
||||
if not host:
|
||||
host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
|
||||
|
||||
if host.startswith(("http://", "https://")):
|
||||
base = host.rstrip("/")
|
||||
else:
|
||||
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
|
||||
scheme = "https" if is_cloud else "http"
|
||||
default_port = 443 if scheme == "https" else 80
|
||||
port_part = "" if port == default_port else f":{port}"
|
||||
base = f"{scheme}://{host}{port_part}"
|
||||
|
||||
return f"{base}{endpoint}"
|
||||
|
||||
|
||||
async def _worker_is_active(worker):
|
||||
"""Ping worker's /prompt endpoint to confirm it's reachable."""
|
||||
url = _build_worker_url(worker, "/prompt")
|
||||
session = await _get_client_session()
|
||||
try:
|
||||
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
|
||||
return resp.status == 200
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def _worker_ws_is_active(worker):
|
||||
"""Ping worker's websocket endpoint to confirm it's reachable."""
|
||||
session = await _get_client_session()
|
||||
url = _build_worker_url(worker, "/distributed/worker_ws")
|
||||
try:
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
|
||||
await ws.close()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
async def _get_worker_ws_connection(worker):
|
||||
"""Return a cached websocket connection to a worker (creating if needed)."""
|
||||
ensure_distributed_state()
|
||||
worker_id = worker.get("id")
|
||||
existing = prompt_server.distributed_worker_ws.get(worker_id)
|
||||
if existing and not existing.closed:
|
||||
return existing
|
||||
|
||||
session = await _get_client_session()
|
||||
url = _build_worker_url(worker, "/distributed/worker_ws")
|
||||
ws = await session.ws_connect(url, heartbeat=20, timeout=10)
|
||||
prompt_server.distributed_worker_ws[worker_id] = ws
|
||||
return ws
|
||||
|
||||
|
||||
async def _close_worker_ws(worker_id):
|
||||
ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
|
||||
if ws and not ws.closed:
|
||||
await ws.close()
|
||||
|
||||
|
||||
async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
|
||||
"""Send the prepared prompt to a worker ComfyUI instance."""
|
||||
url = _build_worker_url(worker, "/prompt")
|
||||
payload = {"prompt": prompt_obj}
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if extra_data:
|
||||
payload["extra_data"] = extra_data
|
||||
|
||||
if use_websocket:
|
||||
request_id = uuid.uuid4().hex
|
||||
ws_payload = {
|
||||
"type": "dispatch_prompt",
|
||||
"request_id": request_id,
|
||||
"prompt": prompt_obj,
|
||||
"workflow": workflow_meta,
|
||||
"client_id": client_id,
|
||||
}
|
||||
try:
|
||||
ws = await _get_worker_ws_connection(worker)
|
||||
await ws.send_json(ws_payload)
|
||||
response = await asyncio.wait_for(ws.receive(), timeout=60)
|
||||
if response.type == aiohttp.WSMsgType.TEXT:
|
||||
data = json.loads(response.data or "{}")
|
||||
if (
|
||||
data.get("type") == "dispatch_ack"
|
||||
and data.get("request_id") == request_id
|
||||
and data.get("ok")
|
||||
):
|
||||
return
|
||||
error_msg = data.get("error") or "Worker rejected websocket dispatch."
|
||||
raise RuntimeError(error_msg)
|
||||
raise RuntimeError("Unexpected websocket response from worker.")
|
||||
except Exception as exc:
|
||||
worker_id = worker.get("id")
|
||||
await _close_worker_ws(worker_id)
|
||||
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
|
||||
|
||||
session = await _get_client_session()
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=60),
|
||||
) as resp:
|
||||
resp.raise_for_status()
|
||||
|
||||
|
||||
async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
|
||||
"""Queue the master prompt through ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
if not valid[0]:
|
||||
raise RuntimeError(f"Invalid prompt: {valid[1]}")
|
||||
|
||||
extra_data = {}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
extra_data["client_id"] = client_id
|
||||
|
||||
sensitive = {}
|
||||
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
return prompt_id
|
||||
|
||||
|
||||
def _apply_participant_overrides(
|
||||
prompt_copy,
|
||||
participant_id,
|
||||
enabled_worker_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
):
|
||||
"""Return a prompt copy with hidden inputs configured for master/worker."""
|
||||
is_master = participant_id == "master"
|
||||
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
|
||||
enabled_json = json.dumps(enabled_worker_ids)
|
||||
|
||||
# Distributed seed nodes
|
||||
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
inputs["is_worker"] = not is_master
|
||||
if not is_master:
|
||||
idx = worker_index_map.get(participant_id, 0)
|
||||
inputs["worker_id"] = f"worker_{idx}"
|
||||
else:
|
||||
inputs["worker_id"] = ""
|
||||
|
||||
# Distributed collectors
|
||||
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
|
||||
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
|
||||
inputs["pass_through"] = True
|
||||
continue
|
||||
|
||||
unique_id = job_id_map.get(node_id, node_id)
|
||||
inputs["multi_job_id"] = unique_id
|
||||
inputs["is_worker"] = not is_master
|
||||
|
||||
if is_master:
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["delegate_only"] = bool(delegate_master)
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs["delegate_only"] = False
|
||||
|
||||
# Distributed upscaler nodes
|
||||
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
|
||||
node = prompt_copy.get(node_id, {})
|
||||
inputs = node.setdefault("inputs", {})
|
||||
unique_id = job_id_map.get(node_id, node_id)
|
||||
inputs["multi_job_id"] = unique_id
|
||||
inputs["is_worker"] = not is_master
|
||||
if is_master:
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
inputs.pop("master_url", None)
|
||||
inputs.pop("worker_id", None)
|
||||
else:
|
||||
inputs["master_url"] = master_url
|
||||
inputs["worker_id"] = participant_id
|
||||
inputs["enabled_worker_ids"] = enabled_json
|
||||
|
||||
return prompt_copy
|
||||
|
||||
|
||||
async def orchestrate_distributed_execution(
|
||||
prompt_obj,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
enabled_worker_ids=None,
|
||||
delegate_master=None,
|
||||
):
|
||||
"""Core orchestration logic for the /distributed/queue endpoint.
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (prompt_id, worker_count)
|
||||
"""
|
||||
ensure_distributed_state()
|
||||
|
||||
config = load_config()
|
||||
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
|
||||
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
|
||||
workers = _resolve_enabled_workers(config, requested_ids)
|
||||
prompt_index = PromptIndex(prompt_obj)
|
||||
|
||||
# Respect master delegate-only configuration
|
||||
if delegate_master is None:
|
||||
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
|
||||
|
||||
if not workers and delegate_master:
|
||||
debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
|
||||
delegate_master = False
|
||||
|
||||
# Filter to active workers
|
||||
active_workers = []
|
||||
for worker in workers:
|
||||
if use_websocket:
|
||||
is_active = await _worker_ws_is_active(worker)
|
||||
else:
|
||||
is_active = await _worker_is_active(worker)
|
||||
if is_active:
|
||||
active_workers.append(worker)
|
||||
else:
|
||||
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
|
||||
|
||||
if not active_workers and delegate_master:
|
||||
debug_log("All workers offline while delegate-only requested; enabling master participation.")
|
||||
delegate_master = False
|
||||
|
||||
enabled_ids = [worker["id"] for worker in active_workers]
|
||||
|
||||
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
|
||||
job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
|
||||
|
||||
if not job_id_map:
|
||||
prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
|
||||
return prompt_id, 0
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
master_url = _resolve_master_url()
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = _apply_participant_overrides(
|
||||
master_prompt,
|
||||
"master",
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
|
||||
if delegate_master:
|
||||
collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
|
||||
upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
|
||||
if upscale_nodes:
|
||||
debug_log(
|
||||
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
|
||||
)
|
||||
elif not collector_ids:
|
||||
debug_log(
|
||||
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
|
||||
)
|
||||
else:
|
||||
master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
|
||||
|
||||
if active_workers:
|
||||
debug_log(
|
||||
"Active distributed workers: "
|
||||
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
|
||||
)
|
||||
worker_payloads = []
|
||||
for worker in active_workers:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
worker_prompt = _apply_participant_overrides(
|
||||
worker_prompt,
|
||||
worker["id"],
|
||||
enabled_ids,
|
||||
job_id_map,
|
||||
master_url,
|
||||
delegate_master,
|
||||
prompt_index,
|
||||
)
|
||||
worker_payloads.append((worker, worker_prompt))
|
||||
|
||||
if worker_payloads:
|
||||
await asyncio.gather(
|
||||
*[
|
||||
_dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
|
||||
for worker, wprompt in worker_payloads
|
||||
]
|
||||
)
|
||||
|
||||
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
|
||||
return prompt_id, len(worker_payloads)
|
||||
File diff suppressed because it is too large
Load Diff
+17
-105
@@ -19,10 +19,9 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
|
||||
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
|
||||
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
|
||||
- Resolves enabled/selected workers.
|
||||
- Probes and dispatches workers through `/distributed/worker_ws` by default.
|
||||
- If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
|
||||
- Pings workers (`GET /prompt`) to include only reachable ones.
|
||||
- Dispatches the workflow to workers (`POST /prompt`).
|
||||
- Queues the master workflow in ComfyUI’s prompt queue.
|
||||
- If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
|
||||
|
||||
### What it does *not* add
|
||||
|
||||
@@ -49,39 +48,27 @@ Queue a workflow for distributed execution.
|
||||
{
|
||||
"prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
|
||||
"workflow": { },
|
||||
"client_id": "external-client",
|
||||
"client_id": "optional",
|
||||
"delegate_master": false,
|
||||
"enabled_worker_ids": ["1", "2"],
|
||||
"workers": ["1", "2"],
|
||||
"auto_prepare": true,
|
||||
"trace_execution_id": "exec_1700000000_ab12cd"
|
||||
"enabled_worker_ids": ["1", "2"]
|
||||
}
|
||||
```
|
||||
|
||||
#### Fields
|
||||
|
||||
- `prompt` (required unless `workflow.prompt` is present, object)
|
||||
- A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
|
||||
- This is not the normal visual workflow export from the ComfyUI editor.
|
||||
- `prompt` (required, object)
|
||||
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
|
||||
- `workflow` (optional, object)
|
||||
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
|
||||
- If you don’t care about UI metadata, you can omit it.
|
||||
- `client_id` (required, string)
|
||||
- `client_id` (optional, string)
|
||||
- Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
|
||||
- `delegate_master` (optional, boolean)
|
||||
- If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
|
||||
- Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
|
||||
- `enabled_worker_ids` (required, array of strings)
|
||||
- The explicit worker IDs to consider for this run.
|
||||
- `workers` (optional, array of strings or objects with `id`)
|
||||
- Transitional alias for `enabled_worker_ids` used by older clients.
|
||||
- `auto_prepare` (optional, boolean)
|
||||
- Kept for wire compatibility.
|
||||
- Backend orchestration always runs with auto-prepare semantics.
|
||||
- If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
|
||||
- `trace_execution_id` (optional, string)
|
||||
- Passed through to orchestration logs.
|
||||
- Server log lines include the marker as `[exec:<trace_execution_id>]`.
|
||||
- `enabled_worker_ids` (optional, array of strings)
|
||||
- If provided, only these worker IDs will be considered.
|
||||
- If omitted, the plugin uses workers marked as enabled in the UI config.
|
||||
|
||||
##### How to get `enabled_worker_ids`
|
||||
|
||||
@@ -106,8 +93,7 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
```json
|
||||
{
|
||||
"prompt_id": "<uuid>",
|
||||
"worker_count": 2,
|
||||
"auto_prepare_supported": true
|
||||
"worker_count": 2
|
||||
}
|
||||
```
|
||||
|
||||
@@ -124,89 +110,24 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
|
||||
|
||||
## Worker requirements (important)
|
||||
|
||||
For a worker to participate, it must be reachable from the master. By default:
|
||||
|
||||
- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
|
||||
|
||||
If `settings.websocket_orchestration=false`:
|
||||
For a worker to participate, it must be reachable from the master:
|
||||
|
||||
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
|
||||
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
|
||||
|
||||
Also, for collector-based flows:
|
||||
|
||||
- Workers will send results back to the master via `POST /distributed/job_complete` (that route must be reachable from workers).
|
||||
|
||||
---
|
||||
|
||||
## Endpoint: `POST /distributed/job_complete`
|
||||
|
||||
Submit one completed worker image back to the master collector queue.
|
||||
|
||||
### URL
|
||||
|
||||
- `http://<master-host>:<master-port>/distributed/job_complete`
|
||||
|
||||
### Request Body
|
||||
|
||||
```json
|
||||
{
|
||||
"job_id": "exec_1234567890_17",
|
||||
"worker_id": "worker-1",
|
||||
"batch_idx": 0,
|
||||
"image": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAA...",
|
||||
"is_last": false
|
||||
}
|
||||
```
|
||||
|
||||
### Canonical envelope (required fields)
|
||||
|
||||
- `job_id` (string, required)
|
||||
- `worker_id` (string, required)
|
||||
- `batch_idx` (integer >= 0, required)
|
||||
- `image` (string, required)
|
||||
- PNG payload as either:
|
||||
- data URL: `data:image/png;base64,...`
|
||||
- raw base64 PNG bytes
|
||||
- `is_last` (boolean, required)
|
||||
|
||||
Legacy multipart/tensor payload formats are no longer accepted on this endpoint.
|
||||
|
||||
### CORS note
|
||||
|
||||
If you call the API from a browser (not from a backend), ensure the master ComfyUI is started with `--enable-cors-header`.
|
||||
|
||||
---
|
||||
|
||||
## Log Endpoints
|
||||
|
||||
### `GET /distributed/worker_log/{worker_id}`
|
||||
|
||||
Read log files for workers launched locally by the master UI process manager.
|
||||
|
||||
- Intended for managed local workers.
|
||||
- Query param: `lines` (optional, default `1000`).
|
||||
|
||||
### `GET /distributed/local_log`
|
||||
|
||||
Read this ComfyUI instance's in-memory runtime log buffer.
|
||||
|
||||
- Available on any ComfyUI-Distributed instance (master or worker).
|
||||
- Query param: `lines` (optional, default `300`, max `3000`).
|
||||
|
||||
### `GET /distributed/remote_worker_log/{worker_id}`
|
||||
|
||||
Proxy endpoint on master that fetches logs from a configured remote/cloud worker's
|
||||
`/distributed/local_log`.
|
||||
|
||||
- Intended for remote/cloud workers in master config.
|
||||
- Query param: `lines` (optional, default `300`, max `3000`).
|
||||
|
||||
---
|
||||
|
||||
## Examples
|
||||
|
||||
### 1) Minimal `curl` request envelope
|
||||
### 1) Minimal `curl`
|
||||
|
||||
```bash
|
||||
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
@@ -214,23 +135,16 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
|
||||
Where `payload.json` contains at least:
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": {
|
||||
"<node_id>": {
|
||||
"class_type": "<node class>",
|
||||
"inputs": {"<required_input>": "<value or connection>"}
|
||||
}
|
||||
},
|
||||
"enabled_worker_ids": [],
|
||||
"client_id": "external-client"
|
||||
"1": {"class_type": "KSampler", "inputs": {} }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
|
||||
|
||||
### 2) Python (`requests`)
|
||||
|
||||
```python
|
||||
@@ -240,7 +154,6 @@ url = "http://127.0.0.1:8188/distributed/queue"
|
||||
payload = {
|
||||
"prompt": {...},
|
||||
"workflow": {...},
|
||||
"client_id": "external-client",
|
||||
"delegate_master": False,
|
||||
"enabled_worker_ids": ["1", "2"],
|
||||
}
|
||||
@@ -258,7 +171,6 @@ const url = "http://127.0.0.1:8188/distributed/queue";
|
||||
const payload = {
|
||||
prompt: {/* ... */},
|
||||
workflow: {/* ... */},
|
||||
client_id: "external-client",
|
||||
delegate_master: false,
|
||||
enabled_worker_ids: ["1", "2"],
|
||||
};
|
||||
|
||||
@@ -1,163 +0,0 @@
|
||||
# Runpod serverless pools (preview)
|
||||
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint/pool, not a permanent
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
|
||||
normal ComfyUI API workflow per logical slot and collects the returned images.
|
||||
The endpoint workers **do not need ComfyUI-Distributed installed**.
|
||||
|
||||
## Setup
|
||||
|
||||
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
|
||||
2. Enter its name, Runpod **endpoint ID** (not a URL), and API key.
|
||||
3. Start with **Jobs per run = 1**, **Concurrent jobs = 1**, and **Timeout = 1800**.
|
||||
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
|
||||
endpoint health only; it does not start GPU inference or prove model/node compatibility.
|
||||
5. Enable the pool when you want queued Distributed workflows to use it.
|
||||
6. Use an image **Distributed Collector** and queue through Distributed as usual.
|
||||
Enable **Delegate Only** if the master should collect without generating locally.
|
||||
|
||||
Keys are write-only. Leave the key field blank to retain a saved key, or select
|
||||
**Clear saved API key** to remove it (disable the pool first). A saved key is
|
||||
required before enabling a pool. Use HTTPS or a trusted local connection when
|
||||
entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
|
||||
users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
|
||||
|
||||
## Endpoint contract
|
||||
|
||||
The endpoint must accept this generic input shape; named application modes are
|
||||
not used:
|
||||
|
||||
```json
|
||||
{
|
||||
"input": {
|
||||
"workflow": {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
|
||||
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}}
|
||||
},
|
||||
"images": []
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
For image inputs, `images` contains `{ "name": "input.png", "image": "<base64>" }`
|
||||
entries and the handler must rewrite `LoadImage` references to the staged files.
|
||||
The adapter transfers local input images; it does not copy models or install nodes.
|
||||
All referenced models and custom nodes must already exist in the endpoint image.
|
||||
|
||||
On completion, Runpod's `output.images` must contain artifact records such as:
|
||||
|
||||
```json
|
||||
{
|
||||
"filename": "example_00001_.png",
|
||||
"subfolder": "",
|
||||
"artifact_type": "output",
|
||||
"artifact_path": "example_00001_.png",
|
||||
"type": "base64",
|
||||
"data": "<base64>"
|
||||
}
|
||||
```
|
||||
|
||||
Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"`
|
||||
with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without
|
||||
Runpod credentials; private-network destinations and unsafe redirects are refused.
|
||||
The handler must preserve injected `SaveImage.filename_prefix` values so results
|
||||
can be matched to the right collector and slot. Duplicate basenames in different
|
||||
subfolders remain different artifact identities.
|
||||
|
||||
A Runpod job marked `COMPLETED` is still rejected if `output.error` or
|
||||
`output.errors` is nonempty, required collector outputs are missing, or image
|
||||
artifacts are invalid. Failures are surfaced rather than silently returning an
|
||||
incomplete serverless batch.
|
||||
|
||||
## Scheduling and cancellation
|
||||
|
||||
- **Jobs per run** adds logical collector participants. Each receives its own
|
||||
Distributed Seed/Value overrides. It is not a physical GPU count.
|
||||
- **Concurrent jobs** bounds this master's outstanding submissions per pool,
|
||||
including retained jobs after a restart. Ambiguous submissions and unconfirmed
|
||||
cancellations keep their capacity reservation until reconciled. Other
|
||||
applications using the same endpoint are unaffected.
|
||||
- **Load balance** selects one participant, so it does not fan out every pool slot.
|
||||
- **Timeout** includes waiting for a local slot, Runpod queue/cold-start time,
|
||||
execution, and output retrieval. Provider policy limits are submitted as well.
|
||||
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
|
||||
- Interrupting a collector or using the panel cancellation action requests
|
||||
cancellation only for jobs submitted by this master. It never purges the endpoint.
|
||||
- A network failure during submission is **not retried automatically**: the job
|
||||
may already be billable. Check the Runpod job list before queuing again.
|
||||
|
||||
## Persistence and recovery
|
||||
|
||||
Credentials live outside `gpu_config.json` and Git in the master ComfyUI user
|
||||
folder under `distributed/serverless/credentials.json`, with restrictive POSIX
|
||||
permissions. They are omitted from public configuration, workflow exports, and
|
||||
browser configuration state. The first version requires a POSIX (Linux/macOS)
|
||||
master for serverless secret storage; existing non-serverless workers remain
|
||||
available on Windows.
|
||||
|
||||
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
|
||||
not image data, workflow payloads, or API keys. A restarted master cannot resume
|
||||
an interrupted collector. It does **not** automatically resubmit jobs. Retained
|
||||
IDs are available through `/distributed/serverless/status/<worker_id>`; the pool's
|
||||
cancel action can request cancellation of known pre-restart jobs. If no job ID
|
||||
was received, check Runpod manually. Cancellation is best-effort and may race
|
||||
with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation.
|
||||
|
||||
Unresolved records protect the pool's endpoint, concurrency, API key, worker type,
|
||||
and deletion even when no local task remains. Disabling a pool is still allowed;
|
||||
it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are
|
||||
closed locally, while interrupted `SUBMITTING` records without IDs become
|
||||
`UNKNOWN` and remain protected.
|
||||
|
||||
### Reconcile or acknowledge recovery
|
||||
|
||||
Recovery uses the same trusted-master API as the panel; it never resubmits a job.
|
||||
|
||||
1. Read `GET /distributed/serverless/status/<worker_id>`. The `jobs` list includes
|
||||
every unresolved record with its local `record_id`, provider `job_id` (if
|
||||
received), and state. These records remain visible if endpoint health fails.
|
||||
2. Request cancellation from the panel, or send
|
||||
`POST /distributed/serverless/reconcile` with the body below to **read** the
|
||||
retained provider IDs. Only a confirmed terminal state releases protection.
|
||||
Network errors, expired IDs, or ambiguous responses do not.
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID"}
|
||||
```
|
||||
|
||||
3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
|
||||
manually first. **Only after verifying that the job is no longer running or
|
||||
queued**, send `POST /distributed/serverless/resolve` for that one record:
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
|
||||
```
|
||||
|
||||
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
|
||||
provider-confirmed cancellation. It cannot resolve a live local task and does
|
||||
not stop remote compute. It releases the record's reservation and may let
|
||||
already-waiting jobs proceed. Never automate this acknowledgment as a retry
|
||||
mechanism; an unverified job can still be billable.
|
||||
|
||||
Pool changes become available once all its local tasks and unresolved records
|
||||
are finished. No credentials or ledger files need to be edited by hand.
|
||||
|
||||
## First-version limits
|
||||
|
||||
- Image collectors only. Audio/video transfer and audio collectors are rejected.
|
||||
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
|
||||
node types are rejected before submission.
|
||||
- Downstream post-collector processing remains local; the remote graph contains
|
||||
only the upstream compute required for image saves.
|
||||
- Inputs must be supported images under ComfyUI's input/output/temp roots.
|
||||
- Raw image input reads are guarded at 10,000,000 bytes total, and the complete
|
||||
serialized `/run` JSON request body (workflow, base64 images, and policy) must
|
||||
also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image
|
||||
budget is lower and depends on workflow size. Large-media/object-storage upload
|
||||
support is not included.
|
||||
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
|
||||
scaling changes, or inbound master callback/tunnel is needed.
|
||||
|
||||
This feature is verified with offline provider/transport tests. A successful
|
||||
health test is not a paid end-to-end generation test; test your deployed handler
|
||||
and model inventory with one slot before increasing concurrency.
|
||||
@@ -1,53 +0,0 @@
|
||||
# Native V3 node API
|
||||
|
||||
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
|
||||
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
|
||||
that API; there is no V1 registration fallback.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- All eight node IDs, display names, categories, visible input order, defaults
|
||||
and output order are retained. Existing execution algorithms remain in the
|
||||
private collector, utilities and upscale modules.
|
||||
- The image/audio dividers explicitly declare ten typed outputs. Their existing
|
||||
frontend extensions still show only the selected number of outputs. This
|
||||
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
|
||||
IMAGE/AUDIO socket indices or the ten returned values.
|
||||
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
|
||||
its existing prompt-input names and Python defaults via `accept_all_inputs`;
|
||||
it does not become visible widgets.
|
||||
- The collector retains list-input handling. Upscale retains its always-changing
|
||||
fingerprint and creates a private runtime helper for each V3 execution, so
|
||||
mutable helper state is not attached to sanitized V3 class clones.
|
||||
- Routes, distributed state and the existing worker startup/shutdown hooks are
|
||||
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
|
||||
Worker mode still suppresses automatic worker launch. `distributed.py` is no
|
||||
longer an entrypoint.
|
||||
|
||||
## Verification
|
||||
|
||||
Run ordinary unit tests from this repository:
|
||||
|
||||
```bash
|
||||
python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
|
||||
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
|
||||
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
|
||||
all node schemas with the V1 fixture, checks all five bundled workflows' node
|
||||
IDs and socket/link contracts, exercises injected worker metadata, image/audio
|
||||
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
|
||||
and decodes an actual preview PNG referenced by executor history.
|
||||
|
||||
The upscale GPU/model boundary is mocked to check argument forwarding and
|
||||
per-execution helper isolation. Full checkpoint inference, browser canvas
|
||||
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
|
||||
workers or model downloads are started; preview files use a temporary scratch
|
||||
directory. The test does not install this branch into a live custom-node folder.
|
||||
@@ -79,7 +79,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
|
||||
|
||||
**On Runpod:**
|
||||
> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
|
||||
> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
|
||||
|
||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
||||
@@ -92,7 +92,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
|
||||
- SAGE_ATTENTION: optional optimisation (set to true/false)
|
||||
5. Deploy your pod.
|
||||
6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
|
||||
7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
|
||||
```
|
||||
# Download from CivitAI
|
||||
comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||
@@ -1,601 +0,0 @@
|
||||
import torch
|
||||
import io
|
||||
import json
|
||||
import asyncio
|
||||
import time
|
||||
import base64
|
||||
|
||||
import aiohttp
|
||||
import server as _server
|
||||
import comfy.model_management
|
||||
from comfy.utils import ProgressBar
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
|
||||
from ..utils.constants import HEARTBEAT_INTERVAL
|
||||
from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
|
||||
from ..utils.network import build_worker_url, get_client_session, probe_worker
|
||||
from ..utils.audio_payload import encode_audio_payload
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
|
||||
prompt_server = _server.PromptServer.instance
|
||||
|
||||
|
||||
class DistributedCollectorNode:
|
||||
INPUT_IS_LIST = True
|
||||
EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"load_balance": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating).",
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": {
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_batch_size": ("INT", {"default": 1, "min": 1, "max": 1024}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"pass_through": ("BOOLEAN", {"default": False}),
|
||||
"delegate_only": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "AUDIO")
|
||||
RETURN_NAMES = ("images", "audio")
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image"
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_list_input(value):
|
||||
"""Unwrap scalar inputs when ComfyUI passes them via INPUT_IS_LIST."""
|
||||
if isinstance(value, (list, tuple)) and len(value) == 1:
|
||||
return value[0]
|
||||
return value
|
||||
|
||||
def _normalize_images_input(self, images):
|
||||
"""Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
|
||||
if isinstance(images, (list, tuple)):
|
||||
if not images:
|
||||
raise ValueError("Collector received an empty image list")
|
||||
if not all(isinstance(image, torch.Tensor) for image in images):
|
||||
raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
|
||||
if len(images) == 1:
|
||||
return ensure_contiguous(images[0])
|
||||
return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
|
||||
return ensure_contiguous(images)
|
||||
|
||||
def _normalize_audio_input(self, audio):
|
||||
"""Collapse ComfyUI list AUDIO inputs into a single AUDIO payload when present."""
|
||||
if not isinstance(audio, (list, tuple)):
|
||||
return audio
|
||||
|
||||
audio_items = [item for item in audio if item is not None]
|
||||
if not audio_items:
|
||||
return None
|
||||
if len(audio_items) == 1:
|
||||
return audio_items[0]
|
||||
|
||||
waveforms = []
|
||||
sample_rate = 44100
|
||||
for item in audio_items:
|
||||
if not isinstance(item, dict):
|
||||
raise TypeError("Collector expected AUDIO list items to be dictionaries")
|
||||
waveform = item.get("waveform")
|
||||
if waveform is None or waveform.numel() == 0:
|
||||
continue
|
||||
waveforms.append(waveform)
|
||||
if sample_rate == 44100:
|
||||
sample_rate = item.get("sample_rate", 44100)
|
||||
|
||||
if not waveforms:
|
||||
return None
|
||||
return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
|
||||
|
||||
def run(self, images=None, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
|
||||
if images is not None:
|
||||
images = self._normalize_images_input(images)
|
||||
audio = self._normalize_audio_input(audio)
|
||||
load_balance = self._unwrap_list_input(load_balance)
|
||||
multi_job_id = self._unwrap_list_input(multi_job_id)
|
||||
is_worker = self._unwrap_list_input(is_worker)
|
||||
master_url = self._unwrap_list_input(master_url)
|
||||
enabled_worker_ids = self._unwrap_list_input(enabled_worker_ids)
|
||||
worker_batch_size = self._unwrap_list_input(worker_batch_size)
|
||||
worker_id = self._unwrap_list_input(worker_id)
|
||||
pass_through = self._unwrap_list_input(pass_through)
|
||||
delegate_only = self._unwrap_list_input(delegate_only)
|
||||
|
||||
remote_only_master = (
|
||||
bool(multi_job_id)
|
||||
and not is_worker
|
||||
and (delegate_only or is_master_delegate_only())
|
||||
)
|
||||
if images is None and audio is None and not remote_only_master:
|
||||
raise ValueError("DistributedCollector requires at least one image or audio input")
|
||||
|
||||
# Create empty audio if not provided
|
||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
|
||||
|
||||
if not multi_job_id or pass_through:
|
||||
if pass_through:
|
||||
debug_log("Collector: pass-through mode enabled, returning images unchanged")
|
||||
return (images, audio if audio is not None else empty_audio)
|
||||
|
||||
# Use async helper to run in server loop
|
||||
result = run_async_in_server_loop(
|
||||
self.execute(
|
||||
images,
|
||||
audio,
|
||||
load_balance,
|
||||
multi_job_id,
|
||||
is_worker,
|
||||
master_url,
|
||||
enabled_worker_ids,
|
||||
worker_batch_size,
|
||||
worker_id,
|
||||
delegate_only,
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _get_serverless_manager():
|
||||
"""Return the serverless manager only for collector jobs with serverless slots."""
|
||||
manager = getattr(prompt_server, "distributed_serverless", None)
|
||||
if manager is not None:
|
||||
return manager
|
||||
from ..api.orchestration.serverless_runtime import get_manager
|
||||
return get_manager()
|
||||
|
||||
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_complete"
|
||||
|
||||
payloads = []
|
||||
batch_size = 0 if image_batch is None else image_batch.shape[0]
|
||||
if batch_size == 0:
|
||||
if encoded_audio is None:
|
||||
raise ValueError("Worker completion requires image or audio data")
|
||||
payloads.append(
|
||||
{
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": 0,
|
||||
"audio": encoded_audio,
|
||||
"is_last": True,
|
||||
}
|
||||
)
|
||||
else:
|
||||
for batch_idx in range(batch_size):
|
||||
img = tensor_to_pil(image_batch[batch_idx:batch_idx+1], 0)
|
||||
byte_io = io.BytesIO()
|
||||
img.save(byte_io, format='PNG', compress_level=0)
|
||||
encoded_image = base64.b64encode(byte_io.getvalue()).decode('utf-8')
|
||||
payload = {
|
||||
"job_id": str(multi_job_id),
|
||||
"worker_id": str(worker_id),
|
||||
"batch_idx": int(batch_idx),
|
||||
"image": f"data:image/png;base64,{encoded_image}",
|
||||
"is_last": bool(batch_idx == batch_size - 1),
|
||||
}
|
||||
if payload["is_last"] and encoded_audio is not None:
|
||||
payload["audio"] = encoded_audio
|
||||
payloads.append(payload)
|
||||
|
||||
for payload in payloads:
|
||||
timeout_seconds = 60 if "image" in payload else 600
|
||||
try:
|
||||
async with session.post(
|
||||
url,
|
||||
json=payload,
|
||||
timeout=aiohttp.ClientTimeout(total=timeout_seconds),
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
except Exception as e:
|
||||
media_type = "image/audio" if "image" in payload else "audio-only"
|
||||
log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
|
||||
debug_log(f"Worker - Full error details: URL={url}")
|
||||
raise
|
||||
|
||||
def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
|
||||
"""Combine audio from master and workers into a single audio output.
|
||||
|
||||
Ordering: master first, then workers in `worker_order` (if provided),
|
||||
then any unexpected worker ids in sorted order.
|
||||
"""
|
||||
audio_pieces = []
|
||||
sample_rate = 44100
|
||||
|
||||
# Add master audio first if present
|
||||
if master_audio is not None:
|
||||
waveform = master_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
sample_rate = master_audio.get("sample_rate", 44100)
|
||||
|
||||
# Add worker audio in configured enabled-worker order first.
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
for worker_id_str in ordered_worker_ids:
|
||||
seen.add(worker_id_str)
|
||||
w_audio = worker_audio.get(worker_id_str)
|
||||
if w_audio is not None:
|
||||
waveform = w_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
# Use first available sample rate
|
||||
if sample_rate == 44100:
|
||||
sample_rate = w_audio.get("sample_rate", 44100)
|
||||
|
||||
# Append any audio from unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_audio.keys()):
|
||||
if worker_id_str in seen:
|
||||
continue
|
||||
w_audio = worker_audio[worker_id_str]
|
||||
if w_audio is not None:
|
||||
waveform = w_audio.get("waveform")
|
||||
if waveform is not None and waveform.numel() > 0:
|
||||
audio_pieces.append(waveform)
|
||||
if sample_rate == 44100:
|
||||
sample_rate = w_audio.get("sample_rate", 44100)
|
||||
|
||||
if not audio_pieces:
|
||||
return empty_audio
|
||||
|
||||
try:
|
||||
# Concatenate along the samples dimension (dim=-1)
|
||||
# Ensure all pieces have same batch and channel dimensions
|
||||
combined_waveform = torch.cat(audio_pieces, dim=-1)
|
||||
debug_log(f"Master - Combined audio: {len(audio_pieces)} pieces, final shape={combined_waveform.shape}")
|
||||
return {"waveform": combined_waveform, "sample_rate": sample_rate}
|
||||
except Exception as e:
|
||||
log(f"[Distributed] Master - Audio combination failed, returning silence: {e}")
|
||||
return empty_audio
|
||||
|
||||
def _store_worker_result(self, worker_images: dict, item: dict) -> int:
|
||||
"""Store one canonical queue item in worker_images in-place.
|
||||
|
||||
Canonical format:
|
||||
- item has 'worker_id', 'image_index', and 'tensor'
|
||||
Returns 1 when stored, otherwise 0.
|
||||
"""
|
||||
worker_id = item['worker_id']
|
||||
tensor = item.get('tensor')
|
||||
image_index = item.get('image_index')
|
||||
if tensor is None or image_index is None:
|
||||
return 0
|
||||
|
||||
worker_images.setdefault(worker_id, {})
|
||||
worker_images[worker_id][image_index] = tensor
|
||||
return 1
|
||||
|
||||
def _reorder_and_combine_tensors(
|
||||
self,
|
||||
worker_images: dict,
|
||||
worker_order: list,
|
||||
master_batch_size: int,
|
||||
images_on_cpu,
|
||||
delegate_mode: bool,
|
||||
fallback_images,
|
||||
):
|
||||
"""Assemble final tensor, or return None when the job contains only audio."""
|
||||
ordered_tensors = []
|
||||
if not delegate_mode and images_on_cpu is not None:
|
||||
for i in range(master_batch_size):
|
||||
ordered_tensors.append(images_on_cpu[i:i+1])
|
||||
|
||||
ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
|
||||
seen = set()
|
||||
for worker_id_str in ordered_worker_ids:
|
||||
seen.add(worker_id_str)
|
||||
if worker_id_str not in worker_images:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
|
||||
# Append any unexpected worker ids deterministically.
|
||||
for worker_id_str in sorted(worker_images.keys()):
|
||||
if worker_id_str in seen:
|
||||
continue
|
||||
for idx in sorted(worker_images[worker_id_str].keys()):
|
||||
ordered_tensors.append(worker_images[worker_id_str][idx])
|
||||
|
||||
cpu_tensors = []
|
||||
for t in ordered_tensors:
|
||||
if t.is_cuda:
|
||||
t = t.cpu()
|
||||
t = ensure_contiguous(t)
|
||||
cpu_tensors.append(t)
|
||||
|
||||
if cpu_tensors:
|
||||
return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
|
||||
if fallback_images is not None:
|
||||
return ensure_contiguous(fallback_images)
|
||||
return None
|
||||
|
||||
async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
|
||||
if is_worker:
|
||||
# Worker mode: send images and audio to master in a single batch
|
||||
image_count = 0 if images is None else images.shape[0]
|
||||
debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
|
||||
await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
else:
|
||||
delegate_mode = delegate_only or is_master_delegate_only()
|
||||
# Master mode: collect images and audio from workers
|
||||
enabled_workers_raw = json.loads(enabled_worker_ids)
|
||||
enabled_workers = []
|
||||
seen_enabled = set()
|
||||
for worker_id in enabled_workers_raw:
|
||||
worker_id_str = str(worker_id)
|
||||
if worker_id_str in seen_enabled:
|
||||
continue
|
||||
seen_enabled.add(worker_id_str)
|
||||
enabled_workers.append(worker_id_str)
|
||||
expected_workers = set(enabled_workers)
|
||||
serverless_workers = {
|
||||
worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
|
||||
}
|
||||
serverless_manager = None
|
||||
|
||||
def get_serverless_manager():
|
||||
nonlocal serverless_manager
|
||||
if serverless_workers and serverless_manager is None:
|
||||
serverless_manager = self._get_serverless_manager()
|
||||
return serverless_manager
|
||||
|
||||
num_workers = len(expected_workers)
|
||||
if num_workers == 0:
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
|
||||
# Create the queue before any expensive local work to avoid job_complete race.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_jobs:
|
||||
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
|
||||
debug_log(f"Master - Initialized queue early for job {multi_job_id}")
|
||||
else:
|
||||
existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
|
||||
debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
|
||||
|
||||
if delegate_mode:
|
||||
master_batch_size = 0
|
||||
images_on_cpu = None
|
||||
master_audio = None
|
||||
debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
|
||||
else:
|
||||
if images is None:
|
||||
images_on_cpu = None
|
||||
master_batch_size = 0
|
||||
else:
|
||||
images_on_cpu = ensure_contiguous(images.cpu())
|
||||
master_batch_size = images.shape[0]
|
||||
master_audio = audio # Keep master's audio for later
|
||||
debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
|
||||
|
||||
|
||||
# Initialize storage for collected images and audio
|
||||
worker_images = {} # Dict to store images by worker_id and index
|
||||
worker_audio = {} # Dict to store audio by worker_id
|
||||
|
||||
# Collect images until all workers report they're done
|
||||
collected_count = 0
|
||||
workers_done = set()
|
||||
|
||||
# Use unified worker timeout from config/UI with simple sliced waits
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
|
||||
last_activity = time.time()
|
||||
|
||||
|
||||
# Get queue size before starting
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
initial_size = q.qsize()
|
||||
|
||||
# NEW: Initialize progress bar for workers (total = num_workers)
|
||||
p = ProgressBar(num_workers)
|
||||
|
||||
def mark_worker_done(done_worker_id):
|
||||
done_worker_id = str(done_worker_id)
|
||||
if done_worker_id not in expected_workers:
|
||||
debug_log(
|
||||
f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
|
||||
)
|
||||
return
|
||||
if done_worker_id in workers_done:
|
||||
debug_log(
|
||||
f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
|
||||
)
|
||||
return
|
||||
workers_done.add(done_worker_id)
|
||||
p.update(1) # +1 per completed expected worker
|
||||
|
||||
def consume_result(result):
|
||||
nonlocal collected_count
|
||||
if "serverless_error" in result:
|
||||
error = result.get("serverless_error")
|
||||
message = str(error).strip() if error is not None else "unspecified serverless failure"
|
||||
participant = result.get("worker_id", "unknown")
|
||||
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
|
||||
|
||||
result_worker_id = str(result["worker_id"])
|
||||
is_last = result.get("is_last", False)
|
||||
collected_count += self._store_worker_result(worker_images, result)
|
||||
debug_log(
|
||||
f"Master - Got canonical result from worker {result_worker_id}, "
|
||||
f"image {result.get('image_index', 0)}, is_last={is_last}"
|
||||
)
|
||||
result_audio = result.get("audio")
|
||||
if result_audio is not None:
|
||||
worker_audio[result_worker_id] = result_audio
|
||||
debug_log(f"Master - Got audio from worker {result_worker_id}")
|
||||
if is_last:
|
||||
mark_worker_done(result_worker_id)
|
||||
|
||||
try:
|
||||
while len(workers_done) < num_workers:
|
||||
# Check for user interruption to abort collection promptly.
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
try:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
|
||||
consume_result(result)
|
||||
|
||||
# Real results refresh the regular-worker timeout baseline.
|
||||
last_activity = time.time()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
if (time.time() - last_activity) < base_timeout:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
continue
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
missing_workers = expected_workers - workers_done
|
||||
|
||||
# A queue delivery can race the timed wait. Process it before
|
||||
# declaring an owned provider slot stopped without a result.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
queue_has_results = current_q is not None and not current_q.empty()
|
||||
if queue_has_results:
|
||||
continue
|
||||
|
||||
missing_serverless = missing_workers & serverless_workers
|
||||
if missing_serverless:
|
||||
manager = get_serverless_manager()
|
||||
if manager is None:
|
||||
raise RuntimeError("Serverless manager is unavailable for pending collector slots")
|
||||
pending_serverless = []
|
||||
for participant_id in sorted(missing_serverless):
|
||||
if manager.is_pending(multi_job_id, participant_id):
|
||||
pending_serverless.append(participant_id)
|
||||
elif manager.owns(multi_job_id, participant_id):
|
||||
raise RuntimeError(
|
||||
f"Serverless participant {participant_id} stopped without delivering a result"
|
||||
)
|
||||
# Provider deadlines are managed by the serverless runtime;
|
||||
# do not apply the regular worker heartbeat timeout to it.
|
||||
if pending_serverless:
|
||||
continue
|
||||
|
||||
missing_regular = missing_workers - serverless_workers
|
||||
elapsed = time.time() - last_activity
|
||||
for missing_worker_id in sorted(missing_regular):
|
||||
log(
|
||||
"Master - Heartbeat timeout: "
|
||||
f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
if missing_regular:
|
||||
log(
|
||||
f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
|
||||
f"(elapsed={elapsed:.1f}s)"
|
||||
)
|
||||
|
||||
# Probe only regular worker /prompt endpoints.
|
||||
any_busy = False
|
||||
if missing_regular:
|
||||
try:
|
||||
cfg = load_config()
|
||||
cfg_workers = cfg.get('workers', [])
|
||||
for wid in sorted(missing_regular):
|
||||
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
|
||||
if not wrec:
|
||||
debug_log(f"Collector probe: worker {wid} not found in config")
|
||||
continue
|
||||
worker_url = build_worker_url(wrec)
|
||||
try:
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
queue_remaining = None
|
||||
if payload is not None:
|
||||
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
|
||||
debug_log(
|
||||
"Collector probe: worker "
|
||||
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
|
||||
)
|
||||
if payload is not None and queue_remaining and queue_remaining > 0:
|
||||
any_busy = True
|
||||
log(
|
||||
f"Master - Probe grace: worker {wid} appears busy "
|
||||
f"(queue_remaining={queue_remaining}). Continuing to wait."
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe failed for worker {wid}: {e}")
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe setup error: {e}")
|
||||
|
||||
if any_busy:
|
||||
last_activity = time.time()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
continue
|
||||
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
remaining_items = []
|
||||
if final_q is not None:
|
||||
while not final_q.empty():
|
||||
try:
|
||||
remaining_items.append(final_q.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
for item in remaining_items:
|
||||
consume_result(item)
|
||||
break
|
||||
|
||||
total_collected = sum(len(imgs) for imgs in worker_images.values())
|
||||
|
||||
# Remove the queue before combining so a late provider callback cannot
|
||||
# resurrect results for a collector that has completed.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
|
||||
|
||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||
serverless_results = set(worker_images) & serverless_workers
|
||||
try:
|
||||
combined = self._reorder_and_combine_tensors(
|
||||
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
||||
)
|
||||
if combined is None:
|
||||
debug_log(f"Master - Job {multi_job_id} complete with audio only")
|
||||
else:
|
||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
return (combined, combined_audio)
|
||||
except Exception as e:
|
||||
if serverless_results:
|
||||
raise RuntimeError(
|
||||
"Serverless result batch shape mismatch while combining participants "
|
||||
f"{sorted(serverless_results)}: {e}"
|
||||
) from e
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Preserve collected audio for regular workers when image assembly fails.
|
||||
return (images, combined_audio)
|
||||
except BaseException:
|
||||
# Remove the queue first, preventing cancelled work from delivering into
|
||||
# a collector that is already unwinding.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
|
||||
if serverless_workers:
|
||||
try:
|
||||
manager = get_serverless_manager()
|
||||
if manager is not None:
|
||||
await manager.cancel(collector_job_id=multi_job_id)
|
||||
except Exception as cancel_error:
|
||||
debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
|
||||
raise
|
||||
@@ -1,270 +0,0 @@
|
||||
import json
|
||||
import math
|
||||
from functools import wraps
|
||||
|
||||
import comfy.samplers
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.async_helpers import run_async_in_server_loop
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
|
||||
from ..upscale.tile_ops import TileOpsMixin
|
||||
from ..upscale.result_collector import ResultCollectorMixin
|
||||
from ..upscale.worker_comms import WorkerCommsMixin
|
||||
from ..upscale.job_state import JobStateMixin
|
||||
from ..upscale.modes.single_gpu import SingleGpuModeMixin
|
||||
from ..upscale.modes.static import StaticModeMixin
|
||||
from ..upscale.modes.dynamic import DynamicModeMixin
|
||||
|
||||
def sync_wrapper(async_func):
|
||||
"""Decorator to wrap async methods for synchronous execution."""
|
||||
@wraps(async_func)
|
||||
def sync_func(self, *args, **kwargs):
|
||||
# Use run_async_in_server_loop for ComfyUI compatibility
|
||||
return run_async_in_server_loop(
|
||||
async_func(self, *args, **kwargs),
|
||||
timeout=600.0 # 10 minute timeout for long operations
|
||||
)
|
||||
return sync_func
|
||||
|
||||
def _parse_enabled_worker_ids(enabled_worker_ids):
|
||||
"""Parse enabled worker IDs from either JSON or list input."""
|
||||
if isinstance(enabled_worker_ids, list):
|
||||
return [str(worker_id) for worker_id in enabled_worker_ids]
|
||||
if not enabled_worker_ids:
|
||||
return []
|
||||
if isinstance(enabled_worker_ids, str):
|
||||
try:
|
||||
parsed = json.loads(enabled_worker_ids)
|
||||
except json.JSONDecodeError:
|
||||
log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
|
||||
return []
|
||||
if isinstance(parsed, list):
|
||||
return [str(wid) for wid in parsed]
|
||||
return []
|
||||
|
||||
class UltimateSDUpscaleDistributed(
|
||||
DynamicModeMixin,
|
||||
StaticModeMixin,
|
||||
SingleGpuModeMixin,
|
||||
ResultCollectorMixin,
|
||||
WorkerCommsMixin,
|
||||
JobStateMixin,
|
||||
TileOpsMixin,
|
||||
):
|
||||
|
||||
"""
|
||||
Distributed version of Ultimate SD Upscale (No Upscale).
|
||||
|
||||
Supports two currently selected processing modes:
|
||||
1. Single GPU: No workers available, process everything locally
|
||||
2. Distributed tile queue: Workers pull tiles from a shared queue
|
||||
|
||||
Features:
|
||||
- Tile-based batch handling for video/image upscaling
|
||||
- Tiled VAE support for memory efficiency
|
||||
- Shared work queue so faster workers can process more tiles
|
||||
- Backward compatible with single-image workflows
|
||||
|
||||
Environment Variables:
|
||||
- COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
|
||||
- COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
|
||||
|
||||
The hidden dynamic_threshold input is retained for workflow compatibility but
|
||||
does not affect the current mode-selection policy.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the node and ensure persistent storage exists."""
|
||||
# Pre-initialize the persistent storage on node creation
|
||||
ensure_tile_jobs_initialized()
|
||||
debug_log("UltimateSDUpscaleDistributed - Node initialized")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"upscaled_image": ("IMAGE",),
|
||||
"model": ("MODEL",),
|
||||
"positive": ("CONDITIONING",),
|
||||
"negative": ("CONDITIONING",),
|
||||
"vae": ("VAE",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
|
||||
"denoise": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"tile_width": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||||
"tile_height": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 8}),
|
||||
"padding": ("INT", {"default": 32, "min": 0, "max": 256, "step": 8}),
|
||||
"mask_blur": ("INT", {"default": 8, "min": 0, "max": 256}),
|
||||
"force_uniform_tiles": ("BOOLEAN", {"default": True}),
|
||||
"tiled_decode": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": ("STRING", {"default": ""}),
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"master_url": ("STRING", {"default": ""}),
|
||||
"enabled_worker_ids": ("STRING", {"default": "[]"}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
"tile_indices": ("STRING", {"default": ""}), # Unused - kept for compatibility
|
||||
"dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "image/upscaling"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
"""Force re-execution."""
|
||||
return float("nan") # Always re-execute
|
||||
|
||||
def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode,
|
||||
multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]",
|
||||
worker_id="", tile_indices="", dynamic_threshold=8):
|
||||
"""Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
|
||||
# Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
|
||||
try:
|
||||
batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
|
||||
except Exception:
|
||||
batch_size = 1
|
||||
# Enforce 4n+1 batches globally for any model when batch > 1 (master only)
|
||||
if not is_worker and batch_size != 1 and (batch_size % 4 != 1):
|
||||
raise ValueError(
|
||||
f"Batch size {batch_size} is not of the form 4n+1. "
|
||||
"This node requires batch sizes of 1 or 4n+1 (1, 5, 9, 13, ...). "
|
||||
"Please adjust the batch size."
|
||||
)
|
||||
if not multi_job_id:
|
||||
# No distributed processing, run single GPU version
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
if is_worker:
|
||||
# Worker mode: process tiles synchronously
|
||||
return self.process_worker(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
else:
|
||||
# Master mode: distribute and collect synchronously
|
||||
return self.process_master(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold)
|
||||
|
||||
def process_worker(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold):
|
||||
"""Unified worker processing - handles both static and dynamic modes."""
|
||||
# Get batch size to determine mode
|
||||
batch_size = upscaled_image.shape[0]
|
||||
|
||||
# Ensure mode consistency across master/workers via shared threshold
|
||||
# Determine mode (must match master's logic)
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
# Compute number of tiles for this image to decide if tile distribution makes sense
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, self.round_to_multiple(tile_width), self.round_to_multiple(tile_height), force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# For USDU-style processing, we want tile distribution whenever workers are available
|
||||
# and there is more than one tile to process, even if batch == 1.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
debug_log(f"USDU Dist Worker - Batch size {batch_size}")
|
||||
|
||||
if mode == "dynamic":
|
||||
return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold)
|
||||
|
||||
# Static mode - enhanced with health monitoring and retry logic
|
||||
return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_workers)
|
||||
|
||||
def process_master(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids,
|
||||
dynamic_threshold):
|
||||
"""Unified master processing with enhanced monitoring and failure handling."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get image dimensions and batch size
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
|
||||
# Calculate all tiles and grid
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
rows = math.ceil(height / tile_height)
|
||||
cols = math.ceil(width / tile_width)
|
||||
log(
|
||||
f"USDU Dist: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({num_tiles_per_image} tiles/image) | Batch {batch_size}"
|
||||
)
|
||||
|
||||
# Parse enabled workers
|
||||
enabled_workers = json.loads(enabled_worker_ids)
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
# Determine processing mode
|
||||
mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
|
||||
# Prefer tile-based static distribution when workers are available and there are multiple tiles,
|
||||
# even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
|
||||
if num_workers > 0 and num_tiles_per_image > 1:
|
||||
mode = "static"
|
||||
|
||||
log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
|
||||
|
||||
if mode == "single_gpu":
|
||||
# No workers, process all tiles locally
|
||||
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
|
||||
|
||||
elif mode == "dynamic":
|
||||
# Dynamic mode for large batches
|
||||
return self.process_master_dynamic(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers)
|
||||
|
||||
# Static mode - enhanced with unified job management
|
||||
return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
|
||||
all_tiles, num_tiles_per_image)
|
||||
|
||||
def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
|
||||
"""Determines processing mode per requested policy:
|
||||
- any workers => prefer static (tile-based) for USDU
|
||||
- no workers => single_gpu
|
||||
"""
|
||||
if num_workers == 0:
|
||||
return "single_gpu"
|
||||
# Default to static when distributed; master/worker may still override if special cases arise
|
||||
return "static"
|
||||
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
@@ -1,354 +0,0 @@
|
||||
import torch
|
||||
import json
|
||||
|
||||
from ..utils.logging import debug_log, log
|
||||
|
||||
|
||||
def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]:
|
||||
"""Return contiguous [start, end) bounds for n_splits chunks."""
|
||||
split_count = max(1, int(n_splits))
|
||||
total = max(0, int(total_items))
|
||||
base, remainder = divmod(total, split_count)
|
||||
|
||||
bounds: list[tuple[int, int]] = []
|
||||
start = 0
|
||||
for idx in range(split_count):
|
||||
size = base + (1 if idx < remainder else 0)
|
||||
end = start + size
|
||||
bounds.append((start, end))
|
||||
start = end
|
||||
return bounds
|
||||
|
||||
|
||||
class DistributedSeed:
|
||||
"""
|
||||
Distributes seed values across multiple GPUs.
|
||||
On master: passes through the original seed.
|
||||
On workers: adds offset based on worker ID.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"seed": ("INT", {
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": False # Widget by default, can be converted to input
|
||||
}),
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("seed",)
|
||||
FUNCTION = "distribute"
|
||||
CATEGORY = "utils"
|
||||
|
||||
def distribute(self, seed, is_worker=False, worker_id=""):
|
||||
if not is_worker:
|
||||
# Master node: pass through original values
|
||||
debug_log(f"Distributor - Master: seed={seed}")
|
||||
return (seed,)
|
||||
else:
|
||||
# Worker node: apply offset based on worker index
|
||||
# Find worker index from enabled_worker_ids
|
||||
try:
|
||||
# Worker IDs are passed as "worker_0", "worker_1", etc.
|
||||
if worker_id.startswith("worker_"):
|
||||
worker_index = int(worker_id.split("_")[1])
|
||||
else:
|
||||
# Fallback: try to parse as direct index
|
||||
worker_index = int(worker_id)
|
||||
|
||||
offset = worker_index + 1
|
||||
new_seed = seed + offset
|
||||
debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
|
||||
return (new_seed,)
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
|
||||
# Fallback: return original seed
|
||||
return (seed,)
|
||||
|
||||
|
||||
# Define ByPassTypeTuple for flexible return types
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
any_type = AnyType("*")
|
||||
|
||||
|
||||
class DistributedValue:
|
||||
"""
|
||||
Outputs a different value per worker.
|
||||
On master: returns default_value.
|
||||
On workers: looks up the worker-specific value from a JSON map,
|
||||
falling back to default_value if not set.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"default_value": ("STRING", {"default": ""}),
|
||||
"worker_values": ("STRING", {"default": "{}"}),
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": ("BOOLEAN", {"default": False}),
|
||||
"worker_id": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
RETURN_NAMES = ("value",)
|
||||
FUNCTION = "distribute"
|
||||
CATEGORY = "utils"
|
||||
|
||||
@staticmethod
|
||||
def _coerce(value, value_type):
|
||||
"""Convert a string value to the requested type."""
|
||||
if value_type == "INT":
|
||||
return int(float(value))
|
||||
if value_type == "FLOAT":
|
||||
return float(value)
|
||||
return value # STRING and COMBO stay as strings
|
||||
|
||||
@staticmethod
|
||||
def _coerce_safe(value, value_type):
|
||||
"""Best-effort coercion with graceful fallback to original value."""
|
||||
try:
|
||||
return DistributedValue._coerce(value, value_type)
|
||||
except (TypeError, ValueError):
|
||||
return value
|
||||
|
||||
def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
|
||||
values = {}
|
||||
value_type = "STRING"
|
||||
|
||||
try:
|
||||
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
|
||||
if not isinstance(values, dict):
|
||||
values = {}
|
||||
except json.JSONDecodeError as e:
|
||||
debug_log(f"DistributedValue - Error parsing worker_values: {e}")
|
||||
values = {}
|
||||
|
||||
value_type = values.get("_type", "STRING")
|
||||
coerced_default = self._coerce_safe(default_value, value_type)
|
||||
|
||||
if not is_worker:
|
||||
debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
try:
|
||||
if worker_id.startswith("worker_"):
|
||||
idx = int(worker_id.split("_")[1])
|
||||
else:
|
||||
idx = int(worker_id)
|
||||
key = str(idx + 1) # worker_0 → key "1" (1-indexed)
|
||||
raw = values.get(key, "")
|
||||
if raw:
|
||||
coerced = self._coerce(raw, value_type)
|
||||
debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
|
||||
return (coerced,)
|
||||
except (ValueError, IndexError) as e:
|
||||
debug_log(f"DistributedValue - Error: {e}")
|
||||
debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
|
||||
return (coerced_default,)
|
||||
|
||||
class DistributedModelName:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", {"default": ""}),
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (any_type,)
|
||||
RETURN_NAMES = ("output",)
|
||||
FUNCTION = "log_input"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "utils"
|
||||
|
||||
def _stringify(self, value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (int, float, bool)):
|
||||
return str(value)
|
||||
try:
|
||||
return json.dumps(value, indent=4)
|
||||
except Exception:
|
||||
return str(value)
|
||||
|
||||
def _update_workflow(self, extra_pnginfo, unique_id, values):
|
||||
if not extra_pnginfo:
|
||||
return
|
||||
info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
|
||||
if not isinstance(info, dict) or "workflow" not in info:
|
||||
return
|
||||
node_id = None
|
||||
if isinstance(unique_id, list) and unique_id:
|
||||
node_id = str(unique_id[0])
|
||||
elif unique_id is not None:
|
||||
node_id = str(unique_id)
|
||||
if not node_id:
|
||||
return
|
||||
workflow = info["workflow"]
|
||||
node = next((x for x in workflow["nodes"] if str(x.get("id")) == node_id), None)
|
||||
if node:
|
||||
node["widgets_values"] = [values]
|
||||
|
||||
def log_input(self, text, unique_id=None, extra_pnginfo=None):
|
||||
values = []
|
||||
if isinstance(text, list):
|
||||
for val in text:
|
||||
values.append(self._stringify(val))
|
||||
else:
|
||||
values.append(self._stringify(text))
|
||||
|
||||
# Keep widget display in workflow metadata if available.
|
||||
self._update_workflow(extra_pnginfo, unique_id, values)
|
||||
|
||||
if isinstance(values, list) and len(values) == 1:
|
||||
return {"ui": {"text": values}, "result": (values[0],)}
|
||||
return {"ui": {"text": values}, "result": (values,)}
|
||||
|
||||
class ByPassTypeTuple(tuple):
|
||||
def __getitem__(self, index):
|
||||
if index > 0:
|
||||
index = 0
|
||||
item = super().__getitem__(index)
|
||||
if isinstance(item, str):
|
||||
return any_type
|
||||
return item
|
||||
|
||||
class ImageBatchDivider:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"divide_by": ("INT", {
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ByPassTypeTuple(("IMAGE", )) # Flexible for variable outputs
|
||||
RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)]))
|
||||
FUNCTION = "divide_batch"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "image"
|
||||
|
||||
def divide_batch(self, images, divide_by):
|
||||
total_splits = max(1, min(int(divide_by), 10))
|
||||
total_frames = images.shape[0]
|
||||
empty_tensor = images[:0]
|
||||
bounds = _chunk_bounds(total_frames, total_splits)
|
||||
outputs = [images[start:end] if end > start else empty_tensor for start, end in bounds]
|
||||
|
||||
while len(outputs) < 10:
|
||||
outputs.append(empty_tensor)
|
||||
|
||||
return tuple(outputs[:10])
|
||||
|
||||
|
||||
class AudioBatchDivider:
|
||||
"""Divides an audio waveform into sequential segments along the time dimension."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"divide_by": ("INT", {
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ByPassTypeTuple(("AUDIO",)) # Flexible for variable outputs
|
||||
RETURN_NAMES = ByPassTypeTuple(tuple([f"audio_{i+1}" for i in range(10)]))
|
||||
FUNCTION = "divide_audio"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "audio"
|
||||
|
||||
def divide_audio(self, audio, divide_by):
|
||||
import torch
|
||||
|
||||
waveform = audio.get("waveform")
|
||||
sample_rate = audio.get("sample_rate", 44100)
|
||||
|
||||
if waveform is None or waveform.numel() == 0:
|
||||
# Return empty audio for all outputs
|
||||
empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": sample_rate}
|
||||
return tuple([empty_audio] * 10)
|
||||
|
||||
total_splits = max(1, min(int(divide_by), 10))
|
||||
total_samples = int(waveform.shape[-1])
|
||||
bounds = _chunk_bounds(total_samples, total_splits)
|
||||
|
||||
outputs = []
|
||||
empty_waveform = waveform[..., :0]
|
||||
for start, end in bounds:
|
||||
split_waveform = waveform[..., start:end] if end > start else empty_waveform
|
||||
outputs.append({
|
||||
"waveform": split_waveform,
|
||||
"sample_rate": sample_rate
|
||||
})
|
||||
|
||||
# Pad with empty audio up to max (10) to match RETURN_TYPES length
|
||||
empty_audio = {
|
||||
"waveform": empty_waveform,
|
||||
"sample_rate": sample_rate
|
||||
}
|
||||
|
||||
while len(outputs) < 10:
|
||||
outputs.append(empty_audio)
|
||||
|
||||
return tuple(outputs)
|
||||
|
||||
|
||||
class DistributedEmptyImage:
|
||||
"""Produces an empty IMAGE batch used when the master delegates all work."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
|
||||
"width": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
|
||||
"channels": ("INT", {"default": 3, "min": 1, "max": 4, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "create"
|
||||
CATEGORY = "image"
|
||||
|
||||
def create(self, height, width, channels):
|
||||
import torch
|
||||
|
||||
shape = (0, height, width, channels)
|
||||
tensor = torch.zeros(shape, dtype=torch.float32)
|
||||
return (tensor,)
|
||||
-181
@@ -1,181 +0,0 @@
|
||||
"""Native V3 schemas with the existing execution algorithms kept intact.
|
||||
|
||||
Runtime objects are private, per-execution helpers, not registered V1 nodes.
|
||||
This avoids sharing mutable instance state through V3's sanitized class clones.
|
||||
"""
|
||||
import comfy.samplers
|
||||
from comfy_api.v0_0_2 import io
|
||||
|
||||
from . import utilities as _utilities
|
||||
from .collector import DistributedCollectorNode as _CollectorRuntime
|
||||
from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
|
||||
|
||||
|
||||
class DistributedSeed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
|
||||
inputs=[io.Int.Input('seed', default=1125899906842, min=0,
|
||||
max=1125899906842624, force_input=False)],
|
||||
outputs=[io.Int.Output(display_name='seed')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed, is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedValue(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedValue', display_name='Distributed Value', category='utils',
|
||||
inputs=[io.String.Input('default_value', default=''),
|
||||
io.String.Input('worker_values', default='{}')],
|
||||
outputs=[io.AnyType.Output(display_name='value')],
|
||||
accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
|
||||
return io.NodeOutput(*_utilities.DistributedValue().distribute(
|
||||
default_value, worker_values, is_worker, worker_id))
|
||||
|
||||
|
||||
class DistributedModelName(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
|
||||
inputs=[io.String.Input('text', default='')],
|
||||
outputs=[io.AnyType.Output(display_name='output')],
|
||||
hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text):
|
||||
result = _utilities.DistributedModelName().log_input(
|
||||
text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
|
||||
return io.NodeOutput(*result['result'], ui=result['ui'])
|
||||
|
||||
|
||||
class ImageBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
|
||||
inputs=[io.Image.Input('images'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of parts to divide the batch into')],
|
||||
# The existing frontend still displays only divide_by sockets.
|
||||
outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, divide_by):
|
||||
return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
|
||||
|
||||
|
||||
class AudioBatchDivider(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
|
||||
inputs=[io.Audio.Input('audio'),
|
||||
io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
|
||||
display_mode=io.NumberDisplay.number,
|
||||
tooltip='Number of sequential time segments to create')],
|
||||
outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, divide_by):
|
||||
return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
|
||||
|
||||
|
||||
class DistributedEmptyImage(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
|
||||
inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('width', default=64, min=1, max=4096, step=1),
|
||||
io.Int.Input('channels', default=3, min=1, max=4, step=1)],
|
||||
outputs=[io.Image.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, height, width, channels):
|
||||
return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
|
||||
|
||||
|
||||
class DistributedCollector(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='DistributedCollector', display_name='Distributed Collector', category='image',
|
||||
inputs=[io.Boolean.Input('load_balance', default=False,
|
||||
tooltip='Run this workflow on one least-busy participant (master included when participating).'),
|
||||
io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True)],
|
||||
outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
|
||||
is_input_list=True, accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
|
||||
is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
|
||||
worker_id='', pass_through=False, delegate_only=False):
|
||||
return io.NodeOutput(*_CollectorRuntime().run(
|
||||
images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
|
||||
is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
|
||||
worker_batch_size=worker_batch_size, worker_id=worker_id,
|
||||
pass_through=pass_through, delegate_only=delegate_only))
|
||||
|
||||
|
||||
class UltimateSDUpscaleDistributed(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id='UltimateSDUpscaleDistributed',
|
||||
display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
|
||||
inputs=[
|
||||
io.Image.Input('upscaled_image'), io.Model.Input('model'),
|
||||
io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
|
||||
io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
|
||||
io.Int.Input('steps', default=20, min=1, max=10000),
|
||||
io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
|
||||
io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
|
||||
io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
|
||||
io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
|
||||
io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
|
||||
io.Int.Input('padding', default=32, min=0, max=256, step=8),
|
||||
io.Int.Input('mask_blur', default=8, min=0, max=256),
|
||||
io.Boolean.Input('force_uniform_tiles', default=True),
|
||||
io.Boolean.Input('tiled_decode', default=False),
|
||||
], outputs=[io.Image.Output()], accept_all_inputs=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, **kwargs):
|
||||
return _UpscaleRuntime.IS_CHANGED(**kwargs)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
|
||||
master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
|
||||
return io.NodeOutput(*_UpscaleRuntime().run(
|
||||
upscaled_image, model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height, padding,
|
||||
mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
|
||||
master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
|
||||
|
||||
|
||||
NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
|
||||
ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed]
|
||||
Generated
-1419
File diff suppressed because it is too large
Load Diff
@@ -1,12 +0,0 @@
|
||||
{
|
||||
"name": "comfyui-distributed-web-tests",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"test:web": "bash ./scripts/test-web.sh",
|
||||
"test:web:watch": "bash ./scripts/test-web.sh --watch"
|
||||
},
|
||||
"devDependencies": {
|
||||
"vitest": "^2.1.9"
|
||||
}
|
||||
}
|
||||
+1
-5
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.5.0"
|
||||
version = "1.2.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
@@ -13,7 +13,3 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
PublisherId = "robertvoy"
|
||||
DisplayName = "ComfyUI-Distributed"
|
||||
Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
pythonpath = ["."]
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Internal extension lifecycle support (not a node provider)."""
|
||||
@@ -1,35 +0,0 @@
|
||||
"""Initialize routes, distributed state and the existing worker lifecycle."""
|
||||
import atexit
|
||||
import os
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.config import CONFIG_FILE, ensure_config_exists
|
||||
from ..utils.logging import debug_log
|
||||
from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
|
||||
from ..upscale.job_store import ensure_tile_jobs_initialized
|
||||
|
||||
_initialized = False
|
||||
|
||||
|
||||
def initialize():
|
||||
"""Called by ComfyExtension.on_load; initialize once per loaded package."""
|
||||
global _initialized
|
||||
if _initialized:
|
||||
return
|
||||
|
||||
ensure_config_exists()
|
||||
from .. import api # noqa: F401 - registers the existing @routes.* handlers
|
||||
from ..api.queue_orchestration import ensure_distributed_state
|
||||
|
||||
ensure_distributed_state(server.PromptServer.instance)
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
if not os.environ.get('COMFYUI_IS_WORKER'):
|
||||
atexit.register(sync_cleanup)
|
||||
delayed_auto_launch()
|
||||
register_async_signals()
|
||||
|
||||
_initialized = True
|
||||
debug_log('Loaded Distributed nodes.')
|
||||
debug_log(f'Config file: {CONFIG_FILE}')
|
||||
@@ -1,29 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd -- "${SCRIPT_DIR}/.." && pwd)"
|
||||
|
||||
export NVM_DIR="${NVM_DIR:-$HOME/.nvm}"
|
||||
if [[ -s "${NVM_DIR}/nvm.sh" ]]; then
|
||||
# shellcheck source=/dev/null
|
||||
. "${NVM_DIR}/nvm.sh"
|
||||
fi
|
||||
|
||||
if ! command -v node >/dev/null 2>&1 || ! command -v npm >/dev/null 2>&1; then
|
||||
echo "[test-web] node/npm are not available." >&2
|
||||
echo "[test-web] Install nvm and Node, or ensure node/npm are on PATH." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ -f "${REPO_ROOT}/.nvmrc" ]] && command -v nvm >/dev/null 2>&1; then
|
||||
nvm use >/dev/null
|
||||
fi
|
||||
|
||||
cd "${REPO_ROOT}"
|
||||
|
||||
if [[ "${1:-}" == "--watch" ]]; then
|
||||
exec npx vitest web/tests
|
||||
fi
|
||||
|
||||
exec npx vitest run web/tests
|
||||
@@ -1,147 +0,0 @@
|
||||
import copy
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None):
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def _load_config_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
|
||||
package_name = "dist_api_config_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = [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
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
default_config = {
|
||||
"workers": [],
|
||||
"master": {"host": ""},
|
||||
"settings": {"debug": False},
|
||||
"tunnel": {},
|
||||
}
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: copy.deepcopy(default_config)
|
||||
config_module.save_config = lambda _cfg: True
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
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
|
||||
|
||||
|
||||
config_routes = _load_config_routes_module()
|
||||
|
||||
|
||||
class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_get_config_returns_core_sections(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIn("workers", response.payload)
|
||||
self.assertIn("master", response.payload)
|
||||
self.assertIn("settings", response.payload)
|
||||
|
||||
async def test_update_config_valid_field_persists(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
|
||||
config_routes, "save_config", return_value=True
|
||||
):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload["status"], "success")
|
||||
self.assertTrue(response.payload["config"]["settings"]["debug"])
|
||||
|
||||
async def test_update_config_unknown_field_returns_400(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"unknown_field": 1}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("unknown_field", " ".join(response.payload.get("error", [])).lower())
|
||||
|
||||
async def test_update_config_wrong_type_returns_400(self):
|
||||
cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
|
||||
with patch.object(config_routes, "load_config", return_value=cfg):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"debug": "true"}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("debug", " ".join(response.payload.get("error", [])).lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,402 +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
|
||||
|
||||
|
||||
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 _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
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,247 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_media_sync_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
|
||||
package_name = "dist_ms_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost{endpoint}"
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in pure-function tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
|
||||
trace_module.trace_debug = lambda *_args, **_kwargs: None
|
||||
trace_module.trace_info = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
class _FormData:
|
||||
def add_field(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.FormData = _FormData
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.media_sync",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
ms = _load_media_sync_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# convert_paths_for_platform
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ConvertPathsForPlatformTests(unittest.TestCase):
|
||||
def test_forward_slash_target_normalises_backslashes(self):
|
||||
obj = {"ckpt_name": "C:\\Models\\model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_backslash_target_normalises_forward_slashes(self):
|
||||
obj = {"ckpt_name": "/models/checkpoints/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("\\", result["ckpt_name"])
|
||||
self.assertNotIn("/", result["ckpt_name"])
|
||||
|
||||
def test_relative_media_paths_always_stay_forward_slash(self):
|
||||
"""Relative image/video/audio paths (Comfy annotated style) must not be backslash-ified."""
|
||||
obj = {"image": "subfolder/my_photo.png"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["image"], "subfolder/my_photo.png")
|
||||
|
||||
def test_relative_audio_paths_stay_forward_slash(self):
|
||||
obj = {"audio": "subfolder/my_track.wav"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["audio"], "subfolder/my_track.wav")
|
||||
|
||||
def test_annotated_relative_media_path_stays_forward_slash(self):
|
||||
obj = {"image": "input/frame.jpg [abc123]"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertIn("/", result["image"])
|
||||
self.assertNotIn("\\", result["image"].split("[")[0])
|
||||
|
||||
def test_non_filename_strings_are_untouched(self):
|
||||
obj = {"prompt": "a beautiful cat", "count": 5}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["prompt"], "a beautiful cat")
|
||||
self.assertEqual(result["count"], 5)
|
||||
|
||||
def test_url_strings_are_untouched(self):
|
||||
obj = {"url": "https://example.com/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertEqual(result["url"], "https://example.com/model.safetensors")
|
||||
|
||||
def test_invalid_separator_returns_obj_unchanged(self):
|
||||
obj = {"ckpt_name": "/models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "|")
|
||||
self.assertEqual(result, obj)
|
||||
|
||||
def test_nested_dict_is_processed_recursively(self):
|
||||
obj = {"node": {"ckpt_name": "C:\\Models\\model.safetensors"}}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["node"]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
|
||||
def test_list_items_are_processed_recursively(self):
|
||||
obj = [{"ckpt_name": "C:\\Models\\model.safetensors"}, "plain string"]
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result[0]["ckpt_name"], "C:/Models/model.safetensors")
|
||||
self.assertEqual(result[1], "plain string")
|
||||
|
||||
def test_non_string_scalar_values_are_untouched(self):
|
||||
obj = {"seed": 42, "enabled": True, "ratio": 1.5}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["seed"], 42)
|
||||
self.assertTrue(result["enabled"])
|
||||
|
||||
def test_absolute_unix_path_to_windows(self):
|
||||
obj = {"lora": "/home/user/loras/my_lora.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "\\")
|
||||
self.assertNotIn("/", result["lora"])
|
||||
|
||||
def test_already_normalised_path_is_idempotent(self):
|
||||
obj = {"ckpt": "C:/Models/model.safetensors"}
|
||||
result = ms.convert_paths_for_platform(obj, "/")
|
||||
self.assertEqual(result["ckpt"], "C:/Models/model.safetensors")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_media_references
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindMediaReferencesTests(unittest.TestCase):
|
||||
def test_finds_file_input_for_load_video(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("1 - Copy.mp4", refs)
|
||||
|
||||
def test_strips_annotation_suffix(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.jpg", refs)
|
||||
self.assertFalse(any("[" in r for r in refs))
|
||||
|
||||
def test_normalises_backslashes_in_path(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "sub\\img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("sub/img.png", refs)
|
||||
|
||||
def test_ignores_non_media_text_inputs(self):
|
||||
prompt = {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_ignores_node_link_values(self):
|
||||
"""Inputs that are [node_id, slot] lists should be ignored."""
|
||||
prompt = {"1": {"class_type": "Anything", "inputs": {"image": ["2", 0]}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, [])
|
||||
|
||||
def test_deduplicates_same_file_across_nodes(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(len(refs), 1)
|
||||
|
||||
def test_returns_sorted_list(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "z_image.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "a_image.jpg"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertEqual(refs, sorted(refs))
|
||||
|
||||
def test_ignores_non_dict_nodes(self):
|
||||
prompt = {"1": "not a node dict", "2": {"class_type": "LoadImage", "inputs": {"image": "img.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("img.png", refs)
|
||||
|
||||
def test_empty_prompt_returns_empty_list(self):
|
||||
self.assertEqual(ms._find_media_references({}), [])
|
||||
|
||||
def test_multiple_media_types_all_found(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "frame.png"}},
|
||||
"2": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}},
|
||||
"3": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}},
|
||||
}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("frame.png", refs)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
self.assertIn("track.wav", refs)
|
||||
|
||||
|
||||
class RewritePromptMediaInputsTests(unittest.TestCase):
|
||||
def test_rewrites_video_file_input_to_worker_path(self):
|
||||
prompt = {
|
||||
"79": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}},
|
||||
}
|
||||
ms._rewrite_prompt_media_inputs(prompt, {"1 - Copy.mp4": "videos/1 - Copy.mp4"})
|
||||
self.assertEqual(prompt["79"]["inputs"]["file"], "videos/1 - Copy.mp4")
|
||||
|
||||
def test_rewrites_audio_input_and_strips_annotation_when_matching(self):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadAudio", "inputs": {"audio": "song.wav [input]"}},
|
||||
}
|
||||
ms._rewrite_prompt_media_inputs(prompt, {"song.wav": "song.wav"})
|
||||
self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,304 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, json_payload=None, post_payload=None, headers=None, query=None):
|
||||
self._json_payload = json_payload
|
||||
self._post_payload = post_payload or {}
|
||||
self.headers = headers or {}
|
||||
self.query = query or {}
|
||||
|
||||
async def json(self):
|
||||
return self._json_payload
|
||||
|
||||
async def post(self):
|
||||
return self._post_payload
|
||||
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
|
||||
def _load_usdu_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
|
||||
package_name = "dist_api_usdu_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.MAX_PAYLOAD_SIZE = 1024
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class BaseJobState:
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
|
||||
job_models_module.BaseJobState = BaseJobState
|
||||
job_models_module.TileJobState = TileJobState
|
||||
job_models_module.ImageJobState = ImageJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
|
||||
parsers_module._parse_tiles_from_form = lambda _data: []
|
||||
sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.usdu_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._TileJobState = TileJobState
|
||||
module._ImageJobState = ImageJobState
|
||||
return module
|
||||
|
||||
|
||||
usdu_routes = _load_usdu_routes_module()
|
||||
|
||||
|
||||
class _UploadField:
|
||||
def __init__(self, data):
|
||||
self.file = io.BytesIO(data)
|
||||
|
||||
|
||||
def _tiny_png_bytes():
|
||||
image = Image.new("RGB", (1, 1), (255, 0, 0))
|
||||
buf = io.BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
usdu_routes._prompt_server_holder["value"] = types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
|
||||
async def test_heartbeat_updates_worker_status(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-1")
|
||||
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-1"})
|
||||
response = await usdu_routes.heartbeat_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertIn("worker-a", job_data.worker_status)
|
||||
|
||||
async def test_heartbeat_missing_fields_returns_400(self):
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a"})
|
||||
response = await usdu_routes.heartbeat_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("missing", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_request_image_dynamic_assigns_next_image(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._ImageJobState("job-2")
|
||||
await job_data.pending_images.put(7)
|
||||
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-2"})
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("image_idx"), 7)
|
||||
self.assertEqual(response.payload.get("estimated_remaining"), 0)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
|
||||
self.assertIn("worker-a", job_data.worker_status)
|
||||
|
||||
async def test_request_image_static_assigns_tile_and_batched_flag(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-3")
|
||||
job_data.batched_static = True
|
||||
await job_data.pending_tasks.put(4)
|
||||
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
||||
|
||||
request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "job-3"})
|
||||
response = await usdu_routes.request_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("tile_idx"), 4)
|
||||
self.assertTrue(response.payload.get("batched_static"))
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
|
||||
|
||||
async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._TileJobState("job-4")
|
||||
prompt_server.distributed_pending_tile_jobs["job-4"] = job_data
|
||||
|
||||
request = _FakeRequest(
|
||||
post_payload={
|
||||
"multi_job_id": "job-4",
|
||||
"worker_id": "worker-a",
|
||||
"batch_size": "0",
|
||||
"is_last": "true",
|
||||
},
|
||||
headers={"content-length": "128"},
|
||||
)
|
||||
|
||||
response = await usdu_routes.submit_tiles_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await job_data.queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-a")
|
||||
self.assertTrue(queued["is_last"])
|
||||
self.assertEqual(queued["tiles"], [])
|
||||
|
||||
async def test_submit_image_enqueues_processed_image_payload(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
job_data = usdu_routes._ImageJobState("job-5")
|
||||
prompt_server.distributed_pending_tile_jobs["job-5"] = job_data
|
||||
|
||||
request = _FakeRequest(
|
||||
post_payload={
|
||||
"multi_job_id": "job-5",
|
||||
"worker_id": "worker-a",
|
||||
"image_idx": "2",
|
||||
"full_image": _UploadField(_tiny_png_bytes()),
|
||||
"is_last": "false",
|
||||
},
|
||||
headers={"content-length": "256"},
|
||||
)
|
||||
|
||||
response = await usdu_routes.submit_image_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
queued = await job_data.queue.get()
|
||||
self.assertEqual(queued["worker_id"], "worker-a")
|
||||
self.assertEqual(queued["image_idx"], 2)
|
||||
self.assertIn("image", queued)
|
||||
|
||||
async def test_job_status_endpoint_reports_ready(self):
|
||||
prompt_server = usdu_routes._prompt_server_holder["value"]
|
||||
prompt_server.distributed_pending_tile_jobs["job-6"] = usdu_routes._TileJobState("job-6")
|
||||
|
||||
request = _FakeRequest(query={"multi_job_id": "job-6"})
|
||||
response = await usdu_routes.job_status_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload.get("ready"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,437 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = payload
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None, match_info=None, query=None):
|
||||
self._payload = payload
|
||||
self.match_info = match_info or {}
|
||||
self.query = query or {}
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
class _FakeHTTPClientResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self._payload = payload
|
||||
self.status = status
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||||
return False
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
async def text(self):
|
||||
return str(self._payload)
|
||||
|
||||
|
||||
class _FakeHTTPClientSession:
|
||||
def __init__(self, payload, status=200):
|
||||
self._payload = payload
|
||||
self._status = status
|
||||
self.calls = []
|
||||
|
||||
def get(self, url, params=None, timeout=None):
|
||||
self.calls.append({"url": url, "params": params, "timeout": timeout})
|
||||
return _FakeHTTPClientResponse(self._payload, status=self._status)
|
||||
|
||||
|
||||
class _DummyWorkerManager:
|
||||
def __init__(self):
|
||||
self.processes = {}
|
||||
|
||||
def launch_worker(self, worker):
|
||||
worker_id = str(worker["id"])
|
||||
self.processes[worker_id] = {
|
||||
"pid": 12345,
|
||||
"log_file": f"/tmp/distributed_worker_{worker_id}.log",
|
||||
"process": None,
|
||||
}
|
||||
return 12345
|
||||
|
||||
def _is_process_running(self, _pid):
|
||||
return False
|
||||
|
||||
def save_processes(self):
|
||||
return None
|
||||
|
||||
def stop_worker(self, _worker_id):
|
||||
return True, "Stopped"
|
||||
|
||||
def get_managed_workers(self):
|
||||
return []
|
||||
|
||||
|
||||
class _ImmediateLoop:
|
||||
async def run_in_executor(self, _executor, func, *args):
|
||||
return func(*args)
|
||||
|
||||
|
||||
def _load_worker_routes_module():
|
||||
module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
|
||||
package_name = "dist_api_worker_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
|
||||
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
detection_module = types.ModuleType(f"{package_name}.workers.detection")
|
||||
detection_module.is_local_worker = lambda *_args, **_kwargs: True
|
||||
detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
|
||||
detection_module.get_machine_id = lambda: "machine-id"
|
||||
detection_module.is_docker_environment = lambda: False
|
||||
detection_module.is_runpod_environment = lambda: False
|
||||
detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
|
||||
sys.modules[f"{package_name}.workers.detection"] = detection_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _WebSocketResponse:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
async def prepare(self, _request):
|
||||
return None
|
||||
|
||||
async def send_json(self, _payload):
|
||||
return None
|
||||
|
||||
def __aiter__(self):
|
||||
async def _empty():
|
||||
if False:
|
||||
yield None
|
||||
return _empty()
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
|
||||
WebSocketResponse=_WebSocketResponse,
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
def post(self, _path):
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
return _decorator
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
created_torch_stub = False
|
||||
if "torch" not in sys.modules:
|
||||
created_torch_stub = True
|
||||
torch_module = types.ModuleType("torch")
|
||||
torch_module.cuda = types.SimpleNamespace(
|
||||
is_available=lambda: False,
|
||||
empty_cache=lambda: None,
|
||||
ipc_collect=lambda: None,
|
||||
current_device=lambda: 0,
|
||||
device_count=lambda: 0,
|
||||
)
|
||||
sys.modules["torch"] = torch_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def _handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
||||
|
||||
network_module.handle_api_error = _handle_api_error
|
||||
network_module.normalize_host = lambda value: value
|
||||
network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"
|
||||
|
||||
async def _probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
|
||||
async def _get_client_session():
|
||||
raise RuntimeError("not used in these tests")
|
||||
|
||||
network_module.get_client_session = _get_client_session
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.CHUNK_SIZE = 8192
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
async def _queue_prompt_payload(*_args, **_kwargs):
|
||||
return "prompt-id"
|
||||
|
||||
async_helpers_module.queue_prompt_payload = _queue_prompt_payload
|
||||
|
||||
class _PromptValidationError(RuntimeError):
|
||||
def __init__(self, message="invalid prompt", validation_error=None, node_errors=None):
|
||||
super().__init__(message)
|
||||
self.validation_error = validation_error if isinstance(validation_error, dict) else {}
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
async_helpers_module.PromptValidationError = _PromptValidationError
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
schemas_module = types.ModuleType(f"{package_name}.api.schemas")
|
||||
|
||||
def _require_fields(data, *fields):
|
||||
missing = []
|
||||
for field in fields:
|
||||
value = data.get(field) if isinstance(data, dict) else None
|
||||
if value is None or (isinstance(value, str) and not value.strip()):
|
||||
missing.append(field)
|
||||
return missing
|
||||
|
||||
def _validate_worker_id(worker_id, config):
|
||||
return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
|
||||
|
||||
schemas_module.require_fields = _require_fields
|
||||
schemas_module.validate_worker_id = _validate_worker_id
|
||||
sys.modules[f"{package_name}.api.schemas"] = schemas_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
# Keep this module's response double independent of aiohttp import order.
|
||||
module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
if created_torch_stub:
|
||||
sys.modules.pop("torch", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_network_info_reports_actual_master_and_available_ports(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": []}), \
|
||||
patch.object(worker_routes, "get_server_port", return_value=8189), \
|
||||
patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["master_port"], 8189)
|
||||
self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
|
||||
self.assertEqual(response.payload["cuda_device_count"], 4)
|
||||
allocate.assert_called_once_with(8189, [], 3)
|
||||
|
||||
async def test_network_info_does_not_reallocate_existing_configuration(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=[]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
|
||||
patch.object(worker_routes, "allocate_worker_ports") as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["local_worker_ports"], [])
|
||||
allocate.assert_not_called()
|
||||
|
||||
async def test_launch_conflict_is_a_clear_client_error(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
|
||||
patch.object(worker_routes, "load_config", return_value=config), \
|
||||
patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
|
||||
response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("master port 8189", response.payload["message"])
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
async def test_launch_worker_valid_id_returns_200(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
||||
request = _FakeRequest({"worker_id": "worker-a"})
|
||||
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes, "load_config", return_value=config
|
||||
), patch.object(
|
||||
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
||||
):
|
||||
response = await worker_routes.launch_worker_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertEqual(response.payload.get("pid"), 12345)
|
||||
|
||||
async def test_launch_worker_unknown_id_returns_404(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
||||
request = _FakeRequest({"worker_id": "missing-worker"})
|
||||
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes, "load_config", return_value=config
|
||||
):
|
||||
response = await worker_routes.launch_worker_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 404)
|
||||
self.assertIn("not found", response.payload.get("message", "").lower())
|
||||
|
||||
async def test_worker_log_returns_content_json(self):
|
||||
manager = _DummyWorkerManager()
|
||||
with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
|
||||
handle.write("line-1\nline-2\nline-3\n")
|
||||
log_path = handle.name
|
||||
|
||||
manager.processes["worker-a"] = {
|
||||
"pid": 9999,
|
||||
"log_file": log_path,
|
||||
"process": None,
|
||||
}
|
||||
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
|
||||
try:
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
|
||||
worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
|
||||
):
|
||||
response = await worker_routes.get_worker_log_endpoint(request)
|
||||
finally:
|
||||
if os.path.exists(log_path):
|
||||
os.remove(log_path)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertIn("content", response.payload)
|
||||
self.assertIn("line-3", response.payload["content"])
|
||||
|
||||
async def test_local_log_reads_memory_buffer(self):
|
||||
request = _FakeRequest(query={"lines": "2"})
|
||||
fake_logs = deque(
|
||||
[
|
||||
{"m": "line-1\n"},
|
||||
{"m": "line-2\n"},
|
||||
{"m": "line-3\n"},
|
||||
],
|
||||
maxlen=300,
|
||||
)
|
||||
app_module = types.ModuleType("app")
|
||||
app_module.__path__ = []
|
||||
logger_module = types.ModuleType("app.logger")
|
||||
logger_module.get_logs = lambda: fake_logs
|
||||
app_module.logger = logger_module
|
||||
|
||||
with patch.dict(sys.modules, {"app": app_module, "app.logger": logger_module}):
|
||||
response = await worker_routes.get_local_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("status"), "success")
|
||||
self.assertEqual(response.payload.get("source"), "memory")
|
||||
self.assertEqual(response.payload.get("entries"), 2)
|
||||
self.assertIn("line-3", response.payload.get("content", ""))
|
||||
|
||||
async def test_remote_worker_log_proxies_to_worker_local_log_endpoint(self):
|
||||
config = {
|
||||
"workers": [
|
||||
{
|
||||
"id": "worker-remote",
|
||||
"name": "Remote Worker",
|
||||
"host": "worker.example.com",
|
||||
"port": 8188,
|
||||
"type": "remote",
|
||||
}
|
||||
]
|
||||
}
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-remote"}, query={"lines": "120"})
|
||||
proxied_payload = {
|
||||
"status": "success",
|
||||
"content": "remote-log-content\n",
|
||||
"entries": 1,
|
||||
"source": "memory",
|
||||
"truncated": False,
|
||||
"lines_shown": 1,
|
||||
}
|
||||
fake_session = _FakeHTTPClientSession(proxied_payload)
|
||||
|
||||
async def _fake_get_client_session():
|
||||
return fake_session
|
||||
|
||||
with patch.object(worker_routes, "load_config", return_value=config), patch.object(
|
||||
worker_routes, "get_client_session", side_effect=_fake_get_client_session
|
||||
):
|
||||
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(response.payload.get("content"), "remote-log-content\n")
|
||||
self.assertEqual(len(fake_session.calls), 1)
|
||||
self.assertEqual(fake_session.calls[0]["params"], {"lines": "120"})
|
||||
self.assertTrue(fake_session.calls[0]["url"].endswith("/distributed/local_log"))
|
||||
|
||||
async def test_remote_worker_log_rejects_local_workers(self):
|
||||
config = {"workers": [{"id": "worker-local", "name": "Local Worker", "port": 8188}]}
|
||||
request = _FakeRequest(match_info={"worker_id": "worker-local"})
|
||||
|
||||
with patch.object(worker_routes, "load_config", return_value=config):
|
||||
response = await worker_routes.get_remote_worker_log_endpoint(request)
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("local", response.payload.get("message", "").lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,2 +0,0 @@
|
||||
# This conftest.py marks the tests/ directory as the pytest collection root,
|
||||
# preventing pytest from traversing into the parent package's __init__.py.
|
||||
-691
@@ -1,691 +0,0 @@
|
||||
{
|
||||
"DistributedCollector": {
|
||||
"input": {
|
||||
"required": {
|
||||
"load_balance": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
|
||||
}
|
||||
]
|
||||
},
|
||||
"optional": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_batch_size": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1024
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"pass_through": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"delegate_only": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"load_balance"
|
||||
],
|
||||
"optional": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_batch_size",
|
||||
"worker_id",
|
||||
"pass_through",
|
||||
"delegate_only"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE",
|
||||
"AUDIO"
|
||||
],
|
||||
"output_name": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": true,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Collector"
|
||||
},
|
||||
"DistributedSeed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"seed"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"INT"
|
||||
],
|
||||
"output_name": [
|
||||
"seed"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Seed"
|
||||
},
|
||||
"DistributedModelName": {
|
||||
"input": {
|
||||
"required": {
|
||||
"text": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"text"
|
||||
],
|
||||
"hidden": [
|
||||
"unique_id",
|
||||
"extra_pnginfo"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"output"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Model Name"
|
||||
},
|
||||
"DistributedValue": {
|
||||
"input": {
|
||||
"required": {
|
||||
"default_value": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"worker_values": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "{}"
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"default_value",
|
||||
"worker_values"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"value"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Value"
|
||||
},
|
||||
"ImageBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"images",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"batch_1",
|
||||
"batch_2",
|
||||
"batch_3",
|
||||
"batch_4",
|
||||
"batch_5",
|
||||
"batch_6",
|
||||
"batch_7",
|
||||
"batch_8",
|
||||
"batch_9",
|
||||
"batch_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "image",
|
||||
"display_name": "Image Batch Divider"
|
||||
},
|
||||
"AudioBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"audio",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"audio_1",
|
||||
"audio_2",
|
||||
"audio_3",
|
||||
"audio_4",
|
||||
"audio_5",
|
||||
"audio_6",
|
||||
"audio_7",
|
||||
"audio_8",
|
||||
"audio_9",
|
||||
"audio_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "audio",
|
||||
"display_name": "Audio Segment Divider"
|
||||
},
|
||||
"DistributedEmptyImage": {
|
||||
"input": {
|
||||
"required": {
|
||||
"height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"channels": [
|
||||
"INT",
|
||||
{
|
||||
"default": 3,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"height",
|
||||
"width",
|
||||
"channels"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Empty Image"
|
||||
},
|
||||
"UltimateSDUpscaleDistributed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"upscaled_image": [
|
||||
"IMAGE"
|
||||
],
|
||||
"model": [
|
||||
"MODEL"
|
||||
],
|
||||
"positive": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"negative": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"vae": [
|
||||
"VAE"
|
||||
],
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 18446744073709551615
|
||||
}
|
||||
],
|
||||
"steps": [
|
||||
"INT",
|
||||
{
|
||||
"default": 20,
|
||||
"min": 1,
|
||||
"max": 10000
|
||||
}
|
||||
],
|
||||
"cfg": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0
|
||||
}
|
||||
],
|
||||
"sampler_name": [
|
||||
[
|
||||
"euler",
|
||||
"euler_cfg_pp",
|
||||
"euler_ancestral",
|
||||
"euler_ancestral_cfg_pp",
|
||||
"heun",
|
||||
"heunpp2",
|
||||
"exp_heun_2_x0",
|
||||
"exp_heun_2_x0_sde",
|
||||
"dpm_2",
|
||||
"dpm_2_ancestral",
|
||||
"lms",
|
||||
"dpm_fast",
|
||||
"dpm_adaptive",
|
||||
"dpmpp_2s_ancestral",
|
||||
"dpmpp_2s_ancestral_cfg_pp",
|
||||
"dpmpp_sde",
|
||||
"dpmpp_sde_gpu",
|
||||
"dpmpp_2m",
|
||||
"dpmpp_2m_cfg_pp",
|
||||
"dpmpp_2m_sde",
|
||||
"dpmpp_2m_sde_gpu",
|
||||
"dpmpp_2m_sde_heun",
|
||||
"dpmpp_2m_sde_heun_gpu",
|
||||
"dpmpp_3m_sde",
|
||||
"dpmpp_3m_sde_gpu",
|
||||
"ddpm",
|
||||
"lcm",
|
||||
"ipndm",
|
||||
"ipndm_v",
|
||||
"deis",
|
||||
"cfgpp_ud10_ab",
|
||||
"res_multistep",
|
||||
"res_multistep_cfg_pp",
|
||||
"res_multistep_ancestral",
|
||||
"res_multistep_ancestral_cfg_pp",
|
||||
"gradient_estimation",
|
||||
"gradient_estimation_cfg_pp",
|
||||
"er_sde",
|
||||
"seeds_2",
|
||||
"seeds_3",
|
||||
"sa_solver",
|
||||
"sa_solver_pece",
|
||||
"ddim",
|
||||
"uni_pc",
|
||||
"uni_pc_bh2"
|
||||
]
|
||||
],
|
||||
"scheduler": [
|
||||
[
|
||||
"simple",
|
||||
"sgm_uniform",
|
||||
"karras",
|
||||
"exponential",
|
||||
"ddim_uniform",
|
||||
"beta",
|
||||
"normal",
|
||||
"linear_quadratic",
|
||||
"kl_optimal"
|
||||
]
|
||||
],
|
||||
"denoise": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}
|
||||
],
|
||||
"tile_width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"tile_height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"padding": [
|
||||
"INT",
|
||||
{
|
||||
"default": 32,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"mask_blur": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 0,
|
||||
"max": 256
|
||||
}
|
||||
],
|
||||
"force_uniform_tiles": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": true
|
||||
}
|
||||
],
|
||||
"tiled_decode": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"tile_indices": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"dynamic_threshold": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 64
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"upscaled_image",
|
||||
"model",
|
||||
"positive",
|
||||
"negative",
|
||||
"vae",
|
||||
"seed",
|
||||
"steps",
|
||||
"cfg",
|
||||
"sampler_name",
|
||||
"scheduler",
|
||||
"denoise",
|
||||
"tile_width",
|
||||
"tile_height",
|
||||
"padding",
|
||||
"mask_blur",
|
||||
"force_uniform_tiles",
|
||||
"tiled_decode"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_id",
|
||||
"tile_indices",
|
||||
"dynamic_threshold"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image/upscaling",
|
||||
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
|
||||
}
|
||||
}
|
||||
@@ -1,252 +0,0 @@
|
||||
"""Exercise the actual loader, schemas and executor in a fresh CPU process.
|
||||
|
||||
No HTTP listener, workers, model downloads or live installation changes.
|
||||
The original contracts were captured from 32ac027 using the same core.
|
||||
"""
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_ROOT = Path(sys.argv[1]).resolve()
|
||||
sys.path.insert(0, str(COMFY_ROOT))
|
||||
os.environ['COMFYUI_IS_WORKER'] = '1'
|
||||
import comfy.cli_args
|
||||
comfy.cli_args.args.cpu = True
|
||||
comfy.cli_args.args.disable_assets = True
|
||||
from app.assets.manager import default_asset_manager
|
||||
from comfy_api.v0_0_2 import io
|
||||
import execution
|
||||
import nodes
|
||||
import server
|
||||
import torch
|
||||
|
||||
BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
|
||||
|
||||
|
||||
def normalized_input(value):
|
||||
kind = value[0]
|
||||
opts = dict(value[1]) if len(value) > 1 else {}
|
||||
if kind == 'STRING':
|
||||
# V3 serializes the same default single-line widget explicitly.
|
||||
opts.setdefault('multiline', False)
|
||||
if kind == 'COMBO':
|
||||
kind = opts['options']
|
||||
opts = {key: val for key, val in opts.items() if key != 'options'}
|
||||
if isinstance(kind, list):
|
||||
# Single-selection is the V1 dropdown default as well.
|
||||
opts.setdefault('multiselect', False)
|
||||
return [kind, opts]
|
||||
|
||||
|
||||
async def map_node(cls, values, extra=None):
|
||||
"""Use the actual input parser and sanitized V3 executor clones."""
|
||||
prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
|
||||
assert not missing, missing
|
||||
# Use a worker thread for synchronous nodes: collector/upscale bridge back
|
||||
# to PromptServer's running loop, just as ComfyUI's prompt worker does.
|
||||
def run():
|
||||
return asyncio.run(execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
|
||||
return await asyncio.to_thread(run)
|
||||
|
||||
|
||||
def check_saved_workflows(mapping):
|
||||
checked = 0
|
||||
occurrences = 0
|
||||
for path in sorted((ROOT / 'workflows').glob('*.json')):
|
||||
workflow = json.loads(path.read_text())
|
||||
local_nodes = {node['id']: node for node in workflow['nodes']}
|
||||
for saved in workflow['nodes']:
|
||||
if saved['type'] not in mapping:
|
||||
continue
|
||||
occurrences += 1
|
||||
info = mapping[saved['type']].GET_NODE_INFO_V1()
|
||||
inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
|
||||
for socket in saved.get('inputs', []):
|
||||
assert socket['name'] in inputs, (path.name, saved['id'], socket)
|
||||
assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
|
||||
for index, output in enumerate(saved.get('outputs', [])):
|
||||
assert output['type'] == info['output'][index], (path.name, index)
|
||||
for link in workflow['links']:
|
||||
if link[1] == saved['id']:
|
||||
assert link[2] < len(info['output']), (path.name, link)
|
||||
target = local_nodes[link[3]]['inputs'][link[4]]
|
||||
assert target['type'] == info['output'][link[2]], (path.name, link)
|
||||
checked += 1
|
||||
assert checked == 5 and occurrences == 10, (checked, occurrences)
|
||||
print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
|
||||
|
||||
|
||||
async def check_execution(mapping, module, prompt_server, asset_manager):
|
||||
result = await map_node(mapping['DistributedSeed'],
|
||||
{'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (126,)
|
||||
values = json.dumps({'_type': 'INT', '3': '17'})
|
||||
result = await map_node(mapping['DistributedValue'],
|
||||
{'default_value': '4', 'worker_values': values,
|
||||
'is_worker': True, 'worker_id': 'worker_2'})
|
||||
assert result[0].result == (17,)
|
||||
metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
|
||||
result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
|
||||
{'extra_pnginfo': metadata})
|
||||
assert result[0].result == ('model.ckpt',)
|
||||
assert result[0].ui == {'text': ['model.ckpt']}
|
||||
assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
|
||||
result = await map_node(mapping['DistributedEmptyImage'],
|
||||
{'height': 8, 'width': 8, 'channels': 3})
|
||||
empty = result[0].result[0]
|
||||
assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
|
||||
images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
|
||||
for count in (1, 3, 10):
|
||||
result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat(result[0].result[:count]), images)
|
||||
audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
|
||||
result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
|
||||
assert len(result[0].result) == 10
|
||||
assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
|
||||
# INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
|
||||
collector = mapping['DistributedCollector']
|
||||
prepared, missing, hidden = execution.get_input_data(
|
||||
{'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
|
||||
assert not missing
|
||||
prepared['images'] = [images[:2], images[2:]]
|
||||
prepared['audio'] = [audio, audio]
|
||||
result = await execution._async_map_node_over_list(
|
||||
'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
|
||||
assert torch.equal(result[0].result[0], images)
|
||||
assert result[0].result[1]['waveform'].shape[-1] == 66
|
||||
# Exercise real async aggregation without an HTTP endpoint or worker.
|
||||
queue = asyncio.Queue()
|
||||
await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
|
||||
prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
|
||||
result = await map_node(collector,
|
||||
{'images': images[:2], 'load_balance': False,
|
||||
'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
|
||||
assert torch.equal(result[0].result[0], images[:3])
|
||||
assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
|
||||
|
||||
# Prove V3's per-execution class clones do not leak helper instance state;
|
||||
# replace only the GPU/model boundary, not the input parser or executor.
|
||||
upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
|
||||
seen = []
|
||||
def fake_upscale(self, *args):
|
||||
assert not hasattr(self, 'acceptance_marker')
|
||||
self.acceptance_marker = True
|
||||
seen.append(args)
|
||||
return (args[0],)
|
||||
inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
|
||||
'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
|
||||
'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
|
||||
'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
|
||||
'force_uniform_tiles': True, 'tiled_decode': False,
|
||||
'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
|
||||
'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
|
||||
'tile_indices': '[2,3]', 'dynamic_threshold': 9}
|
||||
with patch.object(upscale_runtime, 'run', fake_upscale):
|
||||
for _ in range(2):
|
||||
result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
|
||||
assert result[0].result[0] is images
|
||||
assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
|
||||
'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
|
||||
'tile_indices', 'dynamic_threshold')), seen[0]
|
||||
import math
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
|
||||
assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
|
||||
|
||||
graph = {
|
||||
'1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
|
||||
'2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
|
||||
'pass_through': True}},
|
||||
'3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
|
||||
'4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
|
||||
}
|
||||
valid = await execution.validate_prompt('v3-graph', graph, None)
|
||||
assert valid[0], valid
|
||||
invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
|
||||
'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
|
||||
'2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
|
||||
rejected = await execution.validate_prompt('v3-invalid', invalid, None)
|
||||
assert not rejected[0]
|
||||
errors = [error for entry in rejected[3].values() for error in entry['errors']]
|
||||
enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
|
||||
assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
|
||||
import folder_paths
|
||||
scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
|
||||
with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
|
||||
with patch.object(folder_paths, 'temp_directory', temp):
|
||||
executor = execution.PromptExecutor(
|
||||
prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
|
||||
await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
|
||||
assert executor.success, executor.status_messages
|
||||
history = executor.history_result
|
||||
record = history['outputs']['4']['images'][0]
|
||||
preview = Path(temp) / record.get('subfolder', '') / record['filename']
|
||||
with Image.open(preview) as image:
|
||||
assert image.size == (8, 8) and image.mode == 'RGB'
|
||||
assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
|
||||
print('V3_EXECUTION_OK eight nodes; upscale GPU boundary mocked; preview artifact verified')
|
||||
|
||||
|
||||
async def main():
|
||||
asset_manager = default_asset_manager()
|
||||
prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
|
||||
assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
|
||||
assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
|
||||
module = sys.modules[str(ROOT).replace('.', '_x_')]
|
||||
assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
|
||||
extension = await module.comfy_entrypoint()
|
||||
classes = await extension.get_node_list()
|
||||
mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
|
||||
assert len(classes) == len(mapping) == len(BASELINE) == 8
|
||||
assert set(mapping) == set(BASELINE)
|
||||
for node_id, cls in mapping.items():
|
||||
assert issubclass(cls, io.ComfyNode)
|
||||
assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
|
||||
old = dict(BASELINE[node_id])
|
||||
if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
|
||||
# V1's ByPassTypeTuple advertises '*' when indexed, while its
|
||||
# underlying tuple and existing frontend declare IMAGE/AUDIO.
|
||||
# Native V3 declares all ten existing typed sockets explicitly.
|
||||
assert old['output'] == ['*'] * 10
|
||||
old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
|
||||
new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
|
||||
for group in ('required', 'optional'):
|
||||
old_inputs = old['input'].get(group, {})
|
||||
new_inputs = new['input'].get(group, {})
|
||||
assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
|
||||
for name, original in old_inputs.items():
|
||||
assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
|
||||
for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
|
||||
assert old[key] == new[key], (node_id, key, old[key], new[key])
|
||||
# Standard context lives in cls.hidden; orchestrator metadata remains
|
||||
# accepted by its original kwarg name, without creating new widgets.
|
||||
signature = inspect.signature(cls.execute)
|
||||
for name, field in old['input'].get('hidden', {}).items():
|
||||
if isinstance(field, list):
|
||||
assert cls.GET_SCHEMA().accept_all_inputs, node_id
|
||||
assert name in signature.parameters, (node_id, name)
|
||||
assert signature.parameters[name].default == field[1]['default'], (node_id, name)
|
||||
else:
|
||||
assert name in new['input']['hidden'], (node_id, name)
|
||||
expected_hidden = []
|
||||
if node_id == 'DistributedModelName':
|
||||
expected_hidden.extend(['unique_id', 'extra_pnginfo'])
|
||||
if old['output_node']:
|
||||
expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
|
||||
assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
|
||||
print('SCHEMA_PARITY_OK', len(mapping))
|
||||
check_saved_workflows(mapping)
|
||||
await check_execution(mapping, module, prompt_server, asset_manager)
|
||||
print('V3_ACCEPTANCE_OK')
|
||||
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -1,91 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class _PromptQueue:
|
||||
def __init__(self):
|
||||
self.items = []
|
||||
|
||||
def put(self, item):
|
||||
self.items.append(item)
|
||||
|
||||
|
||||
def _load_async_helpers_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
|
||||
package_name = "dist_async_helpers_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
execution_module = types.ModuleType("execution")
|
||||
|
||||
async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
|
||||
return (True, None, ["9"], {})
|
||||
|
||||
execution_module.validate_prompt = _validate_prompt
|
||||
execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
|
||||
sys.modules["execution"] = execution_module
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
trigger_on_prompt=lambda payload: payload,
|
||||
number=12,
|
||||
prompt_queue=_PromptQueue(),
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_loop = lambda: None
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module, prompt_server
|
||||
|
||||
|
||||
async_helpers, prompt_server = _load_async_helpers_module()
|
||||
|
||||
|
||||
class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
|
||||
result = await async_helpers.queue_prompt_payload(
|
||||
{"1": {"class_type": "Node"}},
|
||||
workflow_meta={"id": "workflow-1"},
|
||||
client_id="client-1",
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
|
||||
self.assertIsInstance(result["prompt_id"], str)
|
||||
self.assertTrue(result["prompt_id"])
|
||||
self.assertEqual(result["number"], 12)
|
||||
self.assertEqual(result["node_errors"], {})
|
||||
|
||||
self.assertEqual(prompt_server.number, 13)
|
||||
self.assertEqual(len(prompt_server.prompt_queue.items), 1)
|
||||
queued_item = prompt_server.prompt_queue.items[0]
|
||||
self.assertEqual(queued_item[0], 12)
|
||||
extra_data = queued_item[3]
|
||||
self.assertEqual(extra_data["client_id"], "client-1")
|
||||
self.assertIn("create_time", extra_data)
|
||||
self.assertIsInstance(extra_data["create_time"], int)
|
||||
self.assertGreater(extra_data["create_time"], 0)
|
||||
self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,102 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_utilities_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
package_name = "dist_divider_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.nodes.utilities",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
utils = _load_utilities_module()
|
||||
|
||||
|
||||
class ImageBatchDividerTests(unittest.TestCase):
|
||||
def test_divides_images_into_contiguous_chunks(self):
|
||||
divider = utils.ImageBatchDivider()
|
||||
images = torch.arange(10, dtype=torch.float32).reshape(10, 1, 1, 1)
|
||||
|
||||
outputs = divider.divide_batch(images, 3)
|
||||
|
||||
self.assertEqual(outputs[0].shape[0], 4)
|
||||
self.assertEqual(outputs[1].shape[0], 3)
|
||||
self.assertEqual(outputs[2].shape[0], 3)
|
||||
self.assertEqual(outputs[0][:, 0, 0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
|
||||
self.assertEqual(outputs[1][:, 0, 0, 0].tolist(), [4.0, 5.0, 6.0])
|
||||
self.assertEqual(outputs[2][:, 0, 0, 0].tolist(), [7.0, 8.0, 9.0])
|
||||
|
||||
def test_unused_image_outputs_are_empty(self):
|
||||
divider = utils.ImageBatchDivider()
|
||||
images = torch.arange(4, dtype=torch.float32).reshape(4, 1, 1, 1)
|
||||
|
||||
outputs = divider.divide_batch(images, 2)
|
||||
|
||||
self.assertEqual(len(outputs), 10)
|
||||
for idx in range(2, 10):
|
||||
self.assertEqual(outputs[idx].shape[0], 0)
|
||||
|
||||
|
||||
class AudioBatchDividerTests(unittest.TestCase):
|
||||
def test_divides_audio_samples_into_contiguous_chunks(self):
|
||||
divider = utils.AudioBatchDivider()
|
||||
audio = {
|
||||
"waveform": torch.arange(10, dtype=torch.float32).reshape(1, 1, 10),
|
||||
"sample_rate": 24000,
|
||||
}
|
||||
|
||||
outputs = divider.divide_audio(audio, 3)
|
||||
|
||||
self.assertEqual(outputs[0]["waveform"][0, 0].tolist(), [0.0, 1.0, 2.0, 3.0])
|
||||
self.assertEqual(outputs[1]["waveform"][0, 0].tolist(), [4.0, 5.0, 6.0])
|
||||
self.assertEqual(outputs[2]["waveform"][0, 0].tolist(), [7.0, 8.0, 9.0])
|
||||
|
||||
def test_unused_audio_outputs_are_empty(self):
|
||||
divider = utils.AudioBatchDivider()
|
||||
audio = {
|
||||
"waveform": torch.arange(8, dtype=torch.float32).reshape(1, 1, 8),
|
||||
"sample_rate": 24000,
|
||||
}
|
||||
|
||||
outputs = divider.divide_audio(audio, 2)
|
||||
|
||||
self.assertEqual(len(outputs), 10)
|
||||
for idx in range(2, 10):
|
||||
self.assertEqual(outputs[idx]["waveform"].shape[-1], 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,364 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_collector_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
|
||||
package_name = "dist_collector_list_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
nodes_pkg = types.ModuleType(f"{package_name}.nodes")
|
||||
nodes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.nodes"] = nodes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
class _Routes:
|
||||
def post(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
def get(self, _path):
|
||||
return lambda fn: fn
|
||||
|
||||
prompt_server = types.SimpleNamespace(
|
||||
routes=_Routes(),
|
||||
distributed_jobs_lock=None,
|
||||
distributed_pending_jobs={},
|
||||
)
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_management = types.ModuleType("comfy.model_management")
|
||||
|
||||
class InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_management.InterruptProcessingException = InterruptProcessingException
|
||||
model_management.throw_exception_if_processing_interrupted = lambda: None
|
||||
comfy_module.model_management = model_management
|
||||
|
||||
comfy_utils = types.ModuleType("comfy.utils")
|
||||
|
||||
class ProgressBar:
|
||||
def __init__(self, _total):
|
||||
self.total = _total
|
||||
self.updates = []
|
||||
|
||||
def update(self, value):
|
||||
self.updates.append(value)
|
||||
|
||||
comfy_utils.ProgressBar = ProgressBar
|
||||
comfy_module.utils = comfy_utils
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_management
|
||||
sys.modules["comfy.utils"] = comfy_utils
|
||||
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 0.1
|
||||
config_module.load_config = lambda: {"workers": []}
|
||||
config_module.is_master_delegate_only = lambda: False
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 1.0
|
||||
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"])
|
||||
@@ -1,293 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_config_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
|
||||
package_name = "dist_cfg_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = [str(module_path.parent)]
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[f"{package_name}.constants"] = constants_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
config = _load_config_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _merge_with_defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MergeWithDefaultsTests(unittest.TestCase):
|
||||
def test_non_dict_input_returns_defaults(self):
|
||||
result = config._merge_with_defaults("not a dict", {"key": "default"})
|
||||
self.assertEqual(result, {"key": "default"})
|
||||
|
||||
def test_fills_missing_keys_with_defaults(self):
|
||||
result = config._merge_with_defaults({}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result, {"a": 1, "b": 2})
|
||||
|
||||
def test_loaded_value_overrides_default(self):
|
||||
result = config._merge_with_defaults({"a": 99}, {"a": 1, "b": 2})
|
||||
self.assertEqual(result["a"], 99)
|
||||
self.assertEqual(result["b"], 2)
|
||||
|
||||
def test_nested_dict_merges_recursively(self):
|
||||
defaults = {"settings": {"debug": False, "count": 5}}
|
||||
loaded = {"settings": {"debug": True}}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertTrue(result["settings"]["debug"])
|
||||
self.assertEqual(result["settings"]["count"], 5)
|
||||
|
||||
def test_preserves_unknown_keys_for_forward_compatibility(self):
|
||||
result = config._merge_with_defaults({"extra_key": "extra"}, {"a": 1})
|
||||
self.assertEqual(result["extra_key"], "extra")
|
||||
|
||||
def test_none_loaded_value_overrides_default(self):
|
||||
"""Explicitly set None in config should override non-None default."""
|
||||
result = config._merge_with_defaults({"a": None}, {"a": "default"})
|
||||
self.assertIsNone(result["a"])
|
||||
|
||||
def test_non_dict_nested_loaded_value_replaces_dict_default(self):
|
||||
"""If loaded has a scalar where default has a dict, use the scalar."""
|
||||
defaults = {"settings": {"debug": False}}
|
||||
loaded = {"settings": "flat_string"}
|
||||
result = config._merge_with_defaults(loaded, defaults)
|
||||
self.assertEqual(result["settings"], "flat_string")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# load_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LoadConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_returns_defaults_when_file_missing(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent/path/config.json"):
|
||||
cfg = config.load_config()
|
||||
defaults = config.get_default_config()
|
||||
self.assertEqual(cfg["settings"]["debug"], defaults["settings"]["debug"])
|
||||
self.assertIn("workers", cfg)
|
||||
|
||||
def test_loads_valid_json_file(self):
|
||||
data = {
|
||||
"workers": [{"id": "w1"}],
|
||||
"master": {"host": "test.host"},
|
||||
"settings": {},
|
||||
"tunnel": {},
|
||||
}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertEqual(cfg["master"]["host"], "test.host")
|
||||
self.assertEqual(len(cfg["workers"]), 1)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_merges_loaded_file_with_defaults(self):
|
||||
"""Loaded file with partial settings should be filled in from defaults."""
|
||||
data = {"master": {"host": "h"}, "workers": [], "settings": {"debug": True}, "tunnel": {}}
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
# debug was set to True
|
||||
self.assertTrue(cfg["settings"]["debug"])
|
||||
# auto_launch_workers is a default key and should be present
|
||||
self.assertIn("auto_launch_workers", cfg["settings"])
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_falls_back_to_defaults_on_invalid_json(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
f.write("{invalid json{{")
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg = config.load_config()
|
||||
self.assertIn("settings", cfg)
|
||||
self.assertIn("workers", cfg)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_second_call_returns_cached_object(self):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(config.get_default_config(), f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
cfg2 = config.load_config()
|
||||
self.assertIs(cfg1, cfg2)
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
def test_invalidate_cache_forces_reload(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "first"
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||||
json.dump(data, f)
|
||||
tmp_path = f.name
|
||||
try:
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
cfg1 = config.load_config()
|
||||
config.invalidate_config_cache()
|
||||
data["master"]["host"] = "second"
|
||||
with open(tmp_path, "w", encoding="utf-8") as fh:
|
||||
json.dump(data, fh)
|
||||
cfg2 = config.load_config()
|
||||
self.assertEqual(cfg1["master"]["host"], "first")
|
||||
self.assertEqual(cfg2["master"]["host"], "second")
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# save_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class SaveConfigTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def tearDown(self):
|
||||
config.invalidate_config_cache()
|
||||
|
||||
def test_saves_and_reloads_correctly(self):
|
||||
data = config.get_default_config()
|
||||
data["master"]["host"] = "saved.host"
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
result = config.save_config(data)
|
||||
self.assertTrue(result)
|
||||
loaded = config.load_config()
|
||||
self.assertEqual(loaded["master"]["host"], "saved.host")
|
||||
|
||||
def test_returns_false_when_path_unwritable(self):
|
||||
with patch.object(config, "CONFIG_FILE", "/nonexistent_dir/config.json"):
|
||||
result = config.save_config({})
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_save_invalidates_cache(self):
|
||||
"""After saving, the cache should be cleared so next load re-reads."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
data = config.get_default_config()
|
||||
config.save_config(data)
|
||||
# Cache is now None; load_config should re-read
|
||||
self.assertIsNone(config._config_cache)
|
||||
|
||||
def test_written_file_is_valid_json(self):
|
||||
data = config.get_default_config()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
tmp_path = os.path.join(tmpdir, "config.json")
|
||||
with patch.object(config, "CONFIG_FILE", tmp_path):
|
||||
config.save_config(data)
|
||||
with open(tmp_path, encoding="utf-8") as fh:
|
||||
parsed = json.load(fh)
|
||||
self.assertEqual(parsed["master"], data["master"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_worker_timeout_seconds
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetWorkerTimeoutSecondsTests(unittest.TestCase):
|
||||
def test_returns_configured_value(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 120
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 120)
|
||||
|
||||
def test_clamps_zero_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = 0
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_clamps_negative_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["worker_timeout_seconds"] = -10
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertEqual(config.get_worker_timeout_seconds(), 1)
|
||||
|
||||
def test_falls_back_to_provided_default_when_key_missing(self):
|
||||
cfg = config.get_default_config()
|
||||
# worker_timeout_seconds is not present in default config
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=45)
|
||||
self.assertEqual(result, 45)
|
||||
|
||||
def test_fallback_also_clamped_to_one(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"].pop("worker_timeout_seconds", None)
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
result = config.get_worker_timeout_seconds(default=0)
|
||||
self.assertEqual(result, 1)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_master_delegate_only
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsMasterDelegateOnlyTests(unittest.TestCase):
|
||||
def test_returns_false_by_default(self):
|
||||
cfg = config.get_default_config()
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_true_when_enabled(self):
|
||||
cfg = config.get_default_config()
|
||||
cfg["settings"]["master_delegate_only"] = True
|
||||
with patch.object(config, "load_config", return_value=cfg):
|
||||
self.assertTrue(config.is_master_delegate_only())
|
||||
|
||||
def test_returns_false_on_exception(self):
|
||||
def _raise():
|
||||
raise RuntimeError("config exploded")
|
||||
|
||||
with patch.object(config, "load_config", side_effect=RuntimeError("boom")):
|
||||
self.assertFalse(config.is_master_delegate_only())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,196 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_detection_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
|
||||
package_name = "dist_det_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.normalize_host = lambda value: value
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("network calls not used in these tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
pass
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.detection",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
detection = _load_detection_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_docker_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsDockerEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_dockerenv_file_exists(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=True), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_docker_container_env_var_is_set(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": "1"}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_true_when_platform_node_contains_docker(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-docker-host"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_false_when_none_of_the_signals_are_present(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
def test_docker_node_name_is_case_insensitive(self):
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="My-Docker-Box"):
|
||||
self.assertTrue(detection.is_docker_environment())
|
||||
|
||||
def test_docker_env_var_empty_string_is_falsy(self):
|
||||
"""An empty DOCKER_CONTAINER env var should NOT trigger docker detection."""
|
||||
with patch.object(detection.os.path, "exists", return_value=False), \
|
||||
patch.dict(detection.os.environ, {"DOCKER_CONTAINER": ""}, clear=True), \
|
||||
patch.object(detection.platform, "node", return_value="my-laptop"):
|
||||
self.assertFalse(detection.is_docker_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_runpod_environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsRunpodEnvironmentTests(unittest.TestCase):
|
||||
def test_true_when_runpod_pod_id_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": "pod-abc"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_runpod_api_key_is_set(self):
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_API_KEY": "key-xyz"}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_both_vars_are_set(self):
|
||||
with patch.dict(
|
||||
detection.os.environ,
|
||||
{"RUNPOD_POD_ID": "pod-abc", "RUNPOD_API_KEY": "key-xyz"},
|
||||
clear=True,
|
||||
):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
def test_false_when_neither_var_is_set(self):
|
||||
with patch.dict(detection.os.environ, {}, clear=True):
|
||||
self.assertFalse(detection.is_runpod_environment())
|
||||
|
||||
def test_true_when_pod_id_is_empty_string(self):
|
||||
"""is not None check means even empty string counts as detected."""
|
||||
with patch.dict(detection.os.environ, {"RUNPOD_POD_ID": ""}, clear=True):
|
||||
self.assertTrue(detection.is_runpod_environment())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_local_worker (synchronous paths only)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_true_for_localhost_host(self):
|
||||
result = await detection.is_local_worker({"host": "localhost", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_127_0_0_1(self):
|
||||
result = await detection.is_local_worker({"host": "127.0.0.1", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_for_0_0_0_0(self):
|
||||
result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_true_when_type_is_local(self):
|
||||
result = await detection.is_local_worker({"type": "local", "host": "remote.example.com"})
|
||||
self.assertTrue(result)
|
||||
|
||||
async def test_false_for_remote_host(self):
|
||||
result = await detection.is_local_worker({"host": "remote.example.com", "port": 8188})
|
||||
self.assertFalse(result)
|
||||
|
||||
async def test_true_when_no_host_key(self):
|
||||
"""Missing host defaults to 'localhost'."""
|
||||
result = await detection.is_local_worker({"port": 8188})
|
||||
self.assertTrue(result)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_machine_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class GetMachineIdTests(unittest.TestCase):
|
||||
def test_returns_a_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertIsInstance(result, str)
|
||||
|
||||
def test_returns_non_empty_string(self):
|
||||
result = detection.get_machine_id()
|
||||
self.assertTrue(len(result) > 0)
|
||||
|
||||
def test_stable_across_calls(self):
|
||||
r1 = detection.get_machine_id()
|
||||
r2 = detection.get_machine_id()
|
||||
self.assertEqual(r1, r2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,254 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def _load_dispatch_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
|
||||
|
||||
package_name = "dist_dispatch_testpkg"
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orch_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
|
||||
|
||||
async def _probe_worker(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
|
||||
async def _fake_session():
|
||||
raise RuntimeError("get_client_session should be mocked in these tests")
|
||||
|
||||
network_module.get_client_session = _fake_session
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _ClientTimeout:
|
||||
def __init__(self, total=None):
|
||||
self.total = total
|
||||
|
||||
class _ClientConnectorError(Exception):
|
||||
pass
|
||||
|
||||
class _WSMsgType:
|
||||
TEXT = "TEXT"
|
||||
ERROR = "ERROR"
|
||||
CLOSED = "CLOSED"
|
||||
|
||||
class _TCPConnector:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _ClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.ClientTimeout = _ClientTimeout
|
||||
aiohttp_module.ClientConnectorError = _ClientConnectorError
|
||||
aiohttp_module.WSMsgType = _WSMsgType
|
||||
aiohttp_module.TCPConnector = _TCPConnector
|
||||
aiohttp_module.ClientSession = _ClientSession
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: {"payload": payload, "status": status}
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.dispatch",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
return module
|
||||
|
||||
|
||||
dispatch = _load_dispatch_module()
|
||||
|
||||
|
||||
class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_select_active_workers_filters_offline(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
|
||||
async def fake_probe(worker):
|
||||
return worker["id"] != "w2"
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=3,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_disables_delegate_when_all_offline(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker):
|
||||
return False
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, delegate_master = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=True,
|
||||
probe_concurrency=1,
|
||||
)
|
||||
|
||||
self.assertEqual(active_workers, [])
|
||||
self.assertFalse(delegate_master)
|
||||
|
||||
async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
|
||||
|
||||
async def fake_http_probe(_worker):
|
||||
return False
|
||||
|
||||
async def fake_ws_probe(_worker):
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
|
||||
dispatch,
|
||||
"worker_ws_is_active",
|
||||
side_effect=fake_ws_probe,
|
||||
) as ws_probe:
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=True,
|
||||
delegate_master=False,
|
||||
probe_concurrency=4,
|
||||
)
|
||||
|
||||
self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
|
||||
self.assertEqual(ws_probe.call_count, 2)
|
||||
self.assertEqual(http_probe.call_count, 0)
|
||||
|
||||
async def test_probe_concurrency_is_bounded(self):
|
||||
workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
|
||||
state = {"in_flight": 0, "max_in_flight": 0}
|
||||
|
||||
async def fake_probe(_worker):
|
||||
state["in_flight"] += 1
|
||||
state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
|
||||
await asyncio.sleep(0.01)
|
||||
state["in_flight"] -= 1
|
||||
return True
|
||||
|
||||
with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
|
||||
active_workers, _ = await dispatch.select_active_workers(
|
||||
workers,
|
||||
use_websocket=False,
|
||||
delegate_master=False,
|
||||
probe_concurrency=2,
|
||||
)
|
||||
|
||||
self.assertEqual(len(active_workers), len(workers))
|
||||
self.assertLessEqual(state["max_in_flight"], 2)
|
||||
self.assertGreaterEqual(state["max_in_flight"], 2)
|
||||
|
||||
async def test_select_least_busy_worker_round_robins_idle_workers(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 0, "w2": 0, "w3": 2}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
dispatch._least_busy_rr_index = 0
|
||||
selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
|
||||
|
||||
self.assertEqual(selected1["id"], "w1")
|
||||
self.assertEqual(selected2["id"], "w2")
|
||||
self.assertEqual(selected3["id"], "w1")
|
||||
|
||||
async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
|
||||
workers = [
|
||||
{"id": "w1", "name": "Worker 1"},
|
||||
{"id": "w2", "name": "Worker 2"},
|
||||
{"id": "w3", "name": "Worker 3"},
|
||||
]
|
||||
queue_map = {"w1": 5, "w2": 2, "w3": 4}
|
||||
|
||||
async def fake_probe(worker_url, timeout=3.0):
|
||||
worker_id = worker_url.rsplit("/", 1)[-1]
|
||||
return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
|
||||
|
||||
self.assertEqual(selected["id"], "w2")
|
||||
|
||||
async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
|
||||
workers = [{"id": "w1", "name": "Worker 1"}]
|
||||
|
||||
async def fake_probe(_worker_url, timeout=3.0):
|
||||
return None
|
||||
|
||||
with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
|
||||
dispatch,
|
||||
"probe_worker",
|
||||
side_effect=fake_probe,
|
||||
):
|
||||
selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
|
||||
|
||||
self.assertIsNone(selected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,208 +0,0 @@
|
||||
import json
|
||||
import unittest
|
||||
|
||||
|
||||
class DistributedValueTests(unittest.TestCase):
|
||||
"""Unit tests for the DistributedValue node's distribute() method."""
|
||||
|
||||
def _make_node(self):
|
||||
# Import inline to avoid plugin-level imports
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
|
||||
pkg_name = "dv_test_pkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == pkg_name or mod_name.startswith(f"{pkg_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
# Mock torch if not available
|
||||
if "torch" not in sys.modules:
|
||||
sys.modules["torch"] = MagicMock()
|
||||
|
||||
root_pkg = types.ModuleType(pkg_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[pkg_name] = root_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{pkg_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{pkg_name}.utils"] = utils_pkg
|
||||
|
||||
logging_mod = types.ModuleType(f"{pkg_name}.utils.logging")
|
||||
logging_mod.debug_log = lambda *_a, **_k: None
|
||||
logging_mod.log = lambda *_a, **_k: None
|
||||
sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{pkg_name}.nodes.utilities", module_path
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod.DistributedValue()
|
||||
|
||||
def setUp(self):
|
||||
self.node = self._make_node()
|
||||
|
||||
def test_master_returns_default(self):
|
||||
result = self.node.distribute(
|
||||
default_value="model_a",
|
||||
worker_values="{}",
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, ("model_a",))
|
||||
|
||||
def test_master_coerces_default_int(self):
|
||||
values = json.dumps({"_type": "INT"})
|
||||
result = self.node.distribute(
|
||||
default_value="42",
|
||||
worker_values=values,
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, (42,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_master_coerces_default_float(self):
|
||||
values = json.dumps({"_type": "FLOAT"})
|
||||
result = self.node.distribute(
|
||||
default_value="2.5",
|
||||
worker_values=values,
|
||||
is_worker=False,
|
||||
worker_id="",
|
||||
)
|
||||
self.assertEqual(result, (2.5,))
|
||||
self.assertIsInstance(result[0], float)
|
||||
|
||||
def test_worker_returns_specific_value(self):
|
||||
values = json.dumps({"1": "model_x", "2": "model_y"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("model_x",))
|
||||
|
||||
def test_worker_second_index(self):
|
||||
values = json.dumps({"1": "model_x", "2": "model_y"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_1",
|
||||
)
|
||||
self.assertEqual(result, ("model_y",))
|
||||
|
||||
def test_worker_falls_back_to_default_when_key_missing(self):
|
||||
values = json.dumps({"_type": "INT", "1": "3"})
|
||||
result = self.node.distribute(
|
||||
default_value="9",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_5",
|
||||
)
|
||||
self.assertEqual(result, (9,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_worker_falls_back_to_default_on_empty_value(self):
|
||||
values = json.dumps({"1": ""})
|
||||
result = self.node.distribute(
|
||||
default_value="fallback",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("fallback",))
|
||||
|
||||
def test_worker_falls_back_on_invalid_json(self):
|
||||
result = self.node.distribute(
|
||||
default_value="safe",
|
||||
worker_values="not-json",
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("safe",))
|
||||
|
||||
def test_worker_falls_back_on_invalid_worker_id(self):
|
||||
values = json.dumps({"1": "model_x"})
|
||||
result = self.node.distribute(
|
||||
default_value="safe",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="bad_id",
|
||||
)
|
||||
self.assertEqual(result, ("safe",))
|
||||
|
||||
def test_worker_id_as_direct_integer(self):
|
||||
values = json.dumps({"1": "model_x"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="0",
|
||||
)
|
||||
self.assertEqual(result, ("model_x",))
|
||||
|
||||
def test_type_int_coerces_value(self):
|
||||
values = json.dumps({"_type": "INT", "1": "42"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, (42,))
|
||||
self.assertIsInstance(result[0], int)
|
||||
|
||||
def test_type_float_coerces_value(self):
|
||||
values = json.dumps({"_type": "FLOAT", "1": "3.14"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertAlmostEqual(result[0], 3.14)
|
||||
self.assertIsInstance(result[0], float)
|
||||
|
||||
def test_type_combo_stays_string(self):
|
||||
values = json.dumps({"_type": "COMBO", "1": "model_v2"})
|
||||
result = self.node.distribute(
|
||||
default_value="model_v1",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("model_v2",))
|
||||
self.assertIsInstance(result[0], str)
|
||||
|
||||
def test_type_string_default_stays_string(self):
|
||||
values = json.dumps({"1": "hello"})
|
||||
result = self.node.distribute(
|
||||
default_value="default",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, ("hello",))
|
||||
self.assertIsInstance(result[0], str)
|
||||
|
||||
def test_int_coerce_handles_float_string(self):
|
||||
"""INT coercion of '3.7' should truncate to 3."""
|
||||
values = json.dumps({"_type": "INT", "1": "3.7"})
|
||||
result = self.node.distribute(
|
||||
default_value="0",
|
||||
worker_values=values,
|
||||
is_worker=True,
|
||||
worker_id="worker_0",
|
||||
)
|
||||
self.assertEqual(result, (3,))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,197 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
import unittest
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_job_timeout_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
|
||||
package_name = "dist_job_timeout_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
config_holder = {"value": {"settings": {}, "workers": []}}
|
||||
probe_holder = {"fn": None}
|
||||
prompt_server_holder = {
|
||||
"value": types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
}
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: config_holder["value"]
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 60
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.build_worker_url = lambda worker: f"http://{worker.get('host', '127.0.0.1')}:{worker.get('port', 8188)}"
|
||||
|
||||
async def _probe_worker(url, timeout=2.0):
|
||||
fn = probe_holder["fn"]
|
||||
if fn is None:
|
||||
return None
|
||||
return await fn(url, timeout)
|
||||
|
||||
network_module.probe_worker = _probe_worker
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class BaseJobState:
|
||||
pass
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
|
||||
job_models_module.BaseJobState = BaseJobState
|
||||
job_models_module.ImageJobState = ImageJobState
|
||||
job_models_module.TileJobState = TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.job_timeout", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
module._config_holder = config_holder
|
||||
module._probe_holder = probe_holder
|
||||
module._prompt_server_holder = prompt_server_holder
|
||||
module._ImageJobState = ImageJobState
|
||||
module._TileJobState = TileJobState
|
||||
return module
|
||||
|
||||
|
||||
jt = _load_job_timeout_module()
|
||||
|
||||
|
||||
class JobTimeoutRequeueTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
jt._prompt_server_holder["value"] = types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
jt._config_holder["value"] = {
|
||||
"settings": {"worker_timeout_seconds": 5},
|
||||
"workers": [{"id": "worker-1", "host": "worker.local", "port": 8188}],
|
||||
}
|
||||
|
||||
async def test_requeues_only_incomplete_dynamic_tasks_for_timed_out_worker(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-1")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [0, 1]
|
||||
job_data.completed_images[1] = "done"
|
||||
prompt_server.distributed_pending_tile_jobs["job-1"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-1", total_tasks=2)
|
||||
|
||||
self.assertEqual(requeued, 1)
|
||||
self.assertEqual(await job_data.pending_images.get(), 0)
|
||||
self.assertNotIn("worker-1", job_data.worker_status)
|
||||
self.assertEqual(job_data.assigned_to_workers["worker-1"], [])
|
||||
|
||||
async def test_busy_probe_graces_worker_and_skips_requeue(self):
|
||||
async def _busy_probe(_url, _timeout):
|
||||
return {"exec_info": {"queue_remaining": 3}}
|
||||
|
||||
jt._probe_holder["fn"] = _busy_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-2")
|
||||
old_heartbeat = time.time() - 60.0
|
||||
job_data.worker_status["worker-1"] = old_heartbeat
|
||||
job_data.assigned_to_workers["worker-1"] = [0]
|
||||
prompt_server.distributed_pending_tile_jobs["job-2"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-2", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertIn("worker-1", job_data.worker_status)
|
||||
self.assertGreaterEqual(job_data.worker_status["worker-1"], old_heartbeat)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
async def test_completed_dynamic_task_is_not_requeued(self):
|
||||
async def _offline_probe(_url, _timeout):
|
||||
return None
|
||||
|
||||
jt._probe_holder["fn"] = _offline_probe
|
||||
prompt_server = jt._prompt_server_holder["value"]
|
||||
job_data = jt._ImageJobState("job-3")
|
||||
job_data.worker_status["worker-1"] = time.time() - 60.0
|
||||
job_data.assigned_to_workers["worker-1"] = [7]
|
||||
job_data.completed_images[7] = "complete"
|
||||
prompt_server.distributed_pending_tile_jobs["job-3"] = job_data
|
||||
|
||||
requeued = await jt._check_and_requeue_timed_out_workers("job-3", total_tasks=1)
|
||||
|
||||
self.assertEqual(requeued, 0)
|
||||
self.assertTrue(job_data.pending_images.empty())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,137 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_network_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py"
|
||||
|
||||
package_name = "dist_utils_testpkg"
|
||||
package_module = types.ModuleType(package_name)
|
||||
package_module.__path__ = [] # mark as package
|
||||
sys.modules[package_name] = package_module
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.logging"] = logging_module
|
||||
|
||||
server_module = types.ModuleType("server")
|
||||
server_module.PromptServer = types.SimpleNamespace(
|
||||
instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
|
||||
)
|
||||
sys.modules["server"] = server_module
|
||||
|
||||
if "aiohttp" not in sys.modules:
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
|
||||
class _TCPConnector:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.args = args
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _ClientSession:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
aiohttp_module.TCPConnector = _TCPConnector
|
||||
aiohttp_module.ClientSession = _ClientSession
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: {"payload": payload, "status": status}
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.network", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
network = _load_network_module()
|
||||
|
||||
|
||||
class NetworkHelpersTests(unittest.TestCase):
|
||||
def test_normalize_host_strips_protocol_and_path(self):
|
||||
self.assertEqual(network.normalize_host(" https://example.com/a/b "), "example.com")
|
||||
|
||||
def test_normalize_host_keeps_none(self):
|
||||
self.assertIsNone(network.normalize_host(None))
|
||||
|
||||
def test_build_worker_url_defaults_to_server_address(self):
|
||||
worker = {"id": "w1", "port": 8189}
|
||||
self.assertEqual(network.build_worker_url(worker, "/prompt"), "http://127.0.0.1:8189/prompt")
|
||||
|
||||
def test_build_worker_url_cloud_defaults_to_https(self):
|
||||
worker = {"id": "w2", "host": "foo.proxy.runpod.net", "port": 443}
|
||||
self.assertEqual(network.build_worker_url(worker), "https://foo.proxy.runpod.net")
|
||||
|
||||
def test_build_worker_url_keeps_explicit_scheme(self):
|
||||
worker = {"id": "w3", "host": "https://worker.example.com", "port": 1234}
|
||||
self.assertEqual(network.build_worker_url(worker, "/prompt"), "https://worker.example.com/prompt")
|
||||
|
||||
def test_build_master_url_uses_https_for_cloud_host(self):
|
||||
cfg = {"master": {"host": "demo.proxy.runpod.net"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"https://demo.proxy.runpod.net",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_scheme(self):
|
||||
cfg = {"master": {"host": "https://master.example.com/"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"https://master.example.com",
|
||||
)
|
||||
|
||||
def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
|
||||
cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8188",
|
||||
)
|
||||
|
||||
def test_build_master_url_keeps_explicit_port_in_host(self):
|
||||
cfg = {"master": {"host": "192.168.68.56:8001"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
def test_build_master_url_falls_back_to_server_address(self):
|
||||
cfg = {"master": {"host": "", "port": 8001}}
|
||||
prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
|
||||
self.assertEqual(
|
||||
network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8190",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_uses_loopback_for_local_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://127.0.0.1:8001",
|
||||
)
|
||||
|
||||
def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
|
||||
cfg = {"master": {"host": "192.168.68.56"}}
|
||||
prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
|
||||
worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
|
||||
self.assertEqual(
|
||||
network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
|
||||
"http://192.168.68.56:8001",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,196 +0,0 @@
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
PIL_AVAILABLE = True
|
||||
except ImportError:
|
||||
PIL_AVAILABLE = False
|
||||
|
||||
|
||||
def _load_payload_parsers_module():
|
||||
# payload_parsers.py has no relative imports; only stdlib + PIL
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "payload_parsers.py"
|
||||
spec = importlib.util.spec_from_file_location("upscale_payload_parsers", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
if PIL_AVAILABLE:
|
||||
pp = _load_payload_parsers_module()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_png_bytes(width=64, height=64, color=(128, 64, 32)):
|
||||
"""Return raw PNG bytes for a solid-colour image."""
|
||||
img = PILImage.new("RGB", (width, height), color=color)
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
class _MockFileField:
|
||||
"""Minimal multipart file-field stub."""
|
||||
|
||||
class _MockFile:
|
||||
def __init__(self, data: bytes):
|
||||
self._buf = io.BytesIO(data)
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._buf.read()
|
||||
|
||||
def __init__(self, data: bytes):
|
||||
self.file = self._MockFile(data)
|
||||
|
||||
|
||||
def _make_form(n_tiles, *, padding=None, extra_meta=None, image_color=(128, 64, 32)):
|
||||
"""Build a minimal form-data dict with `n_tiles` tile entries."""
|
||||
image_bytes = _make_png_bytes(color=image_color)
|
||||
metadata = []
|
||||
for i in range(n_tiles):
|
||||
entry = {
|
||||
"tile_idx": i,
|
||||
"x": i * 64,
|
||||
"y": 0,
|
||||
"extracted_width": 64,
|
||||
"extracted_height": 64,
|
||||
}
|
||||
if extra_meta and i < len(extra_meta):
|
||||
entry.update(extra_meta[i])
|
||||
metadata.append(entry)
|
||||
|
||||
form = {"tiles_metadata": json.dumps(metadata)}
|
||||
if padding is not None:
|
||||
form["padding"] = str(padding)
|
||||
for i in range(n_tiles):
|
||||
form[f"tile_{i}"] = _MockFileField(image_bytes)
|
||||
return form
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@unittest.skipUnless(PIL_AVAILABLE, "PIL not installed")
|
||||
class ParseTilesFromFormTests(unittest.TestCase):
|
||||
|
||||
# --- happy paths ---
|
||||
|
||||
def test_single_tile_returns_image_and_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertEqual(len(tiles), 1)
|
||||
tile = tiles[0]
|
||||
self.assertIsInstance(tile["image"], PILImage.Image)
|
||||
self.assertEqual(tile["tile_idx"], 0)
|
||||
self.assertEqual(tile["x"], 0)
|
||||
self.assertEqual(tile["y"], 0)
|
||||
self.assertEqual(tile["extracted_width"], 64)
|
||||
self.assertEqual(tile["extracted_height"], 64)
|
||||
|
||||
def test_multiple_tiles_preserve_count_order_and_coordinates(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(len(tiles), 3)
|
||||
for i, tile in enumerate(tiles):
|
||||
self.assertEqual(tile["tile_idx"], i)
|
||||
self.assertEqual(tiles[1]["x"], 64)
|
||||
self.assertEqual(tiles[2]["x"], 128)
|
||||
|
||||
def test_padding_is_parsed_from_form(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
|
||||
self.assertEqual(tiles[0]["padding"], 16)
|
||||
|
||||
def test_default_padding_is_zero(self):
|
||||
form = _make_form(1)
|
||||
form.pop("padding", None)
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_invalid_padding_string_falls_back_to_zero(self):
|
||||
form = _make_form(1)
|
||||
form["padding"] = "not_a_number"
|
||||
tiles = pp._parse_tiles_from_form(form)
|
||||
self.assertEqual(tiles[0]["padding"], 0)
|
||||
|
||||
def test_optional_batch_idx_included_when_present(self):
|
||||
extra = [{"batch_idx": 2}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["batch_idx"], 2)
|
||||
|
||||
def test_optional_global_idx_included_when_present(self):
|
||||
extra = [{"global_idx": 5}]
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, extra_meta=extra))
|
||||
self.assertEqual(tiles[0]["global_idx"], 5)
|
||||
|
||||
def test_batch_idx_and_global_idx_absent_when_not_in_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertNotIn("batch_idx", tiles[0])
|
||||
self.assertNotIn("global_idx", tiles[0])
|
||||
|
||||
# --- error cases ---
|
||||
|
||||
def test_missing_tiles_metadata_raises_value_error(self):
|
||||
with self.assertRaises(ValueError, msg="Missing tiles_metadata"):
|
||||
pp._parse_tiles_from_form({})
|
||||
|
||||
def test_invalid_json_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": "{not valid json}"}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_list_metadata_raises_value_error(self):
|
||||
form = {"tiles_metadata": json.dumps({"not": "a list"})}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_missing_tile_file_field_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
# tile_0 intentionally omitted
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_tile_field_without_file_attr_raises_value_error(self):
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": "plain string without .file",
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_non_image_bytes_raises_value_error(self):
|
||||
class _BadFileField:
|
||||
class _BadFile:
|
||||
def read(self):
|
||||
return b"this is definitely not image data"
|
||||
file = _BadFile()
|
||||
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": 0, "y": 0}]),
|
||||
"tile_0": _BadFileField(),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
def test_invalid_metadata_value_type_raises_value_error(self):
|
||||
"""Non-integer metadata fields (x, y, etc.) should raise ValueError."""
|
||||
form = {
|
||||
"tiles_metadata": json.dumps([{"tile_idx": 0, "x": "not_int", "y": 0}]),
|
||||
"tile_0": _MockFileField(_make_png_bytes()),
|
||||
}
|
||||
with self.assertRaises(ValueError):
|
||||
pp._parse_tiles_from_form(form)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,151 +0,0 @@
|
||||
import importlib.util
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _load_queue_request_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py"
|
||||
spec = importlib.util.spec_from_file_location("queue_request", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
queue_request = _load_queue_request_module()
|
||||
parse_queue_request_payload = queue_request.parse_queue_request_payload
|
||||
|
||||
|
||||
class QueueRequestPayloadTests(unittest.TestCase):
|
||||
def _base_payload(self):
|
||||
return {
|
||||
"prompt": {"1": {"class_type": "Anything"}},
|
||||
"enabled_worker_ids": ["worker-1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
|
||||
def test_normalizes_enabled_worker_ids(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["enabled_worker_ids"] = ["a", 2, 3]
|
||||
payload_data["delegate_master"] = True
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
|
||||
self.assertTrue(payload.delegate_master)
|
||||
|
||||
def test_supports_legacy_workers_field(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}]
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
|
||||
|
||||
def test_supports_auto_prepare_prompt_fallback(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("prompt", None)
|
||||
payload_data["auto_prepare"] = True
|
||||
payload_data["workflow"] = {
|
||||
"prompt": {
|
||||
"10": {"class_type": "DistributedCollector"},
|
||||
}
|
||||
}
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertIn("10", payload.prompt)
|
||||
self.assertTrue(payload.auto_prepare)
|
||||
|
||||
def test_normalizes_trace_execution_id(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = " exec_123 "
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertEqual(payload.trace_execution_id, "exec_123")
|
||||
|
||||
def test_blank_trace_execution_id_normalizes_to_none(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = " "
|
||||
payload = parse_queue_request_payload(
|
||||
payload_data
|
||||
)
|
||||
self.assertIsNone(payload.trace_execution_id)
|
||||
|
||||
def test_auto_prepare_defaults_true(self):
|
||||
payload = parse_queue_request_payload(self._base_payload())
|
||||
self.assertTrue(payload.auto_prepare)
|
||||
|
||||
def test_workers_field_must_be_list(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
payload_data["workers"] = "worker-a"
|
||||
with self.assertRaisesRegex(ValueError, "Field 'workers' must be a list"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_trace_execution_id_must_be_string(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["trace_execution_id"] = 123
|
||||
with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("prompt", None)
|
||||
payload_data["auto_prepare"] = False
|
||||
payload_data["workflow"] = {
|
||||
"prompt": {"10": {"class_type": "DistributedCollector"}},
|
||||
}
|
||||
payload = parse_queue_request_payload(payload_data)
|
||||
self.assertIn("10", payload.prompt)
|
||||
self.assertFalse(payload.auto_prepare)
|
||||
|
||||
def test_auto_prepare_must_be_boolean(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["auto_prepare"] = "true"
|
||||
with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_delegate_master_type_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["delegate_master"] = "yes"
|
||||
with self.assertRaisesRegex(ValueError, "delegate_master must be a boolean"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_enabled_worker_ids_type_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data["enabled_worker_ids"] = "worker-a"
|
||||
with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_invalid_top_level_payload_raises(self):
|
||||
with self.assertRaisesRegex(ValueError, "Expected a JSON object body"):
|
||||
parse_queue_request_payload(["not", "an", "object"])
|
||||
|
||||
def test_missing_prompt_raises(self):
|
||||
with self.assertRaisesRegex(ValueError, "Field 'prompt' must be an object"):
|
||||
parse_queue_request_payload(
|
||||
{
|
||||
"workflow": {},
|
||||
"enabled_worker_ids": ["worker-1"],
|
||||
"client_id": "client-1",
|
||||
}
|
||||
)
|
||||
|
||||
def test_missing_enabled_worker_ids_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("enabled_worker_ids", None)
|
||||
with self.assertRaisesRegex(ValueError, "enabled_worker_ids required"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
def test_missing_client_id_raises(self):
|
||||
payload_data = self._base_payload()
|
||||
payload_data.pop("client_id", None)
|
||||
with self.assertRaisesRegex(ValueError, "client_id required"):
|
||||
parse_queue_request_payload(payload_data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,167 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
def load_module():
|
||||
spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
client = load_module()
|
||||
|
||||
|
||||
def workflow_for_request_size(target_size, timeout_seconds):
|
||||
"""Build an ASCII workflow whose compact UTF-8 request body has exact size."""
|
||||
workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}}
|
||||
|
||||
def body_size(candidate):
|
||||
payload = {
|
||||
"input": {"workflow": candidate, "images": []},
|
||||
"policy": {
|
||||
"executionTimeout": timeout_seconds * 1000,
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
return len(json.dumps(
|
||||
payload, allow_nan=False, ensure_ascii=False, separators=(",", ":")
|
||||
).encode("utf-8"))
|
||||
|
||||
padding = target_size - body_size(workflow)
|
||||
if padding < 0:
|
||||
raise AssertionError("target is smaller than the request fixture")
|
||||
workflow["1"]["inputs"]["value"] = "x" * padding
|
||||
if body_size(workflow) != target_size:
|
||||
raise AssertionError("request fixture did not reach the exact target size")
|
||||
return workflow
|
||||
|
||||
class Response:
|
||||
def __init__(self, status=200, payload=None):
|
||||
self.status, self.payload = status, payload
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
class Session:
|
||||
def __init__(self, responses):
|
||||
self.responses, self.calls = list(responses), []
|
||||
def request(self, method, url, **kwargs):
|
||||
self.calls.append((method, url, kwargs))
|
||||
response = self.responses.pop(0)
|
||||
if isinstance(response, Exception):
|
||||
raise response
|
||||
return response
|
||||
|
||||
class ClientTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_submit_exact_generic_contract_and_once(self):
|
||||
session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800)
|
||||
self.assertEqual(result, "job-1")
|
||||
method, url, kwargs = session.calls[0]
|
||||
self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run"))
|
||||
self.assertIsInstance(kwargs["data"], bytes)
|
||||
sent_payload = json.loads(kwargs["data"])
|
||||
self.assertEqual(set(sent_payload["input"]), {"workflow", "images"})
|
||||
self.assertEqual(sent_payload["policy"]["ttl"], 1800000)
|
||||
self.assertFalse(kwargs["allow_redirects"])
|
||||
self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value")
|
||||
|
||||
async def test_submit_timeout_is_ambiguous_and_never_retried(self):
|
||||
session = Session([TimeoutError("secret-value in unsafe error")])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
with self.assertRaises(client.AmbiguousSubmission) as ctx:
|
||||
await api.submit({}, [], 30)
|
||||
self.assertNotIn("secret-value", str(ctx.exception))
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
|
||||
async def test_http_submit_failure_not_retried_or_leaked(self):
|
||||
for status in [401, 429, 500, 302]:
|
||||
session = Session([Response(status, {"error": "secret-value"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
with self.assertRaises(client.RunpodError) as ctx:
|
||||
await api.submit({}, [], 30)
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
self.assertNotIn("secret-value", str(ctx.exception))
|
||||
|
||||
async def test_health_zero_workers_is_valid(self):
|
||||
session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
self.assertEqual((await api.health())["workers"]["idle"], 0)
|
||||
|
||||
async def test_invalid_provider_id_no_network(self):
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
client.RunpodClient(session, "https://attacker.invalid", "key")
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
with self.assertRaises(ValueError):
|
||||
await api.status("../other")
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_status_and_cancel_only_target_exact_job(self):
|
||||
session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
await api.status("job-1")
|
||||
await api.cancel("job-1")
|
||||
self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")])
|
||||
|
||||
async def test_missing_id_is_ambiguous(self):
|
||||
api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key")
|
||||
with self.assertRaises(client.AmbiguousSubmission):
|
||||
await api.submit({}, [], 30)
|
||||
|
||||
async def test_large_payload_rejected_before_submit(self):
|
||||
session = Session([])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
with self.assertRaises(ValueError):
|
||||
await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30)
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self):
|
||||
session = Session([Response(payload={"id": "boundary-job"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
timeout_seconds = 37
|
||||
workflow = workflow_for_request_size(10_000_000, timeout_seconds)
|
||||
|
||||
result = await api.submit(workflow, [], timeout_seconds)
|
||||
|
||||
self.assertEqual(result, "boundary-job")
|
||||
sent = session.calls[0][2]["data"]
|
||||
self.assertIsInstance(sent, bytes)
|
||||
self.assertEqual(len(sent), 10_000_000)
|
||||
self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000)
|
||||
|
||||
async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self):
|
||||
session = Session([Response(payload={"id": "must-not-be-used"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
workflow = workflow_for_request_size(10_000_001, 37)
|
||||
workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:]
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "10,000,001") as ctx:
|
||||
await api.submit(workflow, [], 37)
|
||||
|
||||
self.assertNotIn("private graph value", str(ctx.exception))
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self):
|
||||
session = Session([Response(payload={"id": "unicode-job"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}}
|
||||
|
||||
await api.submit(workflow, [], 41)
|
||||
|
||||
sent = session.calls[0][2]["data"]
|
||||
expected = json.dumps({
|
||||
"input": {"workflow": workflow, "images": []},
|
||||
"policy": {"executionTimeout": 41_000, "ttl": 41_000},
|
||||
}, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
self.assertEqual(sent, expected)
|
||||
self.assertGreater(len(sent), len(sent.decode("utf-8")))
|
||||
@@ -1,202 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
|
||||
class _FakeManager:
|
||||
def __init__(self, *, pending=False, owned=True):
|
||||
self.pending = pending
|
||||
self.owned = owned
|
||||
self.pending_calls = []
|
||||
self.owned_calls = []
|
||||
self.cancel_calls = []
|
||||
|
||||
def is_pending(self, collector_job_id, participant_id):
|
||||
self.pending_calls.append((collector_job_id, participant_id))
|
||||
return self.pending
|
||||
|
||||
def owns(self, collector_job_id, participant_id):
|
||||
self.owned_calls.append((collector_job_id, participant_id))
|
||||
return self.owned
|
||||
|
||||
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
|
||||
self.cancel_calls.append({
|
||||
"worker_id": worker_id,
|
||||
"execution_id": execution_id,
|
||||
"collector_job_id": collector_job_id,
|
||||
})
|
||||
return 1
|
||||
|
||||
|
||||
def _set_up_collector(module, job_id, queue, manager=None):
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_pending_jobs = {job_id: queue}
|
||||
if manager is not None:
|
||||
module.prompt_server.distributed_serverless = manager
|
||||
return module.DistributedCollectorNode()
|
||||
|
||||
|
||||
def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-pending"
|
||||
participant = "serverless:pool:1"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
probe_calls = []
|
||||
module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
|
||||
|
||||
async def deliver_after_provider_work():
|
||||
await asyncio.sleep(0.14)
|
||||
queue.put_nowait({
|
||||
"worker_id": participant,
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 2, 2, 2),
|
||||
"is_last": True,
|
||||
})
|
||||
|
||||
async def run_collection():
|
||||
producer = asyncio.create_task(deliver_after_provider_work())
|
||||
result = await collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids=f'["{participant}"]',
|
||||
delegate_only=True,
|
||||
)
|
||||
await producer
|
||||
return result
|
||||
|
||||
images, _audio = asyncio.run(run_collection())
|
||||
|
||||
assert tuple(images.shape) == (1, 2, 2, 2)
|
||||
assert manager.pending_calls
|
||||
assert probe_calls == []
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-error"
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"serverless_error": "Provider execution failed safely",
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.cancel_calls == [{
|
||||
"worker_id": None,
|
||||
"execution_id": None,
|
||||
"collector_job_id": job_id,
|
||||
}]
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-stopped"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=False, owned=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.owned_calls
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_interruption_cancels_serverless_job_and_removes_collector_queue():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-interrupted"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
interrupted = module.comfy.model_management.InterruptProcessingException("stop")
|
||||
module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
|
||||
|
||||
with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-shape-mismatch"
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 3, 2, 2),
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
|
||||
asyncio.run(collector.execute(
|
||||
images=local_images,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
))
|
||||
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-regular-timeout"
|
||||
queue = asyncio.Queue()
|
||||
collector = _set_up_collector(module, job_id, queue)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
|
||||
|
||||
images, _audio = asyncio.run(collector.execute(
|
||||
images=local_images,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
))
|
||||
|
||||
assert torch.equal(images, local_images)
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
@@ -1,398 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
|
||||
def _load_serverless_config_module(*, without_fcntl=False):
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py"
|
||||
module_name = (
|
||||
"dist_serverless_config_no_fcntl_test_module"
|
||||
if without_fcntl
|
||||
else "dist_serverless_config_test_module"
|
||||
)
|
||||
package_name = module_name + "_pkg"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(module_path.parent)]
|
||||
sys.modules[package_name] = package
|
||||
module_name = package_name + ".serverless_config"
|
||||
sys.modules.pop(module_name, None)
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
if without_fcntl:
|
||||
with mock.patch.dict(sys.modules, {"fcntl": None}):
|
||||
spec.loader.exec_module(module)
|
||||
else:
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
serverless_config = _load_serverless_config_module()
|
||||
|
||||
|
||||
class ValidateServerlessWorkerTests(unittest.TestCase):
|
||||
def test_returns_copy_with_runpod_defaults(self):
|
||||
worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}
|
||||
|
||||
result = serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
self.assertEqual(
|
||||
result,
|
||||
{
|
||||
"id": "worker-1",
|
||||
"endpoint_id": "ep_123-abc",
|
||||
"name": "GPU",
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"jobs_per_run": 1,
|
||||
"concurrency": 1,
|
||||
"timeout_seconds": 1800,
|
||||
},
|
||||
)
|
||||
self.assertIsNot(result, worker)
|
||||
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
|
||||
|
||||
def test_accepts_inclusive_numeric_bounds(self):
|
||||
for jobs, concurrency, timeout in ((1, 1, 30), (16, 16, 86400)):
|
||||
with self.subTest(jobs=jobs, concurrency=concurrency, timeout=timeout):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{
|
||||
"id": "worker-1",
|
||||
"endpoint_id": "ep-123",
|
||||
"jobs_per_run": jobs,
|
||||
"concurrency": concurrency,
|
||||
"timeout_seconds": timeout,
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["jobs_per_run"], jobs)
|
||||
self.assertEqual(result["concurrency"], concurrency)
|
||||
self.assertEqual(result["timeout_seconds"], timeout)
|
||||
|
||||
def test_rejects_missing_or_empty_id_and_endpoint(self):
|
||||
base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
|
||||
with self.subTest(field=field, value=value):
|
||||
worker = dict(base)
|
||||
worker[field] = value
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
worker = dict(base)
|
||||
del worker["id"]
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
def test_endpoint_id_must_be_ascii_opaque_identifier(self):
|
||||
for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"):
|
||||
with self.subTest(endpoint_id=endpoint_id):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": endpoint_id}
|
||||
)
|
||||
|
||||
def test_endpoint_id_matches_runpod_client_length_limit(self):
|
||||
valid_id = "e" * 200
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": valid_id}
|
||||
)
|
||||
self.assertEqual(result["endpoint_id"], valid_id)
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "e" * 201}
|
||||
)
|
||||
|
||||
def test_enabled_must_be_boolean_when_supplied(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for enabled in (True, False):
|
||||
with self.subTest(enabled=enabled):
|
||||
self.assertIs(
|
||||
serverless_config.validate_serverless_worker(
|
||||
{**base, "enabled": enabled}
|
||||
)["enabled"],
|
||||
enabled,
|
||||
)
|
||||
for enabled in ("false", 0, 1, None):
|
||||
with self.subTest(enabled=enabled):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{**base, "enabled": enabled}
|
||||
)
|
||||
|
||||
def test_rejects_unknown_provider_and_non_serverless_type(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for extra in ({"provider": "other"}, {"type": "local"}):
|
||||
with self.subTest(extra=extra):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker({**base, **extra})
|
||||
|
||||
def test_numeric_fields_reject_bool_float_and_out_of_range(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
invalid_values = {
|
||||
"jobs_per_run": (True, 1.0, 0, 17),
|
||||
"concurrency": (False, 2.5, 0, 17),
|
||||
"timeout_seconds": (True, 30.0, 29, 86401),
|
||||
}
|
||||
for field, values in invalid_values.items():
|
||||
for value in values:
|
||||
with self.subTest(field=field, value=value):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker({**base, field: value})
|
||||
|
||||
def test_requires_worker_object(self):
|
||||
for worker in (None, [], "worker"):
|
||||
with self.subTest(worker=worker):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
|
||||
class SecretStoreTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self.tempdir.name)
|
||||
self.path = self.root / "private" / "runpod-secrets.json"
|
||||
|
||||
def tearDown(self):
|
||||
self.tempdir.cleanup()
|
||||
|
||||
def test_set_get_has_delete_and_persists_as_json_mapping(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
store.set("worker-1", "secret-value")
|
||||
|
||||
self.assertEqual(store.get("worker-1"), "secret-value")
|
||||
self.assertTrue(store.has("worker-1"))
|
||||
self.assertIsNone(store.get("missing"))
|
||||
self.assertFalse(store.has("missing"))
|
||||
self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"})
|
||||
self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value")
|
||||
|
||||
store.delete("worker-1")
|
||||
self.assertIsNone(store.get("worker-1"))
|
||||
self.assertFalse(store.has("worker-1"))
|
||||
|
||||
def test_rejects_api_keys_with_newlines_or_over_4096_characters(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097):
|
||||
with self.subTest(key_length=len(key)):
|
||||
with self.assertRaises(serverless_config.SecretStoreError):
|
||||
store.set("worker-1", key)
|
||||
store.set("worker-1", "k" * 4096)
|
||||
self.assertEqual(store.get("worker-1"), "k" * 4096)
|
||||
|
||||
def test_new_secret_file_and_directory_have_restrictive_modes(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
store.set("worker-1", "secret-value")
|
||||
|
||||
self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600)
|
||||
self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700)
|
||||
|
||||
def test_missing_store_reads_as_empty_without_creating_files(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
self.assertIsNone(store.get("worker-1"))
|
||||
self.assertFalse(store.has("worker-1"))
|
||||
self.assertFalse(self.path.parent.exists())
|
||||
|
||||
def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self):
|
||||
self.path.parent.mkdir()
|
||||
corrupt = b'{"partial": "NEVER-LEAK-THIS-secret'
|
||||
self.path.write_bytes(corrupt)
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
for operation in (
|
||||
lambda: store.get("worker-1"),
|
||||
lambda: store.set("worker-1", "replacement-secret"),
|
||||
lambda: store.delete("worker-1"),
|
||||
):
|
||||
with self.subTest(operation=operation):
|
||||
with self.assertRaises(Exception) as raised:
|
||||
operation()
|
||||
self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception))
|
||||
self.assertNotIn("replacement-secret", str(raised.exception))
|
||||
self.assertEqual(self.path.read_bytes(), corrupt)
|
||||
|
||||
def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self):
|
||||
for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'):
|
||||
with self.subTest(corrupt=corrupt):
|
||||
self.path.parent.mkdir(exist_ok=True)
|
||||
self.path.write_bytes(corrupt)
|
||||
with self.assertRaises(Exception) as raised:
|
||||
serverless_config.SecretStore(self.path).get("worker")
|
||||
self.assertNotIn("opaque-stored-value", str(raised.exception))
|
||||
|
||||
def test_rejects_symlink_file_and_preserves_link_target(self):
|
||||
target = self.root / "target.json"
|
||||
target.write_text('{"worker-1": "target-secret"}', encoding="utf-8")
|
||||
self.path.parent.mkdir()
|
||||
self.path.symlink_to(target)
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.get("worker-1")
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "replacement-secret")
|
||||
self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}')
|
||||
|
||||
def test_rejects_symlinked_parent_directory(self):
|
||||
real_dir = self.root / "real"
|
||||
real_dir.mkdir()
|
||||
link_dir = self.root / "linked"
|
||||
link_dir.symlink_to(real_dir, target_is_directory=True)
|
||||
store = serverless_config.SecretStore(link_dir / "secrets.json")
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "secret-value")
|
||||
self.assertFalse((real_dir / "secrets.json").exists())
|
||||
|
||||
def test_rejects_parent_traversal_in_store_path(self):
|
||||
store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json")
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "secret-value")
|
||||
self.assertFalse((self.root / "secrets.json").exists())
|
||||
|
||||
def test_threaded_updates_from_independent_instances_are_not_lost(self):
|
||||
workers = 12
|
||||
barrier = threading.Barrier(workers)
|
||||
errors = []
|
||||
|
||||
def write_one(index):
|
||||
try:
|
||||
barrier.wait(timeout=5)
|
||||
serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}")
|
||||
except Exception as exc: # collected to fail in the test thread
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(timeout=10)
|
||||
|
||||
self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish")
|
||||
self.assertEqual(errors, [])
|
||||
saved = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)})
|
||||
|
||||
|
||||
class PublicConfigTests(unittest.TestCase):
|
||||
def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self):
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json")
|
||||
store.set("serverless-1", "stored-api-key")
|
||||
config = {
|
||||
"workers": [
|
||||
{
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"endpoint_id": "ep-123",
|
||||
"api_key": "inline-secret",
|
||||
"clear_api_key": "another-secret",
|
||||
"credential": "credential-secret",
|
||||
"nested": {"runpod_credential": "nested-secret", "safe": "kept"},
|
||||
"name": "remote GPU",
|
||||
},
|
||||
{"id": "serverless-2", "type": "serverless", "name": "not configured"},
|
||||
{"id": "local-1", "type": "local", "api_key": "local-secret"},
|
||||
],
|
||||
"settings": {"mode": "preserved", "credentials_note": "remove-me"},
|
||||
}
|
||||
original = json.loads(json.dumps(config))
|
||||
|
||||
result = serverless_config.public_config(config, store)
|
||||
|
||||
first, second, local = result["workers"]
|
||||
self.assertEqual(first["api_key_configured"], True)
|
||||
self.assertEqual(second["api_key_configured"], False)
|
||||
self.assertNotIn("api_key_configured", local)
|
||||
self.assertNotIn("api_key", first)
|
||||
self.assertNotIn("clear_api_key", first)
|
||||
self.assertNotIn("credential", first)
|
||||
self.assertNotIn("runpod_credential", first["nested"])
|
||||
self.assertEqual(first["nested"]["safe"], "kept")
|
||||
self.assertNotIn("api_key", local)
|
||||
self.assertNotIn("credentials_note", result["settings"])
|
||||
self.assertEqual(first["name"], "remote GPU")
|
||||
self.assertEqual(config, original)
|
||||
self.assertIsNot(result, config)
|
||||
self.assertIsNot(result["workers"], config["workers"])
|
||||
|
||||
def test_redacts_normalized_secret_field_aliases(self):
|
||||
config = {
|
||||
"workers": [{
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"apiKey": "camel-api-key",
|
||||
"APIKEY": "upper-api-key",
|
||||
"clearApiKey": "camel-clear-key",
|
||||
"CLEAR_API_KEY": "upper-clear-key",
|
||||
"runpod_CREDENTIAL": "provider-credential",
|
||||
"nested": {"credential_note": "nested-credential"},
|
||||
"safe": "preserved",
|
||||
}]
|
||||
}
|
||||
|
||||
result = serverless_config.public_config(
|
||||
config, serverless_config.SecretStore(self.root_path())
|
||||
)
|
||||
|
||||
worker = result["workers"][0]
|
||||
self.assertEqual(worker, {
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"nested": {},
|
||||
"safe": "preserved",
|
||||
"api_key_configured": False,
|
||||
})
|
||||
|
||||
def test_handles_non_dict_config_without_mutating_it(self):
|
||||
config = ["safe", {"api_key": "secret"}]
|
||||
result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path()))
|
||||
self.assertEqual(result, ["safe", {}])
|
||||
self.assertEqual(config, ["safe", {"api_key": "secret"}])
|
||||
|
||||
def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self):
|
||||
module = _load_serverless_config_module(without_fcntl=True)
|
||||
store = module.SecretStore(self.root_path())
|
||||
config = {
|
||||
"workers": [
|
||||
{"id": "local-1", "type": "local", "api_key": "inline-secret"}
|
||||
]
|
||||
}
|
||||
|
||||
result = module.public_config(config, store)
|
||||
|
||||
self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]})
|
||||
self.assertEqual(config["workers"][0]["api_key"], "inline-secret")
|
||||
|
||||
def test_secret_operations_fail_clearly_without_fcntl(self):
|
||||
module = _load_serverless_config_module(without_fcntl=True)
|
||||
store = module.SecretStore(self.root_path())
|
||||
unsupported_message = "Serverless credential storage is currently supported on POSIX masters only"
|
||||
|
||||
for operation in (
|
||||
lambda: store.get("worker-1"),
|
||||
lambda: store.set("worker-1", "api-key"),
|
||||
lambda: store.delete("worker-1"),
|
||||
lambda: store.has("worker-1"),
|
||||
):
|
||||
with self.subTest(operation=operation):
|
||||
with self.assertRaisesRegex(module.SecretStoreError, unsupported_message):
|
||||
operation()
|
||||
|
||||
@staticmethod
|
||||
def root_path():
|
||||
return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,183 +0,0 @@
|
||||
"""Offline integration: real graph preparation + orchestration, fake provider."""
|
||||
import asyncio
|
||||
import aiohttp as real_aiohttp
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_serverless_integration"
|
||||
|
||||
|
||||
def module(name, **attrs):
|
||||
value = types.ModuleType(name)
|
||||
value.__dict__.update(attrs)
|
||||
return value
|
||||
|
||||
|
||||
def load_modules():
|
||||
modules = {}
|
||||
for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]:
|
||||
modules[PKG + part] = module(PKG + part, __path__=[str(directory)])
|
||||
ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn))
|
||||
modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps))
|
||||
modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None)
|
||||
modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None)
|
||||
modules[PKG + ".utils.config"] = module("config", load_config=lambda: {})
|
||||
modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock())
|
||||
modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master")
|
||||
modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock())
|
||||
modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), sync_worker_media=AsyncMock())
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py")
|
||||
q = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(q)
|
||||
# Keep imported package namespace only; restore shared Comfy runtime stub.
|
||||
modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)})
|
||||
return q, modules, ps
|
||||
|
||||
|
||||
class IntegrationTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.q, modules, self.ps = load_modules()
|
||||
self.patch = patch.dict(sys.modules, modules)
|
||||
self.patch.start()
|
||||
self.addCleanup(self.patch.stop)
|
||||
self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30}
|
||||
self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}}
|
||||
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
|
||||
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
||||
self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None)
|
||||
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager)
|
||||
sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[]))
|
||||
self.prompt = {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
|
||||
}
|
||||
|
||||
async def test_real_compiler_slots_to_master_collector_without_callbacks(self):
|
||||
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.assertEqual(result[2], 2)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.manager.start.assert_awaited_once()
|
||||
jobs = self.manager.start.call_args.args[2]
|
||||
self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"])
|
||||
self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2)
|
||||
for job in jobs:
|
||||
graph = job["workflow"]
|
||||
self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values()))
|
||||
self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values()))
|
||||
self.assertNotIn("http://master", json.dumps(graph))
|
||||
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
||||
self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"])
|
||||
|
||||
async def test_compile_pack_provider_decode_and_collect_real_components(self):
|
||||
"""Exercise the whole image path, mocking only the provider and Comfy queue."""
|
||||
import base64
|
||||
import io
|
||||
import tempfile
|
||||
import torch
|
||||
from PIL import Image
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
for leaf in ("serverless_media", "serverless_runtime"):
|
||||
name = PKG + ".api.orchestration." + leaf
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py"))
|
||||
real_module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = real_module
|
||||
spec.loader.exec_module(real_module)
|
||||
runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
image_buffer = io.BytesIO()
|
||||
Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG")
|
||||
image_bytes = image_buffer.getvalue()
|
||||
(root / "source.png").write_bytes(image_bytes)
|
||||
paths = module("folder_paths", get_input_directory=lambda: directory,
|
||||
get_output_directory=lambda: directory, get_temp_directory=lambda: directory,
|
||||
get_annotated_filepath=lambda name: str(root / name))
|
||||
submissions = []
|
||||
|
||||
class FakeProvider:
|
||||
async def submit(self, graph, images, timeout_seconds):
|
||||
self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage")
|
||||
submissions.append((graph, images))
|
||||
self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage")
|
||||
assert self_image == images[0]["name"]
|
||||
assert base64.b64decode(images[0]["image"]) == image_bytes
|
||||
return "provider-job-" + str(len(submissions))
|
||||
|
||||
async def status(self, _job_id):
|
||||
filename = self.prefix + "_00001_.png"
|
||||
return {"status": "COMPLETED", "output": {"images": [{
|
||||
"filename": filename, "artifact_path": filename, "artifact_type": "output",
|
||||
"type": "base64", "data": base64.b64encode(image_bytes).decode(),
|
||||
}]}}
|
||||
|
||||
manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None)
|
||||
manager._client = AsyncMock(side_effect=lambda _: FakeProvider())
|
||||
self.ps.distributed_serverless = manager
|
||||
self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}}
|
||||
collector_module = _load_collector_module()
|
||||
collector_module.prompt_server = self.ps
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}):
|
||||
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
||||
inputs = master_graph["2"]["inputs"]
|
||||
images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute(
|
||||
images=None, audio=None, multi_job_id=inputs["multi_job_id"],
|
||||
enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True,
|
||||
), timeout=3)
|
||||
await asyncio.gather(*manager.tasks.values())
|
||||
self.assertEqual(result[2], 2)
|
||||
self.assertEqual(len(submissions), 2)
|
||||
self.assertEqual(tuple(images.shape), (2, 3, 2, 3))
|
||||
self.assertEqual(images.dtype, torch.float32)
|
||||
self.assertTrue(torch.equal(images[0], images[1]))
|
||||
self.assertEqual(float(images[0, 0, 0, 0]), 1.0)
|
||||
self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"])
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
async def test_audio_rejected_before_any_dispatch_or_master_queue(self):
|
||||
self.prompt["2"]["inputs"]["audio"] = ["9", 0]
|
||||
self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}}
|
||||
with self.assertRaises(ValueError):
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.manager.start.assert_not_awaited()
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_disabled_pool_excluded(self):
|
||||
self.pool["enabled"] = False
|
||||
self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), [])
|
||||
|
||||
async def test_master_validation_failure_submits_no_paid_job(self):
|
||||
self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph")
|
||||
with self.assertRaisesRegex(ValueError, "invalid local graph"):
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.manager.start.assert_not_awaited()
|
||||
|
||||
async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self):
|
||||
media_module = sys.modules[PKG + ".api.orchestration.serverless_media"]
|
||||
media_module.prepare_images.return_value = [{
|
||||
"name": "private-image-name.png",
|
||||
"image": "x" * 10_000_001,
|
||||
}]
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "10,000,") as ctx:
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
|
||||
self.assertNotIn("private-image-name.png", str(ctx.exception))
|
||||
self.assertNotIn("x" * 100, str(ctx.exception))
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
self.manager.start.assert_not_awaited()
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
@@ -1,192 +0,0 @@
|
||||
"""Billing-safe lifecycle regressions: local completion is not remote completion."""
|
||||
import asyncio
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from test_serverless_runtime import runtime, Secrets
|
||||
|
||||
|
||||
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
|
||||
self.path = Path(self.temp.name) / "jobs.json"
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
|
||||
self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]}
|
||||
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
|
||||
self.client_patch.start()
|
||||
self.addCleanup(self.client_patch.stop)
|
||||
self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
|
||||
self.decode_patch.start()
|
||||
self.addCleanup(self.decode_patch.stop)
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
|
||||
def second(self, **changes):
|
||||
return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes))
|
||||
|
||||
def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"):
|
||||
self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id,
|
||||
status=status, created_at=1, updated_at=1)
|
||||
|
||||
async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self):
|
||||
cancel_started, release_cancel = asyncio.Event(), asyncio.Event()
|
||||
self.api.submit.side_effect = ["job-1", "job-2"]
|
||||
self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}]
|
||||
async def cancel(_job_id):
|
||||
cancel_started.set()
|
||||
await release_cancel.wait()
|
||||
return {"status": "CANCELLED"}
|
||||
self.api.cancel.side_effect = cancel
|
||||
await self.manager.start("exec", "prompt", [self.prepared, self.second()])
|
||||
try:
|
||||
await asyncio.wait_for(cancel_started.wait(), 1)
|
||||
await asyncio.sleep(0.01)
|
||||
self.assertEqual(self.api.submit.await_count, 1)
|
||||
finally:
|
||||
release_cancel.set()
|
||||
await self.manager.wait_for_execution("exec")
|
||||
self.assertEqual(self.api.submit.await_count, 2)
|
||||
|
||||
async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self):
|
||||
self.api.status.side_effect = runtime.RunpodError("poll failed")
|
||||
self.api.cancel.return_value = {"status": "IN_PROGRESS"}
|
||||
await self.manager.start("first", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("first")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
await self.manager.start("second", "prompt", [self.second()])
|
||||
try:
|
||||
await asyncio.sleep(0.01)
|
||||
self.assertEqual(self.api.submit.await_count, 1)
|
||||
self.api.status.side_effect = None
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": {}}
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
await asyncio.wait_for(self.manager.wait_for_execution("second"), 1)
|
||||
self.assertEqual(self.api.submit.await_count, 2)
|
||||
finally:
|
||||
self.api.cancel.return_value = {"status": "CANCELLED"}
|
||||
await self.manager.cancel()
|
||||
|
||||
async def test_restarted_manager_reserves_known_unresolved_jobs(self):
|
||||
self.retained("IN_PROGRESS")
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertFalse(self.manager.pool_is_active("other"))
|
||||
await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)])
|
||||
await self.manager.wait_for_execution("exec")
|
||||
self.api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT")
|
||||
|
||||
async def test_recovery_closes_never_submitted_waiting_records(self):
|
||||
self.retained("WAITING", job_id=None)
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_recovery_keeps_interrupted_submission_unresolved(self):
|
||||
self.retained("SUBMITTING", job_id=None)
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.api.status.assert_not_awaited()
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self):
|
||||
self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown")
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("exec")
|
||||
record = self.manager.ledger.all()[0]
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False)
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True)
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED")
|
||||
self.api.submit.assert_awaited_once()
|
||||
|
||||
async def test_cancel_already_completed_provider_job_is_terminal(self):
|
||||
self.retained()
|
||||
self.api.cancel.return_value = {"status": "COMPLETED"}
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
|
||||
async def test_reconciliation_failure_never_releases_capacity(self):
|
||||
self.retained()
|
||||
self.api.status.side_effect = runtime.RunpodError("not reachable")
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
|
||||
|
||||
async def test_cancel_before_first_task_step_finalizes_local_record(self):
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
self.assertFalse(self.server.distributed_pending_jobs["collector"].empty())
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_live_record_cannot_be_acknowledged_or_reconciled(self):
|
||||
polling = asyncio.Event()
|
||||
async def status(_job_id):
|
||||
polling.set()
|
||||
await asyncio.Event().wait()
|
||||
self.api.status.side_effect = status
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await asyncio.wait_for(polling.wait(), 1)
|
||||
record = self.manager.ledger.all()[0]
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.assertEqual(self.api.status.await_count, 1)
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
|
||||
async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self):
|
||||
polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def status(_job_id):
|
||||
polling.set()
|
||||
await asyncio.Event().wait()
|
||||
async def cancel(_job_id):
|
||||
cancelling.set()
|
||||
await release.wait()
|
||||
return {"status": "CANCELLED"}
|
||||
self.api.status.side_effect = status
|
||||
self.api.cancel.side_effect = cancel
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await polling.wait()
|
||||
first = asyncio.create_task(self.manager.cancel(worker_id="pool"))
|
||||
await cancelling.wait()
|
||||
second = asyncio.create_task(self.manager.cancel(worker_id="pool"))
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
await asyncio.gather(first, second)
|
||||
self.api.cancel.assert_awaited_once_with("job")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_partial_start_cleanup_does_not_leave_waiting_protection(self):
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.start("exec", "prompt", [self.prepared, self.second(concurrency=2)])
|
||||
await self.manager.cancel(execution_id="exec")
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_public_jobs_never_hides_unresolved_recovery_ids(self):
|
||||
for index in range(22):
|
||||
self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index)
|
||||
for index in range(5):
|
||||
self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index)
|
||||
records = self.manager.public_jobs("pool")
|
||||
self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)})
|
||||
self.assertTrue(all(record["status"] == "UNKNOWN" for record in records))
|
||||
@@ -1,447 +0,0 @@
|
||||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py"
|
||||
|
||||
|
||||
def load_module():
|
||||
spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
media = load_module()
|
||||
|
||||
|
||||
def load_runpod_module():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py"
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
runpod = load_runpod_module()
|
||||
|
||||
|
||||
def png_bytes(mode="RGB", size=(2, 1), color=None):
|
||||
if color is None:
|
||||
color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30)
|
||||
buffer = io.BytesIO()
|
||||
Image.new(mode, size, color).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class FolderPathsFixture:
|
||||
def __init__(self, root):
|
||||
self.input = root / "input"
|
||||
self.output = root / "output"
|
||||
self.temp = root / "temp"
|
||||
for directory in (self.input, self.output, self.temp):
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_input_directory(self):
|
||||
return str(self.input)
|
||||
|
||||
def get_output_directory(self):
|
||||
return str(self.output)
|
||||
|
||||
def get_temp_directory(self):
|
||||
return str(self.temp)
|
||||
|
||||
def get_annotated_filepath(self, filename):
|
||||
value = filename.replace("\\", "/")
|
||||
# Match real ComfyUI folder_paths: root annotations are suffixes,
|
||||
# not directory prefixes (output/foo.png is an input subdirectory).
|
||||
for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)):
|
||||
suffix = f" [{marker}]"
|
||||
if value.endswith(suffix):
|
||||
return str(root / value[:-len(suffix)])
|
||||
return str(self.input / value)
|
||||
|
||||
|
||||
class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.paths = FolderPathsFixture(Path(self.tempdir.name))
|
||||
self.folder_paths = types.ModuleType("folder_paths")
|
||||
self.folder_paths.get_input_directory = self.paths.get_input_directory
|
||||
self.folder_paths.get_output_directory = self.paths.get_output_directory
|
||||
self.folder_paths.get_temp_directory = self.paths.get_temp_directory
|
||||
self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath
|
||||
self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths})
|
||||
self.folder_patch.start()
|
||||
self.addCleanup(self.folder_patch.stop)
|
||||
self.addCleanup(self.tempdir.cleanup)
|
||||
|
||||
async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self):
|
||||
payload = png_bytes()
|
||||
(self.paths.input / "cat.png").write_bytes(payload)
|
||||
workflow = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}},
|
||||
}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 1)
|
||||
self.assertEqual(set(packed[0]), {"name", "image"})
|
||||
self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$")
|
||||
self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload)
|
||||
self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"])
|
||||
self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"])
|
||||
self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged")
|
||||
self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable")
|
||||
|
||||
async def test_resolves_annotated_input_output_and_temp_sources(self):
|
||||
sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp"))
|
||||
workflow = {}
|
||||
for index, (directory, filename) in enumerate(sources, 1):
|
||||
(directory / filename).write_bytes(png_bytes())
|
||||
root = directory.name
|
||||
workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 3)
|
||||
self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()})
|
||||
self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed})
|
||||
|
||||
async def test_directory_prefix_is_not_a_root_annotation(self):
|
||||
(self.paths.input / "output").mkdir()
|
||||
payload = png_bytes(color=(100, 0, 0))
|
||||
(self.paths.input / "output" / "same.png").write_bytes(payload)
|
||||
(self.paths.output / "same.png").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}}
|
||||
packed = await media.prepare_images(workflow)
|
||||
self.assertEqual(base64.b64decode(packed[0]["image"]), payload)
|
||||
|
||||
async def test_rejects_unknown_root_annotation(self):
|
||||
(self.paths.input / "cat.png").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_recognizes_custom_file_image_reference(self):
|
||||
(self.paths.input / "source.bmp").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 1)
|
||||
self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"])
|
||||
|
||||
async def test_rejects_missing_and_dynamic_loadimage_values(self):
|
||||
cases = (
|
||||
{"1": {"class_type": "LoadImage", "inputs": {}}},
|
||||
{"1": {"class_type": "LoadImage"}},
|
||||
{"1": {"class_type": "LoadImage", "inputs": None}},
|
||||
{"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}},
|
||||
)
|
||||
for workflow in cases:
|
||||
with self.subTest(workflow=workflow):
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_missing_source_and_unapproved_extension(self):
|
||||
for value in ("absent.png", "animation.gif", "clip.mp4"):
|
||||
with self.subTest(value=value):
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_unsupported_audio_or_video_references_without_reading(self):
|
||||
for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")):
|
||||
with self.subTest(key=key):
|
||||
workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_parent_traversal_and_absolute_outside_roots(self):
|
||||
outside = Path(self.tempdir.name).parent / "outside.png"
|
||||
outside.write_bytes(png_bytes())
|
||||
self.addCleanup(lambda: outside.unlink(missing_ok=True))
|
||||
for value in ("../outside.png", str(outside)):
|
||||
with self.subTest(value=value):
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_symlink_escape_from_allowed_root(self):
|
||||
outside = Path(self.tempdir.name).parent / "linked-outside.png"
|
||||
outside.write_bytes(png_bytes())
|
||||
self.addCleanup(lambda: outside.unlink(missing_ok=True))
|
||||
(self.paths.input / "link.png").symlink_to(outside)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}}
|
||||
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_individual_and_aggregate_input_size_overages(self):
|
||||
huge = self.paths.input / "large.jpg"
|
||||
with huge.open("wb") as handle:
|
||||
handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
huge.unlink()
|
||||
first = self.paths.input / "first.png"
|
||||
second = self.paths.input / "second.png"
|
||||
first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
|
||||
second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
|
||||
workflow = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}},
|
||||
}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
|
||||
async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self):
|
||||
raw_size = 7 * 1024 * 1024 + 1
|
||||
source = self.paths.input / "large-input.png"
|
||||
source.write_bytes(b"x" * raw_size)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
request_body = runpod.serialize_request(workflow, packed, 30)
|
||||
|
||||
self.assertGreater(raw_size, 7 * 1024 * 1024)
|
||||
self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size)
|
||||
self.assertLessEqual(len(request_body), 10_000_000)
|
||||
|
||||
|
||||
class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase):
|
||||
def manifest(self, prefix="render"):
|
||||
return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}]
|
||||
|
||||
def artifact(self, filename, payload=None, **extra):
|
||||
entry = {
|
||||
"filename": filename,
|
||||
"artifact_type": "output",
|
||||
"artifact_path": f"/safe/output/{filename}",
|
||||
"subfolder": "",
|
||||
"type": "base64",
|
||||
"data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"),
|
||||
}
|
||||
entry.update(extra)
|
||||
return entry
|
||||
|
||||
async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self):
|
||||
output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(list(result), ["job-1"])
|
||||
self.assertEqual(len(result["job-1"]), 2)
|
||||
self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"]))
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
|
||||
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255)
|
||||
|
||||
async def test_groups_images_under_manifest_job_ids(self):
|
||||
manifest = [
|
||||
self.manifest("first")[0],
|
||||
{"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"},
|
||||
]
|
||||
output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]}
|
||||
|
||||
result = await media.collect_images(output, manifest)
|
||||
|
||||
self.assertEqual(set(result), {"job-1", "job-2"})
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
self.assertEqual(len(result["job-2"]), 1)
|
||||
|
||||
async def test_preserves_rgba_alpha(self):
|
||||
output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 4))
|
||||
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255)
|
||||
|
||||
async def test_signed_url_is_downloaded_only_through_injected_downloader(self):
|
||||
output = {"images": [{
|
||||
"filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url",
|
||||
"url": "https://storage.example.test/signed",
|
||||
}]}
|
||||
fetch = AsyncMock(return_value=png_bytes())
|
||||
|
||||
result = await media.collect_images(output, self.manifest(), download=fetch)
|
||||
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
|
||||
fetch.assert_awaited_once_with("https://storage.example.test/signed")
|
||||
|
||||
async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self):
|
||||
output = {"images": [
|
||||
self.artifact("preview_00001_.png"),
|
||||
self.artifact("render_00001_.png"),
|
||||
{"filename": "mesh.obj", "type": "url", "url": "not a fetch target"},
|
||||
]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"):
|
||||
await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest())
|
||||
|
||||
async def test_accepts_compiler_length_flat_safe_prefix(self):
|
||||
prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64
|
||||
manifest = self.manifest(prefix)
|
||||
output = {"images": [self.artifact(f"{prefix}_00001_.png")]}
|
||||
|
||||
result = await media.collect_images(output, manifest)
|
||||
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
|
||||
async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self):
|
||||
duplicate = self.manifest() * 2
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": []}, duplicate)
|
||||
unsafe = self.manifest("../render")
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": []}, unsafe)
|
||||
|
||||
async def test_rejects_duplicate_canonical_artifact_identity(self):
|
||||
item = self.artifact("render_00001_.png")
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"):
|
||||
await media.collect_images({"images": [item, dict(item)]}, self.manifest())
|
||||
|
||||
async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self):
|
||||
wrong_type = self.artifact("render_00001_.png", artifact_type="temp")
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [wrong_type]}, self.manifest())
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest())
|
||||
|
||||
async def test_rejects_provider_errors_even_when_job_completed(self):
|
||||
for output in (
|
||||
{"error": "raw provider secret", "images": []},
|
||||
{"errors": ["raw provider secret"], "images": []},
|
||||
{"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}},
|
||||
):
|
||||
with self.subTest(output=output):
|
||||
with self.assertRaises(media.ServerlessMediaError) as ctx:
|
||||
await media.collect_images(output, self.manifest())
|
||||
self.assertNotIn("raw provider secret", str(ctx.exception))
|
||||
|
||||
async def test_rejects_invalid_base64_and_malformed_image(self):
|
||||
for item in (
|
||||
self.artifact("render_00001_.png", data="!not base64!"),
|
||||
self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()),
|
||||
):
|
||||
with self.subTest(item=item):
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [item]}, self.manifest())
|
||||
|
||||
async def test_validates_shape_consistency_within_collector(self):
|
||||
output = {"images": [
|
||||
self.artifact("render_00001_.png", png_bytes(size=(2, 1))),
|
||||
self.artifact("render_00002_.png", png_bytes(size=(3, 1))),
|
||||
]}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_bounds_output_file_count_and_bytes(self):
|
||||
too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "files"):
|
||||
await media.collect_images(too_many, self.manifest())
|
||||
|
||||
with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4):
|
||||
output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self):
|
||||
output = {"images": [
|
||||
self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()),
|
||||
self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()),
|
||||
]}
|
||||
with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3):
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self):
|
||||
for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"):
|
||||
with self.subTest(url=url):
|
||||
fetch = AsyncMock(return_value=png_bytes())
|
||||
output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images(output, self.manifest(), download=fetch)
|
||||
fetch.assert_not_awaited()
|
||||
|
||||
async def test_public_resolver_rejects_non_global_dns_answers(self):
|
||||
resolver = media._PublicResolver()
|
||||
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))])
|
||||
self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1)
|
||||
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))])
|
||||
with self.assertRaises(OSError):
|
||||
await resolver.resolve("cdn.example.test", 443)
|
||||
|
||||
async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self):
|
||||
class Content:
|
||||
def __init__(self, chunks):
|
||||
self.chunks = iter(chunks)
|
||||
async def iter_chunked(self, _size):
|
||||
for chunk in self.chunks:
|
||||
yield chunk
|
||||
|
||||
class Response:
|
||||
def __init__(self, status, headers=None, chunks=()):
|
||||
self.status, self.headers = status, headers or {}
|
||||
self.content = Content(chunks)
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
class Session:
|
||||
def __init__(self, responses):
|
||||
self.responses, self.calls = list(responses), []
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
def get(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.responses.pop(0)
|
||||
|
||||
session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])])
|
||||
with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session:
|
||||
payload = await media.download_public_artifact("https://storage.example.test/signed")
|
||||
|
||||
self.assertEqual(payload, b"abcdef")
|
||||
self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"])
|
||||
self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls))
|
||||
self.assertFalse(create_session.call_args.kwargs["trust_env"])
|
||||
self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver)
|
||||
self.assertNotIn("Authorization", session.calls[0][1].get("headers", {}))
|
||||
await create_session.call_args.kwargs["connector"].close()
|
||||
|
||||
blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})])
|
||||
with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session:
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.download_public_artifact("https://storage.example.test/signed")
|
||||
self.assertEqual(len(blocked.calls), 1)
|
||||
await blocked_session.call_args.kwargs["connector"].close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,392 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
MODULE_PATH = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "api"
|
||||
/ "orchestration"
|
||||
/ "serverless_prompt.py"
|
||||
)
|
||||
_spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH)
|
||||
assert _spec is not None and _spec.loader is not None
|
||||
_serverless_prompt = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(_serverless_prompt)
|
||||
compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt
|
||||
|
||||
|
||||
def _prepared_prompt():
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
|
||||
"3": {
|
||||
"class_type": "DistributedSeed",
|
||||
"inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"},
|
||||
},
|
||||
"4": {
|
||||
"class_type": "DistributedCollector",
|
||||
"inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True},
|
||||
},
|
||||
"5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}},
|
||||
"8": {"class_type": "UnrelatedNode", "inputs": {}},
|
||||
}
|
||||
|
||||
|
||||
def _run(prompt):
|
||||
return compile_serverless_prompt(prompt)
|
||||
|
||||
|
||||
def _load_prompt_transform():
|
||||
module_path = MODULE_PATH.with_name("prompt_transform.py")
|
||||
package_name = "serverless_prompt_transform_testpkg"
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
for suffix in ("api", "api.orchestration", "utils"):
|
||||
pkg = types.ModuleType(f"{package_name}.{suffix}")
|
||||
pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.{suffix}"] = pkg
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
|
||||
setattr(logging_module, "log", lambda *_args, **_kwargs: None)
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.prompt_transform", module_path
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_compiles_prompt_prepared_by_existing_worker_helpers():
|
||||
prompt_transform = _load_prompt_transform()
|
||||
original = {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
|
||||
"3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
"5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
|
||||
}
|
||||
pruned = prompt_transform.prune_prompt_for_worker(original)
|
||||
prompt_index = prompt_transform.PromptIndex(pruned)
|
||||
prepared = prompt_transform.apply_participant_overrides(
|
||||
pruned,
|
||||
participant_id="worker-a",
|
||||
enabled_worker_ids=["worker-a"],
|
||||
job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"),
|
||||
master_url="http://master.invalid",
|
||||
delegate_master=False,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
graph, manifest = _run(prepared)
|
||||
|
||||
assert graph["2"]["inputs"]["seed"] == 101
|
||||
assert set(graph) == {"1", "2", "7"}
|
||||
assert graph["7"]["class_type"] == "SaveImage"
|
||||
assert manifest[0]["job_id"] == "run_4"
|
||||
assert manifest[0]["collector_id"] == "4"
|
||||
|
||||
|
||||
def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream():
|
||||
original = _prepared_prompt()
|
||||
|
||||
graph, manifest = _run(original)
|
||||
|
||||
assert manifest == [
|
||||
{
|
||||
"collector_id": "4",
|
||||
"job_id": "job/42",
|
||||
"node_id": "9",
|
||||
"kind": "image",
|
||||
"prefix": graph["9"]["inputs"]["filename_prefix"],
|
||||
}
|
||||
]
|
||||
prefix = manifest[0]["prefix"]
|
||||
assert graph["9"] == {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"images": ["2", 0], "filename_prefix": prefix},
|
||||
"_meta": {"title": "Serverless output: collector 4 (image)"},
|
||||
}
|
||||
assert "/" not in prefix and "\\" not in prefix
|
||||
assert graph["2"]["inputs"]["seed"] == 41
|
||||
assert set(graph) == {"1", "2", "9"}
|
||||
assert all("Distributed" not in node["class_type"] for node in graph.values())
|
||||
assert original["3"]["inputs"]["seed"] == 40
|
||||
assert original["4"]["inputs"]["images"] == ["2", 0]
|
||||
|
||||
|
||||
def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts():
|
||||
prompt = _prepared_prompt()
|
||||
prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe")
|
||||
|
||||
first_graph, first_manifest = _run(prompt)
|
||||
second_graph, second_manifest = _run(prompt)
|
||||
|
||||
assert first_manifest == second_manifest
|
||||
assert first_graph == second_graph
|
||||
prefix = first_manifest[0]["prefix"]
|
||||
assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"]
|
||||
assert not any(char in prefix for char in ("/", "\\", ":"))
|
||||
assert first_manifest[0]["job_id"] == "../run/a"
|
||||
|
||||
|
||||
def test_compiles_multiple_disconnected_collectors_in_stable_order():
|
||||
prompt = {
|
||||
"10": {"class_type": "LoadImage", "inputs": {}},
|
||||
"20": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}},
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}},
|
||||
"30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
|
||||
graph, manifest = _run(prompt)
|
||||
|
||||
assert [item["collector_id"] for item in manifest] == ["1", "2"]
|
||||
assert [item["node_id"] for item in manifest] == ["31", "32"]
|
||||
assert graph["31"]["inputs"]["images"] == ["10", 0]
|
||||
assert graph["32"]["inputs"]["images"] == ["20", 0]
|
||||
assert set(graph) == {"10", "20", "31", "32"}
|
||||
|
||||
|
||||
def test_rejects_audio_input_instead_of_emitting_audio_saver():
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadAudio", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="audio.*unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_collector_without_image_input():
|
||||
with pytest.raises(ValueError, match="collector.*image"):
|
||||
_run({"1": {"class_type": "DistributedCollector", "inputs": {}}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_type",
|
||||
["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"],
|
||||
)
|
||||
def test_rejects_unsupported_distributed_and_batch_nodes(class_type):
|
||||
prompt = {
|
||||
"1": {"class_type": class_type, "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
"3": {"class_type": "LoadImage", "inputs": {}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_collector_to_collector_dependency_through_processing_node():
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
|
||||
"3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="collector-to-collector"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"link",
|
||||
[["missing", 0], ["1", -1], ["1", "0"], ["1", True]],
|
||||
)
|
||||
def test_rejects_malformed_api_links(link):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": link}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="link"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_unexpected_output_slot_from_removed_distributed_node():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="output slot"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_distributed_seed_matches_worker_offset_and_master_passthrough():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["3"]["inputs"] == {"seed": 12, "steps": 7}
|
||||
assert "1" not in graph and "2" not in graph
|
||||
|
||||
|
||||
def test_distributed_value_uses_one_based_worker_values_and_type_coercion():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}},
|
||||
"2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["3"]["inputs"] == {"steps": 8, "seed": 3}
|
||||
assert "1" not in graph and "2" not in graph
|
||||
|
||||
|
||||
def test_resolves_stock_primitive_links_for_distributed_seed_and_value():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}},
|
||||
"4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}},
|
||||
"5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}},
|
||||
"7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75}
|
||||
assert not {"1", "2", "3", "4", "5"}.intersection(graph)
|
||||
|
||||
|
||||
def test_rejects_dynamic_linked_seed_or_value_without_executing_it():
|
||||
prompt = {
|
||||
"1": {"class_type": "ArbitraryNode", "inputs": {}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_primitive_resolution_detects_cycles():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}},
|
||||
"2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}},
|
||||
"4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}},
|
||||
"5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="cycle"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_model_name_trivial_text_passthrough_is_inlined():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["2"]["inputs"]["model_name"] == "model.safetensors"
|
||||
assert "1" not in graph
|
||||
|
||||
|
||||
def test_rejects_dynamic_linked_model_name_text():
|
||||
prompt = {
|
||||
"1": {"class_type": "ArbitraryNode", "inputs": {}},
|
||||
"2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_resolves_boolean_and_float_primitive_links():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}},
|
||||
"2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}},
|
||||
"3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}},
|
||||
"4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}},
|
||||
"6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5}
|
||||
assert not {"1", "2", "3", "4"}.intersection(graph)
|
||||
|
||||
|
||||
def test_compiled_graph_matches_golden_prompt_and_manifest():
|
||||
prompt = _prepared_prompt()
|
||||
prompt["3"]["inputs"]["seed"] = ["10", 0]
|
||||
prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}}
|
||||
prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}
|
||||
prompt["12"] = {
|
||||
"class_type": "DistributedValue",
|
||||
"inputs": {
|
||||
"default_value": ["11", 0],
|
||||
"worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}),
|
||||
"is_worker": True,
|
||||
"worker_id": "worker_0",
|
||||
},
|
||||
}
|
||||
prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}}
|
||||
prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}}
|
||||
prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0])
|
||||
|
||||
graph, manifest = _run(prompt)
|
||||
|
||||
expected_prefix = (
|
||||
"serverless_job_42_worker_0_4_image_"
|
||||
"1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244"
|
||||
)
|
||||
assert list(graph) == ["1", "2", "15"]
|
||||
assert graph == {
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {},
|
||||
"_meta": {"title": "source"},
|
||||
},
|
||||
"2": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"seed": 41,
|
||||
"cfg": 2.5,
|
||||
"model_name": "model.safetensors",
|
||||
},
|
||||
},
|
||||
"15": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"images": ["2", 0], "filename_prefix": expected_prefix},
|
||||
"_meta": {"title": "Serverless output: collector 4 (image)"},
|
||||
},
|
||||
}
|
||||
assert manifest == [
|
||||
{
|
||||
"collector_id": "4",
|
||||
"job_id": "job/42",
|
||||
"node_id": "15",
|
||||
"kind": "image",
|
||||
"prefix": expected_prefix,
|
||||
}
|
||||
]
|
||||
@@ -1,608 +0,0 @@
|
||||
import copy
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = copy.deepcopy(payload)
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None, error=None):
|
||||
self.payload = payload
|
||||
self.error = error
|
||||
|
||||
async def json(self):
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.payload
|
||||
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
return lambda handler: handler
|
||||
|
||||
def post(self, _path):
|
||||
return lambda handler: handler
|
||||
|
||||
|
||||
def _load_modules():
|
||||
repo = Path(__file__).resolve().parents[1]
|
||||
package_name = "dist_serverless_routes_testpkg"
|
||||
for name in list(sys.modules):
|
||||
if name == package_name or name.startswith(f"{package_name}."):
|
||||
del sys.modules[name]
|
||||
|
||||
root = types.ModuleType(package_name)
|
||||
root.__path__ = [str(repo)]
|
||||
sys.modules[package_name] = root
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = [str(repo / "api")]
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = [str(repo / "utils")]
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")]
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[logging_module.__name__] = logging_module
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[constants_module.__name__] = constants_module
|
||||
|
||||
def load_file(module_name, path):
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py")
|
||||
serverless_config = load_file(
|
||||
f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py"
|
||||
)
|
||||
|
||||
runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime")
|
||||
runtime.secret_store = None
|
||||
runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
|
||||
runtime.get_secret_store = lambda: runtime.secret_store
|
||||
runtime.get_manager = lambda: runtime.manager
|
||||
sys.modules[runtime.__name__] = runtime
|
||||
|
||||
import aiohttp # Ensure the real dependency is loaded; never replace global aiohttp.
|
||||
previous_server = sys.modules.get("server")
|
||||
server = types.ModuleType("server")
|
||||
server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server
|
||||
|
||||
network = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status)
|
||||
|
||||
network.handle_api_error = handle_api_error
|
||||
network.normalize_host = lambda value: value
|
||||
sys.modules[network.__name__] = network
|
||||
|
||||
routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py")
|
||||
routes.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status)
|
||||
)
|
||||
|
||||
if previous_server is None:
|
||||
sys.modules.pop("server", None)
|
||||
else:
|
||||
sys.modules["server"] = previous_server
|
||||
return config_module, serverless_config, runtime, routes
|
||||
|
||||
|
||||
config_module, serverless_config, runtime_module, config_routes = _load_modules()
|
||||
|
||||
|
||||
class _RunningTask:
|
||||
def done(self):
|
||||
return False
|
||||
|
||||
|
||||
class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self.tempdir.name)
|
||||
config_module.CONFIG_FILE = str(self.root / "gpu_config.json")
|
||||
config_module.invalidate_config_cache()
|
||||
self.config = config_module.get_default_config()
|
||||
self.config["workers"] = []
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json")
|
||||
runtime_module.secret_store = self.secrets
|
||||
runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
|
||||
|
||||
def tearDown(self):
|
||||
config_module.invalidate_config_cache()
|
||||
self.tempdir.cleanup()
|
||||
|
||||
def read_config_file(self):
|
||||
return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self):
|
||||
self.secrets.set("pool-1", "private-api-key")
|
||||
self.config["workers"] = [
|
||||
{"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1",
|
||||
"api_key": "inline-api-key", "name": "GPU"}
|
||||
]
|
||||
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
|
||||
config_module.invalidate_config_cache()
|
||||
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = response.payload["workers"][0]
|
||||
self.assertTrue(worker["api_key_configured"])
|
||||
self.assertNotIn("api_key", worker)
|
||||
self.assertNotIn("private-api-key", json.dumps(response.payload))
|
||||
self.assertNotIn("inline-api-key", json.dumps(response.payload))
|
||||
|
||||
async def test_get_config_redacts_normalized_secret_field_aliases(self):
|
||||
aliases = {
|
||||
"apiKey": "camel-api-key",
|
||||
"APIKEY": "upper-api-key",
|
||||
"clearApiKey": "camel-clear-key",
|
||||
"CLEAR_API_KEY": "upper-clear-key",
|
||||
"runpod_CREDENTIAL": "provider-credential",
|
||||
}
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
**aliases,
|
||||
"nested": {"credential_note": "nested-credential"},
|
||||
}]
|
||||
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
|
||||
config_module.invalidate_config_cache()
|
||||
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = response.payload["workers"][0]
|
||||
for alias, value in aliases.items():
|
||||
self.assertNotIn(alias, worker)
|
||||
self.assertNotIn(value, json.dumps(response.payload))
|
||||
self.assertNotIn("credential_note", worker["nested"])
|
||||
self.assertNotIn("nested-credential", json.dumps(response.payload))
|
||||
|
||||
async def test_bulk_config_rejects_normalized_secret_field_aliases(self):
|
||||
for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"):
|
||||
with self.subTest(secret_field=secret_field):
|
||||
secret = f"private-{secret_field}"
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "local", secret_field: secret}]
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self):
|
||||
secret = "runpod-secret-never-returned"
|
||||
payload = {
|
||||
"worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2,
|
||||
"timeout_seconds": 600, "enabled": True, "api_key": secret,
|
||||
}
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(payload))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload["api_key_configured"])
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(saved["endpoint_id"], "ep-123")
|
||||
self.assertEqual(saved["jobs_per_run"], 3)
|
||||
self.assertEqual(saved["concurrency"], 2)
|
||||
self.assertEqual(saved["timeout_seconds"], 600)
|
||||
for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertEqual(self.secrets.get("pool-1"), secret)
|
||||
|
||||
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3,
|
||||
"timeout_seconds": 900, "enabled": False,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest({"worker_id": "pool-1", "enabled": True})
|
||||
)
|
||||
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertTrue(worker["enabled"])
|
||||
self.assertEqual(worker["endpoint_id"], "ep-123")
|
||||
self.assertEqual(worker["jobs_per_run"], 4)
|
||||
self.assertEqual(worker["concurrency"], 3)
|
||||
self.assertEqual(worker["timeout_seconds"], 900)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
|
||||
self.assertTrue(response.payload["api_key_configured"])
|
||||
|
||||
async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self):
|
||||
secret = "must-not-leak"
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "concurrency": True, "api_key": secret,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_json_decode_error_is_generic_and_does_not_echo_body(self):
|
||||
secret = "body-secret-value"
|
||||
response = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest(error=ValueError(f"bad json near {secret}"))
|
||||
)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
|
||||
async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self):
|
||||
secret = "bulk-secret-never-write"
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "nested": {"api_key": secret},
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertFalse(self.secrets.path.exists())
|
||||
|
||||
async def test_bulk_update_validates_and_returns_redacted_serverless_config(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "name": "Pool", "enabled": True,
|
||||
"host": "obsolete-host", "port": 8188, "cuda_device": 0,
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"])
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for obsolete in ("host", "port", "cuda_device"):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertNotIn("api_key", response.payload["config"]["workers"][0])
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
|
||||
|
||||
async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self):
|
||||
self.secrets.set("pool-1", "bulk-clear-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "local", "port": 8188,
|
||||
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for field in ("provider", "endpoint_id", "api_key_configured"):
|
||||
self.assertNotIn(field, saved)
|
||||
|
||||
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
|
||||
existing = {"id": "local-1", "type": "local", "port": 8188}
|
||||
self.config["workers"] = [existing]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [existing])
|
||||
|
||||
async def test_delete_serverless_worker_clears_saved_key(self):
|
||||
self.secrets.set("pool-1", "delete-me")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123"
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_type_conversion_clears_secret_and_removes_serverless_fields(self):
|
||||
self.secrets.set("pool-1", "old-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188,
|
||||
"cuda_device": 0, "enabled": False,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["type"], "local")
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
|
||||
self.assertNotIn(serverless_field, worker)
|
||||
|
||||
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
"enabled": False,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
|
||||
enabled = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest({"worker_id": "pool-1", "enabled": True})
|
||||
)
|
||||
blocked = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "endpoint_id": "ep-456",
|
||||
}))
|
||||
|
||||
self.assertEqual(enabled.status, 200)
|
||||
self.assertEqual(blocked.status, 409)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
|
||||
|
||||
async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
keep = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "",
|
||||
}))
|
||||
clear = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "clear_api_key": True,
|
||||
}))
|
||||
|
||||
self.assertTrue(keep.payload["api_key_configured"])
|
||||
self.assertFalse(clear.payload["api_key_configured"])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self):
|
||||
with self.subTest(case="inline key"):
|
||||
config = config_module.get_default_config()
|
||||
config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}]
|
||||
self.assertFalse(config_module.save_config(config))
|
||||
with self.subTest(case="invalid serverless record"):
|
||||
config = config_module.get_default_config()
|
||||
config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
|
||||
self.assertFalse(config_module.save_config(config))
|
||||
|
||||
async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
|
||||
credential_change = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "new-key",
|
||||
}))
|
||||
type_change = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0,
|
||||
}))
|
||||
deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(credential_change.status, 409)
|
||||
self.assertEqual(type_change.status, 409)
|
||||
self.assertEqual(deletion.status, 409)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_secret_write_is_rolled_back_when_config_persistence_fails(self):
|
||||
secret = "write-then-rollback"
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"api_key": secret,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertEqual(config_module.load_config()["workers"], [])
|
||||
|
||||
async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "bulk-restore-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "local", "port": 8188}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "prior-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "replacement-key",
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "prior-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "delete-restore-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1")
|
||||
|
||||
async def test_regular_worker_delete_does_not_require_posix_secret_storage(self):
|
||||
self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")):
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_enabling_serverless_without_key_is_rejected(self):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True,
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_bulk_concurrency_change_is_blocked_during_active_run(self):
|
||||
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
updated = copy.deepcopy(self.config["workers"])
|
||||
updated[0]["concurrency"] = 2
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1)
|
||||
|
||||
async def test_non_serverless_update_behavior_is_preserved(self):
|
||||
self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "local-1", "name": "Renamed", "enabled": True,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["name"], "Renamed")
|
||||
self.assertTrue(worker["enabled"])
|
||||
|
||||
async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self):
|
||||
raw = {
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old",
|
||||
}]
|
||||
}
|
||||
normalized = config_module._merge_with_defaults(raw, config_module.get_default_config())
|
||||
worker = normalized["workers"][0]
|
||||
for obsolete in ("host", "port", "cuda_device", "extra_args"):
|
||||
self.assertNotIn(obsolete, worker)
|
||||
|
||||
|
||||
class PersistedPoolGuardTests(_ConfigRouteFixture):
|
||||
"""Use the actual manager/ledger across restart, not a mocked idle result."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
from test_serverless_runtime import runtime
|
||||
from unittest.mock import AsyncMock
|
||||
self.runtime = runtime
|
||||
self.session_factory = AsyncMock()
|
||||
self.ledger_path = self.root / "jobs.json"
|
||||
self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
|
||||
runtime_module.manager = self.manager
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
async def test_retained_jobs_block_destructive_changes_after_restart(self):
|
||||
original = self.read_config_file()
|
||||
for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"):
|
||||
for restart in (False, True):
|
||||
with self.subTest(status=status, restart=restart):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id")
|
||||
if restart:
|
||||
self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
|
||||
runtime_module.manager = self.manager
|
||||
else:
|
||||
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"concurrency": 2}, {"type": "local", "port": 8188}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
|
||||
self.assertEqual(response.status, 409, fields)
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 409)
|
||||
worker = original["workers"][0]
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, concurrency=2)]):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
|
||||
self.assertEqual(response.status, 409, workers)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN")
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused")
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
|
||||
async def test_terminal_or_other_pool_records_do_not_block_deletion(self):
|
||||
self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED")
|
||||
self.manager.ledger.update("other", worker_id="other", status="UNKNOWN")
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,143 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "dist_serverless_runtime_tests"
|
||||
for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = [str(path)]
|
||||
sys.modules[name] = module
|
||||
|
||||
def load(name):
|
||||
spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py"))
|
||||
assert spec and spec.loader
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
runtime = load("serverless_runtime")
|
||||
|
||||
class Secrets:
|
||||
def get(self, worker_id):
|
||||
return "private-test-key"
|
||||
|
||||
class RuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
|
||||
self.manager.poll_seconds = 0
|
||||
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
|
||||
self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
self.temp.cleanup()
|
||||
|
||||
async def test_job_id_persisted_before_status_and_results_injected(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock())
|
||||
async def status(job_id):
|
||||
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
|
||||
return {"status": "COMPLETED", "output": {"images": []}}
|
||||
api.status = status
|
||||
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(api.submit.await_count, 1)
|
||||
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
||||
self.assertEqual(item["worker_id"], "pool:1")
|
||||
self.assertEqual(item["tensor"], "tensor")
|
||||
self.assertTrue(item["is_last"])
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
|
||||
async def test_output_error_marks_failed_and_delivers_explicit_error(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock())
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
record = self.manager.ledger.all()[0]
|
||||
self.assertEqual(record["status"], "FAILED")
|
||||
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
||||
self.assertIn("serverless_error", item)
|
||||
self.assertNotIn("private-test-key", str(record) + str(item))
|
||||
|
||||
async def test_ambiguous_submit_never_retries(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock())
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(api.submit.await_count, 1)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
|
||||
async def test_cancel_targets_owned_job_only(self):
|
||||
polled = asyncio.Event()
|
||||
async def status(_job):
|
||||
polled.set()
|
||||
await asyncio.Event().wait()
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await polled.wait()
|
||||
await self.manager.cancel(worker_id="other-pool")
|
||||
api.cancel.assert_not_awaited()
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
api.cancel.assert_awaited_once_with("owned-job")
|
||||
|
||||
async def test_per_pool_concurrency_bounds_submission(self):
|
||||
gate = asyncio.Event()
|
||||
reached = asyncio.Event()
|
||||
count = 0
|
||||
async def submit(*args):
|
||||
nonlocal count
|
||||
count += 1
|
||||
reached.set()
|
||||
return f"provider-{count}"
|
||||
async def status(_job):
|
||||
await gate.wait()
|
||||
return {"status": "COMPLETED", "output": {"images": []}}
|
||||
api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
|
||||
second = dict(self.prepared, worker=dict(self.worker, id="pool:2"))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
||||
await self.manager.start("execution", "prompt", [self.prepared, second])
|
||||
await reached.wait()
|
||||
await asyncio.sleep(0)
|
||||
self.assertEqual(count, 1)
|
||||
gate.set()
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(count, 2)
|
||||
|
||||
async def test_cancel_persisted_job_after_master_restart(self):
|
||||
self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1)
|
||||
api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api) as client:
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
api.cancel.assert_awaited_once_with("old-job")
|
||||
self.assertEqual(client.call_args.args[1], "old-endpoint")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_master_failure_cancels_job_without_partial_success(self):
|
||||
self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}})
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
|
||||
|
||||
async def test_deadline_includes_local_slot_wait(self):
|
||||
self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
# A blocked slot expires without a billable submission.
|
||||
self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT")
|
||||
@@ -1,98 +0,0 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_serverless_status_tests"
|
||||
|
||||
|
||||
class StatusTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
def mod(name, **kwargs):
|
||||
result = types.ModuleType(name)
|
||||
result.__dict__.update(kwargs)
|
||||
return result
|
||||
self.manager = types.SimpleNamespace(health=AsyncMock(return_value={"ok": True, "queueRemaining": 0, "message": "Ready (scales from zero)"}), public_jobs=lambda _: [], cancel=AsyncMock(return_value=1))
|
||||
self.pool = {"id": "pool", "type": "serverless", "endpoint_id": "ep"}
|
||||
modules = {PKG: mod(PKG, __path__=[str(ROOT)]), PKG + ".api": mod("api", __path__=[str(ROOT/'api')]),
|
||||
PKG + ".utils": mod("utils", __path__=[]),
|
||||
PKG + ".utils.config": mod("config", load_config=lambda: {"workers": [self.pool]}),
|
||||
PKG + ".api.orchestration": mod("orchestration", __path__=[str(ROOT/'api/orchestration')]),
|
||||
PKG + ".api.orchestration.serverless_runtime": mod("runtime", get_manager=lambda: self.manager),
|
||||
"server": mod("server", PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda f: f, post=lambda _: lambda f: f))))}
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec = importlib.util.spec_from_file_location(PKG+'.api.serverless_routes', ROOT/'api/serverless_routes.py')
|
||||
self.routes = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(self.routes)
|
||||
|
||||
async def test_health_only_never_submits_and_supports_scale_to_zero(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.test_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.assertTrue(json.loads(result.text)["ok"])
|
||||
self.manager.health.assert_awaited_once_with(self.pool, fresh=True)
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_unknown_pool_cannot_cancel_other_jobs(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "other"}))
|
||||
result = await self.routes.cancel_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 404)
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_transport_error_does_not_leak_credentials(self):
|
||||
self.manager.health.side_effect = RuntimeError("Bearer secret-token https://signed.invalid/?token=value")
|
||||
request = types.SimpleNamespace(match_info={"worker_id": "pool"})
|
||||
result = await self.routes.serverless_status_endpoint(request)
|
||||
self.assertFalse(json.loads(result.text)["ok"])
|
||||
self.assertNotIn("secret-token", result.text)
|
||||
self.assertNotIn("signed.invalid", result.text)
|
||||
|
||||
async def test_cancel_uses_only_requested_pool(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.cancel_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.manager.cancel.assert_awaited_once_with(worker_id="pool")
|
||||
|
||||
async def test_reconcile_only_reads_the_requested_pools_jobs(self):
|
||||
self.manager.reconcile = AsyncMock(return_value=2)
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.reconcile_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.assertEqual(json.loads(result.text)["resolved"], 2)
|
||||
self.manager.reconcile.assert_awaited_once_with(worker_id="pool")
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_resolve_requires_explicit_boolean_confirmation_and_record_id(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock()
|
||||
for payload in ({"worker_id": "pool", "record_id": "record"},
|
||||
{"worker_id": "pool", "record_id": "record", "provider_checked": "true"},
|
||||
{"worker_id": "pool", "provider_checked": True}):
|
||||
with self.subTest(payload=payload):
|
||||
result = await self.routes.resolve_serverless_endpoint(types.SimpleNamespace(json=AsyncMock(return_value=payload)))
|
||||
self.assertEqual(result.status, 400)
|
||||
self.manager.acknowledge_recovery.assert_not_awaited()
|
||||
|
||||
async def test_resolve_is_scoped_and_does_not_claim_provider_cancellation(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock()
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
|
||||
result = await self.routes.resolve_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.manager.acknowledge_recovery.assert_awaited_once_with("pool", "record", provider_checked=True)
|
||||
self.assertIn("not provider-confirmed", json.loads(result.text)["message"])
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_live_or_mismatched_record_refuses_manual_resolution(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock(side_effect=ValueError("Unresolved inactive job not found for this pool"))
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
|
||||
result = await self.routes.resolve_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 409)
|
||||
|
||||
async def test_failed_health_still_exposes_recovery_ids(self):
|
||||
self.manager.health.side_effect = RuntimeError("private transport error")
|
||||
self.manager.public_jobs = lambda _: [{"record_id": "record", "status": "UNKNOWN"}]
|
||||
result = await self.routes.serverless_status_endpoint(types.SimpleNamespace(match_info={"worker_id": "pool"}))
|
||||
self.assertEqual(json.loads(result.text)["jobs"][0]["record_id"], "record")
|
||||
@@ -1,260 +0,0 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def _load_static_mode_module():
|
||||
module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
|
||||
package_name = "dist_static_mode_testpkg"
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
upscale_pkg = types.ModuleType(f"{package_name}.upscale")
|
||||
upscale_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale"] = upscale_pkg
|
||||
|
||||
modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
|
||||
modes_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.upscale.modes"] = modes_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
created_comfy_stub = False
|
||||
if "comfy" not in sys.modules:
|
||||
created_comfy_stub = True
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
model_mgmt = types.ModuleType("comfy.model_management")
|
||||
|
||||
class _InterruptProcessingException(Exception):
|
||||
pass
|
||||
|
||||
model_mgmt.processing_interrupted = lambda: False
|
||||
model_mgmt.throw_exception_if_processing_interrupted = lambda: None
|
||||
model_mgmt.InterruptProcessingException = _InterruptProcessingException
|
||||
|
||||
comfy_module.model_management = model_mgmt
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.model_management"] = model_mgmt
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
image_module = types.ModuleType(f"{package_name}.utils.image")
|
||||
from PIL import Image as PILImage
|
||||
import numpy as np
|
||||
|
||||
def _tensor_to_pil(img_tensor, batch_index=0):
|
||||
return PILImage.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
|
||||
|
||||
def _pil_to_tensor(image):
|
||||
arr = np.array(image).astype(np.float32) / 255.0
|
||||
return torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
image_module.tensor_to_pil = _tensor_to_pil
|
||||
image_module.pil_to_tensor = _pil_to_tensor
|
||||
sys.modules[f"{package_name}.utils.image"] = image_module
|
||||
|
||||
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
||||
|
||||
def _run_async_in_server_loop(coro, timeout=None):
|
||||
if timeout is not None:
|
||||
return asyncio.run(asyncio.wait_for(coro, timeout=timeout))
|
||||
return asyncio.run(coro)
|
||||
|
||||
async_helpers_module.run_async_in_server_loop = _run_async_in_server_loop
|
||||
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.get_worker_timeout_seconds = lambda: 60
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_INTERVAL = 10.0
|
||||
constants_module.JOB_POLL_INTERVAL = 0.0
|
||||
constants_module.JOB_POLL_MAX_ATTEMPTS = 3
|
||||
constants_module.MAX_BATCH = 20
|
||||
constants_module.TILE_SEND_TIMEOUT = 1.0
|
||||
constants_module.TILE_WAIT_TIMEOUT = 1.0
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
|
||||
job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
|
||||
|
||||
async def _noop(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
job_store_module.ensure_tile_jobs_initialized = lambda: types.SimpleNamespace(
|
||||
distributed_tile_jobs_lock=asyncio.Lock(),
|
||||
distributed_pending_tile_jobs={},
|
||||
)
|
||||
job_store_module.init_static_job_batched = _noop
|
||||
job_store_module._mark_task_completed = _noop
|
||||
job_store_module._cleanup_job = _noop
|
||||
job_store_module._drain_results_queue = _noop
|
||||
job_store_module._get_completed_count = _noop
|
||||
sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
|
||||
|
||||
job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
|
||||
|
||||
class _TileJobState:
|
||||
pass
|
||||
|
||||
job_models_module.TileJobState = _TileJobState
|
||||
sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", module_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_comfy_stub:
|
||||
sys.modules.pop("comfy.model_management", None)
|
||||
sys.modules.pop("comfy", None)
|
||||
|
||||
return module
|
||||
|
||||
|
||||
static_mode = _load_static_mode_module()
|
||||
|
||||
|
||||
class _FakeStaticWorker(static_mode.StaticModeMixin):
|
||||
def __init__(self):
|
||||
self.sent_batches = []
|
||||
self.request_calls = 0
|
||||
self.heartbeat_calls = 0
|
||||
self.job_ready = True
|
||||
self.tile_sequence = [(0, 0, True), (None, 0, True)]
|
||||
|
||||
def round_to_multiple(self, value):
|
||||
return value
|
||||
|
||||
def calculate_tiles(self, _width, _height, _tile_width, _tile_height, _force_uniform_tiles):
|
||||
return [(0, 0)]
|
||||
|
||||
def _poll_job_ready(self, *_args, **_kwargs):
|
||||
return self.job_ready
|
||||
|
||||
async def _request_tile_from_master(self, *_args, **_kwargs):
|
||||
self.request_calls += 1
|
||||
return self.tile_sequence.pop(0)
|
||||
|
||||
async def _send_heartbeat_to_master(self, *_args, **_kwargs):
|
||||
self.heartbeat_calls += 1
|
||||
|
||||
async def send_tiles_batch_to_master(
|
||||
self,
|
||||
processed_tiles,
|
||||
_multi_job_id,
|
||||
_master_url,
|
||||
_padding,
|
||||
_worker_id,
|
||||
is_final_flush=False,
|
||||
):
|
||||
self.sent_batches.append(
|
||||
{
|
||||
"tiles": list(processed_tiles),
|
||||
"is_final_flush": bool(is_final_flush),
|
||||
}
|
||||
)
|
||||
|
||||
def _extract_and_process_tile(self, upscaled_image, *_args, **_kwargs):
|
||||
batch_size = upscaled_image.shape[0]
|
||||
processed_batch = torch.zeros((batch_size, 2, 2, 3), dtype=torch.float32)
|
||||
return processed_batch, 0, 0, 2, 2
|
||||
|
||||
def create_tile_mask(self, *_args, **_kwargs):
|
||||
from PIL import Image
|
||||
return Image.new("L", (4, 4), 255)
|
||||
|
||||
def blend_tile(self, base_image, *_args, **_kwargs):
|
||||
return base_image
|
||||
|
||||
|
||||
def _call_worker_static(fake_worker):
|
||||
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
return fake_worker._process_worker_static_sync(
|
||||
image,
|
||||
model=None,
|
||||
positive=None,
|
||||
negative=None,
|
||||
vae=None,
|
||||
seed=1,
|
||||
steps=1,
|
||||
cfg=1.0,
|
||||
sampler_name="euler",
|
||||
scheduler="normal",
|
||||
denoise=0.5,
|
||||
tile_width=4,
|
||||
tile_height=4,
|
||||
padding=8,
|
||||
mask_blur=4,
|
||||
force_uniform_tiles=True,
|
||||
tiled_decode=False,
|
||||
multi_job_id="job-1",
|
||||
master_url="http://master:8188",
|
||||
worker_id="worker-1",
|
||||
enabled_workers=["worker-1"],
|
||||
)
|
||||
|
||||
|
||||
class StaticModeWorkerFlowTests(unittest.TestCase):
|
||||
def test_worker_static_aborts_when_job_not_ready(self):
|
||||
worker = _FakeStaticWorker()
|
||||
worker.job_ready = False
|
||||
|
||||
result = _call_worker_static(worker)
|
||||
|
||||
self.assertEqual(result[0].shape[0], 1)
|
||||
self.assertEqual(worker.request_calls, 0)
|
||||
self.assertEqual(worker.heartbeat_calls, 0)
|
||||
self.assertEqual(worker.sent_batches, [])
|
||||
|
||||
def test_worker_static_requests_tiles_and_flushes_final_batch(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
_call_worker_static(worker)
|
||||
|
||||
self.assertEqual(worker.request_calls, 2) # one tile, then sentinel
|
||||
self.assertEqual(worker.heartbeat_calls, 1)
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
tiles = worker.sent_batches[0]["tiles"]
|
||||
self.assertEqual(len(tiles), 1)
|
||||
self.assertEqual(tiles[0]["tile_idx"], 0)
|
||||
self.assertEqual(tiles[0]["global_idx"], 0)
|
||||
self.assertEqual(tiles[0]["batch_idx"], 0)
|
||||
|
||||
def test_flush_empty_final_still_sends_completion_signal(self):
|
||||
worker = _FakeStaticWorker()
|
||||
|
||||
returned = worker._flush_tiles_to_master(
|
||||
[],
|
||||
"job-1",
|
||||
"http://master:8188",
|
||||
8,
|
||||
"worker-1",
|
||||
is_final_flush=True,
|
||||
)
|
||||
|
||||
self.assertEqual(returned, [])
|
||||
self.assertEqual(len(worker.sent_batches), 1)
|
||||
self.assertEqual(worker.sent_batches[0]["tiles"], [])
|
||||
self.assertTrue(worker.sent_batches[0]["is_final_flush"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,21 +0,0 @@
|
||||
"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_native_v3_runtime():
|
||||
comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
|
||||
if not comfy_root:
|
||||
pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
|
||||
helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
|
||||
result = subprocess.run(
|
||||
[sys.executable, str(helper), str(Path(comfy_root).resolve())],
|
||||
capture_output=True, text=True, timeout=120,
|
||||
env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
|
||||
)
|
||||
assert result.returncode == 0, result.stdout + '\n' + result.stderr
|
||||
assert 'V3_ACCEPTANCE_OK' in result.stdout
|
||||
@@ -1,113 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import socket
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def load_ports():
|
||||
path = Path(__file__).resolve().parents[1] / "workers" / "ports.py"
|
||||
spec = importlib.util.spec_from_file_location("worker_ports_test", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class WorkerPortTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ports = load_ports()
|
||||
|
||||
def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self):
|
||||
workers = [
|
||||
{"id": "existing", "host": "localhost", "port": 8190, "enabled": False},
|
||||
{"id": "remote", "host": "remote.example", "port": 8192},
|
||||
]
|
||||
with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194])
|
||||
|
||||
def test_local_host_forms_reserve_ports(self):
|
||||
for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]:
|
||||
with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191])
|
||||
|
||||
def test_exhaustion_is_explicit_not_a_partial_allocation(self):
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
with self.assertRaisesRegex(ValueError, "available.*ports"):
|
||||
self.ports.allocate_worker_ports(65534, [], 2)
|
||||
|
||||
def test_launch_rejects_master_port_without_reassigning_configuration(self):
|
||||
worker = {"id": "manual", "port": 8189}
|
||||
with self.assertRaisesRegex(ValueError, "master.*8189"):
|
||||
self.ports.validate_worker_port(worker, 8189, [])
|
||||
self.assertEqual(worker["port"], 8189)
|
||||
|
||||
def test_launch_rejects_another_local_workers_reserved_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "other", "host": None, "port": 8190, "enabled": False}
|
||||
with self.assertRaisesRegex(ValueError, "other"):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_ignores_itself_and_remote_workers_with_same_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "remote", "host": "example.com", "port": 8190}
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_rejects_an_occupied_port(self):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen()
|
||||
port = listener.getsockname()[1]
|
||||
self.assertFalse(self.ports.is_port_available(port))
|
||||
with self.assertRaisesRegex(ValueError, "already in use"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "asyncio does not reuse addresses on Windows")
|
||||
def test_recently_closed_connection_does_not_block_worker_restart(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
port = listener.getsockname()[1]
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=2) as client:
|
||||
accepted, _ = listener.accept()
|
||||
with accepted:
|
||||
accepted.settimeout(2)
|
||||
accepted.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
client.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(accepted.recv(1), b"")
|
||||
self.assertTrue(self.ports.is_port_available(port))
|
||||
self.ports.validate_worker_port({"id": "restart", "port": port}, 1, [])
|
||||
with socket.socket() as restarted:
|
||||
restarted.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
restarted.bind(("127.0.0.1", port))
|
||||
restarted.listen(1)
|
||||
|
||||
def test_active_listener_with_reuseaddr_is_still_a_conflict(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
@unittest.skipUnless(socket.has_ipv6, "IPv6 unavailable")
|
||||
def test_ipv6_only_listener_is_still_a_conflict(self):
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as listener:
|
||||
listener.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
|
||||
try:
|
||||
listener.bind(("::1", 0))
|
||||
except OSError as exc:
|
||||
self.skipTest(f"IPv6 loopback unavailable: {exc}")
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
def test_invalid_ports_fail_clearly(self):
|
||||
for port in [0, 65536, "not-a-port", None, True, 8190.5]:
|
||||
with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,352 +0,0 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import URL, make_url
|
||||
|
||||
|
||||
def _load_process_module(module_filename: str):
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
|
||||
package_name = "dist_proc_testpkg"
|
||||
module_name = module_filename[:-3]
|
||||
|
||||
for mod_name in list(sys.modules):
|
||||
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = [str(module_path.parents[1])]
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
process_pkg = types.ModuleType(f"{package_name}.workers.process")
|
||||
process_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.workers.process"] = process_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
|
||||
process_module = types.ModuleType(f"{package_name}.utils.process")
|
||||
process_module.get_python_executable = lambda: "/usr/bin/test-python"
|
||||
process_module.is_process_alive = lambda _pid: False
|
||||
process_module.terminate_process = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.process"] = process_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}}
|
||||
config_module.save_config = lambda _config: None
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.PROCESS_TERMINATION_TIMEOUT = 1
|
||||
constants_module.PROCESS_WAIT_TIMEOUT = 1
|
||||
constants_module.WORKER_CHECK_INTERVAL = 0.01
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.process.{module_name}",
|
||||
module_path,
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
root_discovery_module = _load_process_module("root_discovery.py")
|
||||
launch_builder_module = _load_process_module("launch_builder.py")
|
||||
|
||||
|
||||
class ComfyRootDiscoveryTests(unittest.TestCase):
|
||||
def test_prefers_loaded_comfyui_module_path(self):
|
||||
discovery = root_discovery_module.ComfyRootDiscovery()
|
||||
server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/opt/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(sys.modules, {"server": server_module}, clear=False), \
|
||||
patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
|
||||
patch.dict(root_discovery_module.os.environ, {}, clear=True):
|
||||
self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
|
||||
|
||||
|
||||
class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
def build_command(self, root, worker=None, runtime=None):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
worker = worker or {"id": "worker-a", "port": 8190}
|
||||
runtime = runtime if runtime is not None else Namespace(database_url=None)
|
||||
with patch.object(builder, "_get_runtime_args", return_value=runtime):
|
||||
return builder.build_launch_command(worker, str(root))
|
||||
|
||||
def test_database_is_unique_and_stable_by_worker_id(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
first = self.build_command(root, {"id": "worker-a", "port": 8190})
|
||||
renamed = self.build_command(root, {
|
||||
"id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3,
|
||||
})
|
||||
second = self.build_command(root, {"id": "worker-b", "port": 8191})
|
||||
database = first[first.index("--database-url") + 1]
|
||||
self.assertEqual(database, renamed[renamed.index("--database-url") + 1])
|
||||
self.assertNotEqual(database, second[second.index("--database-url") + 1])
|
||||
self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/"))
|
||||
self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir())
|
||||
self.assertNotIn("--disable-assets", first)
|
||||
|
||||
def test_database_uses_effective_user_or_base_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra_args, expected in [
|
||||
(f'--user-directory "{root / "custom user"}"', root / "custom user"),
|
||||
(f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"),
|
||||
]:
|
||||
with self.subTest(extra_args=extra_args):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra_args,
|
||||
})
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/"))
|
||||
cmd = self.build_command(root, runtime=Namespace(
|
||||
database_url="sqlite:///master.db", user_directory=str(root / "runtime user"),
|
||||
))
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/"))
|
||||
self.assertNotEqual(database, "sqlite:///master.db")
|
||||
|
||||
def test_relative_database_directories_use_worker_cwd(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory) / "ComfyUI"
|
||||
root.mkdir()
|
||||
(root / "main.py").touch()
|
||||
launcher = Path(directory) / "launcher"
|
||||
launcher.mkdir()
|
||||
cases = [
|
||||
(Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"),
|
||||
(Namespace(database_url=None, base_directory="data"), "", root / "data/user"),
|
||||
(Namespace(database_url=None), "--user-directory=profiles", root / "profiles"),
|
||||
(Namespace(database_url=None), "--base-directory data", root / "data/user"),
|
||||
]
|
||||
for runtime, extra, expected in cases:
|
||||
with self.subTest(runtime=runtime, extra=extra), \
|
||||
patch.object(os, "getcwd", return_value=str(launcher)):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra,
|
||||
}, runtime)
|
||||
database = make_url(cmd[cmd.index("--database-url") + 1]).database
|
||||
self.assertEqual(Path(database).parent, expected / "distributed/workers")
|
||||
self.assertFalse((launcher / "profiles").exists())
|
||||
self.assertFalse((launcher / "data").exists())
|
||||
|
||||
def test_database_url_round_trips_special_characters_without_collisions(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for name in ["user?profile", "user%3Fprofile"]:
|
||||
with self.subTest(name=name):
|
||||
user = root / name
|
||||
runtime = Namespace(database_url=None, user_directory=str(user))
|
||||
# SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail
|
||||
# clearly on those versions rather than silently sharing a DB.
|
||||
sample = (user / "test.db").as_posix()
|
||||
serialized = URL.create("sqlite", database=sample).render_as_string()
|
||||
if make_url(serialized).database != sample:
|
||||
with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"):
|
||||
self.build_command(root, runtime=runtime)
|
||||
self.assertFalse(user.exists())
|
||||
continue
|
||||
databases = []
|
||||
for worker_id in ["worker-a", "worker-b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime)
|
||||
url = cmd[cmd.index("--database-url") + 1]
|
||||
database = make_url(url).database
|
||||
self.assertEqual(Path(database).parent, user / "distributed/workers")
|
||||
engine = create_engine(url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
actual = connection.exec_driver_sql("PRAGMA database_list").one()[2]
|
||||
self.assertEqual(Path(actual), Path(database))
|
||||
finally:
|
||||
engine.dispose()
|
||||
databases.append(database)
|
||||
self.assertNotEqual(*databases)
|
||||
|
||||
def test_explicit_database_and_disabled_assets_are_preserved(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]:
|
||||
with self.subTest(extra=extra):
|
||||
cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_older_comfyui_without_database_flag_keeps_existing_command(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
cmd = self.build_command(root, runtime=Namespace())
|
||||
self.assertNotIn("--database-url", cmd)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_worker_id_cannot_escape_database_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
databases = []
|
||||
for worker_id in ["../outside", "a/b", "a_b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190})
|
||||
path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///"))
|
||||
self.assertTrue(path.is_relative_to(root / "user"))
|
||||
databases.append(path)
|
||||
self.assertEqual(len(set(databases)), 3)
|
||||
|
||||
def test_extra_args_cannot_silently_override_the_configured_port(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--port 8189", "--port=8189"]:
|
||||
with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"):
|
||||
self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
|
||||
def test_inherits_runtime_layout_args_for_desktop(self):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
runtime_args = Namespace(
|
||||
listen="127.0.0.1",
|
||||
base_directory="C:/Users/test/ComfyUI",
|
||||
temp_directory=None,
|
||||
input_directory="C:/Users/test/ComfyUI/input",
|
||||
output_directory="C:/Users/test/ComfyUI/output",
|
||||
user_directory="C:/Users/test/ComfyUI/user",
|
||||
front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
|
||||
extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
|
||||
enable_manager=True,
|
||||
disable_manager_ui=False,
|
||||
enable_manager_legacy_ui=False,
|
||||
windows_standalone_build=True,
|
||||
log_stdout=True,
|
||||
verbose="INFO",
|
||||
enable_cors_header="*",
|
||||
)
|
||||
comfy_module = types.ModuleType("comfy")
|
||||
comfy_cli_args = types.ModuleType("comfy.cli_args")
|
||||
comfy_cli_args.args = runtime_args
|
||||
|
||||
worker_config = {
|
||||
"port": 9001,
|
||||
"extra_args": "--preview-method auto",
|
||||
}
|
||||
|
||||
def fake_exists(path):
|
||||
return path == "/desktop/ComfyUI/main.py"
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
|
||||
clear=False,
|
||||
), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
|
||||
cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
|
||||
|
||||
self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
|
||||
self.assertIn("--listen", cmd)
|
||||
self.assertIn("127.0.0.1", cmd)
|
||||
self.assertIn("--base-directory", cmd)
|
||||
self.assertIn("C:/Users/test/ComfyUI", cmd)
|
||||
self.assertIn("--input-directory", cmd)
|
||||
self.assertIn("--output-directory", cmd)
|
||||
self.assertIn("--user-directory", cmd)
|
||||
self.assertIn("--front-end-root", cmd)
|
||||
self.assertIn("--extra-model-paths-config", cmd)
|
||||
self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
|
||||
self.assertIn("--enable-manager", cmd)
|
||||
self.assertIn("--windows-standalone-build", cmd)
|
||||
self.assertIn("--log-stdout", cmd)
|
||||
self.assertIn("--disable-auto-launch", cmd)
|
||||
self.assertIn("--enable-cors-header", cmd)
|
||||
self.assertIn("*", cmd)
|
||||
self.assertIn("--port", cmd)
|
||||
self.assertIn("9001", cmd)
|
||||
self.assertNotIn("--auto-launch", cmd)
|
||||
|
||||
|
||||
class ProcessLaunchPortTests(unittest.TestCase):
|
||||
def test_launch_before_master_listens_uses_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
for missing in [AttributeError("PromptServer has no port"), None]:
|
||||
with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \
|
||||
patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
|
||||
def test_pre_listen_launch_still_rejects_the_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \
|
||||
patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
|
||||
def test_conflict_is_rejected_before_building_or_spawning(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
def test_available_port_reaches_existing_launch_path(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
self.assertEqual(manager.processes["manual"]["pid"], 1234)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,32 +0,0 @@
|
||||
import copy
|
||||
|
||||
|
||||
def clone_control_chain(control, clone_hint=True):
|
||||
"""Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
|
||||
if control is None:
|
||||
return None
|
||||
new_control = copy.copy(control)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
hint = getattr(control, 'cond_hint_original', None)
|
||||
new_control.cond_hint_original = hint.clone() if hint is not None else None
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
|
||||
|
||||
def clone_conditioning(cond_list, clone_hints=True):
|
||||
"""Clone conditioning without duplicating ControlNet models."""
|
||||
new_cond = []
|
||||
for emb, cond_dict in cond_list:
|
||||
new_emb = emb.clone() if emb is not None else None
|
||||
new_dict = cond_dict.copy()
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict and new_dict['mask'] is not None:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
if 'pooled_output' in new_dict and new_dict['pooled_output'] is not None:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:]
|
||||
new_cond.append([new_emb, new_dict])
|
||||
return new_cond
|
||||
@@ -1,49 +0,0 @@
|
||||
from dataclasses import dataclass, field
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
|
||||
class BaseJobState:
|
||||
"""Marker base class for typed USDU job state containers."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TileJobState(BaseJobState):
|
||||
"""Typed state container for static (tile) USDU jobs."""
|
||||
|
||||
multi_job_id: str
|
||||
mode: str = field(default="static", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_tasks: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_tasks: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageJobState(BaseJobState):
|
||||
"""Typed state container for dynamic (per-image) USDU jobs."""
|
||||
|
||||
multi_job_id: str
|
||||
mode: str = field(default="dynamic", init=False)
|
||||
queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
|
||||
completed_images: dict = field(default_factory=dict)
|
||||
worker_status: dict = field(default_factory=dict)
|
||||
assigned_to_workers: dict = field(default_factory=dict)
|
||||
batch_size: int = 0
|
||||
num_tiles_per_image: int = 0
|
||||
batched_static: bool = False
|
||||
created_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
@property
|
||||
def pending_tasks(self):
|
||||
return self.pending_images
|
||||
|
||||
@property
|
||||
def completed_tasks(self):
|
||||
return self.completed_images
|
||||
@@ -1,135 +0,0 @@
|
||||
import asyncio
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from .job_store import ensure_tile_jobs_initialized
|
||||
from .job_timeout import _check_and_requeue_timed_out_workers as _requeue_usdu
|
||||
from .job_models import ImageJobState, TileJobState
|
||||
|
||||
|
||||
class JobStateMixin:
|
||||
async def _get_job_data(self, multi_job_id):
|
||||
"""Return current job data reference while holding lock briefly."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
return prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
|
||||
async def _get_all_completed_tasks(self, multi_job_id):
|
||||
"""Helper to retrieve all completed tasks from the job data."""
|
||||
job_data = await self._get_job_data(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
return dict(job_data.completed_tasks)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return dict(job_data.completed_images)
|
||||
return {}
|
||||
|
||||
async def _get_next_image_index(self, multi_job_id):
|
||||
"""Get next image index from pending queue for master."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
pending_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
pending_queue = job_data.pending_images
|
||||
|
||||
if pending_queue is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(pending_queue.get(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _get_next_tile_index(self, multi_job_id):
|
||||
"""Get next tile index from pending queue for master in static mode."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
pending_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
|
||||
if pending_queue is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
return await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def _get_total_completed_count(self, multi_job_id):
|
||||
"""Get total count of all completed images (master + workers)."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return len(job_data.completed_images)
|
||||
if isinstance(job_data, TileJobState):
|
||||
return len(job_data.completed_tasks)
|
||||
return 0
|
||||
|
||||
async def _get_all_completed_images(self, multi_job_id):
|
||||
"""Get all completed images."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return job_data.completed_images.copy()
|
||||
return {}
|
||||
|
||||
async def _get_pending_count(self, multi_job_id):
|
||||
"""Get count of pending images in the queue."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
return job_data.pending_images.qsize()
|
||||
if isinstance(job_data, TileJobState):
|
||||
return job_data.pending_tasks.qsize()
|
||||
return 0
|
||||
|
||||
async def _drain_worker_results_queue(self, multi_job_id):
|
||||
"""Drain pending worker results from queue and update completed images."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
worker_queue = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
worker_queue = job_data.queue
|
||||
|
||||
if worker_queue is None:
|
||||
return 0
|
||||
|
||||
drained_results = []
|
||||
while True:
|
||||
try:
|
||||
drained_results.append(worker_queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
if not drained_results:
|
||||
return 0
|
||||
|
||||
collected = 0
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
return 0
|
||||
|
||||
for result in drained_results:
|
||||
worker_id = result.get("worker_id")
|
||||
if "image_idx" in result and "image" in result:
|
||||
image_idx = result["image_idx"]
|
||||
image_pil = result["image"]
|
||||
if image_idx not in job_data.completed_images:
|
||||
job_data.completed_images[image_idx] = image_pil
|
||||
collected += 1
|
||||
debug_log(f"Drained image {image_idx} from worker {worker_id}")
|
||||
|
||||
if collected > 0:
|
||||
debug_log(f"Drained {collected} worker images during retry")
|
||||
|
||||
return collected
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
|
||||
"""Check for timed out workers and requeue their assigned images."""
|
||||
return await _requeue_usdu(multi_job_id, batch_size)
|
||||
@@ -1,180 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
import server
|
||||
|
||||
from ..utils.logging import debug_log
|
||||
from .job_models import BaseJobState, ImageJobState, TileJobState
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
|
||||
def ensure_tile_jobs_initialized():
|
||||
"""Ensure tile job storage is initialized on the server instance."""
|
||||
prompt_server = server.PromptServer.instance
|
||||
if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
|
||||
debug_log("Initializing persistent tile job queue on server instance.")
|
||||
prompt_server.distributed_pending_tile_jobs = {}
|
||||
prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
|
||||
else:
|
||||
invalid_job_ids = [
|
||||
job_id
|
||||
for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
|
||||
if not isinstance(job_data, BaseJobState)
|
||||
]
|
||||
for job_id in invalid_job_ids:
|
||||
debug_log(f"Removing invalid job state for {job_id}")
|
||||
del prompt_server.distributed_pending_tile_jobs[job_id]
|
||||
return prompt_server
|
||||
|
||||
|
||||
async def _init_job_queue(
|
||||
multi_job_id,
|
||||
mode,
|
||||
batch_size=None,
|
||||
num_tiles_per_image=None,
|
||||
all_indices=None,
|
||||
enabled_workers=None,
|
||||
batched_static: bool = False,
|
||||
):
|
||||
"""Unified initialization for job queues in static and dynamic modes."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
debug_log(f"Queue already exists for {multi_job_id}")
|
||||
return
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data = ImageJobState(multi_job_id=multi_job_id)
|
||||
elif mode == 'static':
|
||||
job_data = TileJobState(multi_job_id=multi_job_id)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
job_data.worker_status = {w: time.time() for w in enabled_workers or []}
|
||||
job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}
|
||||
|
||||
if mode == 'dynamic':
|
||||
job_data.batch_size = int(batch_size or 0)
|
||||
pending_queue = job_data.pending_images
|
||||
for i in (all_indices or range(int(batch_size or 0))):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized image queue with {batch_size} pending items")
|
||||
elif mode == 'static':
|
||||
job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
|
||||
job_data.batch_size = int(batch_size or 0)
|
||||
job_data.batched_static = bool(batched_static)
|
||||
# For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
|
||||
pending_queue = job_data.pending_tasks
|
||||
if batched_static and num_tiles_per_image is not None:
|
||||
for i in range(num_tiles_per_image):
|
||||
await pending_queue.put(i)
|
||||
else:
|
||||
total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
|
||||
for i in range(total_tiles):
|
||||
await pending_queue.put(i)
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
|
||||
async def init_dynamic_job(
|
||||
multi_job_id: str,
|
||||
batch_size: int,
|
||||
enabled_workers: List[str],
|
||||
all_indices: Optional[List[int]] = None,
|
||||
):
|
||||
"""Initialize queue for dynamic mode (per-image), with collector fields."""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'dynamic',
|
||||
batch_size=batch_size,
|
||||
all_indices=all_indices or list(range(batch_size)),
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
|
||||
|
||||
|
||||
async def init_static_job_batched(
|
||||
multi_job_id: str,
|
||||
batch_size: int,
|
||||
num_tiles_per_image: int,
|
||||
enabled_workers: List[str],
|
||||
):
|
||||
"""Initialize queue for static mode (batched-per-tile)."""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
enabled_workers=enabled_workers,
|
||||
batched_static=True,
|
||||
)
|
||||
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
q = job_data.queue
|
||||
completed_tasks = job_data.completed_tasks
|
||||
|
||||
collected = 0
|
||||
while True:
|
||||
try:
|
||||
result = q.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
if is_last:
|
||||
if worker_id in job_data.worker_status:
|
||||
del job_data.worker_status[worker_id]
|
||||
|
||||
return collected
|
||||
|
||||
|
||||
async def _get_completed_count(multi_job_id):
|
||||
"""Get count of completed tasks."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, BaseJobState):
|
||||
return len(job_data.completed_tasks)
|
||||
return 0
|
||||
|
||||
|
||||
async def _mark_task_completed(multi_job_id, task_id, result):
|
||||
"""Mark a task as completed."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, BaseJobState):
|
||||
job_data.completed_tasks[task_id] = result
|
||||
|
||||
|
||||
async def _cleanup_job(multi_job_id):
|
||||
"""Cleanup the job data."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
debug_log(f"Cleaned up job {multi_job_id}")
|
||||
@@ -1,150 +0,0 @@
|
||||
import time
|
||||
|
||||
from ..utils.config import load_config
|
||||
from ..utils.constants import HEARTBEAT_TIMEOUT
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import build_worker_url, probe_worker
|
||||
from .job_models import BaseJobState
|
||||
from .job_store import ensure_tile_jobs_initialized
|
||||
|
||||
|
||||
def _find_worker_record(worker_id):
|
||||
"""Return worker config entry by id, or None when missing."""
|
||||
workers = load_config().get("workers", [])
|
||||
return next((w for w in workers if str(w.get("id")) == str(worker_id)), None)
|
||||
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
"""Check timed out workers and requeue their tasks. Returns requeued count."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
current_time = time.time()
|
||||
|
||||
# Allow override via config setting 'worker_timeout_seconds'
|
||||
cfg = load_config()
|
||||
hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT))
|
||||
|
||||
# Snapshot timed-out workers and job details under lock.
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
|
||||
completed_tasks_snapshot = set(job_data.completed_tasks.keys())
|
||||
batched_static_snapshot = bool(job_data.batched_static)
|
||||
num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1)
|
||||
batch_size_snapshot = int(job_data.batch_size or 1)
|
||||
|
||||
timed_out_workers = []
|
||||
for worker, last_heartbeat in list(job_data.worker_status.items()):
|
||||
age = current_time - float(last_heartbeat)
|
||||
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
|
||||
if age > hb_timeout:
|
||||
timed_out_workers.append(
|
||||
{
|
||||
"worker_id": worker,
|
||||
"last_heartbeat": float(last_heartbeat),
|
||||
"assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])),
|
||||
}
|
||||
)
|
||||
|
||||
if not timed_out_workers:
|
||||
return 0
|
||||
|
||||
# Probe outside lock to avoid lock contention on network latency.
|
||||
workers_to_requeue = []
|
||||
workers_graced = []
|
||||
for worker_info in timed_out_workers:
|
||||
worker = worker_info["worker_id"]
|
||||
assigned = worker_info["assigned_tasks"]
|
||||
age = current_time - worker_info["last_heartbeat"]
|
||||
|
||||
incomplete_assigned = 0
|
||||
try:
|
||||
if assigned:
|
||||
if batched_static_snapshot:
|
||||
for task_id in assigned:
|
||||
for b in range(batch_size_snapshot):
|
||||
gidx = b * num_tiles_per_image_snapshot + task_id
|
||||
if gidx not in completed_tasks_snapshot:
|
||||
incomplete_assigned += 1
|
||||
break
|
||||
else:
|
||||
for task_id in assigned:
|
||||
if task_id not in completed_tasks_snapshot:
|
||||
incomplete_assigned += 1
|
||||
debug_log(
|
||||
f"Assigned diagnostics: total_assigned={len(assigned)} "
|
||||
f"incomplete_assigned={incomplete_assigned}"
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
|
||||
|
||||
busy = False
|
||||
probe_queue = None
|
||||
try:
|
||||
worker_record = _find_worker_record(worker)
|
||||
if worker_record:
|
||||
worker_url = build_worker_url(worker_record)
|
||||
debug_log(f"Probing worker {worker} at {worker_url}/prompt")
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
if payload is not None:
|
||||
probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0))
|
||||
busy = probe_queue is not None and probe_queue > 0
|
||||
else:
|
||||
debug_log(f"Probe skipped; worker {worker} not found in config")
|
||||
except Exception as e:
|
||||
debug_log(f"Probe failed for worker {worker}: {e}")
|
||||
finally:
|
||||
debug_log(
|
||||
f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
|
||||
)
|
||||
|
||||
if busy:
|
||||
workers_graced.append(worker)
|
||||
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
|
||||
continue
|
||||
|
||||
log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
|
||||
workers_to_requeue.append((worker, assigned))
|
||||
|
||||
# Re-acquire lock and apply requeue/cleanup decisions.
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return 0
|
||||
|
||||
# Refresh heartbeat for workers that we proved are still busy.
|
||||
for worker in workers_graced:
|
||||
if worker in job_data.worker_status:
|
||||
job_data.worker_status[worker] = current_time
|
||||
|
||||
requeued_count = 0
|
||||
completed_tasks = job_data.completed_tasks
|
||||
batched_static = bool(job_data.batched_static)
|
||||
num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
|
||||
batch_size = int(job_data.batch_size or 1)
|
||||
for worker, assigned_snapshot in workers_to_requeue:
|
||||
# Use current assignments if present, falling back to the snapshot.
|
||||
assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot))
|
||||
for task_id in assigned_tasks:
|
||||
# If batched_static, task_id is a tile_idx; consider it complete only if
|
||||
# all corresponding global_idx entries are present in completed_tasks.
|
||||
if batched_static:
|
||||
all_done = True
|
||||
for b in range(batch_size):
|
||||
gidx = b * num_tiles_per_image + task_id
|
||||
if gidx not in completed_tasks:
|
||||
all_done = False
|
||||
break
|
||||
if not all_done:
|
||||
await job_data.pending_tasks.put(task_id)
|
||||
requeued_count += 1
|
||||
else:
|
||||
if task_id not in completed_tasks:
|
||||
await job_data.pending_tasks.put(task_id)
|
||||
requeued_count += 1
|
||||
job_data.worker_status.pop(worker, None)
|
||||
if worker in job_data.assigned_to_workers:
|
||||
job_data.assigned_to_workers[worker] = []
|
||||
|
||||
return requeued_count
|
||||
@@ -1,313 +0,0 @@
|
||||
import asyncio, torch
|
||||
from PIL import Image
|
||||
import comfy.model_management
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ...utils.async_helpers import run_async_in_server_loop
|
||||
from ...utils.config import get_worker_timeout_seconds
|
||||
from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
|
||||
from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
|
||||
|
||||
|
||||
class DynamicModeMixin:
|
||||
"""
|
||||
Dynamic (per-image queue) USDU mode behaviors for master and worker roles.
|
||||
|
||||
Expected co-mixins on `self`:
|
||||
- TileOpsMixin (`calculate_tiles`, `_slice_conditioning`, `_process_and_blend_tile`).
|
||||
- JobStateMixin (image queue/task completion helpers).
|
||||
- WorkerCommsMixin (`_request_image_from_master`, `_send_full_image_to_master`, `_send_heartbeat_to_master`).
|
||||
"""
|
||||
|
||||
def process_master_dynamic(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers):
|
||||
"""Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
|
||||
# Get batch size and dimensions
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
num_workers = len(enabled_workers)
|
||||
|
||||
log(f"USDU Dist: Image queue distribution | Batch {batch_size} | Workers {num_workers} | Canvas {width}x{height} | Tile {tile_width}x{tile_height}")
|
||||
|
||||
# No fixed share - all images are dynamic
|
||||
all_indices = list(range(batch_size))
|
||||
|
||||
debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
|
||||
|
||||
# Calculate tiles for processing
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
# Initialize job queue for communication
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices),
|
||||
timeout=2.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"UltimateSDUpscale Master - Queue initialization error: {e}")
|
||||
raise RuntimeError(f"Failed to initialize dynamic mode queue: {e}")
|
||||
|
||||
# Convert batch to PIL list
|
||||
result_images = [tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB').copy() for b in range(batch_size)]
|
||||
|
||||
# Process images dynamically with master participating
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
processed_count = 0
|
||||
consecutive_retries = 0
|
||||
max_consecutive_retries = 10
|
||||
|
||||
# Process loop - master pulls from queue and processes synchronously
|
||||
while processed_count < batch_size:
|
||||
# Try to get an image to process
|
||||
image_idx = run_async_in_server_loop(
|
||||
self._get_next_image_index(multi_job_id),
|
||||
timeout=5.0 # Short timeout to allow frequent checks
|
||||
)
|
||||
|
||||
if image_idx is not None:
|
||||
# Reset retry counter and process locally
|
||||
consecutive_retries = 0
|
||||
debug_log(f"Master processing image {image_idx} dynamically")
|
||||
processed_count += 1
|
||||
|
||||
# Process locally
|
||||
single_tensor = upscaled_image[image_idx:image_idx+1]
|
||||
local_image = result_images[image_idx]
|
||||
image_seed = seed
|
||||
|
||||
# Pre-slice conditioning once per image (not per tile)
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
|
||||
|
||||
for tile_idx, pos in enumerate(all_tiles):
|
||||
source_tensor = pil_to_tensor(local_image)
|
||||
if single_tensor.is_cuda:
|
||||
source_tensor = source_tensor.cuda()
|
||||
local_image = self._process_and_blend_tile(
|
||||
tile_idx, pos, source_tensor, local_image,
|
||||
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, width, height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx=image_idx
|
||||
)
|
||||
|
||||
# Yield after each tile to minimize worker downtime
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1)
|
||||
# Note: No per-tile drain here – that's what makes this "per-image"
|
||||
|
||||
result_images[image_idx] = local_image
|
||||
|
||||
# Mark as completed
|
||||
run_async_in_server_loop(
|
||||
self._mark_image_completed(multi_job_id, image_idx, local_image),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
|
||||
drained_count = run_async_in_server_loop(
|
||||
self._drain_worker_results_queue(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
if drained_count > 0:
|
||||
debug_log(f"Drained {drained_count} worker images after master's image {image_idx}")
|
||||
|
||||
# NEW: Log overall progress (includes master's image + any drained workers)
|
||||
completed_now = run_async_in_server_loop(
|
||||
self._get_total_completed_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
|
||||
|
||||
# Yield to allow workers to get new images after completing one
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1)
|
||||
else:
|
||||
# Queue empty: collect any queued worker results to update progress
|
||||
drained_count = run_async_in_server_loop(
|
||||
self._drain_worker_results_queue(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after drain
|
||||
|
||||
# Check for timed out workers and requeue their images
|
||||
requeued_count = run_async_in_server_loop(
|
||||
self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
|
||||
timeout=5.0
|
||||
)
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield after requeue
|
||||
|
||||
if requeued_count > 0:
|
||||
log(f"Requeued {requeued_count} images from timed out workers")
|
||||
consecutive_retries = 0 # Reset since we have work to do
|
||||
continue
|
||||
|
||||
# Now check total completed (includes newly collected)
|
||||
completed_now = run_async_in_server_loop(
|
||||
self._get_total_completed_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
|
||||
log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
|
||||
|
||||
if completed_now >= batch_size:
|
||||
break
|
||||
|
||||
run_async_in_server_loop(self._async_yield(), timeout=0.1) # Yield before pending check
|
||||
|
||||
# Check if there are pending images in the queue (could be requeued)
|
||||
pending_count = run_async_in_server_loop(
|
||||
self._get_pending_count(multi_job_id),
|
||||
timeout=1.0
|
||||
)
|
||||
|
||||
if pending_count > 0:
|
||||
consecutive_retries = 0 # Reset retries since there's work to do
|
||||
continue
|
||||
|
||||
consecutive_retries += 1
|
||||
if consecutive_retries >= max_consecutive_retries:
|
||||
log(f"Max retries ({max_consecutive_retries}) reached. Forcing collection of remaining results.")
|
||||
break # Force exit to collection phase
|
||||
|
||||
debug_log("Waiting for workers")
|
||||
# Use async sleep to allow event loop to process worker requests
|
||||
run_async_in_server_loop(asyncio.sleep(2), timeout=3.0)
|
||||
|
||||
debug_log(f"Master processed {processed_count} images locally")
|
||||
|
||||
# Get all completed images to check what needs to be collected
|
||||
all_completed = run_async_in_server_loop(
|
||||
self._get_all_completed_images(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Calculate how many we still need to collect
|
||||
remaining_to_collect = batch_size - len(all_completed)
|
||||
|
||||
if remaining_to_collect > 0:
|
||||
debug_log(f"Waiting for {remaining_to_collect} more images from workers")
|
||||
# Use the unified worker timeout for the collection phase
|
||||
collection_timeout = float(get_worker_timeout_seconds())
|
||||
collected_images = run_async_in_server_loop(
|
||||
self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
|
||||
timeout=collection_timeout
|
||||
)
|
||||
|
||||
# Merge collected with already completed
|
||||
all_completed.update(collected_images)
|
||||
|
||||
# Update result images with all completed images
|
||||
for idx, processed_img in all_completed.items():
|
||||
if idx < batch_size:
|
||||
result_images[idx] = processed_img
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensor = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if batch_size > 1 else pil_to_tensor(result_images[0])
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
|
||||
log(f"Completed processing all {batch_size} images")
|
||||
return (result_tensor,)
|
||||
|
||||
def process_worker_dynamic(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_worker_ids, dynamic_threshold):
|
||||
"""Worker processing in dynamic mode - processes whole images."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get dimensions and tile grid
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")
|
||||
|
||||
# Keep track of processed images for is_last detection
|
||||
processed_count = 0
|
||||
|
||||
# Poll for job readiness to avoid races during master init
|
||||
max_poll_attempts = 20 # ~20s at 1s sleep
|
||||
if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
|
||||
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
|
||||
return (upscaled_image,)
|
||||
|
||||
# Loop to request and process images
|
||||
while True:
|
||||
# Request an image to process
|
||||
image_idx, estimated_remaining = run_async_in_server_loop(
|
||||
self._request_image_from_master(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_WAIT_TIMEOUT
|
||||
)
|
||||
|
||||
if image_idx is None:
|
||||
debug_log(f"USDU Dist Worker - No more images to process")
|
||||
break
|
||||
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Assigned image {image_idx}")
|
||||
processed_count += 1
|
||||
|
||||
# Determine if this should be marked as last for this worker
|
||||
is_last_for_worker = (estimated_remaining == 0)
|
||||
|
||||
# Extract single image tensor
|
||||
single_tensor = upscaled_image[image_idx:image_idx+1]
|
||||
|
||||
# Convert to PIL for processing
|
||||
local_image = tensor_to_pil(single_tensor, 0).copy()
|
||||
|
||||
# Process all tiles for this image
|
||||
image_seed = seed
|
||||
|
||||
# Pre-slice conditioning once per image (not per tile)
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
|
||||
|
||||
for tile_idx, pos in enumerate(all_tiles):
|
||||
source_tensor = pil_to_tensor(local_image)
|
||||
if single_tensor.is_cuda:
|
||||
source_tensor = source_tensor.cuda()
|
||||
local_image = self._process_and_blend_tile(
|
||||
tile_idx, pos, source_tensor, local_image,
|
||||
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, width, height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx=image_idx
|
||||
)
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Send processed image back to master
|
||||
try:
|
||||
# Use the estimated remaining to determine if this is the last image
|
||||
is_last = is_last_for_worker
|
||||
run_async_in_server_loop(
|
||||
self._send_full_image_to_master(local_image, image_idx, multi_job_id,
|
||||
master_url, worker_id, is_last),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
# Send heartbeat after processing
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
if is_last:
|
||||
break
|
||||
except Exception as e:
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
|
||||
# Continue processing other images
|
||||
|
||||
# Send final is_last signal
|
||||
debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
self._send_worker_complete_signal(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
except Exception as e:
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending completion signal: {e}")
|
||||
|
||||
return (upscaled_image,)
|
||||
@@ -1,72 +0,0 @@
|
||||
import math, torch
|
||||
from PIL import Image
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
|
||||
|
||||
class SingleGpuModeMixin:
|
||||
def process_single_gpu(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode):
|
||||
"""Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get image dimensions and batch size
|
||||
batch_size, height, width, _ = upscaled_image.shape
|
||||
|
||||
# Calculate all tiles
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
|
||||
rows = math.ceil(height / tile_height)
|
||||
cols = math.ceil(width / tile_width)
|
||||
log(
|
||||
f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
|
||||
)
|
||||
|
||||
# Prepare result images list
|
||||
result_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
|
||||
result_images.append(image_pil.copy())
|
||||
|
||||
# Precompute tile masks once
|
||||
tile_masks = []
|
||||
for tx, ty in all_tiles:
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
# Process tiles batched across images
|
||||
for tile_idx, (tx, ty) in enumerate(all_tiles):
|
||||
# Progressive state parity: extract each tile from the current updated image batch.
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
|
||||
# Extract batched tile
|
||||
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
|
||||
source_batch, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
# Process batch
|
||||
region = (x1, y1, x1 + ew, y1 + eh)
|
||||
processed_batch = self.process_tiles_batch(tile_batch, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tiled_decode, region, (width, height))
|
||||
|
||||
# Blend results back into each image using cached mask
|
||||
tile_mask = tile_masks[tile_idx]
|
||||
for b in range(batch_size):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
# Resize back to extracted size
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
|
||||
|
||||
# Convert back to tensor
|
||||
result_tensors = [pil_to_tensor(img) for img in result_images]
|
||||
result_tensor = torch.cat(result_tensors, dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
return (result_tensor,)
|
||||
@@ -1,570 +0,0 @@
|
||||
import asyncio, time, torch
|
||||
from PIL import Image
|
||||
import comfy.model_management
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ...utils.async_helpers import run_async_in_server_loop
|
||||
from ...utils.config import get_worker_timeout_seconds
|
||||
from ...utils.constants import (
|
||||
HEARTBEAT_INTERVAL,
|
||||
JOB_POLL_INTERVAL,
|
||||
JOB_POLL_MAX_ATTEMPTS,
|
||||
MAX_BATCH,
|
||||
TILE_SEND_TIMEOUT,
|
||||
TILE_WAIT_TIMEOUT,
|
||||
)
|
||||
from ..job_store import (
|
||||
ensure_tile_jobs_initialized, init_static_job_batched,
|
||||
_mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count,
|
||||
)
|
||||
from ..job_models import TileJobState
|
||||
|
||||
|
||||
class StaticModeMixin:
|
||||
"""
|
||||
Static (tile-queue) USDU mode behaviors for master and worker roles.
|
||||
|
||||
Expected co-mixins on `self`:
|
||||
- TileOpsMixin (`calculate_tiles`, tile extract/blend helpers).
|
||||
- JobStateMixin (`_get_next_tile_index`, `_get_all_completed_tasks`, requeue checks).
|
||||
- WorkerCommsMixin (`send_tiles_batch_to_master`, `_request_tile_from_master`, `_send_heartbeat_to_master`).
|
||||
"""
|
||||
|
||||
def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS):
|
||||
"""Poll master for job readiness to avoid worker/master initialization race."""
|
||||
for attempt in range(max_attempts):
|
||||
ready = run_async_in_server_loop(
|
||||
self._check_job_status(multi_job_id, master_url),
|
||||
timeout=5.0
|
||||
)
|
||||
if ready:
|
||||
if worker_id:
|
||||
debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
|
||||
else:
|
||||
debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
|
||||
return True
|
||||
time.sleep(JOB_POLL_INTERVAL)
|
||||
return False
|
||||
|
||||
def _extract_and_process_tile(
|
||||
self,
|
||||
upscaled_image,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
):
|
||||
"""Extract one tile position for the whole batch and process it."""
|
||||
tx, ty = all_tiles[tile_id]
|
||||
tile_batch, x1, y1, ew, eh = self.extract_batch_tile_with_padding(
|
||||
upscaled_image, tx, ty, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
region = (x1, y1, x1 + ew, y1 + eh)
|
||||
processed_batch = self.process_tiles_batch(
|
||||
tile_batch, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode,
|
||||
region, (width, height)
|
||||
)
|
||||
return processed_batch, x1, y1, ew, eh
|
||||
|
||||
def _flush_tiles_to_master(
|
||||
self,
|
||||
processed_tiles,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=False,
|
||||
):
|
||||
"""Send accumulated tile payloads to master and return a fresh accumulator."""
|
||||
if not processed_tiles:
|
||||
if is_final_flush:
|
||||
run_async_in_server_loop(
|
||||
self.send_tiles_batch_to_master(
|
||||
[],
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=True,
|
||||
),
|
||||
timeout=TILE_SEND_TIMEOUT,
|
||||
)
|
||||
return processed_tiles
|
||||
run_async_in_server_loop(
|
||||
self.send_tiles_batch_to_master(
|
||||
processed_tiles,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
padding,
|
||||
worker_id,
|
||||
is_final_flush=is_final_flush,
|
||||
),
|
||||
timeout=TILE_SEND_TIMEOUT
|
||||
)
|
||||
return []
|
||||
|
||||
def _master_process_one_tile(
|
||||
self,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
):
|
||||
"""Process one tile_id across the batch and blend into result_images."""
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
|
||||
source_batch,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
tile_mask = tile_masks[tile_id]
|
||||
out_bs = processed_batch.shape[0] if hasattr(processed_batch, "shape") else batch_size
|
||||
processed_items = min(batch_size, out_bs)
|
||||
for b in range(processed_items):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
result_images[b] = self.blend_tile(result_images[b], tile_pil, x1, y1, (ew, eh), tile_mask, padding)
|
||||
global_idx = b * num_tiles_per_image + tile_id
|
||||
run_async_in_server_loop(
|
||||
_mark_task_completed(multi_job_id, global_idx, {'batch_idx': b, 'tile_idx': tile_id}),
|
||||
timeout=5.0
|
||||
)
|
||||
return processed_items
|
||||
|
||||
def _process_worker_static_sync(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, master_url,
|
||||
worker_id, enabled_workers):
|
||||
"""Worker static mode processing with optional dynamic queue pulling."""
|
||||
# Round tile dimensions
|
||||
tile_width = self.round_to_multiple(tile_width)
|
||||
tile_height = self.round_to_multiple(tile_height)
|
||||
|
||||
# Get dimensions and calculate tiles
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
||||
num_tiles_per_image = len(all_tiles)
|
||||
batch_size = upscaled_image.shape[0]
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
|
||||
processed_tiles = []
|
||||
working_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
|
||||
working_images.append(image_pil.copy())
|
||||
tile_masks = []
|
||||
for tx, ty in all_tiles:
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
# Dynamic queue mode (static processing): process batched-per-tile
|
||||
log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}")
|
||||
processed_count = 0
|
||||
|
||||
max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
|
||||
if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
|
||||
log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
|
||||
return (upscaled_image,)
|
||||
|
||||
# Main processing loop - pull tile ids from queue
|
||||
while True:
|
||||
# Request a tile to process
|
||||
tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
|
||||
self._request_tile_from_master(multi_job_id, master_url, worker_id),
|
||||
timeout=TILE_WAIT_TIMEOUT
|
||||
)
|
||||
|
||||
if tile_idx is None:
|
||||
debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
|
||||
break
|
||||
|
||||
# Always batched-per-tile in static mode
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Assigned tile_id {tile_idx}")
|
||||
processed_count += batch_size
|
||||
tile_id = tile_idx
|
||||
source_batch = torch.cat([pil_to_tensor(img) for img in working_images], dim=0)
|
||||
if upscaled_image.is_cuda:
|
||||
source_batch = source_batch.cuda()
|
||||
processed_batch, x1, y1, ew, eh = self._extract_and_process_tile(
|
||||
source_batch,
|
||||
tile_id,
|
||||
all_tiles,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
# Queue results
|
||||
for b in range(batch_size):
|
||||
tile_pil = tensor_to_pil(processed_batch, b)
|
||||
if tile_pil.size != (ew, eh):
|
||||
tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
|
||||
working_images[b] = self.blend_tile(
|
||||
working_images[b],
|
||||
tile_pil,
|
||||
x1,
|
||||
y1,
|
||||
(ew, eh),
|
||||
tile_masks[tile_id],
|
||||
padding,
|
||||
)
|
||||
processed_tiles.append({
|
||||
'tile': processed_batch[b:b+1],
|
||||
'tile_idx': tile_id,
|
||||
'x': x1,
|
||||
'y': y1,
|
||||
'extracted_width': ew,
|
||||
'extracted_height': eh,
|
||||
'padding': padding,
|
||||
'batch_idx': b,
|
||||
'global_idx': b * num_tiles_per_image + tile_id
|
||||
})
|
||||
|
||||
# Send heartbeat
|
||||
try:
|
||||
run_async_in_server_loop(
|
||||
self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
|
||||
timeout=5.0
|
||||
)
|
||||
except Exception as e:
|
||||
debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
|
||||
|
||||
# Send tiles in batches within loop
|
||||
if len(processed_tiles) >= MAX_BATCH:
|
||||
processed_tiles = self._flush_tiles_to_master(
|
||||
processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False
|
||||
)
|
||||
|
||||
# Send any remaining tiles
|
||||
processed_tiles = self._flush_tiles_to_master(
|
||||
processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=True
|
||||
)
|
||||
|
||||
debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
|
||||
return (upscaled_image,)
|
||||
|
||||
async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total):
|
||||
"""Async helper for collection and monitoring in static mode.
|
||||
Returns collected tasks dict. Caller should check if all tasks are complete."""
|
||||
last_progress_log = time.time()
|
||||
progress_interval = 5.0
|
||||
last_heartbeat_check = time.time()
|
||||
last_completed_count = 0
|
||||
|
||||
while True:
|
||||
# Check for user interruption
|
||||
if comfy.model_management.processing_interrupted():
|
||||
log("Processing interrupted by user")
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
|
||||
# Drain any pending results
|
||||
collected_count = await _drain_results_queue(multi_job_id)
|
||||
|
||||
# Check and requeue timed-out workers periodically
|
||||
current_time = time.time()
|
||||
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
|
||||
requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
|
||||
if requeued_count > 0:
|
||||
log(f"Requeued {requeued_count} tasks from timed-out workers")
|
||||
last_heartbeat_check = current_time
|
||||
|
||||
# Get current completion count
|
||||
completed_count = await _get_completed_count(multi_job_id)
|
||||
|
||||
# Progress logging
|
||||
if current_time - last_progress_log >= progress_interval:
|
||||
log(f"Progress: {completed_count}/{expected_total} tasks completed")
|
||||
last_progress_log = current_time
|
||||
|
||||
# Check if all tasks are completed
|
||||
if completed_count >= expected_total:
|
||||
debug_log(f"All {expected_total} tasks completed")
|
||||
break
|
||||
|
||||
# If no active workers remain and there are pending tasks, return for local processing
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, TileJobState):
|
||||
pending_queue = job_data.pending_tasks
|
||||
active_workers = list(job_data.worker_status.keys())
|
||||
if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
|
||||
log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
|
||||
break
|
||||
|
||||
# Wait a bit before next check
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Get all completed tasks for return
|
||||
return await self._get_all_completed_tasks(multi_job_id)
|
||||
|
||||
def _process_master_static_sync(self, upscaled_image, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name, scheduler, denoise,
|
||||
tile_width, tile_height, padding, mask_blur,
|
||||
force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
|
||||
all_tiles, num_tiles_per_image):
|
||||
"""Static mode master processing with optional dynamic queue pulling."""
|
||||
batch_size = upscaled_image.shape[0]
|
||||
_, height, width, _ = upscaled_image.shape
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
|
||||
# Convert batch to PIL list for processing
|
||||
result_images = []
|
||||
for b in range(batch_size):
|
||||
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0)
|
||||
result_images.append(image_pil.copy())
|
||||
|
||||
# Initialize queue: pending queue holds tile ids (batched per tile)
|
||||
log("USDU Dist: Using tile queue distribution")
|
||||
run_async_in_server_loop(
|
||||
init_static_job_batched(multi_job_id, batch_size, num_tiles_per_image, enabled_workers),
|
||||
timeout=10.0
|
||||
)
|
||||
debug_log(
|
||||
f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}"
|
||||
)
|
||||
|
||||
# Precompute masks for all tile positions to avoid repeated Gaussian blur work during blending
|
||||
tile_masks = []
|
||||
for idx, (tx, ty) in enumerate(all_tiles):
|
||||
tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
||||
|
||||
processed_count = 0
|
||||
consecutive_no_tile = 0
|
||||
max_consecutive_no_tile = 2
|
||||
|
||||
while processed_count < total_tiles:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
tile_idx = run_async_in_server_loop(
|
||||
self._get_next_tile_index(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
if tile_idx is not None:
|
||||
consecutive_no_tile = 0
|
||||
tile_id = tile_idx
|
||||
processed_count += self._master_process_one_tile(
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
|
||||
else:
|
||||
consecutive_no_tile += 1
|
||||
if consecutive_no_tile >= max_consecutive_no_tile:
|
||||
debug_log(f"Master processed {processed_count} tiles, moving to collection phase")
|
||||
break
|
||||
time.sleep(0.1)
|
||||
master_processed_count = processed_count
|
||||
|
||||
# Continue processing any remaining tiles while collecting worker results
|
||||
remaining_tiles = total_tiles - master_processed_count
|
||||
if remaining_tiles > 0:
|
||||
debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
|
||||
|
||||
# Collect worker results using async operations
|
||||
try:
|
||||
# Wait until either all tasks are collected or there are no active workers left
|
||||
collected_tasks = run_async_in_server_loop(
|
||||
self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
|
||||
timeout=None
|
||||
)
|
||||
except comfy.model_management.InterruptProcessingException:
|
||||
# Clean up job on interruption
|
||||
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
|
||||
raise
|
||||
|
||||
# Check if we need to process any remaining tasks locally after collection
|
||||
completed_count = len(collected_tasks)
|
||||
if completed_count < total_tiles:
|
||||
log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
|
||||
|
||||
# Process any remaining pending tasks (batched-per-tile)
|
||||
while True:
|
||||
# Check for user interruption
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
# Get next tile_id from pending queue
|
||||
tile_id = run_async_in_server_loop(
|
||||
self._get_next_tile_index(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
if tile_id is None:
|
||||
break
|
||||
|
||||
self._master_process_one_tile(
|
||||
tile_id,
|
||||
all_tiles,
|
||||
upscaled_image,
|
||||
result_images,
|
||||
tile_masks,
|
||||
multi_job_id,
|
||||
batch_size,
|
||||
num_tiles_per_image,
|
||||
tile_width,
|
||||
tile_height,
|
||||
padding,
|
||||
force_uniform_tiles,
|
||||
model,
|
||||
positive,
|
||||
negative,
|
||||
vae,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
denoise,
|
||||
tiled_decode,
|
||||
width,
|
||||
height,
|
||||
)
|
||||
else:
|
||||
# Master processed all tiles
|
||||
collected_tasks = run_async_in_server_loop(
|
||||
self._get_all_completed_tasks(multi_job_id),
|
||||
timeout=5.0
|
||||
)
|
||||
|
||||
# Blend worker tiles synchronously in deterministic tile order.
|
||||
def _sort_key(item):
|
||||
global_idx, tile_data = item
|
||||
batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
|
||||
tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
|
||||
return (tile_idx, batch_idx, global_idx)
|
||||
|
||||
for global_idx, tile_data in sorted(collected_tasks.items(), key=_sort_key):
|
||||
# Skip tiles that don't have tensor data (already processed)
|
||||
if 'tensor' not in tile_data and 'image' not in tile_data:
|
||||
continue
|
||||
|
||||
batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
|
||||
tile_idx = tile_data.get('tile_idx', global_idx % num_tiles_per_image)
|
||||
|
||||
if batch_idx >= batch_size:
|
||||
continue
|
||||
|
||||
# Blend tile synchronously
|
||||
x = tile_data.get('x', 0)
|
||||
y = tile_data.get('y', 0)
|
||||
# Prefer PIL image if present to avoid reconversion
|
||||
if 'image' in tile_data:
|
||||
tile_pil = tile_data['image']
|
||||
else:
|
||||
tile_tensor = tile_data['tensor']
|
||||
tile_pil = tensor_to_pil(tile_tensor, 0)
|
||||
orig_x, orig_y = all_tiles[tile_idx]
|
||||
tile_mask = tile_masks[tile_idx]
|
||||
extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
|
||||
extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
|
||||
result_images[batch_idx] = self.blend_tile(result_images[batch_idx], tile_pil,
|
||||
x, y, (extracted_width, extracted_height), tile_mask, padding)
|
||||
|
||||
try:
|
||||
# Convert back to tensor
|
||||
if batch_size == 1:
|
||||
result_tensor = pil_to_tensor(result_images[0])
|
||||
else:
|
||||
result_tensors = [pil_to_tensor(img) for img in result_images]
|
||||
result_tensor = torch.cat(result_tensors, dim=0)
|
||||
|
||||
if upscaled_image.is_cuda:
|
||||
result_tensor = result_tensor.cuda()
|
||||
|
||||
log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
|
||||
return (result_tensor,)
|
||||
finally:
|
||||
# Cleanup (async operation) - always execute
|
||||
run_async_in_server_loop(_cleanup_job(multi_job_id), timeout=5.0)
|
||||
@@ -1,64 +0,0 @@
|
||||
import io
|
||||
import json
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def _parse_tiles_from_form(data):
|
||||
"""Parse tiles submitted via multipart/form-data into a list of tile dicts."""
|
||||
try:
|
||||
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
|
||||
except Exception:
|
||||
padding = 0
|
||||
|
||||
meta_raw = data.get('tiles_metadata')
|
||||
if meta_raw is None:
|
||||
raise ValueError("Missing tiles_metadata")
|
||||
|
||||
try:
|
||||
metadata = json.loads(meta_raw)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
|
||||
|
||||
if not isinstance(metadata, list):
|
||||
raise ValueError("tiles_metadata must be a list")
|
||||
|
||||
tiles = []
|
||||
for i, meta in enumerate(metadata):
|
||||
file_field = data.get(f'tile_{i}')
|
||||
if file_field is None or not hasattr(file_field, 'file'):
|
||||
raise ValueError(f"Missing tile data for index {i}")
|
||||
|
||||
raw = file_field.file.read()
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid image data for tile {i}: {e}")
|
||||
|
||||
try:
|
||||
tile_info = {
|
||||
'image': img,
|
||||
'tile_idx': int(meta.get('tile_idx', i)),
|
||||
'x': int(meta.get('x', 0)),
|
||||
'y': int(meta.get('y', 0)),
|
||||
'extracted_width': int(meta.get('extracted_width', img.width)),
|
||||
'extracted_height': int(meta.get('extracted_height', img.height)),
|
||||
'padding': int(padding),
|
||||
}
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
|
||||
|
||||
if 'batch_idx' in meta:
|
||||
try:
|
||||
tile_info['batch_idx'] = int(meta['batch_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
if 'global_idx' in meta:
|
||||
try:
|
||||
tile_info['global_idx'] = int(meta['global_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tiles.append(tile_info)
|
||||
|
||||
return tiles
|
||||
@@ -1,198 +0,0 @@
|
||||
import asyncio, time
|
||||
import comfy.model_management
|
||||
import server
|
||||
from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.config import get_worker_timeout_seconds
|
||||
from .job_store import ensure_tile_jobs_initialized, _mark_task_completed
|
||||
from .job_timeout import _check_and_requeue_timed_out_workers
|
||||
from .job_models import BaseJobState, ImageJobState, TileJobState
|
||||
|
||||
|
||||
class ResultCollectorMixin:
|
||||
"""
|
||||
Mixin for master-side result collection in USDU distributed jobs.
|
||||
|
||||
Expected co-mixins/attributes:
|
||||
- JobStateMixin methods for queue/task access.
|
||||
- `self._check_and_requeue_timed_out_workers(...)` coroutine.
|
||||
- `self._async_yield(...)` optional helper from WorkerCommsMixin.
|
||||
"""
|
||||
|
||||
def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
|
||||
"""Log timeout elapsed seconds for each tracked worker and return worker ids."""
|
||||
if not isinstance(job_data, BaseJobState):
|
||||
return []
|
||||
|
||||
worker_status = dict(job_data.worker_status)
|
||||
for worker_id, last_seen in worker_status.items():
|
||||
elapsed = max(0.0, current_time - float(last_seen))
|
||||
log(
|
||||
"UltimateSDUpscale Master - Heartbeat timeout: "
|
||||
f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
return list(worker_status.keys())
|
||||
|
||||
async def _async_collect_results(self, multi_job_id, num_workers, mode='static',
|
||||
remaining_to_collect=None, batch_size=None):
|
||||
"""Unified async helper to collect results from workers (tiles or images)."""
|
||||
# Get the already initialized queue
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id not in prompt_server.distributed_pending_tile_jobs:
|
||||
raise RuntimeError(f"Job queue not initialized for {multi_job_id}")
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if mode == 'dynamic':
|
||||
if not isinstance(job_data, ImageJobState):
|
||||
raise RuntimeError(
|
||||
f"Mode mismatch: expected dynamic, got {getattr(job_data, 'mode', 'unknown')}"
|
||||
)
|
||||
q = job_data.queue
|
||||
completed_images = job_data.completed_images
|
||||
expected_count = remaining_to_collect or batch_size
|
||||
elif mode == 'static':
|
||||
if not isinstance(job_data, TileJobState):
|
||||
raise RuntimeError(
|
||||
f"Mode mismatch: expected static, got {getattr(job_data, 'mode', 'unknown')}"
|
||||
)
|
||||
q = job_data.queue
|
||||
expected_count = len(job_data.completed_tasks) + job_data.pending_tasks.qsize()
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported mode: {mode}")
|
||||
|
||||
item_type = "images" if mode == 'dynamic' else "tiles"
|
||||
debug_log(f"UltimateSDUpscale Master - Starting collection, expecting {expected_count} {item_type} from {num_workers} workers")
|
||||
|
||||
collected_results = {}
|
||||
workers_done = set()
|
||||
# Unify collector/upscaler wait behavior with the UI worker timeout
|
||||
timeout = float(get_worker_timeout_seconds())
|
||||
last_heartbeat_check = time.time()
|
||||
wait_started_at = time.time()
|
||||
collected_count = 0
|
||||
|
||||
while len(workers_done) < num_workers:
|
||||
# Check for user interruption
|
||||
if comfy.model_management.processing_interrupted():
|
||||
log("Processing interrupted by user")
|
||||
raise comfy.model_management.InterruptProcessingException()
|
||||
|
||||
# For dynamic mode with remaining_to_collect, check if we've collected enough
|
||||
if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
|
||||
break
|
||||
|
||||
job_data_snapshot = None
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(current_job_data, BaseJobState):
|
||||
job_data_snapshot = current_job_data
|
||||
|
||||
try:
|
||||
# Shorter poll for dynamic mode, but never exceed the configured timeout
|
||||
wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout)
|
||||
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if mode == 'static':
|
||||
# Handle tiles
|
||||
tiles = result.get('tiles', [])
|
||||
debug_log(
|
||||
f"UltimateSDUpscale Master - Received batch of {len(tiles)} tiles from worker "
|
||||
f"'{worker_id}' (is_last={is_last})"
|
||||
)
|
||||
|
||||
for tile_data in tiles:
|
||||
if 'batch_idx' not in tile_data:
|
||||
log("UltimateSDUpscale Master - Missing batch_idx in tile data, skipping")
|
||||
continue
|
||||
|
||||
tile_idx = tile_data['tile_idx']
|
||||
key = tile_data.get('global_idx', tile_idx)
|
||||
entry = {
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_data['x'],
|
||||
'y': tile_data['y'],
|
||||
'extracted_width': tile_data['extracted_width'],
|
||||
'extracted_height': tile_data['extracted_height'],
|
||||
'padding': tile_data['padding'],
|
||||
'worker_id': worker_id,
|
||||
'batch_idx': tile_data.get('batch_idx', 0),
|
||||
'global_idx': tile_data.get('global_idx', tile_idx),
|
||||
}
|
||||
if 'image' in tile_data:
|
||||
entry['image'] = tile_data['image']
|
||||
elif 'tensor' in tile_data:
|
||||
entry['tensor'] = tile_data['tensor']
|
||||
collected_results[key] = entry
|
||||
|
||||
elif mode == 'dynamic':
|
||||
# Handle full images
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
image_idx = result['image_idx']
|
||||
image_pil = result['image']
|
||||
completed_images[image_idx] = image_pil
|
||||
collected_results[image_idx] = image_pil
|
||||
collected_count += 1
|
||||
debug_log(f"UltimateSDUpscale Master - Received image {image_idx} from worker {worker_id}")
|
||||
|
||||
if is_last:
|
||||
workers_done.add(worker_id)
|
||||
debug_log(f"UltimateSDUpscale Master - Worker {worker_id} completed")
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
current_time = time.time()
|
||||
waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
|
||||
if mode == 'dynamic':
|
||||
# Check for worker timeouts periodically
|
||||
if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
|
||||
# Use the class method to check and requeue
|
||||
requeued = await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)
|
||||
if requeued > 0:
|
||||
log(f"UltimateSDUpscale Master - Requeued {requeued} images from timed out workers")
|
||||
last_heartbeat_check = current_time
|
||||
|
||||
# Check if we've been waiting too long overall
|
||||
if current_time - wait_started_at > timeout:
|
||||
elapsed = current_time - wait_started_at
|
||||
log(
|
||||
"UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
|
||||
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
break
|
||||
else:
|
||||
elapsed = current_time - wait_started_at
|
||||
log(
|
||||
f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; "
|
||||
f"workers={waiting_workers}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
break
|
||||
|
||||
debug_log(f"UltimateSDUpscale Master - Collection complete. Got {len(collected_results)} {item_type} from {len(workers_done)} workers")
|
||||
|
||||
# Clean up job queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
del prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
|
||||
return collected_results if mode == 'static' else completed_images
|
||||
|
||||
async def _async_collect_worker_tiles(self, multi_job_id, num_workers):
|
||||
"""Async helper to collect tiles from workers."""
|
||||
return await self._async_collect_results(multi_job_id, num_workers, mode='static')
|
||||
|
||||
async def _mark_image_completed(self, multi_job_id, image_idx, image_pil):
|
||||
"""Mark an image as completed in the job data."""
|
||||
# Mark the image as completed with the image data
|
||||
await _mark_task_completed(multi_job_id, image_idx, {'image': image_pil})
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
|
||||
if isinstance(job_data, ImageJobState):
|
||||
job_data.completed_images[image_idx] = image_pil
|
||||
|
||||
async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
|
||||
"""Collect remaining processed images from workers."""
|
||||
return await self._async_collect_results(multi_job_id, num_workers, mode='dynamic',
|
||||
remaining_to_collect=remaining_to_collect,
|
||||
batch_size=batch_size)
|
||||
@@ -1,446 +0,0 @@
|
||||
import math, torch
|
||||
from contextlib import nullcontext
|
||||
from PIL import Image, ImageFilter, ImageDraw
|
||||
from typing import List, Tuple
|
||||
import comfy.samplers, comfy.model_management
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.image import tensor_to_pil, pil_to_tensor
|
||||
from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop
|
||||
from ..utils.crop_model_patch import crop_model_cond
|
||||
from .conditioning import clone_conditioning
|
||||
|
||||
|
||||
class TileOpsMixin:
|
||||
def round_to_multiple(self, value: int, multiple: int = 8) -> int:
|
||||
"""Round value to nearest multiple."""
|
||||
return round(value / multiple) * multiple
|
||||
|
||||
def calculate_tiles(self, image_width: int, image_height: int,
|
||||
tile_width: int, tile_height: int, force_uniform_tiles: bool = True) -> List[Tuple[int, int]]:
|
||||
"""Calculate tile positions to match Ultimate SD Upscale.
|
||||
|
||||
Positions are a simple grid starting at (0,0) with steps of
|
||||
`tile_width` and `tile_height`, using ceil(rows/cols) to cover edges.
|
||||
Uniform vs non-uniform affects only crop/resize, not positions.
|
||||
"""
|
||||
rows = math.ceil(image_height / tile_height)
|
||||
cols = math.ceil(image_width / tile_width)
|
||||
tiles: List[Tuple[int, int]] = []
|
||||
for yi in range(rows):
|
||||
for xi in range(cols):
|
||||
tiles.append((xi * tile_width, yi * tile_height))
|
||||
return tiles
|
||||
|
||||
def extract_tile_with_padding(self, image: torch.Tensor, x: int, y: int,
|
||||
tile_width: int, tile_height: int, padding: int,
|
||||
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
|
||||
"""Extract a tile region and resize to match USDU cropping logic.
|
||||
|
||||
Mirrors ComfyUI_UltimateSDUpscale processing:
|
||||
- Build a mask with a white rectangle at the tile rect
|
||||
- Compute crop_region via get_crop_region(mask, padding)
|
||||
- If force_uniform_tiles: expand by crop/aspect ratio, then resize to
|
||||
fixed processing size of round_to_multiple(tile + padding)
|
||||
- Else: target is ceil(crop_size/8)*8 per dimension
|
||||
- Extract the crop and resize to target tile_size
|
||||
Returns the resized tensor and crop origin/size for blending.
|
||||
"""
|
||||
_, h, w, _ = image.shape
|
||||
|
||||
# Create mask and compute initial padded crop region
|
||||
mask = Image.new('L', (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
x1, y1, x2, y2 = get_crop_region(mask, padding)
|
||||
|
||||
# Determine crop + processing size
|
||||
if force_uniform_tiles:
|
||||
process_w = self.round_to_multiple(tile_width + padding, 8)
|
||||
process_h = self.round_to_multiple(tile_height + padding, 8)
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
|
||||
process_ratio = process_w / process_h if process_h != 0 else 1.0
|
||||
if crop_ratio > process_ratio:
|
||||
target_w = crop_w
|
||||
target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
|
||||
else:
|
||||
target_w = round(crop_h * process_ratio)
|
||||
target_h = crop_h
|
||||
(x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
target_w = process_w
|
||||
target_h = process_h
|
||||
else:
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
target_w = max(8, math.ceil(crop_w / 8) * 8)
|
||||
target_h = max(8, math.ceil(crop_h / 8) * 8)
|
||||
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
|
||||
# Actual extracted size before resizing (for blending)
|
||||
extracted_width = x2 - x1
|
||||
extracted_height = y2 - y1
|
||||
|
||||
# Extract tile and resize to processing size
|
||||
tile = image[:, y1:y2, x1:x2, :]
|
||||
tile_pil = tensor_to_pil(tile, 0)
|
||||
if tile_pil.size != (target_w, target_h):
|
||||
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
|
||||
|
||||
tile_tensor = pil_to_tensor(tile_pil)
|
||||
if image.is_cuda:
|
||||
tile_tensor = tile_tensor.cuda()
|
||||
|
||||
return tile_tensor, x1, y1, extracted_width, extracted_height
|
||||
|
||||
def extract_batch_tile_with_padding(self, images: torch.Tensor, x: int, y: int,
|
||||
tile_width: int, tile_height: int, padding: int,
|
||||
force_uniform_tiles: bool) -> Tuple[torch.Tensor, int, int, int, int]:
|
||||
"""Extract a tile region for the entire batch and resize to USDU logic.
|
||||
|
||||
- Computes a single crop region from a mask at (x,y,w,h) with padding
|
||||
- force_uniform_tiles controls target processing size logic
|
||||
- Returns a batched tensor [B,H',W',C] and crop origin/size for blending
|
||||
"""
|
||||
batch, h, w, _ = images.shape
|
||||
|
||||
# Create mask and compute initial padded crop region (same for all images)
|
||||
mask = Image.new('L', (w, h), 0)
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
x1, y1, x2, y2 = get_crop_region(mask, padding)
|
||||
|
||||
# Determine crop + processing size
|
||||
if force_uniform_tiles:
|
||||
process_w = self.round_to_multiple(tile_width + padding, 8)
|
||||
process_h = self.round_to_multiple(tile_height + padding, 8)
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
crop_ratio = crop_w / crop_h if crop_h != 0 else 1.0
|
||||
process_ratio = process_w / process_h if process_h != 0 else 1.0
|
||||
if crop_ratio > process_ratio:
|
||||
target_w = crop_w
|
||||
target_h = round(crop_w / process_ratio) if process_ratio != 0 else crop_h
|
||||
else:
|
||||
target_w = round(crop_h * process_ratio)
|
||||
target_h = crop_h
|
||||
(x1, y1, x2, y2), _ = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
target_w = process_w
|
||||
target_h = process_h
|
||||
else:
|
||||
crop_w = x2 - x1
|
||||
crop_h = y2 - y1
|
||||
target_w = max(8, math.ceil(crop_w / 8) * 8)
|
||||
target_h = max(8, math.ceil(crop_h / 8) * 8)
|
||||
(x1, y1, x2, y2), (target_w, target_h) = expand_crop((x1, y1, x2, y2), w, h, target_w, target_h)
|
||||
|
||||
extracted_width = x2 - x1
|
||||
extracted_height = y2 - y1
|
||||
|
||||
# Slice batch region
|
||||
tiles = images[:, y1:y2, x1:x2, :]
|
||||
|
||||
# Resize each tile to target size
|
||||
resized_tiles = []
|
||||
for i in range(batch):
|
||||
tile_pil = tensor_to_pil(tiles, i)
|
||||
if tile_pil.size != (target_w, target_h):
|
||||
tile_pil = tile_pil.resize((target_w, target_h), Image.LANCZOS)
|
||||
resized_tiles.append(pil_to_tensor(tile_pil))
|
||||
tile_batch = torch.cat(resized_tiles, dim=0)
|
||||
|
||||
if images.is_cuda:
|
||||
tile_batch = tile_batch.cuda()
|
||||
|
||||
return tile_batch, x1, y1, extracted_width, extracted_height
|
||||
|
||||
def process_tile(self, tile_tensor: torch.Tensor, model, positive, negative, vae,
|
||||
seed: int, steps: int, cfg: float, sampler_name: str,
|
||||
scheduler: str, denoise: float, tiled_decode: bool = False,
|
||||
batch_idx: int = 0, region: Tuple[int, int, int, int] = None,
|
||||
image_size: Tuple[int, int] = None) -> torch.Tensor:
|
||||
"""Process a single tile through SD sampling.
|
||||
Note: positive and negative should already be pre-sliced for the current batch_idx."""
|
||||
debug_log(f"[process_tile] Processing tile for batch_idx={batch_idx}, seed={seed}, region={region}")
|
||||
|
||||
|
||||
# Import here to avoid circular dependencies
|
||||
from nodes import common_ksampler, VAEEncode, VAEDecode
|
||||
|
||||
# Try to import tiled VAE nodes if available
|
||||
try:
|
||||
from nodes import VAEEncodeTiled, VAEDecodeTiled
|
||||
tiled_vae_available = True
|
||||
except ImportError:
|
||||
tiled_vae_available = False
|
||||
if tiled_decode:
|
||||
debug_log("Tiled VAE nodes not available, falling back to standard VAE")
|
||||
|
||||
# Convert to PIL and back to ensure clean tensor without gradient tracking
|
||||
tile_pil = tensor_to_pil(tile_tensor, 0)
|
||||
clean_tensor = pil_to_tensor(tile_pil)
|
||||
|
||||
# Ensure tensor is detached and doesn't require gradients
|
||||
clean_tensor = clean_tensor.detach()
|
||||
if hasattr(clean_tensor, 'requires_grad_'):
|
||||
clean_tensor.requires_grad_(False)
|
||||
|
||||
# Move to correct device
|
||||
if tile_tensor.is_cuda:
|
||||
clean_tensor = clean_tensor.cuda()
|
||||
clean_tensor = clean_tensor.detach() # Detach again after device transfer
|
||||
|
||||
# Clone conditioning per tile (shares models, clones hints for cropping)
|
||||
positive_tile = clone_conditioning(positive, clone_hints=True)
|
||||
negative_tile = clone_conditioning(negative, clone_hints=True)
|
||||
|
||||
# Crop conditioning to tile region if provided (assumes hints at image resolution)
|
||||
if region is not None and image_size is not None:
|
||||
init_size = image_size # (width, height) of full image
|
||||
canvas_size = image_size
|
||||
tile_size = (tile_tensor.shape[2], tile_tensor.shape[1]) # (width, height)
|
||||
w_pad = 0 # No extra pad needed; region already includes padding
|
||||
h_pad = 0
|
||||
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
else:
|
||||
# No region cropping needed, use cloned conditioning as-is
|
||||
positive_cropped = positive_tile
|
||||
negative_cropped = negative_tile
|
||||
|
||||
# Encode to latent (always non-tiled, matching original node)
|
||||
latent = VAEEncode().encode(vae, clean_tensor)[0]
|
||||
|
||||
# Sample with model patch cropping parity (ControlNet patch hints)
|
||||
if region is not None and image_size is not None:
|
||||
model_ctx = crop_model_cond(
|
||||
model,
|
||||
region,
|
||||
image_size,
|
||||
image_size,
|
||||
(clean_tensor.shape[2], clean_tensor.shape[1]),
|
||||
)
|
||||
else:
|
||||
model_ctx = nullcontext(model)
|
||||
with model_ctx as model_for_sampling:
|
||||
samples = common_ksampler(
|
||||
model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
|
||||
positive_cropped, negative_cropped, latent, denoise=denoise
|
||||
)[0]
|
||||
|
||||
# Decode back to image
|
||||
if tiled_decode and tiled_vae_available:
|
||||
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
|
||||
else:
|
||||
image = VAEDecode().decode(vae, samples)[0]
|
||||
|
||||
return image
|
||||
|
||||
def process_tiles_batch(self, tile_batch: torch.Tensor, model, positive, negative, vae,
|
||||
seed: int, steps: int, cfg: float, sampler_name: str,
|
||||
scheduler: str, denoise: float, tiled_decode: bool,
|
||||
region: Tuple[int, int, int, int], image_size: Tuple[int, int]) -> torch.Tensor:
|
||||
"""Process a batch of tiles together (USDU behavior).
|
||||
|
||||
tile_batch: [B, H, W, C]
|
||||
Returns image batch tensor [B, H, W, C]
|
||||
"""
|
||||
# Import locally to avoid circular deps
|
||||
from nodes import common_ksampler, VAEEncode, VAEDecode
|
||||
try:
|
||||
from nodes import VAEEncodeTiled, VAEDecodeTiled
|
||||
tiled_vae_available = True
|
||||
except ImportError:
|
||||
tiled_vae_available = False
|
||||
|
||||
# Detach and move device
|
||||
clean = tile_batch.detach()
|
||||
if hasattr(clean, 'requires_grad_'):
|
||||
clean.requires_grad_(False)
|
||||
if tile_batch.is_cuda:
|
||||
clean = clean.cuda().detach()
|
||||
|
||||
# Clone/crop conditioning once for the region
|
||||
positive_tile = clone_conditioning(positive, clone_hints=True)
|
||||
negative_tile = clone_conditioning(negative, clone_hints=True)
|
||||
|
||||
init_size = image_size
|
||||
canvas_size = image_size
|
||||
tile_size = (clean.shape[2], clean.shape[1]) # (W,H)
|
||||
w_pad = 0
|
||||
h_pad = 0
|
||||
positive_cropped = crop_cond(positive_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
negative_cropped = crop_cond(negative_tile, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
|
||||
# Encode -> Sample -> Decode
|
||||
latent = VAEEncode().encode(vae, clean)[0]
|
||||
with crop_model_cond(model, region, image_size, image_size, tile_size) as model_for_sampling:
|
||||
samples = common_ksampler(
|
||||
model_for_sampling, seed, steps, cfg, sampler_name, scheduler,
|
||||
positive_cropped, negative_cropped, latent, denoise=denoise
|
||||
)[0]
|
||||
if tiled_decode and tiled_vae_available:
|
||||
image = VAEDecodeTiled().decode(vae, samples, tile_size=512)[0]
|
||||
else:
|
||||
image = VAEDecode().decode(vae, samples)[0]
|
||||
|
||||
return image
|
||||
|
||||
def create_tile_mask(self, image_width: int, image_height: int,
|
||||
x: int, y: int, tile_width: int, tile_height: int,
|
||||
mask_blur: int) -> Image.Image:
|
||||
"""Create a mask for blending tiles - matches Ultimate SD Upscale approach.
|
||||
|
||||
Creates a black image with a white rectangle at the tile position,
|
||||
then applies blur to create soft edges.
|
||||
"""
|
||||
# Create a full-size mask matching the image dimensions
|
||||
mask = Image.new('L', (image_width, image_height), 0) # Black background
|
||||
|
||||
# Draw white rectangle at tile position
|
||||
draw = ImageDraw.Draw(mask)
|
||||
draw.rectangle([x, y, x + tile_width, y + tile_height], fill=255)
|
||||
|
||||
# Apply blur to soften edges
|
||||
if mask_blur > 0:
|
||||
mask = mask.filter(ImageFilter.GaussianBlur(mask_blur))
|
||||
|
||||
return mask
|
||||
|
||||
def blend_tile(self, base_image: Image.Image, tile_image: Image.Image,
|
||||
x: int, y: int, extracted_size: Tuple[int, int],
|
||||
mask: Image.Image, padding: int) -> Image.Image:
|
||||
"""Blend a processed tile back into the base image using Ultimate SD Upscale's exact approach.
|
||||
|
||||
This follows the exact method from ComfyUI_UltimateSDUpscale/modules/processing.py
|
||||
"""
|
||||
extracted_width, extracted_height = extracted_size
|
||||
|
||||
# Debug logging (uncomment if needed)
|
||||
# debug_log(f"[Blend] Placing tile at ({x}, {y}), size: {extracted_width}x{extracted_height}")
|
||||
|
||||
# Calculate the crop region that was used for extraction
|
||||
crop_region = (x, y, x + extracted_width, y + extracted_height)
|
||||
|
||||
# The mask is already full-size, no need to crop
|
||||
|
||||
# Resize the processed tile back to the extracted size
|
||||
if tile_image.size != (extracted_width, extracted_height):
|
||||
tile_resized = tile_image.resize((extracted_width, extracted_height), Image.LANCZOS)
|
||||
else:
|
||||
tile_resized = tile_image
|
||||
|
||||
# Follow Ultimate SD Upscale blending approach:
|
||||
# Put the tile into position
|
||||
image_tile_only = Image.new('RGBA', base_image.size)
|
||||
image_tile_only.paste(tile_resized, crop_region[:2])
|
||||
|
||||
# Add the mask as an alpha channel
|
||||
# Must make a copy due to the possibility of an edge becoming black
|
||||
temp = image_tile_only.copy()
|
||||
temp.putalpha(mask) # Use the full image mask
|
||||
image_tile_only.paste(temp, image_tile_only)
|
||||
|
||||
# Add back the tile to the initial image according to the mask in the alpha channel
|
||||
result = base_image.convert('RGBA')
|
||||
result.alpha_composite(image_tile_only)
|
||||
|
||||
# Convert back to RGB
|
||||
return result.convert('RGB')
|
||||
|
||||
def _slice_conditioning(self, positive, negative, batch_idx):
|
||||
"""Helper to slice conditioning for a specific batch index."""
|
||||
# Clone and slice conditioning properly, including ControlNet hints
|
||||
positive_sliced = clone_conditioning(positive)
|
||||
negative_sliced = clone_conditioning(negative)
|
||||
|
||||
for cond_list in [positive_sliced, negative_sliced]:
|
||||
for i in range(len(cond_list)):
|
||||
emb, cond_dict = cond_list[i]
|
||||
if emb.shape[0] > 1:
|
||||
cond_list[i][0] = emb[batch_idx:batch_idx+1]
|
||||
if 'control' in cond_dict:
|
||||
control = cond_dict['control']
|
||||
while control is not None:
|
||||
hint = control.cond_hint_original
|
||||
if hint.shape[0] > 1:
|
||||
control.cond_hint_original = hint[batch_idx:batch_idx+1]
|
||||
control = control.previous_controlnet
|
||||
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
|
||||
cond_dict['mask'] = cond_dict['mask'][batch_idx:batch_idx+1]
|
||||
|
||||
return positive_sliced, negative_sliced
|
||||
|
||||
def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
|
||||
model, positive, negative, vae, seed, steps, cfg,
|
||||
sampler_name, scheduler, denoise, tile_width, tile_height,
|
||||
padding, mask_blur, image_width, image_height, force_uniform_tiles,
|
||||
tiled_decode, batch_idx: int = 0):
|
||||
"""Process a single tile and blend it into the result image."""
|
||||
x, y = tile_pos
|
||||
|
||||
# Extract and process tile
|
||||
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
|
||||
upscaled_image, x, y, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
processed_tile = self.process_tile(tile_tensor, model, positive, negative, vae,
|
||||
seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, batch_idx=batch_idx,
|
||||
region=(x1, y1, x1 + ew, y1 + eh), image_size=(image_width, image_height))
|
||||
|
||||
# Convert and blend
|
||||
processed_pil = tensor_to_pil(processed_tile, 0)
|
||||
# Create mask for this specific tile (no cache here; only used in single-tile path)
|
||||
tile_mask = self.create_tile_mask(image_width, image_height, x, y, tile_width, tile_height, mask_blur)
|
||||
# Use extraction position and size for blending
|
||||
result_image = self.blend_tile(result_image, processed_pil,
|
||||
x1, y1, (ew, eh), tile_mask, padding)
|
||||
|
||||
return result_image
|
||||
|
||||
def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
|
||||
model, positive, negative, vae, seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, tile_width, tile_height, padding,
|
||||
width, height, force_uniform_tiles, sliced_conditioning_cache):
|
||||
"""Process a single tile."""
|
||||
# Calculate which image and tile this corresponds to
|
||||
batch_idx = global_idx // num_tiles_per_image
|
||||
tile_idx = global_idx % num_tiles_per_image
|
||||
|
||||
# Skip if batch_idx is out of range
|
||||
if batch_idx >= upscaled_image.shape[0]:
|
||||
debug_log(f"Warning: Calculated batch_idx {batch_idx} exceeds batch size {upscaled_image.shape[0]}")
|
||||
return None
|
||||
|
||||
# Get or create sliced conditioning for this batch index
|
||||
if batch_idx not in sliced_conditioning_cache:
|
||||
positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, batch_idx)
|
||||
sliced_conditioning_cache[batch_idx] = (positive_sliced, negative_sliced)
|
||||
else:
|
||||
positive_sliced, negative_sliced = sliced_conditioning_cache[batch_idx]
|
||||
|
||||
x, y = all_tiles[tile_idx]
|
||||
|
||||
# Extract tile from the specific image in the batch
|
||||
tile_tensor, x1, y1, ew, eh = self.extract_tile_with_padding(
|
||||
upscaled_image[batch_idx:batch_idx+1], x, y, tile_width, tile_height, padding, force_uniform_tiles
|
||||
)
|
||||
|
||||
# Process tile through SD with the exact seed (USDU parity)
|
||||
image_seed = seed
|
||||
processed_tile = self.process_tile(tile_tensor, model, positive_sliced, negative_sliced, vae,
|
||||
image_seed, steps, cfg, sampler_name,
|
||||
scheduler, denoise, tiled_decode, batch_idx=batch_idx,
|
||||
region=(x1, y1, x1 + ew, y1 + eh), image_size=(width, height))
|
||||
|
||||
return {
|
||||
'tile': processed_tile,
|
||||
'global_idx': global_idx,
|
||||
'batch_idx': batch_idx,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x1,
|
||||
'y': y1,
|
||||
'extracted_width': ew,
|
||||
'extracted_height': eh
|
||||
}
|
||||
@@ -1,262 +0,0 @@
|
||||
import asyncio, io, json, time
|
||||
import aiohttp
|
||||
from PIL import Image
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import get_client_session
|
||||
from ..utils.constants import TILE_SEND_TIMEOUT
|
||||
from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, _send_heartbeat_to_master
|
||||
from ..utils.image import tensor_to_pil
|
||||
|
||||
|
||||
class WorkerCommsMixin:
|
||||
async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Proxy heartbeat helper used by worker processing mixins."""
|
||||
await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
|
||||
|
||||
async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url,
|
||||
padding, worker_id, is_final_flush=False):
|
||||
"""Send all processed tiles to master, chunked if large."""
|
||||
if not processed_tiles:
|
||||
if is_final_flush:
|
||||
await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id)
|
||||
return # Early exit if empty
|
||||
|
||||
total_tiles = len(processed_tiles)
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Preparing to send {total_tiles} tiles (size-aware chunks)")
|
||||
|
||||
# Prepare encoded images and sizes to enable size-aware chunking
|
||||
encoded = []
|
||||
for idx, tile_data in enumerate(processed_tiles):
|
||||
img = tensor_to_pil(tile_data['tile'], 0)
|
||||
bio = io.BytesIO()
|
||||
# Keep compression low to balance speed and size; adjust if needed
|
||||
img.save(bio, format='PNG', compress_level=0)
|
||||
raw = bio.getvalue()
|
||||
encoded.append({
|
||||
'bytes': raw,
|
||||
'meta': {
|
||||
'tile_idx': tile_data['tile_idx'],
|
||||
'x': tile_data['x'],
|
||||
'y': tile_data['y'],
|
||||
'extracted_width': tile_data['extracted_width'],
|
||||
'extracted_height': tile_data['extracted_height'],
|
||||
**({'batch_idx': tile_data['batch_idx']} if 'batch_idx' in tile_data else {}),
|
||||
**({'global_idx': tile_data['global_idx']} if 'global_idx' in tile_data else {}),
|
||||
}
|
||||
})
|
||||
|
||||
# Size-aware chunking
|
||||
max_bytes = int(MAX_PAYLOAD_SIZE) - (1024 * 1024) # 1MB headroom
|
||||
i = 0
|
||||
chunk_index = 0
|
||||
while i < total_tiles:
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('padding', str(padding))
|
||||
|
||||
metadata = []
|
||||
used = 0
|
||||
j = i
|
||||
while j < total_tiles:
|
||||
img_bytes = encoded[j]['bytes']
|
||||
meta = encoded[j]['meta']
|
||||
# Rough overhead for fields + JSON
|
||||
overhead = 1024
|
||||
if used + len(img_bytes) + overhead > max_bytes and j > i:
|
||||
break
|
||||
# Accept this tile in this chunk
|
||||
metadata.append(meta)
|
||||
data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
|
||||
used += len(img_bytes) + overhead
|
||||
j += 1
|
||||
|
||||
# Ensure at least one tile per chunk
|
||||
if j == i:
|
||||
# Single oversized tile, send anyway
|
||||
meta = encoded[j]['meta']
|
||||
metadata.append(meta)
|
||||
data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
|
||||
j += 1
|
||||
|
||||
chunk_size = j - i
|
||||
is_chunk_last = (j >= total_tiles)
|
||||
data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
|
||||
data.add_field('batch_size', str(chunk_size))
|
||||
data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
|
||||
|
||||
# Retry logic with exponential backoff
|
||||
max_retries = 5
|
||||
retry_delay = 0.5
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_tiles"
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay = min(retry_delay * 2, 5.0)
|
||||
else:
|
||||
log(f"UltimateSDUpscale Worker - Failed to send chunk {chunk_index} after {max_retries} attempts: {e}")
|
||||
raise
|
||||
|
||||
debug_log(f"Worker[{worker_id[:8]}] - Sent chunk {chunk_index} ({chunk_size} tiles, ~{used/1e6:.2f} MB)")
|
||||
chunk_index += 1
|
||||
i = j
|
||||
|
||||
async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id):
|
||||
"""Send completion signal to master in static mode when no tiles are left."""
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('is_last', 'true')
|
||||
data.add_field('batch_size', '0')
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_tiles"
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Worker {worker_id} sent static completion signal")
|
||||
|
||||
async def _request_work_item_from_master(
|
||||
self,
|
||||
multi_job_id,
|
||||
master_url,
|
||||
worker_id,
|
||||
endpoint="/distributed/request_image",
|
||||
):
|
||||
"""Request one work item from master with retry/backoff and total timeout."""
|
||||
max_retries = 10
|
||||
retry_delay = 0.5
|
||||
start_time = time.monotonic()
|
||||
url = f"{master_url}{endpoint}"
|
||||
|
||||
for attempt in range(max_retries):
|
||||
if time.monotonic() - start_time > 30:
|
||||
log(f"Total request timeout after 30s for worker {worker_id}")
|
||||
return None
|
||||
|
||||
try:
|
||||
session = await get_client_session()
|
||||
async with session.post(url, json={
|
||||
'worker_id': str(worker_id),
|
||||
'multi_job_id': multi_job_id
|
||||
}) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
if response.status == 404:
|
||||
text = await response.text()
|
||||
debug_log(f"Job not found (404), will retry: {text}")
|
||||
await asyncio.sleep(1.0)
|
||||
else:
|
||||
text = await response.text()
|
||||
debug_log(
|
||||
f"Request work item failed ({response.status}) for worker {worker_id}: {text}"
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
if attempt < max_retries - 1:
|
||||
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}")
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay = min(retry_delay * 2, 5.0)
|
||||
else:
|
||||
log(f"Failed to request work item after {max_retries} attempts: {exc}")
|
||||
raise
|
||||
|
||||
return None
|
||||
|
||||
async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Request an image index to process from master in dynamic mode."""
|
||||
data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
|
||||
if not data:
|
||||
return None, 0
|
||||
image_idx = data.get('image_idx')
|
||||
estimated_remaining = data.get('estimated_remaining', 0)
|
||||
return image_idx, estimated_remaining
|
||||
|
||||
async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
|
||||
"""Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
|
||||
data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
|
||||
if not data:
|
||||
return None, 0, False
|
||||
tile_idx = data.get('tile_idx')
|
||||
estimated_remaining = data.get('estimated_remaining', 0)
|
||||
batched_static = data.get('batched_static', False)
|
||||
return tile_idx, estimated_remaining, batched_static
|
||||
|
||||
async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id,
|
||||
master_url, worker_id, is_last):
|
||||
"""Send a processed full image back to master in dynamic mode."""
|
||||
# Serialize image to PNG
|
||||
byte_io = io.BytesIO()
|
||||
image_pil.save(byte_io, format='PNG', compress_level=0)
|
||||
byte_io.seek(0)
|
||||
|
||||
# Prepare form data
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('image_idx', str(image_idx))
|
||||
data.add_field('is_last', str(is_last))
|
||||
data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png',
|
||||
content_type='image/png')
|
||||
|
||||
# Retry logic
|
||||
max_retries = 5
|
||||
retry_delay = 0.5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_image"
|
||||
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Successfully sent image {image_idx} to master")
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
if attempt < max_retries - 1:
|
||||
debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
|
||||
await asyncio.sleep(retry_delay)
|
||||
retry_delay *= 2
|
||||
else:
|
||||
log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
|
||||
raise
|
||||
|
||||
async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
|
||||
"""Send completion signal to master in dynamic mode."""
|
||||
# Send a dummy request with is_last=True
|
||||
data = aiohttp.FormData()
|
||||
data.add_field('multi_job_id', multi_job_id)
|
||||
data.add_field('worker_id', str(worker_id))
|
||||
data.add_field('is_last', 'true')
|
||||
# No image data - just completion signal
|
||||
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/submit_image"
|
||||
|
||||
async with session.post(url, data=data) as response:
|
||||
response.raise_for_status()
|
||||
debug_log(f"Worker {worker_id} sent completion signal")
|
||||
|
||||
async def _check_job_status(self, multi_job_id, master_url):
|
||||
"""Check if job is ready on the master."""
|
||||
try:
|
||||
session = await get_client_session()
|
||||
url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
|
||||
async with session.get(url) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
return data.get('ready', False)
|
||||
return False
|
||||
except Exception as e:
|
||||
debug_log(f"Job status check failed: {e}")
|
||||
return False
|
||||
|
||||
async def _async_yield(self):
|
||||
"""Simple async yield to allow event loop processing."""
|
||||
await asyncio.sleep(0)
|
||||
+1
-100
@@ -3,10 +3,6 @@ Async helper utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
import execution
|
||||
import server
|
||||
from typing import Optional, Any, Coroutine
|
||||
from .network import get_server_loop
|
||||
|
||||
@@ -51,99 +47,4 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
|
||||
|
||||
if error:
|
||||
raise error
|
||||
return result
|
||||
|
||||
|
||||
prompt_server = server.PromptServer.instance
|
||||
|
||||
|
||||
def _summarize_node_errors(node_errors: dict) -> str:
|
||||
if not isinstance(node_errors, dict) or not node_errors:
|
||||
return ""
|
||||
|
||||
parts = []
|
||||
for node_id, entry in node_errors.items():
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
class_type = str(entry.get("class_type") or "UnknownNode")
|
||||
for err in entry.get("errors", []):
|
||||
if not isinstance(err, dict):
|
||||
continue
|
||||
message = str(err.get("message") or "validation error")
|
||||
details = str(err.get("details") or "").strip()
|
||||
parts.append(
|
||||
f"{class_type}#{node_id}: {message}{f' ({details})' if details else ''}"
|
||||
)
|
||||
if len(parts) >= 5:
|
||||
return " | ".join(parts)
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
"""Raised when a prompt fails ComfyUI validation with structured context."""
|
||||
|
||||
def __init__(self, error_payload, node_errors=None):
|
||||
payload = error_payload if isinstance(error_payload, dict) else {
|
||||
"type": "prompt_validation_failed",
|
||||
"message": str(error_payload),
|
||||
"details": "",
|
||||
"extra_info": {},
|
||||
}
|
||||
self.validation_error = dict(payload)
|
||||
self.node_errors = node_errors if isinstance(node_errors, dict) else {}
|
||||
|
||||
if self.node_errors:
|
||||
details = str(self.validation_error.get("details") or "").strip()
|
||||
if not details:
|
||||
summary = _summarize_node_errors(self.node_errors)
|
||||
if summary:
|
||||
self.validation_error["details"] = summary
|
||||
|
||||
merged = dict(self.validation_error)
|
||||
if self.node_errors:
|
||||
merged["node_errors"] = self.node_errors
|
||||
super().__init__(f"Invalid prompt: {merged}")
|
||||
|
||||
|
||||
async def queue_prompt_payload(
|
||||
prompt_obj,
|
||||
workflow_meta=None,
|
||||
client_id=None,
|
||||
include_queue_metadata=False,
|
||||
):
|
||||
"""Validate and queue a prompt via ComfyUI's prompt queue."""
|
||||
payload = {"prompt": prompt_obj}
|
||||
payload = prompt_server.trigger_on_prompt(payload)
|
||||
prompt = payload["prompt"]
|
||||
|
||||
prompt_id = str(uuid.uuid4())
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, None)
|
||||
if not valid[0]:
|
||||
error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
|
||||
node_errors = valid[3] if len(valid) > 3 else {}
|
||||
raise PromptValidationError(error_payload, node_errors)
|
||||
|
||||
extra_data = {"create_time": int(time.time() * 1000)}
|
||||
if workflow_meta:
|
||||
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
|
||||
if client_id:
|
||||
extra_data["client_id"] = client_id
|
||||
|
||||
sensitive = {}
|
||||
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
|
||||
if include_queue_metadata:
|
||||
return {
|
||||
"prompt_id": prompt_id,
|
||||
"number": number,
|
||||
"node_errors": {},
|
||||
}
|
||||
|
||||
return prompt_id
|
||||
return result
|
||||
@@ -1,103 +0,0 @@
|
||||
import base64
|
||||
import binascii
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .image import ensure_contiguous
|
||||
|
||||
|
||||
MAX_AUDIO_PAYLOAD_BYTES = int(
|
||||
os.environ.get("COMFYUI_MAX_AUDIO_PAYLOAD_BYTES", str(256 * 1024 * 1024))
|
||||
)
|
||||
|
||||
|
||||
def encode_audio_payload(audio_payload):
|
||||
"""Serialize an AUDIO dict into JSON-safe canonical envelope payload."""
|
||||
if not isinstance(audio_payload, dict):
|
||||
return None
|
||||
|
||||
waveform = audio_payload.get("waveform")
|
||||
if waveform is None or not isinstance(waveform, torch.Tensor) or waveform.numel() == 0:
|
||||
return None
|
||||
|
||||
sample_rate = audio_payload.get("sample_rate", 44100)
|
||||
try:
|
||||
sample_rate = int(sample_rate)
|
||||
except (TypeError, ValueError):
|
||||
sample_rate = 44100
|
||||
|
||||
waveform_cpu = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()
|
||||
data_bytes = waveform_cpu.numpy().tobytes()
|
||||
if len(data_bytes) > MAX_AUDIO_PAYLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Audio payload too large: {len(data_bytes)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
|
||||
)
|
||||
|
||||
return {
|
||||
"sample_rate": sample_rate,
|
||||
"shape": [int(dim) for dim in waveform_cpu.shape],
|
||||
"dtype": "float32",
|
||||
"data": base64.b64encode(data_bytes).decode("ascii"),
|
||||
}
|
||||
|
||||
|
||||
def decode_audio_payload(audio_payload):
|
||||
"""Decode canonical envelope audio payload into an AUDIO dict."""
|
||||
if audio_payload is None:
|
||||
return None
|
||||
if not isinstance(audio_payload, dict):
|
||||
raise ValueError("Field 'audio' must be an object when provided.")
|
||||
|
||||
encoded = audio_payload.get("data")
|
||||
shape = audio_payload.get("shape")
|
||||
sample_rate = audio_payload.get("sample_rate", 44100)
|
||||
dtype = audio_payload.get("dtype", "float32")
|
||||
|
||||
if not isinstance(encoded, str) or not encoded.strip():
|
||||
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
||||
if not isinstance(shape, list) or len(shape) != 3:
|
||||
raise ValueError("Field 'audio.shape' must be a 3-item list [batch, channels, samples].")
|
||||
if dtype != "float32":
|
||||
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
||||
|
||||
try:
|
||||
shape_tuple = tuple(int(dim) for dim in shape)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
||||
|
||||
if shape_tuple[0] <= 0 or shape_tuple[1] <= 0 or shape_tuple[2] < 0:
|
||||
raise ValueError(
|
||||
"Field 'audio.shape' must be [batch>0, channels>0, samples>=0]."
|
||||
)
|
||||
|
||||
try:
|
||||
sample_rate = int(sample_rate)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.sample_rate' must be an integer.") from exc
|
||||
if sample_rate <= 0:
|
||||
raise ValueError("Field 'audio.sample_rate' must be positive.")
|
||||
|
||||
try:
|
||||
raw = base64.b64decode(encoded, validate=True)
|
||||
except (binascii.Error, ValueError) as exc:
|
||||
raise ValueError("Field 'audio.data' is not valid base64.") from exc
|
||||
|
||||
if len(raw) > MAX_AUDIO_PAYLOAD_BYTES:
|
||||
raise ValueError(
|
||||
f"Field 'audio.data' too large: {len(raw)} bytes exceeds {MAX_AUDIO_PAYLOAD_BYTES}."
|
||||
)
|
||||
|
||||
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
||||
if len(raw) != expected_bytes:
|
||||
raise ValueError(
|
||||
f"Field 'audio.data' byte size mismatch: expected {expected_bytes}, got {len(raw)}."
|
||||
)
|
||||
|
||||
array = np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple)
|
||||
waveform = torch.from_numpy(array.copy())
|
||||
return {
|
||||
"waveform": ensure_contiguous(waveform),
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
@@ -0,0 +1,397 @@
|
||||
"""
|
||||
Lightweight Cloudflare Tunnel manager.
|
||||
|
||||
Responsibilities:
|
||||
- Locate or download the cloudflared binary for the current platform.
|
||||
- Start/stop a quick tunnel pointing to the local ComfyUI server.
|
||||
- Track status and persist minimal state in the shared config.
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import platform
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import stat
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from urllib import request, error as urlerror
|
||||
|
||||
from .config import load_config, save_config
|
||||
from .logging import debug_log, log
|
||||
from .network import get_server_port
|
||||
from .process import is_process_alive, terminate_process
|
||||
|
||||
# Regex to capture the generated public URL from cloudflared output
|
||||
PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
|
||||
|
||||
# Default timeout (seconds) to wait for a tunnel URL before giving up
|
||||
TUNNEL_START_TIMEOUT = 25
|
||||
|
||||
|
||||
def _normalize_host(value):
|
||||
if not value or not isinstance(value, str):
|
||||
return ""
|
||||
value = value.strip()
|
||||
value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
|
||||
return value.split("/")[0]
|
||||
|
||||
|
||||
class CloudflareTunnelManager:
|
||||
def __init__(self):
|
||||
self.process = None
|
||||
self.pid = None
|
||||
self.public_url = None
|
||||
self.last_error = None
|
||||
self.log_file = None
|
||||
self.status = "stopped"
|
||||
self.previous_master_host = None
|
||||
|
||||
self._lock = asyncio.Lock()
|
||||
self._url_event = None
|
||||
self._reader_thread = None
|
||||
self._loop = None
|
||||
self._recent_logs = []
|
||||
|
||||
self.binary_path = None
|
||||
self._restore_state()
|
||||
|
||||
# --- Binary discovery & download ---
|
||||
@property
|
||||
def base_dir(self):
|
||||
return os.path.dirname(os.path.dirname(__file__))
|
||||
|
||||
@property
|
||||
def bin_dir(self):
|
||||
return os.path.join(self.base_dir, "bin")
|
||||
|
||||
def _get_asset_name(self):
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
|
||||
if system == "windows":
|
||||
if "arm" in machine:
|
||||
return "cloudflared-windows-arm64.exe"
|
||||
return "cloudflared-windows-amd64.exe"
|
||||
if system == "darwin":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-darwin-arm64"
|
||||
return "cloudflared-darwin-amd64"
|
||||
if system == "linux":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-linux-arm64"
|
||||
return "cloudflared-linux-amd64"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
|
||||
|
||||
def _download_cloudflared(self):
|
||||
asset = self._get_asset_name()
|
||||
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
|
||||
os.makedirs(self.bin_dir, exist_ok=True)
|
||||
target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
|
||||
|
||||
debug_log(f"Downloading cloudflared from {url}")
|
||||
try:
|
||||
with request.urlopen(url, timeout=30) as resp:
|
||||
with open(target_path, "wb") as f:
|
||||
shutil.copyfileobj(resp, f)
|
||||
except urlerror.URLError as exc:
|
||||
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
|
||||
|
||||
# Make executable
|
||||
st = os.stat(target_path)
|
||||
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
|
||||
debug_log(f"Downloaded cloudflared to {target_path}")
|
||||
return target_path
|
||||
|
||||
def ensure_binary(self):
|
||||
# Allow explicit override
|
||||
env_path = os.environ.get("CLOUDFLARED_PATH")
|
||||
if env_path and os.path.exists(env_path):
|
||||
self.binary_path = env_path
|
||||
return env_path
|
||||
|
||||
# Check cached path
|
||||
cached = self.binary_path
|
||||
if cached and os.path.exists(cached):
|
||||
return cached
|
||||
|
||||
# Local bin directory
|
||||
candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
|
||||
if os.path.exists(candidate):
|
||||
self.binary_path = candidate
|
||||
return candidate
|
||||
|
||||
# System PATH
|
||||
path_binary = shutil.which("cloudflared")
|
||||
if path_binary:
|
||||
self.binary_path = path_binary
|
||||
return path_binary
|
||||
|
||||
# Download on demand
|
||||
self.binary_path = self._download_cloudflared()
|
||||
return self.binary_path
|
||||
|
||||
# --- State helpers ---
|
||||
def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
|
||||
|
||||
if status is not None:
|
||||
tunnel_cfg["status"] = status
|
||||
if public_url is not None:
|
||||
tunnel_cfg["public_url"] = public_url
|
||||
if pid is not None:
|
||||
tunnel_cfg["pid"] = pid
|
||||
if log_file is not None:
|
||||
tunnel_cfg["log_file"] = log_file
|
||||
if previous_host is not None:
|
||||
tunnel_cfg["previous_master_host"] = previous_host
|
||||
if master_host is not None:
|
||||
cfg.setdefault("master", {})["host"] = master_host
|
||||
|
||||
cfg["tunnel"] = tunnel_cfg
|
||||
save_config(cfg)
|
||||
|
||||
def _restore_state(self):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
|
||||
|
||||
self.public_url = tunnel_cfg.get("public_url") or None
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host")
|
||||
self.log_file = tunnel_cfg.get("log_file")
|
||||
pid = tunnel_cfg.get("pid")
|
||||
|
||||
if pid and is_process_alive(pid):
|
||||
self.pid = pid
|
||||
self.status = tunnel_cfg.get("status", "running")
|
||||
debug_log(f"Detected existing cloudflared process (pid={pid})")
|
||||
else:
|
||||
# Clear stale info
|
||||
self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
|
||||
self.status = "stopped"
|
||||
self.pid = None
|
||||
|
||||
# --- Process output handling ---
|
||||
def _append_log(self, line):
|
||||
if self.log_file:
|
||||
try:
|
||||
with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
|
||||
f.write(line + "\n")
|
||||
except Exception as exc: # pragma: no cover - best effort logging
|
||||
debug_log(f"Failed to write tunnel log: {exc}")
|
||||
|
||||
self._recent_logs.append(line)
|
||||
# Keep last ~200 lines to avoid unbounded growth
|
||||
if len(self._recent_logs) > 200:
|
||||
self._recent_logs = self._recent_logs[-200:]
|
||||
|
||||
def _reader(self):
|
||||
assert self.process is not None
|
||||
loop = self._loop
|
||||
for raw_line in iter(self.process.stdout.readline, ""):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
self._append_log(line)
|
||||
match = PUBLIC_URL_PATTERN.search(line)
|
||||
if match and not self.public_url:
|
||||
self.public_url = match.group(1).rstrip("/")
|
||||
self.status = "running"
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
# Capture obvious error strings
|
||||
if "error" in line.lower() and not self.last_error:
|
||||
self.last_error = line
|
||||
|
||||
# Process exited
|
||||
if self.status == "starting" and not self.public_url:
|
||||
self.status = "error"
|
||||
if not self.last_error:
|
||||
self.last_error = "Cloudflare tunnel exited before becoming ready"
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
elif self.status == "running":
|
||||
# Treat unexpected exit as stopped so UI can retry
|
||||
self.status = "stopped"
|
||||
elif self.status == "stopped":
|
||||
# Normal stop path, nothing to do
|
||||
pass
|
||||
|
||||
# --- Public API ---
|
||||
async def start_tunnel(self):
|
||||
async with self._lock:
|
||||
if self.process and self.process.poll() is None:
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.process.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
# If we had a stale pid, ensure it's not running
|
||||
if self.pid and is_process_alive(self.pid):
|
||||
debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
|
||||
await self.stop_tunnel()
|
||||
|
||||
binary = await asyncio.to_thread(self.ensure_binary)
|
||||
port = get_server_port()
|
||||
self.status = "starting"
|
||||
self.last_error = None
|
||||
self.public_url = None
|
||||
self._recent_logs = []
|
||||
|
||||
# Remember current master host so we can restore it if the tunnel stops
|
||||
config = load_config()
|
||||
master_host = (config.get("master") or {}).get("host") or ""
|
||||
tunnel_cfg = config.get("tunnel") or {}
|
||||
if tunnel_cfg.get("previous_master_host"):
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host")
|
||||
else:
|
||||
self.previous_master_host = master_host
|
||||
|
||||
os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
|
||||
timestamp = time.strftime("%Y%m%d-%H%M%S")
|
||||
self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
|
||||
|
||||
cmd = [
|
||||
binary,
|
||||
"tunnel",
|
||||
"--no-autoupdate",
|
||||
"--url",
|
||||
f"http://127.0.0.1:{port}",
|
||||
]
|
||||
|
||||
debug_log(f"Starting cloudflared: {' '.join(cmd)}")
|
||||
try:
|
||||
self.process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
self.status = "error"
|
||||
raise RuntimeError("cloudflared binary not found")
|
||||
except Exception as exc:
|
||||
self.status = "error"
|
||||
raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
|
||||
|
||||
self.pid = self.process.pid
|
||||
self._persist_state(
|
||||
status="starting",
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host
|
||||
)
|
||||
|
||||
# Kick off background reader
|
||||
self._loop = asyncio.get_running_loop()
|
||||
self._url_event = asyncio.Event()
|
||||
self._reader_thread = threading.Thread(target=self._reader, daemon=True)
|
||||
self._reader_thread.start()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
|
||||
except asyncio.TimeoutError:
|
||||
self.last_error = "Timed out waiting for Cloudflare to assign a URL"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(self.last_error)
|
||||
|
||||
if self.status != "running" or not self.public_url:
|
||||
error_msg = self.last_error or "Cloudflare tunnel failed to start"
|
||||
await self.stop_tunnel()
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
|
||||
master_host = _normalize_host(self.public_url)
|
||||
self._persist_state(
|
||||
status="running",
|
||||
public_url=self.public_url,
|
||||
pid=self.pid,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host or "",
|
||||
master_host=master_host
|
||||
)
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": self.pid,
|
||||
"log_file": self.log_file,
|
||||
}
|
||||
|
||||
async def stop_tunnel(self):
|
||||
async with self._lock:
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if not pid:
|
||||
self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
|
||||
self.status = "stopped"
|
||||
return {"status": "stopped"}
|
||||
|
||||
debug_log(f"Stopping cloudflared (pid={pid})")
|
||||
if self.process:
|
||||
terminate_process(self.process, timeout=5)
|
||||
else:
|
||||
try:
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
time.sleep(0.5)
|
||||
except Exception as exc: # pragma: no cover - best effort
|
||||
debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
|
||||
|
||||
config = load_config()
|
||||
tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
|
||||
self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
|
||||
active_url = tunnel_cfg.get("public_url")
|
||||
current_master_host = (config.get("master") or {}).get("host")
|
||||
restore_host = None
|
||||
if active_url:
|
||||
active_host = _normalize_host(active_url)
|
||||
current_host = _normalize_host(current_master_host)
|
||||
if current_host == active_host:
|
||||
restore_host = self.previous_master_host or ""
|
||||
|
||||
self.status = "stopped"
|
||||
self.public_url = None
|
||||
self.pid = None
|
||||
self.process = None
|
||||
self.last_error = None
|
||||
self._url_event = None
|
||||
if self._reader_thread and self._reader_thread.is_alive():
|
||||
self._reader_thread.join(timeout=1)
|
||||
self._reader_thread = None
|
||||
self._persist_state(
|
||||
status="stopped",
|
||||
public_url="",
|
||||
pid=None,
|
||||
log_file=self.log_file,
|
||||
previous_host=self.previous_master_host,
|
||||
master_host=restore_host
|
||||
)
|
||||
return {"status": "stopped"}
|
||||
|
||||
def get_status(self):
|
||||
alive = False
|
||||
pid = self.process.pid if self.process else self.pid
|
||||
if pid:
|
||||
alive = is_process_alive(pid)
|
||||
if not alive and self.status == "running":
|
||||
self.status = "stopped"
|
||||
|
||||
return {
|
||||
"status": self.status,
|
||||
"public_url": self.public_url,
|
||||
"pid": pid,
|
||||
"log_file": self.log_file,
|
||||
"last_error": self.last_error,
|
||||
"binary_path": self.binary_path or shutil.which("cloudflared"),
|
||||
"recent_logs": self._recent_logs[-20:],
|
||||
"previous_master_host": self.previous_master_host,
|
||||
}
|
||||
|
||||
|
||||
# Singleton tunnel manager
|
||||
cloudflare_tunnel_manager = CloudflareTunnelManager()
|
||||
@@ -1,5 +0,0 @@
|
||||
from .tunnel import CloudflareTunnelManager
|
||||
|
||||
cloudflare_tunnel_manager = CloudflareTunnelManager()
|
||||
|
||||
__all__ = ["CloudflareTunnelManager", "cloudflare_tunnel_manager"]
|
||||
@@ -1,83 +0,0 @@
|
||||
"""Cloudflared binary discovery and download helpers."""
|
||||
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import stat
|
||||
from urllib import error as urlerror
|
||||
from urllib import request
|
||||
|
||||
from ..logging import debug_log
|
||||
|
||||
|
||||
def _get_project_root():
|
||||
return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
|
||||
|
||||
def _get_cloudflared_dir():
|
||||
return os.path.join(_get_project_root(), "bin")
|
||||
|
||||
|
||||
def _get_platform_binary_name():
|
||||
system = platform.system().lower()
|
||||
machine = platform.machine().lower()
|
||||
|
||||
if system == "windows":
|
||||
if "arm" in machine:
|
||||
return "cloudflared-windows-arm64.exe"
|
||||
return "cloudflared-windows-amd64.exe"
|
||||
if system == "darwin":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-darwin-arm64"
|
||||
return "cloudflared-darwin-amd64"
|
||||
if system == "linux":
|
||||
if machine in ("arm64", "aarch64"):
|
||||
return "cloudflared-linux-arm64"
|
||||
return "cloudflared-linux-amd64"
|
||||
|
||||
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
|
||||
|
||||
|
||||
def _get_binary_path(bin_dir=None):
|
||||
bin_dir = bin_dir or _get_cloudflared_dir()
|
||||
binary_name = "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared"
|
||||
return os.path.join(bin_dir, binary_name)
|
||||
|
||||
|
||||
def _download_cloudflared():
|
||||
asset = _get_platform_binary_name()
|
||||
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
|
||||
|
||||
bin_dir = _get_cloudflared_dir()
|
||||
os.makedirs(bin_dir, exist_ok=True)
|
||||
target_path = _get_binary_path(bin_dir)
|
||||
|
||||
debug_log(f"Downloading cloudflared from {url}")
|
||||
try:
|
||||
with request.urlopen(url, timeout=30) as resp:
|
||||
with open(target_path, "wb") as f:
|
||||
shutil.copyfileobj(resp, f)
|
||||
except urlerror.URLError as exc:
|
||||
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
|
||||
|
||||
st = os.stat(target_path)
|
||||
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
|
||||
debug_log(f"Downloaded cloudflared to {target_path}")
|
||||
return target_path
|
||||
|
||||
|
||||
def ensure_binary() -> str:
|
||||
"""Return a usable cloudflared binary path, downloading if necessary."""
|
||||
env_path = os.environ.get("CLOUDFLARED_PATH")
|
||||
if env_path and os.path.exists(env_path):
|
||||
return env_path
|
||||
|
||||
local_candidate = _get_binary_path()
|
||||
if os.path.exists(local_candidate):
|
||||
return local_candidate
|
||||
|
||||
path_binary = shutil.which("cloudflared")
|
||||
if path_binary:
|
||||
return path_binary
|
||||
|
||||
return _download_cloudflared()
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Background cloudflared process output reader."""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
import threading
|
||||
|
||||
from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE
|
||||
from ..logging import debug_log
|
||||
|
||||
PUBLIC_URL_PATTERN = re.compile(
|
||||
r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
class ProcessReader:
|
||||
def __init__(self, log_file=None):
|
||||
self._process = None
|
||||
self._thread = None
|
||||
self._loop = None
|
||||
self._url_event = None
|
||||
self._public_url = None
|
||||
self._last_error = None
|
||||
self._recent_logs = []
|
||||
self._log_file = log_file
|
||||
|
||||
def set_log_file(self, log_file):
|
||||
self._log_file = log_file
|
||||
|
||||
def _append_log(self, line):
|
||||
if self._log_file:
|
||||
try:
|
||||
with open(self._log_file, "a", encoding="utf-8", errors="replace") as f:
|
||||
f.write(line + "\n")
|
||||
except Exception as exc: # pragma: no cover
|
||||
debug_log(f"Failed to write tunnel log: {exc}")
|
||||
|
||||
self._recent_logs.append(line)
|
||||
if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE:
|
||||
self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:]
|
||||
|
||||
def _reader(self):
|
||||
process = self._process
|
||||
if process is None:
|
||||
return
|
||||
|
||||
loop = self._loop
|
||||
for raw_line in iter(process.stdout.readline, ""):
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
self._append_log(line)
|
||||
match = PUBLIC_URL_PATTERN.search(line)
|
||||
if match and not self._public_url:
|
||||
self._public_url = match.group(1).rstrip("/")
|
||||
if self._url_event and loop:
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
if "error" in line.lower() and not self._last_error:
|
||||
self._last_error = line
|
||||
|
||||
if self._url_event and loop:
|
||||
if not self._last_error and not self._public_url:
|
||||
self._last_error = "Cloudflare tunnel exited before becoming ready"
|
||||
loop.call_soon_threadsafe(self._url_event.set)
|
||||
|
||||
def start(self, process, loop):
|
||||
self._process = process
|
||||
self._loop = loop
|
||||
self._url_event = asyncio.Event()
|
||||
self._public_url = None
|
||||
self._last_error = None
|
||||
self._recent_logs = []
|
||||
self._thread = threading.Thread(target=self._reader, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
async def wait_for_url(self, timeout):
|
||||
if not self._url_event:
|
||||
return None
|
||||
await asyncio.wait_for(self._url_event.wait(), timeout=timeout)
|
||||
return self._public_url
|
||||
|
||||
def stop(self):
|
||||
if self._thread and self._thread.is_alive():
|
||||
self._thread.join(timeout=1)
|
||||
self._thread = None
|
||||
self._process = None
|
||||
self._loop = None
|
||||
self._url_event = None
|
||||
|
||||
def get_url(self):
|
||||
return self._public_url
|
||||
|
||||
def get_last_error(self):
|
||||
return self._last_error
|
||||
|
||||
def get_recent_logs(self):
|
||||
return list(self._recent_logs)
|
||||
@@ -1,81 +0,0 @@
|
||||
"""Cloudflare tunnel state persistence helpers."""
|
||||
|
||||
from ..config import load_config, save_config
|
||||
from ..network import normalize_host
|
||||
|
||||
|
||||
def _get_tunnel_config(cfg):
|
||||
tunnel_cfg = cfg.get("tunnel", {})
|
||||
if isinstance(tunnel_cfg, dict):
|
||||
return tunnel_cfg
|
||||
return {}
|
||||
|
||||
|
||||
def load_tunnel_state():
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {}
|
||||
return {
|
||||
"status": tunnel_cfg.get("status", "stopped"),
|
||||
"public_url": tunnel_cfg.get("public_url") or None,
|
||||
"pid": tunnel_cfg.get("pid"),
|
||||
"log_file": tunnel_cfg.get("log_file"),
|
||||
"previous_master_host": tunnel_cfg.get("previous_master_host"),
|
||||
"master_host": master_cfg.get("host"),
|
||||
}
|
||||
|
||||
|
||||
def persist_tunnel_state(
|
||||
status=None,
|
||||
public_url=None,
|
||||
pid=None,
|
||||
log_file=None,
|
||||
previous_host=None,
|
||||
master_host=None,
|
||||
):
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
|
||||
if status is not None:
|
||||
tunnel_cfg["status"] = status
|
||||
if public_url is not None:
|
||||
tunnel_cfg["public_url"] = public_url
|
||||
if pid is not None:
|
||||
tunnel_cfg["pid"] = pid
|
||||
if log_file is not None:
|
||||
tunnel_cfg["log_file"] = log_file
|
||||
if previous_host is not None:
|
||||
tunnel_cfg["previous_master_host"] = previous_host
|
||||
if master_host is not None:
|
||||
cfg.setdefault("master", {})["host"] = master_host
|
||||
|
||||
cfg["tunnel"] = tunnel_cfg
|
||||
save_config(cfg)
|
||||
|
||||
|
||||
def clear_tunnel_state(log_file=None, previous_host=None, master_host=None):
|
||||
persist_tunnel_state(
|
||||
status="stopped",
|
||||
public_url="",
|
||||
pid=None,
|
||||
log_file=log_file,
|
||||
previous_host=previous_host,
|
||||
master_host=master_host,
|
||||
)
|
||||
|
||||
|
||||
def resolve_restore_master_host(previous_master_host):
|
||||
"""Determine whether master host should be restored after tunnel stop."""
|
||||
cfg = load_config()
|
||||
tunnel_cfg = _get_tunnel_config(cfg)
|
||||
active_url = tunnel_cfg.get("public_url")
|
||||
current_master_host = (cfg.get("master") or {}).get("host")
|
||||
|
||||
if not active_url:
|
||||
return None
|
||||
|
||||
active_host = normalize_host(active_url)
|
||||
current_host = normalize_host(current_master_host)
|
||||
if current_host == active_host:
|
||||
return previous_master_host or ""
|
||||
return None
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user