import json from contextlib import asynccontextmanager 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 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 @server.PromptServer.instance.routes.get("/distributed/config") async def get_config_endpoint(request): config = load_config() return web.json_response(config) @server.PromptServer.instance.routes.post("/distributed/config") async def update_config_endpoint(request): """Bulk config update with schema validation.""" try: data = await request.json() except Exception as e: return await handle_api_error(request, f"Invalid JSON payload: {e}", 400) if not isinstance(data, dict): return await handle_api_error(request, "Config payload must be an object", 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}: value {value!r} failed validation") continue 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) 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) @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): try: data = await request.json() worker_id = data.get("worker_id") if worker_id is None: return await handle_api_error(request, "Missing worker_id", 400) async with config_transaction() as config: worker_found = False workers = config.setdefault("workers", []) for worker in workers: if worker["id"] == worker_id: _apply_field_patch(worker, data, _WORKER_FIELDS) worker_found = True break 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) else: return await handle_api_error( request, f"Worker {worker_id} not found and missing required fields for creation", 404, ) 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): try: data = await request.json() worker_id = data.get("worker_id") if worker_id is None: return await handle_api_error(request, "Missing worker_id", 400) 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({ "status": "success", "message": f"Worker {removed_worker.get('name', worker_id)} deleted" }) except Exception as e: return await handle_api_error(request, e, 400) @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)