Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
87c0cda9f6 | ||
|
|
371517ae14 | ||
|
|
6eb94a52fe | ||
|
|
a281d03c6d | ||
|
|
d9485265ea | ||
|
|
cd13521d24 | ||
|
|
1093818327 | ||
|
|
6f85129e45 | ||
|
|
a53d88f7fb | ||
|
|
9be342f5fb | ||
|
|
a1faf7e916 |
@@ -38,11 +38,12 @@
|
||||
|
||||
<img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
|
||||
|
||||
ComfyUI Distributed supports three types of workers:
|
||||
ComfyUI Distributed supports four types of workers:
|
||||
|
||||
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
|
||||
- **Remote Workers** - GPUs on other computers in your network
|
||||
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
|
||||
- **Serverless (Runpod)** - Queue-based endpoint pools for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
|
||||
+15
-22
@@ -1,29 +1,22 @@
|
||||
# Import everything needed from the main module
|
||||
from .distributed import (
|
||||
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
"""ComfyUI-Distributed's native V3 extension entrypoint."""
|
||||
from comfy_api.v0_0_2 import ComfyExtension, io
|
||||
|
||||
# Import utilities
|
||||
from .utils.config import ensure_config_exists, CONFIG_FILE
|
||||
from .utils.logging import debug_log
|
||||
from .nodes.v3 import NODES
|
||||
from .runtime.bootstrap import initialize
|
||||
|
||||
# Import distributed upscale nodes
|
||||
from .nodes.distributed_upscale import (
|
||||
NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
|
||||
)
|
||||
WEB_DIRECTORY = './web'
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
|
||||
ensure_config_exists()
|
||||
class DistributedExtension(ComfyExtension):
|
||||
async def on_load(self) -> None:
|
||||
initialize()
|
||||
|
||||
# Merge node mappings
|
||||
NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return list(NODES)
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
async def comfy_entrypoint() -> DistributedExtension:
|
||||
return DistributedExtension()
|
||||
|
||||
|
||||
__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
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
|
||||
|
||||
+372
-69
@@ -1,5 +1,6 @@
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
from aiohttp import web
|
||||
import server
|
||||
@@ -24,6 +25,12 @@ except ImportError:
|
||||
save_config(config)
|
||||
from ..utils.logging import debug_log, log
|
||||
from ..utils.network import handle_api_error, normalize_host
|
||||
from ..utils.secret_fields import is_secret_field_name
|
||||
from ..utils.serverless_config import (
|
||||
ServerlessSecretChanges,
|
||||
public_config,
|
||||
validate_serverless_worker,
|
||||
)
|
||||
|
||||
|
||||
def _positive_int(value):
|
||||
@@ -86,22 +93,160 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
|
||||
target[key] = normalizer(value) if (normalizer and value is not None) else value
|
||||
|
||||
|
||||
_SECRET_FIELDS = {"api_key", "clear_api_key"}
|
||||
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
||||
_SERVERLESS_ONLY_FIELDS = {
|
||||
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
|
||||
}
|
||||
_SERVERLESS_CONFIG_FIELDS = {
|
||||
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
||||
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
|
||||
}
|
||||
|
||||
|
||||
class _ActivePoolCheckError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def _error_response(message, status):
|
||||
return web.json_response({"status": "error", "message": message}, status=status)
|
||||
|
||||
|
||||
def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0):
|
||||
if isinstance(value, dict):
|
||||
for key, item in value.items():
|
||||
field_name = key.casefold() if isinstance(key, str) else ""
|
||||
if is_secret_field_name(field_name):
|
||||
if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS):
|
||||
return True
|
||||
if _contains_secret_fields(
|
||||
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
|
||||
):
|
||||
return True
|
||||
elif isinstance(value, (list, tuple)):
|
||||
return any(
|
||||
_contains_secret_fields(
|
||||
item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
|
||||
)
|
||||
for item in value
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _serverless_secret_store():
|
||||
from .orchestration.serverless_runtime import get_secret_store
|
||||
return get_secret_store()
|
||||
|
||||
|
||||
def _serverless_pool_is_active(worker_id):
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
return get_manager().pool_is_active(worker_id)
|
||||
|
||||
|
||||
def _ensure_pool_idle(worker_id):
|
||||
try:
|
||||
active = _serverless_pool_is_active(worker_id)
|
||||
except Exception:
|
||||
raise _ActivePoolCheckError from None
|
||||
if active:
|
||||
raise _ActivePoolCheckError
|
||||
|
||||
|
||||
def _worker_by_id(workers, worker_id):
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("id") == worker_id:
|
||||
return worker
|
||||
return None
|
||||
|
||||
|
||||
def _worker_type(worker):
|
||||
return worker.get("type", "local") if isinstance(worker, dict) else None
|
||||
|
||||
|
||||
def _strip_sensitive_fields(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: _strip_sensitive_fields(item)
|
||||
for key, item in value.items()
|
||||
if not is_secret_field_name(key)
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [_strip_sensitive_fields(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def _clean_serverless_worker(worker):
|
||||
clean = _strip_sensitive_fields(dict(worker))
|
||||
for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}:
|
||||
clean.pop(field, None)
|
||||
return clean
|
||||
|
||||
|
||||
def _validate_bulk_workers(workers):
|
||||
validated = []
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
||||
clean = _clean_serverless_worker(worker)
|
||||
validated.append(validate_serverless_worker(clean))
|
||||
else:
|
||||
validated.append(worker)
|
||||
return validated
|
||||
|
||||
|
||||
def _clean_bulk_type_conversions(old_workers, new_workers):
|
||||
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
for index, worker in enumerate(new_workers):
|
||||
if not isinstance(worker, dict) or worker.get("id") is None:
|
||||
continue
|
||||
old = old_by_id.get(str(worker["id"]))
|
||||
if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
|
||||
continue
|
||||
converted = dict(worker)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
converted.pop(field, None)
|
||||
new_workers[index] = converted
|
||||
return new_workers
|
||||
|
||||
|
||||
def _bulk_serverless_changes(old_workers, new_workers):
|
||||
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
|
||||
secret_ids = set()
|
||||
for worker_id, old in old_by_id.items():
|
||||
new = new_by_id.get(worker_id)
|
||||
if old.get("type") == "serverless":
|
||||
if new is None or new.get("type") != "serverless":
|
||||
_ensure_pool_idle(worker_id)
|
||||
secret_ids.add(worker_id)
|
||||
elif (old.get("endpoint_id") != new.get("endpoint_id")
|
||||
or old.get("concurrency", 1) != new.get("concurrency", 1)):
|
||||
_ensure_pool_idle(worker_id)
|
||||
elif new is not None and new.get("type") == "serverless":
|
||||
secret_ids.add(worker_id)
|
||||
return secret_ids
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/config")
|
||||
async def get_config_endpoint(request):
|
||||
config = load_config()
|
||||
return web.json_response(config)
|
||||
try:
|
||||
config = load_config()
|
||||
return web.json_response(public_config(config, _serverless_secret_store()))
|
||||
except Exception:
|
||||
return _error_response("Unable to read config", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config")
|
||||
async def update_config_endpoint(request):
|
||||
"""Bulk config update with schema validation."""
|
||||
"""Bulk config update with schema validation and credential redaction."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Invalid JSON payload: {e}", 400)
|
||||
except Exception:
|
||||
return _error_response("Invalid JSON payload", 400)
|
||||
|
||||
if not isinstance(data, dict):
|
||||
return await handle_api_error(request, "Config payload must be an object", 400)
|
||||
return _error_response("Config payload must be an object", 400)
|
||||
if _contains_secret_fields(data):
|
||||
return _error_response("API keys cannot be set through bulk config", 400)
|
||||
|
||||
validated_settings = {}
|
||||
validated_root = {}
|
||||
@@ -118,9 +263,15 @@ async def update_config_endpoint(request):
|
||||
continue
|
||||
|
||||
if validator and not validator(value):
|
||||
errors.append(f"{key}: value {value!r} failed validation")
|
||||
errors.append(f"{key}: failed validation")
|
||||
continue
|
||||
|
||||
if key == "workers":
|
||||
try:
|
||||
value = _validate_bulk_workers(value)
|
||||
except (TypeError, ValueError):
|
||||
return _error_response("Invalid serverless worker configuration", 400)
|
||||
|
||||
if key in _SETTINGS_FIELDS:
|
||||
validated_settings[key] = value
|
||||
else:
|
||||
@@ -133,15 +284,41 @@ async def update_config_endpoint(request):
|
||||
"message": "; ".join(errors),
|
||||
}, status=400)
|
||||
|
||||
store = None
|
||||
response_config = None
|
||||
try:
|
||||
async with config_transaction() as config:
|
||||
settings = config.setdefault("settings", {})
|
||||
settings.update(validated_settings)
|
||||
for key, value in validated_root.items():
|
||||
config[key] = value
|
||||
return web.json_response({"status": "success", "config": config})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e)
|
||||
with ServerlessSecretChanges() as secret_changes:
|
||||
async with config_transaction() as config:
|
||||
if "workers" in validated_root:
|
||||
old_workers = config.get("workers", [])
|
||||
validated_root["workers"] = _clean_bulk_type_conversions(
|
||||
old_workers, validated_root["workers"]
|
||||
)
|
||||
changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"])
|
||||
for worker in validated_root["workers"]:
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False):
|
||||
store = store or _serverless_secret_store()
|
||||
if worker["id"] in changed_ids or not store.has(worker["id"]):
|
||||
return _error_response("Save an API key before enabling this pool", 400)
|
||||
if changed_ids:
|
||||
store = store or _serverless_secret_store()
|
||||
for worker_id in sorted(changed_ids):
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
if prior_secret is not None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
|
||||
settings = config.setdefault("settings", {})
|
||||
settings.update(validated_settings)
|
||||
for key, value in validated_root.items():
|
||||
config[key] = value
|
||||
response_config = public_config(config, store or _serverless_secret_store())
|
||||
response = web.json_response({"status": "success", "config": response_config})
|
||||
secret_changes.commit()
|
||||
return response
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to update config", 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
||||
@@ -164,79 +341,205 @@ async def queue_status_endpoint(request):
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
||||
async def update_worker_endpoint(request):
|
||||
return await _run_with_serverless_secret_changes(_update_worker_endpoint, request)
|
||||
|
||||
|
||||
async def _run_with_serverless_secret_changes(endpoint, request):
|
||||
with ServerlessSecretChanges() as secret_changes:
|
||||
response = await endpoint(request, secret_changes)
|
||||
if response.status < 400:
|
||||
secret_changes.commit()
|
||||
return response
|
||||
|
||||
|
||||
async def _update_worker_endpoint(request, secret_changes):
|
||||
store = None
|
||||
worker_id = None
|
||||
try:
|
||||
data = await request.json()
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
return _error_response("Invalid JSON payload", 400)
|
||||
if not isinstance(data, dict):
|
||||
return _error_response("Worker payload must be an object", 400)
|
||||
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
return _error_response("Missing worker_id", 400)
|
||||
|
||||
if _contains_secret_fields(data, allow_root_serverless_secrets=True):
|
||||
return _error_response("Invalid worker credentials", 400)
|
||||
|
||||
result = None
|
||||
async with config_transaction() as config:
|
||||
worker_found = False
|
||||
workers = config.setdefault("workers", [])
|
||||
existing = _worker_by_id(workers, worker_id)
|
||||
current_type = _worker_type(existing)
|
||||
target_type = data.get("type", current_type or "local")
|
||||
|
||||
for worker in workers:
|
||||
if worker["id"] == worker_id:
|
||||
_apply_field_patch(worker, data, _WORKER_FIELDS)
|
||||
worker_found = True
|
||||
break
|
||||
if target_type == "serverless":
|
||||
unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS
|
||||
if unknown:
|
||||
return _error_response("Unknown serverless worker field", 400)
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
return _error_response("worker_id must be a non-empty string", 400)
|
||||
if "api_key" in data and not isinstance(data["api_key"], str):
|
||||
return _error_response("Invalid worker credentials", 400)
|
||||
if "clear_api_key" in data and type(data["clear_api_key"]) is not bool:
|
||||
return _error_response("Invalid worker credentials", 400)
|
||||
supplied_key = data.get("api_key", "")
|
||||
clear_key = data.get("clear_api_key", False)
|
||||
if supplied_key and clear_key:
|
||||
return _error_response("Choose either api_key or clear_api_key", 400)
|
||||
if "enabled" in data and type(data["enabled"]) is not bool:
|
||||
return _error_response("enabled must be a boolean", 400)
|
||||
if "name" in data and not isinstance(data["name"], str):
|
||||
return _error_response("name must be a string", 400)
|
||||
|
||||
if not worker_found:
|
||||
# If worker not found and all required fields are provided, create new worker
|
||||
if all(key in data for key in ["name", "port", "cuda_device"]):
|
||||
new_worker = {
|
||||
"id": worker_id,
|
||||
"name": data["name"],
|
||||
"host": normalize_host(data.get("host", "localhost")),
|
||||
"port": data["port"],
|
||||
"cuda_device": data["cuda_device"],
|
||||
"enabled": data.get("enabled", False),
|
||||
"extra_args": data.get("extra_args", ""),
|
||||
"type": data.get("type", "local")
|
||||
}
|
||||
workers.append(new_worker)
|
||||
candidate = dict(existing or {})
|
||||
candidate["id"] = worker_id
|
||||
candidate["type"] = "serverless"
|
||||
candidate.setdefault("provider", "runpod")
|
||||
candidate.setdefault("enabled", False)
|
||||
for field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "enabled", "name"):
|
||||
if field in data:
|
||||
candidate[field] = data[field]
|
||||
candidate = _clean_serverless_worker(candidate)
|
||||
validated = validate_serverless_worker(candidate)
|
||||
if "enabled" in validated and type(validated["enabled"]) is not bool:
|
||||
return _error_response("enabled must be a boolean", 400)
|
||||
if "name" in validated and not isinstance(validated["name"], str):
|
||||
return _error_response("name must be a string", 400)
|
||||
|
||||
store = _serverless_secret_store()
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
converting_to_serverless = current_type != "serverless"
|
||||
next_secret = None if converting_to_serverless else prior_secret
|
||||
if supplied_key:
|
||||
next_secret = supplied_key
|
||||
elif clear_key:
|
||||
next_secret = None
|
||||
|
||||
endpoint_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("endpoint_id") != validated.get("endpoint_id")
|
||||
)
|
||||
concurrency_changed = (
|
||||
current_type == "serverless"
|
||||
and existing.get("concurrency", 1) != validated.get("concurrency", 1)
|
||||
)
|
||||
if current_type == "serverless" and (
|
||||
endpoint_changed or concurrency_changed or next_secret != prior_secret
|
||||
):
|
||||
_ensure_pool_idle(worker_id)
|
||||
|
||||
if validated.get("enabled", False) and next_secret is None:
|
||||
return _error_response("Save an API key before enabling this pool; disable it before clearing the key", 400)
|
||||
|
||||
if next_secret != prior_secret:
|
||||
if next_secret is None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
else:
|
||||
secret_changes.set(store, worker_id, next_secret)
|
||||
|
||||
if existing is None:
|
||||
workers.append(validated)
|
||||
else:
|
||||
return await handle_api_error(
|
||||
request,
|
||||
f"Worker {worker_id} not found and missing required fields for creation",
|
||||
404,
|
||||
)
|
||||
existing.clear()
|
||||
existing.update(validated)
|
||||
result = {"status": "success", "api_key_configured": next_secret is not None}
|
||||
|
||||
elif existing is not None:
|
||||
cleaned_existing = _strip_sensitive_fields(existing)
|
||||
existing.clear()
|
||||
existing.update(cleaned_existing)
|
||||
if current_type == "serverless":
|
||||
if _contains_secret_fields(data):
|
||||
return _error_response("Invalid worker credentials", 400)
|
||||
_ensure_pool_idle(worker_id)
|
||||
store = _serverless_secret_store()
|
||||
prior_secret = secret_changes.capture(store, worker_id)
|
||||
if prior_secret is not None:
|
||||
secret_changes.delete(store, worker_id)
|
||||
for field in _SERVERLESS_ONLY_FIELDS:
|
||||
existing.pop(field, None)
|
||||
elif _contains_secret_fields(data):
|
||||
return _error_response("API keys are only accepted for serverless workers", 400)
|
||||
_apply_field_patch(existing, data, _WORKER_FIELDS)
|
||||
result = {"status": "success"}
|
||||
elif all(key in data for key in ["name", "port", "cuda_device"]):
|
||||
if _contains_secret_fields(data):
|
||||
return _error_response("API keys are only accepted for serverless workers", 400)
|
||||
new_worker = {
|
||||
"id": worker_id,
|
||||
"name": data["name"],
|
||||
"host": normalize_host(data.get("host", "localhost")),
|
||||
"port": data["port"],
|
||||
"cuda_device": data["cuda_device"],
|
||||
"enabled": data.get("enabled", False),
|
||||
"extra_args": data.get("extra_args", ""),
|
||||
"type": data.get("type", "local")
|
||||
}
|
||||
workers.append(new_worker)
|
||||
result = {"status": "success"}
|
||||
else:
|
||||
return _error_response(
|
||||
"Worker not found and missing required fields for creation", 404
|
||||
)
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except ValueError:
|
||||
return _error_response("Invalid worker configuration", 400)
|
||||
except Exception:
|
||||
return _error_response("Unable to update worker configuration", 500)
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
||||
async def delete_worker_endpoint(request):
|
||||
return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
|
||||
|
||||
|
||||
async def _delete_worker_endpoint(request, secret_changes):
|
||||
store = None
|
||||
worker_id = None
|
||||
try:
|
||||
data = await request.json()
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception:
|
||||
return _error_response("Invalid JSON payload", 400)
|
||||
if not isinstance(data, dict):
|
||||
return _error_response("Worker payload must be an object", 400)
|
||||
worker_id = data.get("worker_id")
|
||||
|
||||
if worker_id is None:
|
||||
return await handle_api_error(request, "Missing worker_id", 400)
|
||||
|
||||
return _error_response("Missing worker_id", 400)
|
||||
|
||||
result = None
|
||||
async with config_transaction() as config:
|
||||
workers = config.get("workers", [])
|
||||
|
||||
# Find and remove the worker
|
||||
worker_index = -1
|
||||
for i, worker in enumerate(workers):
|
||||
if worker["id"] == worker_id:
|
||||
worker_index = i
|
||||
break
|
||||
|
||||
if worker_index == -1:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
# Remove the worker
|
||||
removed_worker = workers.pop(worker_index)
|
||||
|
||||
return web.json_response({
|
||||
removed = _worker_by_id(workers, worker_id)
|
||||
if removed is None:
|
||||
return _error_response("Worker not found", 404)
|
||||
if _worker_type(removed) == "serverless":
|
||||
_ensure_pool_idle(worker_id)
|
||||
store = _serverless_secret_store()
|
||||
secret_worker_id = str(worker_id)
|
||||
prior_secret = secret_changes.capture(store, secret_worker_id)
|
||||
if prior_secret is not None:
|
||||
secret_changes.delete(store, secret_worker_id)
|
||||
workers.remove(removed)
|
||||
result = {
|
||||
"status": "success",
|
||||
"message": f"Worker {removed_worker.get('name', worker_id)} deleted"
|
||||
})
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 400)
|
||||
"message": "Worker deleted" if _worker_type(removed) == "serverless"
|
||||
else f"Worker {removed.get('name', worker_id)} deleted",
|
||||
}
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to delete worker", 500)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
|
||||
async def update_setting_endpoint(request):
|
||||
|
||||
@@ -55,6 +55,11 @@ async def worker_ws_is_active(worker):
|
||||
|
||||
async def _probe_worker_active(worker, use_websocket, semaphore):
|
||||
async with semaphore:
|
||||
if worker.get("type") == "serverless":
|
||||
from .serverless_runtime import get_manager
|
||||
# Authentication/configuration failures must surface, not silently skip paid work.
|
||||
status = await get_manager().health(worker)
|
||||
return worker, status["ok"]
|
||||
is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
|
||||
return worker, is_active
|
||||
|
||||
@@ -203,6 +208,10 @@ def _extract_queue_remaining(payload):
|
||||
|
||||
async def _probe_worker_queue(worker, semaphore, probe_timeout):
|
||||
async with semaphore:
|
||||
if worker.get("type") == "serverless":
|
||||
from .serverless_runtime import get_manager
|
||||
status = await get_manager().health(worker)
|
||||
return {"worker": worker, "queue_remaining": status["queueRemaining"]}
|
||||
worker_url = build_worker_url(worker)
|
||||
payload = await probe_worker(worker_url, timeout=probe_timeout)
|
||||
if payload is None:
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Small Runpod queue API client. Never retries a potentially billable POST."""
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
|
||||
import aiohttp
|
||||
|
||||
API_BASE = "https://api.runpod.ai/v2"
|
||||
MAX_REQUEST_BYTES = 10_000_000
|
||||
_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z")
|
||||
|
||||
|
||||
class RunpodError(RuntimeError):
|
||||
"""Sanitized provider failure safe to display in the panel."""
|
||||
|
||||
|
||||
class AmbiguousSubmission(RunpodError):
|
||||
"""Submission may have been accepted: never retry it automatically."""
|
||||
|
||||
|
||||
def serialize_request(workflow, images, timeout_seconds):
|
||||
"""Serialize and size-check the exact JSON body used by Runpod ``/run``."""
|
||||
payload = {
|
||||
"input": {"workflow": workflow, "images": images},
|
||||
"policy": {
|
||||
"executionTimeout": timeout_seconds * 1000,
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
body = json.dumps(
|
||||
payload,
|
||||
allow_nan=False,
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
size = len(body)
|
||||
if size > MAX_REQUEST_BYTES:
|
||||
raise ValueError(
|
||||
f"Serialized serverless request body is {size:,} bytes; Runpod allows at most "
|
||||
f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw "
|
||||
"image bytes; reduce image or workflow size."
|
||||
)
|
||||
return body
|
||||
|
||||
|
||||
def checked_id(value):
|
||||
if not isinstance(value, str) or not _ID.fullmatch(value):
|
||||
raise ValueError("Invalid Runpod endpoint or job ID")
|
||||
return value
|
||||
|
||||
|
||||
class RunpodClient:
|
||||
def __init__(self, session, endpoint_id, api_key):
|
||||
self.session = session
|
||||
self.endpoint_id = checked_id(endpoint_id)
|
||||
if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"):
|
||||
raise ValueError("A valid Runpod API key is required")
|
||||
self._api_key = api_key
|
||||
|
||||
async def _request(self, method, path, payload=None, request_body=None):
|
||||
# Base origin is code-owned, never a user-supplied URL. Do not follow redirects
|
||||
# with credentials, or include response bodies/transport errors in exceptions.
|
||||
url = f"{API_BASE}/{self.endpoint_id}/{path}"
|
||||
attempts = 3 if method == "GET" else 1
|
||||
for attempt in range(attempts):
|
||||
try:
|
||||
request_kwargs = {
|
||||
"headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"},
|
||||
"timeout": aiohttp.ClientTimeout(total=30),
|
||||
"allow_redirects": False,
|
||||
}
|
||||
if request_body is None:
|
||||
request_kwargs["json"] = payload
|
||||
else:
|
||||
request_kwargs["data"] = request_body
|
||||
async with self.session.request(
|
||||
method, url,
|
||||
**request_kwargs,
|
||||
) as response:
|
||||
status = response.status
|
||||
if status != 200:
|
||||
if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts:
|
||||
await asyncio.sleep(0.5 * (2 ** attempt))
|
||||
continue
|
||||
if method == "POST" and path == "run" and status >= 500:
|
||||
raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.")
|
||||
raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})")
|
||||
try:
|
||||
result = await response.json()
|
||||
except (ValueError, aiohttp.ContentTypeError):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None
|
||||
raise RunpodError("Runpod returned an invalid JSON response") from None
|
||||
if not isinstance(result, dict):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.")
|
||||
raise RunpodError("Runpod returned an invalid response")
|
||||
return result
|
||||
except (asyncio.TimeoutError, aiohttp.ClientError, OSError):
|
||||
if method == "POST" and path == "run":
|
||||
raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None
|
||||
if method == "GET" and attempt + 1 < attempts:
|
||||
await asyncio.sleep(0.5 * (2 ** attempt))
|
||||
continue
|
||||
raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None
|
||||
raise RunpodError("Runpod request failed")
|
||||
|
||||
async def health(self):
|
||||
return await self._request("GET", "health")
|
||||
|
||||
async def submit(self, workflow, images, timeout_seconds):
|
||||
body = serialize_request(workflow, images, timeout_seconds)
|
||||
result = await self._request("POST", "run", request_body=body)
|
||||
try:
|
||||
return checked_id(result.get("id"))
|
||||
except ValueError:
|
||||
raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None
|
||||
|
||||
async def status(self, job_id):
|
||||
return await self._request("GET", f"status/{checked_id(job_id)}")
|
||||
|
||||
async def cancel(self, job_id):
|
||||
return await self._request("POST", f"cancel/{checked_id(job_id)}")
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Private, atomic job-ID journal. Contains identifiers/status only, never payloads."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import threading
|
||||
|
||||
from .serverless_state import TERMINAL
|
||||
FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"}
|
||||
|
||||
|
||||
class JobLedger:
|
||||
def __init__(self, path):
|
||||
self.path = Path(path)
|
||||
self._lock = threading.RLock()
|
||||
|
||||
def _read(self):
|
||||
if self.path.is_symlink():
|
||||
raise RuntimeError("Unsafe serverless job journal path")
|
||||
if not self.path.exists():
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
if not isinstance(data, dict) or any(not isinstance(v, dict) or set(v) - FIELDS for v in data.values()):
|
||||
raise ValueError()
|
||||
return data
|
||||
except (ValueError, OSError):
|
||||
raise RuntimeError("Serverless job journal is unreadable; preserve it for recovery before submitting more work") from None
|
||||
|
||||
def all(self):
|
||||
with self._lock:
|
||||
return list(self._read().values())
|
||||
|
||||
def update(self, record_id, **fields):
|
||||
if set(fields) - FIELDS:
|
||||
raise ValueError("Unexpected serverless journal field")
|
||||
with self._lock:
|
||||
data = self._read()
|
||||
data.setdefault(record_id, {"record_id": record_id}).update(fields)
|
||||
# Keep every unresolved job, plus a bounded terminal history.
|
||||
completed = sorted((k for k, v in data.items() if v.get("status") in TERMINAL), key=lambda k: data[k].get("updated_at", 0))
|
||||
for key in completed[:-200]:
|
||||
data.pop(key)
|
||||
if self.path.parent.is_symlink():
|
||||
raise RuntimeError("Unsafe serverless journal directory")
|
||||
self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
fd, temporary = tempfile.mkstemp(prefix=".jobs-", dir=self.path.parent)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||
json.dump(data, stream, allow_nan=False)
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
os.replace(temporary, self.path)
|
||||
finally:
|
||||
if os.path.exists(temporary):
|
||||
os.unlink(temporary)
|
||||
return dict(data[record_id])
|
||||
@@ -0,0 +1,601 @@
|
||||
"""Bounded image transport for generic Runpod-backed ComfyUI jobs."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import io
|
||||
import ipaddress
|
||||
import os
|
||||
import re
|
||||
import socket
|
||||
import stat
|
||||
import warnings
|
||||
from pathlib import Path, PurePosixPath
|
||||
from urllib.parse import urljoin, urlsplit
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
|
||||
|
||||
MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024
|
||||
MAX_INPUT_TOTAL_BYTES = 10_000_000
|
||||
MAX_OUTPUT_FILES = 1024
|
||||
MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024
|
||||
MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024
|
||||
MAX_IMAGE_PIXELS = 40_000_000
|
||||
MAX_OUTPUT_ARRAY_BYTES = 256 * 1024 * 1024
|
||||
MAX_REDIRECTS = 3
|
||||
_DOWNLOAD_TIMEOUT_SECONDS = 60
|
||||
_CHUNK_SIZE = 64 * 1024
|
||||
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
_UNSUPPORTED_IMAGE_EXTENSIONS = {".gif", ".tif", ".tiff", ".avif", ".heic", ".ico", ".svg"}
|
||||
_AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus", ".aiff", ".aif", ".wma"}
|
||||
_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
|
||||
_MEDIA_EXTENSIONS = _IMAGE_EXTENSIONS | _UNSUPPORTED_IMAGE_EXTENSIONS | _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS
|
||||
_MEDIA_ANNOTATION = re.compile(r"\s+\[([A-Za-z0-9_-]{1,64})\]$")
|
||||
_PREFIX_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\Z", re.ASCII)
|
||||
_SAFE_FILENAME_RE = re.compile(r"[A-Za-z0-9_.-]{1,255}\Z", re.ASCII)
|
||||
|
||||
|
||||
class ServerlessMediaError(ValueError):
|
||||
"""Sanitized input, artifact, or download validation failure."""
|
||||
|
||||
|
||||
def _folder_roots(folder_paths):
|
||||
roots = {}
|
||||
for kind, getter_name in (
|
||||
("input", "get_input_directory"),
|
||||
("output", "get_output_directory"),
|
||||
("temp", "get_temp_directory"),
|
||||
):
|
||||
getter = getattr(folder_paths, getter_name, None)
|
||||
if not callable(getter):
|
||||
raise ServerlessMediaError("ComfyUI media directories are unavailable")
|
||||
try:
|
||||
roots[kind] = Path(getter()).resolve(strict=True)
|
||||
except Exception:
|
||||
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
|
||||
return roots
|
||||
|
||||
|
||||
def _inside_roots(path, roots):
|
||||
for root in roots.values():
|
||||
try:
|
||||
path.relative_to(root)
|
||||
return True
|
||||
except ValueError:
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _normalize_media_reference(value):
|
||||
if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value:
|
||||
raise ServerlessMediaError("Image input reference is invalid")
|
||||
value = value.replace("\\", "/")
|
||||
explicit_root = None
|
||||
annotation = _MEDIA_ANNOTATION.search(value)
|
||||
if annotation:
|
||||
marker = annotation.group(1).lower()
|
||||
if marker not in {"input", "output", "temp"}:
|
||||
raise ServerlessMediaError("Image input annotation is unsupported")
|
||||
explicit_root = marker
|
||||
value = value[: annotation.start()]
|
||||
if "://" in value:
|
||||
raise ServerlessMediaError("Remote media references are not supported as image inputs")
|
||||
if re.match(r"^[A-Za-z]:", value) or value.startswith("//"):
|
||||
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
|
||||
|
||||
path = PurePosixPath(value)
|
||||
if any(part == ".." for part in path.parts):
|
||||
raise ServerlessMediaError("Image input path traversal is not allowed")
|
||||
if not path.parts or value in {".", ""}:
|
||||
raise ServerlessMediaError("Image input reference is invalid")
|
||||
|
||||
if explicit_root:
|
||||
value = f"{value} [{explicit_root}]"
|
||||
suffix = Path(path.name).suffix.lower()
|
||||
return value, suffix
|
||||
|
||||
|
||||
def _resolve_input_path(value, suffix, folder_paths, roots):
|
||||
if suffix not in _IMAGE_EXTENSIONS:
|
||||
if suffix in _MEDIA_EXTENSIONS:
|
||||
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
|
||||
raise ServerlessMediaError("Unsupported serverless image file type")
|
||||
|
||||
try:
|
||||
source = Path(value)
|
||||
if source.is_absolute():
|
||||
annotation = _MEDIA_ANNOTATION.search(value)
|
||||
unannotated = value[:annotation.start()] if annotation else value
|
||||
resolved = Path(unannotated).resolve(strict=True)
|
||||
else:
|
||||
resolver = getattr(folder_paths, "get_annotated_filepath", None)
|
||||
if not callable(resolver):
|
||||
raise ServerlessMediaError("ComfyUI media path resolver is unavailable")
|
||||
resolved = Path(resolver(value)).resolve(strict=True)
|
||||
except ServerlessMediaError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ServerlessMediaError("Image input file is missing or invalid") from None
|
||||
|
||||
if not _inside_roots(resolved, roots):
|
||||
raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
|
||||
try:
|
||||
mode = resolved.stat().st_mode
|
||||
except OSError:
|
||||
raise ServerlessMediaError("Image input file is missing or invalid") from None
|
||||
if not stat.S_ISREG(mode):
|
||||
raise ServerlessMediaError("Image input is not a regular file")
|
||||
return resolved
|
||||
|
||||
|
||||
def _iter_image_references(workflow):
|
||||
if not isinstance(workflow, dict):
|
||||
raise ServerlessMediaError("Serverless workflow must be an API graph")
|
||||
references = []
|
||||
for node_id, node in workflow.items():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
class_type = node.get("class_type")
|
||||
inputs = node.get("inputs", {})
|
||||
if class_type == "LoadImage" and not isinstance(inputs, dict):
|
||||
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
|
||||
if not isinstance(inputs, dict):
|
||||
continue
|
||||
if class_type == "LoadImage":
|
||||
value = inputs.get("image")
|
||||
if not isinstance(value, str):
|
||||
raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
|
||||
references.append((node_id, inputs, "image", value, True))
|
||||
|
||||
for key in ("image", "file"):
|
||||
if key == "image" and class_type == "LoadImage":
|
||||
continue
|
||||
value = inputs.get(key)
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
normalized, suffix = _normalize_media_reference(value)
|
||||
if suffix in _IMAGE_EXTENSIONS:
|
||||
references.append((node_id, inputs, key, normalized, False))
|
||||
elif suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
|
||||
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
|
||||
elif suffix in _UNSUPPORTED_IMAGE_EXTENSIONS or (key == "image" and suffix):
|
||||
raise ServerlessMediaError("Unsupported serverless image file type")
|
||||
|
||||
for key in ("audio", "video"):
|
||||
value = inputs.get(key)
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
_normalized, suffix = _normalize_media_reference(value)
|
||||
if suffix in _MEDIA_EXTENSIONS:
|
||||
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
|
||||
return references
|
||||
|
||||
|
||||
def _read_limited(path, expected_size):
|
||||
try:
|
||||
descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
|
||||
with os.fdopen(descriptor, "rb") as handle:
|
||||
info = os.fstat(handle.fileno())
|
||||
if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES:
|
||||
raise ServerlessMediaError("Image input exceeds the permitted size")
|
||||
if info.st_size > MAX_INPUT_TOTAL_BYTES:
|
||||
raise ServerlessMediaError(
|
||||
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
|
||||
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
|
||||
)
|
||||
data = handle.read(MAX_INPUT_TOTAL_BYTES + 1)
|
||||
except ServerlessMediaError:
|
||||
raise
|
||||
except OSError:
|
||||
raise ServerlessMediaError("Image input file could not be read safely") from None
|
||||
if len(data) > MAX_INPUT_TOTAL_BYTES:
|
||||
raise ServerlessMediaError(
|
||||
f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
|
||||
"base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
|
||||
)
|
||||
if len(data) != expected_size:
|
||||
raise ServerlessMediaError("Image input changed while being read")
|
||||
return data
|
||||
|
||||
|
||||
async def prepare_images(workflow: dict) -> list[dict]:
|
||||
"""Pack referenced local images and rewrite only their graph input references.
|
||||
|
||||
``workflow`` is expected to be a private API-graph copy owned by the caller.
|
||||
Returned ``name`` values are flat generated names suitable for ``input.images``.
|
||||
"""
|
||||
references = _iter_image_references(workflow)
|
||||
if not references:
|
||||
return []
|
||||
|
||||
# Import only in ComfyUI, keeping this module importable in the serverless client.
|
||||
try:
|
||||
import folder_paths
|
||||
except Exception:
|
||||
raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
|
||||
|
||||
roots = _folder_roots(folder_paths)
|
||||
resolved_refs = []
|
||||
unique_paths = {}
|
||||
total_source_bytes = 0
|
||||
for node_id, inputs, key, value, _is_load_image in references:
|
||||
normalized, suffix = _normalize_media_reference(value)
|
||||
if key == "image" and suffix not in _IMAGE_EXTENSIONS:
|
||||
if suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
|
||||
raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
|
||||
raise ServerlessMediaError("Unsupported serverless image file type")
|
||||
path = _resolve_input_path(normalized, suffix, folder_paths, roots)
|
||||
if path not in unique_paths:
|
||||
try:
|
||||
size = path.stat().st_size
|
||||
except OSError:
|
||||
raise ServerlessMediaError("Image input file is missing or invalid") from None
|
||||
if size < 0 or size > MAX_INPUT_IMAGE_BYTES:
|
||||
raise ServerlessMediaError("Image input exceeds the permitted size")
|
||||
total_source_bytes += size
|
||||
if total_source_bytes > MAX_INPUT_TOTAL_BYTES:
|
||||
raise ServerlessMediaError(
|
||||
f"Total raw serverless image input size is {total_source_bytes:,} bytes; "
|
||||
f"the input read guard is {MAX_INPUT_TOTAL_BYTES:,} bytes. Base64 encoding "
|
||||
"adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
|
||||
)
|
||||
unique_paths[path] = {"size": size, "suffix": suffix}
|
||||
resolved_refs.append((inputs, key, path))
|
||||
|
||||
packed_by_name = {}
|
||||
path_to_name = {}
|
||||
for path, metadata in unique_paths.items():
|
||||
data = _read_limited(path, metadata["size"])
|
||||
digest = hashlib.sha256(data).hexdigest()
|
||||
name = f"serverless_{digest}{metadata['suffix']}"
|
||||
path_to_name[path] = name
|
||||
packed_by_name.setdefault(name, {"name": name, "image": base64.b64encode(data).decode("ascii")})
|
||||
|
||||
for inputs, key, path in resolved_refs:
|
||||
inputs[key] = path_to_name[path]
|
||||
return [packed_by_name[name] for name in sorted(packed_by_name)]
|
||||
|
||||
|
||||
def _validate_artifact_url(url):
|
||||
if not isinstance(url, str) or not url or any(ord(char) < 0x20 or char.isspace() for char in url):
|
||||
raise ServerlessMediaError("Serverless image artifact URL is invalid")
|
||||
try:
|
||||
parsed = urlsplit(url)
|
||||
port = parsed.port
|
||||
except ValueError:
|
||||
raise ServerlessMediaError("Serverless image artifact URL is invalid") from None
|
||||
if (
|
||||
parsed.scheme.lower() != "https"
|
||||
or not parsed.hostname
|
||||
or parsed.username is not None
|
||||
or parsed.password is not None
|
||||
or port not in (None, 443)
|
||||
or parsed.fragment
|
||||
):
|
||||
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
|
||||
host = parsed.hostname.rstrip(".").lower()
|
||||
if host in {"localhost", "metadata", "metadata.google.internal"} or host.endswith((".localhost", ".local", ".internal")):
|
||||
raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
|
||||
try:
|
||||
address = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
return url
|
||||
if not address.is_global:
|
||||
raise ServerlessMediaError("Serverless image artifact URL must resolve to a public address")
|
||||
return url
|
||||
|
||||
|
||||
class _PublicResolver(aiohttp.abc.AbstractResolver):
|
||||
"""Resolve once, reject non-public answers, and pin the checked address."""
|
||||
|
||||
async def _getaddrinfo(self, host, port, family):
|
||||
return await asyncio.get_running_loop().getaddrinfo(
|
||||
host, port, family=family, type=socket.SOCK_STREAM
|
||||
)
|
||||
|
||||
async def resolve(self, host, port=0, family=socket.AF_INET):
|
||||
try:
|
||||
results = await self._getaddrinfo(host, port, family)
|
||||
except (OSError, asyncio.TimeoutError):
|
||||
raise OSError("Public artifact host could not be resolved") from None
|
||||
if not results:
|
||||
raise OSError("Public artifact host could not be resolved")
|
||||
|
||||
resolved = []
|
||||
seen = set()
|
||||
for result_family, _socktype, proto, _canonname, sockaddr in results:
|
||||
try:
|
||||
address = ipaddress.ip_address(sockaddr[0])
|
||||
except (ValueError, TypeError, IndexError):
|
||||
raise OSError("Public artifact host resolved to an invalid address") from None
|
||||
if not address.is_global:
|
||||
raise OSError("Public artifact host resolved to a non-public address")
|
||||
key = (str(address), result_family)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
resolved.append({
|
||||
"hostname": host,
|
||||
"host": str(address),
|
||||
"port": port,
|
||||
"family": result_family,
|
||||
"proto": proto,
|
||||
"flags": 0,
|
||||
})
|
||||
if not resolved:
|
||||
raise OSError("Public artifact host could not be resolved")
|
||||
return resolved
|
||||
|
||||
async def close(self):
|
||||
return None
|
||||
|
||||
|
||||
async def download_public_artifact(url: str, max_bytes: int = MAX_OUTPUT_IMAGE_BYTES) -> bytes:
|
||||
"""Download a signed public HTTPS artifact without provider credentials."""
|
||||
if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or not 0 < max_bytes <= MAX_OUTPUT_IMAGE_BYTES:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
current_url = _validate_artifact_url(url)
|
||||
resolver = _PublicResolver()
|
||||
connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False, limit=4)
|
||||
timeout = aiohttp.ClientTimeout(
|
||||
total=_DOWNLOAD_TIMEOUT_SECONDS,
|
||||
connect=10,
|
||||
sock_read=15,
|
||||
)
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=timeout,
|
||||
trust_env=False,
|
||||
cookie_jar=aiohttp.DummyCookieJar(),
|
||||
) as session:
|
||||
for redirect_count in range(MAX_REDIRECTS + 1):
|
||||
_validate_artifact_url(current_url)
|
||||
async with session.get(current_url, allow_redirects=False) as response:
|
||||
if response.status in {301, 302, 303, 307, 308}:
|
||||
location = response.headers.get("Location")
|
||||
if not location or redirect_count >= MAX_REDIRECTS:
|
||||
raise ServerlessMediaError("Serverless image artifact redirect is invalid")
|
||||
current_url = _validate_artifact_url(urljoin(current_url, location))
|
||||
continue
|
||||
if response.status != 200:
|
||||
raise ServerlessMediaError("Serverless image artifact download failed")
|
||||
payload = bytearray()
|
||||
async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
|
||||
if not isinstance(chunk, bytes):
|
||||
raise ServerlessMediaError("Serverless image artifact response is invalid")
|
||||
if len(payload) + len(chunk) > max_bytes:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
payload.extend(chunk)
|
||||
if not payload:
|
||||
raise ServerlessMediaError("Serverless image artifact is empty")
|
||||
return bytes(payload)
|
||||
except ServerlessMediaError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ServerlessMediaError("Serverless image artifact download failed") from None
|
||||
raise ServerlessMediaError("Serverless image artifact download failed")
|
||||
|
||||
|
||||
def _validated_manifest(manifest):
|
||||
if not isinstance(manifest, list) or len(manifest) > MAX_OUTPUT_FILES:
|
||||
raise ServerlessMediaError("Serverless image manifest is invalid")
|
||||
normalized = []
|
||||
seen_prefixes = set()
|
||||
prefixes = []
|
||||
for item in manifest:
|
||||
if not isinstance(item, dict):
|
||||
raise ServerlessMediaError("Serverless image manifest is invalid")
|
||||
if item.get("kind") != "image":
|
||||
raise ServerlessMediaError("Only image collectors are supported by serverless media transport")
|
||||
for key in ("collector_id", "job_id", "node_id"):
|
||||
value = item.get(key)
|
||||
if not isinstance(value, (str, int)) or isinstance(value, bool) or not str(value):
|
||||
raise ServerlessMediaError("Serverless image manifest is invalid")
|
||||
prefix = item.get("prefix")
|
||||
if not isinstance(prefix, str) or not _PREFIX_RE.fullmatch(prefix) or prefix in {".", ".."}:
|
||||
raise ServerlessMediaError("Serverless image manifest contains an unsafe filename prefix")
|
||||
folded = prefix.casefold()
|
||||
if folded in seen_prefixes:
|
||||
raise ServerlessMediaError("Serverless image manifest prefixes must be unique")
|
||||
seen_prefixes.add(folded)
|
||||
prefixes.append(prefix)
|
||||
normalized.append({
|
||||
"collector_id": str(item["collector_id"]),
|
||||
"job_id": str(item["job_id"]),
|
||||
"node_id": str(item["node_id"]),
|
||||
"prefix": prefix,
|
||||
})
|
||||
for first in prefixes:
|
||||
for second in prefixes:
|
||||
if first != second and second.casefold().startswith(first.casefold() + "_"):
|
||||
raise ServerlessMediaError("Serverless image manifest prefixes are ambiguous")
|
||||
return normalized
|
||||
|
||||
|
||||
def _artifact_filename_matches(filename, prefixes):
|
||||
if not isinstance(filename, str):
|
||||
return None
|
||||
basename = filename.replace("\\", "/").rsplit("/", 1)[-1]
|
||||
matches = [item for item in prefixes if basename.startswith(item["prefix"] + "_")]
|
||||
if len(matches) > 1:
|
||||
raise ServerlessMediaError("Serverless image artifact matches multiple collectors")
|
||||
return matches[0] if matches else None
|
||||
|
||||
|
||||
def _canonical_artifact_identity(entry, filename):
|
||||
artifact_path = entry.get("artifact_path")
|
||||
artifact_type = entry.get("artifact_type", "output")
|
||||
if artifact_path is not None:
|
||||
if not isinstance(artifact_path, str) or not artifact_path or "\x00" in artifact_path:
|
||||
raise ServerlessMediaError("Serverless image artifact identity is invalid")
|
||||
normalized_path = os.path.normcase(os.path.normpath(artifact_path.replace("\\", "/")))
|
||||
else:
|
||||
subfolder = entry.get("subfolder", "")
|
||||
if not isinstance(subfolder, str) or "\x00" in subfolder:
|
||||
raise ServerlessMediaError("Serverless image artifact identity is invalid")
|
||||
normalized_path = os.path.normcase(os.path.normpath(f"{subfolder}/{filename}"))
|
||||
return str(artifact_type), normalized_path
|
||||
|
||||
|
||||
def _decode_base64_data(value, max_bytes=MAX_OUTPUT_IMAGE_BYTES):
|
||||
if not isinstance(value, str) or not value:
|
||||
raise ServerlessMediaError("Serverless image artifact data is invalid")
|
||||
max_bytes = min(max_bytes, MAX_OUTPUT_IMAGE_BYTES)
|
||||
max_encoded = ((max_bytes + 2) // 3) * 4
|
||||
if len(value) > max_encoded:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
padding = len(value) - len(value.rstrip("="))
|
||||
estimated_size = (len(value) // 4) * 3 - padding
|
||||
if estimated_size > max_bytes:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
try:
|
||||
payload = base64.b64decode(value, validate=True)
|
||||
except (binascii.Error, ValueError):
|
||||
raise ServerlessMediaError("Serverless image artifact data is invalid") from None
|
||||
if not payload:
|
||||
raise ServerlessMediaError("Serverless image artifact is empty")
|
||||
if len(payload) > max_bytes:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
return payload
|
||||
|
||||
|
||||
def _decode_image_array(payload, total_array_bytes):
|
||||
try:
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("error", Image.DecompressionBombWarning)
|
||||
image = Image.open(io.BytesIO(payload))
|
||||
width, height = image.size
|
||||
if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
|
||||
raise ServerlessMediaError("Serverless image dimensions exceed the permitted limit")
|
||||
has_alpha = image.mode in {"RGBA", "LA"} or "transparency" in image.info
|
||||
mode = "RGBA" if has_alpha else "RGB"
|
||||
output_bytes = width * height * (4 if has_alpha else 3) * 4
|
||||
if total_array_bytes + output_bytes > MAX_OUTPUT_ARRAY_BYTES:
|
||||
raise ServerlessMediaError("Decoded serverless images exceed the permitted memory limit")
|
||||
converted = image.convert(mode)
|
||||
converted.load()
|
||||
array = np.array(converted, dtype=np.float32, copy=True)
|
||||
np.multiply(array, 1.0 / 255.0, out=array)
|
||||
return array, output_bytes
|
||||
except ServerlessMediaError:
|
||||
raise
|
||||
except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
|
||||
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
|
||||
except Exception:
|
||||
raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
|
||||
|
||||
|
||||
def _image_output(output):
|
||||
"""Unwrap handler output and reject explicit execution errors."""
|
||||
if not isinstance(output, dict):
|
||||
raise ServerlessMediaError("Serverless job output is invalid")
|
||||
nested = output.get("output")
|
||||
data = nested if isinstance(nested, dict) else output
|
||||
for candidate in (output, data):
|
||||
if candidate.get("error") or candidate.get("errors"):
|
||||
raise ServerlessMediaError("Serverless job reported an execution error")
|
||||
|
||||
return data
|
||||
|
||||
|
||||
async def _select_image_payloads(data, expected, download):
|
||||
"""Match artifact identities and fetch bounded bytes, without image decoding."""
|
||||
artifacts = data.get("images", [])
|
||||
if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES:
|
||||
raise ServerlessMediaError("Serverless job returned too many image files")
|
||||
|
||||
selected = {item["prefix"]: [] for item in expected}
|
||||
identities = set()
|
||||
total_payload_bytes = 0
|
||||
for entry in artifacts:
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
filename = entry.get("filename")
|
||||
manifest_entry = _artifact_filename_matches(filename, expected)
|
||||
if manifest_entry is None:
|
||||
continue
|
||||
normalized_filename = filename.replace("\\", "/")
|
||||
if "/" in normalized_filename or not _SAFE_FILENAME_RE.fullmatch(filename):
|
||||
raise ServerlessMediaError("Serverless image artifact filename is invalid")
|
||||
if Path(filename).suffix.lower() not in _IMAGE_EXTENSIONS:
|
||||
raise ServerlessMediaError("Serverless image artifact has an unsupported file type")
|
||||
if "artifact_type" in entry and entry["artifact_type"] != "output":
|
||||
raise ServerlessMediaError("Expected serverless image artifact was not an output image")
|
||||
identity = _canonical_artifact_identity(entry, filename)
|
||||
if identity in identities:
|
||||
raise ServerlessMediaError("Duplicate serverless image artifact identity")
|
||||
identities.add(identity)
|
||||
match = re.match(rf"^{re.escape(manifest_entry['prefix'])}_(\d+)(?:_|\.)", filename)
|
||||
if match is None:
|
||||
raise ServerlessMediaError("Expected serverless image filename counter is invalid")
|
||||
artifact_type = entry.get("type")
|
||||
remaining_bytes = MAX_OUTPUT_TOTAL_BYTES - total_payload_bytes
|
||||
if remaining_bytes <= 0:
|
||||
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
|
||||
if artifact_type == "base64":
|
||||
payload = _decode_base64_data(entry.get("data"), remaining_bytes)
|
||||
elif artifact_type in {"gcs_url", "s3_url", "url"}:
|
||||
url = entry.get("data") or entry.get("url")
|
||||
url = _validate_artifact_url(url)
|
||||
try:
|
||||
if download is None:
|
||||
payload = await download_public_artifact(
|
||||
url, max_bytes=min(MAX_OUTPUT_IMAGE_BYTES, remaining_bytes)
|
||||
)
|
||||
else:
|
||||
payload = await download(url)
|
||||
except ServerlessMediaError:
|
||||
raise
|
||||
except Exception:
|
||||
raise ServerlessMediaError("Serverless image artifact download failed") from None
|
||||
if not isinstance(payload, bytes) or not payload:
|
||||
raise ServerlessMediaError("Serverless image artifact response is invalid")
|
||||
if len(payload) > MAX_OUTPUT_IMAGE_BYTES:
|
||||
raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
|
||||
else:
|
||||
raise ServerlessMediaError("Serverless image artifact transport type is unsupported")
|
||||
total_payload_bytes += len(payload)
|
||||
if total_payload_bytes > MAX_OUTPUT_TOTAL_BYTES:
|
||||
raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
|
||||
selected[manifest_entry["prefix"]].append((int(match.group(1)), filename, identity, payload))
|
||||
|
||||
if any(not selected[item["prefix"]] for item in expected):
|
||||
raise ServerlessMediaError("Serverless job is missing an expected collector image")
|
||||
return selected
|
||||
|
||||
|
||||
def _decode_image_groups(selected, expected):
|
||||
"""Decode in stable counter order, bounding array memory across collectors."""
|
||||
result = {}
|
||||
total_array_bytes = 0
|
||||
for manifest_entry in expected:
|
||||
prefix = manifest_entry["prefix"]
|
||||
rows = sorted(selected[prefix], key=lambda row: (row[0], row[1], row[2]))
|
||||
arrays = []
|
||||
shape = None
|
||||
for _counter, _filename, _identity, payload in rows:
|
||||
array, added_bytes = _decode_image_array(payload, total_array_bytes)
|
||||
total_array_bytes += added_bytes
|
||||
if shape is None:
|
||||
shape = array.shape
|
||||
elif array.shape != shape:
|
||||
raise ServerlessMediaError("Images from one serverless collector have inconsistent dimensions")
|
||||
arrays.append(array)
|
||||
result.setdefault(manifest_entry["job_id"], []).extend(arrays)
|
||||
return result
|
||||
|
||||
|
||||
async def collect_images(output: dict, manifest: list, download=None) -> dict[str, list[np.ndarray]]:
|
||||
"""Validate and decode expected image artifacts, grouped by serverless job ID."""
|
||||
data = _image_output(output)
|
||||
expected = _validated_manifest(manifest)
|
||||
if not expected:
|
||||
return {}
|
||||
selected = await _select_image_payloads(data, expected, download)
|
||||
return _decode_image_groups(selected, expected)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""Compile prepared distributed prompts into worker-side ComfyUI image jobs.
|
||||
|
||||
This module is deliberately pure: it only copies and rewrites an API prompt
|
||||
mapping and never imports or executes ComfyUI nodes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
|
||||
|
||||
_SUPPORTED_DISTRIBUTED = {
|
||||
"DistributedCollector",
|
||||
"DistributedSeed",
|
||||
"DistributedValue",
|
||||
"DistributedModelName",
|
||||
}
|
||||
_REMOVED_SCALAR_NODES = {
|
||||
"DistributedSeed",
|
||||
"DistributedValue",
|
||||
"DistributedModelName",
|
||||
}
|
||||
_PRIMITIVE_NODES = {
|
||||
"PrimitiveBoolean": False,
|
||||
"PrimitiveFloat": 0.0,
|
||||
"PrimitiveInt": 0,
|
||||
"PrimitiveString": "",
|
||||
}
|
||||
_REMOVED_NODE_OUTPUTS = {
|
||||
"DistributedCollector": {0, 1},
|
||||
"DistributedSeed": {0},
|
||||
"DistributedValue": {0},
|
||||
"DistributedModelName": {0},
|
||||
}
|
||||
_SEED_DEFAULT = 1125899906842
|
||||
|
||||
|
||||
def _is_link(value):
|
||||
return isinstance(value, list) and len(value) == 2
|
||||
|
||||
|
||||
def _is_unsupported_distributed(class_type):
|
||||
if not isinstance(class_type, str) or class_type in _SUPPORTED_DISTRIBUTED:
|
||||
return False
|
||||
lowered = class_type.lower()
|
||||
return "distributed" in lowered or "batchdivider" in lowered
|
||||
|
||||
|
||||
def _safe_component(value):
|
||||
cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_-")
|
||||
return cleaned[:24] or "empty"
|
||||
|
||||
|
||||
def _filename_prefix(job_id, worker_id, collector_id, kind):
|
||||
identity = json.dumps(
|
||||
[job_id, worker_id, collector_id, kind],
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
parts = (
|
||||
"serverless",
|
||||
_safe_component(job_id),
|
||||
_safe_component(worker_id),
|
||||
_safe_component(collector_id),
|
||||
kind,
|
||||
digest,
|
||||
)
|
||||
return "_".join(parts)
|
||||
|
||||
|
||||
def _worker_index(worker_id):
|
||||
"""Mirror nodes.utilities worker-index parsing; None means fallback."""
|
||||
try:
|
||||
if worker_id.startswith("worker_"):
|
||||
return int(worker_id.split("_")[1])
|
||||
return int(worker_id)
|
||||
except (ValueError, IndexError):
|
||||
return None
|
||||
|
||||
|
||||
def _coerce(value, value_type):
|
||||
if value_type == "INT":
|
||||
return int(float(value))
|
||||
if value_type == "FLOAT":
|
||||
return float(value)
|
||||
return value
|
||||
|
||||
|
||||
def _coerce_safe(value, value_type):
|
||||
try:
|
||||
return _coerce(value, value_type)
|
||||
except (TypeError, ValueError):
|
||||
return value
|
||||
|
||||
|
||||
def _stringify_model_name(value):
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, (int, float, bool)):
|
||||
return str(value)
|
||||
try:
|
||||
return json.dumps(value, indent=4)
|
||||
except (TypeError, ValueError):
|
||||
return str(value)
|
||||
|
||||
|
||||
def _model_name_result(text):
|
||||
if isinstance(text, list):
|
||||
values = [_stringify_model_name(value) for value in text]
|
||||
return values[0] if len(values) == 1 else values
|
||||
return _stringify_model_name(text)
|
||||
|
||||
|
||||
def _validate_prepared_prompt(prepared_prompt):
|
||||
"""Validate graph shape, links, collectors, and supported topology."""
|
||||
if not isinstance(prepared_prompt, dict):
|
||||
raise ValueError("prepared prompt must be an API graph mapping")
|
||||
|
||||
source_nodes = {}
|
||||
for raw_id, node in prepared_prompt.items():
|
||||
node_id = str(raw_id)
|
||||
if node_id in source_nodes:
|
||||
raise ValueError("prompt has duplicate node IDs after string normalization")
|
||||
if not isinstance(node, dict):
|
||||
raise ValueError(f"prompt node {node_id!r} must be an object")
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str):
|
||||
raise ValueError(f"prompt node {node_id!r} is missing a class_type")
|
||||
if _is_unsupported_distributed(class_type):
|
||||
raise ValueError(f"unsupported distributed or batch node: {class_type}")
|
||||
inputs = node.get("inputs", {})
|
||||
if not isinstance(inputs, dict):
|
||||
raise ValueError(f"prompt node {node_id!r} inputs must be an object")
|
||||
source_nodes[node_id] = node
|
||||
|
||||
if not source_nodes:
|
||||
raise ValueError("prepared prompt contains no nodes")
|
||||
|
||||
# Validate every API link before any graph rewriting. Slots on the nodes
|
||||
# removed by this compiler have a known schema and are checked exactly.
|
||||
for node_id, node in source_nodes.items():
|
||||
for input_name, value in node.get("inputs", {}).items():
|
||||
if not _is_link(value):
|
||||
continue
|
||||
upstream_id, output_slot = value
|
||||
upstream_id = str(upstream_id)
|
||||
if upstream_id not in source_nodes:
|
||||
raise ValueError(
|
||||
f"invalid link at node {node_id!r} input {input_name!r}: "
|
||||
f"source {upstream_id!r} does not exist"
|
||||
)
|
||||
if isinstance(output_slot, bool) or not isinstance(output_slot, int) or output_slot < 0:
|
||||
raise ValueError(
|
||||
f"invalid link at node {node_id!r} input {input_name!r}: "
|
||||
"output slot must be a nonnegative integer"
|
||||
)
|
||||
upstream_class = source_nodes[upstream_id]["class_type"]
|
||||
valid_slots = _REMOVED_NODE_OUTPUTS.get(upstream_class)
|
||||
if valid_slots is not None and output_slot not in valid_slots:
|
||||
raise ValueError(
|
||||
f"invalid output slot {output_slot} for removed node "
|
||||
f"{upstream_class} {upstream_id!r}"
|
||||
)
|
||||
|
||||
collector_ids = sorted(
|
||||
(node_id for node_id, node in source_nodes.items() if node["class_type"] == "DistributedCollector"),
|
||||
key=str,
|
||||
)
|
||||
if not collector_ids:
|
||||
raise ValueError("prepared prompt contains no DistributedCollector nodes")
|
||||
|
||||
# Only image collectors are supported by the current generic handler.
|
||||
for collector_id in collector_ids:
|
||||
inputs = source_nodes[collector_id].get("inputs", {})
|
||||
if "audio" in inputs:
|
||||
raise ValueError("collector audio input is unsupported by this image-only compiler")
|
||||
if not _is_link(inputs.get("images")):
|
||||
raise ValueError(f"collector {collector_id!r} is missing a linked image input")
|
||||
source_id = str(inputs["images"][0])
|
||||
source_class = source_nodes[source_id]["class_type"]
|
||||
if source_class in _REMOVED_SCALAR_NODES:
|
||||
raise ValueError(f"collector {collector_id!r} image input must link to an image node")
|
||||
|
||||
# A collector cannot be used, directly or through processing, as another
|
||||
# collector's input: this first compiler handles disconnected jobs only.
|
||||
for collector_id in collector_ids:
|
||||
pending = [str(source_nodes[collector_id]["inputs"]["images"][0])]
|
||||
seen = set()
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in seen:
|
||||
continue
|
||||
seen.add(node_id)
|
||||
if source_nodes[node_id]["class_type"] == "DistributedCollector":
|
||||
raise ValueError("collector-to-collector dependence is unsupported")
|
||||
for value in source_nodes[node_id].get("inputs", {}).values():
|
||||
if _is_link(value):
|
||||
pending.append(str(value[0]))
|
||||
|
||||
return source_nodes, collector_ids
|
||||
|
||||
|
||||
def _make_scalar_resolver(source_nodes):
|
||||
"""Build pure resolvers for distributed scalar nodes and stock primitives."""
|
||||
# Resolve only stock primitive nodes. A visited set makes malformed
|
||||
# primitive-to-primitive cycles fail clearly instead of recursing forever.
|
||||
def resolve_primitive(node_id, visiting=None):
|
||||
node_id = str(node_id)
|
||||
node = source_nodes[node_id]
|
||||
class_type = node["class_type"]
|
||||
if class_type not in _PRIMITIVE_NODES:
|
||||
raise ValueError("linked value must be literal or supported PrimitiveInt/Float/String/Boolean")
|
||||
if visiting is None:
|
||||
visiting = set()
|
||||
if node_id in visiting:
|
||||
raise ValueError(f"cycle while resolving primitive value at node {node_id!r}")
|
||||
visiting.add(node_id)
|
||||
value = node.get("inputs", {}).get("value", _PRIMITIVE_NODES[class_type])
|
||||
if _is_link(value):
|
||||
parent_id, output_slot = str(value[0]), value[1]
|
||||
if output_slot != 0:
|
||||
raise ValueError(f"invalid output slot {output_slot} for primitive node {parent_id!r}")
|
||||
if source_nodes[parent_id]["class_type"] not in _PRIMITIVE_NODES:
|
||||
raise ValueError("linked primitive value must come from a stock Primitive node")
|
||||
result = resolve_primitive(parent_id, visiting)
|
||||
else:
|
||||
result = copy.deepcopy(value)
|
||||
visiting.remove(node_id)
|
||||
return result
|
||||
|
||||
def resolve_literal_or_primitive(value, context):
|
||||
if _is_link(value):
|
||||
source_id, output_slot = str(value[0]), value[1]
|
||||
if output_slot != 0 or source_nodes[source_id]["class_type"] not in _PRIMITIVE_NODES:
|
||||
raise ValueError(f"{context} must be a literal or supported Primitive link")
|
||||
try:
|
||||
return resolve_primitive(source_id)
|
||||
except ValueError as exc:
|
||||
if "cycle" in str(exc):
|
||||
raise
|
||||
raise ValueError(f"{context} must be a literal or supported Primitive link") from None
|
||||
return copy.deepcopy(value)
|
||||
|
||||
def resolve_hidden(node_id, name, default):
|
||||
value = source_nodes[node_id].get("inputs", {}).get(name, default)
|
||||
if _is_link(value):
|
||||
raise ValueError(f"{name} on distributed node {node_id!r} must be a literal")
|
||||
return copy.deepcopy(value)
|
||||
|
||||
resolved_distributed = {}
|
||||
|
||||
def resolve_distributed(node_id, output_slot):
|
||||
node_id = str(node_id)
|
||||
class_type = source_nodes[node_id]["class_type"]
|
||||
if output_slot != 0:
|
||||
raise ValueError(f"invalid output slot {output_slot} for {class_type} {node_id!r}")
|
||||
cache_key = (node_id, output_slot)
|
||||
if cache_key in resolved_distributed:
|
||||
return copy.deepcopy(resolved_distributed[cache_key])
|
||||
inputs = source_nodes[node_id].get("inputs", {})
|
||||
if class_type == "DistributedSeed":
|
||||
seed = resolve_literal_or_primitive(inputs.get("seed", _SEED_DEFAULT), f"seed on {node_id!r}")
|
||||
is_worker = resolve_hidden(node_id, "is_worker", False)
|
||||
worker_id = resolve_hidden(node_id, "worker_id", "")
|
||||
if is_worker:
|
||||
worker_index = _worker_index(worker_id)
|
||||
if worker_index is not None:
|
||||
seed = seed + worker_index + 1
|
||||
result = seed
|
||||
elif class_type == "DistributedValue":
|
||||
default_value = resolve_literal_or_primitive(
|
||||
inputs.get("default_value", ""), f"default_value on {node_id!r}"
|
||||
)
|
||||
worker_values = resolve_literal_or_primitive(
|
||||
inputs.get("worker_values", "{}"), f"worker_values on {node_id!r}"
|
||||
)
|
||||
is_worker = resolve_hidden(node_id, "is_worker", False)
|
||||
worker_id = resolve_hidden(node_id, "worker_id", "")
|
||||
try:
|
||||
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
|
||||
except json.JSONDecodeError:
|
||||
values = {}
|
||||
if not isinstance(values, dict):
|
||||
values = {}
|
||||
value_type = values.get("_type", "STRING")
|
||||
fallback = _coerce_safe(default_value, value_type)
|
||||
result = fallback
|
||||
if is_worker:
|
||||
worker_index = _worker_index(worker_id)
|
||||
if worker_index is not None:
|
||||
raw_value = values.get(str(worker_index + 1), "")
|
||||
if raw_value:
|
||||
try:
|
||||
result = _coerce(raw_value, value_type)
|
||||
except (ValueError, IndexError):
|
||||
result = fallback
|
||||
elif class_type == "DistributedModelName":
|
||||
text = resolve_literal_or_primitive(
|
||||
inputs.get("text", ""), f"DistributedModelName text on {node_id!r}"
|
||||
)
|
||||
result = _model_name_result(text)
|
||||
else:
|
||||
raise ValueError(f"unsupported distributed node: {class_type}")
|
||||
resolved_distributed[cache_key] = copy.deepcopy(result)
|
||||
return result
|
||||
|
||||
return resolve_distributed
|
||||
|
||||
|
||||
def _build_savers(source_nodes, collector_ids):
|
||||
"""Construct stable SaveImage nodes and their collector metadata."""
|
||||
# The prepared collector values define the artifact identity. The fallback
|
||||
# for absent hidden values is stable and retains the collector's job scope.
|
||||
collector_specs = []
|
||||
max_numeric_id = 0
|
||||
for node_id in source_nodes:
|
||||
try:
|
||||
max_numeric_id = max(max_numeric_id, int(node_id))
|
||||
except ValueError:
|
||||
pass
|
||||
next_numeric_id = max_numeric_id + 1
|
||||
for collector_id in collector_ids:
|
||||
inputs = source_nodes[collector_id].get("inputs", {})
|
||||
raw_job_id = inputs.get("multi_job_id", collector_id)
|
||||
if _is_link(raw_job_id):
|
||||
raise ValueError(f"collector {collector_id!r} multi_job_id must be a literal")
|
||||
job_id = str(collector_id if raw_job_id in (None, "") else raw_job_id)
|
||||
raw_worker_id = inputs.get("worker_id", "")
|
||||
if _is_link(raw_worker_id):
|
||||
raise ValueError(f"collector {collector_id!r} worker_id must be a literal")
|
||||
worker_id = "" if raw_worker_id is None else str(raw_worker_id)
|
||||
saver_id = str(next_numeric_id)
|
||||
next_numeric_id += 1
|
||||
prefix = _filename_prefix(job_id, worker_id, collector_id, "image")
|
||||
collector_specs.append(
|
||||
{
|
||||
"collector_id": collector_id,
|
||||
"job_id": job_id,
|
||||
"node_id": saver_id,
|
||||
"kind": "image",
|
||||
"prefix": prefix,
|
||||
"source_link": copy.deepcopy(inputs["images"]),
|
||||
}
|
||||
)
|
||||
|
||||
saver_nodes = {}
|
||||
for spec in collector_specs:
|
||||
saver_nodes[spec["node_id"]] = {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {
|
||||
"images": spec["source_link"],
|
||||
"filename_prefix": spec["prefix"],
|
||||
},
|
||||
"_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"},
|
||||
}
|
||||
|
||||
return collector_specs, saver_nodes
|
||||
|
||||
|
||||
def _prune_graph(source_nodes, collector_specs, saver_nodes, resolve_distributed):
|
||||
"""Keep saver ancestors, replacing distributed scalar links with literals."""
|
||||
all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()}
|
||||
all_nodes.update(copy.deepcopy(saver_nodes))
|
||||
needed = set()
|
||||
pending = [spec["node_id"] for spec in collector_specs]
|
||||
while pending:
|
||||
node_id = pending.pop()
|
||||
if node_id in needed:
|
||||
continue
|
||||
needed.add(node_id)
|
||||
node = all_nodes[node_id]
|
||||
if node["class_type"] == "DistributedCollector":
|
||||
raise ValueError("collector-to-collector dependence is unsupported")
|
||||
for input_name, value in list(node.get("inputs", {}).items()):
|
||||
if not _is_link(value):
|
||||
continue
|
||||
upstream_id, output_slot = str(value[0]), value[1]
|
||||
upstream_class = source_nodes[upstream_id]["class_type"]
|
||||
if upstream_class in _REMOVED_SCALAR_NODES:
|
||||
node["inputs"][input_name] = resolve_distributed(upstream_id, output_slot)
|
||||
else:
|
||||
pending.append(upstream_id)
|
||||
|
||||
compiled = {
|
||||
node_id: copy.deepcopy(node)
|
||||
for node_id, node in all_nodes.items()
|
||||
if node_id in needed
|
||||
}
|
||||
if any("distributed" in node["class_type"].lower() for node in compiled.values()):
|
||||
raise ValueError("compiled prompt still contains a distributed dependency")
|
||||
|
||||
return compiled
|
||||
|
||||
|
||||
def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]:
|
||||
"""Compile a prepared worker prompt to an image-only serverless graph.
|
||||
|
||||
Every DistributedCollector's original ``images`` input is connected to a
|
||||
new SaveImage node. Only those savers and their upstream dependencies are
|
||||
retained. DistributedSeed/Value/ModelName links are resolved as scalar
|
||||
literals; unsupported node execution is never attempted.
|
||||
"""
|
||||
source_nodes, collector_ids = _validate_prepared_prompt(prepared_prompt)
|
||||
resolve_distributed = _make_scalar_resolver(source_nodes)
|
||||
collector_specs, saver_nodes = _build_savers(source_nodes, collector_ids)
|
||||
compiled = _prune_graph(
|
||||
source_nodes,
|
||||
collector_specs,
|
||||
saver_nodes,
|
||||
resolve_distributed,
|
||||
)
|
||||
manifest = [
|
||||
{key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")}
|
||||
for spec in collector_specs
|
||||
]
|
||||
return compiled, manifest
|
||||
@@ -0,0 +1,375 @@
|
||||
"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
|
||||
from .serverless_ledger import JobLedger
|
||||
from .serverless_state import (
|
||||
PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
|
||||
occupies_capacity, terminal_provider_status,
|
||||
)
|
||||
|
||||
|
||||
async def decode_results(output, manifest):
|
||||
from .serverless_media import collect_images
|
||||
import torch
|
||||
images = await collect_images(output, manifest)
|
||||
return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
|
||||
|
||||
|
||||
def state_directory():
|
||||
import folder_paths
|
||||
return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
|
||||
|
||||
|
||||
def get_secret_store():
|
||||
from ...utils.serverless_config import SecretStore
|
||||
return SecretStore(state_directory() / "credentials.json")
|
||||
|
||||
|
||||
def get_manager():
|
||||
import server
|
||||
from ...utils.network import get_client_session
|
||||
ps = server.PromptServer.instance
|
||||
manager = getattr(ps, "distributed_serverless", None)
|
||||
if manager is None:
|
||||
manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
|
||||
ps.distributed_serverless = manager
|
||||
return manager
|
||||
|
||||
|
||||
class ServerlessManager:
|
||||
def __init__(self, prompt_server, secrets, ledger_path, session_factory):
|
||||
self.server = prompt_server
|
||||
self.secrets = secrets
|
||||
self.ledger = JobLedger(ledger_path)
|
||||
self.session_factory = session_factory
|
||||
self.tasks = {}
|
||||
self.runs = {}
|
||||
self.pool_limits = {}
|
||||
self.health_cache = {}
|
||||
self.health_locks = {}
|
||||
self.poll_seconds = 2.0
|
||||
self.capacity_changed = asyncio.Condition()
|
||||
self.recovery_lock = asyncio.Lock()
|
||||
self._recover_local_records()
|
||||
|
||||
def _recover_local_records(self):
|
||||
# WAITING is committed before admission; it cannot have reached Runpod.
|
||||
# SUBMITTING is an intent record, so a lost ID must remain reserved.
|
||||
for record in self.ledger.all():
|
||||
if record.get("status") == "WAITING" and not record.get("job_id"):
|
||||
self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time())
|
||||
elif record.get("status") == "SUBMITTING" and not record.get("job_id"):
|
||||
self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time())
|
||||
|
||||
def _task_is_active(self, record_id):
|
||||
task = self.tasks.get(record_id)
|
||||
return task is not None and not task.done()
|
||||
|
||||
def pool_is_active(self, worker_id):
|
||||
"""Protect both live local work and unresolved pre-restart provider work."""
|
||||
for rid, run in self.runs.items():
|
||||
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
|
||||
if str(pool_id) == str(worker_id) and self._task_is_active(rid):
|
||||
return True
|
||||
return any(belongs_to_pool(record, worker_id) and is_unresolved(record)
|
||||
for record in self.ledger.all())
|
||||
|
||||
async def _reserve_provider_capacity(self, run):
|
||||
pool_id = run["worker"].get("pool_id", run["worker"]["id"])
|
||||
async with self.capacity_changed:
|
||||
while True:
|
||||
self._check_master(run)
|
||||
outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record)
|
||||
for record in self.ledger.all())
|
||||
if outstanding < run["worker"]["concurrency"]:
|
||||
# Claim atomically before any await that could submit a billed job.
|
||||
self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time())
|
||||
return
|
||||
await self.capacity_changed.wait()
|
||||
|
||||
async def _notify_capacity(self):
|
||||
async with self.capacity_changed:
|
||||
self.capacity_changed.notify_all()
|
||||
|
||||
async def _client(self, worker):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
key = self.secrets.get(pool_id)
|
||||
return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
|
||||
|
||||
async def health(self, worker, *, fresh=False):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
cache_key = (pool_id, worker["endpoint_id"])
|
||||
lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
|
||||
async with lock:
|
||||
cached = self.health_cache.get(cache_key)
|
||||
if cached and not fresh and time.monotonic() - cached[0] < 10:
|
||||
return cached[1]
|
||||
client = await self._client(worker)
|
||||
data = await client.health()
|
||||
jobs = data.get("jobs", {})
|
||||
if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
|
||||
raise RunpodError("Runpod returned an invalid endpoint health record")
|
||||
try:
|
||||
pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
|
||||
except (ValueError, TypeError):
|
||||
raise RunpodError("Runpod returned invalid queue counts") from None
|
||||
result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
|
||||
self.health_cache[cache_key] = (time.monotonic(), result)
|
||||
return result
|
||||
|
||||
def _semaphore(self, worker):
|
||||
pool_id = worker.get("pool_id", worker["id"])
|
||||
limit = worker["concurrency"]
|
||||
current = self.pool_limits.get(pool_id)
|
||||
if current is None:
|
||||
current = (limit, asyncio.Semaphore(limit))
|
||||
self.pool_limits[pool_id] = current
|
||||
elif current[0] != limit:
|
||||
if self.pool_is_active(pool_id):
|
||||
raise ValueError("Wait for this serverless pool's current jobs before changing concurrency")
|
||||
current = (limit, asyncio.Semaphore(limit))
|
||||
self.pool_limits[pool_id] = current
|
||||
return current[1]
|
||||
|
||||
def ensure_capacity(self, count):
|
||||
if sum(not task.done() for task in self.tasks.values()) + count > 64:
|
||||
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
|
||||
|
||||
async def start(self, execution_id, prompt_id, prepared_jobs):
|
||||
self.ensure_capacity(len(prepared_jobs))
|
||||
new_runs = []
|
||||
# Journal every local slot before starting any billed submission.
|
||||
for prepared in prepared_jobs:
|
||||
worker = prepared["worker"]
|
||||
semaphore = self._semaphore(worker)
|
||||
record_id = uuid.uuid4().hex
|
||||
now = time.time()
|
||||
self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
|
||||
worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
|
||||
endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
|
||||
run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
|
||||
provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
|
||||
self.runs[record_id] = run
|
||||
new_runs.append((record_id, run, semaphore))
|
||||
for record_id, run, semaphore in new_runs:
|
||||
self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
|
||||
# Bound in-memory completed tracking as well as the on-disk journal.
|
||||
done = [rid for rid, task in self.tasks.items() if task.done()]
|
||||
for rid in done[:-200]:
|
||||
self.tasks.pop(rid, None)
|
||||
self.runs.pop(rid, None)
|
||||
|
||||
def is_pending(self, collector_job_id, participant_id):
|
||||
return any(not task.done() and run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"])
|
||||
for rid, task in self.tasks.items() for run in [self.runs[rid]])
|
||||
|
||||
def owns(self, collector_job_id, participant_id):
|
||||
return any(run["worker"]["id"] == participant_id and
|
||||
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
|
||||
|
||||
async def _deliver(self, run, images=None, error=None):
|
||||
async with self.server.distributed_jobs_lock:
|
||||
for binding in run["manifest"]:
|
||||
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if queue is None:
|
||||
continue # A cancelled/finished collector must never be resurrected.
|
||||
if error:
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
|
||||
else:
|
||||
batch = images[binding["job_id"]]
|
||||
for index, tensor in enumerate(batch):
|
||||
queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
|
||||
"tensor": tensor, "is_last": index == len(batch) - 1})
|
||||
|
||||
async def _cancel_provider(self, run):
|
||||
if run["provider_job_id"] and run["client"]:
|
||||
try:
|
||||
result = await run["client"].cancel(run["provider_job_id"])
|
||||
return terminal_provider_status(result)
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def _check_master(self, run):
|
||||
queue = getattr(self.server, "prompt_queue", None)
|
||||
if queue is None:
|
||||
return
|
||||
history = queue.get_history(run["prompt_id"])
|
||||
record = history.get(run["prompt_id"], {})
|
||||
if record.get("status", {}).get("status_str") == "error":
|
||||
raise RunpodError("Master workflow failed; serverless work stopped")
|
||||
if not record and hasattr(queue, "get_current_queue"):
|
||||
running, pending = queue.get_current_queue()
|
||||
if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
|
||||
raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
|
||||
|
||||
async def _run(self, run, semaphore):
|
||||
rid = run["record_id"]
|
||||
submitted = False
|
||||
provider_terminal = False
|
||||
try:
|
||||
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
|
||||
async with semaphore:
|
||||
await self._reserve_provider_capacity(run)
|
||||
client = await self._client(run["worker"])
|
||||
run["client"] = client
|
||||
submitted = True
|
||||
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
||||
run["provider_job_id"] = job_id
|
||||
# Persist immediately: all subsequent actions use this exact ID.
|
||||
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
|
||||
while True:
|
||||
self._check_master(run)
|
||||
response = await client.status(job_id)
|
||||
status = response.get("status")
|
||||
if not isinstance(status, str) or status not in PROVIDER_STATES:
|
||||
raise RunpodError("Runpod returned an unknown job state")
|
||||
provider_terminal = status in PROVIDER_TERMINAL
|
||||
if status == "COMPLETED":
|
||||
output = response.get("output")
|
||||
if not isinstance(output, dict) or output.get("error") or output.get("errors"):
|
||||
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
|
||||
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
|
||||
images = await decode_results(output, run["manifest"])
|
||||
await self._deliver(run, images=images)
|
||||
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
|
||||
return
|
||||
self.ledger.update(rid, status=status, updated_at=time.time())
|
||||
if provider_terminal:
|
||||
raise RunpodError(f"Serverless job ended with {status}")
|
||||
await asyncio.sleep(self.poll_seconds)
|
||||
except AmbiguousSubmission as exc:
|
||||
self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time())
|
||||
await self._deliver(run, error=str(exc))
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError) as exc:
|
||||
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
|
||||
if submitted and not run["provider_job_id"]:
|
||||
status = "UNKNOWN"
|
||||
error = "Submission interrupted before its job ID was received; check Runpod before retrying"
|
||||
elif isinstance(exc, asyncio.TimeoutError):
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
|
||||
error = "Serverless job exceeded its queue/execution deadline"
|
||||
else:
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED")
|
||||
error = "Serverless work stopped" if terminal or not submitted else "Serverless cancellation requested; provider confirmation is pending"
|
||||
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
|
||||
await self._deliver(run, error=error)
|
||||
except Exception as exc:
|
||||
# If polling/decoding failed, stop only our own provider job, never the pool.
|
||||
terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
|
||||
error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract"
|
||||
status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED")
|
||||
self.ledger.update(rid, status=status, error=error, updated_at=time.time())
|
||||
await self._deliver(run, error=error)
|
||||
finally:
|
||||
# Completed tracking retains IDs/bindings only, not potentially large
|
||||
# base64 inputs, workflow metadata, or a client holding credentials.
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
run["client"] = None
|
||||
# A released local semaphore alone cannot release a durable reservation.
|
||||
await self._notify_capacity()
|
||||
|
||||
async def wait_for_execution(self, execution_id):
|
||||
tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id]
|
||||
await asyncio.gather(*tasks, return_exceptions=False)
|
||||
|
||||
def _detached_records(self, worker_id=None, execution_id=None):
|
||||
return [record for record in self.ledger.all()
|
||||
if is_unresolved(record) and not self._task_is_active(record["record_id"])
|
||||
and (worker_id is None or belongs_to_pool(record, worker_id))
|
||||
and (execution_id is None or record.get("execution_id") == execution_id)]
|
||||
|
||||
async def _recover_provider_record(self, record, *, cancel=False):
|
||||
"""Reconcile an exact saved ID, never submit or resume a workflow."""
|
||||
if not record.get("job_id"):
|
||||
return False
|
||||
status = None
|
||||
try:
|
||||
client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]})
|
||||
response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"]))
|
||||
status = terminal_provider_status(response)
|
||||
except Exception:
|
||||
pass # Unreachable/expired is not proof of provider termination.
|
||||
if status is not None or cancel:
|
||||
self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time())
|
||||
return status is not None
|
||||
|
||||
async def reconcile(self, worker_id=None):
|
||||
"""Read terminal states for detached jobs and release only proven capacity."""
|
||||
resolved = 0
|
||||
async with self.recovery_lock:
|
||||
for record in self._detached_records(worker_id):
|
||||
resolved += await self._recover_provider_record(record)
|
||||
await self._notify_capacity()
|
||||
return resolved
|
||||
|
||||
async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False):
|
||||
"""Explicit operator escape hatch for lost IDs/expired provider records.
|
||||
|
||||
This records a human's verification, not a cancellation or provider
|
||||
status. It is never called automatically or allowed for live tasks.
|
||||
"""
|
||||
if provider_checked is not True:
|
||||
raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it")
|
||||
async with self.recovery_lock:
|
||||
record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None)
|
||||
if record is None:
|
||||
raise ValueError("Unresolved inactive job not found for this pool")
|
||||
self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(),
|
||||
error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed")
|
||||
await self._notify_capacity()
|
||||
|
||||
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
|
||||
targets = {}
|
||||
for rid, task in self.tasks.items():
|
||||
run = self.runs[rid]
|
||||
if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
|
||||
continue
|
||||
if execution_id and run["execution_id"] != execution_id:
|
||||
continue
|
||||
if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
|
||||
continue
|
||||
if not task.done():
|
||||
# Repeated requests must not interrupt the first cancellation's cleanup.
|
||||
if not task.cancelling():
|
||||
task.cancel()
|
||||
targets[rid] = task
|
||||
if targets:
|
||||
await asyncio.gather(*targets.values(), return_exceptions=True)
|
||||
recovered = 0
|
||||
async with self.recovery_lock:
|
||||
for record in self._detached_records(worker_id, execution_id):
|
||||
rid = record["record_id"]
|
||||
# A task cancelled before its first step never enters _run's handlers.
|
||||
# A failed start() may also leave a journalled but unstarted local slot.
|
||||
if record.get("status") == "WAITING" and not record.get("job_id"):
|
||||
run = self.runs.get(rid)
|
||||
if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])):
|
||||
self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
|
||||
run.pop("workflow", None)
|
||||
run.pop("images", None)
|
||||
await self._deliver(run, error="Serverless job cancelled before submission")
|
||||
continue
|
||||
# Do not send two cancellation requests for the same live job in
|
||||
# one action. Its unresolved reservation survives for later recovery.
|
||||
if rid in targets or collector_job_id or not record.get("job_id"):
|
||||
continue
|
||||
await self._recover_provider_record(record, cancel=True)
|
||||
recovered += 1
|
||||
await self._notify_capacity()
|
||||
return len(targets) + recovered
|
||||
|
||||
def public_jobs(self, worker_id):
|
||||
records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)]
|
||||
# Never hide unresolved IDs behind the bounded completed history.
|
||||
visible = [record for record in records if is_unresolved(record)]
|
||||
visible.extend(record for record in records if not is_unresolved(record))
|
||||
unresolved_count = sum(is_unresolved(record) for record in records)
|
||||
return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")}
|
||||
for record in visible[:max(20, unresolved_count)]]
|
||||
@@ -0,0 +1,29 @@
|
||||
"""Shared rules for provider capacity and durable recovery protection.
|
||||
|
||||
A local task ending is not evidence that provider work ended. RESOLVED is an
|
||||
explicit operator acknowledgment, never a provider-reported state.
|
||||
"""
|
||||
|
||||
PROVIDER_TERMINAL = frozenset({"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"})
|
||||
TERMINAL = PROVIDER_TERMINAL | {"RESOLVED"}
|
||||
PROVIDER_STATES = PROVIDER_TERMINAL | {"IN_QUEUE", "IN_PROGRESS"}
|
||||
|
||||
|
||||
def terminal_provider_status(response):
|
||||
status = response.get("status") if isinstance(response, dict) else None
|
||||
return status if isinstance(status, str) and status in PROVIDER_TERMINAL else None
|
||||
|
||||
|
||||
def is_unresolved(record):
|
||||
return record.get("status") not in TERMINAL
|
||||
|
||||
|
||||
def occupies_capacity(record):
|
||||
# WAITING is journalled before admission and cannot have been submitted.
|
||||
# Every other unresolved state is conservatively reserved, including a
|
||||
# submission whose provider ID was lost and an unconfirmed cancellation.
|
||||
return is_unresolved(record) and (record.get("status") != "WAITING" or bool(record.get("job_id")))
|
||||
|
||||
|
||||
def belongs_to_pool(record, worker_id):
|
||||
return str(record.get("worker_id")) == str(worker_id)
|
||||
+64
-11
@@ -74,6 +74,14 @@ def _resolve_enabled_workers(config, requested_ids=None):
|
||||
elif not worker.get("enabled", False):
|
||||
continue
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
from ..utils.serverless_config import validate_serverless_worker
|
||||
pool = validate_serverless_worker(worker)
|
||||
for slot in range(pool["jobs_per_run"]):
|
||||
workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
|
||||
name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
|
||||
continue
|
||||
|
||||
raw_port = worker.get("port", worker.get("listen_port", 8188))
|
||||
try:
|
||||
port = int(raw_port or 8188)
|
||||
@@ -154,6 +162,21 @@ async def _prepare_worker_payload(
|
||||
"""Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
|
||||
async with worker_prep_semaphore:
|
||||
worker_prompt = prompt_index.copy_prompt()
|
||||
if worker.get("type") == "serverless":
|
||||
from .orchestration.serverless_prompt import compile_serverless_prompt
|
||||
from .orchestration.serverless_media import prepare_images
|
||||
from .orchestration.runpod_client import serialize_request
|
||||
worker_prompt = apply_participant_overrides(
|
||||
prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
|
||||
job_id_map, "", delegate_master, prompt_index,
|
||||
)
|
||||
workflow, manifest = compile_serverless_prompt(worker_prompt)
|
||||
images = await prepare_images(workflow)
|
||||
# Fail during preparation, before regular workers or the master are
|
||||
# dispatched. submit() uses this same serializer for the sent bytes.
|
||||
serialize_request(workflow, images, worker["timeout_seconds"])
|
||||
return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images}
|
||||
|
||||
worker_master_url = build_master_callback_url(
|
||||
worker,
|
||||
config=config,
|
||||
@@ -330,9 +353,6 @@ async def orchestrate_distributed_execution(
|
||||
queue_result.get("node_errors", {}),
|
||||
)
|
||||
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
master_prompt = prompt_index.copy_prompt()
|
||||
master_prompt = apply_participant_overrides(
|
||||
master_prompt,
|
||||
@@ -387,7 +407,19 @@ async def orchestrate_distributed_execution(
|
||||
]
|
||||
)
|
||||
|
||||
if worker_payloads:
|
||||
serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"]
|
||||
regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"]
|
||||
manager = None
|
||||
if serverless_jobs:
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
manager = get_manager()
|
||||
manager.ensure_capacity(len(serverless_jobs))
|
||||
|
||||
# Do not leave collector queues behind when offline compilation/media validation fails.
|
||||
for job_id in job_id_map.values():
|
||||
await _ensure_distributed_queue(job_id)
|
||||
|
||||
if regular_payloads:
|
||||
await asyncio.gather(
|
||||
*[
|
||||
dispatch_worker_prompt(
|
||||
@@ -398,17 +430,38 @@ async def orchestrate_distributed_execution(
|
||||
use_websocket=use_websocket,
|
||||
trace_execution_id=execution_trace_id,
|
||||
)
|
||||
for worker, wprompt in worker_payloads
|
||||
for worker, wprompt in regular_payloads
|
||||
]
|
||||
)
|
||||
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
try:
|
||||
queue_result = await queue_prompt_payload(
|
||||
master_prompt,
|
||||
workflow_meta,
|
||||
client_id,
|
||||
include_queue_metadata=True,
|
||||
)
|
||||
except Exception:
|
||||
if serverless_jobs:
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
for job_id in job_id_map.values():
|
||||
prompt_server.distributed_pending_jobs.pop(job_id, None)
|
||||
raise
|
||||
prompt_id = queue_result["prompt_id"]
|
||||
if manager is not None:
|
||||
try:
|
||||
await manager.start(execution_trace_id, prompt_id, serverless_jobs)
|
||||
except Exception:
|
||||
# Wake queued collectors with an explicit error if journal/capacity setup failed.
|
||||
await manager.cancel(execution_id=execution_trace_id)
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
for prepared in serverless_jobs:
|
||||
for binding in prepared["manifest"]:
|
||||
q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
|
||||
if q is not None:
|
||||
q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
|
||||
"serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
|
||||
raise
|
||||
prompt_number = queue_result["number"]
|
||||
node_errors = queue_result.get("node_errors", {})
|
||||
trace_debug(
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Master-only Runpod pool monitoring. Health checks never submit inference."""
|
||||
from aiohttp import web
|
||||
import server
|
||||
|
||||
from ..utils.config import load_config
|
||||
from .orchestration.serverless_runtime import get_manager
|
||||
from .orchestration.runpod_client import RunpodError
|
||||
|
||||
|
||||
def _pool(worker_id):
|
||||
return next((w for w in load_config().get("workers", [])
|
||||
if w.get("id") == worker_id and w.get("type") == "serverless"), None)
|
||||
|
||||
|
||||
async def _status(worker_id, fresh=False):
|
||||
worker = _pool(worker_id)
|
||||
if worker is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
jobs = []
|
||||
try:
|
||||
manager = get_manager()
|
||||
jobs = manager.public_jobs(worker_id)
|
||||
result = await manager.health(worker, fresh=fresh)
|
||||
return web.json_response(dict(result, jobs=jobs))
|
||||
except Exception as exc:
|
||||
message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
|
||||
return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/test")
|
||||
async def test_serverless_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
||||
return await _status(worker_id, fresh=True)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
|
||||
async def serverless_status_endpoint(request):
|
||||
return await _status(request.match_info["worker_id"])
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
|
||||
async def cancel_serverless_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
|
||||
if _pool(worker_id) is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
try:
|
||||
manager = get_manager()
|
||||
count = await manager.cancel(worker_id=worker_id)
|
||||
return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
|
||||
"message": "Cancellation requested for this master's jobs only"})
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
|
||||
|
||||
|
||||
async def _recover(request, *, acknowledge=False):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get("worker_id")
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
raise ValueError()
|
||||
record_id = data.get("record_id")
|
||||
if acknowledge:
|
||||
if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
|
||||
raise ValueError()
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
|
||||
if _pool(worker_id) is None:
|
||||
return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
|
||||
try:
|
||||
manager = get_manager()
|
||||
if acknowledge:
|
||||
await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
|
||||
result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
|
||||
else:
|
||||
result = {"resolved": await manager.reconcile(worker_id=worker_id)}
|
||||
return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
|
||||
except ValueError:
|
||||
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this pool"}, status=409)
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
|
||||
async def reconcile_serverless_endpoint(request):
|
||||
return await _recover(request)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
|
||||
async def resolve_serverless_endpoint(request):
|
||||
return await _recover(request, acknowledge=True)
|
||||
+20
-1
@@ -16,12 +16,14 @@ from ..utils.logging import debug_log, log
|
||||
from ..utils.network import (
|
||||
build_worker_url,
|
||||
get_client_session,
|
||||
get_server_port,
|
||||
handle_api_error,
|
||||
normalize_host,
|
||||
probe_worker,
|
||||
)
|
||||
from ..utils.constants import CHUNK_SIZE
|
||||
from ..workers import get_worker_manager
|
||||
from ..workers.ports import allocate_worker_ports
|
||||
from .schemas import require_fields, validate_worker_id
|
||||
from ..workers.detection import (
|
||||
get_machine_id,
|
||||
@@ -277,6 +279,13 @@ def _get_cuda_info():
|
||||
def _collect_network_info_sync():
|
||||
"""Collect network/cuda info in a worker thread to avoid blocking route handlers."""
|
||||
cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
|
||||
device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
|
||||
master_port = get_server_port()
|
||||
config = load_config()
|
||||
worker_ports = []
|
||||
if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
|
||||
worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
|
||||
worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
|
||||
hostname = socket.gethostname()
|
||||
all_ips = get_network_ips()
|
||||
recommended_ip = get_recommended_ip(all_ips)
|
||||
@@ -285,7 +294,9 @@ def _collect_network_info_sync():
|
||||
"all_ips": all_ips,
|
||||
"recommended_ip": recommended_ip,
|
||||
"cuda_device": cuda_device,
|
||||
"cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
|
||||
"cuda_device_count": device_count,
|
||||
"master_port": master_port,
|
||||
"local_worker_ports": worker_ports,
|
||||
}
|
||||
|
||||
|
||||
@@ -449,6 +460,9 @@ async def launch_worker_endpoint(request):
|
||||
if not worker:
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
return await handle_api_error(request, "Serverless pools are managed by Runpod, not local processes", 400)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = worker_id
|
||||
|
||||
@@ -483,6 +497,8 @@ async def launch_worker_endpoint(request):
|
||||
"message": f"Worker {worker['name']} launched",
|
||||
"log_file": log_file
|
||||
})
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
|
||||
|
||||
@@ -541,6 +557,9 @@ async def get_local_worker_status_endpoint(request):
|
||||
worker_statuses = {}
|
||||
|
||||
for worker in config.get("workers", []):
|
||||
# Serverless pools have no host, port, or local process.
|
||||
if worker.get("type") == "serverless":
|
||||
continue
|
||||
# Only check local workers
|
||||
host = normalize_host(worker.get("host")) or ""
|
||||
if not host or host in ["localhost", "127.0.0.1"]:
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
# conftest.py — project-level pytest configuration.
|
||||
#
|
||||
# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
|
||||
# (from .distributed import ...) that fail when pytest tries to import it as a
|
||||
# (from .nodes.v3 import ...) that fail when pytest tries to import it as a
|
||||
# standalone module during Package.setup() for the root package node.
|
||||
#
|
||||
# Fix: patch Package.setup() to skip the root-package's __init__.py import.
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
"""
|
||||
ComfyUI-Distributed: thin entry point.
|
||||
All implementation lives in workers/, nodes/, api/.
|
||||
"""
|
||||
import atexit
|
||||
import os
|
||||
|
||||
import server
|
||||
|
||||
from .utils.config import ensure_config_exists
|
||||
from .utils.logging import debug_log
|
||||
from .utils.network import cleanup_client_session
|
||||
from .workers import get_worker_manager
|
||||
from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
|
||||
from .upscale.job_store import ensure_tile_jobs_initialized
|
||||
from .nodes import (
|
||||
NODE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS,
|
||||
ImageBatchDivider,
|
||||
DistributedCollectorNode,
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from . import api # noqa: F401 - triggers all @routes.* registrations
|
||||
from .api.queue_orchestration import ensure_distributed_state
|
||||
|
||||
ensure_config_exists()
|
||||
|
||||
# Aiohttp session cleanup
|
||||
async def _cleanup_session():
|
||||
await cleanup_client_session()
|
||||
|
||||
|
||||
atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup
|
||||
|
||||
# Initialize distributed job state on prompt_server
|
||||
prompt_server = server.PromptServer.instance
|
||||
ensure_distributed_state(prompt_server)
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
# Worker startup
|
||||
if not os.environ.get('COMFYUI_IS_WORKER'):
|
||||
atexit.register(sync_cleanup)
|
||||
delayed_auto_launch()
|
||||
register_async_signals()
|
||||
@@ -0,0 +1,163 @@
|
||||
# Runpod serverless pools (preview)
|
||||
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint/pool, not a permanent
|
||||
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
|
||||
normal ComfyUI API workflow per logical slot and collects the returned images.
|
||||
The endpoint workers **do not need ComfyUI-Distributed installed**.
|
||||
|
||||
## Setup
|
||||
|
||||
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
|
||||
2. Enter its name, Runpod **endpoint ID** (not a URL), and API key.
|
||||
3. Start with **Jobs per run = 1**, **Concurrent jobs = 1**, and **Timeout = 1800**.
|
||||
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
|
||||
endpoint health only; it does not start GPU inference or prove model/node compatibility.
|
||||
5. Enable the pool when you want queued Distributed workflows to use it.
|
||||
6. Use an image **Distributed Collector** and queue through Distributed as usual.
|
||||
Enable **Delegate Only** if the master should collect without generating locally.
|
||||
|
||||
Keys are write-only. Leave the key field blank to retain a saved key, or select
|
||||
**Clear saved API key** to remove it (disable the pool first). A saved key is
|
||||
required before enabling a pool. Use HTTPS or a trusted local connection when
|
||||
entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
|
||||
users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
|
||||
|
||||
## Endpoint contract
|
||||
|
||||
The endpoint must accept this generic input shape; named application modes are
|
||||
not used:
|
||||
|
||||
```json
|
||||
{
|
||||
"input": {
|
||||
"workflow": {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
|
||||
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}}
|
||||
},
|
||||
"images": []
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
For image inputs, `images` contains `{ "name": "input.png", "image": "<base64>" }`
|
||||
entries and the handler must rewrite `LoadImage` references to the staged files.
|
||||
The adapter transfers local input images; it does not copy models or install nodes.
|
||||
All referenced models and custom nodes must already exist in the endpoint image.
|
||||
|
||||
On completion, Runpod's `output.images` must contain artifact records such as:
|
||||
|
||||
```json
|
||||
{
|
||||
"filename": "example_00001_.png",
|
||||
"subfolder": "",
|
||||
"artifact_type": "output",
|
||||
"artifact_path": "example_00001_.png",
|
||||
"type": "base64",
|
||||
"data": "<base64>"
|
||||
}
|
||||
```
|
||||
|
||||
Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"`
|
||||
with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without
|
||||
Runpod credentials; private-network destinations and unsafe redirects are refused.
|
||||
The handler must preserve injected `SaveImage.filename_prefix` values so results
|
||||
can be matched to the right collector and slot. Duplicate basenames in different
|
||||
subfolders remain different artifact identities.
|
||||
|
||||
A Runpod job marked `COMPLETED` is still rejected if `output.error` or
|
||||
`output.errors` is nonempty, required collector outputs are missing, or image
|
||||
artifacts are invalid. Failures are surfaced rather than silently returning an
|
||||
incomplete serverless batch.
|
||||
|
||||
## Scheduling and cancellation
|
||||
|
||||
- **Jobs per run** adds logical collector participants. Each receives its own
|
||||
Distributed Seed/Value overrides. It is not a physical GPU count.
|
||||
- **Concurrent jobs** bounds this master's outstanding submissions per pool,
|
||||
including retained jobs after a restart. Ambiguous submissions and unconfirmed
|
||||
cancellations keep their capacity reservation until reconciled. Other
|
||||
applications using the same endpoint are unaffected.
|
||||
- **Load balance** selects one participant, so it does not fan out every pool slot.
|
||||
- **Timeout** includes waiting for a local slot, Runpod queue/cold-start time,
|
||||
execution, and output retrieval. Provider policy limits are submitted as well.
|
||||
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
|
||||
- Interrupting a collector or using the panel cancellation action requests
|
||||
cancellation only for jobs submitted by this master. It never purges the endpoint.
|
||||
- A network failure during submission is **not retried automatically**: the job
|
||||
may already be billable. Check the Runpod job list before queuing again.
|
||||
|
||||
## Persistence and recovery
|
||||
|
||||
Credentials live outside `gpu_config.json` and Git in the master ComfyUI user
|
||||
folder under `distributed/serverless/credentials.json`, with restrictive POSIX
|
||||
permissions. They are omitted from public configuration, workflow exports, and
|
||||
browser configuration state. The first version requires a POSIX (Linux/macOS)
|
||||
master for serverless secret storage; existing non-serverless workers remain
|
||||
available on Windows.
|
||||
|
||||
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
|
||||
not image data, workflow payloads, or API keys. A restarted master cannot resume
|
||||
an interrupted collector. It does **not** automatically resubmit jobs. Retained
|
||||
IDs are available through `/distributed/serverless/status/<worker_id>`; the pool's
|
||||
cancel action can request cancellation of known pre-restart jobs. If no job ID
|
||||
was received, check Runpod manually. Cancellation is best-effort and may race
|
||||
with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation.
|
||||
|
||||
Unresolved records protect the pool's endpoint, concurrency, API key, worker type,
|
||||
and deletion even when no local task remains. Disabling a pool is still allowed;
|
||||
it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are
|
||||
closed locally, while interrupted `SUBMITTING` records without IDs become
|
||||
`UNKNOWN` and remain protected.
|
||||
|
||||
### Reconcile or acknowledge recovery
|
||||
|
||||
Recovery uses the same trusted-master API as the panel; it never resubmits a job.
|
||||
|
||||
1. Read `GET /distributed/serverless/status/<worker_id>`. The `jobs` list includes
|
||||
every unresolved record with its local `record_id`, provider `job_id` (if
|
||||
received), and state. These records remain visible if endpoint health fails.
|
||||
2. Request cancellation from the panel, or send
|
||||
`POST /distributed/serverless/reconcile` with the body below to **read** the
|
||||
retained provider IDs. Only a confirmed terminal state releases protection.
|
||||
Network errors, expired IDs, or ambiguous responses do not.
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID"}
|
||||
```
|
||||
|
||||
3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
|
||||
manually first. **Only after verifying that the job is no longer running or
|
||||
queued**, send `POST /distributed/serverless/resolve` for that one record:
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
|
||||
```
|
||||
|
||||
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
|
||||
provider-confirmed cancellation. It cannot resolve a live local task and does
|
||||
not stop remote compute. It releases the record's reservation and may let
|
||||
already-waiting jobs proceed. Never automate this acknowledgment as a retry
|
||||
mechanism; an unverified job can still be billable.
|
||||
|
||||
Pool changes become available once all its local tasks and unresolved records
|
||||
are finished. No credentials or ledger files need to be edited by hand.
|
||||
|
||||
## First-version limits
|
||||
|
||||
- Image collectors only. Audio/video transfer and audio collectors are rejected.
|
||||
- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
|
||||
node types are rejected before submission.
|
||||
- Downstream post-collector processing remains local; the remote graph contains
|
||||
only the upstream compute required for image saves.
|
||||
- Inputs must be supported images under ComfyUI's input/output/temp roots.
|
||||
- Raw image input reads are guarded at 10,000,000 bytes total, and the complete
|
||||
serialized `/run` JSON request body (workflow, base64 images, and policy) must
|
||||
also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image
|
||||
budget is lower and depends on workflow size. Large-media/object-storage upload
|
||||
support is not included.
|
||||
- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
|
||||
scaling changes, or inbound master callback/tunnel is needed.
|
||||
|
||||
This feature is verified with offline provider/transport tests. A successful
|
||||
health test is not a paid end-to-end generation test; test your deployed handler
|
||||
and model inventory with one slot before increasing concurrency.
|
||||
@@ -0,0 +1,53 @@
|
||||
# Native V3 node API
|
||||
|
||||
The package registers one `ComfyExtension` through `comfy_entrypoint()` and
|
||||
imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
|
||||
that API; there is no V1 registration fallback.
|
||||
|
||||
## Compatibility
|
||||
|
||||
- All eight node IDs, display names, categories, visible input order, defaults
|
||||
and output order are retained. Existing execution algorithms remain in the
|
||||
private collector, utilities and upscale modules.
|
||||
- The image/audio dividers explicitly declare ten typed outputs. Their existing
|
||||
frontend extensions still show only the selected number of outputs. This
|
||||
replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
|
||||
IMAGE/AUDIO socket indices or the ten returned values.
|
||||
- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
|
||||
its existing prompt-input names and Python defaults via `accept_all_inputs`;
|
||||
it does not become visible widgets.
|
||||
- The collector retains list-input handling. Upscale retains its always-changing
|
||||
fingerprint and creates a private runtime helper for each V3 execution, so
|
||||
mutable helper state is not attached to sanitized V3 class clones.
|
||||
- Routes, distributed state and the existing worker startup/shutdown hooks are
|
||||
initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
|
||||
Worker mode still suppresses automatic worker launch. `distributed.py` is no
|
||||
longer an entrypoint.
|
||||
|
||||
## Verification
|
||||
|
||||
Run ordinary unit tests from this repository:
|
||||
|
||||
```bash
|
||||
python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
|
||||
|
||||
```bash
|
||||
COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
|
||||
/path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
|
||||
```
|
||||
|
||||
The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
|
||||
parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
|
||||
all node schemas with the V1 fixture, checks all five bundled workflows' node
|
||||
IDs and socket/link contracts, exercises injected worker metadata, image/audio
|
||||
lists, collector aggregation and divider outputs, rejects invalid upscale enums,
|
||||
and decodes an actual preview PNG referenced by executor history.
|
||||
|
||||
The upscale GPU/model boundary is mocked to check argument forwarding and
|
||||
per-execution helper isolation. Full checkpoint inference, browser canvas
|
||||
acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
|
||||
workers or model downloads are started; preview files use a temporary scratch
|
||||
directory. The test does not install this branch into a live custom-node folder.
|
||||
+1
-31
@@ -1,31 +1 @@
|
||||
from .utilities import (
|
||||
DistributedSeed,
|
||||
DistributedModelName,
|
||||
DistributedValue,
|
||||
ImageBatchDivider,
|
||||
AudioBatchDivider,
|
||||
DistributedEmptyImage,
|
||||
AnyType,
|
||||
ByPassTypeTuple,
|
||||
any_type,
|
||||
)
|
||||
from .collector import DistributedCollectorNode
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DistributedCollector": DistributedCollectorNode,
|
||||
"DistributedSeed": DistributedSeed,
|
||||
"DistributedModelName": DistributedModelName,
|
||||
"DistributedValue": DistributedValue,
|
||||
"ImageBatchDivider": ImageBatchDivider,
|
||||
"AudioBatchDivider": AudioBatchDivider,
|
||||
"DistributedEmptyImage": DistributedEmptyImage,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DistributedCollector": "Distributed Collector",
|
||||
"DistributedSeed": "Distributed Seed",
|
||||
"DistributedModelName": "Distributed Model Name",
|
||||
"DistributedValue": "Distributed Value",
|
||||
"ImageBatchDivider": "Image Batch Divider",
|
||||
"AudioBatchDivider": "Audio Segment Divider",
|
||||
"DistributedEmptyImage": "Distributed Empty Image",
|
||||
}
|
||||
"""Private execution helpers; public registration lives in nodes.v3."""
|
||||
|
||||
+157
-112
@@ -151,6 +151,15 @@ class DistributedCollectorNode:
|
||||
)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _get_serverless_manager():
|
||||
"""Return the serverless manager only for collector jobs with serverless slots."""
|
||||
manager = getattr(prompt_server, "distributed_serverless", None)
|
||||
if manager is not None:
|
||||
return manager
|
||||
from ..api.orchestration.serverless_runtime import get_manager
|
||||
return get_manager()
|
||||
|
||||
async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
|
||||
"""Send an image batch, optionally with audio, or an audio-only completion."""
|
||||
encoded_audio = encode_audio_payload(audio)
|
||||
@@ -339,6 +348,17 @@ class DistributedCollectorNode:
|
||||
seen_enabled.add(worker_id_str)
|
||||
enabled_workers.append(worker_id_str)
|
||||
expected_workers = set(enabled_workers)
|
||||
serverless_workers = {
|
||||
worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
|
||||
}
|
||||
serverless_manager = None
|
||||
|
||||
def get_serverless_manager():
|
||||
nonlocal serverless_manager
|
||||
if serverless_workers and serverless_manager is None:
|
||||
serverless_manager = self._get_serverless_manager()
|
||||
return serverless_manager
|
||||
|
||||
num_workers = len(expected_workers)
|
||||
if num_workers == 0:
|
||||
return (images, audio if audio is not None else self.EMPTY_AUDIO)
|
||||
@@ -405,152 +425,177 @@ class DistributedCollectorNode:
|
||||
workers_done.add(done_worker_id)
|
||||
p.update(1) # +1 per completed expected worker
|
||||
|
||||
def consume_result(result):
|
||||
nonlocal collected_count
|
||||
if "serverless_error" in result:
|
||||
error = result.get("serverless_error")
|
||||
message = str(error).strip() if error is not None else "unspecified serverless failure"
|
||||
participant = result.get("worker_id", "unknown")
|
||||
raise RuntimeError(f"Serverless participant {participant} failed: {message}")
|
||||
|
||||
result_worker_id = str(result["worker_id"])
|
||||
is_last = result.get("is_last", False)
|
||||
collected_count += self._store_worker_result(worker_images, result)
|
||||
debug_log(
|
||||
f"Master - Got canonical result from worker {result_worker_id}, "
|
||||
f"image {result.get('image_index', 0)}, is_last={is_last}"
|
||||
)
|
||||
result_audio = result.get("audio")
|
||||
if result_audio is not None:
|
||||
worker_audio[result_worker_id] = result_audio
|
||||
debug_log(f"Master - Got audio from worker {result_worker_id}")
|
||||
if is_last:
|
||||
mark_worker_done(result_worker_id)
|
||||
|
||||
try:
|
||||
while len(workers_done) < num_workers:
|
||||
# Check for user interruption to abort collection promptly
|
||||
# Check for user interruption to abort collection promptly.
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
try:
|
||||
# Get the queue again each time to ensure we have the right reference
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
current_size = q.qsize()
|
||||
|
||||
result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
count = self._store_worker_result(worker_images, result)
|
||||
collected_count += count
|
||||
debug_log(
|
||||
f"Master - Got canonical result from worker {worker_id}, "
|
||||
f"image {result.get('image_index', 0)}, is_last={is_last}"
|
||||
)
|
||||
consume_result(result)
|
||||
|
||||
# Collect audio data if present
|
||||
result_audio = result.get('audio')
|
||||
if result_audio is not None:
|
||||
worker_audio[worker_id] = result_audio
|
||||
debug_log(f"Master - Got audio from worker {worker_id}")
|
||||
|
||||
# Record activity and refresh timeout baseline
|
||||
# Real results refresh the regular-worker timeout baseline.
|
||||
last_activity = time.time()
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
|
||||
if is_last:
|
||||
mark_worker_done(worker_id)
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# If we still have time, continue polling; otherwise handle timeout
|
||||
if (time.time() - last_activity) < base_timeout:
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
continue
|
||||
# Re-check for user interruption after timeout expiry
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
missing_workers = set(str(w) for w in enabled_workers) - workers_done
|
||||
missing_workers = expected_workers - workers_done
|
||||
|
||||
# A queue delivery can race the timed wait. Process it before
|
||||
# declaring an owned provider slot stopped without a result.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
queue_has_results = current_q is not None and not current_q.empty()
|
||||
if queue_has_results:
|
||||
continue
|
||||
|
||||
missing_serverless = missing_workers & serverless_workers
|
||||
if missing_serverless:
|
||||
manager = get_serverless_manager()
|
||||
if manager is None:
|
||||
raise RuntimeError("Serverless manager is unavailable for pending collector slots")
|
||||
pending_serverless = []
|
||||
for participant_id in sorted(missing_serverless):
|
||||
if manager.is_pending(multi_job_id, participant_id):
|
||||
pending_serverless.append(participant_id)
|
||||
elif manager.owns(multi_job_id, participant_id):
|
||||
raise RuntimeError(
|
||||
f"Serverless participant {participant_id} stopped without delivering a result"
|
||||
)
|
||||
# Provider deadlines are managed by the serverless runtime;
|
||||
# do not apply the regular worker heartbeat timeout to it.
|
||||
if pending_serverless:
|
||||
continue
|
||||
|
||||
missing_regular = missing_workers - serverless_workers
|
||||
elapsed = time.time() - last_activity
|
||||
for missing_worker_id in sorted(missing_workers):
|
||||
for missing_worker_id in sorted(missing_regular):
|
||||
log(
|
||||
"Master - Heartbeat timeout: "
|
||||
f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
|
||||
)
|
||||
log(
|
||||
f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
|
||||
f"(elapsed={elapsed:.1f}s)"
|
||||
)
|
||||
if missing_regular:
|
||||
log(
|
||||
f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
|
||||
f"(elapsed={elapsed:.1f}s)"
|
||||
)
|
||||
|
||||
# Probe missing workers' /prompt endpoints to check if they are actively processing
|
||||
# Probe only regular worker /prompt endpoints.
|
||||
any_busy = False
|
||||
try:
|
||||
cfg = load_config()
|
||||
cfg_workers = cfg.get('workers', [])
|
||||
for wid in list(missing_workers):
|
||||
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
|
||||
if not wrec:
|
||||
debug_log(f"Collector probe: worker {wid} not found in config")
|
||||
continue
|
||||
worker_url = build_worker_url(wrec)
|
||||
try:
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
queue_remaining = None
|
||||
if payload is not None:
|
||||
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
|
||||
debug_log(
|
||||
"Collector probe: worker "
|
||||
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
|
||||
)
|
||||
if payload is not None and queue_remaining and queue_remaining > 0:
|
||||
any_busy = True
|
||||
log(
|
||||
f"Master - Probe grace: worker {wid} appears busy "
|
||||
f"(queue_remaining={queue_remaining}). Continuing to wait."
|
||||
if missing_regular:
|
||||
try:
|
||||
cfg = load_config()
|
||||
cfg_workers = cfg.get('workers', [])
|
||||
for wid in sorted(missing_regular):
|
||||
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
|
||||
if not wrec:
|
||||
debug_log(f"Collector probe: worker {wid} not found in config")
|
||||
continue
|
||||
worker_url = build_worker_url(wrec)
|
||||
try:
|
||||
payload = await probe_worker(worker_url, timeout=2.0)
|
||||
queue_remaining = None
|
||||
if payload is not None:
|
||||
queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
|
||||
debug_log(
|
||||
"Collector probe: worker "
|
||||
f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe failed for worker {wid}: {e}")
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe setup error: {e}")
|
||||
if payload is not None and queue_remaining and queue_remaining > 0:
|
||||
any_busy = True
|
||||
log(
|
||||
f"Master - Probe grace: worker {wid} appears busy "
|
||||
f"(queue_remaining={queue_remaining}). Continuing to wait."
|
||||
)
|
||||
break
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe failed for worker {wid}: {e}")
|
||||
except Exception as e:
|
||||
debug_log(f"Collector probe setup error: {e}")
|
||||
|
||||
if any_busy:
|
||||
# Refresh last_activity and continue waiting
|
||||
last_activity = time.time()
|
||||
# Refresh base timeout in case the user changed it in UI
|
||||
base_timeout = float(get_worker_timeout_seconds())
|
||||
continue
|
||||
|
||||
# Check queue size again with lock
|
||||
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
final_q = prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
final_size = final_q.qsize()
|
||||
|
||||
# Try to drain any remaining items
|
||||
remaining_items = []
|
||||
final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
|
||||
remaining_items = []
|
||||
if final_q is not None:
|
||||
while not final_q.empty():
|
||||
try:
|
||||
item = final_q.get_nowait()
|
||||
remaining_items.append(item)
|
||||
remaining_items.append(final_q.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
if remaining_items:
|
||||
# Process them
|
||||
for item in remaining_items:
|
||||
worker_id = item['worker_id']
|
||||
is_last = item.get('is_last', False)
|
||||
|
||||
collected_count += self._store_worker_result(worker_images, item)
|
||||
|
||||
if is_last:
|
||||
mark_worker_done(worker_id)
|
||||
else:
|
||||
log(f"Master - Queue {multi_job_id} no longer exists!")
|
||||
for item in remaining_items:
|
||||
consume_result(item)
|
||||
break
|
||||
except comfy.model_management.InterruptProcessingException:
|
||||
# Cleanup queue on interruption and re-raise to abort prompt cleanly
|
||||
|
||||
total_collected = sum(len(imgs) for imgs in worker_images.values())
|
||||
|
||||
# Remove the queue before combining so a late provider callback cannot
|
||||
# resurrect results for a collector that has completed.
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
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
|
||||
|
||||
total_collected = sum(len(imgs) for imgs in worker_images.values())
|
||||
|
||||
# Clean up job queue
|
||||
async with prompt_server.distributed_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_jobs:
|
||||
del prompt_server.distributed_pending_jobs[multi_job_id]
|
||||
|
||||
combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
|
||||
try:
|
||||
combined = self._reorder_and_combine_tensors(
|
||||
worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
|
||||
)
|
||||
if combined is None:
|
||||
debug_log(f"Master - Job {multi_job_id} complete with audio only")
|
||||
else:
|
||||
debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
|
||||
f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
|
||||
|
||||
return (combined, combined_audio)
|
||||
except Exception as e:
|
||||
log(f"Master - Error combining images: {e}")
|
||||
# Preserve collected audio even when image assembly fails.
|
||||
return (images, combined_audio)
|
||||
|
||||
@@ -268,12 +268,3 @@ class UltimateSDUpscaleDistributed(
|
||||
|
||||
# Ensure initialization before registering routes
|
||||
ensure_tile_jobs_initialized()
|
||||
|
||||
# Node registration
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
|
||||
}
|
||||
|
||||
+181
@@ -0,0 +1,181 @@
|
||||
"""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]
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.4.7"
|
||||
version = "1.5.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Internal extension lifecycle support (not a node provider)."""
|
||||
@@ -0,0 +1,35 @@
|
||||
"""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}')
|
||||
@@ -38,17 +38,10 @@ def _load_config_routes_module():
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = []
|
||||
utils_pkg.__path__ = [str(module_path.parents[1] / "utils")]
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
created_aiohttp_stub = False
|
||||
if "aiohttp" not in sys.modules:
|
||||
created_aiohttp_stub = True
|
||||
aiohttp_module = types.ModuleType("aiohttp")
|
||||
aiohttp_module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
sys.modules["aiohttp"] = aiohttp_module
|
||||
import aiohttp # Load the real dependency; isolate the fake response at the tested module.
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
@@ -96,8 +89,14 @@ def _load_config_routes_module():
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
class _EmptySecretStore:
|
||||
def has(self, _worker_id):
|
||||
return False
|
||||
|
||||
module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
module._serverless_secret_store = lambda: _EmptySecretStore()
|
||||
|
||||
return module
|
||||
|
||||
|
||||
@@ -205,6 +205,10 @@ def _load_job_routes_module():
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
# Keep response assertions independent of real aiohttp import ordering.
|
||||
module.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
||||
)
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
|
||||
@@ -164,26 +164,11 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class FindMediaReferencesTests(unittest.TestCase):
|
||||
def test_finds_image_input(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("photo.png", refs)
|
||||
|
||||
def test_finds_video_input(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("clip.mp4", refs)
|
||||
|
||||
def test_finds_file_input_for_load_video(self):
|
||||
prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("1 - Copy.mp4", refs)
|
||||
|
||||
def test_finds_audio_input(self):
|
||||
prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
self.assertIn("track.wav", refs)
|
||||
|
||||
def test_strips_annotation_suffix(self):
|
||||
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
|
||||
refs = ms._find_media_references(prompt)
|
||||
|
||||
@@ -106,7 +106,7 @@ def _load_worker_routes_module():
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
|
||||
workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
@@ -213,6 +213,7 @@ def _load_worker_routes_module():
|
||||
raise RuntimeError("not used in these tests")
|
||||
|
||||
network_module.get_client_session = _get_client_session
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
@@ -256,6 +257,8 @@ def _load_worker_routes_module():
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
assert spec is not None and spec.loader is not None
|
||||
spec.loader.exec_module(module)
|
||||
# Keep this module's response double independent of aiohttp import order.
|
||||
module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
|
||||
|
||||
if created_aiohttp_stub:
|
||||
sys.modules.pop("aiohttp", None)
|
||||
@@ -269,6 +272,38 @@ worker_routes = _load_worker_routes_module()
|
||||
|
||||
|
||||
class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_network_info_reports_actual_master_and_available_ports(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": []}), \
|
||||
patch.object(worker_routes, "get_server_port", return_value=8189), \
|
||||
patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["master_port"], 8189)
|
||||
self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
|
||||
self.assertEqual(response.payload["cuda_device_count"], 4)
|
||||
allocate.assert_called_once_with(8189, [], 3)
|
||||
|
||||
async def test_network_info_does_not_reallocate_existing_configuration(self):
|
||||
with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
|
||||
patch.object(worker_routes, "get_network_ips", return_value=[]), \
|
||||
patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
|
||||
patch.object(worker_routes, "allocate_worker_ports") as allocate:
|
||||
response = await worker_routes.get_network_info_endpoint(_FakeRequest())
|
||||
self.assertEqual(response.payload["local_worker_ports"], [])
|
||||
allocate.assert_not_called()
|
||||
|
||||
async def test_launch_conflict_is_a_clear_client_error(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
|
||||
with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
|
||||
patch.object(worker_routes, "load_config", return_value=config), \
|
||||
patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
|
||||
response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertIn("master port 8189", response.payload["message"])
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
async def test_launch_worker_valid_id_returns_200(self):
|
||||
manager = _DummyWorkerManager()
|
||||
config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
|
||||
|
||||
+691
@@ -0,0 +1,691 @@
|
||||
{
|
||||
"DistributedCollector": {
|
||||
"input": {
|
||||
"required": {
|
||||
"load_balance": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false,
|
||||
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
|
||||
}
|
||||
]
|
||||
},
|
||||
"optional": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_batch_size": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 1024
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"pass_through": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"delegate_only": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"load_balance"
|
||||
],
|
||||
"optional": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_batch_size",
|
||||
"worker_id",
|
||||
"pass_through",
|
||||
"delegate_only"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE",
|
||||
"AUDIO"
|
||||
],
|
||||
"output_name": [
|
||||
"images",
|
||||
"audio"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": true,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Collector"
|
||||
},
|
||||
"DistributedSeed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 1125899906842,
|
||||
"min": 0,
|
||||
"max": 1125899906842624,
|
||||
"forceInput": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"seed"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"INT"
|
||||
],
|
||||
"output_name": [
|
||||
"seed"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Seed"
|
||||
},
|
||||
"DistributedModelName": {
|
||||
"input": {
|
||||
"required": {
|
||||
"text": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"unique_id": "UNIQUE_ID",
|
||||
"extra_pnginfo": "EXTRA_PNGINFO"
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"text"
|
||||
],
|
||||
"hidden": [
|
||||
"unique_id",
|
||||
"extra_pnginfo"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"output"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Model Name"
|
||||
},
|
||||
"DistributedValue": {
|
||||
"input": {
|
||||
"required": {
|
||||
"default_value": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"worker_values": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "{}"
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"default_value",
|
||||
"worker_values"
|
||||
],
|
||||
"hidden": [
|
||||
"is_worker",
|
||||
"worker_id"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"value"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "utils",
|
||||
"display_name": "Distributed Value"
|
||||
},
|
||||
"ImageBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"images": [
|
||||
"IMAGE"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of parts to divide the batch into"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"images",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"batch_1",
|
||||
"batch_2",
|
||||
"batch_3",
|
||||
"batch_4",
|
||||
"batch_5",
|
||||
"batch_6",
|
||||
"batch_7",
|
||||
"batch_8",
|
||||
"batch_9",
|
||||
"batch_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "image",
|
||||
"display_name": "Image Batch Divider"
|
||||
},
|
||||
"AudioBatchDivider": {
|
||||
"input": {
|
||||
"required": {
|
||||
"audio": [
|
||||
"AUDIO"
|
||||
],
|
||||
"divide_by": [
|
||||
"INT",
|
||||
{
|
||||
"default": 2,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"tooltip": "Number of sequential time segments to create"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"audio",
|
||||
"divide_by"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*",
|
||||
"*"
|
||||
],
|
||||
"output_name": [
|
||||
"audio_1",
|
||||
"audio_2",
|
||||
"audio_3",
|
||||
"audio_4",
|
||||
"audio_5",
|
||||
"audio_6",
|
||||
"audio_7",
|
||||
"audio_8",
|
||||
"audio_9",
|
||||
"audio_10"
|
||||
],
|
||||
"output_is_list": [
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": true,
|
||||
"category": "audio",
|
||||
"display_name": "Audio Segment Divider"
|
||||
},
|
||||
"DistributedEmptyImage": {
|
||||
"input": {
|
||||
"required": {
|
||||
"height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 64,
|
||||
"min": 1,
|
||||
"max": 4096,
|
||||
"step": 1
|
||||
}
|
||||
],
|
||||
"channels": [
|
||||
"INT",
|
||||
{
|
||||
"default": 3,
|
||||
"min": 1,
|
||||
"max": 4,
|
||||
"step": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"height",
|
||||
"width",
|
||||
"channels"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image",
|
||||
"display_name": "Distributed Empty Image"
|
||||
},
|
||||
"UltimateSDUpscaleDistributed": {
|
||||
"input": {
|
||||
"required": {
|
||||
"upscaled_image": [
|
||||
"IMAGE"
|
||||
],
|
||||
"model": [
|
||||
"MODEL"
|
||||
],
|
||||
"positive": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"negative": [
|
||||
"CONDITIONING"
|
||||
],
|
||||
"vae": [
|
||||
"VAE"
|
||||
],
|
||||
"seed": [
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 18446744073709551615
|
||||
}
|
||||
],
|
||||
"steps": [
|
||||
"INT",
|
||||
{
|
||||
"default": 20,
|
||||
"min": 1,
|
||||
"max": 10000
|
||||
}
|
||||
],
|
||||
"cfg": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 8.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0
|
||||
}
|
||||
],
|
||||
"sampler_name": [
|
||||
[
|
||||
"euler",
|
||||
"euler_cfg_pp",
|
||||
"euler_ancestral",
|
||||
"euler_ancestral_cfg_pp",
|
||||
"heun",
|
||||
"heunpp2",
|
||||
"exp_heun_2_x0",
|
||||
"exp_heun_2_x0_sde",
|
||||
"dpm_2",
|
||||
"dpm_2_ancestral",
|
||||
"lms",
|
||||
"dpm_fast",
|
||||
"dpm_adaptive",
|
||||
"dpmpp_2s_ancestral",
|
||||
"dpmpp_2s_ancestral_cfg_pp",
|
||||
"dpmpp_sde",
|
||||
"dpmpp_sde_gpu",
|
||||
"dpmpp_2m",
|
||||
"dpmpp_2m_cfg_pp",
|
||||
"dpmpp_2m_sde",
|
||||
"dpmpp_2m_sde_gpu",
|
||||
"dpmpp_2m_sde_heun",
|
||||
"dpmpp_2m_sde_heun_gpu",
|
||||
"dpmpp_3m_sde",
|
||||
"dpmpp_3m_sde_gpu",
|
||||
"ddpm",
|
||||
"lcm",
|
||||
"ipndm",
|
||||
"ipndm_v",
|
||||
"deis",
|
||||
"cfgpp_ud10_ab",
|
||||
"res_multistep",
|
||||
"res_multistep_cfg_pp",
|
||||
"res_multistep_ancestral",
|
||||
"res_multistep_ancestral_cfg_pp",
|
||||
"gradient_estimation",
|
||||
"gradient_estimation_cfg_pp",
|
||||
"er_sde",
|
||||
"seeds_2",
|
||||
"seeds_3",
|
||||
"sa_solver",
|
||||
"sa_solver_pece",
|
||||
"ddim",
|
||||
"uni_pc",
|
||||
"uni_pc_bh2"
|
||||
]
|
||||
],
|
||||
"scheduler": [
|
||||
[
|
||||
"simple",
|
||||
"sgm_uniform",
|
||||
"karras",
|
||||
"exponential",
|
||||
"ddim_uniform",
|
||||
"beta",
|
||||
"normal",
|
||||
"linear_quadratic",
|
||||
"kl_optimal"
|
||||
]
|
||||
],
|
||||
"denoise": [
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.5,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}
|
||||
],
|
||||
"tile_width": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"tile_height": [
|
||||
"INT",
|
||||
{
|
||||
"default": 512,
|
||||
"min": 64,
|
||||
"max": 2048,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"padding": [
|
||||
"INT",
|
||||
{
|
||||
"default": 32,
|
||||
"min": 0,
|
||||
"max": 256,
|
||||
"step": 8
|
||||
}
|
||||
],
|
||||
"mask_blur": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 0,
|
||||
"max": 256
|
||||
}
|
||||
],
|
||||
"force_uniform_tiles": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": true
|
||||
}
|
||||
],
|
||||
"tiled_decode": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
]
|
||||
},
|
||||
"hidden": {
|
||||
"multi_job_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"is_worker": [
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": false
|
||||
}
|
||||
],
|
||||
"master_url": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"enabled_worker_ids": [
|
||||
"STRING",
|
||||
{
|
||||
"default": "[]"
|
||||
}
|
||||
],
|
||||
"worker_id": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"tile_indices": [
|
||||
"STRING",
|
||||
{
|
||||
"default": ""
|
||||
}
|
||||
],
|
||||
"dynamic_threshold": [
|
||||
"INT",
|
||||
{
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 64
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"input_order": {
|
||||
"required": [
|
||||
"upscaled_image",
|
||||
"model",
|
||||
"positive",
|
||||
"negative",
|
||||
"vae",
|
||||
"seed",
|
||||
"steps",
|
||||
"cfg",
|
||||
"sampler_name",
|
||||
"scheduler",
|
||||
"denoise",
|
||||
"tile_width",
|
||||
"tile_height",
|
||||
"padding",
|
||||
"mask_blur",
|
||||
"force_uniform_tiles",
|
||||
"tiled_decode"
|
||||
],
|
||||
"hidden": [
|
||||
"multi_job_id",
|
||||
"is_worker",
|
||||
"master_url",
|
||||
"enabled_worker_ids",
|
||||
"worker_id",
|
||||
"tile_indices",
|
||||
"dynamic_threshold"
|
||||
]
|
||||
},
|
||||
"output": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_name": [
|
||||
"IMAGE"
|
||||
],
|
||||
"output_is_list": [
|
||||
false
|
||||
],
|
||||
"is_input_list": false,
|
||||
"output_node": false,
|
||||
"category": "image/upscaling",
|
||||
"display_name": "Ultimate SD Upscale Distributed (No Upscale)"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
"""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())
|
||||
@@ -18,7 +18,7 @@ def _load_config_module():
|
||||
del sys.modules[mod_name]
|
||||
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
root_pkg.__path__ = [str(module_path.parent)]
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.logging")
|
||||
|
||||
@@ -86,27 +86,25 @@ class ParseTilesFromFormTests(unittest.TestCase):
|
||||
|
||||
# --- happy paths ---
|
||||
|
||||
def test_single_tile_returns_one_entry(self):
|
||||
def test_single_tile_returns_image_and_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertEqual(len(tiles), 1)
|
||||
|
||||
def test_multiple_tiles_all_returned(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(len(tiles), 3)
|
||||
|
||||
def test_tile_image_is_pil_image(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
self.assertIsInstance(tiles[0]["image"], PILImage.Image)
|
||||
|
||||
def test_tile_metadata_fields_are_parsed(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1))
|
||||
tile = tiles[0]
|
||||
self.assertIsInstance(tile["image"], PILImage.Image)
|
||||
self.assertEqual(tile["tile_idx"], 0)
|
||||
self.assertEqual(tile["x"], 0)
|
||||
self.assertEqual(tile["y"], 0)
|
||||
self.assertEqual(tile["extracted_width"], 64)
|
||||
self.assertEqual(tile["extracted_height"], 64)
|
||||
|
||||
def test_multiple_tiles_preserve_count_order_and_coordinates(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(len(tiles), 3)
|
||||
for i, tile in enumerate(tiles):
|
||||
self.assertEqual(tile["tile_idx"], i)
|
||||
self.assertEqual(tiles[1]["x"], 64)
|
||||
self.assertEqual(tiles[2]["x"], 128)
|
||||
|
||||
def test_padding_is_parsed_from_form(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
|
||||
self.assertEqual(tiles[0]["padding"], 16)
|
||||
@@ -138,16 +136,6 @@ class ParseTilesFromFormTests(unittest.TestCase):
|
||||
self.assertNotIn("batch_idx", tiles[0])
|
||||
self.assertNotIn("global_idx", tiles[0])
|
||||
|
||||
def test_tile_indices_match_metadata_order(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
for i, tile in enumerate(tiles):
|
||||
self.assertEqual(tile["tile_idx"], i)
|
||||
|
||||
def test_x_coordinates_reflect_metadata(self):
|
||||
tiles = pp._parse_tiles_from_form(_make_form(3))
|
||||
self.assertEqual(tiles[1]["x"], 64)
|
||||
self.assertEqual(tiles[2]["x"], 128)
|
||||
|
||||
# --- error cases ---
|
||||
|
||||
def test_missing_tiles_metadata_raises_value_error(self):
|
||||
|
||||
@@ -302,8 +302,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
||||
self.assertNotIn("1", result)
|
||||
self.assertNotIn("2", result)
|
||||
|
||||
def test_removes_dangling_upstream_refs(self):
|
||||
"""Collector must not retain dangling refs to pruned upstream nodes."""
|
||||
def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
|
||||
"""Collector must point to exactly one valid placeholder after pruning."""
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
collector_inputs = result["3"].get("inputs", {})
|
||||
@@ -314,10 +314,6 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
||||
self.assertNotEqual(source_id, "2")
|
||||
self.assertIn(source_id, result)
|
||||
self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
|
||||
|
||||
def test_injects_empty_image_placeholder(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
|
||||
self.assertEqual(len(empty_nodes), 1)
|
||||
placeholder_id = empty_nodes[0][0]
|
||||
@@ -384,13 +380,6 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
|
||||
self.assertEqual(result["15"], prompt["15"])
|
||||
self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
|
||||
|
||||
def test_does_not_preserve_non_primitive_upstream_for_collector(self):
|
||||
prompt = _delegate_prompt()
|
||||
result = pt.prepare_delegate_master_prompt(prompt, ["3"])
|
||||
|
||||
self.assertNotIn("2", result)
|
||||
self.assertNotEqual(result["3"]["inputs"]["images"], ["2", 0])
|
||||
|
||||
def test_preserves_load_image_for_switch_alternate_required_input(self):
|
||||
"""Delegate-only master keeps LoadImage inputs needed by switches."""
|
||||
prompt = {
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
def load_module():
|
||||
spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
client = load_module()
|
||||
|
||||
|
||||
def workflow_for_request_size(target_size, timeout_seconds):
|
||||
"""Build an ASCII workflow whose compact UTF-8 request body has exact size."""
|
||||
workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}}
|
||||
|
||||
def body_size(candidate):
|
||||
payload = {
|
||||
"input": {"workflow": candidate, "images": []},
|
||||
"policy": {
|
||||
"executionTimeout": timeout_seconds * 1000,
|
||||
"ttl": timeout_seconds * 1000,
|
||||
},
|
||||
}
|
||||
return len(json.dumps(
|
||||
payload, allow_nan=False, ensure_ascii=False, separators=(",", ":")
|
||||
).encode("utf-8"))
|
||||
|
||||
padding = target_size - body_size(workflow)
|
||||
if padding < 0:
|
||||
raise AssertionError("target is smaller than the request fixture")
|
||||
workflow["1"]["inputs"]["value"] = "x" * padding
|
||||
if body_size(workflow) != target_size:
|
||||
raise AssertionError("request fixture did not reach the exact target size")
|
||||
return workflow
|
||||
|
||||
class Response:
|
||||
def __init__(self, status=200, payload=None):
|
||||
self.status, self.payload = status, payload
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *args):
|
||||
return False
|
||||
async def json(self):
|
||||
return self.payload
|
||||
|
||||
class Session:
|
||||
def __init__(self, responses):
|
||||
self.responses, self.calls = list(responses), []
|
||||
def request(self, method, url, **kwargs):
|
||||
self.calls.append((method, url, kwargs))
|
||||
response = self.responses.pop(0)
|
||||
if isinstance(response, Exception):
|
||||
raise response
|
||||
return response
|
||||
|
||||
class ClientTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_submit_exact_generic_contract_and_once(self):
|
||||
session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800)
|
||||
self.assertEqual(result, "job-1")
|
||||
method, url, kwargs = session.calls[0]
|
||||
self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run"))
|
||||
self.assertIsInstance(kwargs["data"], bytes)
|
||||
sent_payload = json.loads(kwargs["data"])
|
||||
self.assertEqual(set(sent_payload["input"]), {"workflow", "images"})
|
||||
self.assertEqual(sent_payload["policy"]["ttl"], 1800000)
|
||||
self.assertFalse(kwargs["allow_redirects"])
|
||||
self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value")
|
||||
|
||||
async def test_submit_timeout_is_ambiguous_and_never_retried(self):
|
||||
session = Session([TimeoutError("secret-value in unsafe error")])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
with self.assertRaises(client.AmbiguousSubmission) as ctx:
|
||||
await api.submit({}, [], 30)
|
||||
self.assertNotIn("secret-value", str(ctx.exception))
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
|
||||
async def test_http_submit_failure_not_retried_or_leaked(self):
|
||||
for status in [401, 429, 500, 302]:
|
||||
session = Session([Response(status, {"error": "secret-value"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "secret-value")
|
||||
with self.assertRaises(client.RunpodError) as ctx:
|
||||
await api.submit({}, [], 30)
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
self.assertNotIn("secret-value", str(ctx.exception))
|
||||
|
||||
async def test_health_zero_workers_is_valid(self):
|
||||
session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
self.assertEqual((await api.health())["workers"]["idle"], 0)
|
||||
|
||||
async def test_invalid_provider_id_no_network(self):
|
||||
session = Session([])
|
||||
with self.assertRaises(ValueError):
|
||||
client.RunpodClient(session, "https://attacker.invalid", "key")
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
with self.assertRaises(ValueError):
|
||||
await api.status("../other")
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_status_and_cancel_only_target_exact_job(self):
|
||||
session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
await api.status("job-1")
|
||||
await api.cancel("job-1")
|
||||
self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")])
|
||||
|
||||
async def test_missing_id_is_ambiguous(self):
|
||||
api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key")
|
||||
with self.assertRaises(client.AmbiguousSubmission):
|
||||
await api.submit({}, [], 30)
|
||||
|
||||
async def test_large_payload_rejected_before_submit(self):
|
||||
session = Session([])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
with self.assertRaises(ValueError):
|
||||
await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30)
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self):
|
||||
session = Session([Response(payload={"id": "boundary-job"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
timeout_seconds = 37
|
||||
workflow = workflow_for_request_size(10_000_000, timeout_seconds)
|
||||
|
||||
result = await api.submit(workflow, [], timeout_seconds)
|
||||
|
||||
self.assertEqual(result, "boundary-job")
|
||||
sent = session.calls[0][2]["data"]
|
||||
self.assertIsInstance(sent, bytes)
|
||||
self.assertEqual(len(sent), 10_000_000)
|
||||
self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000)
|
||||
|
||||
async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self):
|
||||
session = Session([Response(payload={"id": "must-not-be-used"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
workflow = workflow_for_request_size(10_000_001, 37)
|
||||
workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:]
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "10,000,001") as ctx:
|
||||
await api.submit(workflow, [], 37)
|
||||
|
||||
self.assertNotIn("private graph value", str(ctx.exception))
|
||||
self.assertFalse(session.calls)
|
||||
|
||||
async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self):
|
||||
session = Session([Response(payload={"id": "unicode-job"})])
|
||||
api = client.RunpodClient(session, "endpoint123", "key")
|
||||
workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}}
|
||||
|
||||
await api.submit(workflow, [], 41)
|
||||
|
||||
sent = session.calls[0][2]["data"]
|
||||
expected = json.dumps({
|
||||
"input": {"workflow": workflow, "images": []},
|
||||
"policy": {"executionTimeout": 41_000, "ttl": 41_000},
|
||||
}, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
self.assertEqual(sent, expected)
|
||||
self.assertGreater(len(sent), len(sent.decode("utf-8")))
|
||||
@@ -0,0 +1,202 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
|
||||
class _FakeManager:
|
||||
def __init__(self, *, pending=False, owned=True):
|
||||
self.pending = pending
|
||||
self.owned = owned
|
||||
self.pending_calls = []
|
||||
self.owned_calls = []
|
||||
self.cancel_calls = []
|
||||
|
||||
def is_pending(self, collector_job_id, participant_id):
|
||||
self.pending_calls.append((collector_job_id, participant_id))
|
||||
return self.pending
|
||||
|
||||
def owns(self, collector_job_id, participant_id):
|
||||
self.owned_calls.append((collector_job_id, participant_id))
|
||||
return self.owned
|
||||
|
||||
async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
|
||||
self.cancel_calls.append({
|
||||
"worker_id": worker_id,
|
||||
"execution_id": execution_id,
|
||||
"collector_job_id": collector_job_id,
|
||||
})
|
||||
return 1
|
||||
|
||||
|
||||
def _set_up_collector(module, job_id, queue, manager=None):
|
||||
module.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
||||
module.prompt_server.distributed_pending_jobs = {job_id: queue}
|
||||
if manager is not None:
|
||||
module.prompt_server.distributed_serverless = manager
|
||||
return module.DistributedCollectorNode()
|
||||
|
||||
|
||||
def test_serverless_pending_waits_past_normal_worker_timeout():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-pending"
|
||||
participant = "serverless:pool:1"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
probe_calls = []
|
||||
module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
|
||||
|
||||
async def deliver_after_provider_work():
|
||||
await asyncio.sleep(0.14)
|
||||
queue.put_nowait({
|
||||
"worker_id": participant,
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 2, 2, 2),
|
||||
"is_last": True,
|
||||
})
|
||||
|
||||
async def run_collection():
|
||||
producer = asyncio.create_task(deliver_after_provider_work())
|
||||
result = await collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids=f'["{participant}"]',
|
||||
delegate_only=True,
|
||||
)
|
||||
await producer
|
||||
return result
|
||||
|
||||
images, _audio = asyncio.run(run_collection())
|
||||
|
||||
assert tuple(images.shape) == (1, 2, 2, 2)
|
||||
assert manager.pending_calls
|
||||
assert probe_calls == []
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-error"
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"serverless_error": "Provider execution failed safely",
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.cancel_calls == [{
|
||||
"worker_id": None,
|
||||
"execution_id": None,
|
||||
"collector_job_id": job_id,
|
||||
}]
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-stopped"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=False, owned=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.owned_calls
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_interruption_cancels_serverless_job_and_removes_collector_queue():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-interrupted"
|
||||
queue = asyncio.Queue()
|
||||
manager = _FakeManager(pending=True)
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
interrupted = module.comfy.model_management.InterruptProcessingException("stop")
|
||||
module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
|
||||
|
||||
with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
|
||||
asyncio.run(collector.execute(
|
||||
images=None,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
delegate_only=True,
|
||||
))
|
||||
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-shape-mismatch"
|
||||
queue = asyncio.Queue()
|
||||
queue.put_nowait({
|
||||
"worker_id": "serverless:pool:1",
|
||||
"image_index": 0,
|
||||
"tensor": torch.ones(1, 3, 2, 2),
|
||||
"is_last": True,
|
||||
})
|
||||
manager = _FakeManager()
|
||||
collector = _set_up_collector(module, job_id, queue, manager)
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
|
||||
with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
|
||||
asyncio.run(collector.execute(
|
||||
images=local_images,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["serverless:pool:1"]',
|
||||
))
|
||||
|
||||
assert manager.cancel_calls[0]["collector_job_id"] == job_id
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
|
||||
|
||||
def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
|
||||
module = _load_collector_module()
|
||||
job_id = "collector-regular-timeout"
|
||||
queue = asyncio.Queue()
|
||||
collector = _set_up_collector(module, job_id, queue)
|
||||
module.get_worker_timeout_seconds = lambda: 0.01
|
||||
module.HEARTBEAT_INTERVAL = 0.1
|
||||
local_images = torch.zeros(1, 2, 2, 2)
|
||||
collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
|
||||
|
||||
images, _audio = asyncio.run(collector.execute(
|
||||
images=local_images,
|
||||
audio=None,
|
||||
multi_job_id=job_id,
|
||||
enabled_worker_ids='["worker-a"]',
|
||||
))
|
||||
|
||||
assert torch.equal(images, local_images)
|
||||
assert job_id not in module.prompt_server.distributed_pending_jobs
|
||||
@@ -0,0 +1,398 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
|
||||
def _load_serverless_config_module(*, without_fcntl=False):
|
||||
module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py"
|
||||
module_name = (
|
||||
"dist_serverless_config_no_fcntl_test_module"
|
||||
if without_fcntl
|
||||
else "dist_serverless_config_test_module"
|
||||
)
|
||||
package_name = module_name + "_pkg"
|
||||
package = types.ModuleType(package_name)
|
||||
package.__path__ = [str(module_path.parent)]
|
||||
sys.modules[package_name] = package
|
||||
module_name = package_name + ".serverless_config"
|
||||
sys.modules.pop(module_name, None)
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
if without_fcntl:
|
||||
with mock.patch.dict(sys.modules, {"fcntl": None}):
|
||||
spec.loader.exec_module(module)
|
||||
else:
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
serverless_config = _load_serverless_config_module()
|
||||
|
||||
|
||||
class ValidateServerlessWorkerTests(unittest.TestCase):
|
||||
def test_returns_copy_with_runpod_defaults(self):
|
||||
worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}
|
||||
|
||||
result = serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
self.assertEqual(
|
||||
result,
|
||||
{
|
||||
"id": "worker-1",
|
||||
"endpoint_id": "ep_123-abc",
|
||||
"name": "GPU",
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"jobs_per_run": 1,
|
||||
"concurrency": 1,
|
||||
"timeout_seconds": 1800,
|
||||
},
|
||||
)
|
||||
self.assertIsNot(result, worker)
|
||||
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
|
||||
|
||||
def test_accepts_inclusive_numeric_bounds(self):
|
||||
for jobs, concurrency, timeout in ((1, 1, 30), (16, 16, 86400)):
|
||||
with self.subTest(jobs=jobs, concurrency=concurrency, timeout=timeout):
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{
|
||||
"id": "worker-1",
|
||||
"endpoint_id": "ep-123",
|
||||
"jobs_per_run": jobs,
|
||||
"concurrency": concurrency,
|
||||
"timeout_seconds": timeout,
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["jobs_per_run"], jobs)
|
||||
self.assertEqual(result["concurrency"], concurrency)
|
||||
self.assertEqual(result["timeout_seconds"], timeout)
|
||||
|
||||
def test_rejects_missing_or_empty_id_and_endpoint(self):
|
||||
base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
|
||||
with self.subTest(field=field, value=value):
|
||||
worker = dict(base)
|
||||
worker[field] = value
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
worker = dict(base)
|
||||
del worker["id"]
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
def test_endpoint_id_must_be_ascii_opaque_identifier(self):
|
||||
for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"):
|
||||
with self.subTest(endpoint_id=endpoint_id):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": endpoint_id}
|
||||
)
|
||||
|
||||
def test_endpoint_id_matches_runpod_client_length_limit(self):
|
||||
valid_id = "e" * 200
|
||||
result = serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": valid_id}
|
||||
)
|
||||
self.assertEqual(result["endpoint_id"], valid_id)
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{"id": "worker-1", "endpoint_id": "e" * 201}
|
||||
)
|
||||
|
||||
def test_enabled_must_be_boolean_when_supplied(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for enabled in (True, False):
|
||||
with self.subTest(enabled=enabled):
|
||||
self.assertIs(
|
||||
serverless_config.validate_serverless_worker(
|
||||
{**base, "enabled": enabled}
|
||||
)["enabled"],
|
||||
enabled,
|
||||
)
|
||||
for enabled in ("false", 0, 1, None):
|
||||
with self.subTest(enabled=enabled):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(
|
||||
{**base, "enabled": enabled}
|
||||
)
|
||||
|
||||
def test_rejects_unknown_provider_and_non_serverless_type(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
for extra in ({"provider": "other"}, {"type": "local"}):
|
||||
with self.subTest(extra=extra):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker({**base, **extra})
|
||||
|
||||
def test_numeric_fields_reject_bool_float_and_out_of_range(self):
|
||||
base = {"id": "worker-1", "endpoint_id": "ep-123"}
|
||||
invalid_values = {
|
||||
"jobs_per_run": (True, 1.0, 0, 17),
|
||||
"concurrency": (False, 2.5, 0, 17),
|
||||
"timeout_seconds": (True, 30.0, 29, 86401),
|
||||
}
|
||||
for field, values in invalid_values.items():
|
||||
for value in values:
|
||||
with self.subTest(field=field, value=value):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker({**base, field: value})
|
||||
|
||||
def test_requires_worker_object(self):
|
||||
for worker in (None, [], "worker"):
|
||||
with self.subTest(worker=worker):
|
||||
with self.assertRaises(ValueError):
|
||||
serverless_config.validate_serverless_worker(worker)
|
||||
|
||||
|
||||
class SecretStoreTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self.tempdir.name)
|
||||
self.path = self.root / "private" / "runpod-secrets.json"
|
||||
|
||||
def tearDown(self):
|
||||
self.tempdir.cleanup()
|
||||
|
||||
def test_set_get_has_delete_and_persists_as_json_mapping(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
store.set("worker-1", "secret-value")
|
||||
|
||||
self.assertEqual(store.get("worker-1"), "secret-value")
|
||||
self.assertTrue(store.has("worker-1"))
|
||||
self.assertIsNone(store.get("missing"))
|
||||
self.assertFalse(store.has("missing"))
|
||||
self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"})
|
||||
self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value")
|
||||
|
||||
store.delete("worker-1")
|
||||
self.assertIsNone(store.get("worker-1"))
|
||||
self.assertFalse(store.has("worker-1"))
|
||||
|
||||
def test_rejects_api_keys_with_newlines_or_over_4096_characters(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097):
|
||||
with self.subTest(key_length=len(key)):
|
||||
with self.assertRaises(serverless_config.SecretStoreError):
|
||||
store.set("worker-1", key)
|
||||
store.set("worker-1", "k" * 4096)
|
||||
self.assertEqual(store.get("worker-1"), "k" * 4096)
|
||||
|
||||
def test_new_secret_file_and_directory_have_restrictive_modes(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
store.set("worker-1", "secret-value")
|
||||
|
||||
self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600)
|
||||
self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700)
|
||||
|
||||
def test_missing_store_reads_as_empty_without_creating_files(self):
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
self.assertIsNone(store.get("worker-1"))
|
||||
self.assertFalse(store.has("worker-1"))
|
||||
self.assertFalse(self.path.parent.exists())
|
||||
|
||||
def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self):
|
||||
self.path.parent.mkdir()
|
||||
corrupt = b'{"partial": "NEVER-LEAK-THIS-secret'
|
||||
self.path.write_bytes(corrupt)
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
for operation in (
|
||||
lambda: store.get("worker-1"),
|
||||
lambda: store.set("worker-1", "replacement-secret"),
|
||||
lambda: store.delete("worker-1"),
|
||||
):
|
||||
with self.subTest(operation=operation):
|
||||
with self.assertRaises(Exception) as raised:
|
||||
operation()
|
||||
self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception))
|
||||
self.assertNotIn("replacement-secret", str(raised.exception))
|
||||
self.assertEqual(self.path.read_bytes(), corrupt)
|
||||
|
||||
def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self):
|
||||
for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'):
|
||||
with self.subTest(corrupt=corrupt):
|
||||
self.path.parent.mkdir(exist_ok=True)
|
||||
self.path.write_bytes(corrupt)
|
||||
with self.assertRaises(Exception) as raised:
|
||||
serverless_config.SecretStore(self.path).get("worker")
|
||||
self.assertNotIn("opaque-stored-value", str(raised.exception))
|
||||
|
||||
def test_rejects_symlink_file_and_preserves_link_target(self):
|
||||
target = self.root / "target.json"
|
||||
target.write_text('{"worker-1": "target-secret"}', encoding="utf-8")
|
||||
self.path.parent.mkdir()
|
||||
self.path.symlink_to(target)
|
||||
store = serverless_config.SecretStore(self.path)
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.get("worker-1")
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "replacement-secret")
|
||||
self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}')
|
||||
|
||||
def test_rejects_symlinked_parent_directory(self):
|
||||
real_dir = self.root / "real"
|
||||
real_dir.mkdir()
|
||||
link_dir = self.root / "linked"
|
||||
link_dir.symlink_to(real_dir, target_is_directory=True)
|
||||
store = serverless_config.SecretStore(link_dir / "secrets.json")
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "secret-value")
|
||||
self.assertFalse((real_dir / "secrets.json").exists())
|
||||
|
||||
def test_rejects_parent_traversal_in_store_path(self):
|
||||
store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json")
|
||||
|
||||
with self.assertRaises(Exception):
|
||||
store.set("worker-1", "secret-value")
|
||||
self.assertFalse((self.root / "secrets.json").exists())
|
||||
|
||||
def test_threaded_updates_from_independent_instances_are_not_lost(self):
|
||||
workers = 12
|
||||
barrier = threading.Barrier(workers)
|
||||
errors = []
|
||||
|
||||
def write_one(index):
|
||||
try:
|
||||
barrier.wait(timeout=5)
|
||||
serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}")
|
||||
except Exception as exc: # collected to fail in the test thread
|
||||
errors.append(exc)
|
||||
|
||||
threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join(timeout=10)
|
||||
|
||||
self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish")
|
||||
self.assertEqual(errors, [])
|
||||
saved = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)})
|
||||
|
||||
|
||||
class PublicConfigTests(unittest.TestCase):
|
||||
def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self):
|
||||
with tempfile.TemporaryDirectory() as tempdir:
|
||||
store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json")
|
||||
store.set("serverless-1", "stored-api-key")
|
||||
config = {
|
||||
"workers": [
|
||||
{
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"provider": "runpod",
|
||||
"endpoint_id": "ep-123",
|
||||
"api_key": "inline-secret",
|
||||
"clear_api_key": "another-secret",
|
||||
"credential": "credential-secret",
|
||||
"nested": {"runpod_credential": "nested-secret", "safe": "kept"},
|
||||
"name": "remote GPU",
|
||||
},
|
||||
{"id": "serverless-2", "type": "serverless", "name": "not configured"},
|
||||
{"id": "local-1", "type": "local", "api_key": "local-secret"},
|
||||
],
|
||||
"settings": {"mode": "preserved", "credentials_note": "remove-me"},
|
||||
}
|
||||
original = json.loads(json.dumps(config))
|
||||
|
||||
result = serverless_config.public_config(config, store)
|
||||
|
||||
first, second, local = result["workers"]
|
||||
self.assertEqual(first["api_key_configured"], True)
|
||||
self.assertEqual(second["api_key_configured"], False)
|
||||
self.assertNotIn("api_key_configured", local)
|
||||
self.assertNotIn("api_key", first)
|
||||
self.assertNotIn("clear_api_key", first)
|
||||
self.assertNotIn("credential", first)
|
||||
self.assertNotIn("runpod_credential", first["nested"])
|
||||
self.assertEqual(first["nested"]["safe"], "kept")
|
||||
self.assertNotIn("api_key", local)
|
||||
self.assertNotIn("credentials_note", result["settings"])
|
||||
self.assertEqual(first["name"], "remote GPU")
|
||||
self.assertEqual(config, original)
|
||||
self.assertIsNot(result, config)
|
||||
self.assertIsNot(result["workers"], config["workers"])
|
||||
|
||||
def test_redacts_normalized_secret_field_aliases(self):
|
||||
config = {
|
||||
"workers": [{
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"apiKey": "camel-api-key",
|
||||
"APIKEY": "upper-api-key",
|
||||
"clearApiKey": "camel-clear-key",
|
||||
"CLEAR_API_KEY": "upper-clear-key",
|
||||
"runpod_CREDENTIAL": "provider-credential",
|
||||
"nested": {"credential_note": "nested-credential"},
|
||||
"safe": "preserved",
|
||||
}]
|
||||
}
|
||||
|
||||
result = serverless_config.public_config(
|
||||
config, serverless_config.SecretStore(self.root_path())
|
||||
)
|
||||
|
||||
worker = result["workers"][0]
|
||||
self.assertEqual(worker, {
|
||||
"id": "serverless-1",
|
||||
"type": "serverless",
|
||||
"nested": {},
|
||||
"safe": "preserved",
|
||||
"api_key_configured": False,
|
||||
})
|
||||
|
||||
def test_handles_non_dict_config_without_mutating_it(self):
|
||||
config = ["safe", {"api_key": "secret"}]
|
||||
result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path()))
|
||||
self.assertEqual(result, ["safe", {}])
|
||||
self.assertEqual(config, ["safe", {"api_key": "secret"}])
|
||||
|
||||
def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self):
|
||||
module = _load_serverless_config_module(without_fcntl=True)
|
||||
store = module.SecretStore(self.root_path())
|
||||
config = {
|
||||
"workers": [
|
||||
{"id": "local-1", "type": "local", "api_key": "inline-secret"}
|
||||
]
|
||||
}
|
||||
|
||||
result = module.public_config(config, store)
|
||||
|
||||
self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]})
|
||||
self.assertEqual(config["workers"][0]["api_key"], "inline-secret")
|
||||
|
||||
def test_secret_operations_fail_clearly_without_fcntl(self):
|
||||
module = _load_serverless_config_module(without_fcntl=True)
|
||||
store = module.SecretStore(self.root_path())
|
||||
unsupported_message = "Serverless credential storage is currently supported on POSIX masters only"
|
||||
|
||||
for operation in (
|
||||
lambda: store.get("worker-1"),
|
||||
lambda: store.set("worker-1", "api-key"),
|
||||
lambda: store.delete("worker-1"),
|
||||
lambda: store.has("worker-1"),
|
||||
):
|
||||
with self.subTest(operation=operation):
|
||||
with self.assertRaisesRegex(module.SecretStoreError, unsupported_message):
|
||||
operation()
|
||||
|
||||
@staticmethod
|
||||
def root_path():
|
||||
return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Offline integration: real graph preparation + orchestration, fake provider."""
|
||||
import asyncio
|
||||
import aiohttp as real_aiohttp
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_serverless_integration"
|
||||
|
||||
|
||||
def module(name, **attrs):
|
||||
value = types.ModuleType(name)
|
||||
value.__dict__.update(attrs)
|
||||
return value
|
||||
|
||||
|
||||
def load_modules():
|
||||
modules = {}
|
||||
for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]:
|
||||
modules[PKG + part] = module(PKG + part, __path__=[str(directory)])
|
||||
ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn))
|
||||
modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps))
|
||||
modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None)
|
||||
modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None)
|
||||
modules[PKG + ".utils.config"] = module("config", load_config=lambda: {})
|
||||
modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock())
|
||||
modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master")
|
||||
modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock())
|
||||
modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), sync_worker_media=AsyncMock())
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py")
|
||||
q = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(q)
|
||||
# Keep imported package namespace only; restore shared Comfy runtime stub.
|
||||
modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)})
|
||||
return q, modules, ps
|
||||
|
||||
|
||||
class IntegrationTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.q, modules, self.ps = load_modules()
|
||||
self.patch = patch.dict(sys.modules, modules)
|
||||
self.patch.start()
|
||||
self.addCleanup(self.patch.stop)
|
||||
self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30}
|
||||
self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}}
|
||||
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
|
||||
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
||||
self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None)
|
||||
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager)
|
||||
sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[]))
|
||||
self.prompt = {
|
||||
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
|
||||
}
|
||||
|
||||
async def test_real_compiler_slots_to_master_collector_without_callbacks(self):
|
||||
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.assertEqual(result[2], 2)
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.manager.start.assert_awaited_once()
|
||||
jobs = self.manager.start.call_args.args[2]
|
||||
self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"])
|
||||
self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2)
|
||||
for job in jobs:
|
||||
graph = job["workflow"]
|
||||
self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values()))
|
||||
self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values()))
|
||||
self.assertNotIn("http://master", json.dumps(graph))
|
||||
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
||||
self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"])
|
||||
|
||||
async def test_compile_pack_provider_decode_and_collect_real_components(self):
|
||||
"""Exercise the whole image path, mocking only the provider and Comfy queue."""
|
||||
import base64
|
||||
import io
|
||||
import tempfile
|
||||
import torch
|
||||
from PIL import Image
|
||||
from test_collector_list_inputs import _load_collector_module
|
||||
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
for leaf in ("serverless_media", "serverless_runtime"):
|
||||
name = PKG + ".api.orchestration." + leaf
|
||||
spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py"))
|
||||
real_module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = real_module
|
||||
spec.loader.exec_module(real_module)
|
||||
runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
image_buffer = io.BytesIO()
|
||||
Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG")
|
||||
image_bytes = image_buffer.getvalue()
|
||||
(root / "source.png").write_bytes(image_bytes)
|
||||
paths = module("folder_paths", get_input_directory=lambda: directory,
|
||||
get_output_directory=lambda: directory, get_temp_directory=lambda: directory,
|
||||
get_annotated_filepath=lambda name: str(root / name))
|
||||
submissions = []
|
||||
|
||||
class FakeProvider:
|
||||
async def submit(self, graph, images, timeout_seconds):
|
||||
self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage")
|
||||
submissions.append((graph, images))
|
||||
self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage")
|
||||
assert self_image == images[0]["name"]
|
||||
assert base64.b64decode(images[0]["image"]) == image_bytes
|
||||
return "provider-job-" + str(len(submissions))
|
||||
|
||||
async def status(self, _job_id):
|
||||
filename = self.prefix + "_00001_.png"
|
||||
return {"status": "COMPLETED", "output": {"images": [{
|
||||
"filename": filename, "artifact_path": filename, "artifact_type": "output",
|
||||
"type": "base64", "data": base64.b64encode(image_bytes).decode(),
|
||||
}]}}
|
||||
|
||||
manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None)
|
||||
manager._client = AsyncMock(side_effect=lambda _: FakeProvider())
|
||||
self.ps.distributed_serverless = manager
|
||||
self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}}
|
||||
collector_module = _load_collector_module()
|
||||
collector_module.prompt_server = self.ps
|
||||
sys.modules["aiohttp"] = real_aiohttp
|
||||
with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}):
|
||||
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
||||
inputs = master_graph["2"]["inputs"]
|
||||
images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute(
|
||||
images=None, audio=None, multi_job_id=inputs["multi_job_id"],
|
||||
enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True,
|
||||
), timeout=3)
|
||||
await asyncio.gather(*manager.tasks.values())
|
||||
self.assertEqual(result[2], 2)
|
||||
self.assertEqual(len(submissions), 2)
|
||||
self.assertEqual(tuple(images.shape), (2, 3, 2, 3))
|
||||
self.assertEqual(images.dtype, torch.float32)
|
||||
self.assertTrue(torch.equal(images[0], images[1]))
|
||||
self.assertEqual(float(images[0, 0, 0, 0]), 1.0)
|
||||
self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"])
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
|
||||
async def test_audio_rejected_before_any_dispatch_or_master_queue(self):
|
||||
self.prompt["2"]["inputs"]["audio"] = ["9", 0]
|
||||
self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}}
|
||||
with self.assertRaises(ValueError):
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.manager.start.assert_not_awaited()
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
|
||||
async def test_disabled_pool_excluded(self):
|
||||
self.pool["enabled"] = False
|
||||
self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), [])
|
||||
|
||||
async def test_master_validation_failure_submits_no_paid_job(self):
|
||||
self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph")
|
||||
with self.assertRaisesRegex(ValueError, "invalid local graph"):
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
self.manager.start.assert_not_awaited()
|
||||
|
||||
async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self):
|
||||
media_module = sys.modules[PKG + ".api.orchestration.serverless_media"]
|
||||
media_module.prepare_images.return_value = [{
|
||||
"name": "private-image-name.png",
|
||||
"image": "x" * 10_000_001,
|
||||
}]
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "10,000,") as ctx:
|
||||
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
||||
|
||||
self.assertNotIn("private-image-name.png", str(ctx.exception))
|
||||
self.assertNotIn("x" * 100, str(ctx.exception))
|
||||
self.q.dispatch_worker_prompt.assert_not_awaited()
|
||||
self.q.queue_prompt_payload.assert_not_awaited()
|
||||
self.manager.start.assert_not_awaited()
|
||||
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
||||
@@ -0,0 +1,192 @@
|
||||
"""Billing-safe lifecycle regressions: local completion is not remote completion."""
|
||||
import asyncio
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from test_serverless_runtime import runtime, Secrets
|
||||
|
||||
|
||||
class LifecycleTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
|
||||
self.path = Path(self.temp.name) / "jobs.json"
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
|
||||
self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]}
|
||||
self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
|
||||
self.client_patch.start()
|
||||
self.addCleanup(self.client_patch.stop)
|
||||
self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
|
||||
self.decode_patch.start()
|
||||
self.addCleanup(self.decode_patch.stop)
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
|
||||
def second(self, **changes):
|
||||
return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes))
|
||||
|
||||
def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"):
|
||||
self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id,
|
||||
status=status, created_at=1, updated_at=1)
|
||||
|
||||
async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self):
|
||||
cancel_started, release_cancel = asyncio.Event(), asyncio.Event()
|
||||
self.api.submit.side_effect = ["job-1", "job-2"]
|
||||
self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}]
|
||||
async def cancel(_job_id):
|
||||
cancel_started.set()
|
||||
await release_cancel.wait()
|
||||
return {"status": "CANCELLED"}
|
||||
self.api.cancel.side_effect = cancel
|
||||
await self.manager.start("exec", "prompt", [self.prepared, self.second()])
|
||||
try:
|
||||
await asyncio.wait_for(cancel_started.wait(), 1)
|
||||
await asyncio.sleep(0.01)
|
||||
self.assertEqual(self.api.submit.await_count, 1)
|
||||
finally:
|
||||
release_cancel.set()
|
||||
await self.manager.wait_for_execution("exec")
|
||||
self.assertEqual(self.api.submit.await_count, 2)
|
||||
|
||||
async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self):
|
||||
self.api.status.side_effect = runtime.RunpodError("poll failed")
|
||||
self.api.cancel.return_value = {"status": "IN_PROGRESS"}
|
||||
await self.manager.start("first", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("first")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
await self.manager.start("second", "prompt", [self.second()])
|
||||
try:
|
||||
await asyncio.sleep(0.01)
|
||||
self.assertEqual(self.api.submit.await_count, 1)
|
||||
self.api.status.side_effect = None
|
||||
self.api.status.return_value = {"status": "COMPLETED", "output": {}}
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
await asyncio.wait_for(self.manager.wait_for_execution("second"), 1)
|
||||
self.assertEqual(self.api.submit.await_count, 2)
|
||||
finally:
|
||||
self.api.cancel.return_value = {"status": "CANCELLED"}
|
||||
await self.manager.cancel()
|
||||
|
||||
async def test_restarted_manager_reserves_known_unresolved_jobs(self):
|
||||
self.retained("IN_PROGRESS")
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertFalse(self.manager.pool_is_active("other"))
|
||||
await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)])
|
||||
await self.manager.wait_for_execution("exec")
|
||||
self.api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT")
|
||||
|
||||
async def test_recovery_closes_never_submitted_waiting_records(self):
|
||||
self.retained("WAITING", job_id=None)
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_recovery_keeps_interrupted_submission_unresolved(self):
|
||||
self.retained("SUBMITTING", job_id=None)
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.api.status.assert_not_awaited()
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self):
|
||||
self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown")
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("exec")
|
||||
record = self.manager.ledger.all()[0]
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False)
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True)
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED")
|
||||
self.api.submit.assert_awaited_once()
|
||||
|
||||
async def test_cancel_already_completed_provider_job_is_terminal(self):
|
||||
self.retained()
|
||||
self.api.cancel.return_value = {"status": "COMPLETED"}
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
|
||||
async def test_reconciliation_failure_never_releases_capacity(self):
|
||||
self.retained()
|
||||
self.api.status.side_effect = runtime.RunpodError("not reachable")
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
|
||||
|
||||
async def test_cancel_before_first_task_step_finalizes_local_record(self):
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
self.assertFalse(self.server.distributed_pending_jobs["collector"].empty())
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_live_record_cannot_be_acknowledged_or_reconciled(self):
|
||||
polling = asyncio.Event()
|
||||
async def status(_job_id):
|
||||
polling.set()
|
||||
await asyncio.Event().wait()
|
||||
self.api.status.side_effect = status
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await asyncio.wait_for(polling.wait(), 1)
|
||||
record = self.manager.ledger.all()[0]
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
|
||||
await self.manager.reconcile(worker_id="pool")
|
||||
self.assertEqual(self.api.status.await_count, 1)
|
||||
self.assertTrue(self.manager.pool_is_active("pool"))
|
||||
|
||||
async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self):
|
||||
polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
||||
async def status(_job_id):
|
||||
polling.set()
|
||||
await asyncio.Event().wait()
|
||||
async def cancel(_job_id):
|
||||
cancelling.set()
|
||||
await release.wait()
|
||||
return {"status": "CANCELLED"}
|
||||
self.api.status.side_effect = status
|
||||
self.api.cancel.side_effect = cancel
|
||||
await self.manager.start("exec", "prompt", [self.prepared])
|
||||
await polling.wait()
|
||||
first = asyncio.create_task(self.manager.cancel(worker_id="pool"))
|
||||
await cancelling.wait()
|
||||
second = asyncio.create_task(self.manager.cancel(worker_id="pool"))
|
||||
await asyncio.sleep(0)
|
||||
release.set()
|
||||
await asyncio.gather(first, second)
|
||||
self.api.cancel.assert_awaited_once_with("job")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_partial_start_cleanup_does_not_leave_waiting_protection(self):
|
||||
with self.assertRaises(ValueError):
|
||||
await self.manager.start("exec", "prompt", [self.prepared, self.second(concurrency=2)])
|
||||
await self.manager.cancel(execution_id="exec")
|
||||
self.assertFalse(self.manager.pool_is_active("pool"))
|
||||
self.api.submit.assert_not_awaited()
|
||||
|
||||
async def test_public_jobs_never_hides_unresolved_recovery_ids(self):
|
||||
for index in range(22):
|
||||
self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index)
|
||||
for index in range(5):
|
||||
self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index)
|
||||
records = self.manager.public_jobs("pool")
|
||||
self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)})
|
||||
self.assertTrue(all(record["status"] == "UNKNOWN" for record in records))
|
||||
@@ -0,0 +1,447 @@
|
||||
import base64
|
||||
import importlib.util
|
||||
import io
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from PIL import Image
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py"
|
||||
|
||||
|
||||
def load_module():
|
||||
spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
media = load_module()
|
||||
|
||||
|
||||
def load_runpod_module():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py"
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
runpod = load_runpod_module()
|
||||
|
||||
|
||||
def png_bytes(mode="RGB", size=(2, 1), color=None):
|
||||
if color is None:
|
||||
color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30)
|
||||
buffer = io.BytesIO()
|
||||
Image.new(mode, size, color).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class FolderPathsFixture:
|
||||
def __init__(self, root):
|
||||
self.input = root / "input"
|
||||
self.output = root / "output"
|
||||
self.temp = root / "temp"
|
||||
for directory in (self.input, self.output, self.temp):
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_input_directory(self):
|
||||
return str(self.input)
|
||||
|
||||
def get_output_directory(self):
|
||||
return str(self.output)
|
||||
|
||||
def get_temp_directory(self):
|
||||
return str(self.temp)
|
||||
|
||||
def get_annotated_filepath(self, filename):
|
||||
value = filename.replace("\\", "/")
|
||||
# Match real ComfyUI folder_paths: root annotations are suffixes,
|
||||
# not directory prefixes (output/foo.png is an input subdirectory).
|
||||
for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)):
|
||||
suffix = f" [{marker}]"
|
||||
if value.endswith(suffix):
|
||||
return str(root / value[:-len(suffix)])
|
||||
return str(self.input / value)
|
||||
|
||||
|
||||
class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.paths = FolderPathsFixture(Path(self.tempdir.name))
|
||||
self.folder_paths = types.ModuleType("folder_paths")
|
||||
self.folder_paths.get_input_directory = self.paths.get_input_directory
|
||||
self.folder_paths.get_output_directory = self.paths.get_output_directory
|
||||
self.folder_paths.get_temp_directory = self.paths.get_temp_directory
|
||||
self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath
|
||||
self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths})
|
||||
self.folder_patch.start()
|
||||
self.addCleanup(self.folder_patch.stop)
|
||||
self.addCleanup(self.tempdir.cleanup)
|
||||
|
||||
async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self):
|
||||
payload = png_bytes()
|
||||
(self.paths.input / "cat.png").write_bytes(payload)
|
||||
workflow = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}},
|
||||
"3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}},
|
||||
}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 1)
|
||||
self.assertEqual(set(packed[0]), {"name", "image"})
|
||||
self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$")
|
||||
self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload)
|
||||
self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"])
|
||||
self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"])
|
||||
self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged")
|
||||
self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable")
|
||||
|
||||
async def test_resolves_annotated_input_output_and_temp_sources(self):
|
||||
sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp"))
|
||||
workflow = {}
|
||||
for index, (directory, filename) in enumerate(sources, 1):
|
||||
(directory / filename).write_bytes(png_bytes())
|
||||
root = directory.name
|
||||
workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 3)
|
||||
self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()})
|
||||
self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed})
|
||||
|
||||
async def test_directory_prefix_is_not_a_root_annotation(self):
|
||||
(self.paths.input / "output").mkdir()
|
||||
payload = png_bytes(color=(100, 0, 0))
|
||||
(self.paths.input / "output" / "same.png").write_bytes(payload)
|
||||
(self.paths.output / "same.png").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}}
|
||||
packed = await media.prepare_images(workflow)
|
||||
self.assertEqual(base64.b64decode(packed[0]["image"]), payload)
|
||||
|
||||
async def test_rejects_unknown_root_annotation(self):
|
||||
(self.paths.input / "cat.png").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_recognizes_custom_file_image_reference(self):
|
||||
(self.paths.input / "source.bmp").write_bytes(png_bytes())
|
||||
workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
|
||||
self.assertEqual(len(packed), 1)
|
||||
self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"])
|
||||
|
||||
async def test_rejects_missing_and_dynamic_loadimage_values(self):
|
||||
cases = (
|
||||
{"1": {"class_type": "LoadImage", "inputs": {}}},
|
||||
{"1": {"class_type": "LoadImage"}},
|
||||
{"1": {"class_type": "LoadImage", "inputs": None}},
|
||||
{"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}},
|
||||
)
|
||||
for workflow in cases:
|
||||
with self.subTest(workflow=workflow):
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_missing_source_and_unapproved_extension(self):
|
||||
for value in ("absent.png", "animation.gif", "clip.mp4"):
|
||||
with self.subTest(value=value):
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_unsupported_audio_or_video_references_without_reading(self):
|
||||
for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")):
|
||||
with self.subTest(key=key):
|
||||
workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_parent_traversal_and_absolute_outside_roots(self):
|
||||
outside = Path(self.tempdir.name).parent / "outside.png"
|
||||
outside.write_bytes(png_bytes())
|
||||
self.addCleanup(lambda: outside.unlink(missing_ok=True))
|
||||
for value in ("../outside.png", str(outside)):
|
||||
with self.subTest(value=value):
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_symlink_escape_from_allowed_root(self):
|
||||
outside = Path(self.tempdir.name).parent / "linked-outside.png"
|
||||
outside.write_bytes(png_bytes())
|
||||
self.addCleanup(lambda: outside.unlink(missing_ok=True))
|
||||
(self.paths.input / "link.png").symlink_to(outside)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}}
|
||||
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
async def test_rejects_individual_and_aggregate_input_size_overages(self):
|
||||
huge = self.paths.input / "large.jpg"
|
||||
with huge.open("wb") as handle:
|
||||
handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
huge.unlink()
|
||||
first = self.paths.input / "first.png"
|
||||
second = self.paths.input / "second.png"
|
||||
first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
|
||||
second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
|
||||
workflow = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}},
|
||||
"2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}},
|
||||
}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.prepare_images(workflow)
|
||||
|
||||
|
||||
async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self):
|
||||
raw_size = 7 * 1024 * 1024 + 1
|
||||
source = self.paths.input / "large-input.png"
|
||||
source.write_bytes(b"x" * raw_size)
|
||||
workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}}
|
||||
|
||||
packed = await media.prepare_images(workflow)
|
||||
request_body = runpod.serialize_request(workflow, packed, 30)
|
||||
|
||||
self.assertGreater(raw_size, 7 * 1024 * 1024)
|
||||
self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size)
|
||||
self.assertLessEqual(len(request_body), 10_000_000)
|
||||
|
||||
|
||||
class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase):
|
||||
def manifest(self, prefix="render"):
|
||||
return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}]
|
||||
|
||||
def artifact(self, filename, payload=None, **extra):
|
||||
entry = {
|
||||
"filename": filename,
|
||||
"artifact_type": "output",
|
||||
"artifact_path": f"/safe/output/{filename}",
|
||||
"subfolder": "",
|
||||
"type": "base64",
|
||||
"data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"),
|
||||
}
|
||||
entry.update(extra)
|
||||
return entry
|
||||
|
||||
async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self):
|
||||
output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(list(result), ["job-1"])
|
||||
self.assertEqual(len(result["job-1"]), 2)
|
||||
self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"]))
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
|
||||
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255)
|
||||
|
||||
async def test_groups_images_under_manifest_job_ids(self):
|
||||
manifest = [
|
||||
self.manifest("first")[0],
|
||||
{"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"},
|
||||
]
|
||||
output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]}
|
||||
|
||||
result = await media.collect_images(output, manifest)
|
||||
|
||||
self.assertEqual(set(result), {"job-1", "job-2"})
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
self.assertEqual(len(result["job-2"]), 1)
|
||||
|
||||
async def test_preserves_rgba_alpha(self):
|
||||
output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 4))
|
||||
self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255)
|
||||
|
||||
async def test_signed_url_is_downloaded_only_through_injected_downloader(self):
|
||||
output = {"images": [{
|
||||
"filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url",
|
||||
"url": "https://storage.example.test/signed",
|
||||
}]}
|
||||
fetch = AsyncMock(return_value=png_bytes())
|
||||
|
||||
result = await media.collect_images(output, self.manifest(), download=fetch)
|
||||
|
||||
self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
|
||||
fetch.assert_awaited_once_with("https://storage.example.test/signed")
|
||||
|
||||
async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self):
|
||||
output = {"images": [
|
||||
self.artifact("preview_00001_.png"),
|
||||
self.artifact("render_00001_.png"),
|
||||
{"filename": "mesh.obj", "type": "url", "url": "not a fetch target"},
|
||||
]}
|
||||
|
||||
result = await media.collect_images(output, self.manifest())
|
||||
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"):
|
||||
await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest())
|
||||
|
||||
async def test_accepts_compiler_length_flat_safe_prefix(self):
|
||||
prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64
|
||||
manifest = self.manifest(prefix)
|
||||
output = {"images": [self.artifact(f"{prefix}_00001_.png")]}
|
||||
|
||||
result = await media.collect_images(output, manifest)
|
||||
|
||||
self.assertEqual(len(result["job-1"]), 1)
|
||||
|
||||
async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self):
|
||||
duplicate = self.manifest() * 2
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": []}, duplicate)
|
||||
unsafe = self.manifest("../render")
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": []}, unsafe)
|
||||
|
||||
async def test_rejects_duplicate_canonical_artifact_identity(self):
|
||||
item = self.artifact("render_00001_.png")
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"):
|
||||
await media.collect_images({"images": [item, dict(item)]}, self.manifest())
|
||||
|
||||
async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self):
|
||||
wrong_type = self.artifact("render_00001_.png", artifact_type="temp")
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [wrong_type]}, self.manifest())
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest())
|
||||
|
||||
async def test_rejects_provider_errors_even_when_job_completed(self):
|
||||
for output in (
|
||||
{"error": "raw provider secret", "images": []},
|
||||
{"errors": ["raw provider secret"], "images": []},
|
||||
{"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}},
|
||||
):
|
||||
with self.subTest(output=output):
|
||||
with self.assertRaises(media.ServerlessMediaError) as ctx:
|
||||
await media.collect_images(output, self.manifest())
|
||||
self.assertNotIn("raw provider secret", str(ctx.exception))
|
||||
|
||||
async def test_rejects_invalid_base64_and_malformed_image(self):
|
||||
for item in (
|
||||
self.artifact("render_00001_.png", data="!not base64!"),
|
||||
self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()),
|
||||
):
|
||||
with self.subTest(item=item):
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images({"images": [item]}, self.manifest())
|
||||
|
||||
async def test_validates_shape_consistency_within_collector(self):
|
||||
output = {"images": [
|
||||
self.artifact("render_00001_.png", png_bytes(size=(2, 1))),
|
||||
self.artifact("render_00002_.png", png_bytes(size=(3, 1))),
|
||||
]}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_bounds_output_file_count_and_bytes(self):
|
||||
too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "files"):
|
||||
await media.collect_images(too_many, self.manifest())
|
||||
|
||||
with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4):
|
||||
output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]}
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self):
|
||||
output = {"images": [
|
||||
self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()),
|
||||
self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()),
|
||||
]}
|
||||
with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3):
|
||||
with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
|
||||
await media.collect_images(output, self.manifest())
|
||||
|
||||
async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self):
|
||||
for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"):
|
||||
with self.subTest(url=url):
|
||||
fetch = AsyncMock(return_value=png_bytes())
|
||||
output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]}
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.collect_images(output, self.manifest(), download=fetch)
|
||||
fetch.assert_not_awaited()
|
||||
|
||||
async def test_public_resolver_rejects_non_global_dns_answers(self):
|
||||
resolver = media._PublicResolver()
|
||||
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))])
|
||||
self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1)
|
||||
resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))])
|
||||
with self.assertRaises(OSError):
|
||||
await resolver.resolve("cdn.example.test", 443)
|
||||
|
||||
async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self):
|
||||
class Content:
|
||||
def __init__(self, chunks):
|
||||
self.chunks = iter(chunks)
|
||||
async def iter_chunked(self, _size):
|
||||
for chunk in self.chunks:
|
||||
yield chunk
|
||||
|
||||
class Response:
|
||||
def __init__(self, status, headers=None, chunks=()):
|
||||
self.status, self.headers = status, headers or {}
|
||||
self.content = Content(chunks)
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
|
||||
class Session:
|
||||
def __init__(self, responses):
|
||||
self.responses, self.calls = list(responses), []
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
async def __aexit__(self, *_args):
|
||||
return False
|
||||
def get(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.responses.pop(0)
|
||||
|
||||
session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])])
|
||||
with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session:
|
||||
payload = await media.download_public_artifact("https://storage.example.test/signed")
|
||||
|
||||
self.assertEqual(payload, b"abcdef")
|
||||
self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"])
|
||||
self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls))
|
||||
self.assertFalse(create_session.call_args.kwargs["trust_env"])
|
||||
self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver)
|
||||
self.assertNotIn("Authorization", session.calls[0][1].get("headers", {}))
|
||||
await create_session.call_args.kwargs["connector"].close()
|
||||
|
||||
blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})])
|
||||
with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session:
|
||||
with self.assertRaises(media.ServerlessMediaError):
|
||||
await media.download_public_artifact("https://storage.example.test/signed")
|
||||
self.assertEqual(len(blocked.calls), 1)
|
||||
await blocked_session.call_args.kwargs["connector"].close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,392 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
MODULE_PATH = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "api"
|
||||
/ "orchestration"
|
||||
/ "serverless_prompt.py"
|
||||
)
|
||||
_spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH)
|
||||
assert _spec is not None and _spec.loader is not None
|
||||
_serverless_prompt = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(_serverless_prompt)
|
||||
compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt
|
||||
|
||||
|
||||
def _prepared_prompt():
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
|
||||
"3": {
|
||||
"class_type": "DistributedSeed",
|
||||
"inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"},
|
||||
},
|
||||
"4": {
|
||||
"class_type": "DistributedCollector",
|
||||
"inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True},
|
||||
},
|
||||
"5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
"7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}},
|
||||
"8": {"class_type": "UnrelatedNode", "inputs": {}},
|
||||
}
|
||||
|
||||
|
||||
def _run(prompt):
|
||||
return compile_serverless_prompt(prompt)
|
||||
|
||||
|
||||
def _load_prompt_transform():
|
||||
module_path = MODULE_PATH.with_name("prompt_transform.py")
|
||||
package_name = "serverless_prompt_transform_testpkg"
|
||||
root_pkg = types.ModuleType(package_name)
|
||||
root_pkg.__path__ = []
|
||||
sys.modules[package_name] = root_pkg
|
||||
for suffix in ("api", "api.orchestration", "utils"):
|
||||
pkg = types.ModuleType(f"{package_name}.{suffix}")
|
||||
pkg.__path__ = []
|
||||
sys.modules[f"{package_name}.{suffix}"] = pkg
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
|
||||
setattr(logging_module, "log", lambda *_args, **_kwargs: None)
|
||||
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.api.orchestration.prompt_transform", module_path
|
||||
)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_compiles_prompt_prepared_by_existing_worker_helpers():
|
||||
prompt_transform = _load_prompt_transform()
|
||||
original = {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
|
||||
"3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
"5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
|
||||
}
|
||||
pruned = prompt_transform.prune_prompt_for_worker(original)
|
||||
prompt_index = prompt_transform.PromptIndex(pruned)
|
||||
prepared = prompt_transform.apply_participant_overrides(
|
||||
pruned,
|
||||
participant_id="worker-a",
|
||||
enabled_worker_ids=["worker-a"],
|
||||
job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"),
|
||||
master_url="http://master.invalid",
|
||||
delegate_master=False,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
graph, manifest = _run(prepared)
|
||||
|
||||
assert graph["2"]["inputs"]["seed"] == 101
|
||||
assert set(graph) == {"1", "2", "7"}
|
||||
assert graph["7"]["class_type"] == "SaveImage"
|
||||
assert manifest[0]["job_id"] == "run_4"
|
||||
assert manifest[0]["collector_id"] == "4"
|
||||
|
||||
|
||||
def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream():
|
||||
original = _prepared_prompt()
|
||||
|
||||
graph, manifest = _run(original)
|
||||
|
||||
assert manifest == [
|
||||
{
|
||||
"collector_id": "4",
|
||||
"job_id": "job/42",
|
||||
"node_id": "9",
|
||||
"kind": "image",
|
||||
"prefix": graph["9"]["inputs"]["filename_prefix"],
|
||||
}
|
||||
]
|
||||
prefix = manifest[0]["prefix"]
|
||||
assert graph["9"] == {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"images": ["2", 0], "filename_prefix": prefix},
|
||||
"_meta": {"title": "Serverless output: collector 4 (image)"},
|
||||
}
|
||||
assert "/" not in prefix and "\\" not in prefix
|
||||
assert graph["2"]["inputs"]["seed"] == 41
|
||||
assert set(graph) == {"1", "2", "9"}
|
||||
assert all("Distributed" not in node["class_type"] for node in graph.values())
|
||||
assert original["3"]["inputs"]["seed"] == 40
|
||||
assert original["4"]["inputs"]["images"] == ["2", 0]
|
||||
|
||||
|
||||
def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts():
|
||||
prompt = _prepared_prompt()
|
||||
prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe")
|
||||
|
||||
first_graph, first_manifest = _run(prompt)
|
||||
second_graph, second_manifest = _run(prompt)
|
||||
|
||||
assert first_manifest == second_manifest
|
||||
assert first_graph == second_graph
|
||||
prefix = first_manifest[0]["prefix"]
|
||||
assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"]
|
||||
assert not any(char in prefix for char in ("/", "\\", ":"))
|
||||
assert first_manifest[0]["job_id"] == "../run/a"
|
||||
|
||||
|
||||
def test_compiles_multiple_disconnected_collectors_in_stable_order():
|
||||
prompt = {
|
||||
"10": {"class_type": "LoadImage", "inputs": {}},
|
||||
"20": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}},
|
||||
"1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}},
|
||||
"30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}},
|
||||
}
|
||||
|
||||
graph, manifest = _run(prompt)
|
||||
|
||||
assert [item["collector_id"] for item in manifest] == ["1", "2"]
|
||||
assert [item["node_id"] for item in manifest] == ["31", "32"]
|
||||
assert graph["31"]["inputs"]["images"] == ["10", 0]
|
||||
assert graph["32"]["inputs"]["images"] == ["20", 0]
|
||||
assert set(graph) == {"10", "20", "31", "32"}
|
||||
|
||||
|
||||
def test_rejects_audio_input_instead_of_emitting_audio_saver():
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadAudio", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="audio.*unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_collector_without_image_input():
|
||||
with pytest.raises(ValueError, match="collector.*image"):
|
||||
_run({"1": {"class_type": "DistributedCollector", "inputs": {}}})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_type",
|
||||
["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"],
|
||||
)
|
||||
def test_rejects_unsupported_distributed_and_batch_nodes(class_type):
|
||||
prompt = {
|
||||
"1": {"class_type": class_type, "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
"3": {"class_type": "LoadImage", "inputs": {}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="unsupported"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_collector_to_collector_dependency_through_processing_node():
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
|
||||
"3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="collector-to-collector"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"link",
|
||||
[["missing", 0], ["1", -1], ["1", "0"], ["1", True]],
|
||||
)
|
||||
def test_rejects_malformed_api_links(link):
|
||||
prompt = {
|
||||
"1": {"class_type": "LoadImage", "inputs": {}},
|
||||
"2": {"class_type": "DistributedCollector", "inputs": {"images": link}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="link"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_rejects_unexpected_output_slot_from_removed_distributed_node():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="output slot"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_distributed_seed_matches_worker_offset_and_master_passthrough():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["3"]["inputs"] == {"seed": 12, "steps": 7}
|
||||
assert "1" not in graph and "2" not in graph
|
||||
|
||||
|
||||
def test_distributed_value_uses_one_based_worker_values_and_type_coercion():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}},
|
||||
"2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["3"]["inputs"] == {"steps": 8, "seed": 3}
|
||||
assert "1" not in graph and "2" not in graph
|
||||
|
||||
|
||||
def test_resolves_stock_primitive_links_for_distributed_seed_and_value():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}},
|
||||
"4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}},
|
||||
"5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}},
|
||||
"7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75}
|
||||
assert not {"1", "2", "3", "4", "5"}.intersection(graph)
|
||||
|
||||
|
||||
def test_rejects_dynamic_linked_seed_or_value_without_executing_it():
|
||||
prompt = {
|
||||
"1": {"class_type": "ArbitraryNode", "inputs": {}},
|
||||
"2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_primitive_resolution_detects_cycles():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}},
|
||||
"2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}},
|
||||
"4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}},
|
||||
"5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="cycle"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_model_name_trivial_text_passthrough_is_inlined():
|
||||
prompt = {
|
||||
"1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}},
|
||||
"2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}},
|
||||
"3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["2"]["inputs"]["model_name"] == "model.safetensors"
|
||||
assert "1" not in graph
|
||||
|
||||
|
||||
def test_rejects_dynamic_linked_model_name_text():
|
||||
prompt = {
|
||||
"1": {"class_type": "ArbitraryNode", "inputs": {}},
|
||||
"2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}},
|
||||
"3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}},
|
||||
"4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"):
|
||||
_run(prompt)
|
||||
|
||||
|
||||
def test_resolves_boolean_and_float_primitive_links():
|
||||
prompt = {
|
||||
"1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}},
|
||||
"2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}},
|
||||
"3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}},
|
||||
"4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}},
|
||||
"5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}},
|
||||
"6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
|
||||
graph, _manifest = _run(prompt)
|
||||
|
||||
assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5}
|
||||
assert not {"1", "2", "3", "4"}.intersection(graph)
|
||||
|
||||
|
||||
def test_compiled_graph_matches_golden_prompt_and_manifest():
|
||||
prompt = _prepared_prompt()
|
||||
prompt["3"]["inputs"]["seed"] = ["10", 0]
|
||||
prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}}
|
||||
prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}
|
||||
prompt["12"] = {
|
||||
"class_type": "DistributedValue",
|
||||
"inputs": {
|
||||
"default_value": ["11", 0],
|
||||
"worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}),
|
||||
"is_worker": True,
|
||||
"worker_id": "worker_0",
|
||||
},
|
||||
}
|
||||
prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}}
|
||||
prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}}
|
||||
prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0])
|
||||
|
||||
graph, manifest = _run(prompt)
|
||||
|
||||
expected_prefix = (
|
||||
"serverless_job_42_worker_0_4_image_"
|
||||
"1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244"
|
||||
)
|
||||
assert list(graph) == ["1", "2", "15"]
|
||||
assert graph == {
|
||||
"1": {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {},
|
||||
"_meta": {"title": "source"},
|
||||
},
|
||||
"2": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["1", 0],
|
||||
"seed": 41,
|
||||
"cfg": 2.5,
|
||||
"model_name": "model.safetensors",
|
||||
},
|
||||
},
|
||||
"15": {
|
||||
"class_type": "SaveImage",
|
||||
"inputs": {"images": ["2", 0], "filename_prefix": expected_prefix},
|
||||
"_meta": {"title": "Serverless output: collector 4 (image)"},
|
||||
},
|
||||
}
|
||||
assert manifest == [
|
||||
{
|
||||
"collector_id": "4",
|
||||
"job_id": "job/42",
|
||||
"node_id": "15",
|
||||
"kind": "image",
|
||||
"prefix": expected_prefix,
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,608 @@
|
||||
import copy
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload, status=200):
|
||||
self.payload = copy.deepcopy(payload)
|
||||
self.status = status
|
||||
|
||||
|
||||
class _FakeRequest:
|
||||
def __init__(self, payload=None, error=None):
|
||||
self.payload = payload
|
||||
self.error = error
|
||||
|
||||
async def json(self):
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
return self.payload
|
||||
|
||||
|
||||
class _Routes:
|
||||
def get(self, _path):
|
||||
return lambda handler: handler
|
||||
|
||||
def post(self, _path):
|
||||
return lambda handler: handler
|
||||
|
||||
|
||||
def _load_modules():
|
||||
repo = Path(__file__).resolve().parents[1]
|
||||
package_name = "dist_serverless_routes_testpkg"
|
||||
for name in list(sys.modules):
|
||||
if name == package_name or name.startswith(f"{package_name}."):
|
||||
del sys.modules[name]
|
||||
|
||||
root = types.ModuleType(package_name)
|
||||
root.__path__ = [str(repo)]
|
||||
sys.modules[package_name] = root
|
||||
api_pkg = types.ModuleType(f"{package_name}.api")
|
||||
api_pkg.__path__ = [str(repo / "api")]
|
||||
sys.modules[f"{package_name}.api"] = api_pkg
|
||||
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
||||
utils_pkg.__path__ = [str(repo / "utils")]
|
||||
sys.modules[f"{package_name}.utils"] = utils_pkg
|
||||
orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
|
||||
orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")]
|
||||
sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
|
||||
|
||||
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
||||
logging_module.log = lambda *_args, **_kwargs: None
|
||||
logging_module.debug_log = lambda *_args, **_kwargs: None
|
||||
sys.modules[logging_module.__name__] = logging_module
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.HEARTBEAT_TIMEOUT = 30
|
||||
sys.modules[constants_module.__name__] = constants_module
|
||||
|
||||
def load_file(module_name, path):
|
||||
spec = importlib.util.spec_from_file_location(module_name, path)
|
||||
assert spec is not None and spec.loader is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py")
|
||||
serverless_config = load_file(
|
||||
f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py"
|
||||
)
|
||||
|
||||
runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime")
|
||||
runtime.secret_store = None
|
||||
runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
|
||||
runtime.get_secret_store = lambda: runtime.secret_store
|
||||
runtime.get_manager = lambda: runtime.manager
|
||||
sys.modules[runtime.__name__] = runtime
|
||||
|
||||
import aiohttp # Ensure the real dependency is loaded; never replace global aiohttp.
|
||||
previous_server = sys.modules.get("server")
|
||||
server = types.ModuleType("server")
|
||||
server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
|
||||
sys.modules["server"] = server
|
||||
|
||||
network = types.ModuleType(f"{package_name}.utils.network")
|
||||
|
||||
async def handle_api_error(_request, error, status=500):
|
||||
return _FakeResponse({"status": "error", "message": str(error)}, status)
|
||||
|
||||
network.handle_api_error = handle_api_error
|
||||
network.normalize_host = lambda value: value
|
||||
sys.modules[network.__name__] = network
|
||||
|
||||
routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py")
|
||||
routes.web = types.SimpleNamespace(
|
||||
json_response=lambda payload, status=200: _FakeResponse(payload, status)
|
||||
)
|
||||
|
||||
if previous_server is None:
|
||||
sys.modules.pop("server", None)
|
||||
else:
|
||||
sys.modules["server"] = previous_server
|
||||
return config_module, serverless_config, runtime, routes
|
||||
|
||||
|
||||
config_module, serverless_config, runtime_module, config_routes = _load_modules()
|
||||
|
||||
|
||||
class _RunningTask:
|
||||
def done(self):
|
||||
return False
|
||||
|
||||
|
||||
class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase):
|
||||
def setUp(self):
|
||||
self.tempdir = tempfile.TemporaryDirectory()
|
||||
self.root = Path(self.tempdir.name)
|
||||
config_module.CONFIG_FILE = str(self.root / "gpu_config.json")
|
||||
config_module.invalidate_config_cache()
|
||||
self.config = config_module.get_default_config()
|
||||
self.config["workers"] = []
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json")
|
||||
runtime_module.secret_store = self.secrets
|
||||
runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
|
||||
|
||||
def tearDown(self):
|
||||
config_module.invalidate_config_cache()
|
||||
self.tempdir.cleanup()
|
||||
|
||||
def read_config_file(self):
|
||||
return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
class ServerlessConfigRouteTests(_ConfigRouteFixture):
|
||||
async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self):
|
||||
self.secrets.set("pool-1", "private-api-key")
|
||||
self.config["workers"] = [
|
||||
{"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1",
|
||||
"api_key": "inline-api-key", "name": "GPU"}
|
||||
]
|
||||
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
|
||||
config_module.invalidate_config_cache()
|
||||
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = response.payload["workers"][0]
|
||||
self.assertTrue(worker["api_key_configured"])
|
||||
self.assertNotIn("api_key", worker)
|
||||
self.assertNotIn("private-api-key", json.dumps(response.payload))
|
||||
self.assertNotIn("inline-api-key", json.dumps(response.payload))
|
||||
|
||||
async def test_get_config_redacts_normalized_secret_field_aliases(self):
|
||||
aliases = {
|
||||
"apiKey": "camel-api-key",
|
||||
"APIKEY": "upper-api-key",
|
||||
"clearApiKey": "camel-clear-key",
|
||||
"CLEAR_API_KEY": "upper-clear-key",
|
||||
"runpod_CREDENTIAL": "provider-credential",
|
||||
}
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
**aliases,
|
||||
"nested": {"credential_note": "nested-credential"},
|
||||
}]
|
||||
Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
|
||||
config_module.invalidate_config_cache()
|
||||
|
||||
response = await config_routes.get_config_endpoint(_FakeRequest())
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = response.payload["workers"][0]
|
||||
for alias, value in aliases.items():
|
||||
self.assertNotIn(alias, worker)
|
||||
self.assertNotIn(value, json.dumps(response.payload))
|
||||
self.assertNotIn("credential_note", worker["nested"])
|
||||
self.assertNotIn("nested-credential", json.dumps(response.payload))
|
||||
|
||||
async def test_bulk_config_rejects_normalized_secret_field_aliases(self):
|
||||
for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"):
|
||||
with self.subTest(secret_field=secret_field):
|
||||
secret = f"private-{secret_field}"
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "local", secret_field: secret}]
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self):
|
||||
secret = "runpod-secret-never-returned"
|
||||
payload = {
|
||||
"worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2,
|
||||
"timeout_seconds": 600, "enabled": True, "api_key": secret,
|
||||
}
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(payload))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload["api_key_configured"])
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(saved["endpoint_id"], "ep-123")
|
||||
self.assertEqual(saved["jobs_per_run"], 3)
|
||||
self.assertEqual(saved["concurrency"], 2)
|
||||
self.assertEqual(saved["timeout_seconds"], 600)
|
||||
for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertEqual(self.secrets.get("pool-1"), secret)
|
||||
|
||||
async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3,
|
||||
"timeout_seconds": 900, "enabled": False,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest({"worker_id": "pool-1", "enabled": True})
|
||||
)
|
||||
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertTrue(worker["enabled"])
|
||||
self.assertEqual(worker["endpoint_id"], "ep-123")
|
||||
self.assertEqual(worker["jobs_per_run"], 4)
|
||||
self.assertEqual(worker["concurrency"], 3)
|
||||
self.assertEqual(worker["timeout_seconds"], 900)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
|
||||
self.assertTrue(response.payload["api_key_configured"])
|
||||
|
||||
async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self):
|
||||
secret = "must-not-leak"
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "concurrency": True, "api_key": secret,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_json_decode_error_is_generic_and_does_not_echo_body(self):
|
||||
secret = "body-secret-value"
|
||||
response = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest(error=ValueError(f"bad json near {secret}"))
|
||||
)
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
|
||||
async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self):
|
||||
secret = "bulk-secret-never-write"
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "nested": {"api_key": secret},
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertNotIn(secret, json.dumps(response.payload))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertFalse(self.secrets.path.exists())
|
||||
|
||||
async def test_bulk_update_validates_and_returns_redacted_serverless_config(self):
|
||||
self.secrets.set("pool-1", "saved-secret")
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "name": "Pool", "enabled": True,
|
||||
"host": "obsolete-host", "port": 8188, "cuda_device": 0,
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"])
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for obsolete in ("host", "port", "cuda_device"):
|
||||
self.assertNotIn(obsolete, saved)
|
||||
self.assertNotIn("api_key", response.payload["config"]["workers"][0])
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
|
||||
|
||||
async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self):
|
||||
self.secrets.set("pool-1", "bulk-clear-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "local", "port": 8188,
|
||||
"provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
|
||||
}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
saved = self.read_config_file()["workers"][0]
|
||||
for field in ("provider", "endpoint_id", "api_key_configured"):
|
||||
self.assertNotIn(field, saved)
|
||||
|
||||
async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
|
||||
existing = {"id": "local-1", "type": "local", "port": 8188}
|
||||
self.config["workers"] = [existing]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [existing])
|
||||
|
||||
async def test_delete_serverless_worker_clears_saved_key(self):
|
||||
self.secrets.set("pool-1", "delete-me")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123"
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_type_conversion_clears_secret_and_removes_serverless_fields(self):
|
||||
self.secrets.set("pool-1", "old-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
|
||||
"endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188,
|
||||
"cuda_device": 0, "enabled": False,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["type"], "local")
|
||||
for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
|
||||
self.assertNotIn(serverless_field, worker)
|
||||
|
||||
async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
"enabled": False,
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
|
||||
enabled = await config_routes.update_worker_endpoint(
|
||||
_FakeRequest({"worker_id": "pool-1", "enabled": True})
|
||||
)
|
||||
blocked = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "endpoint_id": "ep-456",
|
||||
}))
|
||||
|
||||
self.assertEqual(enabled.status, 200)
|
||||
self.assertEqual(blocked.status, 409)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
|
||||
|
||||
async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
keep = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "",
|
||||
}))
|
||||
clear = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "clear_api_key": True,
|
||||
}))
|
||||
|
||||
self.assertTrue(keep.payload["api_key_configured"])
|
||||
self.assertFalse(clear.payload["api_key_configured"])
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self):
|
||||
with self.subTest(case="inline key"):
|
||||
config = config_module.get_default_config()
|
||||
config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}]
|
||||
self.assertFalse(config_module.save_config(config))
|
||||
with self.subTest(case="invalid serverless record"):
|
||||
config = config_module.get_default_config()
|
||||
config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
|
||||
self.assertFalse(config_module.save_config(config))
|
||||
|
||||
async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self):
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
|
||||
credential_change = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "new-key",
|
||||
}))
|
||||
type_change = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0,
|
||||
}))
|
||||
deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(credential_change.status, 409)
|
||||
self.assertEqual(type_change.status, 409)
|
||||
self.assertEqual(deletion.status, 409)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_secret_write_is_rolled_back_when_config_persistence_fails(self):
|
||||
secret = "write-then-rollback"
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"api_key": secret,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
self.assertEqual(config_module.load_config()["workers"], [])
|
||||
|
||||
async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "bulk-restore-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({
|
||||
"workers": [{"id": "pool-1", "type": "local", "port": 8188}]
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "prior-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "api_key": "replacement-key",
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "prior-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
|
||||
|
||||
async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self):
|
||||
self.secrets.set("pool-1", "delete-restore-key")
|
||||
self.config["workers"] = [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
with patch.object(config_module, "save_config", return_value=False):
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key")
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1")
|
||||
|
||||
async def test_regular_worker_delete_does_not_require_posix_secret_storage(self):
|
||||
self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")):
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_enabling_serverless_without_key_is_rejected(self):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True,
|
||||
}))
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.read_config_file()["workers"], [])
|
||||
|
||||
async def test_bulk_concurrency_change_is_blocked_during_active_run(self):
|
||||
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
runtime_module.manager = types.SimpleNamespace(
|
||||
tasks={"run-1": _RunningTask()},
|
||||
runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
|
||||
pool_is_active=lambda worker_id: worker_id == "pool-1",
|
||||
)
|
||||
updated = copy.deepcopy(self.config["workers"])
|
||||
updated[0]["concurrency"] = 2
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1)
|
||||
|
||||
async def test_non_serverless_update_behavior_is_preserved(self):
|
||||
self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({
|
||||
"worker_id": "local-1", "name": "Renamed", "enabled": True,
|
||||
}))
|
||||
|
||||
self.assertEqual(response.status, 200)
|
||||
worker = self.read_config_file()["workers"][0]
|
||||
self.assertEqual(worker["name"], "Renamed")
|
||||
self.assertTrue(worker["enabled"])
|
||||
|
||||
async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self):
|
||||
raw = {
|
||||
"workers": [{
|
||||
"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
|
||||
"host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old",
|
||||
}]
|
||||
}
|
||||
normalized = config_module._merge_with_defaults(raw, config_module.get_default_config())
|
||||
worker = normalized["workers"][0]
|
||||
for obsolete in ("host", "port", "cuda_device", "extra_args"):
|
||||
self.assertNotIn(obsolete, worker)
|
||||
|
||||
|
||||
class PersistedPoolGuardTests(_ConfigRouteFixture):
|
||||
"""Use the actual manager/ledger across restart, not a mocked idle result."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
from test_serverless_runtime import runtime
|
||||
from unittest.mock import AsyncMock
|
||||
self.runtime = runtime
|
||||
self.session_factory = AsyncMock()
|
||||
self.ledger_path = self.root / "jobs.json"
|
||||
self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
|
||||
runtime_module.manager = self.manager
|
||||
self.secrets.set("pool-1", "saved-key")
|
||||
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
|
||||
self.assertTrue(config_module.save_config(self.config))
|
||||
|
||||
async def test_retained_jobs_block_destructive_changes_after_restart(self):
|
||||
original = self.read_config_file()
|
||||
for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"):
|
||||
for restart in (False, True):
|
||||
with self.subTest(status=status, restart=restart):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id")
|
||||
if restart:
|
||||
self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
|
||||
runtime_module.manager = self.manager
|
||||
else:
|
||||
self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
|
||||
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"concurrency": 2}, {"type": "local", "port": 8188}):
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
|
||||
self.assertEqual(response.status, 409, fields)
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 409)
|
||||
worker = original["workers"][0]
|
||||
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, concurrency=2)]):
|
||||
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
|
||||
self.assertEqual(response.status, 409, workers)
|
||||
self.assertEqual(self.read_config_file(), original)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
self.session_factory.assert_not_awaited()
|
||||
|
||||
async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self):
|
||||
self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN")
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused")
|
||||
response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True}))
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual(self.secrets.get("pool-1"), "saved-key")
|
||||
|
||||
async def test_terminal_or_other_pool_records_do_not_block_deletion(self):
|
||||
self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED")
|
||||
self.manager.ledger.update("other", worker_id="other", status="UNKNOWN")
|
||||
response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone(self.secrets.get("pool-1"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,143 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PACKAGE = "dist_serverless_runtime_tests"
|
||||
for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]:
|
||||
module = types.ModuleType(name)
|
||||
module.__path__ = [str(path)]
|
||||
sys.modules[name] = module
|
||||
|
||||
def load(name):
|
||||
spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py"))
|
||||
assert spec and spec.loader
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
runtime = load("serverless_runtime")
|
||||
|
||||
class Secrets:
|
||||
def get(self, worker_id):
|
||||
return "private-test-key"
|
||||
|
||||
class RuntimeTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
|
||||
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
|
||||
self.manager.poll_seconds = 0
|
||||
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
|
||||
self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.manager.cancel()
|
||||
self.temp.cleanup()
|
||||
|
||||
async def test_job_id_persisted_before_status_and_results_injected(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock())
|
||||
async def status(job_id):
|
||||
self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
|
||||
return {"status": "COMPLETED", "output": {"images": []}}
|
||||
api.status = status
|
||||
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(api.submit.await_count, 1)
|
||||
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
||||
self.assertEqual(item["worker_id"], "pool:1")
|
||||
self.assertEqual(item["tensor"], "tensor")
|
||||
self.assertTrue(item["is_last"])
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
|
||||
|
||||
async def test_output_error_marks_failed_and_delivers_explicit_error(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock())
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
record = self.manager.ledger.all()[0]
|
||||
self.assertEqual(record["status"], "FAILED")
|
||||
item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
|
||||
self.assertIn("serverless_error", item)
|
||||
self.assertNotIn("private-test-key", str(record) + str(item))
|
||||
|
||||
async def test_ambiguous_submit_never_retries(self):
|
||||
api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock())
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(api.submit.await_count, 1)
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
|
||||
|
||||
async def test_cancel_targets_owned_job_only(self):
|
||||
polled = asyncio.Event()
|
||||
async def status(_job):
|
||||
polled.set()
|
||||
await asyncio.Event().wait()
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await polled.wait()
|
||||
await self.manager.cancel(worker_id="other-pool")
|
||||
api.cancel.assert_not_awaited()
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
api.cancel.assert_awaited_once_with("owned-job")
|
||||
|
||||
async def test_per_pool_concurrency_bounds_submission(self):
|
||||
gate = asyncio.Event()
|
||||
reached = asyncio.Event()
|
||||
count = 0
|
||||
async def submit(*args):
|
||||
nonlocal count
|
||||
count += 1
|
||||
reached.set()
|
||||
return f"provider-{count}"
|
||||
async def status(_job):
|
||||
await gate.wait()
|
||||
return {"status": "COMPLETED", "output": {"images": []}}
|
||||
api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
|
||||
second = dict(self.prepared, worker=dict(self.worker, id="pool:2"))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
|
||||
await self.manager.start("execution", "prompt", [self.prepared, second])
|
||||
await reached.wait()
|
||||
await asyncio.sleep(0)
|
||||
self.assertEqual(count, 1)
|
||||
gate.set()
|
||||
await self.manager.wait_for_execution("execution")
|
||||
self.assertEqual(count, 2)
|
||||
|
||||
async def test_cancel_persisted_job_after_master_restart(self):
|
||||
self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1)
|
||||
api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api) as client:
|
||||
await self.manager.cancel(worker_id="pool")
|
||||
api.cancel.assert_awaited_once_with("old-job")
|
||||
self.assertEqual(client.call_args.args[1], "old-endpoint")
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
|
||||
|
||||
async def test_master_failure_cancels_job_without_partial_success(self):
|
||||
self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}})
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
|
||||
|
||||
async def test_deadline_includes_local_slot_wait(self):
|
||||
self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
|
||||
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
|
||||
# A blocked slot expires without a billable submission.
|
||||
self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0))
|
||||
with patch.object(runtime, "RunpodClient", return_value=api):
|
||||
await self.manager.start("execution", "prompt", [self.prepared])
|
||||
await self.manager.wait_for_execution("execution")
|
||||
api.submit.assert_not_awaited()
|
||||
self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT")
|
||||
@@ -0,0 +1,98 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
PKG = "dist_serverless_status_tests"
|
||||
|
||||
|
||||
class StatusTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
def mod(name, **kwargs):
|
||||
result = types.ModuleType(name)
|
||||
result.__dict__.update(kwargs)
|
||||
return result
|
||||
self.manager = types.SimpleNamespace(health=AsyncMock(return_value={"ok": True, "queueRemaining": 0, "message": "Ready (scales from zero)"}), public_jobs=lambda _: [], cancel=AsyncMock(return_value=1))
|
||||
self.pool = {"id": "pool", "type": "serverless", "endpoint_id": "ep"}
|
||||
modules = {PKG: mod(PKG, __path__=[str(ROOT)]), PKG + ".api": mod("api", __path__=[str(ROOT/'api')]),
|
||||
PKG + ".utils": mod("utils", __path__=[]),
|
||||
PKG + ".utils.config": mod("config", load_config=lambda: {"workers": [self.pool]}),
|
||||
PKG + ".api.orchestration": mod("orchestration", __path__=[str(ROOT/'api/orchestration')]),
|
||||
PKG + ".api.orchestration.serverless_runtime": mod("runtime", get_manager=lambda: self.manager),
|
||||
"server": mod("server", PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda f: f, post=lambda _: lambda f: f))))}
|
||||
with patch.dict(sys.modules, modules):
|
||||
spec = importlib.util.spec_from_file_location(PKG+'.api.serverless_routes', ROOT/'api/serverless_routes.py')
|
||||
self.routes = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(self.routes)
|
||||
|
||||
async def test_health_only_never_submits_and_supports_scale_to_zero(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.test_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.assertTrue(json.loads(result.text)["ok"])
|
||||
self.manager.health.assert_awaited_once_with(self.pool, fresh=True)
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_unknown_pool_cannot_cancel_other_jobs(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "other"}))
|
||||
result = await self.routes.cancel_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 404)
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_transport_error_does_not_leak_credentials(self):
|
||||
self.manager.health.side_effect = RuntimeError("Bearer secret-token https://signed.invalid/?token=value")
|
||||
request = types.SimpleNamespace(match_info={"worker_id": "pool"})
|
||||
result = await self.routes.serverless_status_endpoint(request)
|
||||
self.assertFalse(json.loads(result.text)["ok"])
|
||||
self.assertNotIn("secret-token", result.text)
|
||||
self.assertNotIn("signed.invalid", result.text)
|
||||
|
||||
async def test_cancel_uses_only_requested_pool(self):
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.cancel_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.manager.cancel.assert_awaited_once_with(worker_id="pool")
|
||||
|
||||
async def test_reconcile_only_reads_the_requested_pools_jobs(self):
|
||||
self.manager.reconcile = AsyncMock(return_value=2)
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
|
||||
result = await self.routes.reconcile_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.assertEqual(json.loads(result.text)["resolved"], 2)
|
||||
self.manager.reconcile.assert_awaited_once_with(worker_id="pool")
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_resolve_requires_explicit_boolean_confirmation_and_record_id(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock()
|
||||
for payload in ({"worker_id": "pool", "record_id": "record"},
|
||||
{"worker_id": "pool", "record_id": "record", "provider_checked": "true"},
|
||||
{"worker_id": "pool", "provider_checked": True}):
|
||||
with self.subTest(payload=payload):
|
||||
result = await self.routes.resolve_serverless_endpoint(types.SimpleNamespace(json=AsyncMock(return_value=payload)))
|
||||
self.assertEqual(result.status, 400)
|
||||
self.manager.acknowledge_recovery.assert_not_awaited()
|
||||
|
||||
async def test_resolve_is_scoped_and_does_not_claim_provider_cancellation(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock()
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
|
||||
result = await self.routes.resolve_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 200)
|
||||
self.manager.acknowledge_recovery.assert_awaited_once_with("pool", "record", provider_checked=True)
|
||||
self.assertIn("not provider-confirmed", json.loads(result.text)["message"])
|
||||
self.manager.cancel.assert_not_awaited()
|
||||
|
||||
async def test_live_or_mismatched_record_refuses_manual_resolution(self):
|
||||
self.manager.acknowledge_recovery = AsyncMock(side_effect=ValueError("Unresolved inactive job not found for this pool"))
|
||||
request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
|
||||
result = await self.routes.resolve_serverless_endpoint(request)
|
||||
self.assertEqual(result.status, 409)
|
||||
|
||||
async def test_failed_health_still_exposes_recovery_ids(self):
|
||||
self.manager.health.side_effect = RuntimeError("private transport error")
|
||||
self.manager.public_jobs = lambda _: [{"record_id": "record", "status": "UNKNOWN"}]
|
||||
result = await self.routes.serverless_status_endpoint(types.SimpleNamespace(match_info={"worker_id": "pool"}))
|
||||
self.assertEqual(json.loads(result.text)["jobs"][0]["record_id"], "record")
|
||||
@@ -0,0 +1,21 @@
|
||||
"""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
|
||||
@@ -0,0 +1,113 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import socket
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
def load_ports():
|
||||
path = Path(__file__).resolve().parents[1] / "workers" / "ports.py"
|
||||
spec = importlib.util.spec_from_file_location("worker_ports_test", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
class WorkerPortTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ports = load_ports()
|
||||
|
||||
def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self):
|
||||
workers = [
|
||||
{"id": "existing", "host": "localhost", "port": 8190, "enabled": False},
|
||||
{"id": "remote", "host": "remote.example", "port": 8192},
|
||||
]
|
||||
with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194])
|
||||
|
||||
def test_local_host_forms_reserve_ports(self):
|
||||
for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]:
|
||||
with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191])
|
||||
|
||||
def test_exhaustion_is_explicit_not_a_partial_allocation(self):
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
with self.assertRaisesRegex(ValueError, "available.*ports"):
|
||||
self.ports.allocate_worker_ports(65534, [], 2)
|
||||
|
||||
def test_launch_rejects_master_port_without_reassigning_configuration(self):
|
||||
worker = {"id": "manual", "port": 8189}
|
||||
with self.assertRaisesRegex(ValueError, "master.*8189"):
|
||||
self.ports.validate_worker_port(worker, 8189, [])
|
||||
self.assertEqual(worker["port"], 8189)
|
||||
|
||||
def test_launch_rejects_another_local_workers_reserved_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "other", "host": None, "port": 8190, "enabled": False}
|
||||
with self.assertRaisesRegex(ValueError, "other"):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_ignores_itself_and_remote_workers_with_same_port(self):
|
||||
worker = {"id": "manual", "port": 8190}
|
||||
other = {"id": "remote", "host": "example.com", "port": 8190}
|
||||
with patch.object(self.ports, "is_port_available", return_value=True):
|
||||
self.ports.validate_worker_port(worker, 8189, [worker, other])
|
||||
|
||||
def test_launch_rejects_an_occupied_port(self):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen()
|
||||
port = listener.getsockname()[1]
|
||||
self.assertFalse(self.ports.is_port_available(port))
|
||||
with self.assertRaisesRegex(ValueError, "already in use"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
@unittest.skipIf(os.name == "nt", "asyncio does not reuse addresses on Windows")
|
||||
def test_recently_closed_connection_does_not_block_worker_restart(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
port = listener.getsockname()[1]
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=2) as client:
|
||||
accepted, _ = listener.accept()
|
||||
with accepted:
|
||||
accepted.settimeout(2)
|
||||
accepted.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
client.shutdown(socket.SHUT_WR)
|
||||
self.assertEqual(accepted.recv(1), b"")
|
||||
self.assertTrue(self.ports.is_port_available(port))
|
||||
self.ports.validate_worker_port({"id": "restart", "port": port}, 1, [])
|
||||
with socket.socket() as restarted:
|
||||
restarted.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
restarted.bind(("127.0.0.1", port))
|
||||
restarted.listen(1)
|
||||
|
||||
def test_active_listener_with_reuseaddr_is_still_a_conflict(self):
|
||||
with socket.socket() as listener:
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind(("127.0.0.1", 0))
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
@unittest.skipUnless(socket.has_ipv6, "IPv6 unavailable")
|
||||
def test_ipv6_only_listener_is_still_a_conflict(self):
|
||||
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as listener:
|
||||
listener.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
|
||||
try:
|
||||
listener.bind(("::1", 0))
|
||||
except OSError as exc:
|
||||
self.skipTest(f"IPv6 loopback unavailable: {exc}")
|
||||
listener.listen(1)
|
||||
self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
|
||||
|
||||
def test_invalid_ports_fail_clearly(self):
|
||||
for port in [0, 65536, "not-a-port", None, True, 8190.5]:
|
||||
with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"):
|
||||
self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,11 +1,16 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from argparse import Namespace
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.engine import URL, make_url
|
||||
|
||||
|
||||
def _load_process_module(module_filename: str):
|
||||
module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
|
||||
@@ -21,7 +26,7 @@ def _load_process_module(module_filename: str):
|
||||
sys.modules[package_name] = root_pkg
|
||||
|
||||
workers_pkg = types.ModuleType(f"{package_name}.workers")
|
||||
workers_pkg.__path__ = []
|
||||
workers_pkg.__path__ = [str(module_path.parents[1])]
|
||||
sys.modules[f"{package_name}.workers"] = workers_pkg
|
||||
|
||||
process_pkg = types.ModuleType(f"{package_name}.workers.process")
|
||||
@@ -39,8 +44,23 @@ def _load_process_module(module_filename: str):
|
||||
|
||||
process_module = types.ModuleType(f"{package_name}.utils.process")
|
||||
process_module.get_python_executable = lambda: "/usr/bin/test-python"
|
||||
process_module.is_process_alive = lambda _pid: False
|
||||
process_module.terminate_process = lambda *_args, **_kwargs: None
|
||||
sys.modules[f"{package_name}.utils.process"] = process_module
|
||||
|
||||
config_module = types.ModuleType(f"{package_name}.utils.config")
|
||||
config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}}
|
||||
config_module.save_config = lambda _config: None
|
||||
sys.modules[f"{package_name}.utils.config"] = config_module
|
||||
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
||||
constants_module.PROCESS_TERMINATION_TIMEOUT = 1
|
||||
constants_module.PROCESS_WAIT_TIMEOUT = 1
|
||||
constants_module.WORKER_CHECK_INTERVAL = 0.01
|
||||
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
||||
network_module = types.ModuleType(f"{package_name}.utils.network")
|
||||
network_module.get_server_port = lambda: 8189
|
||||
sys.modules[f"{package_name}.utils.network"] = network_module
|
||||
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
f"{package_name}.workers.process.{module_name}",
|
||||
module_path,
|
||||
@@ -70,6 +90,145 @@ class ComfyRootDiscoveryTests(unittest.TestCase):
|
||||
|
||||
|
||||
class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
def build_command(self, root, worker=None, runtime=None):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
worker = worker or {"id": "worker-a", "port": 8190}
|
||||
runtime = runtime if runtime is not None else Namespace(database_url=None)
|
||||
with patch.object(builder, "_get_runtime_args", return_value=runtime):
|
||||
return builder.build_launch_command(worker, str(root))
|
||||
|
||||
def test_database_is_unique_and_stable_by_worker_id(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
first = self.build_command(root, {"id": "worker-a", "port": 8190})
|
||||
renamed = self.build_command(root, {
|
||||
"id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3,
|
||||
})
|
||||
second = self.build_command(root, {"id": "worker-b", "port": 8191})
|
||||
database = first[first.index("--database-url") + 1]
|
||||
self.assertEqual(database, renamed[renamed.index("--database-url") + 1])
|
||||
self.assertNotEqual(database, second[second.index("--database-url") + 1])
|
||||
self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/"))
|
||||
self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir())
|
||||
self.assertNotIn("--disable-assets", first)
|
||||
|
||||
def test_database_uses_effective_user_or_base_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra_args, expected in [
|
||||
(f'--user-directory "{root / "custom user"}"', root / "custom user"),
|
||||
(f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"),
|
||||
]:
|
||||
with self.subTest(extra_args=extra_args):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra_args,
|
||||
})
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/"))
|
||||
cmd = self.build_command(root, runtime=Namespace(
|
||||
database_url="sqlite:///master.db", user_directory=str(root / "runtime user"),
|
||||
))
|
||||
database = cmd[cmd.index("--database-url") + 1]
|
||||
self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/"))
|
||||
self.assertNotEqual(database, "sqlite:///master.db")
|
||||
|
||||
def test_relative_database_directories_use_worker_cwd(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory) / "ComfyUI"
|
||||
root.mkdir()
|
||||
(root / "main.py").touch()
|
||||
launcher = Path(directory) / "launcher"
|
||||
launcher.mkdir()
|
||||
cases = [
|
||||
(Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"),
|
||||
(Namespace(database_url=None, base_directory="data"), "", root / "data/user"),
|
||||
(Namespace(database_url=None), "--user-directory=profiles", root / "profiles"),
|
||||
(Namespace(database_url=None), "--base-directory data", root / "data/user"),
|
||||
]
|
||||
for runtime, extra, expected in cases:
|
||||
with self.subTest(runtime=runtime, extra=extra), \
|
||||
patch.object(os, "getcwd", return_value=str(launcher)):
|
||||
cmd = self.build_command(root, {
|
||||
"id": "worker-a", "port": 8190, "extra_args": extra,
|
||||
}, runtime)
|
||||
database = make_url(cmd[cmd.index("--database-url") + 1]).database
|
||||
self.assertEqual(Path(database).parent, expected / "distributed/workers")
|
||||
self.assertFalse((launcher / "profiles").exists())
|
||||
self.assertFalse((launcher / "data").exists())
|
||||
|
||||
def test_database_url_round_trips_special_characters_without_collisions(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for name in ["user?profile", "user%3Fprofile"]:
|
||||
with self.subTest(name=name):
|
||||
user = root / name
|
||||
runtime = Namespace(database_url=None, user_directory=str(user))
|
||||
# SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail
|
||||
# clearly on those versions rather than silently sharing a DB.
|
||||
sample = (user / "test.db").as_posix()
|
||||
serialized = URL.create("sqlite", database=sample).render_as_string()
|
||||
if make_url(serialized).database != sample:
|
||||
with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"):
|
||||
self.build_command(root, runtime=runtime)
|
||||
self.assertFalse(user.exists())
|
||||
continue
|
||||
databases = []
|
||||
for worker_id in ["worker-a", "worker-b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime)
|
||||
url = cmd[cmd.index("--database-url") + 1]
|
||||
database = make_url(url).database
|
||||
self.assertEqual(Path(database).parent, user / "distributed/workers")
|
||||
engine = create_engine(url)
|
||||
try:
|
||||
with engine.connect() as connection:
|
||||
actual = connection.exec_driver_sql("PRAGMA database_list").one()[2]
|
||||
self.assertEqual(Path(actual), Path(database))
|
||||
finally:
|
||||
engine.dispose()
|
||||
databases.append(database)
|
||||
self.assertNotEqual(*databases)
|
||||
|
||||
def test_explicit_database_and_disabled_assets_are_preserved(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]:
|
||||
with self.subTest(extra=extra):
|
||||
cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_older_comfyui_without_database_flag_keeps_existing_command(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
cmd = self.build_command(root, runtime=Namespace())
|
||||
self.assertNotIn("--database-url", cmd)
|
||||
self.assertFalse((root / "user").exists())
|
||||
|
||||
def test_worker_id_cannot_escape_database_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
databases = []
|
||||
for worker_id in ["../outside", "a/b", "a_b"]:
|
||||
cmd = self.build_command(root, {"id": worker_id, "port": 8190})
|
||||
path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///"))
|
||||
self.assertTrue(path.is_relative_to(root / "user"))
|
||||
databases.append(path)
|
||||
self.assertEqual(len(set(databases)), 3)
|
||||
|
||||
def test_extra_args_cannot_silently_override_the_configured_port(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
(root / "main.py").touch()
|
||||
for extra in ["--port 8189", "--port=8189"]:
|
||||
with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"):
|
||||
self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
|
||||
|
||||
def test_inherits_runtime_layout_args_for_desktop(self):
|
||||
builder = launch_builder_module.LaunchCommandBuilder()
|
||||
runtime_args = Namespace(
|
||||
@@ -130,5 +289,64 @@ class LaunchCommandBuilderTests(unittest.TestCase):
|
||||
self.assertNotIn("--auto-launch", cmd)
|
||||
|
||||
|
||||
class ProcessLaunchPortTests(unittest.TestCase):
|
||||
def test_launch_before_master_listens_uses_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
for missing in [AttributeError("PromptServer has no port"), None]:
|
||||
with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \
|
||||
patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
|
||||
def test_pre_listen_launch_still_rejects_the_configured_master_port(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
cli_args = types.ModuleType("comfy.cli_args")
|
||||
cli_args.args = Namespace(port=8189)
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
|
||||
patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \
|
||||
patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
|
||||
def test_conflict_is_rejected_before_building_or_spawning(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
|
||||
with patch.object(lifecycle_module.subprocess, "Popen") as spawn:
|
||||
with self.assertRaisesRegex(ValueError, "master port 8189"):
|
||||
lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
|
||||
spawn.assert_not_called()
|
||||
self.assertEqual(manager.processes, {})
|
||||
|
||||
def test_available_port_reaches_existing_launch_path(self):
|
||||
lifecycle_module = _load_process_module("lifecycle.py")
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
manager = types.SimpleNamespace(
|
||||
find_comfy_root=lambda: directory,
|
||||
build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"],
|
||||
processes={}, save_processes=lambda: None,
|
||||
)
|
||||
with patch.object(lifecycle_module, "validate_worker_port") as validate, \
|
||||
patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
|
||||
pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
|
||||
validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
|
||||
spawn.assert_called_once()
|
||||
self.assertEqual(pid, 1234)
|
||||
self.assertEqual(manager.processes["manual"]["pid"], 1234)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
+65
-7
@@ -2,10 +2,12 @@
|
||||
Configuration management for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import copy
|
||||
import os
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from .logging import log
|
||||
from .secret_fields import is_secret_field_name as _is_secret_field_name
|
||||
|
||||
# Import defaults for timeout fallbacks
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
@@ -15,6 +17,52 @@ _config_cache = None
|
||||
_config_mtime = 0.0
|
||||
_config_lock = asyncio.Lock()
|
||||
|
||||
_SERVERLESS_OBSOLETE_WORKER_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
||||
|
||||
|
||||
def _contains_serverless_secret(value):
|
||||
if isinstance(value, dict):
|
||||
return any(
|
||||
(isinstance(key, str) and _is_secret_field_name(key))
|
||||
or _contains_serverless_secret(item)
|
||||
for key, item in value.items()
|
||||
)
|
||||
if isinstance(value, (list, tuple)):
|
||||
return any(_contains_serverless_secret(item) for item in value)
|
||||
return False
|
||||
|
||||
|
||||
def _normalize_serverless_workers(config):
|
||||
"""Remove obsolete local-worker fields without adding local defaults."""
|
||||
if not isinstance(config, dict):
|
||||
return config
|
||||
workers = config.get("workers")
|
||||
if isinstance(workers, list):
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
||||
for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
|
||||
worker.pop(field, None)
|
||||
return config
|
||||
|
||||
|
||||
def _validated_config_for_save(config):
|
||||
"""Return a safe normalized copy ready for gpu_config.json."""
|
||||
if _contains_serverless_secret(config):
|
||||
raise ValueError("API keys must be stored outside the config file")
|
||||
saved = _normalize_serverless_workers(copy.deepcopy(config))
|
||||
workers = saved.get("workers") if isinstance(saved, dict) else None
|
||||
if isinstance(workers, list) and any(
|
||||
isinstance(worker, dict) and worker.get("type") == "serverless" for worker in workers
|
||||
):
|
||||
from .serverless_config import validate_serverless_worker
|
||||
for index, worker in enumerate(workers):
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
||||
saved_worker = validate_serverless_worker(worker)
|
||||
for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
|
||||
saved_worker.pop(field, None)
|
||||
workers[index] = saved_worker
|
||||
return saved
|
||||
|
||||
|
||||
def _config_path():
|
||||
return CONFIG_FILE
|
||||
@@ -62,7 +110,7 @@ def _merge_with_defaults(data, defaults):
|
||||
if key not in merged:
|
||||
merged[key] = value
|
||||
|
||||
return merged
|
||||
return _normalize_serverless_workers(merged)
|
||||
|
||||
|
||||
def invalidate_config_cache():
|
||||
@@ -97,11 +145,12 @@ def load_config():
|
||||
return _config_cache
|
||||
|
||||
def save_config(config):
|
||||
"""Saves the configuration to file."""
|
||||
"""Saves normalized configuration without allowing API keys in the config file."""
|
||||
tmp_path = f"{_config_path()}.tmp"
|
||||
try:
|
||||
saved_config = _validated_config_for_save(config)
|
||||
with open(tmp_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(config, f, indent=2)
|
||||
json.dump(saved_config, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
os.replace(tmp_path, _config_path())
|
||||
@@ -121,12 +170,21 @@ async def config_transaction():
|
||||
"""Acquire config lock, yield loaded config, and save if changed."""
|
||||
async with _config_lock:
|
||||
config = load_config()
|
||||
original_config = copy.deepcopy(config)
|
||||
original_snapshot = json.dumps(config, sort_keys=True)
|
||||
yield config
|
||||
try:
|
||||
yield config
|
||||
except BaseException:
|
||||
if isinstance(config, dict):
|
||||
config.clear()
|
||||
config.update(original_config)
|
||||
raise
|
||||
updated_snapshot = json.dumps(config, sort_keys=True)
|
||||
if updated_snapshot != original_snapshot:
|
||||
if not save_config(config):
|
||||
raise RuntimeError("Failed to save config")
|
||||
if updated_snapshot != original_snapshot and not save_config(config):
|
||||
if isinstance(config, dict):
|
||||
config.clear()
|
||||
config.update(original_config)
|
||||
raise RuntimeError("Failed to save config")
|
||||
|
||||
def ensure_config_exists():
|
||||
"""Creates default config file if it doesn't exist. Used by __init__.py"""
|
||||
|
||||
@@ -87,6 +87,8 @@ def _split_host_and_port(host):
|
||||
|
||||
def build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
if worker.get("type") == "serverless":
|
||||
raise ValueError("Serverless pools have no ComfyUI worker URL")
|
||||
host = (worker.get("host") or "").strip()
|
||||
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
|
||||
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
"""Shared classification for credential-bearing configuration fields."""
|
||||
|
||||
|
||||
def is_secret_field_name(key: object) -> bool:
|
||||
"""Return whether a field name identifies an API key or credential."""
|
||||
if not isinstance(key, str):
|
||||
return False
|
||||
normalized = "".join(character for character in key.casefold() if character.isalnum())
|
||||
return normalized in {"apikey", "clearapikey"} or "credential" in normalized
|
||||
@@ -0,0 +1,441 @@
|
||||
"""Validation and private storage for Runpod serverless worker settings.
|
||||
|
||||
Secret storage is supported only on POSIX masters. The module remains
|
||||
importable elsewhere so ordinary worker configuration and redaction continue
|
||||
to work, but secret-store operations fail explicitly without POSIX locking.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import stat
|
||||
import threading
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
|
||||
from .secret_fields import is_secret_field_name
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # Windows and other platforms without POSIX file locking.
|
||||
fcntl = None
|
||||
|
||||
|
||||
class ServerlessConfigError(ValueError):
|
||||
"""A serverless worker configuration is invalid."""
|
||||
|
||||
|
||||
class SecretStoreError(RuntimeError):
|
||||
"""The secret store could not be accessed safely."""
|
||||
|
||||
|
||||
_ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII)
|
||||
_UNSUPPORTED_SECRET_STORE_MESSAGE = (
|
||||
"Serverless credential storage is currently supported on POSIX masters only"
|
||||
)
|
||||
_NUMERIC_FIELDS = {
|
||||
"jobs_per_run": (1, 16, 1),
|
||||
"concurrency": (1, 16, 1),
|
||||
"timeout_seconds": (30, 86400, 1800),
|
||||
}
|
||||
_THREAD_LOCKS: dict[str, threading.Lock] = {}
|
||||
_THREAD_LOCKS_GUARD = threading.Lock()
|
||||
|
||||
|
||||
def validate_serverless_worker(worker: dict) -> dict:
|
||||
"""Return a validated copy of a Runpod serverless worker record.
|
||||
|
||||
Provider credentials are intentionally not required here; they are stored
|
||||
separately by :class:`SecretStore`.
|
||||
"""
|
||||
if not isinstance(worker, dict):
|
||||
raise ServerlessConfigError("worker must be an object")
|
||||
|
||||
try:
|
||||
result = copy.deepcopy(worker)
|
||||
except Exception:
|
||||
raise ServerlessConfigError("worker must be copyable") from None
|
||||
|
||||
worker_id = result.get("id")
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
raise ServerlessConfigError("id must be a non-empty string")
|
||||
|
||||
endpoint_id = result.get("endpoint_id")
|
||||
if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id):
|
||||
raise ServerlessConfigError("endpoint_id must be an opaque Runpod identifier")
|
||||
|
||||
if result.get("type", "serverless") != "serverless":
|
||||
raise ServerlessConfigError("type must be serverless")
|
||||
if result.get("provider", "runpod") != "runpod":
|
||||
raise ServerlessConfigError("provider is not supported")
|
||||
if "enabled" in result and type(result["enabled"]) is not bool:
|
||||
raise ServerlessConfigError("enabled must be a boolean")
|
||||
|
||||
for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
|
||||
value = result.get(field, default)
|
||||
if type(value) is not int or not minimum <= value <= maximum:
|
||||
raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}")
|
||||
result[field] = value
|
||||
|
||||
result["type"] = "serverless"
|
||||
result["provider"] = "runpod"
|
||||
return result
|
||||
|
||||
|
||||
def _thread_lock_for(path: str) -> threading.Lock:
|
||||
with _THREAD_LOCKS_GUARD:
|
||||
lock = _THREAD_LOCKS.get(path)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_THREAD_LOCKS[path] = lock
|
||||
return lock
|
||||
|
||||
|
||||
class SecretStore:
|
||||
"""Atomically persist worker-id to API-key mappings in a private JSON file."""
|
||||
|
||||
def __init__(self, path: Path):
|
||||
self.path = Path(path)
|
||||
|
||||
def get(self, worker_id: str) -> str | None:
|
||||
self._require_supported()
|
||||
worker_id = self._validate_worker_id(worker_id)
|
||||
try:
|
||||
with self._locked(create_parent=False) as (parent_fd, filename):
|
||||
data = self._read_mapping(parent_fd, filename)
|
||||
return data.get(worker_id)
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
|
||||
def set(self, worker_id: str, key: str) -> None:
|
||||
self._require_supported()
|
||||
worker_id = self._validate_worker_id(worker_id)
|
||||
if (
|
||||
not isinstance(key, str)
|
||||
or not key
|
||||
or len(key) > 4096
|
||||
or "\r" in key
|
||||
or "\n" in key
|
||||
):
|
||||
raise SecretStoreError("secret must be 1 to 4096 characters and contain no CR or LF")
|
||||
try:
|
||||
with self._locked(create_parent=True) as (parent_fd, filename):
|
||||
data = self._read_mapping(parent_fd, filename)
|
||||
data[worker_id] = key
|
||||
self._write_mapping(parent_fd, filename, data)
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
|
||||
def delete(self, worker_id: str) -> None:
|
||||
self._require_supported()
|
||||
worker_id = self._validate_worker_id(worker_id)
|
||||
try:
|
||||
with self._locked(create_parent=False) as (parent_fd, filename):
|
||||
data = self._read_mapping(parent_fd, filename)
|
||||
if worker_id in data:
|
||||
del data[worker_id]
|
||||
self._write_mapping(parent_fd, filename, data)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
|
||||
def has(self, worker_id: str) -> bool:
|
||||
return self.get(worker_id) is not None
|
||||
|
||||
@staticmethod
|
||||
def _require_supported() -> None:
|
||||
# Never substitute a weaker Windows store: all operations retain the
|
||||
# dir_fd/O_NOFOLLOW and flock protections implemented below.
|
||||
if os.name != "posix" or fcntl is None:
|
||||
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
|
||||
|
||||
@staticmethod
|
||||
def _validate_worker_id(worker_id: str) -> str:
|
||||
if not isinstance(worker_id, str) or not worker_id.strip():
|
||||
raise SecretStoreError("worker_id must be a non-empty string")
|
||||
return worker_id
|
||||
|
||||
def _absolute_path(self) -> Path:
|
||||
path = self.path
|
||||
if ".." in path.parts:
|
||||
raise SecretStoreError("unsafe secret store path")
|
||||
if not path.is_absolute():
|
||||
path = Path.cwd() / path
|
||||
if len(path.parts) < 2 or path.name in ("", ".", ".."):
|
||||
raise SecretStoreError("unsafe secret store path")
|
||||
return path
|
||||
|
||||
@contextmanager
|
||||
def _locked(self, create_parent: bool) -> Iterator[tuple[int, str]]:
|
||||
lock_api = fcntl
|
||||
if os.name != "posix" or lock_api is None:
|
||||
raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
|
||||
path = self._absolute_path()
|
||||
filename = path.name
|
||||
lock_name = f".{filename}.lock"
|
||||
thread_lock = _thread_lock_for(str(path))
|
||||
|
||||
with thread_lock:
|
||||
parent_fd = self._open_parent(path, create=create_parent)
|
||||
lock_fd = None
|
||||
try:
|
||||
flags = os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW
|
||||
lock_fd = os.open(lock_name, flags, 0o600, dir_fd=parent_fd)
|
||||
lock_stat = os.fstat(lock_fd)
|
||||
if not stat.S_ISREG(lock_stat.st_mode):
|
||||
raise SecretStoreError("unsafe secret store lock")
|
||||
os.fchmod(lock_fd, 0o600)
|
||||
lock_api.flock(lock_fd, lock_api.LOCK_EX)
|
||||
self._check_target(parent_fd, filename)
|
||||
yield parent_fd, filename
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
finally:
|
||||
if lock_fd is not None:
|
||||
try:
|
||||
lock_api.flock(lock_fd, lock_api.LOCK_UN)
|
||||
except OSError:
|
||||
pass
|
||||
os.close(lock_fd)
|
||||
os.close(parent_fd)
|
||||
|
||||
@staticmethod
|
||||
def _open_parent(path: Path, create: bool) -> int:
|
||||
parts = path.parts
|
||||
current_fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
|
||||
try:
|
||||
for part in parts[1:-1]:
|
||||
created = False
|
||||
try:
|
||||
next_fd = os.open(
|
||||
part,
|
||||
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
||||
dir_fd=current_fd,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
if not create:
|
||||
raise
|
||||
try:
|
||||
os.mkdir(part, 0o700, dir_fd=current_fd)
|
||||
created = True
|
||||
except FileExistsError:
|
||||
pass
|
||||
next_fd = os.open(
|
||||
part,
|
||||
os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
||||
dir_fd=current_fd,
|
||||
)
|
||||
if created:
|
||||
os.fchmod(next_fd, 0o700)
|
||||
os.close(current_fd)
|
||||
current_fd = next_fd
|
||||
return current_fd
|
||||
except Exception:
|
||||
os.close(current_fd)
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _check_target(parent_fd: int, filename: str) -> None:
|
||||
try:
|
||||
target_stat = os.stat(filename, dir_fd=parent_fd, follow_symlinks=False)
|
||||
except FileNotFoundError:
|
||||
return
|
||||
if not stat.S_ISREG(target_stat.st_mode):
|
||||
raise SecretStoreError("unsafe secret store file")
|
||||
|
||||
@classmethod
|
||||
def _read_mapping(cls, parent_fd: int, filename: str) -> dict[str, str]:
|
||||
cls._check_target(parent_fd, filename)
|
||||
try:
|
||||
fd = os.open(
|
||||
filename,
|
||||
os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
|
||||
dir_fd=parent_fd,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
return {}
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
|
||||
try:
|
||||
file_stat = os.fstat(fd)
|
||||
if not stat.S_ISREG(file_stat.st_mode):
|
||||
raise SecretStoreError("unsafe secret store file")
|
||||
os.fchmod(fd, 0o600)
|
||||
with os.fdopen(fd, "rb") as stream:
|
||||
raw = stream.read()
|
||||
try:
|
||||
mapping = json.loads(raw.decode("utf-8"), object_pairs_hook=cls._unique_object)
|
||||
except Exception:
|
||||
raise SecretStoreError("secret store is corrupt") from None
|
||||
if not isinstance(mapping, dict) or any(
|
||||
not isinstance(worker_id, str)
|
||||
or not worker_id.strip()
|
||||
or not isinstance(key, str)
|
||||
or not key
|
||||
for worker_id, key in mapping.items()
|
||||
):
|
||||
raise SecretStoreError("secret store is corrupt")
|
||||
return mapping
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
finally:
|
||||
try:
|
||||
os.close(fd)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
||||
result: dict[str, object] = {}
|
||||
for key, value in pairs:
|
||||
if key in result:
|
||||
raise ValueError("duplicate key")
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _write_mapping(parent_fd: int, filename: str, mapping: dict[str, str]) -> None:
|
||||
payload = json.dumps(mapping, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
|
||||
temp_name = f".{filename}.{secrets.token_hex(12)}.tmp"
|
||||
temp_fd = None
|
||||
try:
|
||||
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW
|
||||
temp_fd = os.open(temp_name, flags, 0o600, dir_fd=parent_fd)
|
||||
os.fchmod(temp_fd, 0o600)
|
||||
with os.fdopen(temp_fd, "wb") as stream:
|
||||
temp_fd = None
|
||||
stream.write(payload)
|
||||
stream.flush()
|
||||
os.fsync(stream.fileno())
|
||||
SecretStore._check_target(parent_fd, filename)
|
||||
os.replace(temp_name, filename, src_dir_fd=parent_fd, dst_dir_fd=parent_fd)
|
||||
os.fsync(parent_fd)
|
||||
except SecretStoreError:
|
||||
raise
|
||||
except Exception:
|
||||
raise SecretStoreError("Unable to access secret store") from None
|
||||
finally:
|
||||
if temp_fd is not None:
|
||||
try:
|
||||
os.close(temp_fd)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
os.unlink(temp_name, dir_fd=parent_fd)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
class ServerlessSecretChanges:
|
||||
"""Apply serverless secret changes and restore them unless committed.
|
||||
|
||||
Use this around the associated config transaction so a failed config save
|
||||
compensates any secret-store writes or deletions made along the way.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._originals: dict[tuple[int, str], tuple[SecretStore, str, str | None]] = {}
|
||||
self._mutations: list[tuple[int, str]] = []
|
||||
self._committed = False
|
||||
|
||||
def __enter__(self) -> "ServerlessSecretChanges":
|
||||
return self
|
||||
|
||||
def __exit__(self, _exc_type, _exc_value, _traceback) -> bool:
|
||||
if not self._committed:
|
||||
self.rollback()
|
||||
return False
|
||||
|
||||
def capture(self, store: SecretStore, worker_id: str) -> str | None:
|
||||
"""Capture and return the original key once per store/worker pair."""
|
||||
mutation_id = (id(store), worker_id)
|
||||
if mutation_id not in self._originals:
|
||||
self._originals[mutation_id] = (store, worker_id, store.get(worker_id))
|
||||
return self._originals[mutation_id][2]
|
||||
|
||||
def set(self, store: SecretStore, worker_id: str, key: str) -> None:
|
||||
self._prepare_mutation(store, worker_id)
|
||||
store.set(worker_id, key)
|
||||
|
||||
def delete(self, store: SecretStore, worker_id: str) -> None:
|
||||
self._prepare_mutation(store, worker_id)
|
||||
store.delete(worker_id)
|
||||
|
||||
def commit(self) -> None:
|
||||
"""Keep the secret changes after the corresponding config save."""
|
||||
self._committed = True
|
||||
|
||||
def rollback(self) -> None:
|
||||
"""Best-effort restoration without masking the original route error."""
|
||||
for mutation_id in reversed(self._mutations):
|
||||
store, worker_id, prior_secret = self._originals[mutation_id]
|
||||
try:
|
||||
if prior_secret is None:
|
||||
store.delete(worker_id)
|
||||
else:
|
||||
store.set(worker_id, prior_secret)
|
||||
except Exception:
|
||||
pass
|
||||
self._mutations.clear()
|
||||
|
||||
def _prepare_mutation(self, store: SecretStore, worker_id: str) -> None:
|
||||
self.capture(store, worker_id)
|
||||
mutation_id = (id(store), worker_id)
|
||||
if mutation_id not in self._mutations:
|
||||
# Record before applying: a store operation may fail after replacing
|
||||
# its file, so the exception path must still attempt compensation.
|
||||
self._mutations.append(mutation_id)
|
||||
|
||||
|
||||
def _redact_credentials(value):
|
||||
if isinstance(value, dict):
|
||||
redacted = {}
|
||||
for key, item in value.items():
|
||||
if is_secret_field_name(key):
|
||||
continue
|
||||
redacted[key] = _redact_credentials(item)
|
||||
return redacted
|
||||
if isinstance(value, list):
|
||||
return [_redact_credentials(item) for item in value]
|
||||
if isinstance(value, tuple):
|
||||
return tuple(_redact_credentials(item) for item in value)
|
||||
return value
|
||||
|
||||
|
||||
def public_config(config, store: SecretStore):
|
||||
"""Return a deep-copied public configuration with credentials removed."""
|
||||
result = _redact_credentials(copy.deepcopy(config))
|
||||
if not isinstance(result, dict):
|
||||
return result
|
||||
|
||||
workers = result.get("workers")
|
||||
if isinstance(workers, list):
|
||||
for worker in workers:
|
||||
if isinstance(worker, dict) and worker.get("type") == "serverless":
|
||||
worker_id = worker.get("id")
|
||||
configured = (
|
||||
store.has(worker_id)
|
||||
if isinstance(worker_id, str) and worker_id.strip()
|
||||
else False
|
||||
)
|
||||
worker["api_key_configured"] = configured
|
||||
return result
|
||||
+30
-4
@@ -7,7 +7,7 @@ export function createApiClient(baseUrl) {
|
||||
const request = async (
|
||||
endpoint,
|
||||
options = {},
|
||||
{ retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
|
||||
{ retries = TIMEOUTS.MAX_RETRIES, retry = true, redactErrors = false } = {},
|
||||
) => {
|
||||
const maxAttempts = retry ? retries : 1;
|
||||
let lastError;
|
||||
@@ -36,7 +36,7 @@ export function createApiClient(baseUrl) {
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
|
||||
console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${redactErrors ? "[redacted]" : error.message}`);
|
||||
if (attempt < maxAttempts - 1) {
|
||||
await new Promise(resolve => setTimeout(resolve, delay));
|
||||
delay *= 2; // Exponential backoff
|
||||
@@ -86,9 +86,35 @@ export function createApiClient(baseUrl) {
|
||||
},
|
||||
|
||||
async updateWorker(workerId, data) {
|
||||
return request('/distributed/config/update_worker', {
|
||||
const apiKey = data?.api_key;
|
||||
try {
|
||||
return await request('/distributed/config/update_worker', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId, ...data })
|
||||
}, { retry: false, redactErrors: Boolean(apiKey) });
|
||||
} catch (error) {
|
||||
if (apiKey && typeof error?.message === "string") {
|
||||
error.message = error.message.split(apiKey).join("[redacted]");
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
},
|
||||
|
||||
async testServerlessWorker(workerId) {
|
||||
return request('/distributed/serverless/test', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId, ...data })
|
||||
body: JSON.stringify({ worker_id: workerId }),
|
||||
}, { retry: false });
|
||||
},
|
||||
|
||||
async getServerlessStatus(workerId) {
|
||||
return request(`/distributed/serverless/status/${encodeURIComponent(workerId)}`);
|
||||
},
|
||||
|
||||
async cancelServerlessWorker(workerId) {
|
||||
return request('/distributed/serverless/cancel', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({ worker_id: workerId }),
|
||||
}, { retry: false });
|
||||
},
|
||||
|
||||
|
||||
+5
-4
@@ -104,7 +104,8 @@ export const PULSE_ANIMATION_CSS = `
|
||||
transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease;
|
||||
}
|
||||
.worker-settings.expanded {
|
||||
max-height: 500px;
|
||||
/* Let long forms and wrapped labels fit; the sidebar owns scrolling. */
|
||||
max-height: none;
|
||||
opacity: 1;
|
||||
padding: 12px;
|
||||
margin-top: 8px;
|
||||
@@ -130,7 +131,7 @@ export const PULSE_ANIMATION_CSS = `
|
||||
`;
|
||||
|
||||
export const UI_STYLES = {
|
||||
statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
|
||||
statusDot: "display: inline-block; flex: 0 0 10px; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
|
||||
controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;",
|
||||
formGroup: "display: flex; flex-direction: column; gap: 5px;",
|
||||
formLabel: "font-size: 12px; color: var(--dist-label-text, #ccc); font-weight: 500;",
|
||||
@@ -146,12 +147,12 @@ export const UI_STYLES = {
|
||||
// Column styles
|
||||
columnBase: "display: flex; align-items: center; justify-content: center;",
|
||||
checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; cursor: default;",
|
||||
contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
|
||||
contentColumn: "flex: 1; min-width: 0; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
|
||||
iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: var(--dist-placeholder-add-color, #666);",
|
||||
|
||||
// Row and content styles
|
||||
infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
|
||||
workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;",
|
||||
workerContent: "display: flex; align-items: center; gap: 10px; flex: 1; min-width: 0;",
|
||||
|
||||
// Form and controls styles
|
||||
buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
|
||||
|
||||
@@ -235,6 +235,17 @@
|
||||
border-radius: 4px;
|
||||
}
|
||||
|
||||
.dist-worker-info {
|
||||
min-width: 0;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
|
||||
.dist-serverless-status {
|
||||
display: block;
|
||||
margin-top: 4px;
|
||||
line-height: 1.4;
|
||||
}
|
||||
|
||||
.dist-worker-info__title {
|
||||
color: var(--dist-card-title) !important;
|
||||
font-size: 1.03em;
|
||||
|
||||
+9
-12
@@ -2,7 +2,6 @@ import { app } from "/scripts/app.js";
|
||||
import { ENDPOINTS } from "./constants.js";
|
||||
|
||||
const NODE_CLASS = "DistributedValue";
|
||||
const CONVERTED_WIDGET = "converted-widget";
|
||||
const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
|
||||
const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
|
||||
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
|
||||
@@ -42,15 +41,13 @@ function getDynamicWorkerWidgets(node) {
|
||||
return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
|
||||
}
|
||||
|
||||
function hideWidgetForGood(node, widget, suffix = "") {
|
||||
function hideWidgetForGood(widget) {
|
||||
if (!widget) return;
|
||||
if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
|
||||
|
||||
widget.origType = widget.type;
|
||||
widget.origComputeSize = widget.computeSize;
|
||||
widget.origSerializeValue = widget.serializeValue;
|
||||
widget.computeSize = () => [0, -4];
|
||||
widget.type = `${CONVERTED_WIDGET}${suffix}`;
|
||||
// These are serialized backing fields, not widgets converted into sockets.
|
||||
// Keep their native type/serializer and mutate the shared options in place
|
||||
// so both the classic canvas and Nodes 2.0 suppress the whole widget row.
|
||||
widget.options ??= {};
|
||||
widget.options.hidden = true;
|
||||
|
||||
// Hide any attached DOM element (multiline widgets).
|
||||
if (widget.element) widget.element.style.display = "none";
|
||||
@@ -58,14 +55,14 @@ function hideWidgetForGood(node, widget, suffix = "") {
|
||||
|
||||
if (widget.linkedWidgets) {
|
||||
for (const linked of widget.linkedWidgets) {
|
||||
hideWidgetForGood(node, linked, `:${widget.name}`);
|
||||
hideWidgetForGood(linked);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function hideRawWidgets(node) {
|
||||
hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
|
||||
hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
|
||||
hideWidgetForGood(getRawDefaultWidget(node));
|
||||
hideWidgetForGood(getRawWorkerValuesWidget(node));
|
||||
}
|
||||
|
||||
function removeDynamicDefaultWidget(node) {
|
||||
|
||||
@@ -112,6 +112,21 @@ export async function performPreflightCheck(extension, workers) {
|
||||
const startTime = Date.now();
|
||||
|
||||
const checkPromises = workers.map(async (worker) => {
|
||||
if (worker.type === "serverless") {
|
||||
try {
|
||||
const providerStatus = await extension.api.getServerlessStatus(worker.id);
|
||||
if (providerStatus?.ok) {
|
||||
extension.log(`Serverless worker ${worker.name} is ready to scale`, "debug");
|
||||
return { worker, active: true };
|
||||
}
|
||||
extension.log(`Serverless worker ${worker.name} is not ready: ${providerStatus?.message || "unavailable"}`, "debug");
|
||||
return { worker, active: false };
|
||||
} catch (error) {
|
||||
extension.log(`Serverless worker ${worker.name} status check failed: ${error.message}`, "debug");
|
||||
return { worker, active: false };
|
||||
}
|
||||
}
|
||||
|
||||
const workerUrl = getWorkerUrl(extension, worker);
|
||||
|
||||
extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug");
|
||||
|
||||
+10
-1
@@ -42,6 +42,15 @@ export async function detectMasterIP(extension) {
|
||||
if (shouldAutoPopulate) {
|
||||
extension.log(`Auto-populating workers based on ${extension.cudaDeviceCount} CUDA devices (excluding master on CUDA ${extension.masterCudaDevice})`, "info");
|
||||
|
||||
const workerCount = extension.cudaDeviceCount - (
|
||||
Number.isInteger(extension.masterCudaDevice) &&
|
||||
extension.masterCudaDevice >= 0 &&
|
||||
extension.masterCudaDevice < extension.cudaDeviceCount ? 1 : 0
|
||||
);
|
||||
const workerPorts = data.local_worker_ports;
|
||||
if (!Array.isArray(workerPorts) || workerPorts.length !== workerCount) {
|
||||
throw new Error("Could not allocate enough available local worker ports; check the server log");
|
||||
}
|
||||
const newWorkers = [];
|
||||
let workerNum = 1;
|
||||
let portOffset = 0;
|
||||
@@ -56,7 +65,7 @@ export async function detectMasterIP(extension) {
|
||||
id: generateUUID(),
|
||||
name: `Worker ${workerNum}`,
|
||||
host: isRunpod ? null : "localhost",
|
||||
port: 8189 + portOffset,
|
||||
port: workerPorts[portOffset],
|
||||
cuda_device: i,
|
||||
enabled: true,
|
||||
extra_args: isRunpod ? "--listen" : "",
|
||||
|
||||
@@ -18,6 +18,19 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) {
|
||||
const workerState = extension.state.getWorker(workerId);
|
||||
const isLaunching = Boolean(workerState?.launching);
|
||||
|
||||
if (worker.type === "serverless") {
|
||||
const providerStatus = workerState?.serverlessStatus || {};
|
||||
if (newStatus.online && newStatus.processing) {
|
||||
extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, `Runpod active (${newStatus.queueCount || 0} queued)`, false);
|
||||
} else if (newStatus.online) {
|
||||
extension.ui.updateStatusDot(workerId, STATUS_COLORS.ONLINE_GREEN, providerStatus.message || "Runpod ready to scale", false);
|
||||
} else if (worker.enabled) {
|
||||
extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, providerStatus.message || "Runpod endpoint unavailable", false);
|
||||
}
|
||||
updateWorkerControls(extension, workerId);
|
||||
return true;
|
||||
}
|
||||
|
||||
if (isLaunching && !newStatus.online) {
|
||||
extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true);
|
||||
} else if (newStatus.online && newStatus.processing) {
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import vm from "node:vm";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
// Evaluate this repository's browser entrypoint without needing a ComfyUI server.
|
||||
// The real-renderer acceptance additionally checks the input rows in Chromium.
|
||||
const source = readFileSync(new URL("../distributedValue.js", import.meta.url), "utf8")
|
||||
.replace(/^import .*;\r?\n/gm, "");
|
||||
|
||||
function setup({ connected = false, targetType = "STRING" } = {}) {
|
||||
let extension;
|
||||
const listeners = new Map();
|
||||
const tasks = [];
|
||||
const workers = [
|
||||
{ id: "gpu-a", name: "GPU A", enabled: true },
|
||||
{ id: "gpu-b", name: "GPU B", enabled: true },
|
||||
{ id: "off", name: "Disabled", enabled: false },
|
||||
];
|
||||
const target = {
|
||||
inputs: [{ name: "value" }],
|
||||
widgets: [{ name: "value", type: targetType === "INT" ? "number" : "string", options: { step: 1, precision: 0 } }],
|
||||
};
|
||||
const graph = { links: { 1: { target_id: 2, target_slot: 0 } }, getNodeById: () => target };
|
||||
const raw = ["default_value", "worker_values"].map((name, index) => ({
|
||||
name,
|
||||
type: "text",
|
||||
value: index === 0 ? "saved default" : '{"1":"saved worker"}',
|
||||
options: { multiline: false },
|
||||
computeSize: vi.fn(() => [200, 20]),
|
||||
serializeValue: vi.fn(function () { return this.value; }),
|
||||
}));
|
||||
const node = {
|
||||
comfyClass: "DistributedValue", graph, widgets: [...raw], size: [200, 100],
|
||||
inputs: raw.map(w => ({ name: w.name, widget: { name: w.name }, link: null })),
|
||||
outputs: [{ links: connected ? [1] : [] }],
|
||||
computeSize: () => [200, 100], setSize: vi.fn(), setDirtyCanvas: vi.fn(),
|
||||
addWidget(type, name, value, callback, options) {
|
||||
const widget = { type, name, value, callback, options };
|
||||
this.widgets.push(widget);
|
||||
return widget;
|
||||
},
|
||||
configure(data) {
|
||||
data.widgets_values.forEach((value, i) => { this.widgets[i].value = value; });
|
||||
return "configured";
|
||||
},
|
||||
};
|
||||
vm.runInNewContext(source, {
|
||||
app: { graph, registerExtension: ext => { extension = ext; } },
|
||||
ENDPOINTS: { CONFIG: "/distributed/config" },
|
||||
fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
|
||||
window: { addEventListener: (name, handler) => listeners.set(name, handler) },
|
||||
setTimeout: callback => { tasks.push(callback); }, console,
|
||||
}, { filename: "distributedValue.js" });
|
||||
return { node, raw, extension, listeners, tasks, flush: () => { while (tasks.length) tasks.shift()(); } };
|
||||
}
|
||||
|
||||
describe("Distributed Value internal widget visibility", () => {
|
||||
it("hides raw fields natively without converting their types or replacing their options", async () => {
|
||||
const h = setup();
|
||||
const options = h.raw.map(w => w.options);
|
||||
const sizes = h.raw.map(w => w.computeSize);
|
||||
const serializers = h.raw.map(w => w.serializeValue);
|
||||
const slots = h.node.inputs;
|
||||
await h.extension.nodeCreated(h.node);
|
||||
|
||||
for (const [i, widget] of h.raw.entries()) {
|
||||
expect(widget.type).toBe("text");
|
||||
expect(widget.options).toBe(options[i]);
|
||||
expect(widget.options.hidden).toBe(true);
|
||||
expect(widget.computeSize).toBe(sizes[i]);
|
||||
expect(widget.serializeValue).toBe(serializers[i]);
|
||||
expect(h.node.widgets[i]).toBe(widget);
|
||||
}
|
||||
expect(h.node.inputs).toBe(slots);
|
||||
expect(h.node.widgets.slice(2).map(w => w.label)).toEqual(["default_value", "GPU A", "GPU B"]);
|
||||
expect(h.node.widgets.slice(2).every(w => !w.options.hidden)).toBe(true);
|
||||
});
|
||||
|
||||
it("keeps typed edits in the original serialized fields across rebuild/configure", async () => {
|
||||
const h = setup({ connected: true, targetType: "INT" });
|
||||
await h.extension.nodeCreated(h.node);
|
||||
h.node.widgets.find(w => w.name === "_dv_default").callback(42);
|
||||
h.node.widgets.find(w => w.name === "_dv_worker_1").callback(7);
|
||||
h.node.widgets.find(w => w.name === "_dv_worker_2").callback(9);
|
||||
const values = h.raw.map(w => w.serializeValue());
|
||||
expect(values[0]).toBe(42);
|
||||
expect(JSON.parse(values[1])).toMatchObject({
|
||||
_type: "INT", 1: "7", 2: "9", _by_worker_id: { "gpu-a": "7", "gpu-b": "9" },
|
||||
});
|
||||
expect(h.node.configure({ widgets_values: values })).toBe("configured");
|
||||
h.flush();
|
||||
expect(h.raw.map(w => w.serializeValue())).toEqual(values);
|
||||
expect(h.raw.every(w => w.options.hidden && w.type === "text")).toBe(true);
|
||||
expect(h.node.widgets.slice(2).map(w => w.value)).toEqual([42, 7, 9]);
|
||||
});
|
||||
|
||||
it("reapplies visibility on worker refresh without duplicating raw or dynamic fields", async () => {
|
||||
const h = setup();
|
||||
await h.extension.nodeCreated(h.node);
|
||||
h.raw[0].options.hidden = false;
|
||||
await h.listeners.get("distributed:workers-changed")({ detail: { workers: [
|
||||
{ id: "gpu-b", name: "Renamed B", enabled: true },
|
||||
] } });
|
||||
expect(h.raw.every(w => w.options.hidden)).toBe(true);
|
||||
expect(h.node.widgets.map(w => w.name)).toEqual(["default_value", "worker_values", "_dv_default", "_dv_worker_1"]);
|
||||
expect(h.node.widgets[3].label).toBe("Renamed B");
|
||||
});
|
||||
|
||||
it("hides linked backing widgets and handles absent options without altering serializers", async () => {
|
||||
const h = setup();
|
||||
delete h.raw[1].options;
|
||||
const linked = { type: "text", options: {}, serializeValue: vi.fn(() => "linked") };
|
||||
h.raw[0].linkedWidgets = [linked];
|
||||
await h.extension.nodeCreated(h.node);
|
||||
expect(h.raw[1].options.hidden).toBe(true);
|
||||
expect(linked.options.hidden).toBe(true);
|
||||
expect(linked.type).toBe("text");
|
||||
expect(linked.serializeValue()).toBe("linked");
|
||||
});
|
||||
});
|
||||
@@ -1,15 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { buildWorkerWebSocketUrl } from "../urlUtils.js";
|
||||
|
||||
|
||||
describe("execution decision helpers", () => {
|
||||
it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
|
||||
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
|
||||
"ws://worker.local:8188/distributed/worker_ws"
|
||||
);
|
||||
expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
|
||||
"wss://worker.example.com/distributed/worker_ws"
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,65 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import { detectMasterIP } from "../masterDetection.js";
|
||||
|
||||
describe("automatic local worker ports", () => {
|
||||
let originalWindow;
|
||||
beforeEach(() => {
|
||||
originalWindow = globalThis.window;
|
||||
globalThis.window = { location: { hostname: "localhost", port: "443" } };
|
||||
});
|
||||
afterEach(() => { globalThis.window = originalWindow; });
|
||||
|
||||
function makeExtension(overrides = {}) {
|
||||
return {
|
||||
config: { workers: [], settings: {}, master: {} },
|
||||
log: vi.fn(),
|
||||
api: {
|
||||
getNetworkInfo: vi.fn().mockResolvedValue({
|
||||
cuda_device: 0, cuda_device_count: 3,
|
||||
master_port: 8189, local_worker_ports: [8191, 8193],
|
||||
...overrides,
|
||||
}),
|
||||
updateMaster: vi.fn().mockResolvedValue({}),
|
||||
updateWorker: vi.fn().mockResolvedValue({}),
|
||||
updateSetting: vi.fn().mockResolvedValue({}),
|
||||
},
|
||||
ui: { updateMasterDisplay: vi.fn() },
|
||||
app: {},
|
||||
loadConfig: vi.fn().mockResolvedValue(),
|
||||
};
|
||||
}
|
||||
|
||||
it("uses server-allocated ports, not hardcoded ports or the browser proxy port", async () => {
|
||||
const extension = makeExtension();
|
||||
await detectMasterIP(extension);
|
||||
expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
|
||||
expect(extension.config.workers.map(worker => worker.cuda_device)).toEqual([1, 2]);
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it("does not overwrite an existing manually configured worker", async () => {
|
||||
const extension = makeExtension();
|
||||
const worker = { id: "manual", host: "localhost", port: 8189 };
|
||||
extension.config.workers = [worker];
|
||||
await detectMasterIP(extension);
|
||||
expect(extension.config.workers).toEqual([worker]);
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("fails without partially saving workers if not enough ports are available", async () => {
|
||||
const extension = makeExtension({ local_worker_ports: [8191] });
|
||||
await detectMasterIP(extension);
|
||||
expect(extension.config.workers).toEqual([]);
|
||||
expect(extension.api.updateWorker).not.toHaveBeenCalled();
|
||||
expect(extension.api.updateSetting).not.toHaveBeenCalled();
|
||||
expect(extension.log).toHaveBeenCalledWith(expect.stringContaining("port"), "error");
|
||||
});
|
||||
|
||||
it("preserves Runpod worker host and launch settings", async () => {
|
||||
globalThis.window.location.hostname = "pod-8189.proxy.runpod.net";
|
||||
const extension = makeExtension();
|
||||
await detectMasterIP(extension);
|
||||
expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
|
||||
expect(extension.config.workers.every(worker => worker.host === null && worker.extra_args === "--listen")).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,507 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client" } }), { virtual: true });
|
||||
|
||||
import { createApiClient } from "../apiClient.js";
|
||||
import { createWorkerSettingsForm } from "../ui/settingsForm.js";
|
||||
import { DistributedUI } from "../ui.js";
|
||||
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
|
||||
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
|
||||
import { isRemoteWorker, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
|
||||
import { performPreflightCheck } from "../executionUtils.js";
|
||||
|
||||
function createDocumentMock() {
|
||||
const elements = new Map();
|
||||
const makeElement = (tagName) => {
|
||||
const element = {
|
||||
tagName,
|
||||
children: [],
|
||||
dataset: {},
|
||||
style: {},
|
||||
setAttribute(name, value) { this[name] = value; },
|
||||
classList: {
|
||||
values: new Set(),
|
||||
add(...names) { names.forEach((name) => this.values.add(name)); },
|
||||
remove(...names) { names.forEach((name) => this.values.delete(name)); },
|
||||
toggle(name, force) {
|
||||
const enabled = force ?? !this.values.has(name);
|
||||
if (enabled) this.values.add(name);
|
||||
else this.values.delete(name);
|
||||
return enabled;
|
||||
},
|
||||
},
|
||||
appendChild(child) { this.children.push(child); child.parentNode = this; return child; },
|
||||
dispatchEvent() { this.onchange?.({ target: this }); },
|
||||
};
|
||||
Object.defineProperty(element, "id", {
|
||||
get() { return this._id || ""; },
|
||||
set(value) { this._id = value; elements.set(value, this); },
|
||||
});
|
||||
return element;
|
||||
};
|
||||
return {
|
||||
elements,
|
||||
document: {
|
||||
createElement: makeElement,
|
||||
getElementById: (id) => elements.get(id) || null,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function createFormUi() {
|
||||
return {
|
||||
createFormGroup(label, value, id, type = "text", placeholder = "") {
|
||||
const group = document.createElement("div");
|
||||
const input = document.createElement("input");
|
||||
input.id = id;
|
||||
input.value = String(value ?? "");
|
||||
input.type = type;
|
||||
input.placeholder = placeholder;
|
||||
group.appendChild(input);
|
||||
return { group, input, label };
|
||||
},
|
||||
createButton(text, onclick) {
|
||||
const button = document.createElement("button");
|
||||
button.textContent = text;
|
||||
button.onclick = onclick;
|
||||
return button;
|
||||
},
|
||||
createButtonGroup(buttons) {
|
||||
const group = document.createElement("div");
|
||||
buttons.forEach((button) => group.appendChild(button));
|
||||
return group;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function makeSaveExtension(worker, updateWorker = vi.fn().mockResolvedValue({})) {
|
||||
return {
|
||||
config: { workers: [worker] },
|
||||
api: { updateWorker },
|
||||
state: { updateWorker: vi.fn(), setWorkerExpanded: vi.fn() },
|
||||
app: { extensionManager: { toast: { add: vi.fn() } } },
|
||||
panelElement: null,
|
||||
};
|
||||
}
|
||||
|
||||
describe("Runpod serverless worker panel", () => {
|
||||
let originalDocument;
|
||||
let originalWindow;
|
||||
let elements;
|
||||
|
||||
beforeEach(() => {
|
||||
originalDocument = globalThis.document;
|
||||
originalWindow = globalThis.window;
|
||||
const mock = createDocumentMock();
|
||||
elements = mock.elements;
|
||||
globalThis.document = mock.document;
|
||||
globalThis.window = {
|
||||
location: { hostname: "localhost", protocol: "http:", port: "8188", origin: "http://localhost:8188" },
|
||||
};
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
globalThis.document = originalDocument;
|
||||
globalThis.window = originalWindow;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("classifies Runpod serverless as remote for worker allocation", () => {
|
||||
expect(isRemoteWorker({}, { type: "serverless" })).toBe(true);
|
||||
});
|
||||
|
||||
it("initializes a saved serverless form before it is attached to the document", () => {
|
||||
// Real document lookups cannot see a form while its card is being built.
|
||||
vi.spyOn(document, "getElementById").mockReturnValue(null);
|
||||
const worker = {
|
||||
id: "saved", name: "Runpod", type: "serverless", endpoint_id: "ep_123",
|
||||
jobs_per_run: 2, concurrency: 1, timeout_seconds: 900, api_key_configured: true,
|
||||
};
|
||||
const form = createWorkerSettingsForm(createFormUi(), {}, worker);
|
||||
|
||||
expect(form.parentNode).toBeUndefined();
|
||||
expect(elements.get("worker-type-saved").value).toBe("serverless");
|
||||
for (const field of ["host", "port", "cuda", "args"]) {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
|
||||
}
|
||||
expect(elements.get("runpod-text-saved").style.display).toBe("none");
|
||||
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "concurrency", "timeout-seconds", "serverless-test"]) {
|
||||
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
|
||||
}
|
||||
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
|
||||
expect(elements.get("jobs-per-run-saved").value).toBe("2");
|
||||
expect(elements.get("concurrency-saved").value).toBe("1");
|
||||
expect(elements.get("timeout-seconds-saved").value).toBe("900");
|
||||
expect(elements.get("api-key-saved").value).toBe("");
|
||||
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
|
||||
});
|
||||
|
||||
it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
|
||||
vi.spyOn(document, "getElementById").mockReturnValue(null);
|
||||
createWorkerSettingsForm(createFormUi(), {}, {
|
||||
id: "cloud", name: "Cloud", type: "cloud", host: "worker.example", port: 8443,
|
||||
});
|
||||
const type = elements.get("worker-type-cloud");
|
||||
expect(type.value).toBe("cloud");
|
||||
expect(elements.get("port-cloud").value).toBe("8443");
|
||||
expect(elements.get("host-cloud").value).toBe("worker.example");
|
||||
expect(elements.get("runpod-text-cloud").style.display).toBe("block");
|
||||
|
||||
type.value = "remote";
|
||||
type.dispatchEvent();
|
||||
type.value = "cloud";
|
||||
type.dispatchEvent();
|
||||
expect(elements.get("port-cloud").value).toBe("443");
|
||||
});
|
||||
|
||||
it("shows only serverless settings and keeps credentials write-only", () => {
|
||||
const worker = {
|
||||
id: "w1", name: "Runpod", type: "remote", host: "worker.example", port: 8188,
|
||||
cuda_device: 0, extra_args: "--test", api_key_configured: true,
|
||||
};
|
||||
const form = createWorkerSettingsForm(createFormUi(), {}, worker);
|
||||
const type = elements.get("worker-type-w1");
|
||||
type.value = "serverless";
|
||||
type.dispatchEvent();
|
||||
|
||||
expect(type.children.map((option) => [option.value, option.textContent])).toContainEqual([
|
||||
"serverless", "Serverless (Runpod)",
|
||||
]);
|
||||
for (const id of ["host-group-w1", "port-group-w1", "cuda-group-w1", "args-group-w1"]) {
|
||||
expect(elements.get(id).style.display).toBe("none");
|
||||
}
|
||||
for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "concurrency-w1", "timeout-seconds-w1"]) {
|
||||
expect(elements.has(id)).toBe(true);
|
||||
}
|
||||
for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "concurrency-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) {
|
||||
expect(elements.get(id).style.display).toBe("flex");
|
||||
}
|
||||
expect(elements.has("serverless-help-w1")).toBe(false);
|
||||
expect(elements.get("api-key-w1").type).toBe("password");
|
||||
expect(elements.get("api-key-w1").value).toBe("");
|
||||
expect(elements.get("api-key-w1").placeholder).toMatch(/leave blank/i);
|
||||
expect(elements.get("jobs-per-run-w1").value).toBe("1");
|
||||
expect(elements.get("jobs-per-run-w1").min).toBe("1");
|
||||
expect(elements.get("jobs-per-run-w1").max).toBe("16");
|
||||
expect(elements.get("concurrency-w1").value).toBe("1");
|
||||
expect(elements.get("concurrency-w1").min).toBe("1");
|
||||
expect(elements.get("concurrency-w1").max).toBe("16");
|
||||
expect(elements.get("timeout-seconds-w1").value).toBe("1800");
|
||||
expect(elements.get("timeout-seconds-w1").min).toBe("30");
|
||||
expect(elements.get("timeout-seconds-w1").max).toBe("86400");
|
||||
expect(form.children.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("saves fixed Runpod settings without enabling or retaining the API key", async () => {
|
||||
const worker = { id: "w2", name: "Old", type: "remote", enabled: false, api_key_configured: false };
|
||||
const extension = makeSaveExtension(worker);
|
||||
elements.set("name-w2", { value: "Runpod worker" });
|
||||
elements.set("worker-type-w2", { value: "serverless" });
|
||||
elements.set("endpoint-id-w2", { value: "ep_123" });
|
||||
elements.set("api-key-w2", { value: "secret-token" });
|
||||
elements.set("clear-api-key-w2", { checked: false });
|
||||
elements.set("jobs-per-run-w2", { value: "4" });
|
||||
elements.set("concurrency-w2", { value: "3" });
|
||||
elements.set("timeout-seconds-w2", { value: "2400" });
|
||||
|
||||
await saveWorkerSettings(extension, "w2");
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
|
||||
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
api_key: "secret-token", jobs_per_run: 4, concurrency: 3, timeout_seconds: 2400, enabled: false,
|
||||
});
|
||||
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
|
||||
expect(worker).not.toHaveProperty("api_key");
|
||||
expect(JSON.stringify(extension.config)).not.toContain("secret-token");
|
||||
expect(elements.get("api-key-w2").value).toBe("");
|
||||
expect(extension.state.updateWorker).toHaveBeenCalledWith("w2", { enabled: false });
|
||||
expect(JSON.stringify(extension.state.updateWorker.mock.calls)).not.toContain("secret-token");
|
||||
});
|
||||
|
||||
it("omits a blank API key so the saved key is preserved", async () => {
|
||||
const worker = { id: "w-blank", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
|
||||
const extension = makeSaveExtension(worker);
|
||||
elements.set("name-w-blank", { value: "Runpod" });
|
||||
elements.set("worker-type-w-blank", { value: "serverless" });
|
||||
elements.set("endpoint-id-w-blank", { value: "ep_123" });
|
||||
elements.set("api-key-w-blank", { value: "" });
|
||||
elements.set("clear-api-key-w-blank", { checked: false });
|
||||
elements.set("jobs-per-run-w-blank", { value: "1" });
|
||||
elements.set("concurrency-w-blank", { value: "1" });
|
||||
elements.set("timeout-seconds-w-blank", { value: "1800" });
|
||||
|
||||
await saveWorkerSettings(extension, "w-blank");
|
||||
|
||||
const payload = extension.api.updateWorker.mock.calls[0][1];
|
||||
expect(payload).not.toHaveProperty("api_key");
|
||||
expect(payload).not.toHaveProperty("clear_api_key");
|
||||
expect(worker.api_key_configured).toBe(true);
|
||||
expect(worker).not.toHaveProperty("api_key");
|
||||
});
|
||||
|
||||
it("omits a blank key to preserve it and sends an explicit clear flag when requested", async () => {
|
||||
const worker = { id: "w3", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
|
||||
const extension = makeSaveExtension(worker);
|
||||
elements.set("name-w3", { value: "Runpod" });
|
||||
elements.set("worker-type-w3", { value: "serverless" });
|
||||
elements.set("endpoint-id-w3", { value: "ep_123" });
|
||||
elements.set("api-key-w3", { value: "" });
|
||||
elements.set("clear-api-key-w3", { checked: true });
|
||||
elements.set("jobs-per-run-w3", { value: "1" });
|
||||
elements.set("concurrency-w3", { value: "1" });
|
||||
elements.set("timeout-seconds-w3", { value: "1800" });
|
||||
|
||||
await saveWorkerSettings(extension, "w3");
|
||||
|
||||
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
|
||||
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
|
||||
clear_api_key: true, jobs_per_run: 1, concurrency: 1, timeout_seconds: 1800, enabled: true,
|
||||
});
|
||||
expect(worker.api_key_configured).toBe(false);
|
||||
expect(worker).not.toHaveProperty("api_key");
|
||||
});
|
||||
|
||||
it.each([
|
||||
[undefined, "Checking status…"],
|
||||
[{ ok: true, queueRemaining: 0, jobs: [] }, "Ready"],
|
||||
[{ ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] }, "Queued"],
|
||||
[{ ok: true, jobs: [{ status: "IN_PROGRESS" }] }, "Running"],
|
||||
[{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_QUEUE" }] }, "Running"],
|
||||
[{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_PROGRESS" }] }, "Running · 2 jobs"],
|
||||
[{ ok: true, jobs: [{ status: "DOWNLOADING" }] }, "Downloading result"],
|
||||
[{ ok: true, jobs: [{ status: "CANCEL_REQUESTED" }] }, "Cancelling"],
|
||||
[{ ok: true, queueRemaining: 1 }, "Active · 1 job"],
|
||||
[{ ok: true, queueRemaining: 3 }, "Active · 3 jobs"],
|
||||
[{ ok: true, jobs: [{ status: "COMPLETED" }] }, "Ready"],
|
||||
[{ ok: true, jobs: [{ status: "FAILED" }] }, "Last job failed"],
|
||||
[{ ok: false, message: "Check the saved endpoint ID and API key" }, "Unavailable"],
|
||||
])("formats provider state as one plain-language summary: %j", (status, expected) => {
|
||||
expect(formatServerlessStatus(status)).toBe(expected);
|
||||
});
|
||||
|
||||
it("keeps status above a consistent endpoint info box and retains expanded settings", () => {
|
||||
const ui = new DistributedUI();
|
||||
const worker = {
|
||||
id: "compact", name: "Used Car Inventory", type: "serverless", enabled: true,
|
||||
endpoint_id: "ep_private_id", jobs_per_run: 1, concurrency: 1,
|
||||
};
|
||||
const extension = {
|
||||
config: { workers: [worker] },
|
||||
state: {
|
||||
isWorkerExpanded: () => false,
|
||||
getWorker: () => ({ serverlessStatus: { ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] } }),
|
||||
},
|
||||
};
|
||||
const card = ui.renderEntityCard("worker", worker, extension);
|
||||
const collect = (element) => [element, ...(element.children || []).flatMap(collect)];
|
||||
const info = collect(card).find((element) => element.classList.values.has("dist-worker-info"));
|
||||
expect(info.innerHTML).not.toContain("Runpod Serverless");
|
||||
expect(info.innerHTML).not.toMatch(/ep_private_id|jobs\/run|slots/);
|
||||
expect(elements.get("serverless-status-compact").textContent).toBe("Queued");
|
||||
expect(elements.get("endpoint-id-compact").value).toBe("ep_private_id");
|
||||
expect(Number(elements.get("jobs-per-run-compact").value)).toBe(1);
|
||||
expect(Number(elements.get("concurrency-compact").value)).toBe(1);
|
||||
const box = elements.get("controls-compact").children[0].children[0];
|
||||
expect(box.textContent).toBe("Runpod Serverless endpoint");
|
||||
expect(box.classList.values.has("dist-info-box")).toBe(true);
|
||||
expect(box.style.cssText).toBe(ui.styles.infoBox);
|
||||
expect(elements.get("status-compact").style.cssText).toContain("flex: 0 0 10px");
|
||||
});
|
||||
|
||||
it("uses provider status instead of a direct ComfyUI URL probe", async () => {
|
||||
const statusElement = { textContent: "" };
|
||||
elements.set("serverless-status-w4", statusElement);
|
||||
const status = {
|
||||
ok: true, queueRemaining: 0, message: "Ready to scale",
|
||||
jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w4", execution_id: "exec-1" }],
|
||||
};
|
||||
const extension = {
|
||||
config: { workers: [{ id: "w4", name: "Runpod", type: "serverless", enabled: true }] },
|
||||
api: { getServerlessStatus: vi.fn().mockResolvedValue(status), probeWorker: vi.fn() },
|
||||
statusCheckAbortController: null,
|
||||
state: {
|
||||
setWorkerStatus: vi.fn(),
|
||||
updateWorker: vi.fn(),
|
||||
getWorkerStatus: vi.fn().mockReturnValue({ online: true, processing: false, queueCount: 0 }),
|
||||
getWorker: vi.fn().mockReturnValue({}),
|
||||
isWorkerLaunching: vi.fn().mockReturnValue(false),
|
||||
},
|
||||
updateWorkerCard: vi.fn().mockReturnValue(true),
|
||||
log: vi.fn(),
|
||||
};
|
||||
|
||||
await checkWorkerStatus(extension, extension.config.workers[0]);
|
||||
|
||||
expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w4");
|
||||
expect(extension.api.probeWorker).not.toHaveBeenCalled();
|
||||
expect(extension.state.setWorkerStatus).toHaveBeenCalledWith("w4", {
|
||||
online: true, processing: true, queueCount: 0,
|
||||
});
|
||||
expect(statusElement.textContent).toBe("Running");
|
||||
expect(statusElement.title).toBe("Ready to scale");
|
||||
});
|
||||
|
||||
it("considers a healthy zero-worker Runpod pool ready and never builds a worker URL", async () => {
|
||||
const worker = { id: "w-ready", name: "Runpod", type: "serverless", enabled: true };
|
||||
const extension = {
|
||||
api: {
|
||||
getServerlessStatus: vi.fn().mockResolvedValue({ ok: true, queueRemaining: 0, message: "Ready to scale" }),
|
||||
probeWorker: vi.fn(),
|
||||
},
|
||||
log: vi.fn(),
|
||||
};
|
||||
|
||||
await expect(performPreflightCheck(extension, [worker])).resolves.toEqual([worker]);
|
||||
expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w-ready");
|
||||
expect(extension.api.probeWorker).not.toHaveBeenCalled();
|
||||
expect(getWorkerUrl(extension, worker)).toBeNull();
|
||||
});
|
||||
|
||||
it("does not render or invoke direct ComfyUI launch, stop, or log controls", async () => {
|
||||
const ui = new DistributedUI();
|
||||
const worker = { id: "w-controls", name: "Runpod", type: "serverless", enabled: true };
|
||||
const controls = ui.createControlsSection({ dynamic: true }, worker, {}, true);
|
||||
|
||||
const extension = {
|
||||
config: { workers: [worker] },
|
||||
api: {
|
||||
launchWorker: vi.fn(), stopWorker: vi.fn(), getRemoteWorkerLog: vi.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
const box = controls.children[0].children[0];
|
||||
expect(box.textContent).toBe("Runpod Serverless endpoint");
|
||||
expect(box.tagName).toBe("div");
|
||||
expect(controls.children[0].children).toHaveLength(1);
|
||||
await launchWorker(extension, worker.id);
|
||||
await stopWorker(extension, worker.id);
|
||||
await viewWorkerLog(extension, worker.id, true);
|
||||
expect(extension.api.launchWorker).not.toHaveBeenCalled();
|
||||
expect(extension.api.stopWorker).not.toHaveBeenCalled();
|
||||
expect(extension.api.getRemoteWorkerLog).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("uses the serverless status, test, and cancel routes with the frozen request shapes", async () => {
|
||||
const originalFetch = globalThis.fetch;
|
||||
const statusPayload = {
|
||||
ok: true,
|
||||
queueRemaining: 0,
|
||||
message: "Ready to scale",
|
||||
jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w5", execution_id: "exec-1" }],
|
||||
};
|
||||
globalThis.fetch = vi.fn()
|
||||
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint test passed" }) })
|
||||
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue(statusPayload) })
|
||||
.mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true }) });
|
||||
try {
|
||||
const client = createApiClient("http://localhost:8188");
|
||||
await client.testServerlessWorker("w5");
|
||||
await expect(client.getServerlessStatus("w5")).resolves.toEqual(statusPayload);
|
||||
await client.cancelServerlessWorker("w5");
|
||||
|
||||
expect(globalThis.fetch).toHaveBeenNthCalledWith(1, "http://localhost:8188/distributed/serverless/test", expect.objectContaining({
|
||||
method: "POST", body: JSON.stringify({ worker_id: "w5" }),
|
||||
}));
|
||||
expect(globalThis.fetch).toHaveBeenNthCalledWith(2, "http://localhost:8188/distributed/serverless/status/w5", expect.any(Object));
|
||||
expect(globalThis.fetch).toHaveBeenNthCalledWith(3, "http://localhost:8188/distributed/serverless/cancel", expect.objectContaining({
|
||||
method: "POST", body: JSON.stringify({ worker_id: "w5" }),
|
||||
}));
|
||||
} finally {
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
|
||||
it("tests only the saved Runpod endpoint and reports the backend message", async () => {
|
||||
const button = document.createElement("button");
|
||||
button.id = "test-serverless-w6";
|
||||
button.textContent = "Test Saved Endpoint";
|
||||
const extension = {
|
||||
api: {
|
||||
testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
|
||||
queueDistributed: vi.fn(),
|
||||
},
|
||||
app: { extensionManager: { toast: { add: vi.fn() } } },
|
||||
};
|
||||
|
||||
await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
|
||||
|
||||
expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
|
||||
expect(extension.api.queueDistributed).not.toHaveBeenCalled();
|
||||
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
|
||||
severity: "success", detail: "Endpoint reachable",
|
||||
}));
|
||||
expect(button.textContent).toBe("Test Saved Endpoint");
|
||||
expect(button.disabled).toBe(false);
|
||||
});
|
||||
|
||||
it("redacts API-key-bearing worker update errors from logs and user-visible exceptions", async () => {
|
||||
const originalFetch = globalThis.fetch;
|
||||
const logSpy = vi.spyOn(console, "log").mockImplementation(() => {});
|
||||
const secret = "never-log-this-runpod-key";
|
||||
globalThis.fetch = vi.fn().mockResolvedValue({
|
||||
ok: false,
|
||||
status: 400,
|
||||
json: vi.fn().mockResolvedValue({ message: `Rejected key ${secret}` }),
|
||||
});
|
||||
try {
|
||||
const client = createApiClient("http://localhost:8188");
|
||||
await expect(client.updateWorker("w-secret", { api_key: secret })).rejects.toThrow("[redacted]");
|
||||
expect(logSpy.mock.calls.flat().join(" ")).not.toContain(secret);
|
||||
} finally {
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
|
||||
it("never sends VRAM operations to serverless workers", async () => {
|
||||
const originalFetch = globalThis.fetch;
|
||||
globalThis.fetch = vi.fn();
|
||||
vi.useFakeTimers();
|
||||
const button = {
|
||||
textContent: "Clear Worker VRAM", style: { cssText: "" }, disabled: false,
|
||||
classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
|
||||
};
|
||||
const extension = { enabledWorkers: [{ id: "s-vram", name: "Runpod", type: "serverless", enabled: true }] };
|
||||
try {
|
||||
await handleClearMemory(extension, button);
|
||||
expect(globalThis.fetch).not.toHaveBeenCalled();
|
||||
} finally {
|
||||
vi.runOnlyPendingTimers();
|
||||
vi.useRealTimers();
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
|
||||
it("cancels each enabled serverless worker once without changing normal worker interrupts", async () => {
|
||||
const originalFetch = globalThis.fetch;
|
||||
globalThis.fetch = vi.fn().mockResolvedValue({ ok: true });
|
||||
vi.useFakeTimers();
|
||||
const button = {
|
||||
textContent: "Interrupt Workers", style: { cssText: "" }, disabled: false,
|
||||
classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
|
||||
};
|
||||
const workers = [
|
||||
{ id: "s1", name: "Serverless 1", type: "serverless", enabled: true },
|
||||
{ id: "s2", name: "Serverless 2", type: "serverless", enabled: true },
|
||||
{ id: "s3", name: "Disabled", type: "serverless", enabled: false },
|
||||
{ id: "local", name: "Local", type: "local", enabled: true, port: 8189 },
|
||||
];
|
||||
const extension = {
|
||||
enabledWorkers: workers.filter((worker) => worker.enabled),
|
||||
api: { cancelServerlessWorker: vi.fn().mockResolvedValue({ ok: true }) },
|
||||
log: vi.fn(),
|
||||
};
|
||||
try {
|
||||
await handleInterruptWorkers(extension, button);
|
||||
|
||||
expect(extension.api.cancelServerlessWorker).toHaveBeenCalledTimes(2);
|
||||
expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(1, "s1");
|
||||
expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(2, "s2");
|
||||
expect(globalThis.fetch).toHaveBeenCalledTimes(1);
|
||||
expect(globalThis.fetch.mock.calls[0][0]).toBe("http://localhost:8189/interrupt");
|
||||
} finally {
|
||||
vi.runOnlyPendingTimers();
|
||||
vi.useRealTimers();
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -106,11 +106,6 @@ describe("buildWorkerWebSocketUrl", () => {
|
||||
"wss://worker.example.com/distributed/worker_ws"
|
||||
);
|
||||
});
|
||||
|
||||
it("always appends /distributed/worker_ws", () => {
|
||||
const url = buildWorkerWebSocketUrl("http://worker.local:8188");
|
||||
expect(url.endsWith("/distributed/worker_ws")).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { getWorkerUrl } from "../workerLifecycle.js";
|
||||
|
||||
|
||||
describe("workerLifecycle URL construction", () => {
|
||||
describe("workerLifecycle URL wiring", () => {
|
||||
let originalWindow;
|
||||
|
||||
beforeEach(() => {
|
||||
@@ -22,31 +22,9 @@ describe("workerLifecycle URL construction", () => {
|
||||
globalThis.window = originalWindow;
|
||||
});
|
||||
|
||||
it("builds local worker URL with explicit local port", () => {
|
||||
// URL variants are covered in urlUtils.test.js; keep the wrapper wiring here.
|
||||
it("forwards the worker and endpoint using window.location", () => {
|
||||
const worker = { id: "w1", port: 8189, type: "local" };
|
||||
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
|
||||
});
|
||||
|
||||
it("builds remote worker URL with host:port", () => {
|
||||
const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
|
||||
expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
|
||||
});
|
||||
|
||||
it("builds cloud worker URL as https", () => {
|
||||
const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
|
||||
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
|
||||
});
|
||||
|
||||
it("rewrites runpod proxy hostname for local worker ports", () => {
|
||||
globalThis.window = {
|
||||
location: {
|
||||
hostname: "podabc.proxy.runpod.net",
|
||||
protocol: "https:",
|
||||
port: "",
|
||||
origin: "https://podabc.proxy.runpod.net",
|
||||
},
|
||||
};
|
||||
const worker = { id: "w4", port: 8189, type: "local" };
|
||||
expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
|
||||
});
|
||||
});
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { PULSE_ANIMATION_CSS } from "../constants.js";
|
||||
|
||||
describe("worker settings layout", () => {
|
||||
it("lets expanded forms grow instead of clipping the final controls", () => {
|
||||
const expanded = PULSE_ANIMATION_CSS.match(/\.worker-settings\.expanded\s*\{([^}]+)\}/)?.[1];
|
||||
expect(expanded).toBeDefined();
|
||||
expect(expanded).toMatch(/max-height:\s*none\s*;/);
|
||||
});
|
||||
|
||||
it("still hides collapsed settings", () => {
|
||||
const collapsed = PULSE_ANIMATION_CSS.match(/\.worker-settings\s*\{([^}]+)\}/)?.[1];
|
||||
expect(collapsed).toMatch(/max-height:\s*0\s*;/);
|
||||
expect(collapsed).toMatch(/overflow:\s*hidden\s*;/);
|
||||
});
|
||||
});
|
||||
@@ -7,6 +7,14 @@ import { createLogModal } from './ui/logModal.js';
|
||||
import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js';
|
||||
import { isRemoteWorker } from './workerSettings.js';
|
||||
|
||||
const escapeWorkerInfo = (value) => String(value ?? "").replace(/[&<>"']/g, (character) => ({
|
||||
"&": "&",
|
||||
"<": "<",
|
||||
">": ">",
|
||||
'"': """,
|
||||
"'": "'",
|
||||
}[character]));
|
||||
|
||||
const cardConfigs = {
|
||||
master: {
|
||||
checkbox: {
|
||||
@@ -58,7 +66,12 @@ const cardConfigs = {
|
||||
infoText: (data, extension) => {
|
||||
const isRemote = isRemoteWorker(extension, data);
|
||||
const isCloud = data.type === 'cloud';
|
||||
|
||||
const isServerless = data.type === 'serverless';
|
||||
|
||||
if (isServerless) {
|
||||
return `<span class="dist-worker-info__title">${escapeWorkerInfo(data.name)}</span>`;
|
||||
}
|
||||
|
||||
if (isCloud) {
|
||||
// For cloud workers, don't show port (it's always 443)
|
||||
return `<span class="dist-worker-info__title">${data.name}</span><br><span class="dist-worker-info__meta">${data.host}</span>`;
|
||||
@@ -590,7 +603,10 @@ export class DistributedUI {
|
||||
}
|
||||
controlsWrapper.appendChild(badge);
|
||||
} else if (config.dynamic && data) {
|
||||
if (isRemote) {
|
||||
if (data.type === 'serverless') {
|
||||
// Match the master/remote footer without exposing direct-worker actions.
|
||||
controlsWrapper.appendChild(this.createInfoBox("Runpod Serverless endpoint"));
|
||||
} else if (isRemote) {
|
||||
const isCloud = data.type === 'cloud';
|
||||
const workerTypeText = isCloud ? "Cloud worker" : "Remote worker";
|
||||
const workerTypeBadge = this.createInfoBox(workerTypeText);
|
||||
|
||||
+12
-1
@@ -1,4 +1,4 @@
|
||||
import { updateWorkerControls, toggleWorkerExpanded } from "../workerLifecycle.js";
|
||||
import { updateWorkerControls, toggleWorkerExpanded, formatServerlessStatus } from "../workerLifecycle.js";
|
||||
import { isRemoteWorker } from "../workerSettings.js";
|
||||
|
||||
export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
|
||||
@@ -74,6 +74,17 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
|
||||
infoSpan.innerHTML = config.infoText(data, extension);
|
||||
workerContent.appendChild(infoSpan);
|
||||
|
||||
if (isWorker && data.type === "serverless") {
|
||||
const statusMessage = document.createElement("small");
|
||||
statusMessage.id = `serverless-status-${data.id}`;
|
||||
statusMessage.className = "dist-worker-info__meta dist-serverless-status";
|
||||
statusMessage.setAttribute("aria-live", "polite");
|
||||
const status = extension.state.getWorker(data.id)?.serverlessStatus;
|
||||
statusMessage.textContent = data.enabled ? formatServerlessStatus(status) : "Disabled";
|
||||
statusMessage.title = status?.message || "";
|
||||
infoSpan.appendChild(statusMessage);
|
||||
}
|
||||
|
||||
infoRow.appendChild(workerContent);
|
||||
|
||||
let settingsArrow;
|
||||
|
||||
+128
-45
@@ -1,5 +1,5 @@
|
||||
import { BUTTON_STYLES } from '../constants.js';
|
||||
import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings } from '../workerSettings.js';
|
||||
import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings, testServerlessWorker } from '../workerSettings.js';
|
||||
|
||||
export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const form = document.createElement("div");
|
||||
@@ -33,9 +33,14 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
cloudOption.value = "cloud";
|
||||
cloudOption.textContent = "Cloud";
|
||||
|
||||
const serverlessOption = document.createElement("option");
|
||||
serverlessOption.value = "serverless";
|
||||
serverlessOption.textContent = "Serverless (Runpod)";
|
||||
|
||||
typeSelect.appendChild(localOption);
|
||||
typeSelect.appendChild(remoteOption);
|
||||
typeSelect.appendChild(cloudOption);
|
||||
typeSelect.appendChild(serverlessOption);
|
||||
|
||||
const runpodText = document.createElement("a");
|
||||
runpodText.id = `runpod-text-${worker.id}`;
|
||||
@@ -44,50 +49,61 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
runpodText.textContent = "Deploy Cloud Worker with Runpod";
|
||||
runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;";
|
||||
|
||||
const createOnChangeHandler = () => {
|
||||
return (e) => {
|
||||
const workerType = e.target.value;
|
||||
const hostGroup = document.getElementById(`host-group-${worker.id}`);
|
||||
const hostInput = document.getElementById(`host-${worker.id}`);
|
||||
const portGroup = document.getElementById(`port-group-${worker.id}`);
|
||||
const portInput = document.getElementById(`port-${worker.id}`);
|
||||
const cudaGroup = document.getElementById(`cuda-group-${worker.id}`);
|
||||
const argsGroup = document.getElementById(`args-group-${worker.id}`);
|
||||
const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`);
|
||||
const updateWorkerType = (workerType, applyDefaults = false) => {
|
||||
// The form is initialized before its card is attached to document.
|
||||
// Use this form's own elements, not document-wide ID lookups.
|
||||
const hostInput = hostGroup.input;
|
||||
const portInput = portGroup.input;
|
||||
const serverlessGroups = [
|
||||
endpointGroup.group,
|
||||
apiKeyGroup.group,
|
||||
clearApiKeyGroup,
|
||||
jobsPerRunGroup.group,
|
||||
concurrencyGroup.group,
|
||||
timeoutGroup.group,
|
||||
testServerlessGroup,
|
||||
];
|
||||
|
||||
if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) {
|
||||
return;
|
||||
}
|
||||
const hideServerless = () => serverlessGroups.forEach((group) => { group.style.display = "none"; });
|
||||
|
||||
if (workerType === "local") {
|
||||
hostGroup.style.display = "none";
|
||||
portGroup.style.display = "flex";
|
||||
cudaGroup.style.display = "flex";
|
||||
argsGroup.style.display = "flex";
|
||||
runpodTextElem.style.display = "none";
|
||||
} else if (workerType === "remote") {
|
||||
hostGroup.style.display = "flex";
|
||||
portGroup.style.display = "flex";
|
||||
cudaGroup.style.display = "none";
|
||||
argsGroup.style.display = "none";
|
||||
runpodTextElem.style.display = "none";
|
||||
hostInput.placeholder = "e.g., 192.168.1.100";
|
||||
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
|
||||
hostInput.value = "";
|
||||
}
|
||||
} else if (workerType === "cloud") {
|
||||
hostGroup.style.display = "flex";
|
||||
portGroup.style.display = "flex";
|
||||
cudaGroup.style.display = "none";
|
||||
argsGroup.style.display = "none";
|
||||
runpodTextElem.style.display = "block";
|
||||
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
|
||||
portInput.value = "443";
|
||||
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
|
||||
hostInput.value = "";
|
||||
}
|
||||
if (workerType === "local") {
|
||||
hideServerless();
|
||||
hostGroup.group.style.display = "none";
|
||||
portGroup.group.style.display = "flex";
|
||||
cudaGroup.group.style.display = "flex";
|
||||
argsGroup.group.style.display = "flex";
|
||||
runpodText.style.display = "none";
|
||||
} else if (workerType === "remote") {
|
||||
hideServerless();
|
||||
hostGroup.group.style.display = "flex";
|
||||
portGroup.group.style.display = "flex";
|
||||
cudaGroup.group.style.display = "none";
|
||||
argsGroup.group.style.display = "none";
|
||||
runpodText.style.display = "none";
|
||||
hostInput.placeholder = "e.g., 192.168.1.100";
|
||||
if (applyDefaults && (hostInput.value === "localhost" || hostInput.value === "127.0.0.1")) {
|
||||
hostInput.value = "";
|
||||
}
|
||||
};
|
||||
} else if (workerType === "cloud") {
|
||||
hideServerless();
|
||||
hostGroup.group.style.display = "flex";
|
||||
portGroup.group.style.display = "flex";
|
||||
cudaGroup.group.style.display = "none";
|
||||
argsGroup.group.style.display = "none";
|
||||
runpodText.style.display = "block";
|
||||
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
|
||||
if (applyDefaults) portInput.value = "443";
|
||||
if (applyDefaults && (hostInput.value === "localhost" || hostInput.value === "127.0.0.1")) {
|
||||
hostInput.value = "";
|
||||
}
|
||||
} else if (workerType === "serverless") {
|
||||
hostGroup.group.style.display = "none";
|
||||
portGroup.group.style.display = "none";
|
||||
cudaGroup.group.style.display = "none";
|
||||
argsGroup.group.style.display = "none";
|
||||
runpodText.style.display = "none";
|
||||
serverlessGroups.forEach((group) => { group.style.display = "flex"; });
|
||||
}
|
||||
};
|
||||
|
||||
typeGroup.appendChild(typeLabel);
|
||||
@@ -114,6 +130,71 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
argsGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex";
|
||||
form.appendChild(argsGroup.group);
|
||||
|
||||
const endpointGroup = ui.createFormGroup("Endpoint ID:", worker.endpoint_id || "", `endpoint-id-${worker.id}`, "text", "Runpod endpoint ID");
|
||||
endpointGroup.group.id = `endpoint-id-group-${worker.id}`;
|
||||
endpointGroup.group.style.display = "none";
|
||||
form.appendChild(endpointGroup.group);
|
||||
|
||||
const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key");
|
||||
apiKeyGroup.group.id = `api-key-group-${worker.id}`;
|
||||
apiKeyGroup.input.autocomplete = "new-password";
|
||||
apiKeyGroup.input.spellcheck = false;
|
||||
apiKeyGroup.group.style.display = "none";
|
||||
form.appendChild(apiKeyGroup.group);
|
||||
|
||||
const clearApiKeyGroup = document.createElement("div");
|
||||
clearApiKeyGroup.id = `clear-api-key-group-${worker.id}`;
|
||||
clearApiKeyGroup.style.cssText = "display: none; align-items: center; gap: 8px;";
|
||||
const clearApiKeyInput = document.createElement("input");
|
||||
clearApiKeyInput.type = "checkbox";
|
||||
clearApiKeyInput.id = `clear-api-key-${worker.id}`;
|
||||
const clearApiKeyLabel = document.createElement("label");
|
||||
clearApiKeyLabel.htmlFor = clearApiKeyInput.id;
|
||||
clearApiKeyLabel.textContent = "Clear saved API key (disable pool first)";
|
||||
clearApiKeyGroup.appendChild(clearApiKeyInput);
|
||||
clearApiKeyGroup.appendChild(clearApiKeyLabel);
|
||||
clearApiKeyInput.onchange = () => {
|
||||
if (clearApiKeyInput.checked) apiKeyGroup.input.value = "";
|
||||
apiKeyGroup.input.disabled = clearApiKeyInput.checked;
|
||||
};
|
||||
form.appendChild(clearApiKeyGroup);
|
||||
|
||||
const jobsPerRunGroup = ui.createFormGroup("Jobs per run:", worker.jobs_per_run ?? 1, `jobs-per-run-${worker.id}`, "number");
|
||||
jobsPerRunGroup.group.id = `jobs-per-run-group-${worker.id}`;
|
||||
jobsPerRunGroup.input.min = "1";
|
||||
jobsPerRunGroup.input.max = "16";
|
||||
jobsPerRunGroup.input.step = "1";
|
||||
jobsPerRunGroup.input.title = "Maximum jobs submitted in one run";
|
||||
jobsPerRunGroup.group.style.display = "none";
|
||||
form.appendChild(jobsPerRunGroup.group);
|
||||
|
||||
const concurrencyGroup = ui.createFormGroup("Concurrent jobs (slots):", worker.concurrency ?? 1, `concurrency-${worker.id}`, "number");
|
||||
concurrencyGroup.group.id = `concurrency-group-${worker.id}`;
|
||||
concurrencyGroup.input.min = "1";
|
||||
concurrencyGroup.input.max = "16";
|
||||
concurrencyGroup.input.step = "1";
|
||||
concurrencyGroup.input.title = "Maximum concurrent jobs per serverless worker";
|
||||
concurrencyGroup.group.style.display = "none";
|
||||
form.appendChild(concurrencyGroup.group);
|
||||
|
||||
const timeoutGroup = ui.createFormGroup("Timeout (seconds):", worker.timeout_seconds ?? 1800, `timeout-seconds-${worker.id}`, "number");
|
||||
timeoutGroup.group.id = `timeout-seconds-group-${worker.id}`;
|
||||
timeoutGroup.input.min = "30";
|
||||
timeoutGroup.input.max = "86400";
|
||||
timeoutGroup.input.step = "1";
|
||||
timeoutGroup.group.style.display = "none";
|
||||
form.appendChild(timeoutGroup.group);
|
||||
|
||||
const testServerlessButton = ui.createButton("Test Saved Endpoint", () => testServerlessWorker(extension, worker.id));
|
||||
testServerlessButton.id = `test-serverless-${worker.id}`;
|
||||
testServerlessButton.title = "Test the saved Runpod endpoint without submitting inference";
|
||||
testServerlessButton.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl;
|
||||
const testServerlessGroup = document.createElement("div");
|
||||
testServerlessGroup.id = `serverless-test-group-${worker.id}`;
|
||||
testServerlessGroup.style.cssText = "display: none;";
|
||||
testServerlessGroup.appendChild(testServerlessButton);
|
||||
form.appendChild(testServerlessGroup);
|
||||
|
||||
const saveBtn = ui.createButton("Save", () => saveWorkerSettings(extension, worker.id), "background-color: #4a7c4a;");
|
||||
saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success;
|
||||
|
||||
@@ -126,9 +207,11 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
const buttonGroup = ui.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;");
|
||||
form.appendChild(buttonGroup);
|
||||
|
||||
typeSelect.onchange = createOnChangeHandler();
|
||||
typeSelect.onchange = (event) => updateWorkerType(event.target.value, true);
|
||||
|
||||
if (worker.type === "cloud") {
|
||||
if (worker.type === "serverless") {
|
||||
typeSelect.value = "serverless";
|
||||
} else if (worker.type === "cloud") {
|
||||
typeSelect.value = "cloud";
|
||||
runpodText.style.display = "block";
|
||||
} else if (isRemoteWorker(extension, worker)) {
|
||||
@@ -137,7 +220,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
typeSelect.value = "local";
|
||||
}
|
||||
|
||||
typeSelect.dispatchEvent(new Event('change'));
|
||||
updateWorkerType(typeSelect.value);
|
||||
|
||||
return form;
|
||||
}
|
||||
|
||||
+106
-1
@@ -41,6 +41,41 @@ function setButtonVisibility(button, visible) {
|
||||
button.style.display = visible ? "" : "none";
|
||||
}
|
||||
|
||||
export function formatServerlessStatus(status) {
|
||||
if (!status) return "Checking status…";
|
||||
if (!status.ok) return "Unavailable";
|
||||
const jobs = (Array.isArray(status.jobs) ? status.jobs : [])
|
||||
.map((job) => String(job?.status || "").trim().toUpperCase());
|
||||
// Prefer current work over old ledger entries, and avoid raw provider enums.
|
||||
const phases = [
|
||||
["Downloading result", ["DOWNLOADING"]],
|
||||
["Cancelling", ["CANCEL_REQUESTED"]],
|
||||
["Running", ["RUNNING", "PROCESSING", "IN_PROGRESS"]],
|
||||
["Starting", ["STARTING"]],
|
||||
["Submitting", ["SUBMITTING"]],
|
||||
["Queued", ["WAITING", "IN_QUEUE", "QUEUED", "PENDING"]],
|
||||
];
|
||||
for (const [label, states] of phases) {
|
||||
const count = jobs.filter((state) => states.includes(state)).length;
|
||||
if (count) return count > 1 ? `${label} · ${count} jobs` : label;
|
||||
}
|
||||
// Health counts queued AND running jobs, not just the waiting queue.
|
||||
const remaining = Number(status.queueRemaining);
|
||||
if (Number.isFinite(remaining) && remaining > 0) {
|
||||
return `Active · ${remaining} ${remaining === 1 ? "job" : "jobs"}`;
|
||||
}
|
||||
if (jobs[0] === "FAILED") return "Last job failed";
|
||||
if (jobs[0] === "TIMED_OUT") return "Last job timed out";
|
||||
return "Ready";
|
||||
}
|
||||
|
||||
function updateServerlessStatusMessage(workerId, status) {
|
||||
const messageElement = document.getElementById(`serverless-status-${workerId}`);
|
||||
if (!messageElement) return;
|
||||
messageElement.textContent = formatServerlessStatus(status);
|
||||
messageElement.title = status?.message || "";
|
||||
}
|
||||
|
||||
export async function checkAllWorkerStatuses(extension) {
|
||||
if (_statusCheckRunning || !extension.panelElement) {
|
||||
return;
|
||||
@@ -136,11 +171,60 @@ export async function checkMasterStatus(extension) {
|
||||
|
||||
// Helper to build worker URL
|
||||
export function getWorkerUrl(extension, worker, endpoint = "") {
|
||||
if (worker?.type === "serverless") {
|
||||
return null;
|
||||
}
|
||||
return buildWorkerUrl(worker, endpoint, window.location);
|
||||
}
|
||||
|
||||
export async function checkWorkerStatus(extension, worker) {
|
||||
// Assume caller ensured enabled; proceed with check
|
||||
if (worker.type === "serverless") {
|
||||
try {
|
||||
const result = await extension.api.getServerlessStatus(worker.id);
|
||||
const queueRemaining = Number(result?.queueRemaining);
|
||||
const normalizedQueue = Number.isFinite(queueRemaining) ? Math.max(0, queueRemaining) : 0;
|
||||
const normalizedStatus = {
|
||||
ok: Boolean(result?.ok),
|
||||
queueRemaining: normalizedQueue,
|
||||
message: typeof result?.message === "string" ? result.message : "",
|
||||
jobs: (Array.isArray(result?.jobs) ? result.jobs : []).map((job) => ({
|
||||
job_id: String(job?.job_id || ""),
|
||||
status: String(job?.status || ""),
|
||||
worker_id: String(job?.worker_id || ""),
|
||||
execution_id: String(job?.execution_id || ""),
|
||||
})),
|
||||
};
|
||||
const hasActiveJobs = normalizedStatus.jobs.some((job) =>
|
||||
/^(waiting|submitting|in_queue|queued|pending|running|processing|starting|in_progress|downloading|cancel_requested)$/i.test(job.status.trim())
|
||||
);
|
||||
const status = {
|
||||
online: normalizedStatus.ok,
|
||||
processing: normalizedQueue > 0 || hasActiveJobs,
|
||||
queueCount: normalizedQueue,
|
||||
};
|
||||
extension.state.updateWorker(worker.id, { serverlessStatus: normalizedStatus });
|
||||
extension.state.setWorkerStatus(worker.id, status);
|
||||
updateServerlessStatusMessage(worker.id, normalizedStatus);
|
||||
const updatedInPlace = extension.updateWorkerCard?.(worker.id, status);
|
||||
if (!updatedInPlace) {
|
||||
updateWorkerControls(extension, worker.id);
|
||||
}
|
||||
} catch (error) {
|
||||
if (error?.name === "AbortError") return;
|
||||
const failedStatus = { ok: false, queueRemaining: 0, message: error?.message || "Runpod status unavailable", jobs: [] };
|
||||
extension.state.updateWorker(worker.id, { serverlessStatus: failedStatus });
|
||||
extension.state.setWorkerStatus(worker.id, { online: false, processing: false, queueCount: 0 });
|
||||
updateServerlessStatusMessage(worker.id, failedStatus);
|
||||
const updatedInPlace = extension.updateWorkerCard?.(worker.id, extension.state.getWorkerStatus(worker.id));
|
||||
if (!updatedInPlace) {
|
||||
updateWorkerControls(extension, worker.id);
|
||||
}
|
||||
extension.log(`Serverless worker ${worker.id} status check failed: ${error.message}`, "debug");
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const workerUrl = getWorkerUrl(extension, worker);
|
||||
|
||||
try {
|
||||
@@ -206,6 +290,9 @@ export async function checkWorkerStatus(extension, worker) {
|
||||
|
||||
export async function launchWorker(extension, workerId) {
|
||||
const worker = extension.config.workers.find((w) => w.id === workerId);
|
||||
if (!worker || worker.type === "serverless") {
|
||||
return;
|
||||
}
|
||||
|
||||
// If worker is disabled, enable it first
|
||||
if (!worker.enabled) {
|
||||
@@ -273,6 +360,9 @@ export async function launchWorker(extension, workerId) {
|
||||
|
||||
export async function stopWorker(extension, workerId) {
|
||||
const worker = extension.config.workers.find((w) => w.id === workerId);
|
||||
if (!worker || worker.type === "serverless") {
|
||||
return;
|
||||
}
|
||||
const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1];
|
||||
|
||||
// Provide immediate feedback
|
||||
@@ -349,8 +439,12 @@ export async function loadManagedWorkers(extension) {
|
||||
try {
|
||||
const result = await extension.api.getManagedWorkers();
|
||||
|
||||
// Check for launching workers
|
||||
const serverlessWorkerIds = new Set(
|
||||
(extension.config?.workers || []).filter((worker) => worker.type === "serverless").map((worker) => String(worker.id))
|
||||
);
|
||||
// Managed-worker lifecycle state is ComfyUI-process-only; serverless is provider-managed.
|
||||
for (const [workerId, info] of Object.entries(result.managed_workers)) {
|
||||
if (serverlessWorkerIds.has(String(workerId))) continue;
|
||||
extension.state.setWorkerManaged(workerId, info);
|
||||
|
||||
// If worker is marked as launching, add to launchingWorkers set
|
||||
@@ -386,6 +480,13 @@ export function updateWorkerControls(extension, workerId) {
|
||||
const stopBtn = document.getElementById(`stop-${workerId}`);
|
||||
const logBtn = document.getElementById(`log-${workerId}`);
|
||||
|
||||
if (worker.type === "serverless") {
|
||||
setButtonVisibility(launchBtn, false);
|
||||
setButtonVisibility(stopBtn, false);
|
||||
setButtonVisibility(logBtn, false);
|
||||
return;
|
||||
}
|
||||
|
||||
if (isRemoteWorker(extension, worker)) {
|
||||
setButtonVisibility(launchBtn, false);
|
||||
setButtonVisibility(stopBtn, false);
|
||||
@@ -438,6 +539,9 @@ export function updateWorkerControls(extension, workerId) {
|
||||
|
||||
export async function viewWorkerLog(extension, workerId, isRemote = false) {
|
||||
const worker = extension.config.workers.find((w) => w.id === workerId);
|
||||
if (worker?.type === "serverless") {
|
||||
return;
|
||||
}
|
||||
const isRemoteLog = isRemote || (worker ? isRemoteWorker(extension, worker) : false);
|
||||
const managedInfo = extension.state.getWorker(workerId).managed;
|
||||
if (!isRemoteLog && !managedInfo?.log_file) {
|
||||
@@ -498,6 +602,7 @@ export async function refreshLog(extension, workerId, silent = false) {
|
||||
|
||||
try {
|
||||
const worker = extension.config.workers.find((w) => w.id === workerId);
|
||||
if (worker?.type === "serverless") return;
|
||||
const isRemoteLog = worker ? isRemoteWorker(extension, worker) : false;
|
||||
const data = isRemoteLog
|
||||
? await extension.api.getRemoteWorkerLog(workerId, 300)
|
||||
|
||||
+149
-2
@@ -18,7 +18,7 @@ export function isRemoteWorker(extension, worker) {
|
||||
const workerType = String(worker?.type || "").toLowerCase();
|
||||
|
||||
// Explicit type always wins over host heuristics.
|
||||
if (workerType === "cloud" || workerType === "remote") {
|
||||
if (workerType === "cloud" || workerType === "remote" || workerType === "serverless") {
|
||||
return true;
|
||||
}
|
||||
if (workerType === "local") {
|
||||
@@ -37,6 +37,134 @@ export function isCloudWorker(extension, worker) {
|
||||
return worker.type === "cloud";
|
||||
}
|
||||
|
||||
function showServerlessValidationError(extension, detail) {
|
||||
extension.app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "Validation Error",
|
||||
detail,
|
||||
life: 3000,
|
||||
});
|
||||
}
|
||||
|
||||
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
|
||||
const endpointId = String(document.getElementById(`endpoint-id-${workerId}`).value || "").trim();
|
||||
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
|
||||
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
|
||||
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
|
||||
const concurrency = Number(document.getElementById(`concurrency-${workerId}`).value);
|
||||
const timeoutSeconds = Number(document.getElementById(`timeout-seconds-${workerId}`).value);
|
||||
|
||||
if (!name.trim()) {
|
||||
showServerlessValidationError(extension, "Worker name is required");
|
||||
return;
|
||||
}
|
||||
if (!endpointId) {
|
||||
showServerlessValidationError(extension, "Endpoint ID is required for Runpod serverless workers");
|
||||
return;
|
||||
}
|
||||
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
|
||||
showServerlessValidationError(extension, "Jobs per run must be an integer between 1 and 16");
|
||||
return;
|
||||
}
|
||||
if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 16) {
|
||||
showServerlessValidationError(extension, "Concurrency must be an integer between 1 and 16");
|
||||
return;
|
||||
}
|
||||
if (!Number.isInteger(timeoutSeconds) || timeoutSeconds < 30 || timeoutSeconds > 86400) {
|
||||
showServerlessValidationError(extension, "Timeout must be an integer between 30 and 86400 seconds");
|
||||
return;
|
||||
}
|
||||
|
||||
const updateData = {
|
||||
name: name.trim(),
|
||||
type: "serverless",
|
||||
provider: "runpod",
|
||||
endpoint_id: endpointId,
|
||||
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
|
||||
...(clearApiKey ? { clear_api_key: true } : {}),
|
||||
jobs_per_run: jobsPerRun,
|
||||
concurrency,
|
||||
timeout_seconds: timeoutSeconds,
|
||||
// Never enable a serverless worker as a side effect of saving credentials.
|
||||
enabled: worker.enabled,
|
||||
};
|
||||
|
||||
try {
|
||||
await extension.api.updateWorker(workerId, updateData);
|
||||
document.getElementById(`api-key-${workerId}`).value = "";
|
||||
|
||||
worker.name = name.trim();
|
||||
worker.type = "serverless";
|
||||
worker.provider = "runpod";
|
||||
worker.endpoint_id = endpointId;
|
||||
worker.jobs_per_run = jobsPerRun;
|
||||
worker.concurrency = concurrency;
|
||||
worker.timeout_seconds = timeoutSeconds;
|
||||
worker.api_key_configured = clearApiKey
|
||||
? false
|
||||
: apiKey.length > 0 || Boolean(worker.api_key_configured);
|
||||
// The credential is write-only: never copy it into extension config/state.
|
||||
delete worker.api_key;
|
||||
delete worker.host;
|
||||
delete worker.port;
|
||||
delete worker.cuda_device;
|
||||
delete worker.extra_args;
|
||||
|
||||
extension.state.updateWorker(workerId, { enabled: worker.enabled });
|
||||
emitWorkersChanged(extension);
|
||||
extension.app.extensionManager.toast.add({
|
||||
severity: "success",
|
||||
summary: "Settings Saved",
|
||||
detail: `Worker ${worker.name} settings updated`,
|
||||
life: 3000,
|
||||
});
|
||||
extension.state.setWorkerExpanded(workerId, false);
|
||||
if (extension.panelElement) {
|
||||
renderSidebarContent(extension, extension.panelElement);
|
||||
}
|
||||
} catch (error) {
|
||||
extension.app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "Save Failed",
|
||||
detail: error.message,
|
||||
life: 5000,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
export async function testServerlessWorker(extension, workerId) {
|
||||
const button = document.getElementById(`test-serverless-${workerId}`);
|
||||
const originalText = button?.textContent;
|
||||
if (button) {
|
||||
button.disabled = true;
|
||||
button.textContent = "Testing...";
|
||||
}
|
||||
try {
|
||||
// This endpoint only checks the saved endpoint configuration; it does not enqueue inference.
|
||||
const result = await extension.api.testServerlessWorker(workerId);
|
||||
extension.app.extensionManager.toast.add({
|
||||
severity: result?.ok ? "success" : "error",
|
||||
summary: result?.ok ? "Runpod Endpoint Ready" : "Runpod Endpoint Test Failed",
|
||||
detail: result?.message || (result?.ok ? "Saved endpoint is reachable" : "Endpoint test failed"),
|
||||
life: 4000,
|
||||
});
|
||||
return result;
|
||||
} catch (error) {
|
||||
extension.app.extensionManager.toast.add({
|
||||
severity: "error",
|
||||
summary: "Runpod Endpoint Test Failed",
|
||||
detail: error.message,
|
||||
life: 5000,
|
||||
});
|
||||
return null;
|
||||
} finally {
|
||||
if (button) {
|
||||
button.disabled = false;
|
||||
button.textContent = originalText;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function saveWorkerSettings(extension, workerId) {
|
||||
const worker = extension.config.workers.find((w) => w.id === workerId);
|
||||
if (!worker) {
|
||||
@@ -44,8 +172,12 @@ export async function saveWorkerSettings(extension, workerId) {
|
||||
}
|
||||
|
||||
// Get form values
|
||||
const name = document.getElementById(`name-${workerId}`).value;
|
||||
const workerType = document.getElementById(`worker-type-${workerId}`).value;
|
||||
if (workerType === "serverless") {
|
||||
await saveServerlessWorkerSettings(extension, worker, workerId, document.getElementById(`name-${workerId}`).value);
|
||||
return;
|
||||
}
|
||||
const name = document.getElementById(`name-${workerId}`).value;
|
||||
const isRemote = workerType === "remote" || workerType === "cloud";
|
||||
const isCloud = workerType === "cloud";
|
||||
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
|
||||
@@ -209,6 +341,21 @@ export function cancelWorkerSettings(extension, workerId) {
|
||||
document.getElementById(`port-${workerId}`).value = worker.port;
|
||||
document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0;
|
||||
document.getElementById(`args-${workerId}`).value = worker.extra_args || "";
|
||||
const endpointInput = document.getElementById(`endpoint-id-${workerId}`);
|
||||
if (endpointInput) endpointInput.value = worker.endpoint_id || "";
|
||||
const apiKeyInput = document.getElementById(`api-key-${workerId}`);
|
||||
if (apiKeyInput) {
|
||||
apiKeyInput.value = "";
|
||||
apiKeyInput.disabled = false;
|
||||
}
|
||||
const clearApiKeyInput = document.getElementById(`clear-api-key-${workerId}`);
|
||||
if (clearApiKeyInput) clearApiKeyInput.checked = false;
|
||||
const jobsPerRunInput = document.getElementById(`jobs-per-run-${workerId}`);
|
||||
if (jobsPerRunInput) jobsPerRunInput.value = worker.jobs_per_run ?? 1;
|
||||
const concurrencyInput = document.getElementById(`concurrency-${workerId}`);
|
||||
if (concurrencyInput) concurrencyInput.value = worker.concurrency ?? 1;
|
||||
const timeoutInput = document.getElementById(`timeout-seconds-${workerId}`);
|
||||
if (timeoutInput) timeoutInput.value = worker.timeout_seconds ?? 1800;
|
||||
|
||||
// Reset remote checkbox
|
||||
const remoteCheckbox = document.getElementById(`remote-${workerId}`);
|
||||
|
||||
+20
-10
@@ -19,12 +19,17 @@ export async function handleWorkerOperation(extension, button, operation, succes
|
||||
setButtonStateClass("btn--working");
|
||||
|
||||
try {
|
||||
const urlsToProcess = extension.enabledWorkers.map(w => ({
|
||||
name: w.name,
|
||||
url: getWorkerUrl(extension, w)
|
||||
const workersToProcess = extension.enabledWorkers.filter((worker) =>
|
||||
operation.endpoint !== ENDPOINTS.CLEAR_MEMORY || worker.type !== "serverless"
|
||||
);
|
||||
const targets = workersToProcess.map((worker) => ({
|
||||
name: worker.name,
|
||||
worker,
|
||||
serverless: worker.type === "serverless",
|
||||
url: worker.type === "serverless" ? null : getWorkerUrl(extension, worker),
|
||||
}));
|
||||
|
||||
if (urlsToProcess.length === 0) {
|
||||
|
||||
if (targets.length === 0) {
|
||||
button.textContent = "No Workers";
|
||||
setButtonStateClass("btn--error");
|
||||
setTimeout(() => {
|
||||
@@ -37,14 +42,19 @@ export async function handleWorkerOperation(extension, button, operation, succes
|
||||
return;
|
||||
}
|
||||
|
||||
const promises = urlsToProcess.map(target =>
|
||||
fetch(`${target.url}${operation.endpoint}`, {
|
||||
method: 'POST',
|
||||
const promises = targets.map((target) => {
|
||||
if (target.serverless && operation.name === "Interrupt") {
|
||||
return extension.api.cancelServerlessWorker(target.worker.id)
|
||||
.then((result) => ({ ok: result?.ok !== false, name: target.name }))
|
||||
.catch(() => ({ ok: false, name: target.name }));
|
||||
}
|
||||
return fetch(`${target.url}${operation.endpoint}`, {
|
||||
method: 'POST',
|
||||
mode: 'cors'
|
||||
})
|
||||
.then(response => ({ ok: response.ok, name: target.name }))
|
||||
.catch(() => ({ ok: false, name: target.name }))
|
||||
);
|
||||
.catch(() => ({ ok: false, name: target.name }));
|
||||
});
|
||||
|
||||
const results = await Promise.all(promises);
|
||||
const failures = results.filter(r => !r.ok);
|
||||
|
||||
@@ -10,6 +10,8 @@ from ..utils.logging import debug_log
|
||||
|
||||
async def is_local_worker(worker_config):
|
||||
"""Check if a worker is running on the same machine as the master."""
|
||||
if worker_config.get('type') == 'serverless':
|
||||
return False
|
||||
host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
|
||||
if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local':
|
||||
return True
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Host-local worker port selection and launch preflight (no ComfyUI imports)."""
|
||||
import errno
|
||||
import os
|
||||
import socket
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
|
||||
def _port_number(value):
|
||||
try:
|
||||
port = int(value)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("Worker port must be an integer between 1 and 65535") from None
|
||||
if isinstance(value, (bool, float)) or not 1 <= port <= 65535:
|
||||
raise ValueError("Worker port must be an integer between 1 and 65535")
|
||||
return port
|
||||
|
||||
|
||||
def _local_config(worker):
|
||||
local_hosts = ("localhost", "127.0.0.1", "0.0.0.0", "::1", "::")
|
||||
host = (worker.get("host") or "localhost").strip().lower()
|
||||
if worker.get("type") == "local" or host in local_hosts:
|
||||
return True
|
||||
try:
|
||||
host = urlsplit(host if "://" in host else "//" + host).hostname
|
||||
except ValueError:
|
||||
return False
|
||||
return host in local_hosts
|
||||
|
||||
|
||||
def _reserved_ports(workers, exclude_id=None):
|
||||
reserved = {}
|
||||
for worker in workers:
|
||||
if exclude_id is not None and str(worker.get("id")) == str(exclude_id):
|
||||
continue
|
||||
if not _local_config(worker):
|
||||
continue
|
||||
try:
|
||||
port = _port_number(worker.get("port"))
|
||||
except ValueError:
|
||||
continue
|
||||
reserved[port] = worker.get("name") or worker.get("id") or "another local worker"
|
||||
return reserved
|
||||
|
||||
|
||||
def is_port_available(port):
|
||||
"""Match asyncio's address reuse while rejecting active IPv4/IPv6 listeners."""
|
||||
addresses = [(socket.AF_INET, "0.0.0.0")]
|
||||
if socket.has_ipv6:
|
||||
addresses.append((socket.AF_INET6, "::"))
|
||||
for family, address in addresses:
|
||||
try:
|
||||
with socket.socket(family, socket.SOCK_STREAM) as probe:
|
||||
if os.name != "nt":
|
||||
# Like asyncio, allow a restart while old connections are
|
||||
# in TIME_WAIT. Do not enable SO_REUSEPORT.
|
||||
probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
elif hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
|
||||
# Windows SO_REUSEADDR can steal an active listener's port.
|
||||
probe.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
|
||||
probe.bind((address, port))
|
||||
probe.listen(1)
|
||||
except OSError as exc:
|
||||
if family == socket.AF_INET6 and exc.errno in (errno.EAFNOSUPPORT, errno.EADDRNOTAVAIL):
|
||||
continue
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def allocate_worker_ports(master_port, workers, count):
|
||||
"""Select ports above the actual master, skipping reservations/listeners.
|
||||
|
||||
This is a snapshot, not a reservation. Launch must validate again because
|
||||
another process can claim a suggested port before the user starts a worker.
|
||||
"""
|
||||
master_port = _port_number(master_port)
|
||||
reserved = _reserved_ports(workers)
|
||||
selected = []
|
||||
if count == 0:
|
||||
return selected
|
||||
for port in range(master_port + 1, 65536):
|
||||
if port not in reserved and is_port_available(port):
|
||||
selected.append(port)
|
||||
if len(selected) == count:
|
||||
return selected
|
||||
raise ValueError(f"Not enough available worker ports above master port {master_port}; need {count}")
|
||||
|
||||
|
||||
def validate_worker_port(worker, master_port, workers):
|
||||
"""Reject conflicts without silently changing a configured worker port."""
|
||||
port = _port_number(worker.get("port"))
|
||||
if port == _port_number(master_port):
|
||||
raise ValueError(f"Worker port {port} conflicts with the master port {master_port}; choose another worker port")
|
||||
reserved = _reserved_ports(workers, exclude_id=worker.get("id"))
|
||||
if port in reserved:
|
||||
raise ValueError(f"Worker port {port} is assigned to local worker {reserved[port]}; choose another worker port")
|
||||
if not is_port_available(port):
|
||||
raise ValueError(f"Worker port {port} is already in use on this host; choose another worker port")
|
||||
@@ -1,4 +1,5 @@
|
||||
import glob
|
||||
import hashlib
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
@@ -87,6 +88,51 @@ class LaunchCommandBuilder:
|
||||
return wt_path
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _arg_value(cmd, flag):
|
||||
"""Return the last value, matching argparse's override semantics."""
|
||||
value = None
|
||||
for index, arg in enumerate(cmd):
|
||||
if arg == flag and index + 1 < len(cmd):
|
||||
value = cmd[index + 1]
|
||||
elif arg.startswith(flag + "="):
|
||||
value = arg.split("=", 1)[1]
|
||||
return value
|
||||
|
||||
def _add_worker_database(self, cmd, worker_config, comfy_root):
|
||||
args = self._get_runtime_args()
|
||||
# The parsed namespace advertises support without launching ComfyUI
|
||||
# (or importing another checkout's argument parser).
|
||||
if args is None or not hasattr(args, "database_url"):
|
||||
return
|
||||
if any(arg.split("=", 1)[0] in ("--database-url", "--disable-assets") for arg in cmd):
|
||||
return
|
||||
|
||||
user_directory = self._arg_value(cmd, "--user-directory")
|
||||
base_directory = self._arg_value(cmd, "--base-directory")
|
||||
if not user_directory:
|
||||
user_directory = os.path.join(base_directory or comfy_root, "user")
|
||||
if not os.path.isabs(user_directory):
|
||||
# Popen uses comfy_root as cwd, which can differ from the master's.
|
||||
user_directory = os.path.join(comfy_root, user_directory)
|
||||
database_directory = os.path.abspath(os.path.join(user_directory, "distributed", "workers"))
|
||||
# Hash the immutable ID: names/GPU/ports may change, and IDs must not
|
||||
# become filesystem paths or collide after filename sanitization.
|
||||
worker_id = str(worker_config["id"])
|
||||
filename = hashlib.sha256(worker_id.encode("utf-8")).hexdigest() + ".db"
|
||||
database_path = os.path.join(database_directory, filename).replace(os.sep, "/")
|
||||
# Use the same URL codec as ComfyUI; raw '?' and '%xx' can otherwise
|
||||
# truncate/decode the path and make distinct workers share a database.
|
||||
from sqlalchemy.engine import URL, make_url
|
||||
database_url = URL.create("sqlite", database=database_path).render_as_string()
|
||||
if make_url(database_url).database != database_path:
|
||||
raise ValueError(
|
||||
"Installed SQLAlchemy cannot safely encode the worker database path; "
|
||||
"use a user directory without '?' or an explicit --database-url with a safe path"
|
||||
)
|
||||
os.makedirs(database_directory, exist_ok=True)
|
||||
cmd.extend(["--database-url", database_url])
|
||||
|
||||
def build_launch_command(self, worker_config, comfy_root):
|
||||
"""Build the command to launch a worker."""
|
||||
main_py = os.path.join(comfy_root, "main.py")
|
||||
@@ -141,4 +187,7 @@ class LaunchCommandBuilder:
|
||||
raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {forbidden}")
|
||||
cmd.extend(extra_args_list)
|
||||
|
||||
if int(self._arg_value(cmd, "--port")) != int(worker_config["port"]):
|
||||
raise ValueError("extra_args must not override the configured worker port; set the port in the UI/config instead")
|
||||
self._add_worker_database(cmd, worker_config, comfy_root)
|
||||
return cmd
|
||||
|
||||
@@ -7,7 +7,9 @@ import time
|
||||
from ...utils.config import load_config, save_config
|
||||
from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL
|
||||
from ...utils.logging import debug_log, log
|
||||
from ...utils.network import get_server_port
|
||||
from ...utils.process import get_python_executable, is_process_alive, terminate_process
|
||||
from ..ports import validate_worker_port
|
||||
|
||||
try:
|
||||
import psutil
|
||||
@@ -28,6 +30,17 @@ class ProcessLifecycle:
|
||||
"""Launch a worker process with logging."""
|
||||
_ = show_window # Kept for API compatibility.
|
||||
comfy_root = self._manager.find_comfy_root()
|
||||
config = load_config()
|
||||
try:
|
||||
master_port = get_server_port()
|
||||
except AttributeError:
|
||||
master_port = None
|
||||
if master_port is None:
|
||||
# The auto-launch timer can fire during custom-node imports, before
|
||||
# PromptServer assigns its port. Still reserve the configured port.
|
||||
from comfy.cli_args import args
|
||||
master_port = args.port
|
||||
validate_worker_port(worker_config, master_port, config.get("workers", []))
|
||||
|
||||
env = os.environ.copy()
|
||||
env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))
|
||||
|
||||
@@ -22,6 +22,8 @@ class WorkerProcessManager:
|
||||
return self._launch_builder.build_launch_command(worker_config, comfy_root)
|
||||
|
||||
def launch_worker(self, worker_config, show_window=False):
|
||||
if worker_config.get("type") == "serverless":
|
||||
raise ValueError("Serverless pools cannot be launched as local processes")
|
||||
return self._lifecycle.launch_worker(worker_config, show_window=show_window)
|
||||
|
||||
def stop_worker(self, worker_id):
|
||||
|
||||
@@ -35,6 +35,8 @@ def auto_launch_workers():
|
||||
launched_count = 0
|
||||
|
||||
for worker in workers:
|
||||
if worker.get('type') == 'serverless':
|
||||
continue
|
||||
if worker.get('enabled', False):
|
||||
worker_id = worker.get('id')
|
||||
worker_name = worker.get('name', f'Worker {worker_id}')
|
||||
|
||||
Reference in New Issue
Block a user