From 5eacd17aeb61265e21312d2e377224b0abb78c14 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Tue, 24 Feb 2026 01:38:41 +0000 Subject: [PATCH] refactor: implement REFACTOR_PLAN_2 across distributed stack --- api/__init__.py | 1 + api/config_routes.py | 228 +++++----- api/job_routes.py | 8 +- api/orchestration/dispatch.py | 34 +- api/orchestration/media_sync.py | 26 +- api/queue_orchestration.py | 63 +-- api/schemas.py | 18 + api/usdu_routes.py | 228 ++++++++++ api/worker_routes.py | 30 +- distributed.py | 6 +- nodes/collector.py | 1 - nodes/distributed_upscale.py | 2 +- upscale/conditioning.py | 32 ++ upscale/job_state.py | 6 +- upscale/job_store.py | 180 ++++++++ upscale/job_timeout.py | 150 +++++++ upscale/modes/dynamic.py | 110 ++--- upscale/modes/static.py | 31 +- upscale/payload_parsers.py | 64 +++ upscale/result_collector.py | 62 +-- upscale/tile_ops.py | 2 +- upscale/worker_comms.py | 4 + utils/cloudflare.py | 388 ----------------- utils/cloudflare/__init__.py | 5 + utils/cloudflare/binary.py | 83 ++++ utils/cloudflare/process_reader.py | 97 +++++ utils/cloudflare/state.py | 81 ++++ utils/cloudflare/tunnel.py | 207 +++++++++ utils/config.py | 26 +- utils/constants.py | 11 + utils/network.py | 10 +- utils/trace_logger.py | 13 + utils/usdu_managment.py | 653 ++--------------------------- web/apiClient.js | 138 ++++-- web/distributed.css | 171 ++++++++ web/executionUtils.js | 10 +- web/ui.js | 233 ++-------- web/ui/logModal.js | 192 +++++++++ web/workerLifecycle.js | 147 +++---- web/workerUtils.js | 31 ++ workers/detection.py | 34 +- 41 files changed, 2103 insertions(+), 1713 deletions(-) create mode 100644 api/usdu_routes.py create mode 100644 upscale/conditioning.py create mode 100644 upscale/job_store.py create mode 100644 upscale/job_timeout.py create mode 100644 upscale/payload_parsers.py delete mode 100644 utils/cloudflare.py create mode 100644 utils/cloudflare/__init__.py create mode 100644 utils/cloudflare/binary.py create mode 100644 utils/cloudflare/process_reader.py create mode 100644 utils/cloudflare/state.py create mode 100644 utils/cloudflare/tunnel.py create mode 100644 utils/trace_logger.py create mode 100644 web/ui/logModal.js diff --git a/api/__init__.py b/api/__init__.py index 0c01bef..d20f665 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -2,3 +2,4 @@ from . import config_routes # noqa: F401 from . import tunnel_routes # noqa: F401 from . import worker_routes # noqa: F401 from . import job_routes # noqa: F401 +from . import usdu_routes # noqa: F401 diff --git a/api/config_routes.py b/api/config_routes.py index dd39fb6..8b22876 100644 --- a/api/config_routes.py +++ b/api/config_routes.py @@ -3,7 +3,7 @@ import json from aiohttp import web import server -from ..utils.config import load_config, save_config +from ..utils.config import config_transaction, load_config from ..utils.logging import debug_log, log from ..utils.network import handle_api_error, normalize_host @@ -37,6 +37,36 @@ _SETTINGS_FIELDS = { "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): @@ -55,8 +85,8 @@ async def update_config_endpoint(request): if not isinstance(data, dict): return await handle_api_error(request, "Config payload must be an object", 400) - config = load_config() - settings = config.setdefault("settings", {}) + validated_settings = {} + validated_root = {} errors = [] for key, value in data.items(): @@ -74,16 +104,22 @@ async def update_config_endpoint(request): continue if key in _SETTINGS_FIELDS: - settings[key] = value + validated_settings[key] = value else: - config[key] = value + validated_root[key] = value if errors: - return web.json_response({"error": errors}, status=400) + return await handle_api_error(request, errors, 400) - if save_config(config): - return web.json_response({"status": "success", "config": config}) - return await handle_api_error(request, "Failed to save config") + 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}") @@ -93,7 +129,7 @@ async def queue_status_endpoint(request): job_id = request.match_info['job_id'] # Import to ensure initialization - from ..nodes.distributed_upscale import ensure_tile_jobs_initialized + from ..upscale.job_store import ensure_tile_jobs_initialized prompt_server = ensure_tile_jobs_initialized() async with prompt_server.distributed_tile_jobs_lock: @@ -112,72 +148,39 @@ async def update_worker_endpoint(request): if worker_id is None: return await handle_api_error(request, "Missing worker_id", 400) - - config = load_config() - worker_found = False - - for worker in config.get("workers", []): - if worker["id"] == worker_id: - # Update all provided fields - if "enabled" in data: - worker["enabled"] = data["enabled"] - if "name" in data: - worker["name"] = data["name"] - if "port" in data: - worker["port"] = data["port"] - - # Handle host field - remove it if None - if "host" in data: - if data["host"] is None: - worker.pop("host", None) - else: - worker["host"] = normalize_host(data["host"]) - - # Handle cuda_device field - remove it if None - if "cuda_device" in data: - if data["cuda_device"] is None: - worker.pop("cuda_device", None) - else: - worker["cuda_device"] = data["cuda_device"] - - # Handle extra_args field - remove it if None - if "extra_args" in data: - if data["extra_args"] is None: - worker.pop("extra_args", None) - else: - worker["extra_args"] = data["extra_args"] - - # Handle type field - if "type" in data: - worker["type"] = data["type"] - - 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") - } - if "workers" not in config: - config["workers"] = [] - config["workers"].append(new_worker) - worker_found = True - else: - return await handle_api_error(request, f"Worker {worker_id} not found and missing required fields for creation", 404) - - if save_config(config): + + 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"}) - else: - return await handle_api_error(request, "Failed to save config") except Exception as e: return await handle_api_error(request, e, 400) @@ -190,29 +193,26 @@ async def delete_worker_endpoint(request): if worker_id is None: return await handle_api_error(request, "Missing worker_id", 400) - config = load_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) - - if save_config(config): + 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" }) - else: - return await handle_api_error(request, "Failed to save config") except Exception as e: return await handle_api_error(request, e, 400) @@ -229,16 +229,13 @@ async def update_setting_endpoint(request): if key not in _SETTINGS_FIELDS: return await handle_api_error(request, f"Unknown setting: {key}", 400) - config = load_config() - if 'settings' not in config: - config['settings'] = {} - - config['settings'][key] = value + async with config_transaction() as config: + if 'settings' not in config: + config['settings'] = {} + + config['settings'][key] = value - if save_config(config): return web.json_response({"status": "success", "message": f"Setting '{key}' updated."}) - else: - return await handle_api_error(request, "Failed to save config") except Exception as e: return await handle_api_error(request, e, 400) @@ -248,28 +245,11 @@ async def update_master_endpoint(request): try: data = await request.json() - config = load_config() - if 'master' not in config: - config['master'] = {} - - # Update all provided fields - if "name" in data: - config['master']['name'] = data['name'] - if "host" in data: - if data["host"] is None: - config['master'].pop('host', None) - else: - config['master']['host'] = normalize_host(data["host"]) - if "port" in data: - config['master']['port'] = data['port'] - if "cuda_device" in data: - config['master']['cuda_device'] = data['cuda_device'] - if "extra_args" in data: - config['master']['extra_args'] = data['extra_args'] - - if save_config(config): + 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."}) - else: - return await handle_api_error(request, "Failed to save config") except Exception as e: return await handle_api_error(request, e, 400) diff --git a/api/job_routes.py b/api/job_routes.py index ad93e3c..4c149b2 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -8,7 +8,6 @@ import time from aiohttp import web import server -import server as _server import torch from PIL import Image @@ -17,10 +16,10 @@ from ..utils.image import pil_to_tensor, ensure_contiguous from ..utils.network import handle_api_error from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY from ..utils.async_helpers import queue_prompt_payload -from .queue_orchestration import orchestrate_distributed_execution +from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution from .queue_request import parse_queue_request_payload -prompt_server = _server.PromptServer.instance +prompt_server = server.PromptServer.instance # Canonical worker result envelope accepted by POST /distributed/job_complete: # { "job_id": str, "worker_id": str, "batch_idx": int, "image": , "is_last": bool } @@ -136,6 +135,7 @@ async def prepare_job_endpoint(request): if not multi_job_id: return await handle_api_error(request, "Missing multi_job_id", 400) + ensure_distributed_state() async with prompt_server.distributed_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_jobs: prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() @@ -285,7 +285,7 @@ async def job_complete_endpoint(request): if not isinstance(is_last, bool): errors.append("is_last: expected boolean") if errors: - return web.json_response({"error": errors}, status=400) + return await handle_api_error(request, errors, 400) tensor = _decode_canonical_png_tensor(image_payload) multi_job_id = job_id.strip() diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index fd54a5a..def77f3 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -6,18 +6,8 @@ import aiohttp from ...utils.logging import debug_log, log from ...utils.network import build_worker_url, get_client_session, probe_worker - - -def _trace_prefix(trace_execution_id): - return f"[Distributed][exec:{trace_execution_id}]" - - -def _trace_debug(trace_execution_id, message): - debug_log(f"{_trace_prefix(trace_execution_id)} {message}") - - -def _trace_info(trace_execution_id, message): - log(f"{_trace_prefix(trace_execution_id)} {message}") +from ...utils.trace_logger import trace_debug, trace_info +from ..schemas import parse_positive_int async def worker_is_active(worker): @@ -45,14 +35,6 @@ async def worker_ws_is_active(worker): return False -def _coerce_positive_int(value, default): - try: - parsed = int(value) - except (TypeError, ValueError): - return max(1, int(default)) - return max(1, parsed) - - async def _probe_worker_active(worker, use_websocket, semaphore): async with semaphore: is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker)) @@ -120,7 +102,7 @@ async def dispatch_worker_prompt( except Exception as exc: worker_id = worker.get("id") if trace_execution_id: - _trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}") + trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}") else: log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}") raise @@ -142,11 +124,11 @@ async def select_active_workers( probe_concurrency=8, ): """Probe workers and return (active_workers, updated_delegate_master).""" - probe_limit = _coerce_positive_int(probe_concurrency, 8) + probe_limit = parse_positive_int(probe_concurrency, 8) probe_semaphore = asyncio.Semaphore(probe_limit) if trace_execution_id and workers: - _trace_debug( + trace_debug( trace_execution_id, f"Probing {len(workers)} workers with probe_concurrency={probe_limit}", ) @@ -164,19 +146,19 @@ async def select_active_workers( active_workers.append(worker) else: if trace_execution_id: - _trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.") + trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.") else: log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") if trace_execution_id and workers: - _trace_debug( + trace_debug( trace_execution_id, f"Worker probe complete: active={len(active_workers)}/{len(workers)}", ) if not active_workers and delegate_master: if trace_execution_id: - _trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.") + trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.") else: debug_log("All workers offline while delegate-only requested; enabling master participation.") delegate_master = False diff --git a/api/orchestration/media_sync.py b/api/orchestration/media_sync.py index ffaadfb..936675c 100644 --- a/api/orchestration/media_sync.py +++ b/api/orchestration/media_sync.py @@ -8,6 +8,7 @@ import aiohttp from ...utils.logging import debug_log, log from ...utils.network import build_worker_url, get_client_session +from ...utils.trace_logger import trace_debug, trace_info LIKELY_FILENAME_RE = re.compile( @@ -20,19 +21,6 @@ IMAGE_OR_VIDEO_RE = re.compile( re.IGNORECASE, ) - -def _trace_prefix(trace_execution_id): - return f"[Distributed][exec:{trace_execution_id}]" - - -def _trace_debug(trace_execution_id, message): - debug_log(f"{_trace_prefix(trace_execution_id)} {message}") - - -def _trace_info(trace_execution_id, message): - log(f"{_trace_prefix(trace_execution_id)} {message}") - - def convert_paths_for_platform(obj, target_separator): """Recursively normalize likely file paths for the worker platform separator.""" if target_separator not in ("/", "\\"): @@ -114,7 +102,7 @@ async def fetch_worker_path_separator(worker, trace_execution_id=None): return separator if separator in ("/", "\\") else None except Exception as exc: if trace_execution_id: - _trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}") + trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}") else: debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}") return None @@ -138,7 +126,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_ return False except Exception as exc: if trace_execution_id: - _trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") + trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") else: debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") @@ -180,13 +168,13 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None): except FileNotFoundError: missing += 1 if trace_execution_id: - _trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.") + trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.") else: log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.") continue except Exception as exc: if trace_execution_id: - _trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}") + trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}") else: log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}") continue @@ -206,7 +194,7 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None): skipped += 1 except Exception as exc: if trace_execution_id: - _trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") + trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") else: log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") @@ -215,6 +203,6 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None): f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}" ) if trace_execution_id: - _trace_debug(trace_execution_id, summary) + trace_debug(trace_execution_id, summary) else: debug_log(f"[Distributed] {summary}") diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 16ccd16..00a964e 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -14,6 +14,8 @@ from ..utils.constants import ( ) from ..utils.logging import debug_log, log from ..utils.network import build_master_url +from ..utils.trace_logger import trace_debug +from .schemas import parse_positive_float, parse_positive_int from .orchestration.dispatch import dispatch_worker_prompt, select_active_workers from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media from .orchestration.prompt_transform import ( @@ -33,20 +35,17 @@ def _generate_execution_trace_id(): return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" -def _trace_prefix(trace_execution_id): - return f"[Distributed][exec:{trace_execution_id}]" - - -def _trace_debug(trace_execution_id, message): - debug_log(f"{_trace_prefix(trace_execution_id)} {message}") - - -def ensure_distributed_state(): +def ensure_distributed_state(server_instance=None): """Ensure prompt_server has the state used by distributed queue orchestration.""" - if not hasattr(prompt_server, "distributed_pending_jobs"): - prompt_server.distributed_pending_jobs = {} - if not hasattr(prompt_server, "distributed_jobs_lock"): - prompt_server.distributed_jobs_lock = asyncio.Lock() + ps = server_instance or prompt_server + if not hasattr(ps, "distributed_pending_jobs"): + ps.distributed_pending_jobs = {} + if not hasattr(ps, "distributed_jobs_lock"): + ps.distributed_jobs_lock = asyncio.Lock() + + +# Initialize top-level distributed queue state at module import time. +ensure_distributed_state() async def _ensure_distributed_queue(job_id): @@ -90,40 +89,22 @@ def _resolve_enabled_workers(config, requested_ids=None): return workers -def _parse_positive_int(value, default): - """Parse positive integer setting values with fallback.""" - try: - parsed = int(value) - except (TypeError, ValueError): - return default - return max(1, parsed) - - -def _parse_positive_float(value, default): - """Parse positive float setting values with fallback.""" - try: - parsed = float(value) - except (TypeError, ValueError): - return default - return max(1.0, parsed) - - def _resolve_orchestration_limits(config): """Resolve bounded concurrency/timeouts for worker preparation pipeline.""" settings = (config or {}).get("settings", {}) or {} - worker_probe_concurrency = _parse_positive_int( + worker_probe_concurrency = parse_positive_int( settings.get("worker_probe_concurrency"), ORCHESTRATION_WORKER_PROBE_CONCURRENCY, ) - worker_prep_concurrency = _parse_positive_int( + worker_prep_concurrency = parse_positive_int( settings.get("worker_prep_concurrency"), ORCHESTRATION_WORKER_PREP_CONCURRENCY, ) - media_sync_concurrency = _parse_positive_int( + media_sync_concurrency = parse_positive_int( settings.get("media_sync_concurrency"), ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, ) - media_sync_timeout_seconds = _parse_positive_float( + media_sync_timeout_seconds = parse_positive_float( settings.get("media_sync_timeout_seconds"), ORCHESTRATION_MEDIA_SYNC_TIMEOUT, ) @@ -176,7 +157,7 @@ async def _prepare_worker_payload( timeout=media_sync_timeout_seconds, ) except asyncio.TimeoutError: - _trace_debug( + trace_debug( trace_execution_id, ( f"Media sync timed out after {media_sync_timeout_seconds:.1f}s " @@ -214,7 +195,7 @@ async def orchestrate_distributed_execution( requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None workers = _resolve_enabled_workers(config, requested_ids) prompt_index = PromptIndex(prompt_obj) - _trace_debug( + trace_debug( execution_trace_id, ( f"Orchestration start: requested_workers={len(workers)}, " @@ -232,7 +213,7 @@ async def orchestrate_distributed_execution( delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False)) if not workers and delegate_master: - _trace_debug( + trace_debug( execution_trace_id, "Delegate-only requested but no workers are enabled. Falling back to master execution.", ) @@ -252,7 +233,7 @@ async def orchestrate_distributed_execution( job_id_map = generate_job_id_map(prompt_index, discovery_prefix) if not job_id_map: - _trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") + trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id) return prompt_id, 0 @@ -286,7 +267,7 @@ async def orchestrate_distributed_execution( master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids) if active_workers: - _trace_debug( + trace_debug( execution_trace_id, "Active distributed workers: " + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers), @@ -329,7 +310,7 @@ async def orchestrate_distributed_execution( ) prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id) - _trace_debug( + trace_debug( execution_trace_id, f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}", ) diff --git a/api/schemas.py b/api/schemas.py index 506a1a7..feaa3ec 100644 --- a/api/schemas.py +++ b/api/schemas.py @@ -34,3 +34,21 @@ def validate_positive_int(value, field_name: str) -> str | None: if parsed <= 0: return f"Field '{field_name}' must be a positive integer." return None + + +def parse_positive_int(value, default: int) -> int: + """Parse value as positive int, returning default on failure.""" + try: + parsed = int(value) + except (TypeError, ValueError): + return max(1, int(default)) + return max(1, parsed) + + +def parse_positive_float(value, default: float) -> float: + """Parse value as positive float, returning default on failure.""" + try: + parsed = float(value) + except (TypeError, ValueError): + return max(0.0, float(default)) + return max(0.0, parsed) diff --git a/api/usdu_routes.py b/api/usdu_routes.py new file mode 100644 index 0000000..2735ad4 --- /dev/null +++ b/api/usdu_routes.py @@ -0,0 +1,228 @@ +import asyncio +import io +import time + +from aiohttp import web +from PIL import Image +import server + +from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState +from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized +from ..upscale.payload_parsers import _parse_tiles_from_form +from ..utils.logging import debug_log +from ..utils.network import handle_api_error + + +@server.PromptServer.instance.routes.post("/distributed/heartbeat") +async def heartbeat_endpoint(request): + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if isinstance(job_data, BaseJobState): + job_data.worker_status[worker_id] = time.time() + debug_log(f"Heartbeat from worker {worker_id}") + return web.json_response({"status": "success"}) + return await handle_api_error(request, "Worker status tracking not available", 400) + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/submit_tiles") +async def submit_tiles_endpoint(request): + """Endpoint for workers to submit processed tiles in static mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + batch_size = int(data.get('batch_size', 0)) + + # Handle completion signal + if batch_size == 0 and is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, TileJobState): + return await handle_api_error(request, "Job not configured for tile submissions", 400) + await job_data.queue.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [], + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + + try: + tiles = _parse_tiles_from_form(data) + except ValueError as e: + return await handle_api_error(request, str(e), 400) + + # Submit tiles to queue + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, TileJobState): + return await handle_api_error(request, "Job not configured for tile submissions", 400) + + q = job_data.queue + if batch_size > 0 or len(tiles) > 0: + await q.put({ + 'worker_id': worker_id, + 'tiles': tiles, + 'is_last': is_last, + }) + debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})") + else: + await q.put({ + 'worker_id': worker_id, + 'is_last': True, + 'tiles': [], + }) + + return web.json_response({"status": "success"}) + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/submit_image") +async def submit_image_endpoint(request): + """Endpoint for workers to submit processed images in dynamic mode.""" + try: + content_length = request.headers.get('content-length') + if content_length and int(content_length) > MAX_PAYLOAD_SIZE: + return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) + + data = await request.post() + multi_job_id = data.get('multi_job_id') + worker_id = data.get('worker_id') + is_last = data.get('is_last', 'False').lower() == 'true' + + if multi_job_id is None or worker_id is None: + return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + + # Handle image submission + if 'full_image' in data and 'image_idx' in data: + image_idx = int(data.get('image_idx')) + img_data = data['full_image'].file.read() + img = Image.open(io.BytesIO(img_data)).convert("RGB") + + debug_log(f"Received full image {image_idx} from worker {worker_id}") + + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, ImageJobState): + return await handle_api_error(request, "Job not configured for image submissions", 400) + await job_data.queue.put({ + 'worker_id': worker_id, + 'image_idx': image_idx, + 'image': img, + 'is_last': is_last, + }) + return web.json_response({"status": "success"}) + + # Handle completion signal (no image data) + elif is_last: + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, ImageJobState): + return await handle_api_error(request, "Job not configured for image submissions", 400) + await job_data.queue.put({ + 'worker_id': worker_id, + 'is_last': True, + }) + debug_log(f"Received completion signal from worker {worker_id}") + return web.json_response({"status": "success"}) + else: + return await handle_api_error(request, "Missing image data or invalid request", 400) + + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.post("/distributed/request_image") +async def request_image_endpoint(request): + """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode).""" + try: + data = await request.json() + worker_id = data.get('worker_id') + multi_job_id = data.get('multi_job_id') + + if not worker_id or not multi_job_id: + return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) + + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] + if not isinstance(job_data, BaseJobState): + return await handle_api_error(request, "Invalid job data structure", 500) + + mode = job_data.mode + if isinstance(job_data, ImageJobState): + pending_queue = job_data.pending_images + elif isinstance(job_data, TileJobState): + pending_queue = job_data.pending_tasks + else: + return await handle_api_error(request, "Invalid job configuration", 400) + + try: + task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) + job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx) + job_data.worker_status[worker_id] = time.time() + remaining = pending_queue.qsize() + + if mode == 'dynamic': + debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}") + return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining}) + debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}") + return web.json_response({ + "tile_idx": task_idx, + "estimated_remaining": remaining, + "batched_static": job_data.batched_static, + }) + except asyncio.TimeoutError: + if mode == 'dynamic': + return web.json_response({"image_idx": None}) + return web.json_response({"tile_idx": None}) + return await handle_api_error(request, "Job not found", 404) + except Exception as e: + return await handle_api_error(request, e, 500) + + +@server.PromptServer.instance.routes.get("/distributed/job_status") +async def job_status_endpoint(request): + """Endpoint to check if a job is ready.""" + multi_job_id = request.query.get('multi_job_id') + if not multi_job_id: + return web.json_response({"ready": False}) + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None) + return web.json_response({"ready": ready}) diff --git a/api/worker_routes.py b/api/worker_routes.py index fe5d42e..905c8ed 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -18,12 +18,9 @@ from ..utils.constants import CHUNK_SIZE from ..workers import get_worker_manager from .schemas import require_fields, validate_worker_id from ..workers.detection import ( - is_local_worker, - is_same_physical_host, get_machine_id, is_docker_environment, is_runpod_environment, - get_comms_channel, ) from ..utils.async_helpers import queue_prompt_payload @@ -317,13 +314,7 @@ async def get_network_info_endpoint(request): "message": "Auto-detected network configuration" }) except Exception as e: - return web.json_response({ - "status": "error", - "message": str(e), - "hostname": "unknown", - "all_ips": [], - "recommended_ip": None - }) + return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/system_info") async def get_system_info_endpoint(request): @@ -347,10 +338,7 @@ async def get_system_info_endpoint(request): "runpod_pod_id": os.environ.get('RUNPOD_POD_ID') }) except Exception as e: - return web.json_response({ - "status": "error", - "message": str(e) - }, status=500) + return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.post("/distributed/launch_worker") async def launch_worker_endpoint(request): @@ -389,12 +377,7 @@ async def launch_worker_endpoint(request): is_running = wm._is_process_running(proc_info['pid']) if is_running: - return web.json_response({ - "status": "error", - "message": "Worker already running (managed by UI)", - "pid": proc_info['pid'], - "log_file": proc_info.get('log_file') - }, status=409) + return await handle_api_error(request, "Worker already running (managed by UI)", 409) else: # Process is dead, remove it del wm.processes[worker_id_str] @@ -438,8 +421,11 @@ async def stop_worker_endpoint(request): if success: return web.json_response({"status": "success", "message": message}) else: - return web.json_response({"status": "error", "message": message}, - status=404 if "not managed" in message else 409) + return await handle_api_error( + request, + message, + 404 if "not managed" in message else 409, + ) except Exception as e: return await handle_api_error(request, e, 400) diff --git a/distributed.py b/distributed.py index 222b777..9c02da5 100644 --- a/distributed.py +++ b/distributed.py @@ -2,7 +2,6 @@ ComfyUI-Distributed: thin entry point. All implementation lives in workers/, nodes/, api/. """ -import asyncio import atexit import os @@ -28,6 +27,7 @@ from .nodes import ( any_type, ) from . import api # noqa: F401 - triggers all @routes.* registrations +from .api.queue_orchestration import ensure_distributed_state ensure_config_exists() @@ -40,9 +40,7 @@ atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup # Initialize distributed job state on prompt_server prompt_server = server.PromptServer.instance -if not hasattr(prompt_server, 'distributed_pending_jobs'): - prompt_server.distributed_pending_jobs = {} - prompt_server.distributed_jobs_lock = asyncio.Lock() +ensure_distributed_state(prompt_server) if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): prompt_server.distributed_pending_tile_jobs = {} diff --git a/nodes/collector.py b/nodes/collector.py index e075210..bca55ac 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -17,7 +17,6 @@ from ..utils.constants import HEARTBEAT_INTERVAL from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous from ..utils.network import build_worker_url, get_client_session, probe_worker from ..utils.async_helpers import run_async_in_server_loop -from ..workers.detection import is_local_worker, get_comms_channel prompt_server = _server.PromptServer.instance diff --git a/nodes/distributed_upscale.py b/nodes/distributed_upscale.py index 0ce0ff7..2e92e44 100644 --- a/nodes/distributed_upscale.py +++ b/nodes/distributed_upscale.py @@ -6,7 +6,7 @@ import comfy.samplers from ..utils.logging import debug_log, log from ..utils.async_helpers import run_async_in_server_loop -from ..utils.usdu_managment import ensure_tile_jobs_initialized +from ..upscale.job_store import ensure_tile_jobs_initialized from ..upscale.tile_ops import TileOpsMixin from ..upscale.result_collector import ResultCollectorMixin diff --git a/upscale/conditioning.py b/upscale/conditioning.py new file mode 100644 index 0000000..3e776bc --- /dev/null +++ b/upscale/conditioning.py @@ -0,0 +1,32 @@ +import copy + + +def clone_control_chain(control, clone_hint=True): + """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" + if control is None: + return None + new_control = copy.copy(control) + if clone_hint and hasattr(control, 'cond_hint_original'): + hint = getattr(control, 'cond_hint_original', None) + new_control.cond_hint_original = hint.clone() if hint is not None else None + if hasattr(control, 'previous_controlnet'): + new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) + return new_control + + +def clone_conditioning(cond_list, clone_hints=True): + """Clone conditioning without duplicating ControlNet models.""" + new_cond = [] + for emb, cond_dict in cond_list: + new_emb = emb.clone() if emb is not None else None + new_dict = cond_dict.copy() + if 'control' in new_dict: + new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) + if 'mask' in new_dict and new_dict['mask'] is not None: + new_dict['mask'] = new_dict['mask'].clone() + if 'pooled_output' in new_dict and new_dict['pooled_output'] is not None: + new_dict['pooled_output'] = new_dict['pooled_output'].clone() + if 'area' in new_dict: + new_dict['area'] = new_dict['area'][:] + new_cond.append([new_emb, new_dict]) + return new_cond diff --git a/upscale/job_state.py b/upscale/job_state.py index 6cdf09e..51ed1f9 100644 --- a/upscale/job_state.py +++ b/upscale/job_state.py @@ -1,10 +1,8 @@ import asyncio from ..utils.logging import debug_log -from ..utils.usdu_managment import ( - _check_and_requeue_timed_out_workers as _requeue_usdu, - ensure_tile_jobs_initialized, -) +from .job_store import ensure_tile_jobs_initialized +from .job_timeout import _check_and_requeue_timed_out_workers as _requeue_usdu from .job_models import ImageJobState, TileJobState diff --git a/upscale/job_store.py b/upscale/job_store.py new file mode 100644 index 0000000..c027a33 --- /dev/null +++ b/upscale/job_store.py @@ -0,0 +1,180 @@ +import asyncio +import os +import time +from typing import List, Optional + +import server + +from ..utils.logging import debug_log +from .job_models import BaseJobState, ImageJobState, TileJobState + +# Configure maximum payload size (50MB default, configurable via environment variable) +MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) + + +def ensure_tile_jobs_initialized(): + """Ensure tile job storage is initialized on the server instance.""" + prompt_server = server.PromptServer.instance + if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): + debug_log("Initializing persistent tile job queue on server instance.") + prompt_server.distributed_pending_tile_jobs = {} + prompt_server.distributed_tile_jobs_lock = asyncio.Lock() + else: + invalid_job_ids = [ + job_id + for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() + if not isinstance(job_data, BaseJobState) + ] + for job_id in invalid_job_ids: + debug_log(f"Removing invalid job state for {job_id}") + del prompt_server.distributed_pending_tile_jobs[job_id] + return prompt_server + + +async def _init_job_queue( + multi_job_id, + mode, + batch_size=None, + num_tiles_per_image=None, + all_indices=None, + enabled_workers=None, + batched_static: bool = False, +): + """Unified initialization for job queues in static and dynamic modes.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + debug_log(f"Queue already exists for {multi_job_id}") + return + + if mode == 'dynamic': + job_data = ImageJobState(multi_job_id=multi_job_id) + elif mode == 'static': + job_data = TileJobState(multi_job_id=multi_job_id) + else: + raise ValueError(f"Unknown mode: {mode}") + + job_data.worker_status = {w: time.time() for w in enabled_workers or []} + job_data.assigned_to_workers = {w: [] for w in enabled_workers or []} + + if mode == 'dynamic': + job_data.batch_size = int(batch_size or 0) + pending_queue = job_data.pending_images + for i in (all_indices or range(int(batch_size or 0))): + await pending_queue.put(i) + debug_log(f"Initialized image queue with {batch_size} pending items") + elif mode == 'static': + job_data.num_tiles_per_image = int(num_tiles_per_image or 0) + job_data.batch_size = int(batch_size or 0) + job_data.batched_static = bool(batched_static) + # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1] + pending_queue = job_data.pending_tasks + if batched_static and num_tiles_per_image is not None: + for i in range(num_tiles_per_image): + await pending_queue.put(i) + else: + total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0) + for i in range(total_tiles): + await pending_queue.put(i) + + prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data + + +async def init_dynamic_job( + multi_job_id: str, + batch_size: int, + enabled_workers: List[str], + all_indices: Optional[List[int]] = None, +): + """Initialize queue for dynamic mode (per-image), with collector fields.""" + await _init_job_queue( + multi_job_id, + 'dynamic', + batch_size=batch_size, + all_indices=all_indices or list(range(batch_size)), + enabled_workers=enabled_workers, + ) + debug_log(f"Job {multi_job_id} initialized with {batch_size} images") + + +async def init_static_job_batched( + multi_job_id: str, + batch_size: int, + num_tiles_per_image: int, + enabled_workers: List[str], +): + """Initialize queue for static mode (batched-per-tile).""" + await _init_job_queue( + multi_job_id, + 'static', + batch_size=batch_size, + num_tiles_per_image=num_tiles_per_image, + enabled_workers=enabled_workers, + batched_static=True, + ) + + +async def _drain_results_queue(multi_job_id): + """Drain pending results from queue and update completed_tasks. Returns count drained.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not isinstance(job_data, BaseJobState): + return 0 + q = job_data.queue + completed_tasks = job_data.completed_tasks + + collected = 0 + while True: + try: + result = q.get_nowait() + except asyncio.QueueEmpty: + break + + worker_id = result['worker_id'] + is_last = result.get('is_last', False) + + if 'image_idx' in result and 'image' in result: + task_id = result['image_idx'] + if task_id not in completed_tasks: + completed_tasks[task_id] = result['image'] + collected += 1 + elif 'tiles' in result: + for tile_data in result['tiles']: + task_id = tile_data.get('global_idx', tile_data['tile_idx']) + if task_id not in completed_tasks: + completed_tasks[task_id] = tile_data + collected += 1 + if is_last: + if worker_id in job_data.worker_status: + del job_data.worker_status[worker_id] + + return collected + + +async def _get_completed_count(multi_job_id): + """Get count of completed tasks.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if isinstance(job_data, BaseJobState): + return len(job_data.completed_tasks) + return 0 + + +async def _mark_task_completed(multi_job_id, task_id, result): + """Mark a task as completed.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if isinstance(job_data, BaseJobState): + job_data.completed_tasks[task_id] = result + + +async def _cleanup_job(multi_job_id): + """Cleanup the job data.""" + prompt_server = ensure_tile_jobs_initialized() + async with prompt_server.distributed_tile_jobs_lock: + if multi_job_id in prompt_server.distributed_pending_tile_jobs: + del prompt_server.distributed_pending_tile_jobs[multi_job_id] + debug_log(f"Cleaned up job {multi_job_id}") diff --git a/upscale/job_timeout.py b/upscale/job_timeout.py new file mode 100644 index 0000000..37baa9c --- /dev/null +++ b/upscale/job_timeout.py @@ -0,0 +1,150 @@ +import time + +from ..utils.config import load_config +from ..utils.constants import HEARTBEAT_TIMEOUT +from ..utils.logging import debug_log, log +from ..utils.network import build_worker_url, probe_worker +from .job_models import BaseJobState +from .job_store import ensure_tile_jobs_initialized + + +def _find_worker_record(worker_id): + """Return worker config entry by id, or None when missing.""" + workers = load_config().get("workers", []) + return next((w for w in workers if str(w.get("id")) == str(worker_id)), None) + + +async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): + """Check timed out workers and requeue their tasks. Returns requeued count.""" + prompt_server = ensure_tile_jobs_initialized() + current_time = time.time() + + # Allow override via config setting 'worker_timeout_seconds' + cfg = load_config() + hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT)) + + # Snapshot timed-out workers and job details under lock. + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not isinstance(job_data, BaseJobState): + return 0 + + completed_tasks_snapshot = set(job_data.completed_tasks.keys()) + batched_static_snapshot = bool(job_data.batched_static) + num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1) + batch_size_snapshot = int(job_data.batch_size or 1) + + timed_out_workers = [] + for worker, last_heartbeat in list(job_data.worker_status.items()): + age = current_time - float(last_heartbeat) + debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s") + if age > hb_timeout: + timed_out_workers.append( + { + "worker_id": worker, + "last_heartbeat": float(last_heartbeat), + "assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])), + } + ) + + if not timed_out_workers: + return 0 + + # Probe outside lock to avoid lock contention on network latency. + workers_to_requeue = [] + workers_graced = [] + for worker_info in timed_out_workers: + worker = worker_info["worker_id"] + assigned = worker_info["assigned_tasks"] + age = current_time - worker_info["last_heartbeat"] + + incomplete_assigned = 0 + try: + if assigned: + if batched_static_snapshot: + for task_id in assigned: + for b in range(batch_size_snapshot): + gidx = b * num_tiles_per_image_snapshot + task_id + if gidx not in completed_tasks_snapshot: + incomplete_assigned += 1 + break + else: + for task_id in assigned: + if task_id not in completed_tasks_snapshot: + incomplete_assigned += 1 + debug_log( + f"Assigned diagnostics: total_assigned={len(assigned)} " + f"incomplete_assigned={incomplete_assigned}" + ) + except Exception as e: + debug_log(f"Assigned diagnostics failed for worker {worker}: {e}") + + busy = False + probe_queue = None + try: + worker_record = _find_worker_record(worker) + if worker_record: + worker_url = build_worker_url(worker_record) + debug_log(f"Probing worker {worker} at {worker_url}/prompt") + payload = await probe_worker(worker_url, timeout=2.0) + if payload is not None: + probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0)) + busy = probe_queue is not None and probe_queue > 0 + else: + debug_log(f"Probe skipped; worker {worker} not found in config") + except Exception as e: + debug_log(f"Probe failed for worker {worker}: {e}") + finally: + debug_log( + f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}" + ) + + if busy: + workers_graced.append(worker) + debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue") + continue + + log(f"Worker {worker} heartbeat timed out after {age:.1f}s") + workers_to_requeue.append((worker, assigned)) + + # Re-acquire lock and apply requeue/cleanup decisions. + async with prompt_server.distributed_tile_jobs_lock: + job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if not isinstance(job_data, BaseJobState): + return 0 + + # Refresh heartbeat for workers that we proved are still busy. + for worker in workers_graced: + if worker in job_data.worker_status: + job_data.worker_status[worker] = current_time + + requeued_count = 0 + completed_tasks = job_data.completed_tasks + batched_static = bool(job_data.batched_static) + num_tiles_per_image = int(job_data.num_tiles_per_image or 1) + batch_size = int(job_data.batch_size or 1) + for worker, assigned_snapshot in workers_to_requeue: + # Use current assignments if present, falling back to the snapshot. + assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot)) + for task_id in assigned_tasks: + # If batched_static, task_id is a tile_idx; consider it complete only if + # all corresponding global_idx entries are present in completed_tasks. + if batched_static: + all_done = True + for b in range(batch_size): + gidx = b * num_tiles_per_image + task_id + if gidx not in completed_tasks: + all_done = False + break + if not all_done: + await job_data.pending_tasks.put(task_id) + requeued_count += 1 + else: + if task_id not in completed_tasks: + await job_data.pending_tasks.put(task_id) + requeued_count += 1 + job_data.worker_status.pop(worker, None) + if worker in job_data.assigned_to_workers: + job_data.assigned_to_workers[worker] = [] + + return requeued_count diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py index a09ff36..d7d9a96 100644 --- a/upscale/modes/dynamic.py +++ b/upscale/modes/dynamic.py @@ -6,13 +6,19 @@ from ...utils.image import tensor_to_pil, pil_to_tensor from ...utils.async_helpers import run_async_in_server_loop from ...utils.config import get_worker_timeout_seconds from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT -from ...utils.usdu_managment import ( - ensure_tile_jobs_initialized, init_dynamic_job, - clone_conditioning, _send_heartbeat_to_master, -) +from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job class DynamicModeMixin: + """ + Dynamic (per-image queue) USDU mode behaviors for master and worker roles. + + Expected co-mixins on `self`: + - TileOpsMixin (`calculate_tiles`, `_slice_conditioning`, `_process_and_blend_tile`). + - JobStateMixin (image queue/task completion helpers). + - WorkerCommsMixin (`_request_image_from_master`, `_send_full_image_to_master`, `_send_heartbeat_to_master`). + """ + def process_master_dynamic(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, @@ -71,22 +77,7 @@ class DynamicModeMixin: image_seed = seed + image_idx * 1000 # Pre-slice conditioning once per image (not per tile) - positive_sliced = clone_conditioning(positive) - negative_sliced = clone_conditioning(negative) - for cond_list in [positive_sliced, negative_sliced]: - for i in range(len(cond_list)): - emb, cond_dict = cond_list[i] - if emb.shape[0] > 1: - cond_list[i][0] = emb[image_idx:image_idx+1] - if 'control' in cond_dict: - control = cond_dict['control'] - while control is not None: - hint = control.cond_hint_original - if hint.shape[0] > 1: - control.cond_hint_original = hint[image_idx:image_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx) for tile_idx, pos in enumerate(all_tiles): local_image = self._process_and_blend_tile( @@ -268,55 +259,40 @@ class DynamicModeMixin: image_seed = seed + image_idx * 1000 # Pre-slice conditioning once per image (not per tile) - positive_sliced = clone_conditioning(positive) - negative_sliced = clone_conditioning(negative) - for cond_list in [positive_sliced, negative_sliced]: - for i in range(len(cond_list)): - emb, cond_dict = cond_list[i] - if emb.shape[0] > 1: - cond_list[i][0] = emb[image_idx:image_idx+1] - if 'control' in cond_dict: - control = cond_dict['control'] - while control is not None: - hint = control.cond_hint_original - if hint.shape[0] > 1: - control.cond_hint_original = hint[image_idx:image_idx+1] - control = control.previous_controlnet - if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1: - cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1] + positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx) - for tile_idx, pos in enumerate(all_tiles): - local_image = self._process_and_blend_tile( - tile_idx, pos, single_tensor, local_image, - model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, - sampler_name, scheduler, denoise, tile_width, tile_height, - padding, mask_blur, width, height, force_uniform_tiles, - tiled_decode, batch_idx=image_idx - ) - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) + for tile_idx, pos in enumerate(all_tiles): + local_image = self._process_and_blend_tile( + tile_idx, pos, single_tensor, local_image, + model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg, + sampler_name, scheduler, denoise, tile_width, tile_height, + padding, mask_blur, width, height, force_uniform_tiles, + tiled_decode, batch_idx=image_idx + ) + run_async_in_server_loop( + self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) - # Send processed image back to master - try: - # Use the estimated remaining to determine if this is the last image - is_last = is_last_for_worker - run_async_in_server_loop( - self._send_full_image_to_master(local_image, image_idx, multi_job_id, - master_url, worker_id, is_last), - timeout=TILE_SEND_TIMEOUT - ) - # Send heartbeat after processing - run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), - timeout=5.0 - ) - if is_last: - break - except Exception as e: - log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}") - # Continue processing other images + # Send processed image back to master + try: + # Use the estimated remaining to determine if this is the last image + is_last = is_last_for_worker + run_async_in_server_loop( + self._send_full_image_to_master(local_image, image_idx, multi_job_id, + master_url, worker_id, is_last), + timeout=TILE_SEND_TIMEOUT + ) + # Send heartbeat after processing + run_async_in_server_loop( + self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + timeout=5.0 + ) + if is_last: + break + except Exception as e: + log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}") + # Continue processing other images # Send final is_last signal debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal") diff --git a/upscale/modes/static.py b/upscale/modes/static.py index 54960a4..2c61eed 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -5,17 +5,32 @@ from ...utils.logging import debug_log, log from ...utils.image import tensor_to_pil, pil_to_tensor from ...utils.async_helpers import run_async_in_server_loop from ...utils.config import get_worker_timeout_seconds -from ...utils.constants import HEARTBEAT_INTERVAL, MAX_BATCH, TILE_SEND_TIMEOUT, TILE_WAIT_TIMEOUT -from ...utils.usdu_managment import ( +from ...utils.constants import ( + HEARTBEAT_INTERVAL, + JOB_POLL_INTERVAL, + JOB_POLL_MAX_ATTEMPTS, + MAX_BATCH, + TILE_SEND_TIMEOUT, + TILE_WAIT_TIMEOUT, +) +from ..job_store import ( ensure_tile_jobs_initialized, init_static_job_batched, - _mark_task_completed, _cleanup_job, _drain_results_queue, - _send_heartbeat_to_master, _get_completed_count, + _mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count, ) from ..job_models import TileJobState class StaticModeMixin: - def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=20): + """ + Static (tile-queue) USDU mode behaviors for master and worker roles. + + Expected co-mixins on `self`: + - TileOpsMixin (`calculate_tiles`, tile extract/blend helpers). + - JobStateMixin (`_get_next_tile_index`, `_get_all_completed_tasks`, requeue checks). + - WorkerCommsMixin (`send_tiles_batch_to_master`, `_request_tile_from_master`, `_send_heartbeat_to_master`). + """ + + def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS): """Poll master for job readiness to avoid worker/master initialization race.""" for attempt in range(max_attempts): ready = run_async_in_server_loop( @@ -28,7 +43,7 @@ class StaticModeMixin: else: debug_log(f"Job {multi_job_id} ready after {attempt} attempts") return True - time.sleep(1.0) + time.sleep(JOB_POLL_INTERVAL) return False def _extract_and_process_tile( @@ -194,7 +209,7 @@ class StaticModeMixin: log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}") processed_count = 0 - max_poll_attempts = 20 + max_poll_attempts = JOB_POLL_MAX_ATTEMPTS if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts): log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") return (upscaled_image,) @@ -254,7 +269,7 @@ class StaticModeMixin: # Send heartbeat try: run_async_in_server_loop( - _send_heartbeat_to_master(multi_job_id, master_url, worker_id), + self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), timeout=5.0 ) except Exception as e: diff --git a/upscale/payload_parsers.py b/upscale/payload_parsers.py new file mode 100644 index 0000000..6dc64e4 --- /dev/null +++ b/upscale/payload_parsers.py @@ -0,0 +1,64 @@ +import io +import json + +from PIL import Image + + +def _parse_tiles_from_form(data): + """Parse tiles submitted via multipart/form-data into a list of tile dicts.""" + try: + padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0 + except Exception: + padding = 0 + + meta_raw = data.get('tiles_metadata') + if meta_raw is None: + raise ValueError("Missing tiles_metadata") + + try: + metadata = json.loads(meta_raw) + except Exception as e: + raise ValueError(f"Invalid tiles_metadata JSON: {e}") + + if not isinstance(metadata, list): + raise ValueError("tiles_metadata must be a list") + + tiles = [] + for i, meta in enumerate(metadata): + file_field = data.get(f'tile_{i}') + if file_field is None or not hasattr(file_field, 'file'): + raise ValueError(f"Missing tile data for index {i}") + + raw = file_field.file.read() + try: + img = Image.open(io.BytesIO(raw)).convert("RGB") + except Exception as e: + raise ValueError(f"Invalid image data for tile {i}: {e}") + + try: + tile_info = { + 'image': img, + 'tile_idx': int(meta.get('tile_idx', i)), + 'x': int(meta.get('x', 0)), + 'y': int(meta.get('y', 0)), + 'extracted_width': int(meta.get('extracted_width', img.width)), + 'extracted_height': int(meta.get('extracted_height', img.height)), + 'padding': int(padding), + } + except Exception as e: + raise ValueError(f"Invalid metadata values for tile {i}: {e}") + + if 'batch_idx' in meta: + try: + tile_info['batch_idx'] = int(meta['batch_idx']) + except Exception: + pass + if 'global_idx' in meta: + try: + tile_info['global_idx'] = int(meta['global_idx']) + except Exception: + pass + + tiles.append(tile_info) + + return tiles diff --git a/upscale/result_collector.py b/upscale/result_collector.py index fde2a5e..5ceb79c 100644 --- a/upscale/result_collector.py +++ b/upscale/result_collector.py @@ -1,17 +1,38 @@ import asyncio, time import comfy.model_management import server -from ..utils.constants import HEARTBEAT_INTERVAL +from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL from ..utils.logging import debug_log, log from ..utils.config import get_worker_timeout_seconds -from ..utils.usdu_managment import ( - ensure_tile_jobs_initialized, _mark_task_completed, - _check_and_requeue_timed_out_workers, -) +from .job_store import ensure_tile_jobs_initialized, _mark_task_completed +from .job_timeout import _check_and_requeue_timed_out_workers from .job_models import BaseJobState, ImageJobState, TileJobState class ResultCollectorMixin: + """ + Mixin for master-side result collection in USDU distributed jobs. + + Expected co-mixins/attributes: + - JobStateMixin methods for queue/task access. + - `self._check_and_requeue_timed_out_workers(...)` coroutine. + - `self._async_yield(...)` optional helper from WorkerCommsMixin. + """ + + def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]: + """Log timeout elapsed seconds for each tracked worker and return worker ids.""" + if not isinstance(job_data, BaseJobState): + return [] + + worker_status = dict(job_data.worker_status) + for worker_id, last_seen in worker_status.items(): + elapsed = max(0.0, current_time - float(last_seen)) + log( + "UltimateSDUpscale Master - Heartbeat timeout: " + f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s" + ) + return list(worker_status.keys()) + async def _async_collect_results(self, multi_job_id, num_workers, mode='static', remaining_to_collect=None, batch_size=None): """Unified async helper to collect results from workers (tiles or images).""" @@ -59,10 +80,16 @@ class ResultCollectorMixin: # For dynamic mode with remaining_to_collect, check if we've collected enough if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect: break - + + job_data_snapshot = None + async with prompt_server.distributed_tile_jobs_lock: + current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) + if isinstance(current_job_data, BaseJobState): + job_data_snapshot = current_job_data + try: # Shorter poll for dynamic mode, but never exceed the configured timeout - wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout) + wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout) result = await asyncio.wait_for(q.get(), timeout=wait_timeout) worker_id = result['worker_id'] is_last = result.get('is_last', False) @@ -115,6 +142,7 @@ class ResultCollectorMixin: except asyncio.TimeoutError: current_time = time.time() + waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id) if mode == 'dynamic': # Check for worker timeouts periodically if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL: @@ -126,16 +154,6 @@ class ResultCollectorMixin: # Check if we've been waiting too long overall if current_time - wait_started_at > timeout: - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {} - waiting_workers = list(worker_status.keys()) - for worker_id, last_seen in worker_status.items(): - elapsed_for_worker = max(0.0, current_time - float(last_seen)) - log( - "UltimateSDUpscale Master - Heartbeat timeout: " - f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s" - ) elapsed = current_time - wait_started_at log( "UltimateSDUpscale Master - Heartbeat timeout while waiting for images; " @@ -143,16 +161,6 @@ class ResultCollectorMixin: ) break else: - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {} - waiting_workers = list(worker_status.keys()) - for worker_id, last_seen in worker_status.items(): - elapsed_for_worker = max(0.0, current_time - float(last_seen)) - log( - "UltimateSDUpscale Master - Heartbeat timeout: " - f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s" - ) elapsed = current_time - wait_started_at log( f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; " diff --git a/upscale/tile_ops.py b/upscale/tile_ops.py index 50acf86..3a51f6b 100644 --- a/upscale/tile_ops.py +++ b/upscale/tile_ops.py @@ -5,7 +5,7 @@ import comfy.samplers, comfy.model_management from ..utils.logging import debug_log, log from ..utils.image import tensor_to_pil, pil_to_tensor from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop -from ..utils.usdu_managment import clone_conditioning +from .conditioning import clone_conditioning class TileOpsMixin: diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py index 7c8ad5c..ddd0004 100644 --- a/upscale/worker_comms.py +++ b/upscale/worker_comms.py @@ -9,6 +9,10 @@ from ..utils.image import tensor_to_pil class WorkerCommsMixin: + async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id): + """Proxy heartbeat helper used by worker processing mixins.""" + await _send_heartbeat_to_master(multi_job_id, master_url, worker_id) + async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False): """Send all processed tiles to master, chunked if large.""" diff --git a/utils/cloudflare.py b/utils/cloudflare.py deleted file mode 100644 index df87b3c..0000000 --- a/utils/cloudflare.py +++ /dev/null @@ -1,388 +0,0 @@ -""" -Lightweight Cloudflare Tunnel manager. - -Responsibilities: -- Locate or download the cloudflared binary for the current platform. -- Start/stop a quick tunnel pointing to the local ComfyUI server. -- Track status and persist minimal state in the shared config. -""" -import asyncio -import os -import platform -import re -import shutil -import signal -import stat -import subprocess -import threading -import time -from urllib import request, error as urlerror - -from .config import load_config, save_config -from .logging import debug_log, log -from .network import get_server_port, normalize_host -from .process import is_process_alive, terminate_process - -# Regex to capture the generated public URL from cloudflared output -PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE) - -# Default timeout (seconds) to wait for a tunnel URL before giving up -TUNNEL_START_TIMEOUT = 25 - -class CloudflareTunnelManager: - def __init__(self): - self.process = None - self.pid = None - self.public_url = None - self.last_error = None - self.log_file = None - self.status = "stopped" - self.previous_master_host = None - - self._lock = asyncio.Lock() - self._url_event = None - self._reader_thread = None - self._loop = None - self._recent_logs = [] - - self.binary_path = None - self._restore_state() - - # --- Binary discovery & download --- - @property - def base_dir(self): - return os.path.dirname(os.path.dirname(__file__)) - - @property - def bin_dir(self): - return os.path.join(self.base_dir, "bin") - - def _get_asset_name(self): - system = platform.system().lower() - machine = platform.machine().lower() - - if system == "windows": - if "arm" in machine: - return "cloudflared-windows-arm64.exe" - return "cloudflared-windows-amd64.exe" - if system == "darwin": - if machine in ("arm64", "aarch64"): - return "cloudflared-darwin-arm64" - return "cloudflared-darwin-amd64" - if system == "linux": - if machine in ("arm64", "aarch64"): - return "cloudflared-linux-arm64" - return "cloudflared-linux-amd64" - - raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}") - - def _download_cloudflared(self): - asset = self._get_asset_name() - url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}" - os.makedirs(self.bin_dir, exist_ok=True) - target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared") - - debug_log(f"Downloading cloudflared from {url}") - try: - with request.urlopen(url, timeout=30) as resp: - with open(target_path, "wb") as f: - shutil.copyfileobj(resp, f) - except urlerror.URLError as exc: - raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc - - # Make executable - st = os.stat(target_path) - os.chmod(target_path, st.st_mode | stat.S_IEXEC) - debug_log(f"Downloaded cloudflared to {target_path}") - return target_path - - def ensure_binary(self): - # Allow explicit override - env_path = os.environ.get("CLOUDFLARED_PATH") - if env_path and os.path.exists(env_path): - self.binary_path = env_path - return env_path - - # Check cached path - cached = self.binary_path - if cached and os.path.exists(cached): - return cached - - # Local bin directory - candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared") - if os.path.exists(candidate): - self.binary_path = candidate - return candidate - - # System PATH - path_binary = shutil.which("cloudflared") - if path_binary: - self.binary_path = path_binary - return path_binary - - # Download on demand - self.binary_path = self._download_cloudflared() - return self.binary_path - - # --- State helpers --- - def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None): - cfg = load_config() - tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {} - - if status is not None: - tunnel_cfg["status"] = status - if public_url is not None: - tunnel_cfg["public_url"] = public_url - if pid is not None: - tunnel_cfg["pid"] = pid - if log_file is not None: - tunnel_cfg["log_file"] = log_file - if previous_host is not None: - tunnel_cfg["previous_master_host"] = previous_host - if master_host is not None: - cfg.setdefault("master", {})["host"] = master_host - - cfg["tunnel"] = tunnel_cfg - save_config(cfg) - - def _restore_state(self): - cfg = load_config() - tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {} - - self.public_url = tunnel_cfg.get("public_url") or None - self.previous_master_host = tunnel_cfg.get("previous_master_host") - self.log_file = tunnel_cfg.get("log_file") - pid = tunnel_cfg.get("pid") - - if pid and is_process_alive(pid): - self.pid = pid - self.status = tunnel_cfg.get("status", "running") - debug_log(f"Detected existing cloudflared process (pid={pid})") - else: - # Clear stale info - self._persist_state(status="stopped", public_url="", pid=None, log_file=None) - self.status = "stopped" - self.pid = None - - # --- Process output handling --- - def _append_log(self, line): - if self.log_file: - try: - with open(self.log_file, "a", encoding="utf-8", errors="replace") as f: - f.write(line + "\n") - except Exception as exc: # pragma: no cover - best effort logging - debug_log(f"Failed to write tunnel log: {exc}") - - self._recent_logs.append(line) - # Keep last ~200 lines to avoid unbounded growth - if len(self._recent_logs) > 200: - self._recent_logs = self._recent_logs[-200:] - - def _reader(self): - assert self.process is not None - loop = self._loop - for raw_line in iter(self.process.stdout.readline, ""): - line = raw_line.strip() - if not line: - continue - - self._append_log(line) - match = PUBLIC_URL_PATTERN.search(line) - if match and not self.public_url: - self.public_url = match.group(1).rstrip("/") - self.status = "running" - if self._url_event and loop: - loop.call_soon_threadsafe(self._url_event.set) - - # Capture obvious error strings - if "error" in line.lower() and not self.last_error: - self.last_error = line - - # Process exited - if self.status == "starting" and not self.public_url: - self.status = "error" - if not self.last_error: - self.last_error = "Cloudflare tunnel exited before becoming ready" - if self._url_event and loop: - loop.call_soon_threadsafe(self._url_event.set) - elif self.status == "running": - # Treat unexpected exit as stopped so UI can retry - self.status = "stopped" - elif self.status == "stopped": - # Normal stop path, nothing to do - pass - - # --- Public API --- - async def start_tunnel(self): - async with self._lock: - if self.process and self.process.poll() is None: - return { - "status": self.status, - "public_url": self.public_url, - "pid": self.process.pid, - "log_file": self.log_file, - } - - # If we had a stale pid, ensure it's not running - if self.pid and is_process_alive(self.pid): - debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one") - await self.stop_tunnel() - - binary = await asyncio.to_thread(self.ensure_binary) - port = get_server_port() - self.status = "starting" - self.last_error = None - self.public_url = None - self._recent_logs = [] - - # Remember current master host so we can restore it if the tunnel stops - config = load_config() - master_host = (config.get("master") or {}).get("host") or "" - tunnel_cfg = config.get("tunnel") or {} - if tunnel_cfg.get("previous_master_host"): - self.previous_master_host = tunnel_cfg.get("previous_master_host") - else: - self.previous_master_host = master_host - - os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True) - timestamp = time.strftime("%Y%m%d-%H%M%S") - self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log") - - cmd = [ - binary, - "tunnel", - "--no-autoupdate", - "--url", - f"http://127.0.0.1:{port}", - ] - - debug_log(f"Starting cloudflared: {' '.join(cmd)}") - try: - self.process = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, - ) - except FileNotFoundError: - self.status = "error" - raise RuntimeError("cloudflared binary not found") - except Exception as exc: - self.status = "error" - raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc - - self.pid = self.process.pid - self._persist_state( - status="starting", - pid=self.pid, - log_file=self.log_file, - previous_host=self.previous_master_host - ) - - # Kick off background reader - self._loop = asyncio.get_running_loop() - self._url_event = asyncio.Event() - self._reader_thread = threading.Thread(target=self._reader, daemon=True) - self._reader_thread.start() - - try: - await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT) - except asyncio.TimeoutError: - self.last_error = "Timed out waiting for Cloudflare to assign a URL" - await self.stop_tunnel() - raise RuntimeError(self.last_error) - - if self.status != "running" or not self.public_url: - error_msg = self.last_error or "Cloudflare tunnel failed to start" - await self.stop_tunnel() - raise RuntimeError(error_msg) - - debug_log(f"Cloudflare tunnel ready at {self.public_url}") - master_host = normalize_host(self.public_url) - self._persist_state( - status="running", - public_url=self.public_url, - pid=self.pid, - log_file=self.log_file, - previous_host=self.previous_master_host or "", - master_host=master_host - ) - return { - "status": self.status, - "public_url": self.public_url, - "pid": self.pid, - "log_file": self.log_file, - } - - async def stop_tunnel(self): - async with self._lock: - pid = self.process.pid if self.process else self.pid - if not pid: - self._persist_state(status="stopped", public_url=None, pid=None, log_file=None) - self.status = "stopped" - return {"status": "stopped"} - - debug_log(f"Stopping cloudflared (pid={pid})") - if self.process: - terminate_process(self.process, timeout=5) - else: - try: - os.kill(pid, signal.SIGTERM) - time.sleep(0.5) - except Exception as exc: # pragma: no cover - best effort - debug_log(f"Error stopping cloudflared pid {pid}: {exc}") - - config = load_config() - tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {} - self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host) - active_url = tunnel_cfg.get("public_url") - current_master_host = (config.get("master") or {}).get("host") - restore_host = None - if active_url: - active_host = normalize_host(active_url) - current_host = normalize_host(current_master_host) - if current_host == active_host: - restore_host = self.previous_master_host or "" - - self.status = "stopped" - self.public_url = None - self.pid = None - self.process = None - self.last_error = None - self._url_event = None - if self._reader_thread and self._reader_thread.is_alive(): - self._reader_thread.join(timeout=1) - self._reader_thread = None - self._persist_state( - status="stopped", - public_url="", - pid=None, - log_file=self.log_file, - previous_host=self.previous_master_host, - master_host=restore_host - ) - return {"status": "stopped"} - - def get_status(self): - alive = False - pid = self.process.pid if self.process else self.pid - if pid: - alive = is_process_alive(pid) - if not alive and self.status == "running": - self.status = "stopped" - - return { - "status": self.status, - "public_url": self.public_url, - "pid": pid, - "log_file": self.log_file, - "last_error": self.last_error, - "binary_path": self.binary_path or shutil.which("cloudflared"), - "recent_logs": self._recent_logs[-20:], - "previous_master_host": self.previous_master_host, - } - - -# Singleton tunnel manager -cloudflare_tunnel_manager = CloudflareTunnelManager() diff --git a/utils/cloudflare/__init__.py b/utils/cloudflare/__init__.py new file mode 100644 index 0000000..95a2199 --- /dev/null +++ b/utils/cloudflare/__init__.py @@ -0,0 +1,5 @@ +from .tunnel import CloudflareTunnelManager + +cloudflare_tunnel_manager = CloudflareTunnelManager() + +__all__ = ["CloudflareTunnelManager", "cloudflare_tunnel_manager"] diff --git a/utils/cloudflare/binary.py b/utils/cloudflare/binary.py new file mode 100644 index 0000000..4966345 --- /dev/null +++ b/utils/cloudflare/binary.py @@ -0,0 +1,83 @@ +"""Cloudflared binary discovery and download helpers.""" + +import os +import platform +import shutil +import stat +from urllib import error as urlerror +from urllib import request + +from ..logging import debug_log + + +def _get_project_root(): + return os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + + +def _get_cloudflared_dir(): + return os.path.join(_get_project_root(), "bin") + + +def _get_platform_binary_name(): + system = platform.system().lower() + machine = platform.machine().lower() + + if system == "windows": + if "arm" in machine: + return "cloudflared-windows-arm64.exe" + return "cloudflared-windows-amd64.exe" + if system == "darwin": + if machine in ("arm64", "aarch64"): + return "cloudflared-darwin-arm64" + return "cloudflared-darwin-amd64" + if system == "linux": + if machine in ("arm64", "aarch64"): + return "cloudflared-linux-arm64" + return "cloudflared-linux-amd64" + + raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}") + + +def _get_binary_path(bin_dir=None): + bin_dir = bin_dir or _get_cloudflared_dir() + binary_name = "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared" + return os.path.join(bin_dir, binary_name) + + +def _download_cloudflared(): + asset = _get_platform_binary_name() + url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}" + + bin_dir = _get_cloudflared_dir() + os.makedirs(bin_dir, exist_ok=True) + target_path = _get_binary_path(bin_dir) + + debug_log(f"Downloading cloudflared from {url}") + try: + with request.urlopen(url, timeout=30) as resp: + with open(target_path, "wb") as f: + shutil.copyfileobj(resp, f) + except urlerror.URLError as exc: + raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc + + st = os.stat(target_path) + os.chmod(target_path, st.st_mode | stat.S_IEXEC) + debug_log(f"Downloaded cloudflared to {target_path}") + return target_path + + +def ensure_binary() -> str: + """Return a usable cloudflared binary path, downloading if necessary.""" + env_path = os.environ.get("CLOUDFLARED_PATH") + if env_path and os.path.exists(env_path): + return env_path + + local_candidate = _get_binary_path() + if os.path.exists(local_candidate): + return local_candidate + + path_binary = shutil.which("cloudflared") + if path_binary: + return path_binary + + return _download_cloudflared() diff --git a/utils/cloudflare/process_reader.py b/utils/cloudflare/process_reader.py new file mode 100644 index 0000000..cf6de9f --- /dev/null +++ b/utils/cloudflare/process_reader.py @@ -0,0 +1,97 @@ +"""Background cloudflared process output reader.""" + +import asyncio +import re +import threading + +from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE +from ..logging import debug_log + +PUBLIC_URL_PATTERN = re.compile( + r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", + re.IGNORECASE, +) +class ProcessReader: + def __init__(self, log_file=None): + self._process = None + self._thread = None + self._loop = None + self._url_event = None + self._public_url = None + self._last_error = None + self._recent_logs = [] + self._log_file = log_file + + def set_log_file(self, log_file): + self._log_file = log_file + + def _append_log(self, line): + if self._log_file: + try: + with open(self._log_file, "a", encoding="utf-8", errors="replace") as f: + f.write(line + "\n") + except Exception as exc: # pragma: no cover + debug_log(f"Failed to write tunnel log: {exc}") + + self._recent_logs.append(line) + if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE: + self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:] + + def _reader(self): + process = self._process + if process is None: + return + + loop = self._loop + for raw_line in iter(process.stdout.readline, ""): + line = raw_line.strip() + if not line: + continue + + self._append_log(line) + match = PUBLIC_URL_PATTERN.search(line) + if match and not self._public_url: + self._public_url = match.group(1).rstrip("/") + if self._url_event and loop: + loop.call_soon_threadsafe(self._url_event.set) + + if "error" in line.lower() and not self._last_error: + self._last_error = line + + if self._url_event and loop: + if not self._last_error and not self._public_url: + self._last_error = "Cloudflare tunnel exited before becoming ready" + loop.call_soon_threadsafe(self._url_event.set) + + def start(self, process, loop): + self._process = process + self._loop = loop + self._url_event = asyncio.Event() + self._public_url = None + self._last_error = None + self._recent_logs = [] + self._thread = threading.Thread(target=self._reader, daemon=True) + self._thread.start() + + async def wait_for_url(self, timeout): + if not self._url_event: + return None + await asyncio.wait_for(self._url_event.wait(), timeout=timeout) + return self._public_url + + def stop(self): + if self._thread and self._thread.is_alive(): + self._thread.join(timeout=1) + self._thread = None + self._process = None + self._loop = None + self._url_event = None + + def get_url(self): + return self._public_url + + def get_last_error(self): + return self._last_error + + def get_recent_logs(self): + return list(self._recent_logs) diff --git a/utils/cloudflare/state.py b/utils/cloudflare/state.py new file mode 100644 index 0000000..4ec08b5 --- /dev/null +++ b/utils/cloudflare/state.py @@ -0,0 +1,81 @@ +"""Cloudflare tunnel state persistence helpers.""" + +from ..config import load_config, save_config +from ..network import normalize_host + + +def _get_tunnel_config(cfg): + tunnel_cfg = cfg.get("tunnel", {}) + if isinstance(tunnel_cfg, dict): + return tunnel_cfg + return {} + + +def load_tunnel_state(): + cfg = load_config() + tunnel_cfg = _get_tunnel_config(cfg) + master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {} + return { + "status": tunnel_cfg.get("status", "stopped"), + "public_url": tunnel_cfg.get("public_url") or None, + "pid": tunnel_cfg.get("pid"), + "log_file": tunnel_cfg.get("log_file"), + "previous_master_host": tunnel_cfg.get("previous_master_host"), + "master_host": master_cfg.get("host"), + } + + +def persist_tunnel_state( + status=None, + public_url=None, + pid=None, + log_file=None, + previous_host=None, + master_host=None, +): + cfg = load_config() + tunnel_cfg = _get_tunnel_config(cfg) + + if status is not None: + tunnel_cfg["status"] = status + if public_url is not None: + tunnel_cfg["public_url"] = public_url + if pid is not None: + tunnel_cfg["pid"] = pid + if log_file is not None: + tunnel_cfg["log_file"] = log_file + if previous_host is not None: + tunnel_cfg["previous_master_host"] = previous_host + if master_host is not None: + cfg.setdefault("master", {})["host"] = master_host + + cfg["tunnel"] = tunnel_cfg + save_config(cfg) + + +def clear_tunnel_state(log_file=None, previous_host=None, master_host=None): + persist_tunnel_state( + status="stopped", + public_url="", + pid=None, + log_file=log_file, + previous_host=previous_host, + master_host=master_host, + ) + + +def resolve_restore_master_host(previous_master_host): + """Determine whether master host should be restored after tunnel stop.""" + cfg = load_config() + tunnel_cfg = _get_tunnel_config(cfg) + active_url = tunnel_cfg.get("public_url") + current_master_host = (cfg.get("master") or {}).get("host") + + if not active_url: + return None + + active_host = normalize_host(active_url) + current_host = normalize_host(current_master_host) + if current_host == active_host: + return previous_master_host or "" + return None diff --git a/utils/cloudflare/tunnel.py b/utils/cloudflare/tunnel.py new file mode 100644 index 0000000..936819d --- /dev/null +++ b/utils/cloudflare/tunnel.py @@ -0,0 +1,207 @@ +"""Cloudflare tunnel lifecycle manager.""" + +import asyncio +import os +import shutil +import signal +import subprocess +import time + +from ..constants import TUNNEL_START_TIMEOUT +from ..logging import debug_log +from ..network import get_server_port, normalize_host +from ..process import is_process_alive, terminate_process +from .binary import ensure_binary +from .process_reader import ProcessReader +from .state import clear_tunnel_state, load_tunnel_state, persist_tunnel_state, resolve_restore_master_host + + +class CloudflareTunnelManager: + def __init__(self): + self.process = None + self.pid = None + self.public_url = None + self.last_error = None + self.log_file = None + self.status = "stopped" + self.previous_master_host = None + + self._lock = asyncio.Lock() + self._reader = ProcessReader() + self.binary_path = None + + self._restore_state() + + @property + def base_dir(self): + return os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + + def _restore_state(self): + state = load_tunnel_state() + + self.public_url = state.get("public_url") or None + self.previous_master_host = state.get("previous_master_host") + self.log_file = state.get("log_file") + pid = state.get("pid") + + if pid and is_process_alive(pid): + self.pid = pid + self.status = state.get("status") or "running" + debug_log(f"Detected existing cloudflared process (pid={pid})") + else: + clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host) + self.status = "stopped" + self.pid = None + + async def start_tunnel(self): + async with self._lock: + if self.process and self.process.poll() is None: + return { + "status": self.status, + "public_url": self.public_url, + "pid": self.process.pid, + "log_file": self.log_file, + } + + if self.pid and is_process_alive(self.pid): + debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one") + await self.stop_tunnel() + + binary = await asyncio.to_thread(ensure_binary) + self.binary_path = binary + port = get_server_port() + self.status = "starting" + self.last_error = None + self.public_url = None + + state = load_tunnel_state() + master_host = state.get("master_host") or "" + if state.get("previous_master_host"): + self.previous_master_host = state.get("previous_master_host") + else: + self.previous_master_host = master_host + + os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True) + timestamp = time.strftime("%Y%m%d-%H%M%S") + self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log") + + cmd = [ + binary, + "tunnel", + "--no-autoupdate", + "--url", + f"http://127.0.0.1:{port}", + ] + + debug_log(f"Starting cloudflared: {' '.join(cmd)}") + try: + self.process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + except FileNotFoundError: + self.status = "error" + raise RuntimeError("cloudflared binary not found") + except Exception as exc: + self.status = "error" + raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc + + self.pid = self.process.pid + persist_tunnel_state( + status="starting", + pid=self.pid, + log_file=self.log_file, + previous_host=self.previous_master_host, + ) + + loop = asyncio.get_running_loop() + self._reader.set_log_file(self.log_file) + self._reader.start(self.process, loop) + + try: + await self._reader.wait_for_url(timeout=TUNNEL_START_TIMEOUT) + except asyncio.TimeoutError: + self.last_error = "Timed out waiting for Cloudflare to assign a URL" + await self.stop_tunnel() + raise RuntimeError(self.last_error) + + public_url = self._reader.get_url() + if not public_url: + self.last_error = self._reader.get_last_error() or "Cloudflare tunnel failed to start" + await self.stop_tunnel() + raise RuntimeError(self.last_error) + + self.public_url = public_url + self.status = "running" + debug_log(f"Cloudflare tunnel ready at {self.public_url}") + + persist_tunnel_state( + status="running", + public_url=self.public_url, + pid=self.pid, + log_file=self.log_file, + previous_host=self.previous_master_host or "", + master_host=normalize_host(self.public_url), + ) + return { + "status": self.status, + "public_url": self.public_url, + "pid": self.pid, + "log_file": self.log_file, + } + + async def stop_tunnel(self): + async with self._lock: + pid = self.process.pid if self.process else self.pid + if not pid: + clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host) + self.status = "stopped" + return {"status": "stopped"} + + debug_log(f"Stopping cloudflared (pid={pid})") + if self.process: + terminate_process(self.process, timeout=5) + else: + try: + os.kill(pid, signal.SIGTERM) + time.sleep(0.5) + except Exception as exc: # pragma: no cover + debug_log(f"Error stopping cloudflared pid {pid}: {exc}") + + restore_host = resolve_restore_master_host(self.previous_master_host) + + self.status = "stopped" + self.public_url = None + self.pid = None + self.process = None + self.last_error = None + self._reader.stop() + + clear_tunnel_state( + log_file=self.log_file, + previous_host=self.previous_master_host, + master_host=restore_host, + ) + return {"status": "stopped"} + + def get_status(self): + alive = False + pid = self.process.pid if self.process else self.pid + if pid: + alive = is_process_alive(pid) + if not alive and self.status == "running": + self.status = "stopped" + + return { + "status": self.status, + "public_url": self.public_url, + "pid": pid, + "log_file": self.log_file, + "last_error": self.last_error or self._reader.get_last_error(), + "binary_path": self.binary_path or shutil.which("cloudflared"), + "recent_logs": self._reader.get_recent_logs()[-20:], + "previous_master_host": self.previous_master_host, + } diff --git a/utils/config.py b/utils/config.py index 02fef34..c580415 100644 --- a/utils/config.py +++ b/utils/config.py @@ -1,8 +1,10 @@ """ Configuration management for ComfyUI-Distributed. """ +import asyncio import os import json +from contextlib import asynccontextmanager from .logging import log # Import defaults for timeout fallbacks @@ -11,6 +13,7 @@ from .constants import HEARTBEAT_TIMEOUT CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json") _config_cache = None _config_mtime = 0.0 +_config_lock = asyncio.Lock() def _config_path(): @@ -95,15 +98,36 @@ def load_config(): def save_config(config): """Saves the configuration to file.""" + tmp_path = f"{_config_path()}.tmp" try: - with open(_config_path(), 'w', encoding='utf-8') as f: + with open(tmp_path, 'w', encoding='utf-8') as f: json.dump(config, f, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, _config_path()) invalidate_config_cache() return True except Exception as e: + try: + os.unlink(tmp_path) + except OSError: + pass log(f"Error saving config: {e}") return False + +@asynccontextmanager +async def config_transaction(): + """Acquire config lock, yield loaded config, and save if changed.""" + async with _config_lock: + config = load_config() + original_snapshot = json.dumps(config, sort_keys=True) + yield config + updated_snapshot = json.dumps(config, sort_keys=True) + if updated_snapshot != original_snapshot: + if not save_config(config): + raise RuntimeError("Failed to save config") + def ensure_config_exists(): """Creates default config file if it doesn't exist. Used by __init__.py""" if not os.path.exists(_config_path()): diff --git a/utils/constants.py b/utils/constants.py index 0216583..caca4f9 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -13,6 +13,10 @@ PROCESS_TERMINATION_TIMEOUT = 5.0 WORKER_CHECK_INTERVAL = 2.0 STATUS_CHECK_INTERVAL = 5.0 +# Cloudflare tunnel +TUNNEL_START_TIMEOUT = float(os.environ.get("TUNNEL_START_TIMEOUT", "25")) +CLOUDFLARE_LOG_BUFFER_SIZE = 200 + # Network CHUNK_SIZE = 8192 LOG_TAIL_BYTES = 65536 # 64KB @@ -42,6 +46,13 @@ MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10')) # Heartbeat/check interval in seconds HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s) +# USDU result collection +DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0 + +# Static mode job poll loop +JOB_POLL_INTERVAL = 1.0 +JOB_POLL_MAX_ATTEMPTS = 20 + # Orchestration pipeline ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int( os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8') diff --git a/utils/network.py b/utils/network.py index 1a23bd4..ce5c7a7 100644 --- a/utils/network.py +++ b/utils/network.py @@ -34,8 +34,14 @@ async def cleanup_client_session(): async def handle_api_error(request, error, status=500): """Standardized error response handler.""" - debug_log(f"API Error: {error}") - return web.json_response({"status": "error", "message": str(error)}, status=status) + if isinstance(error, list): + messages = [str(item) for item in error] + debug_log(f"API Error [{status}]: {messages}") + return web.json_response({"errors": messages}, status=status) + + message = str(error) + debug_log(f"API Error [{status}]: {message}") + return web.json_response({"error": message}, status=status) def get_server_port(): """Get the ComfyUI server port.""" diff --git a/utils/trace_logger.py b/utils/trace_logger.py new file mode 100644 index 0000000..e723529 --- /dev/null +++ b/utils/trace_logger.py @@ -0,0 +1,13 @@ +from .logging import debug_log, log + + +def trace_prefix(trace_execution_id: str) -> str: + return f"[Distributed][exec:{trace_execution_id}]" + + +def trace_debug(trace_execution_id: str, message: str) -> None: + debug_log(f"{trace_prefix(trace_execution_id)} {message}") + + +def trace_info(trace_execution_id: str, message: str) -> None: + log(f"{trace_prefix(trace_execution_id)} {message}") diff --git a/utils/usdu_managment.py b/utils/usdu_managment.py index bee033e..a4109f3 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_managment.py @@ -1,343 +1,32 @@ -import asyncio -import time -import json -import copy -import os -import io -from aiohttp import web -import server -from PIL import Image +"""Backward-compatibility shim for USDU helpers. -# Import from other utilities -from .logging import debug_log, log -from .network import build_worker_url, handle_api_error, get_client_session, probe_worker -from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState -# We avoid converting to tensors on the master for tiles; blending uses PIL -from typing import List, Optional +Route handlers and job logic now live in: +- upscale.job_store +- upscale.job_timeout +- upscale.payload_parsers +- api.usdu_routes +""" -# Configure maximum payload size (50MB default, configurable via environment variable) -MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024))) +from ..upscale.conditioning import clone_conditioning, clone_control_chain +from ..upscale.job_store import ( + MAX_PAYLOAD_SIZE, + _cleanup_job, + _drain_results_queue, + _get_completed_count, + _init_job_queue, + _mark_task_completed, + ensure_tile_jobs_initialized, + init_dynamic_job, + init_static_job_batched, +) +from ..upscale.job_timeout import _check_and_requeue_timed_out_workers +from ..upscale.payload_parsers import _parse_tiles_from_form +from .logging import debug_log +from .network import get_client_session -# Import HEARTBEAT_TIMEOUT from constants -from .constants import HEARTBEAT_TIMEOUT -from .config import load_config - - -def _parse_tiles_from_form(data): - """Parse tiles submitted via multipart/form-data into a list of tile dicts. - - Expects the following fields in the aiohttp form data: - - 'tiles_metadata': JSON list with per-tile metadata items containing at least - 'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional - 'batch_idx' and 'global_idx' are included when available. - - 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG). - - 'padding': integer padding used during extraction (optional; defaults 0). - - Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y', - 'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx', - plus 'padding'. - """ - try: - # Parse padding if present - padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0 - except Exception: - padding = 0 - - # Parse tiles metadata (JSON list) - meta_raw = data.get('tiles_metadata') - if meta_raw is None: - raise ValueError("Missing tiles_metadata") - - try: - metadata = json.loads(meta_raw) - except Exception as e: - raise ValueError(f"Invalid tiles_metadata JSON: {e}") - - if not isinstance(metadata, list): - raise ValueError("tiles_metadata must be a list") - - tiles = [] - # Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ... - for i, meta in enumerate(metadata): - file_field = data.get(f'tile_{i}') - if file_field is None or not hasattr(file_field, 'file'): - raise ValueError(f"Missing tile data for index {i}") - - # Read image bytes and decode to PIL - raw = file_field.file.read() - try: - img = Image.open(io.BytesIO(raw)).convert("RGB") - except Exception as e: - raise ValueError(f"Invalid image data for tile {i}: {e}") - - # Build tile dictionary (store PIL only; master blends via PIL) - try: - tile_info = { - 'image': img, - 'tile_idx': int(meta.get('tile_idx', i)), - 'x': int(meta.get('x', 0)), - 'y': int(meta.get('y', 0)), - 'extracted_width': int(meta.get('extracted_width', img.width)), - 'extracted_height': int(meta.get('extracted_height', img.height)), - 'padding': int(padding), - } - except Exception as e: - raise ValueError(f"Invalid metadata values for tile {i}: {e}") - - # Optional fields - if 'batch_idx' in meta: - try: - tile_info['batch_idx'] = int(meta['batch_idx']) - except Exception: - pass - if 'global_idx' in meta: - try: - tile_info['global_idx'] = int(meta['global_idx']) - except Exception: - pass - - tiles.append(tile_info) - - return tiles - -async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None): - """Initialize queue for dynamic mode (per-image), with collector fields. - - - Creates pending image queue entries - - Initializes dynamic completion/status tracking - """ - await _init_job_queue( - multi_job_id, - 'dynamic', - batch_size=batch_size, - all_indices=all_indices or list(range(batch_size)), - enabled_workers=enabled_workers, - ) - debug_log(f"Job {multi_job_id} initialized with {batch_size} images") - - -async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]): - """Initialize queue for static mode (batched-per-tile). - - - Populates pending tile queue with tile ids [0..num_tiles_per_image-1] - """ - await _init_job_queue( - multi_job_id, - 'static', - batch_size=batch_size, - num_tiles_per_image=num_tiles_per_image, - enabled_workers=enabled_workers, - batched_static=True, - ) - # Initialization handled by master; avoid duplicate init logs here - -async def _init_job_queue( - multi_job_id, - mode, - batch_size=None, - num_tiles_per_image=None, - all_indices=None, - enabled_workers=None, - batched_static: bool = False, -): - """Unified initialization for job queues in static and dynamic modes.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - debug_log(f"Queue already exists for {multi_job_id}") - return - - if mode == 'dynamic': - job_data = ImageJobState(multi_job_id=multi_job_id) - elif mode == 'static': - job_data = TileJobState(multi_job_id=multi_job_id) - else: - raise ValueError(f"Unknown mode: {mode}") - - job_data.worker_status = {w: time.time() for w in enabled_workers or []} - job_data.assigned_to_workers = {w: [] for w in enabled_workers or []} - - if mode == 'dynamic': - job_data.batch_size = int(batch_size or 0) - pending_queue = job_data.pending_images - for i in (all_indices or range(int(batch_size or 0))): - await pending_queue.put(i) - debug_log(f"Initialized image queue with {batch_size} pending items") - elif mode == 'static': - job_data.num_tiles_per_image = int(num_tiles_per_image or 0) - job_data.batch_size = int(batch_size or 0) - job_data.batched_static = bool(batched_static) - # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1] - pending_queue = job_data.pending_tasks - if batched_static and num_tiles_per_image is not None: - for i in range(num_tiles_per_image): - await pending_queue.put(i) - else: - total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0) - for i in range(total_tiles): - await pending_queue.put(i) - - prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data - -async def _drain_results_queue(multi_job_id): - """Drain pending results from queue and update completed_tasks. Returns count drained. - - Uses non-blocking get_nowait to avoid await timeouts and reduce latency. - """ - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if not isinstance(job_data, BaseJobState): - return 0 - q = job_data.queue - completed_tasks = job_data.completed_tasks - - collected = 0 - while True: - try: - result = q.get_nowait() - except asyncio.QueueEmpty: - break - - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - - if 'image_idx' in result and 'image' in result: - task_id = result['image_idx'] - if task_id not in completed_tasks: - completed_tasks[task_id] = result['image'] - collected += 1 - elif 'tiles' in result: - for tile_data in result['tiles']: - task_id = tile_data.get('global_idx', tile_data['tile_idx']) - if task_id not in completed_tasks: - completed_tasks[task_id] = tile_data - collected += 1 - if is_last: - # Track worker completion - if worker_id in job_data.worker_status: - del job_data.worker_status[worker_id] - - return collected - -async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks): - """Check timed out workers and requeue their tasks. Returns requeued count.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if not job_data: - return 0 - - current_time = time.time() - if not isinstance(job_data, BaseJobState): - return 0 - - requeued_count = 0 - completed_tasks = job_data.completed_tasks - - # Allow override via config setting 'worker_timeout_seconds' - cfg = load_config() - hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT)) - - for worker, last_heartbeat in list(job_data.worker_status.items()): - age = current_time - last_heartbeat - debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s") - if age > hb_timeout: - # Busy-only grace policy: require positive signal from worker (/prompt) - # We also log assignment state for diagnostics but do not grace on it alone. - assigned = job_data.assigned_to_workers.get(worker, []) - incomplete_assigned = 0 - try: - if assigned: - batched_static = bool(job_data.batched_static) - if batched_static: - num_tiles_per_image = int(job_data.num_tiles_per_image or 1) - batch_size = int(job_data.batch_size or 1) - for task_id in assigned: - for b in range(batch_size): - gidx = b * num_tiles_per_image + task_id - if gidx not in completed_tasks: - incomplete_assigned += 1 - break - else: - for task_id in assigned: - if task_id not in completed_tasks: - incomplete_assigned += 1 - debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}") - except Exception as e: - debug_log(f"Assigned diagnostics failed for worker {worker}: {e}") - - busy = False - probe_queue = None - try: - cfg_workers = load_config().get('workers', []) - wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None) - if wrec: - worker_url = build_worker_url(wrec) - debug_log(f"Probing worker {worker} at {worker_url}/prompt") - payload = await probe_worker(worker_url, timeout=2.0) - if payload is not None: - probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0)) - busy = probe_queue is not None and probe_queue > 0 - except Exception as e: - debug_log(f"Probe failed for worker {worker}: {e}") - finally: - debug_log( - f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}" - ) - - if busy: - job_data.worker_status[worker] = current_time - debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue") - continue - - log(f"Worker {worker} heartbeat timed out after {age:.1f}s") - for task_id in job_data.assigned_to_workers.get(worker, []): - # If batched_static, task_id is a tile_idx; consider it complete only if - # all corresponding global_idx entries are present in completed_tasks. - batched_static = bool(job_data.batched_static) - if batched_static: - num_tiles_per_image = int(job_data.num_tiles_per_image or 1) - batch_size = int(job_data.batch_size or 1) - # Check all global indices for this tile across the batch - all_done = True - for b in range(batch_size): - gidx = b * num_tiles_per_image + task_id - if gidx not in completed_tasks: - all_done = False - break - if not all_done: - await job_data.pending_tasks.put(task_id) - requeued_count += 1 - else: - if task_id not in completed_tasks: - await job_data.pending_tasks.put(task_id) - requeued_count += 1 - job_data.worker_status.pop(worker, None) - if worker in job_data.assigned_to_workers: - job_data.assigned_to_workers[worker] = [] - - return requeued_count - -async def _get_completed_count(multi_job_id): - """Get count of completed tasks.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if isinstance(job_data, BaseJobState): - return len(job_data.completed_tasks) - return 0 - -async def _mark_task_completed(multi_job_id, task_id, result): - """Mark a task as completed.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if isinstance(job_data, BaseJobState): - job_data.completed_tasks[task_id] = result async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id): - """Send heartbeat to master.""" + """Send heartbeat to master from worker-side processing loops.""" try: data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} session = await get_client_session() @@ -347,282 +36,20 @@ async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id): except Exception as e: debug_log(f"Heartbeat failed: {e}") -async def _cleanup_job(multi_job_id): - """Cleanup the job data.""" - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - del prompt_server.distributed_pending_tile_jobs[multi_job_id] - debug_log(f"Cleaned up job {multi_job_id}") -# API Endpoints (generalized) - -@server.PromptServer.instance.routes.post("/distributed/heartbeat") -async def heartbeat_endpoint(request): - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if isinstance(job_data, BaseJobState): - job_data.worker_status[worker_id] = time.time() - debug_log(f"Heartbeat from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Worker status tracking not available", 400) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.post("/distributed/submit_tiles") -async def submit_tiles_endpoint(request): - """Endpoint for workers to submit processed tiles in static mode.""" - try: - content_length = request.headers.get('content-length') - if content_length and int(content_length) > MAX_PAYLOAD_SIZE: - return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) - - data = await request.post() - multi_job_id = data.get('multi_job_id') - worker_id = data.get('worker_id') - is_last = data.get('is_last', 'False').lower() == 'true' - - if multi_job_id is None or worker_id is None: - return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - - batch_size = int(data.get('batch_size', 0)) - tiles = [] - - # Handle completion signal - if batch_size == 0 and is_last: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, TileJobState): - return await handle_api_error(request, "Job not configured for tile submissions", 400) - await job_data.queue.put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) - - try: - tiles = _parse_tiles_from_form(data) - except ValueError as e: - return await handle_api_error(request, str(e), 400) - - # Submit tiles to queue - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, TileJobState): - return await handle_api_error(request, "Job not configured for tile submissions", 400) - - q = job_data.queue - if batch_size > 0 or len(tiles) > 0: - await q.put({ - 'worker_id': worker_id, - 'tiles': tiles, - 'is_last': is_last - }) - debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})") - else: - await q.put({ - 'worker_id': worker_id, - 'is_last': True, - 'tiles': [] - }) - - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.post("/distributed/submit_image") -async def submit_image_endpoint(request): - """Endpoint for workers to submit processed images in dynamic mode.""" - try: - content_length = request.headers.get('content-length') - if content_length and int(content_length) > MAX_PAYLOAD_SIZE: - return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413) - - data = await request.post() - multi_job_id = data.get('multi_job_id') - worker_id = data.get('worker_id') - is_last = data.get('is_last', 'False').lower() == 'true' - - if multi_job_id is None or worker_id is None: - return await handle_api_error(request, "Missing multi_job_id or worker_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - - # Handle image submission - if 'full_image' in data and 'image_idx' in data: - image_idx = int(data.get('image_idx')) - img_data = data['full_image'].file.read() - img = Image.open(io.BytesIO(img_data)).convert("RGB") - - debug_log(f"Received full image {image_idx} from worker {worker_id}") - - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, ImageJobState): - return await handle_api_error(request, "Job not configured for image submissions", 400) - await job_data.queue.put({ - 'worker_id': worker_id, - 'image_idx': image_idx, - 'image': img, - 'is_last': is_last - }) - return web.json_response({"status": "success"}) - - # Handle completion signal (no image data) - elif is_last: - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, ImageJobState): - return await handle_api_error(request, "Job not configured for image submissions", 400) - await job_data.queue.put({ - 'worker_id': worker_id, - 'is_last': True, - }) - debug_log(f"Received completion signal from worker {worker_id}") - return web.json_response({"status": "success"}) - else: - return await handle_api_error(request, "Missing image data or invalid request", 400) - - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -# Helper functions for shallow copying conditioning without duplicating models -def clone_control_chain(control, clone_hint=True): - """Shallow copy the ControlNet chain, optionally cloning hints but sharing models.""" - if control is None: - return None - new_control = copy.copy(control) # Shallow copy (shares model) - if clone_hint and hasattr(control, 'cond_hint_original'): - hint = getattr(control, 'cond_hint_original', None) - new_control.cond_hint_original = hint.clone() if hint is not None else None - if hasattr(control, 'previous_controlnet'): - new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint) - return new_control - -def clone_conditioning(cond_list, clone_hints=True): - """Clone conditioning without duplicating ControlNet models.""" - new_cond = [] - for emb, cond_dict in cond_list: - new_emb = emb.clone() if emb is not None else None - new_dict = cond_dict.copy() - if 'control' in new_dict: - new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints) - if 'mask' in new_dict: - if new_dict['mask'] is not None: - new_dict['mask'] = new_dict['mask'].clone() - # Handle other potential fields if needed - if 'pooled_output' in new_dict: - if new_dict['pooled_output'] is not None: - new_dict['pooled_output'] = new_dict['pooled_output'].clone() - if 'area' in new_dict: - new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple - new_cond.append([new_emb, new_dict]) - return new_cond - -# Helper function to ensure persistent state is initialized -def ensure_tile_jobs_initialized(): - """Ensure tile job storage is initialized on the server instance.""" - prompt_server = server.PromptServer.instance - if not hasattr(prompt_server, 'distributed_pending_tile_jobs'): - debug_log("Initializing persistent tile job queue on server instance.") - prompt_server.distributed_pending_tile_jobs = {} - prompt_server.distributed_tile_jobs_lock = asyncio.Lock() - else: - invalid_job_ids = [ - job_id - for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items() - if not isinstance(job_data, BaseJobState) - ] - for job_id in invalid_job_ids: - debug_log(f"Removing invalid job state for {job_id}") - del prompt_server.distributed_pending_tile_jobs[job_id] - return prompt_server - -# API Endpoint for tile completion -@server.PromptServer.instance.routes.post("/distributed/request_image") -async def request_image_endpoint(request): - """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode).""" - try: - data = await request.json() - worker_id = data.get('worker_id') - multi_job_id = data.get('multi_job_id') - - if not worker_id or not multi_job_id: - return await handle_api_error(request, "Missing worker_id or multi_job_id", 400) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id] - if not isinstance(job_data, BaseJobState): - return await handle_api_error(request, "Invalid job data structure", 500) - - mode = job_data.mode - if isinstance(job_data, ImageJobState): - pending_queue = job_data.pending_images - elif isinstance(job_data, TileJobState): - pending_queue = job_data.pending_tasks - else: - return await handle_api_error(request, "Invalid job configuration", 400) - - try: - task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - # Track assigned task - job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx) - # Update worker heartbeat - job_data.worker_status[worker_id] = time.time() - # Get estimated remaining count - remaining = pending_queue.qsize() # Approximate - - # Return appropriate response based on mode - if mode == 'dynamic': - debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}") - return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining}) - else: # static - debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}") - return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.batched_static}) - except asyncio.TimeoutError: - if mode == 'dynamic': - return web.json_response({"image_idx": None}) - else: - return web.json_response({"tile_idx": None}) - else: - return await handle_api_error(request, "Job not found", 404) - except Exception as e: - return await handle_api_error(request, e, 500) - -@server.PromptServer.instance.routes.get("/distributed/job_status") -async def job_status_endpoint(request): - """Endpoint to check if a job is ready.""" - multi_job_id = request.query.get('multi_job_id') - if not multi_job_id: - return web.json_response({"ready": False}) - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None) - return web.json_response({"ready": ready}) +__all__ = [ + "MAX_PAYLOAD_SIZE", + "_check_and_requeue_timed_out_workers", + "_cleanup_job", + "_drain_results_queue", + "_get_completed_count", + "_init_job_queue", + "_mark_task_completed", + "_parse_tiles_from_form", + "_send_heartbeat_to_master", + "clone_conditioning", + "clone_control_chain", + "ensure_tile_jobs_initialized", + "init_dynamic_job", + "init_static_job_batched", +] diff --git a/web/apiClient.js b/web/apiClient.js index 6cfcda9..7c1292e 100644 --- a/web/apiClient.js +++ b/web/apiClient.js @@ -4,27 +4,40 @@ import { normalizeWorkerUrl } from './urlUtils.js'; export function createApiClient(baseUrl) { const normalizedBaseUrl = normalizeWorkerUrl(baseUrl); - const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => { + const request = async ( + endpoint, + options = {}, + { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {}, + ) => { + const maxAttempts = retry ? retries : 1; let lastError; let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff - for (let attempt = 0; attempt < retries; attempt++) { + for (let attempt = 0; attempt < maxAttempts; attempt++) { try { + const headers = { + 'Content-Type': 'application/json', + ...(options.headers || {}), + }; const response = await fetch(`${normalizedBaseUrl}${endpoint}`, { - headers: { 'Content-Type': 'application/json' }, - ...options + ...options, + headers, }); if (!response.ok) { - const error = await response.json().catch(() => ({ message: 'Request failed' })); - throw new Error(error.message || `HTTP ${response.status}`); + const error = await response.json().catch(() => ({})); + const message = error.message + || error.error + || (Array.isArray(error.errors) ? error.errors.join('; ') : null) + || `HTTP ${response.status}`; + throw new Error(message); } return await response.json(); } catch (error) { lastError = error; - console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`); - if (attempt < retries - 1) { + console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`); + if (attempt < maxAttempts - 1) { await new Promise(resolve => setTimeout(resolve, delay)); delay *= 2; // Exponential backoff } @@ -32,6 +45,39 @@ export function createApiClient(baseUrl) { } throw lastError; }; + + const requestUrl = async ( + url, + options = {}, + { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {}, + ) => { + const maxAttempts = retry ? retries : 1; + let lastError; + let delay = TIMEOUTS.RETRY_DELAY; + + for (let attempt = 0; attempt < maxAttempts; attempt++) { + try { + const response = await fetch(url, options); + if (!response.ok) { + const error = await response.json().catch(() => ({})); + const message = error.message + || error.error + || (Array.isArray(error.errors) ? error.errors.join('; ') : null) + || `HTTP ${response.status}`; + throw new Error(message); + } + return await response.json(); + } catch (error) { + lastError = error; + console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${url} - ${error.message}`); + if (attempt < maxAttempts - 1) { + await new Promise(resolve => setTimeout(resolve, delay)); + delay *= 2; + } + } + } + throw lastError; + }; return { // Config endpoints @@ -43,28 +89,28 @@ export function createApiClient(baseUrl) { return request('/distributed/config/update_worker', { method: 'POST', body: JSON.stringify({ worker_id: workerId, ...data }) - }); + }, { retry: false }); }, async deleteWorker(workerId) { return request('/distributed/config/delete_worker', { method: 'POST', body: JSON.stringify({ worker_id: workerId }) - }); + }, { retry: false }); }, async updateSetting(key, value) { return request('/distributed/config/update_setting', { method: 'POST', body: JSON.stringify({ key, value }) - }); + }, { retry: false }); }, async updateMaster(data) { return request('/distributed/config/update_master', { method: 'POST', body: JSON.stringify(data) - }); + }, { retry: false }); }, // Worker management endpoints @@ -72,14 +118,14 @@ export function createApiClient(baseUrl) { return request('/distributed/launch_worker', { method: 'POST', body: JSON.stringify({ worker_id: workerId }) - }); + }, { retry: false }); }, async stopWorker(workerId) { return request('/distributed/stop_worker', { method: 'POST', body: JSON.stringify({ worker_id: workerId }) - }); + }, { retry: false }); }, async getManagedWorkers() { @@ -94,51 +140,59 @@ export function createApiClient(baseUrl) { return request('/distributed/worker/clear_launching', { method: 'POST', body: JSON.stringify({ worker_id: workerId }) - }); + }, { retry: false }); }, async queueDistributed(payload) { return request('/distributed/queue', { method: 'POST', + headers: { + ...(payload?.trace_execution_id + ? { 'X-Idempotency-Key': payload.trace_execution_id } + : {}), + }, body: JSON.stringify(payload) - }); + }, { retry: false }); }, - async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK) { + async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK, signal = null) { const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl); - const response = await fetch(`${normalizedWorkerUrl}/prompt`, { - method: 'GET', - mode: 'cors', - cache: 'no-store', - signal: AbortSignal.timeout(timeoutMs) - }); + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), timeoutMs); + const effectiveSignal = signal + ? AbortSignal.any([controller.signal, signal]) + : controller.signal; + try { + const response = await fetch(`${normalizedWorkerUrl}/prompt`, { + method: 'GET', + mode: 'cors', + cache: 'no-store', + signal: effectiveSignal, + }); - if (!response.ok) { - return { ok: false, status: response.status, queueRemaining: null }; + if (!response.ok) { + return { ok: false, status: response.status, queueRemaining: null }; + } + + const data = await response.json().catch(() => ({})); + return { + ok: true, + status: response.status, + queueRemaining: data.exec_info?.queue_remaining || 0, + }; + } finally { + clearTimeout(timeoutId); } - - const data = await response.json().catch(() => ({})); - return { - ok: true, - status: response.status, - queueRemaining: data.exec_info?.queue_remaining || 0, - }; }, async dispatchToWorker(workerUrl, promptPayload) { const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl); - const response = await fetch(`${normalizedWorkerUrl}/prompt`, { + return requestUrl(`${normalizedWorkerUrl}/prompt`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, mode: 'cors', body: JSON.stringify(promptPayload), - }); - - if (!response.ok) { - throw new Error(`Worker returned ${response.status}`); - } - - return response.json(); + }, { retry: false }); }, // Network info @@ -179,14 +233,14 @@ export function createApiClient(baseUrl) { return request('/distributed/tunnel/start', { method: 'POST', body: JSON.stringify({}) - }); + }, { retry: false }); }, async stopTunnel() { return request('/distributed/tunnel/stop', { method: 'POST', body: JSON.stringify({}) - }); + }, { retry: false }); }, async getTunnelStatus() { diff --git a/web/distributed.css b/web/distributed.css index 118fd82..0024e65 100644 --- a/web/distributed.css +++ b/web/distributed.css @@ -5,6 +5,13 @@ --btn-working: #666; --btn-success: #3a6a3a; --btn-error: #6a3a3a; + --tunnel-enable: #665533; + --tunnel-disable: #7c4a4a; + --master-badge-fallback-bg: #243024; + --master-badge-fallback-text: #6bd06b; + --master-badge-fallback-border: #335533; + --master-badge-delegate-bg: #3a3a3a; + --master-badge-delegate-text: #ffcc66; } .btn--stop { @@ -30,3 +37,167 @@ .btn--error { background-color: var(--btn-error) !important; } + +.master-info-badge--fallback { + background-color: var(--master-badge-fallback-bg) !important; + color: var(--master-badge-fallback-text) !important; + border: 1px solid var(--master-badge-fallback-border) !important; +} + +.master-info-badge--delegate { + background-color: var(--master-badge-delegate-bg) !important; + color: var(--master-badge-delegate-text) !important; +} + +.entity-card-content--hovered { + background-color: #333 !important; +} + +.placeholder-card--blueprint { + border-color: #555 !important; + background-color: rgba(255, 255, 255, 0.02) !important; +} + +.placeholder-card--blueprint.is-hovered { + border-color: #777 !important; + background-color: rgba(255, 255, 255, 0.05) !important; +} + +.placeholder-card--add { + border-color: #444 !important; + background-color: transparent !important; +} + +.placeholder-card--add.is-hovered { + border-color: #666 !important; + background-color: rgba(255, 255, 255, 0.02) !important; +} + +.placeholder-column--blueprint { + color: #777 !important; +} + +.placeholder-column--blueprint.is-hovered { + color: #999 !important; +} + +.placeholder-column--add { + color: #555 !important; + border-color: #444 !important; +} + +.placeholder-column--add.is-hovered { + color: #888 !important; + border-color: #666 !important; +} + +.tunnel-button--enable { + background-color: var(--tunnel-enable) !important; +} + +.tunnel-button--disable { + background-color: var(--tunnel-disable) !important; +} + +.tunnel-status--enable { + background-color: var(--tunnel-enable) !important; +} + +.tunnel-status--disable { + background-color: var(--tunnel-disable) !important; +} + +.is-hidden { + display: none !important; +} + +.settings-arrow--expanded { + transform: rotate(90deg) !important; +} + +.log-modal { + position: fixed; + top: 0; + left: 0; + width: 100%; + height: 100%; + background: rgba(0, 0, 0, 0.8); + display: flex; + align-items: center; + justify-content: center; + z-index: 10000; +} + +.log-modal__content { + background: #1e1e1e; + border-radius: 8px; + width: 90%; + max-width: 1200px; + height: 80%; + display: flex; + flex-direction: column; + border: 1px solid #444; +} + +.log-modal__header { + padding: 15px 20px; + border-bottom: 1px solid #444; + display: flex; + justify-content: space-between; + align-items: center; +} + +.log-modal__title { + margin: 0; + color: #fff; +} + +.log-modal__header-buttons { + display: flex; + gap: 20px; + align-items: center; +} + +.log-modal__refresh { + display: flex; + align-items: center; + gap: 4px; +} + +.log-modal__refresh-input { + cursor: pointer; +} + +.log-modal__refresh-label { + font-size: 12px; + color: #ccc; + cursor: pointer; + white-space: nowrap; +} + +.log-modal__close { + background-color: #c04c4c !important; + padding: 5px 10px !important; + font-size: 14px !important; + font-weight: bold !important; +} + +.log-modal__body { + flex: 1; + overflow: auto; + padding: 15px; + font-family: 'Consolas', 'Monaco', 'Courier New', monospace; + font-size: 12px; + line-height: 1.4; + color: #ddd; + background: #0d0d0d; + white-space: pre-wrap; + word-wrap: break-word; +} + +.log-modal__status { + padding: 10px 20px; + border-top: 1px solid #444; + font-size: 11px; + color: #888; +} diff --git a/web/executionUtils.js b/web/executionUtils.js index 03f6197..c6cfb18 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -1,5 +1,5 @@ import { api } from "../../scripts/api.js"; -import { findNodesByClass } from './workerUtils.js'; +import { applyProbeResultToWorkerDot, findNodesByClass } from './workerUtils.js'; import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js'; import { markSkipDispatch, selectLeastBusyWorker } from './executionDecisionUtils.js'; import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js'; @@ -232,13 +232,7 @@ export async function performPreflightCheck(extension, workers) { // Update UI status indicators for inactive workers results.filter(r => !r.active).forEach(r => { - const statusDot = document.getElementById(`status-${r.worker.id}`); - if (statusDot) { - // Remove pulsing animation once status is determined - statusDot.classList.remove('status-pulsing', 'worker-status--online', 'worker-status--unknown'); - statusDot.classList.add('worker-status--offline'); - statusDot.title = "Offline - Cannot connect"; - } + applyProbeResultToWorkerDot(r.worker.id, { ok: false }); }); return activeWorkers; diff --git a/web/ui.js b/web/ui.js index 1d51197..3185d7e 100644 --- a/web/ui.js +++ b/web/ui.js @@ -3,7 +3,8 @@ import { createButtonHelper as createButtonHelperFn } from './ui/buttonHelpers.j import { showCloudflareWarning as showCloudflareWarningFn } from './ui/cloudflareWarning.js'; import { createWorkerSettingsForm as createWorkerSettingsFormFn } from './ui/settingsForm.js'; import { renderEntityCard as renderEntityCardFn } from './ui/entityCard.js'; -import { launchWorker, startLogAutoRefresh, stopLogAutoRefresh, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js'; +import { createLogModal } from './ui/logModal.js'; +import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js'; import { isRemoteWorker } from './workerSettings.js'; const cardConfigs = { @@ -377,184 +378,34 @@ export class DistributedUI { [STATUS_COLORS.PROCESSING_YELLOW]: "worker-status--processing", }; - const statusClass = colorClassMap[color]; - if (statusClass) { - statusDot.classList.add(statusClass); - statusDot.style.backgroundColor = ""; - } else { - statusDot.style.backgroundColor = color; - } + const statusClass = colorClassMap[color] || "worker-status--unknown"; + statusDot.classList.add(statusClass); + statusDot.style.backgroundColor = ""; statusDot.title = title; statusDot.classList.toggle('status-pulsing', pulsing); } showLogModal(extension, workerId, logData) { - // Remove any existing modal - const existingModal = document.getElementById('distributed-log-modal'); - if (existingModal) { - existingModal.remove(); + if (this._logModal) { + this._logModal.unmount(); + this._logModal = null; } - + const worker = extension.config.workers.find(w => w.id === workerId); const workerName = worker?.name || `Worker ${workerId}`; - - // Create modal container - const modal = document.createElement('div'); - modal.id = 'distributed-log-modal'; - modal.style.cssText = ` - position: fixed; - top: 0; - left: 0; - width: 100%; - height: 100%; - background: rgba(0, 0, 0, 0.8); - display: flex; - align-items: center; - justify-content: center; - z-index: 10000; - `; - - // Create modal content - const content = document.createElement('div'); - content.style.cssText = ` - background: #1e1e1e; - border-radius: 8px; - width: 90%; - max-width: 1200px; - height: 80%; - display: flex; - flex-direction: column; - border: 1px solid #444; - `; - - // Header - const header = document.createElement('div'); - header.style.cssText = ` - padding: 15px 20px; - border-bottom: 1px solid #444; - display: flex; - justify-content: space-between; - align-items: center; - `; - - const title = document.createElement('h3'); - title.textContent = `${workerName} - Log Viewer`; - title.style.cssText = 'margin: 0; color: #fff;'; - - const headerButtons = document.createElement('div'); - headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;'; - - // Auto-refresh container - const refreshContainer = document.createElement('div'); - refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;'; - - // Auto-refresh checkbox - const refreshCheckbox = document.createElement('input'); - refreshCheckbox.type = 'checkbox'; - refreshCheckbox.id = 'log-auto-refresh'; - refreshCheckbox.checked = true; // Enabled by default - refreshCheckbox.style.cssText = 'cursor: pointer;'; - refreshCheckbox.onchange = (e) => { - if (e.target.checked) { - startLogAutoRefresh(extension, workerId); - } else { - stopLogAutoRefresh(extension); - } - }; - - const refreshLabel = document.createElement('label'); - refreshLabel.htmlFor = 'log-auto-refresh'; - refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;'; - refreshLabel.textContent = 'Auto-refresh'; - - // Add checkbox and label to container - refreshContainer.appendChild(refreshCheckbox); - refreshContainer.appendChild(refreshLabel); - - // Close button - const closeBtn = this.createButton('✕', - () => { - stopLogAutoRefresh(extension); - modal.remove(); - }, - 'background-color: #c04c4c;'); - closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;'; - - headerButtons.appendChild(refreshContainer); - headerButtons.appendChild(closeBtn); - - header.appendChild(title); - header.appendChild(headerButtons); - - // Log content area - const logContainer = document.createElement('div'); - logContainer.style.cssText = ` - flex: 1; - overflow: auto; - padding: 15px; - font-family: 'Consolas', 'Monaco', 'Courier New', monospace; - font-size: 12px; - line-height: 1.4; - color: #ddd; - background: #0d0d0d; - white-space: pre-wrap; - word-wrap: break-word; - `; - logContainer.id = 'distributed-log-content'; - logContainer.textContent = logData.content; - - // Auto-scroll to bottom - setTimeout(() => { - logContainer.scrollTop = logContainer.scrollHeight; - }, 0); - - // Status bar - const statusBar = document.createElement('div'); - statusBar.style.cssText = ` - padding: 10px 20px; - border-top: 1px solid #444; - font-size: 11px; - color: #888; - `; - statusBar.textContent = `Log file: ${logData.log_file}`; - if (logData.truncated) { - statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`; - } - - // Assemble modal - content.appendChild(header); - content.appendChild(logContainer); - content.appendChild(statusBar); - modal.appendChild(content); - - // Close on background click - modal.onclick = (e) => { - if (e.target === modal) { - stopLogAutoRefresh(extension); - modal.remove(); - } - }; - - // Close on Escape key - const handleEscape = (e) => { - if (e.key === 'Escape') { - stopLogAutoRefresh(extension); - modal.remove(); - document.removeEventListener('keydown', handleEscape); - } - }; - document.addEventListener('keydown', handleEscape); - - document.body.appendChild(modal); - - // Start auto-refresh - startLogAutoRefresh(extension, workerId); - } - formatFileSize(bytes) { - if (bytes < 1024) return bytes + ' B'; - if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; - return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; + const modal = createLogModal(); + this._logModal = modal; + modal.mount(document.body, { + workerName, + logData, + fetchLog: async () => extension.api.getWorkerLog(workerId, 1000), + onClose: () => { + if (this._logModal === modal) { + this._logModal = null; + } + }, + }); } createWorkerSettingsForm(extension, worker) { @@ -712,17 +563,15 @@ export class DistributedUI { let message; const badge = document.createElement("div"); badge.style.cssText = this.styles.infoBox; + badge.classList.add("master-info-badge"); if (fallbackActive) { message = "No workers selected. Master fallback execution active."; badge.textContent = message; - badge.style.backgroundColor = "#243024"; - badge.style.color = "#6bd06b"; - badge.style.border = "1px solid #335533"; + badge.classList.add("master-info-badge--fallback"); } else if (!participationEnabled) { message = "Master disabled: running as orchestrator only."; badge.textContent = message; - badge.style.backgroundColor = "#3a3a3a"; - badge.style.color = "#ffcc66"; + badge.classList.add("master-info-badge--delegate"); } else { message = "Master participating in workflows."; badge.textContent = message; @@ -822,9 +671,10 @@ export class DistributedUI { settingsForm.appendChild(hostResult.group); // Cloudflare tunnel toggle (simple button inside master settings) - const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target), "background-color: #665533;"); + const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target)); tunnelBtn.id = "cloudflare-tunnel-button"; - tunnelBtn.style.cssText = BUTTON_STYLES.base + " background-color: #665533; margin: 4px 0 -5px 0;"; + tunnelBtn.style.cssText = BUTTON_STYLES.base + " margin: 4px 0 -5px 0;"; + tunnelBtn.classList.add("tunnel-button", "tunnel-button--enable"); settingsForm.appendChild(tunnelBtn); extension.tunnelElements = { button: tunnelBtn }; extension.updateTunnelUIElements(); @@ -890,30 +740,19 @@ export class DistributedUI { } addPlaceholderHover(card, leftColumn, entityType) { + const cardTypeClass = entityType === 'blueprint' ? 'placeholder-card--blueprint' : 'placeholder-card--add'; + const columnTypeClass = entityType === 'blueprint' ? 'placeholder-column--blueprint' : 'placeholder-column--add'; + card.classList.add('placeholder-card', cardTypeClass); + leftColumn.classList.add('placeholder-column', columnTypeClass); + card.onmouseover = () => { - if (entityType === 'blueprint') { - card.style.borderColor = "#777"; - card.style.backgroundColor = "rgba(255, 255, 255, 0.05)"; - leftColumn.style.color = "#999"; - } else { - card.style.borderColor = "#666"; - card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - leftColumn.style.color = "#888"; - leftColumn.style.borderColor = "#666"; - } + card.classList.add('is-hovered'); + leftColumn.classList.add('is-hovered'); }; card.onmouseout = () => { - if (entityType === 'blueprint') { - card.style.borderColor = "#555"; - card.style.backgroundColor = "rgba(255, 255, 255, 0.02)"; - leftColumn.style.color = "#777"; - } else { - card.style.borderColor = "#444"; - card.style.backgroundColor = "transparent"; - leftColumn.style.color = "#555"; - leftColumn.style.borderColor = "#444"; - } + card.classList.remove('is-hovered'); + leftColumn.classList.remove('is-hovered'); }; } diff --git a/web/ui/logModal.js b/web/ui/logModal.js new file mode 100644 index 0000000..80b7532 --- /dev/null +++ b/web/ui/logModal.js @@ -0,0 +1,192 @@ +import { TIMEOUTS } from '../constants.js'; + +function formatFileSize(bytes) { + if (bytes < 1024) return bytes + ' B'; + if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB'; + return (bytes / (1024 * 1024)).toFixed(1) + ' MB'; +} + +export function createLogModal() { + let _modalEl = null; + let _keydownHandler = null; + let _refreshTimer = null; + let _fetching = false; + let _fetchLog = null; + let _onClose = null; + let _logContentEl = null; + let _statusBarEl = null; + let _refreshCheckbox = null; + + const updateLogView = (logData) => { + if (!_logContentEl || !_statusBarEl || !logData) { + return; + } + + const shouldAutoScroll = + _logContentEl.scrollTop + _logContentEl.clientHeight >= _logContentEl.scrollHeight - 50; + _logContentEl.textContent = logData.content || ''; + if (shouldAutoScroll) { + _logContentEl.scrollTop = _logContentEl.scrollHeight; + } + + let statusText = `Log file: ${logData.log_file || 'unknown'}`; + if (logData.truncated) { + statusText += ` (showing last ${logData.lines_shown} lines of ${formatFileSize(logData.file_size || 0)})`; + } + _statusBarEl.textContent = statusText; + }; + + const refreshLog = async () => { + if (_fetching || !_fetchLog || !_refreshCheckbox?.checked) { + return; + } + _fetching = true; + try { + const data = await _fetchLog(); + if (data) { + updateLogView(data); + } + } catch (_error) { + // Keep the modal open and continue retrying on next interval. + } finally { + _fetching = false; + } + }; + + const stopRefresh = () => { + if (_refreshTimer) { + clearInterval(_refreshTimer); + _refreshTimer = null; + } + }; + + const startRefresh = () => { + stopRefresh(); + _refreshTimer = setInterval(() => { + refreshLog(); + }, TIMEOUTS.LOG_REFRESH); + }; + + const unmount = () => { + stopRefresh(); + + if (_keydownHandler) { + document.removeEventListener('keydown', _keydownHandler); + _keydownHandler = null; + } + + if (_modalEl) { + _modalEl.remove(); + _modalEl = null; + } + + const onClose = _onClose; + _onClose = null; + if (onClose) { + onClose(); + } + }; + + const mount = (container, { workerName, logData, onClose, fetchLog }) => { + _onClose = onClose; + _fetchLog = fetchLog; + + const modal = document.createElement('div'); + modal.id = 'distributed-log-modal'; + modal.className = 'log-modal'; + + const content = document.createElement('div'); + content.className = 'log-modal__content'; + + const header = document.createElement('div'); + header.className = 'log-modal__header'; + + const title = document.createElement('h3'); + title.className = 'log-modal__title'; + title.textContent = `${workerName} - Log Viewer`; + + const headerButtons = document.createElement('div'); + headerButtons.className = 'log-modal__header-buttons'; + + const refreshContainer = document.createElement('div'); + refreshContainer.className = 'log-modal__refresh'; + + const refreshCheckbox = document.createElement('input'); + refreshCheckbox.type = 'checkbox'; + refreshCheckbox.id = 'log-auto-refresh'; + refreshCheckbox.className = 'log-modal__refresh-input'; + refreshCheckbox.checked = true; + + const refreshLabel = document.createElement('label'); + refreshLabel.htmlFor = 'log-auto-refresh'; + refreshLabel.className = 'log-modal__refresh-label'; + refreshLabel.textContent = 'Auto-refresh'; + + refreshContainer.appendChild(refreshCheckbox); + refreshContainer.appendChild(refreshLabel); + + const closeBtn = document.createElement('button'); + closeBtn.className = 'distributed-button log-modal__close'; + closeBtn.textContent = '✕'; + + headerButtons.appendChild(refreshContainer); + headerButtons.appendChild(closeBtn); + + header.appendChild(title); + header.appendChild(headerButtons); + + const logContainer = document.createElement('div'); + logContainer.className = 'log-modal__body'; + logContainer.id = 'distributed-log-content'; + + const statusBar = document.createElement('div'); + statusBar.className = 'log-modal__status'; + + content.appendChild(header); + content.appendChild(logContainer); + content.appendChild(statusBar); + modal.appendChild(content); + + closeBtn.addEventListener('click', unmount); + modal.addEventListener('click', (event) => { + if (event.target === modal) { + unmount(); + } + }); + + _keydownHandler = (event) => { + if (event.key === 'Escape') { + unmount(); + } + }; + document.addEventListener('keydown', _keydownHandler); + + _modalEl = modal; + _logContentEl = logContainer; + _statusBarEl = statusBar; + _refreshCheckbox = refreshCheckbox; + + refreshCheckbox.addEventListener('change', () => { + if (refreshCheckbox.checked) { + refreshLog(); + } + }); + + container.appendChild(modal); + updateLogView(logData); + + requestAnimationFrame(() => { + if (_logContentEl) { + _logContentEl.scrollTop = _logContentEl.scrollHeight; + } + }); + + startRefresh(); + }; + + return { + mount, + unmount, + update: updateLogView, + }; +} diff --git a/web/workerLifecycle.js b/web/workerLifecycle.js index 1b54af4..875c353 100644 --- a/web/workerLifecycle.js +++ b/web/workerLifecycle.js @@ -1,6 +1,7 @@ -import { TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js'; +import { TIMEOUTS, STATUS_COLORS } from './constants.js'; import { buildWorkerUrl, normalizeWorkerUrl } from './urlUtils.js'; import { isRemoteWorker } from './workerSettings.js'; +import { applyProbeResultToWorkerDot } from './workerUtils.js'; export { normalizeWorkerUrl }; @@ -74,40 +75,45 @@ export async function checkAllWorkerStatuses(extension) { export async function checkMasterStatus(extension) { try { - const response = await fetch(`${window.location.origin}${ENDPOINTS.PROMPT}`, { - method: "GET", - signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK), - }); + const signal = extension.statusCheckAbortController?.signal || null; + const probeResult = await extension.api.probeWorker( + window.location.origin, + TIMEOUTS.STATUS_CHECK, + signal, + ); + if (!probeResult.ok) { + throw new Error(`HTTP ${probeResult.status}`); + } - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - const isProcessing = queueRemaining > 0; + const queueRemaining = probeResult.queueRemaining || 0; + const isProcessing = queueRemaining > 0; - // Update master status in state - extension.state.setMasterStatus(isProcessing ? "processing" : "online"); + // Update master status in state + extension.state.setMasterStatus(isProcessing ? "processing" : "online"); - // Update master status dot - const statusDot = document.getElementById("master-status"); - if (statusDot) { - if (!extension.isMasterParticipating()) { - if (isProcessing) { - setStatusDotClass(statusDot, "worker-status--processing"); - statusDot.title = `Orchestrating (${queueRemaining} in queue)`; - } else { - setStatusDotClass(statusDot, "worker-status--unknown"); - statusDot.title = "Master orchestrator only"; - } - } else if (isProcessing) { + // Update master status dot + const statusDot = document.getElementById("master-status"); + if (statusDot) { + if (!extension.isMasterParticipating()) { + if (isProcessing) { setStatusDotClass(statusDot, "worker-status--processing"); - statusDot.title = `Processing (${queueRemaining} in queue)`; + statusDot.title = `Orchestrating (${queueRemaining} in queue)`; } else { - setStatusDotClass(statusDot, "worker-status--online"); - statusDot.title = "Online"; + setStatusDotClass(statusDot, "worker-status--unknown"); + statusDot.title = "Master orchestrator only"; } + } else if (isProcessing) { + setStatusDotClass(statusDot, "worker-status--processing"); + statusDot.title = `Processing (${queueRemaining} in queue)`; + } else { + setStatusDotClass(statusDot, "worker-status--online"); + statusDot.title = "Online"; } } } catch (error) { + if (error?.name === "AbortError") { + return; + } // Master is always online (we're running on it), so keep it green const statusDot = document.getElementById("master-status"); if (statusDot) { @@ -127,47 +133,36 @@ export function getWorkerUrl(extension, worker, endpoint = "") { export async function checkWorkerStatus(extension, worker) { // Assume caller ensured enabled; proceed with check - const url = getWorkerUrl(extension, worker, ENDPOINTS.PROMPT); + const workerUrl = getWorkerUrl(extension, worker); try { - // Combine timeout with abort controller signal - const timeoutSignal = AbortSignal.timeout(TIMEOUTS.STATUS_CHECK); - const signal = extension.statusCheckAbortController - ? AbortSignal.any([timeoutSignal, extension.statusCheckAbortController.signal]) - : timeoutSignal; - - const response = await fetch(url, { - method: "GET", - mode: "cors", + const signal = extension.statusCheckAbortController?.signal || null; + const probeResult = await extension.api.probeWorker( + workerUrl, + TIMEOUTS.STATUS_CHECK, signal, + ); + if (!probeResult.ok) { + throw new Error(`HTTP ${probeResult.status}`); + } + + const queueRemaining = probeResult.queueRemaining || 0; + const isProcessing = queueRemaining > 0; + + // Update status + extension.state.setWorkerStatus(worker.id, { + online: true, + processing: isProcessing, + queueCount: queueRemaining, }); - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - const isProcessing = queueRemaining > 0; + // Update status dot based on probe result + applyProbeResultToWorkerDot(worker.id, probeResult); - // Update status - extension.state.setWorkerStatus(worker.id, { - online: true, - processing: isProcessing, - queueCount: queueRemaining, - }); - - // Update status dot based on processing state - if (isProcessing) { - extension.ui.updateStatusDot(worker.id, STATUS_COLORS.PROCESSING_YELLOW, `Online - Processing (${queueRemaining} in queue)`, false); - } else { - extension.ui.updateStatusDot(worker.id, STATUS_COLORS.ONLINE_GREEN, "Online - Idle", false); - } - - // Clear launching state since worker is now online - if (extension.state.isWorkerLaunching(worker.id)) { - extension.state.setWorkerLaunching(worker.id, false); - clearLaunchingFlag(extension, worker.id); - } - } else { - throw new Error(`HTTP ${response.status}`); + // Clear launching state since worker is now online + if (extension.state.isWorkerLaunching(worker.id)) { + extension.state.setWorkerLaunching(worker.id, false); + clearLaunchingFlag(extension, worker.id); } } catch (error) { // Don't process aborted requests @@ -186,8 +181,8 @@ export async function checkWorkerStatus(extension, worker) { if (extension.state.isWorkerLaunching(worker.id)) { extension.ui.updateStatusDot(worker.id, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true); } else if (worker.enabled) { - // Only update to red if not currently launching AND still enabled - extension.ui.updateStatusDot(worker.id, STATUS_COLORS.OFFLINE_RED, "Offline - Cannot connect", false); + // Only update to red if not currently launching AND still enabled. + applyProbeResultToWorkerDot(worker.id, { ok: false }); } // If disabled, don't update the dot (leave it gray) @@ -394,34 +389,34 @@ export function updateWorkerControls(extension, workerId) { // Show log button immediately if we have log file info (even if worker is still starting) if (managedInfo?.log_file && logBtn) { - logBtn.style.display = ""; + logBtn.classList.remove("is-hidden"); setButtonClass(logBtn, "btn--log"); } else if (logBtn && !managedInfo) { - logBtn.style.display = "none"; + logBtn.classList.add("is-hidden"); } if (status?.online || managedInfo) { // Worker is running or we just launched it - launchBtn.style.display = "none"; // Hide launch button when running + launchBtn.classList.add("is-hidden"); if (managedInfo) { // Only show stop button if we manage this worker - stopBtn.style.display = ""; + stopBtn.classList.remove("is-hidden"); stopBtn.disabled = false; stopBtn.textContent = "Stop"; setButtonClass(stopBtn, "btn--stop"); } else { // Hide stop button for workers launched outside UI - stopBtn.style.display = "none"; + stopBtn.classList.add("is-hidden"); } } else { // Worker is not running - launchBtn.style.display = ""; // Show launch button + launchBtn.classList.remove("is-hidden"); launchBtn.disabled = false; launchBtn.textContent = "Launch"; setButtonClass(launchBtn, "btn--launch"); - stopBtn.style.display = "none"; // Hide stop button when not running + stopBtn.classList.add("is-hidden"); } } @@ -545,21 +540,13 @@ export function toggleWorkerExpanded(extension, workerId) { extension.state.setWorkerExpanded(workerId, false); settingsDiv.classList.remove("expanded"); if (settingsArrow) { - settingsArrow.style.transform = "rotate(0deg)"; + settingsArrow.classList.remove("settings-arrow--expanded"); } - // Animate padding to 0 - settingsDiv.style.padding = "0 12px"; - settingsDiv.style.marginTop = "0"; - settingsDiv.style.marginBottom = "0"; } else { extension.state.setWorkerExpanded(workerId, true); settingsDiv.classList.add("expanded"); if (settingsArrow) { - settingsArrow.style.transform = "rotate(90deg)"; + settingsArrow.classList.add("settings-arrow--expanded"); } - // Animate padding to full - settingsDiv.style.padding = "12px"; - settingsDiv.style.marginTop = "8px"; - settingsDiv.style.marginBottom = "8px"; } } diff --git a/web/workerUtils.js b/web/workerUtils.js index 465c5e1..a7d9e90 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -92,3 +92,34 @@ export function findNodesByClass(apiPrompt, className) { .filter(([, nodeData]) => nodeData.class_type === className) .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); } + + +export function applyProbeResultToWorkerDot(workerId, probeResult) { + const dot = document.getElementById(`status-${workerId}`); + if (!dot) { + return; + } + + dot.classList.remove( + 'worker-status--online', + 'worker-status--offline', + 'worker-status--processing', + 'worker-status--unknown', + 'status-pulsing', + ); + + if (!probeResult || !probeResult.ok) { + dot.classList.add('worker-status--offline'); + dot.title = 'Offline - Cannot connect'; + return; + } + + if ((probeResult.queueRemaining || 0) > 0) { + dot.classList.add('worker-status--processing'); + dot.title = `Processing (${probeResult.queueRemaining} queued)`; + return; + } + + dot.classList.add('worker-status--online'); + dot.title = 'Online - Idle'; +} diff --git a/workers/detection.py b/workers/detection.py index 77aba9c..ec98731 100644 --- a/workers/detection.py +++ b/workers/detection.py @@ -1,16 +1,11 @@ -import asyncio import os -import socket -import json import platform import uuid -from multiprocessing import Queue import aiohttp from ..utils.network import normalize_host, get_client_session -from ..utils.logging import debug_log, log -from ..utils.config import load_config +from ..utils.logging import debug_log async def is_local_worker(worker_config): @@ -76,30 +71,3 @@ def is_runpod_environment(): """Check if running in Runpod environment.""" return (os.environ.get('RUNPOD_POD_ID') is not None or os.environ.get('RUNPOD_API_KEY') is not None) - -async def get_comms_channel(worker_id, worker_config): - """Get communication channel for a worker (HTTP URL or IPC Queue).""" - if await is_local_worker(worker_config): - comm_mode = worker_config.get('communicationMode', 'http') - - if comm_mode == 'ipc': - # Use multiprocessing Queue for IPC - from . import get_worker_manager - wm = get_worker_manager() - if worker_id not in wm.queues: - wm.queues[worker_id] = Queue() - return wm.queues[worker_id] - elif is_docker_environment(): - # Docker: use host.docker.internal - return f"http://host.docker.internal:{worker_config['port']}" - else: - # Local: use loopback - return f"http://127.0.0.1:{worker_config['port']}" - elif worker_config.get('type') == 'cloud' and is_runpod_environment(): - # Runpod same-host optimization (if detected) - host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' - return f"http://{host}:{worker_config['port']}" - else: - # Remote worker: use configured endpoint - host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' - return f"http://{host}:{worker_config['port']}"