591 lines
23 KiB
Python
591 lines
23 KiB
Python
import json
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from aiohttp import web
|
|
import server
|
|
|
|
try:
|
|
from ..utils.config import config_transaction, load_config, save_config
|
|
except ImportError:
|
|
from ..utils.config import load_config
|
|
|
|
try:
|
|
from ..utils.config import save_config
|
|
except ImportError:
|
|
def save_config(_config):
|
|
return True
|
|
|
|
@asynccontextmanager
|
|
async def config_transaction():
|
|
config = load_config()
|
|
original_snapshot = json.dumps(config, sort_keys=True)
|
|
yield config
|
|
if json.dumps(config, sort_keys=True) != original_snapshot:
|
|
save_config(config)
|
|
from ..utils.logging import debug_log, log
|
|
from ..utils.network import handle_api_error, normalize_host
|
|
from ..utils.secret_fields import is_secret_field_name
|
|
from ..utils.serverless_config import (
|
|
ServerlessSecretChanges,
|
|
public_config,
|
|
validate_serverless_worker,
|
|
)
|
|
|
|
|
|
def _positive_int(value):
|
|
return value > 0
|
|
|
|
|
|
CONFIG_SCHEMA = {
|
|
"workers": (list, None),
|
|
"master": (dict, None),
|
|
"settings": (dict, None),
|
|
"tunnel": (dict, None),
|
|
"managed_processes": (dict, None),
|
|
"worker_timeout_seconds": (int, _positive_int),
|
|
"debug": (bool, None),
|
|
"auto_launch_workers": (bool, None),
|
|
"stop_workers_on_master_exit": (bool, None),
|
|
"master_delegate_only": (bool, None),
|
|
"websocket_orchestration": (bool, None),
|
|
"has_auto_populated_workers": (bool, None),
|
|
}
|
|
|
|
_SETTINGS_FIELDS = {
|
|
"worker_timeout_seconds",
|
|
"debug",
|
|
"auto_launch_workers",
|
|
"stop_workers_on_master_exit",
|
|
"master_delegate_only",
|
|
"websocket_orchestration",
|
|
"has_auto_populated_workers",
|
|
}
|
|
|
|
_WORKER_FIELDS = [
|
|
("enabled", None, False),
|
|
("name", None, False),
|
|
("port", None, False),
|
|
("host", normalize_host, True),
|
|
("cuda_device", None, True),
|
|
("extra_args", None, True),
|
|
("type", None, False),
|
|
]
|
|
|
|
_MASTER_FIELDS = [
|
|
("name", None, False),
|
|
("host", normalize_host, True),
|
|
("port", None, False),
|
|
("cuda_device", None, False),
|
|
("extra_args", None, False),
|
|
]
|
|
|
|
|
|
def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
|
|
"""Apply a partial update to a target dict based on field rules."""
|
|
for key, normalizer, remove_on_none in field_rules:
|
|
if key not in data:
|
|
continue
|
|
value = data[key]
|
|
if value is None and remove_on_none:
|
|
target.pop(key, None)
|
|
else:
|
|
target[key] = normalizer(value) if (normalizer and value is not None) else value
|
|
|
|
|
|
_SECRET_FIELDS = {"api_key", "clear_api_key"}
|
|
_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
|
|
_SERVERLESS_ONLY_FIELDS = {
|
|
"provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
|
|
"media_profile",
|
|
}
|
|
_SERVERLESS_CONFIG_FIELDS = {
|
|
"worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
|
|
"timeout_seconds", "enabled", "name", "api_key", "clear_api_key", "media_profile",
|
|
}
|
|
|
|
|
|
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":
|
|
continue
|
|
if worker.get("type") == "serverless":
|
|
# Older clients do not know this field. Only an explicit selection
|
|
# may change the endpoint contract, including on bulk saves.
|
|
if "media_profile" not in worker and "media_profile" in old:
|
|
worker["media_profile"] = old["media_profile"]
|
|
continue
|
|
converted = dict(worker)
|
|
for field in _SERVERLESS_ONLY_FIELDS:
|
|
converted.pop(field, None)
|
|
new_workers[index] = converted
|
|
return new_workers
|
|
|
|
|
|
def _bulk_serverless_changes(old_workers, new_workers):
|
|
old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
|
|
new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
|
|
secret_ids = set()
|
|
for worker_id, old in old_by_id.items():
|
|
new = new_by_id.get(worker_id)
|
|
if old.get("type") == "serverless":
|
|
if new is None or new.get("type") != "serverless":
|
|
_ensure_pool_idle(worker_id)
|
|
secret_ids.add(worker_id)
|
|
elif (
|
|
old.get("endpoint_id") != new.get("endpoint_id")
|
|
or old.get("media_profile", "legacy-image-v1") != new.get("media_profile", "legacy-image-v1")
|
|
):
|
|
_ensure_pool_idle(worker_id)
|
|
elif new is not None and new.get("type") == "serverless":
|
|
secret_ids.add(worker_id)
|
|
return secret_ids
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/distributed/config")
|
|
async def get_config_endpoint(request):
|
|
try:
|
|
config = load_config()
|
|
return web.json_response(public_config(config, _serverless_secret_store()))
|
|
except Exception:
|
|
return _error_response("Unable to read config", 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/config")
|
|
async def update_config_endpoint(request):
|
|
"""Bulk config update with schema validation and credential redaction."""
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return _error_response("Invalid JSON payload", 400)
|
|
|
|
if not isinstance(data, dict):
|
|
return _error_response("Config payload must be an object", 400)
|
|
if _contains_secret_fields(data):
|
|
return _error_response("API keys cannot be set through bulk config", 400)
|
|
|
|
validated_settings = {}
|
|
validated_root = {}
|
|
errors = []
|
|
|
|
for key, value in data.items():
|
|
if key not in CONFIG_SCHEMA:
|
|
errors.append(f"Unknown field: {key}")
|
|
continue
|
|
|
|
expected_type, validator = CONFIG_SCHEMA[key]
|
|
if not isinstance(value, expected_type):
|
|
errors.append(f"{key}: expected {expected_type.__name__}")
|
|
continue
|
|
|
|
if validator and not validator(value):
|
|
errors.append(f"{key}: failed validation")
|
|
continue
|
|
|
|
if key == "workers":
|
|
try:
|
|
value = _validate_bulk_workers(value)
|
|
except (TypeError, ValueError):
|
|
return _error_response("Invalid serverless worker configuration", 400)
|
|
|
|
if key in _SETTINGS_FIELDS:
|
|
validated_settings[key] = value
|
|
else:
|
|
validated_root[key] = value
|
|
|
|
if errors:
|
|
return web.json_response({
|
|
"status": "error",
|
|
"error": errors,
|
|
"message": "; ".join(errors),
|
|
}, status=400)
|
|
|
|
store = None
|
|
response_config = None
|
|
try:
|
|
with ServerlessSecretChanges() as secret_changes:
|
|
async with config_transaction() as config:
|
|
if "workers" in validated_root:
|
|
old_workers = config.get("workers", [])
|
|
validated_root["workers"] = _clean_bulk_type_conversions(
|
|
old_workers, validated_root["workers"]
|
|
)
|
|
changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"])
|
|
for worker in validated_root["workers"]:
|
|
if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False):
|
|
store = store or _serverless_secret_store()
|
|
if worker["id"] in changed_ids or not store.has(worker["id"]):
|
|
return _error_response("Save an API key before enabling this worker", 400)
|
|
if changed_ids:
|
|
store = store or _serverless_secret_store()
|
|
for worker_id in sorted(changed_ids):
|
|
prior_secret = secret_changes.capture(store, worker_id)
|
|
if prior_secret is not None:
|
|
secret_changes.delete(store, worker_id)
|
|
|
|
settings = config.setdefault("settings", {})
|
|
settings.update(validated_settings)
|
|
for key, value in validated_root.items():
|
|
config[key] = value
|
|
response_config = public_config(config, store or _serverless_secret_store())
|
|
response = web.json_response({"status": "success", "config": response_config})
|
|
secret_changes.commit()
|
|
return response
|
|
except _ActivePoolCheckError:
|
|
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
|
|
except Exception:
|
|
return _error_response("Unable to update config", 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
|
|
async def queue_status_endpoint(request):
|
|
"""Check if a job queue is initialized."""
|
|
try:
|
|
job_id = request.match_info['job_id']
|
|
|
|
# Import to ensure initialization
|
|
from ..upscale.job_store import ensure_tile_jobs_initialized
|
|
prompt_server = ensure_tile_jobs_initialized()
|
|
|
|
async with prompt_server.distributed_tile_jobs_lock:
|
|
exists = job_id in prompt_server.distributed_pending_tile_jobs
|
|
|
|
debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
|
|
return web.json_response({"exists": exists, "job_id": job_id})
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 500)
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/config/update_worker")
|
|
async def update_worker_endpoint(request):
|
|
return await _run_with_serverless_secret_changes(_update_worker_endpoint, request)
|
|
|
|
|
|
async def _run_with_serverless_secret_changes(endpoint, request):
|
|
with ServerlessSecretChanges() as secret_changes:
|
|
response = await endpoint(request, secret_changes)
|
|
if response.status < 400:
|
|
secret_changes.commit()
|
|
return response
|
|
|
|
|
|
async def _update_worker_endpoint(request, secret_changes):
|
|
store = None
|
|
worker_id = None
|
|
try:
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return _error_response("Invalid JSON payload", 400)
|
|
if not isinstance(data, dict):
|
|
return _error_response("Worker payload must be an object", 400)
|
|
|
|
worker_id = data.get("worker_id")
|
|
if worker_id is None:
|
|
return _error_response("Missing worker_id", 400)
|
|
|
|
if _contains_secret_fields(data, allow_root_serverless_secrets=True):
|
|
return _error_response("Invalid worker credentials", 400)
|
|
|
|
result = None
|
|
async with config_transaction() as config:
|
|
workers = config.setdefault("workers", [])
|
|
existing = _worker_by_id(workers, worker_id)
|
|
current_type = _worker_type(existing)
|
|
target_type = data.get("type", current_type or "local")
|
|
|
|
if target_type == "serverless":
|
|
unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS
|
|
if unknown:
|
|
return _error_response("Unknown serverless worker field", 400)
|
|
if not isinstance(worker_id, str) or not worker_id.strip():
|
|
return _error_response("worker_id must be a non-empty string", 400)
|
|
if "api_key" in data and not isinstance(data["api_key"], str):
|
|
return _error_response("Invalid worker credentials", 400)
|
|
if "clear_api_key" in data and type(data["clear_api_key"]) is not bool:
|
|
return _error_response("Invalid worker credentials", 400)
|
|
supplied_key = data.get("api_key", "")
|
|
clear_key = data.get("clear_api_key", False)
|
|
if supplied_key and clear_key:
|
|
return _error_response("Choose either api_key or clear_api_key", 400)
|
|
if "enabled" in data and type(data["enabled"]) is not bool:
|
|
return _error_response("enabled must be a boolean", 400)
|
|
if "name" in data and not isinstance(data["name"], str):
|
|
return _error_response("name must be a string", 400)
|
|
|
|
candidate = dict(existing or {})
|
|
candidate["id"] = worker_id
|
|
candidate["type"] = "serverless"
|
|
candidate.setdefault("provider", "runpod")
|
|
candidate.setdefault("enabled", False)
|
|
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
|
|
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name", "media_profile"):
|
|
if field in data:
|
|
candidate[field] = data[field]
|
|
candidate = _clean_serverless_worker(candidate)
|
|
validated = validate_serverless_worker(candidate)
|
|
if "enabled" in validated and type(validated["enabled"]) is not bool:
|
|
return _error_response("enabled must be a boolean", 400)
|
|
if "name" in validated and not isinstance(validated["name"], str):
|
|
return _error_response("name must be a string", 400)
|
|
|
|
store = _serverless_secret_store()
|
|
prior_secret = secret_changes.capture(store, worker_id)
|
|
converting_to_serverless = current_type != "serverless"
|
|
next_secret = None if converting_to_serverless else prior_secret
|
|
if supplied_key:
|
|
next_secret = supplied_key
|
|
elif clear_key:
|
|
next_secret = None
|
|
|
|
endpoint_changed = (
|
|
current_type == "serverless"
|
|
and existing.get("endpoint_id") != validated.get("endpoint_id")
|
|
)
|
|
media_profile_changed = (
|
|
current_type == "serverless"
|
|
and existing.get("media_profile", "legacy-image-v1") != validated.get("media_profile", "legacy-image-v1")
|
|
)
|
|
if current_type == "serverless" and (
|
|
endpoint_changed or media_profile_changed or next_secret != prior_secret
|
|
):
|
|
_ensure_pool_idle(worker_id)
|
|
|
|
if validated.get("enabled", False) and next_secret is None:
|
|
return _error_response("Save an API key before enabling this worker; turn the worker off before clearing its key", 400)
|
|
|
|
if next_secret != prior_secret:
|
|
if next_secret is None:
|
|
secret_changes.delete(store, worker_id)
|
|
else:
|
|
secret_changes.set(store, worker_id, next_secret)
|
|
|
|
if existing is None:
|
|
workers.append(validated)
|
|
else:
|
|
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 worker", 409)
|
|
except ValueError:
|
|
return _error_response("Invalid worker configuration", 400)
|
|
except Exception:
|
|
return _error_response("Unable to update worker configuration", 500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
|
|
async def delete_worker_endpoint(request):
|
|
return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
|
|
|
|
|
|
async def _delete_worker_endpoint(request, secret_changes):
|
|
store = None
|
|
worker_id = None
|
|
try:
|
|
try:
|
|
data = await request.json()
|
|
except Exception:
|
|
return _error_response("Invalid JSON payload", 400)
|
|
if not isinstance(data, dict):
|
|
return _error_response("Worker payload must be an object", 400)
|
|
worker_id = data.get("worker_id")
|
|
if worker_id is None:
|
|
return _error_response("Missing worker_id", 400)
|
|
|
|
result = None
|
|
async with config_transaction() as config:
|
|
workers = config.get("workers", [])
|
|
removed = _worker_by_id(workers, worker_id)
|
|
if removed is None:
|
|
return _error_response("Worker not found", 404)
|
|
if _worker_type(removed) == "serverless":
|
|
_ensure_pool_idle(worker_id)
|
|
store = _serverless_secret_store()
|
|
secret_worker_id = str(worker_id)
|
|
prior_secret = secret_changes.capture(store, secret_worker_id)
|
|
if prior_secret is not None:
|
|
secret_changes.delete(store, secret_worker_id)
|
|
workers.remove(removed)
|
|
result = {
|
|
"status": "success",
|
|
"message": "Worker deleted" if _worker_type(removed) == "serverless"
|
|
else f"Worker {removed.get('name', worker_id)} deleted",
|
|
}
|
|
|
|
return web.json_response(result)
|
|
except _ActivePoolCheckError:
|
|
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
|
|
except Exception:
|
|
return _error_response("Unable to delete worker", 500)
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/config/update_setting")
|
|
async def update_setting_endpoint(request):
|
|
"""Updates a specific key in the settings object."""
|
|
try:
|
|
data = await request.json()
|
|
key = data.get("key")
|
|
value = data.get("value")
|
|
|
|
if not key or value is None:
|
|
return await handle_api_error(request, "Missing 'key' or 'value' in request", 400)
|
|
if key not in _SETTINGS_FIELDS:
|
|
return await handle_api_error(request, f"Unknown setting: {key}", 400)
|
|
|
|
async with config_transaction() as config:
|
|
if 'settings' not in config:
|
|
config['settings'] = {}
|
|
|
|
config['settings'][key] = value
|
|
|
|
return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 400)
|
|
|
|
@server.PromptServer.instance.routes.post("/distributed/config/update_master")
|
|
async def update_master_endpoint(request):
|
|
"""Updates master configuration."""
|
|
try:
|
|
data = await request.json()
|
|
|
|
async with config_transaction() as config:
|
|
if 'master' not in config:
|
|
config['master'] = {}
|
|
_apply_field_patch(config['master'], data, _MASTER_FIELDS)
|
|
|
|
return web.json_response({"status": "success", "message": "Master configuration updated."})
|
|
except Exception as e:
|
|
return await handle_api_error(request, e, 400)
|