From 3058b11ff2896ce180a55b9bdb88ebc521e87eb6 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 7 Mar 2026 05:22:09 +0000 Subject: [PATCH] Remove list splitter/collector feature and refactor distributed modules - Remove DistributedListSplitter and DistributedListCollector nodes - Remove /distributed/init_list_queue and /distributed/request_list_item endpoints - Clean up all references from prompt_transform, queue_orchestration, constants, and tests - Refactor and reorganize distributed modules for quality improvements --- .claude/skills/desloppify/SKILL.md | 151 +++++++ README.md | 6 +- __init__.py | 30 +- api/__init__.py | 26 +- api/config_routes.py | 160 ++++--- api/endpoint_policy.py | 25 ++ api/job_routes.py | 94 ++-- api/orchestration/__init__.py | 11 +- api/orchestration/dispatch.py | 48 +- api/orchestration/media_sync.py | 7 +- api/orchestration/prompt_transform.py | 103 +---- api/queue_orchestration.py | 48 +- api/queue_request.py | 35 +- api/request_guards.py | 15 + api/schemas.py | 76 +++- api/tunnel_routes.py | 62 +-- api/usdu_routes.py | 188 ++++---- api/worker_routes.py | 75 +++- bootstrap/__init__.py | 5 + bootstrap/entrypoint.py | 66 +++ conftest.py | 26 +- distributed.py | 62 +-- docs/comfyui-distributed-api.md | 8 +- nodes/__init__.py | 6 - nodes/branch_collector.py | 14 +- nodes/collector.py | 12 +- nodes/context_kwargs.py | 25 +- nodes/distributed_upscale.py | 45 +- nodes/hidden_inputs.py | 13 +- nodes/list_collector.py | 238 ---------- nodes/list_splitter.py | 220 --------- nodes/utilities.py | 55 +-- pyproject.toml | 5 +- tests/api/harness.py | 163 +++++++ tests/api/test_auth_guard_integration.py | 423 ++++++++++++++++++ tests/api/test_config_routes.py | 81 ++-- tests/api/test_job_routes.py | 137 +++--- tests/api/test_media_sync.py | 67 ++- tests/api/test_queue_orchestration.py | 66 ++- tests/api/test_tunnel_routes.py | 46 +- tests/api/test_usdu_routes.py | 114 ++--- tests/api/test_worker_routes.py | 108 ++--- tests/conftest.py | 2 - tests/unit/nodes/test_distributed_upscale.py | 9 + tests/unit/nodes/test_distributed_value.py | 22 + tests/unit/nodes/test_list_collector.py | 237 ---------- tests/unit/nodes/test_list_splitter.py | 261 ----------- tests/unit/system/test_config.py | 9 +- tests/unit/system/test_distributed.py | 153 ++----- tests/unit/system/test_logging.py | 54 +-- tests/unit/system/test_prompt_transform.py | 98 +--- tests/unit/system/test_queue_request.py | 31 +- tests/unit/upscale/test_dynamic.py | 31 +- tests/unit/upscale/test_job_store.py | 32 +- tests/unit/upscale/test_single_gpu.py | 1 + tests/unit/upscale/test_static.py | 25 +- upscale/job_state.py | 32 ++ upscale/job_store.py | 32 +- upscale/mode_contexts.py | 238 ++++++++++ upscale/modes/dynamic.py | 143 ++++-- upscale/modes/single_gpu.py | 20 +- upscale/modes/static.py | 133 ++++-- upscale/result_collector.py | 23 +- upscale/tile_ops.py | 57 +++ upscale/worker_comms.py | 182 +++++--- utils/async_helpers.py | 10 +- utils/auth.py | 36 ++ utils/cloudflare/process_reader.py | 5 +- utils/config.py | 45 +- utils/logging.py | 47 +- utils/network.py | 42 +- utils/parsing.py | 34 ++ utils/process.py | 22 +- utils/runtime_state.py | 39 ++ .../{usdu_managment.py => usdu_management.py} | 44 +- utils/usdu_utils.py | 18 +- utils/worker_ids.py | 96 +++- web/constants.js | 2 - web/executionUtils.js | 5 +- workers/__init__.py | 8 +- workers/process/lifecycle.py | 7 + workers/process/persistence.py | 4 +- workers/process_manager.py | 12 +- workers/startup.py | 2 +- workers/worker_monitor.py | 55 +-- 85 files changed, 2971 insertions(+), 2552 deletions(-) create mode 100644 .claude/skills/desloppify/SKILL.md create mode 100644 api/endpoint_policy.py create mode 100644 api/request_guards.py create mode 100644 bootstrap/__init__.py create mode 100644 bootstrap/entrypoint.py delete mode 100644 nodes/list_collector.py delete mode 100644 nodes/list_splitter.py create mode 100644 tests/api/harness.py create mode 100644 tests/api/test_auth_guard_integration.py delete mode 100644 tests/conftest.py delete mode 100644 tests/unit/nodes/test_list_collector.py delete mode 100644 tests/unit/nodes/test_list_splitter.py create mode 100644 upscale/mode_contexts.py create mode 100644 utils/auth.py create mode 100644 utils/parsing.py create mode 100644 utils/runtime_state.py rename utils/{usdu_managment.py => usdu_management.py} (58%) diff --git a/.claude/skills/desloppify/SKILL.md b/.claude/skills/desloppify/SKILL.md new file mode 100644 index 0000000..1828136 --- /dev/null +++ b/.claude/skills/desloppify/SKILL.md @@ -0,0 +1,151 @@ + + +--- +name: desloppify +description: > + Codebase health scanner and technical debt tracker. Use when the user asks + about code quality, technical debt, dead code, large files, god classes, + duplicate functions, code smells, naming issues, import cycles, or coupling + problems. Also use when asked for a health score, what to fix next, or to + create a cleanup plan. Supports 28 languages. +allowed-tools: Bash(desloppify *) +--- + +# Desloppify + +## 1. Your Job + +Improve code quality by maximising the **strict score** honestly. + +**The main thing you do is run `desloppify next`** — it tells you exactly what to fix and how. Fix it, resolve it, run `next` again. Keep going. + +Follow the scan output's **INSTRUCTIONS FOR AGENTS** — don't substitute your own analysis. + +## 2. The Workflow + +Two loops. The **outer loop** rescans periodically to measure progress. +The **inner loop** is where you spend most of your time: fixing issues one by one. + +### Outer loop — scan and check + +```bash +desloppify scan --path . # analyse the codebase +desloppify status # check scores — are we at target? +``` +If not at target, work the inner loop. Rescan periodically — especially after clearing a cluster or batch of related fixes. Issues cascade-resolve and new ones may surface. + +### Inner loop — fix issues + +Repeat until the queue is clear: + +``` +1. desloppify next ← tells you exactly what to fix next +2. Fix the issue in code +3. Resolve it (next shows you the exact command including required attestation) +``` + +Score may temporarily drop after fixes — cascade effects are normal, keep going. +If `next` suggests an auto-fixer, run `desloppify fix --dry-run` to preview, then apply. + +**To be strategic**, use `plan` to shape what `next` gives you: +```bash +desloppify plan # see the full ordered queue +desloppify plan move top # reorder — what unblocks the most? +desloppify plan cluster create # group related issues to batch-fix +desloppify plan focus # scope next to one cluster +desloppify plan defer # push low-value items aside +desloppify plan skip # hide from next +desloppify plan done # mark complete +desloppify plan reopen # reopen +``` + +### Subjective reviews + +The scan will prompt you when a subjective review is needed — just follow its instructions. +If you need to trigger one manually: +```bash +desloppify review --run-batches --runner codex --parallel --scan-after-import +``` + +### Other useful commands + +```bash +desloppify next --count 5 # top 5 priorities +desloppify next --cluster # drill into a cluster +desloppify show # filter by file/detector/ID +desloppify show --status open # all open findings +desloppify plan skip --permanent "" --note "reason" # accept debt (lowers strict score) +desloppify scan --path . --reset-subjective # reset subjective baseline to 0 +``` + +## 3. Reference + +### How scoring works + +Overall score = **40% mechanical** + **60% subjective**. + +- **Mechanical (40%)**: auto-detected issues — duplication, dead code, smells, unused imports, security. Fixed by changing code and rescanning. +- **Subjective (60%)**: design quality review — naming, error handling, abstractions, clarity. Starts at **0%** until reviewed. The scan will prompt you when a review is needed. +- **Strict score** is the north star: wontfix items count as open. The gap between overall and strict is your wontfix debt. +- **Score types**: overall (lenient), strict (wontfix counts), objective (mechanical only), verified (confirmed fixes only). + +### Subjective reviews in detail + +- **Preferred**: `desloppify review --run-batches --runner codex --parallel --scan-after-import` — does everything in one command. +- **Manual path**: `desloppify review --prepare` → review per dimension → `desloppify review --import file.json`. +- Import first, fix after — import creates tracked state entries for correlation. +- Integrity: reviewers score from evidence only. Scores hitting exact targets trigger auto-reset. +- Even moderate scores (60-80) dramatically improve overall health. +- Stale dimensions auto-surface in `next` — just follow the queue. + +### Key concepts + +- **Tiers**: T1 auto-fix → T2 quick manual → T3 judgment call → T4 major refactor. +- **Auto-clusters**: related findings are auto-grouped in `next`. Drill in with `next --cluster `. +- **Zones**: production/script (scored), test/config/generated/vendor (not scored). Fix with `zone set`. +- **Wontfix cost**: widens the lenient↔strict gap. Challenge past decisions when the gap grows. +- Score can temporarily drop after fixes (cascade effects are normal). + +## 4. Escalate Tool Issues Upstream + +When desloppify itself appears wrong or inconsistent: + +1. Capture a minimal repro (`command`, `path`, `expected`, `actual`). +2. Open a GitHub issue in `peteromallet/desloppify`. +3. If you can fix it safely, open a PR linked to that issue. +4. If unsure whether it is tool bug vs user workflow, issue first, PR second. + +## Prerequisite + +`command -v desloppify >/dev/null 2>&1 && echo "desloppify: installed" || echo "NOT INSTALLED — run: pip install --upgrade git+https://github.com/peteromallet/desloppify.git"` + + + +## Claude Code Overlay + +Use Claude subagents for subjective scoring work that should be context-isolated. + +### Subjective review + +1. **Preferred**: `desloppify review --run-batches --runner codex --parallel --scan-after-import` — does everything in one command. +2. **Claude cloud path**: `desloppify review --external-start --external-runner claude` → use generated `claude_launch_prompt.md` + `review_result.template.json` → run printed `desloppify review --external-submit --session-id --import `. +3. **Manual path**: split dimensions across N subagents (one message, multiple Task calls), merge outputs, then `desloppify review --import findings.json`. + +For the manual path: +- Read `dimension_prompts` from `query.json` for dimension definitions. +- Give each agent the codebase path, dimensions, and output format. Let agents decide what to read. +- Each agent writes output to a separate file. Merge assessments (average overlaps) and findings. +- Import first, fix after — import creates tracked state for correlation. + +### Subagent rules + +1. Prefer delegating review tasks to a project subagent in `.claude/agents/`. +2. Set `context: fork` so prior chat context does not leak into scoring. +3. For blind reviews, consume `.desloppify/review_packet_blind.json` instead of full `query.json`. +4. Score from evidence only; do not anchor to target thresholds. When mixed, score lower. +5. Return machine-readable JSON matching the format in the base skill doc. For `--external-submit`, include `session` from the generated template. +6. `findings` MUST match `query.system_prompt` exactly. Use `"findings": []` when no defects found. +7. Import is fail-closed: invalid findings abort unless `--allow-partial` is passed. + + + diff --git a/README.md b/README.md index 6a3b0f9..6762b86 100644 --- a/README.md +++ b/README.md @@ -157,9 +157,9 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work Control your distributed cluster programmatically without opening the browser. -* **Endpoint:** `POST /distributed/queue` -* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID. -* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md) +* **Endpoint:** `POST /distributed/queue` +* **Functionality:** Accepts a distributed queue payload (`prompt`, `client_id`, `enabled_worker_ids`, optional `workflow`/`delegate_master`/`trace_execution_id`), dispatches to healthy workers, and returns `{prompt_id, worker_count}`. +* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md) > **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN). diff --git a/__init__.py b/__init__.py index 2029bdc..5df260c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,29 +1,11 @@ -# Import everything needed from the main module -from .distributed import ( - NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS -) +"""ComfyUI-Distributed package entrypoint.""" +from __future__ import annotations -# Import utilities -from .utils.config import ensure_config_exists, CONFIG_FILE -from .utils.logging import debug_log - -# Import distributed upscale nodes -from .nodes.distributed_upscale import ( - NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS -) +from .bootstrap.entrypoint import build_node_mappings, initialize_runtime WEB_DIRECTORY = "./web" -ensure_config_exists() +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings() +initialize_runtime() -# Merge node mappings -NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS} -NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS} - -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] - -debug_log("Loaded Distributed nodes.") -debug_log(f"Config file: {CONFIG_FILE}") -debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}") +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/api/__init__.py b/api/__init__.py index d20f665..a2f2d2c 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -1,5 +1,21 @@ -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 +"""Distributed API package with explicit route bootstrap.""" +from __future__ import annotations + +from importlib import import_module + +_ROUTE_MODULES = ( + "config_routes", + "tunnel_routes", + "worker_routes", + "job_routes", + "usdu_routes", +) + + +def bootstrap_routes() -> None: + """Import route modules once so aiohttp decorators register endpoints.""" + for module_name in _ROUTE_MODULES: + import_module(f"{__name__}.{module_name}") + + +__all__ = ["bootstrap_routes"] diff --git a/api/config_routes.py b/api/config_routes.py index 677245c..dc40291 100644 --- a/api/config_routes.py +++ b/api/config_routes.py @@ -1,31 +1,12 @@ from __future__ import annotations -import json -from contextlib import asynccontextmanager - from aiohttp import web import server -try: - from ..utils.config import config_transaction, load_config, save_config -except ImportError: - from ..utils.config import load_config - - try: - from ..utils.config import save_config - except ImportError: - def save_config(_config): - return True - - @asynccontextmanager - async def config_transaction(): - config = load_config() - original_snapshot = json.dumps(config, sort_keys=True) - yield config - if json.dumps(config, sort_keys=True) != original_snapshot: - save_config(config) -from ..utils.logging import debug_log, log +from ..utils.config import config_transaction, load_config +from ..utils.logging import debug_log from ..utils.network import handle_api_error, normalize_host +from .endpoint_policy import run_authorized_endpoint def _positive_int(value: int) -> bool: @@ -88,68 +69,103 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None: target[key] = normalizer(value) if (normalizer and value is not None) else value +def _validate_workers_payload(workers_value: object) -> list[str]: + """Validate top-level workers payload shape before persisting config.""" + if not isinstance(workers_value, list): + return ["workers: expected list"] + + errors: list[str] = [] + seen_ids: set[str] = set() + for index, worker in enumerate(workers_value): + if not isinstance(worker, dict): + errors.append(f"workers[{index}]: expected object") + continue + + worker_id = str(worker.get("id", "")).strip() + if not worker_id: + errors.append(f"workers[{index}].id: expected non-empty string") + elif worker_id in seen_ids: + errors.append(f"workers[{index}].id: duplicate id '{worker_id}'") + else: + seen_ids.add(worker_id) + + if "port" in worker and worker.get("port") is not None: + try: + port = int(worker.get("port")) + except (TypeError, ValueError): + errors.append(f"workers[{index}].port: expected integer") + else: + if port <= 0: + errors.append(f"workers[{index}].port: expected positive integer") + + return errors + + @server.PromptServer.instance.routes.get("/distributed/config") async def get_config_endpoint(request: web.Request) -> web.StreamResponse: - config = load_config() - return web.json_response(config) + async def _operation() -> web.StreamResponse: + return web.json_response(load_config()) + + return await run_authorized_endpoint(request, _operation) @server.PromptServer.instance.routes.post("/distributed/config") async def update_config_endpoint(request: web.Request) -> web.StreamResponse: """Bulk config update with schema validation.""" - try: - data = await request.json() - except Exception as e: - return await handle_api_error(request, f"Invalid JSON payload: {e}", 400) + async def _operation() -> web.StreamResponse: + try: + data = await request.json() + except Exception as exc: + return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400) - if not isinstance(data, dict): - return await handle_api_error(request, "Config payload must be an object", 400) + if not isinstance(data, dict): + return await handle_api_error(request, "Config payload must be an object", 400) - validated_settings = {} - validated_root = {} - errors = [] + validated_settings = {} + validated_root = {} + errors = [] - for key, value in data.items(): - if key not in CONFIG_SCHEMA: - errors.append(f"Unknown field: {key}") - continue + for key, value in data.items(): + if key not in CONFIG_SCHEMA: + errors.append(f"Unknown field: {key}") + continue - expected_type, validator = CONFIG_SCHEMA[key] - if not isinstance(value, expected_type): - errors.append(f"{key}: expected {expected_type.__name__}") - continue + expected_type, validator = CONFIG_SCHEMA[key] + if not isinstance(value, expected_type): + errors.append(f"{key}: expected {expected_type.__name__}") + continue - if validator and not validator(value): - errors.append(f"{key}: value {value!r} failed validation") - continue + if validator and not validator(value): + errors.append(f"{key}: value {value!r} failed validation") + continue - if key in _SETTINGS_FIELDS: - validated_settings[key] = value - else: - validated_root[key] = value + if key == "workers": + errors.extend(_validate_workers_payload(value)) + if errors: + continue - if errors: - return web.json_response({ - "status": "error", - "error": errors, - "message": "; ".join(errors), - }, status=400) + if key in _SETTINGS_FIELDS: + validated_settings[key] = value + else: + validated_root[key] = value + + if errors: + return await handle_api_error(request, errors, 400) - try: async with config_transaction() as config: settings = config.setdefault("settings", {}) settings.update(validated_settings) for key, value in validated_root.items(): config[key] = value return web.json_response({"status": "success", "config": config}) - except Exception as e: - return await handle_api_error(request, e) + + return await run_authorized_endpoint(request, _operation) @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}") async def queue_status_endpoint(request: web.Request) -> web.StreamResponse: """Check if a job queue is initialized.""" - try: + async def _operation() -> web.StreamResponse: job_id = request.match_info['job_id'] # Import to ensure initialization @@ -161,12 +177,12 @@ async def queue_status_endpoint(request: web.Request) -> web.StreamResponse: debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}") return web.json_response({"exists": exists, "job_id": job_id}) - except Exception as e: - return await handle_api_error(request, e, 500) + + return await run_authorized_endpoint(request, _operation, unexpected_status=500) @server.PromptServer.instance.routes.post("/distributed/config/update_worker") async def update_worker_endpoint(request: web.Request) -> web.StreamResponse: - try: + async def _operation() -> web.StreamResponse: data = await request.json() worker_id = data.get("worker_id") @@ -205,12 +221,12 @@ async def update_worker_endpoint(request: web.Request) -> web.StreamResponse: ) return web.json_response({"status": "success"}) - except Exception as e: - return await handle_api_error(request, e, 400) + + return await run_authorized_endpoint(request, _operation, unexpected_status=500) @server.PromptServer.instance.routes.post("/distributed/config/delete_worker") async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse: - try: + async def _operation() -> web.StreamResponse: data = await request.json() worker_id = data.get("worker_id") @@ -237,13 +253,13 @@ async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse: "status": "success", "message": f"Worker {removed_worker.get('name', worker_id)} deleted" }) - except Exception as e: - return await handle_api_error(request, e, 400) + + return await run_authorized_endpoint(request, _operation, unexpected_status=500) @server.PromptServer.instance.routes.post("/distributed/config/update_setting") async def update_setting_endpoint(request: web.Request) -> web.StreamResponse: """Updates a specific key in the settings object.""" - try: + async def _operation() -> web.StreamResponse: data = await request.json() key = data.get("key") value = data.get("value") @@ -260,13 +276,13 @@ async def update_setting_endpoint(request: web.Request) -> web.StreamResponse: config['settings'][key] = value return web.json_response({"status": "success", "message": f"Setting '{key}' updated."}) - except Exception as e: - return await handle_api_error(request, e, 400) + + return await run_authorized_endpoint(request, _operation, unexpected_status=500) @server.PromptServer.instance.routes.post("/distributed/config/update_master") async def update_master_endpoint(request: web.Request) -> web.StreamResponse: """Updates master configuration.""" - try: + async def _operation() -> web.StreamResponse: data = await request.json() async with config_transaction() as config: @@ -275,5 +291,5 @@ async def update_master_endpoint(request: web.Request) -> web.StreamResponse: _apply_field_patch(config['master'], data, _MASTER_FIELDS) return web.json_response({"status": "success", "message": "Master configuration updated."}) - except Exception as e: - return await handle_api_error(request, e, 400) + + return await run_authorized_endpoint(request, _operation, unexpected_status=500) diff --git a/api/endpoint_policy.py b/api/endpoint_policy.py new file mode 100644 index 0000000..5802a31 --- /dev/null +++ b/api/endpoint_policy.py @@ -0,0 +1,25 @@ +"""Shared endpoint policy helpers for distributed API routes.""" +from __future__ import annotations + +from collections.abc import Awaitable, Callable + +from aiohttp import web + +from ..utils.network import handle_api_error +from .request_guards import authorization_error_or_none + + +async def run_authorized_endpoint( + request: web.Request, + operation: Callable[[], Awaitable[web.StreamResponse]], + *, + unexpected_status: int = 500, +) -> web.StreamResponse: + """Run endpoint operation with shared auth guard and fallback error mapping.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error + try: + return await operation() + except Exception as exc: + return await handle_api_error(request, exc, unexpected_status) diff --git a/api/job_routes.py b/api/job_routes.py index 78bd753..d06a100 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -18,25 +18,20 @@ from ..utils.logging import debug_log from ..utils.image import pil_to_tensor, ensure_contiguous from ..utils.network import handle_api_error from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY -try: - from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution -except ImportError: - from .queue_orchestration import orchestrate_distributed_execution - - def ensure_distributed_state() -> None: - prompt_server_instance = server.PromptServer.instance - if not hasattr(prompt_server_instance, "distributed_jobs_lock"): - prompt_server_instance.distributed_jobs_lock = asyncio.Lock() - if not hasattr(prompt_server_instance, "distributed_pending_jobs"): - prompt_server_instance.distributed_pending_jobs = {} +from ..utils.runtime_state import ensure_distributed_runtime_state +from .request_guards import authorization_error_or_none +from .schemas import require_bool_literal +from .queue_orchestration import orchestrate_distributed_execution from .queue_request import parse_queue_request_payload -prompt_server = server.PromptServer.instance - # Canonical worker result envelope accepted by POST /distributed/job_complete: # { "job_id": str, "worker_id": str, "batch_idx": int, "image": , "is_last": bool } +def _runtime_state(): + return ensure_distributed_runtime_state() + + def _decode_image_sync(image_path: str) -> dict[str, Any]: """Decode image/video file and compute hash in a threadpool worker.""" import base64 @@ -148,16 +143,19 @@ def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]: @server.PromptServer.instance.routes.post("/distributed/prepare_job") async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse: + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() multi_job_id = data.get('multi_job_id') if not multi_job_id: return await handle_api_error(request, "Missing multi_job_id", 400) - ensure_distributed_state() - async with prompt_server.distributed_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_jobs: - prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() + runtime_state = _runtime_state() + async with runtime_state.distributed_jobs_lock: + if multi_job_id not in runtime_state.distributed_pending_jobs: + runtime_state.distributed_pending_jobs[multi_job_id] = asyncio.Queue() debug_log(f"Prepared queue for job {multi_job_id}") return web.json_response({"status": "success"}) @@ -167,7 +165,11 @@ async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/clear_memory") async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse: debug_log("Received request to clear VRAM.") + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: + warnings: list[str] = [] # Use ComfyUI's prompt server queue system like the /free endpoint does if hasattr(server.PromptServer.instance, 'prompt_queue'): server.PromptServer.instance.prompt_queue.set_flag("unload_models", True) @@ -184,12 +186,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse: try: mm.unload_all_models() except AttributeError as e: - debug_log(f"Warning during model unload: {e}") + warning = f"Model unload warning: {e}" + warnings.append(warning) + debug_log(warning) try: mm.soft_empty_cache() except Exception as e: - debug_log(f"Warning during cache clear: {e}") + warning = f"Cache clear warning: {e}" + warnings.append(warning) + debug_log(warning) for _ in range(3): gc.collect() @@ -198,6 +204,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse: torch.cuda.empty_cache() torch.cuda.ipc_collect() + if warnings: + debug_log("VRAM cleared with warnings.") + return web.json_response( + { + "status": "partial", + "message": "GPU memory cleared with warnings.", + "warnings": warnings, + } + ) + debug_log("VRAM cleared successfully.") return web.json_response({"status": "success", "message": "GPU memory cleared."}) except Exception as e: @@ -206,13 +222,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() - debug_log(f"Partial VRAM clear completed with warning: {e}") - return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"}) + debug_log(f"VRAM clear failed: {e}") + return await handle_api_error(request, f"GPU memory clear failed: {e}", 500) @server.PromptServer.instance.routes.post("/distributed/queue") async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse: """Queue a distributed workflow, mirroring the UI orchestration pipeline.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: raw_payload = await request.json() except Exception as exc: @@ -235,7 +254,6 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse return web.json_response({ "prompt_id": prompt_id, "worker_count": worker_count, - "auto_prepare_supported": True, }) except Exception as exc: return await handle_api_error(request, exc, 500) @@ -243,6 +261,9 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse @server.PromptServer.instance.routes.post("/distributed/load_image") async def load_image_endpoint(request: web.Request) -> web.StreamResponse: """Load an image or video file and return it as base64 data with hash.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() image_path = data.get("image_path") @@ -260,6 +281,9 @@ async def load_image_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/check_file") async def check_file_endpoint(request: web.Request) -> web.StreamResponse: """Check if a file exists and matches the given hash.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() filename = data.get("filename") @@ -277,6 +301,9 @@ async def check_file_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/job_complete") async def job_complete_endpoint(request: web.Request) -> web.StreamResponse: + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() except Exception as exc: @@ -291,7 +318,7 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse: batch_idx = data.get("batch_idx") image_payload = data.get("image") audio_payload = data.get("audio") - is_last = data.get("is_last") + is_last_raw = data.get("is_last") errors = [] if not isinstance(job_id, str) or not job_id.strip(): @@ -304,8 +331,12 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse: errors.append("image: expected non-empty base64 PNG string") if audio_payload is not None and not isinstance(audio_payload, dict): errors.append("audio: expected object when provided") - if not isinstance(is_last, bool): - errors.append("is_last: expected boolean") + + try: + is_last = require_bool_literal(is_last_raw, field_name="is_last") + except ValueError as exc: + errors.append(str(exc)) + is_last = False if errors: return await handle_api_error(request, errors, 400) @@ -314,12 +345,21 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse: multi_job_id = job_id.strip() worker_id = worker_id.strip() + runtime_state = _runtime_state() + allowed_workers = runtime_state.distributed_job_allowed_workers.get(multi_job_id) + if allowed_workers is not None and worker_id not in allowed_workers: + return await handle_api_error( + request, + f"Unauthorized worker_id for job {multi_job_id}", + 403, + ) + pending = None queue_size = 0 deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD) while pending is None: - async with prompt_server.distributed_jobs_lock: - pending = prompt_server.distributed_pending_jobs.get(multi_job_id) + async with runtime_state.distributed_jobs_lock: + pending = runtime_state.distributed_pending_jobs.get(multi_job_id) if pending is not None: queue_item = { "tensor": tensor, diff --git a/api/orchestration/__init__.py b/api/orchestration/__init__.py index 589b461..23a466b 100644 --- a/api/orchestration/__init__.py +++ b/api/orchestration/__init__.py @@ -1,4 +1,7 @@ -# Orchestration helpers split out from queue_orchestration.py: -# - prompt_transform.py: graph pruning + hidden input overrides -# - media_sync.py: remote media/path normalization -# - dispatch.py: worker probe + prompt dispatch +"""Orchestration helpers used by distributed queue execution.""" + +__all__ = [ + "dispatch", + "media_sync", + "prompt_transform", +] diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 34c0f16..9a2b596 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -5,26 +5,12 @@ from typing import Any import aiohttp +from ...utils.auth import distributed_auth_headers +from ...utils.config import load_config from ...utils.logging import debug_log, log from ...utils.network import build_worker_url, get_client_session, probe_worker -try: - from ...utils.trace_logger import trace_debug, trace_info -except ImportError: - def _trace_fallback(_execution_id: str, message: str) -> None: - debug_log(message) - - trace_debug = _trace_fallback - trace_info = _trace_fallback - -try: - from ..schemas import parse_positive_int -except ImportError: - def parse_positive_int(value, default): - try: - parsed = int(value) - return parsed if parsed > 0 else default - except (TypeError, ValueError): - return default +from ...utils.trace_logger import trace_debug, trace_info +from ..schemas import coerce_positive_int _least_busy_rr_index = 0 @@ -39,8 +25,9 @@ async def worker_ws_is_active(worker: dict[str, Any]) -> bool: """Ping worker's websocket endpoint to confirm it's reachable.""" session = await get_client_session() url = build_worker_url(worker, "/distributed/worker_ws") + headers = distributed_auth_headers(load_config()) try: - ws = await session.ws_connect(url, heartbeat=20, timeout=3) + ws = await session.ws_connect(url, heartbeat=20, timeout=3, headers=headers) await ws.close() return True except asyncio.TimeoutError: @@ -60,7 +47,13 @@ async def _probe_worker_active(worker, use_websocket, semaphore): return worker, is_active -async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): +async def _dispatch_via_websocket( + worker_url, + payload, + client_id, + headers: dict[str, str], + timeout=60.0, +): """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close.""" request_id = uuid.uuid4().hex ws_payload = { @@ -74,7 +67,12 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): ws_url = f"{ws_url}/distributed/worker_ws" session = await get_client_session() - async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws: + async with session.ws_connect( + ws_url, + heartbeat=20, + timeout=timeout, + headers=headers, + ) as ws: await ws.send_json(ws_payload) async for msg in ws: if msg.type == aiohttp.WSMsgType.TEXT: @@ -115,6 +113,7 @@ async def dispatch_worker_prompt( payload["extra_data"] = extra_data if use_websocket: + ws_headers = distributed_auth_headers(load_config()) try: await _dispatch_via_websocket( worker_url, @@ -123,6 +122,7 @@ async def dispatch_worker_prompt( "workflow": workflow_meta, }, client_id, + ws_headers, ) return except Exception as exc: @@ -150,7 +150,7 @@ async def select_active_workers( probe_concurrency: int = 8, ) -> tuple[list[dict[str, Any]], bool]: """Probe workers and return (active_workers, updated_delegate_master).""" - probe_limit = parse_positive_int(probe_concurrency, 8) + probe_limit = coerce_positive_int(probe_concurrency, 8) probe_semaphore = asyncio.Semaphore(probe_limit) if trace_execution_id and workers: @@ -233,7 +233,7 @@ async def select_least_busy_worker( if not workers: return None - probe_limit = parse_positive_int(probe_concurrency, 8) + probe_limit = coerce_positive_int(probe_concurrency, 8) probe_semaphore = asyncio.Semaphore(probe_limit) statuses = await asyncio.gather( *[ @@ -279,7 +279,7 @@ async def rank_workers_by_load( if not workers: return [] - probe_limit = parse_positive_int(probe_concurrency, 8) + probe_limit = coerce_positive_int(probe_concurrency, 8) probe_semaphore = asyncio.Semaphore(probe_limit) statuses = await asyncio.gather( *[ diff --git a/api/orchestration/media_sync.py b/api/orchestration/media_sync.py index fc61c6e..27bd22d 100644 --- a/api/orchestration/media_sync.py +++ b/api/orchestration/media_sync.py @@ -7,6 +7,8 @@ from typing import Any import aiohttp +from ...utils.auth import distributed_auth_headers +from ...utils.config import load_config 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 @@ -132,8 +134,9 @@ async def fetch_worker_path_separator( """Best-effort fetch of a worker's path separator from /distributed/system_info.""" url = build_worker_url(worker, "/distributed/system_info") session = await get_client_session() + headers = distributed_auth_headers(load_config()) try: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp: + async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp: if resp.status != 200: return None payload = await resp.json() @@ -150,6 +153,7 @@ async def fetch_worker_path_separator( async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None): """Upload one media file to worker iff missing or hash-mismatched.""" session = await get_client_session() + headers = distributed_auth_headers(load_config()) normalized = filename.replace("\\", "/") check_url = build_worker_url(worker, "/distributed/check_file") @@ -157,6 +161,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_ async with session.post( check_url, json={"filename": normalized, "hash": file_hash}, + headers=headers, timeout=aiohttp.ClientTimeout(total=6), ) as resp: if resp.status == 200: diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 0111585..673f7f8 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -332,11 +332,10 @@ def prune_prompt_for_branch_worker( def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]: """Prune worker prompt to distributed nodes and their upstream dependencies.""" collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector") - list_collector_ids = find_nodes_by_class(prompt_obj, "DistributedListCollector") branch_collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector") upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed") branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch") - distributed_ids = collector_ids + list_collector_ids + branch_collector_ids + upscale_ids + branch_ids + distributed_ids = collector_ids + branch_collector_ids + upscale_ids + branch_ids if not distributed_ids: return prompt_obj @@ -439,8 +438,6 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str job_map = {} distributed_nodes = ( prompt_index.nodes_for_class("DistributedCollector") - + prompt_index.nodes_for_class("DistributedListSplitter") - + prompt_index.nodes_for_class("DistributedListCollector") + prompt_index.nodes_for_class("DistributedBranchCollector") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") + prompt_index.nodes_for_class("DistributedBranch") @@ -450,7 +447,7 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str return job_map -def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map): +def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id): """Configure DistributedSeed nodes for master or worker role.""" for node_id in prompt_index.nodes_for_class("DistributedSeed"): node = prompt_copy.get(node_id) @@ -461,7 +458,7 @@ def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, w if is_master: inputs["worker_id"] = "" else: - inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}" + inputs["worker_id"] = str(participant_id) def _override_collector_nodes( @@ -524,7 +521,7 @@ def _override_upscale_nodes( inputs["worker_id"] = participant_id -def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map): +def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id): """Configure DistributedValue nodes for master or worker role.""" for node_id in prompt_index.nodes_for_class("DistributedValue"): node = prompt_copy.get(node_id) @@ -535,74 +532,7 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, if is_master: inputs["worker_id"] = "" else: - inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}" - - -def _override_list_splitter_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - enabled_worker_ids, - job_id_map, - master_url, - delegate_master, -): - """Configure participant chunk mapping for DistributedListSplitter nodes.""" - participants = _resolve_participants(enabled_worker_ids, delegate_master) - total_participants = max(1, len(participants)) - participant_id = str(participant_id) - - if participant_id in participants: - participant_index = participants.index(participant_id) - else: - participant_index = 0 - - for node_id in prompt_index.nodes_for_class("DistributedListSplitter"): - node = prompt_copy.get(node_id) - if not isinstance(node, dict): - continue - inputs = node.setdefault("inputs", {}) - inputs["participant_index"] = participant_index - inputs["total_participants"] = total_participants - inputs["multi_job_id"] = job_id_map.get(node_id, node_id) - inputs["is_worker"] = not is_master - if is_master: - inputs.pop("master_url", None) - inputs["worker_id"] = "master" - else: - inputs["master_url"] = master_url - inputs["worker_id"] = participant_id - - -def _override_list_collector_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, - delegate_master, -): - """Configure DistributedListCollector nodes for master/worker role.""" - for node_id in prompt_index.nodes_for_class("DistributedListCollector"): - node = prompt_copy.get(node_id) - if not isinstance(node, dict): - continue - - inputs = node.setdefault("inputs", {}) - inputs["multi_job_id"] = job_id_map.get(node_id, node_id) - inputs["is_worker"] = not is_master - inputs["enabled_worker_ids"] = enabled_json - if is_master: - inputs["delegate_only"] = bool(delegate_master) - inputs.pop("master_url", None) - inputs.pop("worker_id", None) - else: - inputs["master_url"] = master_url inputs["worker_id"] = str(participant_id) - inputs["delegate_only"] = False def _override_branch_nodes( @@ -740,21 +670,10 @@ def apply_participant_overrides( ) -> dict[str, Any]: """Return a prompt copy with hidden inputs configured for master/worker.""" is_master = participant_id == "master" - worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)} enabled_json = json.dumps(enabled_worker_ids) - _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map) - _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map) - _override_list_splitter_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - enabled_worker_ids, - job_id_map, - master_url, - delegate_master, - ) + _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id) + _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id) _override_collector_nodes( prompt_copy, prompt_index, @@ -765,16 +684,6 @@ def apply_participant_overrides( enabled_json, delegate_master, ) - _override_list_collector_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, - delegate_master, - ) _override_upscale_nodes( prompt_copy, prompt_index, diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 52c2ec0..9f8a7b2 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -3,8 +3,6 @@ import time import uuid from typing import Any -import server - from ..utils.async_helpers import queue_prompt_payload from ..utils.config import load_config from ..utils.constants import ( @@ -15,8 +13,9 @@ from ..utils.constants import ( ) from ..utils.logging import debug_log, log from ..utils.network import build_master_url +from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance from ..utils.trace_logger import trace_debug -from .schemas import parse_positive_float, parse_positive_int +from .schemas import coerce_positive_float, coerce_positive_int from .orchestration.dispatch import ( dispatch_worker_prompt, rank_workers_by_load, @@ -33,33 +32,21 @@ from .orchestration.prompt_transform import ( prune_prompt_for_worker, ) - -prompt_server = server.PromptServer.instance - - def _generate_execution_trace_id(): return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" def ensure_distributed_state(server_instance=None): """Ensure prompt_server has the state used by distributed queue orchestration.""" - ps = server_instance or prompt_server - if not hasattr(ps, "distributed_pending_jobs"): - ps.distributed_pending_jobs = {} - if not hasattr(ps, "distributed_jobs_lock"): - ps.distributed_jobs_lock = asyncio.Lock() - - -# Initialize top-level distributed queue state at module import time. -ensure_distributed_state() + ensure_distributed_runtime_state(server_instance) async def _ensure_distributed_queue(job_id): """Ensure a queue exists for the given distributed job ID.""" - ensure_distributed_state() - async with prompt_server.distributed_jobs_lock: - if job_id not in prompt_server.distributed_pending_jobs: - prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue() + runtime_state = ensure_distributed_runtime_state() + async with runtime_state.distributed_jobs_lock: + if job_id not in runtime_state.distributed_pending_jobs: + runtime_state.distributed_pending_jobs[job_id] = asyncio.Queue() def _resolve_enabled_workers(config, requested_ids=None): @@ -98,19 +85,19 @@ def _resolve_enabled_workers(config, requested_ids=None): 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 = coerce_positive_int( settings.get("worker_probe_concurrency"), ORCHESTRATION_WORKER_PROBE_CONCURRENCY, ) - worker_prep_concurrency = parse_positive_int( + worker_prep_concurrency = coerce_positive_int( settings.get("worker_prep_concurrency"), ORCHESTRATION_WORKER_PREP_CONCURRENCY, ) - media_sync_concurrency = parse_positive_int( + media_sync_concurrency = coerce_positive_int( settings.get("media_sync_concurrency"), ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, ) - media_sync_timeout_seconds = parse_positive_float( + media_sync_timeout_seconds = coerce_positive_float( settings.get("media_sync_timeout_seconds"), ORCHESTRATION_MEDIA_SYNC_TIMEOUT, ) @@ -272,13 +259,20 @@ async def _select_execution_workers( def _distributed_queue_nodes(prompt_index): return ( prompt_index.nodes_for_class("DistributedCollector") - + prompt_index.nodes_for_class("DistributedListCollector") + prompt_index.nodes_for_class("DistributedBranchCollector") + prompt_index.nodes_for_class("DistributedBranch") + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed") ) +def _register_job_allowed_workers(job_id_map, enabled_ids): + runtime_state = ensure_distributed_runtime_state() + allowed = {str(worker_id) for worker_id in (enabled_ids or []) if str(worker_id).strip()} + for job_id in job_id_map.values(): + if job_id: + runtime_state.distributed_job_allowed_workers[str(job_id)] = set(allowed) + + async def _ensure_job_queues(prompt_index, job_id_map): for node_id in _distributed_queue_nodes(prompt_index): job_id = job_id_map.get(node_id) @@ -309,7 +303,6 @@ def _build_master_prompt( collector_ids = ( find_nodes_by_class(master_prompt, "DistributedCollector") - + find_nodes_by_class(master_prompt, "DistributedListCollector") + find_nodes_by_class(master_prompt, "DistributedBranchCollector") ) upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") @@ -405,7 +398,7 @@ async def orchestrate_distributed_execution( config = load_config() use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False)) - master_url = build_master_url(config=config, prompt_server_instance=prompt_server) + master_url = build_master_url(config=config, prompt_server_instance=get_prompt_server_instance()) ( worker_probe_concurrency, worker_prep_concurrency, @@ -463,6 +456,7 @@ async def orchestrate_distributed_execution( prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id) return prompt_id, 0 + _register_job_allowed_workers(job_id_map, enabled_ids) await _ensure_job_queues(prompt_index, job_id_map) master_prompt = _build_master_prompt( diff --git a/api/queue_request.py b/api/queue_request.py index 366ea8f..6452962 100644 --- a/api/queue_request.py +++ b/api/queue_request.py @@ -1,16 +1,20 @@ from dataclasses import dataclass -from typing import Any, Dict, List, Optional +from typing import Any + +try: + from ..utils.worker_ids import require_enabled_worker_ids +except ImportError: # pragma: no cover - supports direct module loading in isolated tests. + from utils.worker_ids import require_enabled_worker_ids @dataclass(frozen=True) class QueueRequestPayload: - prompt: Dict[str, Any] - workflow_meta: Any + prompt: dict[str, Any] + workflow_meta: dict[str, Any] | None client_id: str - delegate_master: Optional[bool] - enabled_worker_ids: List[str] - auto_prepare: bool - trace_execution_id: Optional[str] + delegate_master: bool | None + enabled_worker_ids: list[str] + trace_execution_id: str | None def parse_queue_request_payload(data: Any) -> QueueRequestPayload: @@ -18,11 +22,6 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload: if not isinstance(data, dict): raise ValueError("Expected a JSON object body") - auto_prepare_raw = data.get("auto_prepare", True) - if not isinstance(auto_prepare_raw, bool): - raise ValueError("auto_prepare must be a boolean when provided") - auto_prepare = auto_prepare_raw - prompt = data.get("prompt") # Auto-prepare is always on server-side; keep the field for wire compatibility. if prompt is None: @@ -35,6 +34,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload: if not isinstance(prompt, dict): raise ValueError("Field 'prompt' must be an object") + workflow_meta = data.get("workflow") + if workflow_meta is not None and not isinstance(workflow_meta, dict): + raise ValueError("Field 'workflow' must be an object when provided") + enabled_ids_raw = data.get("enabled_worker_ids") workers_field = data.get("workers") if enabled_ids_raw is None and workers_field is not None: @@ -51,7 +54,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload: else: if not isinstance(enabled_ids_raw, list): raise ValueError("enabled_worker_ids must be a list of worker IDs") - enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()] + try: + enabled_ids = require_enabled_worker_ids(enabled_ids_raw) + except ValueError as exc: + raise ValueError(f"enabled_worker_ids invalid: {exc}") from exc delegate_master = data.get("delegate_master") if delegate_master is not None and not isinstance(delegate_master, bool): @@ -70,10 +76,9 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload: return QueueRequestPayload( prompt=prompt, - workflow_meta=data.get("workflow"), + workflow_meta=workflow_meta, client_id=client_id, delegate_master=delegate_master, enabled_worker_ids=enabled_ids, - auto_prepare=auto_prepare, trace_execution_id=trace_execution_id, ) diff --git a/api/request_guards.py b/api/request_guards.py new file mode 100644 index 0000000..43e0d50 --- /dev/null +++ b/api/request_guards.py @@ -0,0 +1,15 @@ +"""Shared request guards for distributed API endpoints.""" +from __future__ import annotations + +from aiohttp import web + +from ..utils.config import load_config +from ..utils.network import handle_api_error +from .schemas import is_authorized_request + + +async def authorization_error_or_none(request: web.Request) -> web.StreamResponse | None: + """Return a 403 response when distributed API auth fails, else None.""" + if is_authorized_request(request, load_config()): + return None + return await handle_api_error(request, "Unauthorized", 403) diff --git a/api/schemas.py b/api/schemas.py index feaa3ec..a159ad3 100644 --- a/api/schemas.py +++ b/api/schemas.py @@ -1,3 +1,29 @@ +from __future__ import annotations + +from typing import Any + +from ..utils.auth import ( + AUTH_HEADER_NAME, + distributed_auth_headers, + is_authorized_request, +) +from ..utils.parsing import coerce_bool, require_bool_literal + +__all__ = [ + "AUTH_HEADER_NAME", + "coerce_bool", + "coerce_positive_float", + "coerce_positive_int", + "distributed_auth_headers", + "is_authorized_request", + "require_bool_literal", + "require_fields", + "require_worker_id", + "require_positive_int", + "validate_worker_id", +] + + def require_fields(data: dict, *fields) -> list[str]: """Return field names that are missing or empty in a JSON object.""" if not isinstance(data, dict): @@ -18,26 +44,43 @@ def require_fields(data: dict, *fields) -> list[str]: return missing -def validate_worker_id(worker_id: str, config: dict) -> bool: +def validate_worker_id(worker_id: str, config: dict[str, Any]) -> bool: """Return True when worker_id exists in config['workers'].""" - worker_id_str = str(worker_id) + try: + require_worker_id(worker_id, config) + return True + except ValueError: + return False + + +def require_worker_id(worker_id: str, config: dict[str, Any], field_name: str = "worker_id") -> str: + """Strict worker-id validator for request boundaries.""" + worker_id_str = str(worker_id).strip() + if not worker_id_str: + raise ValueError(f"Field '{field_name}' must be a non-empty worker id.") + workers = (config or {}).get("workers", []) - return any(str(worker.get("id")) == worker_id_str for worker in workers) + for worker in workers: + if not isinstance(worker, dict): + continue + if str(worker.get("id")).strip() == worker_id_str: + return worker_id_str + raise ValueError(f"Worker {worker_id_str} not found.") -def validate_positive_int(value, field_name: str) -> str | None: - """Validate positive integers and return an error string when invalid.""" +def require_positive_int(value, field_name: str = "value") -> int: + """Strict positive-int parser for request boundary validation.""" try: parsed = int(value) - except (TypeError, ValueError): - return f"Field '{field_name}' must be a positive integer." + except (TypeError, ValueError) as exc: + raise ValueError(f"Field '{field_name}' must be a positive integer.") from exc if parsed <= 0: - return f"Field '{field_name}' must be a positive integer." - return None + raise ValueError(f"Field '{field_name}' must be a positive integer.") + return parsed -def parse_positive_int(value, default: int) -> int: - """Parse value as positive int, returning default on failure.""" +def coerce_positive_int(value, default: int) -> int: + """Permissive positive-int coercion used for non-boundary defaults.""" try: parsed = int(value) except (TypeError, ValueError): @@ -45,10 +88,13 @@ def parse_positive_int(value, default: int) -> int: return max(1, parsed) -def parse_positive_float(value, default: float) -> float: - """Parse value as positive float, returning default on failure.""" +def coerce_positive_float(value, default: float) -> float: + """Permissive positive-float coercion used for non-boundary defaults.""" + fallback = float(default) + if fallback <= 0.0: + fallback = 0.1 try: parsed = float(value) except (TypeError, ValueError): - return max(0.0, float(default)) - return max(0.0, parsed) + return fallback + return parsed if parsed > 0.0 else fallback diff --git a/api/tunnel_routes.py b/api/tunnel_routes.py index 89158a4..703e275 100644 --- a/api/tunnel_routes.py +++ b/api/tunnel_routes.py @@ -3,51 +3,61 @@ from __future__ import annotations from aiohttp import web import server +from .endpoint_policy import run_authorized_endpoint from ..utils.cloudflare import cloudflare_tunnel_manager from ..utils.config import load_config -from ..utils.logging import debug_log, log -from ..utils.network import handle_api_error @server.PromptServer.instance.routes.get("/distributed/tunnel/status") async def tunnel_status_endpoint(request: web.Request) -> web.StreamResponse: """Return Cloudflare tunnel status and last known details.""" - try: + + async def _operation() -> web.StreamResponse: status = cloudflare_tunnel_manager.get_status() config = load_config() master_host = (config.get("master") or {}).get("host") - return web.json_response({ - "status": "success", - "tunnel": status, - "master_host": master_host - }) - except Exception as e: - return await handle_api_error(request, e, 500) + return web.json_response( + { + "status": "success", + "tunnel": status, + "master_host": master_host, + } + ) + + return await run_authorized_endpoint(request, _operation) + @server.PromptServer.instance.routes.post("/distributed/tunnel/start") async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse: """Start a Cloudflare tunnel pointing at the current ComfyUI server.""" - try: + + async def _operation() -> web.StreamResponse: result = await cloudflare_tunnel_manager.start_tunnel() config = load_config() - return web.json_response({ - "status": "success", - "tunnel": result, - "master_host": (config.get("master") or {}).get("host") - }) - except Exception as e: - return await handle_api_error(request, e, 500) + return web.json_response( + { + "status": "success", + "tunnel": result, + "master_host": (config.get("master") or {}).get("host"), + } + ) + + return await run_authorized_endpoint(request, _operation) + @server.PromptServer.instance.routes.post("/distributed/tunnel/stop") async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse: """Stop the managed Cloudflare tunnel if running.""" - try: + + async def _operation() -> web.StreamResponse: result = await cloudflare_tunnel_manager.stop_tunnel() config = load_config() - return web.json_response({ - "status": "success", - "tunnel": result, - "master_host": (config.get("master") or {}).get("host") - }) - except Exception as e: - return await handle_api_error(request, e, 500) + return web.json_response( + { + "status": "success", + "tunnel": result, + "master_host": (config.get("master") or {}).get("host"), + } + ) + + return await run_authorized_endpoint(request, _operation) diff --git a/api/usdu_routes.py b/api/usdu_routes.py index abd06c5..f011af3 100644 --- a/api/usdu_routes.py +++ b/api/usdu_routes.py @@ -8,15 +8,50 @@ from aiohttp import web from PIL import Image import server +from .request_guards import authorization_error_or_none +from .schemas import require_bool_literal from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState -from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized, init_dynamic_job +from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job from ..upscale.payload_parsers import parse_tiles_from_form from ..utils.logging import debug_log from ..utils.network import handle_api_error +from ..utils.usdu_management import MAX_PAYLOAD_SIZE + + +def _parse_int_field(value, field_name: str, *, minimum: int | None = None) -> int: + try: + parsed = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be an integer") from exc + if minimum is not None and parsed < minimum: + raise ValueError(f"{field_name} must be >= {minimum}") + return parsed + + +def _parse_bool_field( + value, + field_name: str, + *, + default: bool = False, +) -> bool: + if value is None: + return default + return require_bool_literal(value, field_name=field_name) + + +def _is_worker_allowed(job_data: BaseJobState, worker_id: str) -> bool: + known_workers = {str(wid) for wid in getattr(job_data, "worker_status", {}).keys()} + known_workers.update(str(wid) for wid in getattr(job_data, "assigned_to_workers", {}).keys()) + if not known_workers: + return True + return str(worker_id) in known_workers @server.PromptServer.instance.routes.post("/distributed/heartbeat") async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse: + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() worker_id = data.get('worker_id') @@ -30,6 +65,8 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse: 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): + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) job_data.worker_status[worker_id] = time.time() debug_log(f"Heartbeat from worker {worker_id}") return web.json_response({"status": "success"}) @@ -42,22 +79,36 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/submit_tiles") async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse: """Endpoint for workers to submit processed tiles in static mode.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error 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) + if content_length: + try: + payload_size = int(content_length) + except ValueError: + return await handle_api_error(request, "Invalid content-length header", 400) + if payload_size > 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' + try: + is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False) + except ValueError as e: + return await handle_api_error(request, str(e), 400) 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)) + try: + batch_size = _parse_int_field(data.get('batch_size', 0), "batch_size", minimum=0) + except ValueError as e: + return await handle_api_error(request, str(e), 400) # Handle completion signal if batch_size == 0 and is_last: @@ -66,6 +117,8 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse: 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) + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) await job_data.queue.put({ 'worker_id': worker_id, 'is_last': True, @@ -85,6 +138,8 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse: 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) + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) q = job_data.queue if batch_size > 0 or len(tiles) > 0: @@ -110,15 +165,26 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/submit_image") async def submit_image_endpoint(request: web.Request) -> web.StreamResponse: """Endpoint for workers to submit processed images in dynamic mode.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error 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) + if content_length: + try: + payload_size = int(content_length) + except ValueError: + return await handle_api_error(request, "Invalid content-length header", 400) + if payload_size > 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' + try: + is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False) + except ValueError as e: + return await handle_api_error(request, str(e), 400) 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) @@ -127,7 +193,10 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse: # Handle image submission if 'full_image' in data and 'image_idx' in data: - image_idx = int(data.get('image_idx')) + try: + image_idx = _parse_int_field(data.get('image_idx'), "image_idx", minimum=0) + except ValueError as e: + return await handle_api_error(request, str(e), 400) img_data = data['full_image'].file.read() img = Image.open(io.BytesIO(img_data)).convert("RGB") @@ -138,6 +207,8 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse: 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) + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) await job_data.queue.put({ 'worker_id': worker_id, 'image_idx': image_idx, @@ -153,6 +224,8 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse: 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) + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) await job_data.queue.put({ 'worker_id': worker_id, 'is_last': True, @@ -170,6 +243,9 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/request_image") async def request_image_endpoint(request: web.Request) -> web.StreamResponse: """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode).""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: data = await request.json() worker_id = data.get('worker_id') @@ -192,6 +268,8 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse: pending_queue = job_data.pending_tasks else: return await handle_api_error(request, "Invalid job configuration", 400) + if not _is_worker_allowed(job_data, worker_id): + return await handle_api_error(request, "Unauthorized worker_id", 403) try: task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) @@ -201,99 +279,33 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse: 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}) + return web.json_response( + { + "kind": "image", + "task_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, + "kind": "tile", + "task_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 web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0}) 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/init_list_queue") -async def init_list_queue_endpoint(request: web.Request) -> web.StreamResponse: - """Initialize a shared queue of list indices for dynamic DistributedListSplitter mode.""" - try: - data = await request.json() - multi_job_id = data.get("multi_job_id") - list_size = data.get("list_size", 0) - enabled_workers = data.get("enabled_workers", []) - - if not multi_job_id: - return await handle_api_error(request, "Missing multi_job_id", 400) - - try: - list_size = max(0, int(list_size)) - except (TypeError, ValueError): - return await handle_api_error(request, "list_size must be an integer", 400) - - worker_ids = [str(worker_id) for worker_id in (enabled_workers or []) if str(worker_id).strip()] - await init_dynamic_job( - multi_job_id=str(multi_job_id), - batch_size=list_size, - enabled_workers=worker_ids, - all_indices=list(range(list_size)), - ) - - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id)) - remaining = ( - job_data.pending_images.qsize() - if isinstance(job_data, ImageJobState) - else 0 - ) - - return web.json_response({"status": "success", "remaining": remaining}) - except Exception as e: - return await handle_api_error(request, e, 500) - - -@server.PromptServer.instance.routes.post("/distributed/request_list_item") -async def request_list_item_endpoint(request: web.Request) -> web.StreamResponse: - """Allow workers to pull one list item index for dynamic list distribution.""" - 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: - job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id)) - if not isinstance(job_data, BaseJobState): - return await handle_api_error(request, "Job not found", 404) - - pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None) - if pending_queue is None: - return await handle_api_error(request, "Job not configured for list item requests", 400) - - try: - item_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1) - except asyncio.TimeoutError: - return web.json_response({"item_idx": None}) - - worker_id = str(worker_id) - job_data.assigned_to_workers.setdefault(worker_id, []).append(item_idx) - job_data.worker_status[worker_id] = time.time() - remaining = pending_queue.qsize() - return web.json_response({"item_idx": int(item_idx), "estimated_remaining": remaining}) - 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: web.Request) -> web.StreamResponse: """Endpoint to check if a job is ready.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error multi_job_id = request.query.get('multi_job_id') if not multi_job_id: return web.json_response({"ready": False}) diff --git a/api/worker_routes.py b/api/worker_routes.py index 7d32d46..8869e75 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -25,27 +25,27 @@ from ..utils.network import ( ) from ..utils.constants import CHUNK_SIZE from ..workers import get_worker_manager -from .schemas import require_fields, validate_worker_id +from .request_guards import authorization_error_or_none +from .schemas import ( + distributed_auth_headers, + require_fields, + require_worker_id, +) from ..workers.detection import ( get_machine_id, is_docker_environment, is_runpod_environment, ) -try: - from ..utils.async_helpers import PromptValidationError, queue_prompt_payload -except ImportError: - from ..utils.async_helpers import queue_prompt_payload - - class PromptValidationError(RuntimeError): - def __init__(self, message, validation_error=None, node_errors=None): - super().__init__(str(message)) - self.validation_error = validation_error if isinstance(validation_error, dict) else {} - self.node_errors = node_errors if isinstance(node_errors, dict) else {} +from ..utils.async_helpers import PromptValidationError, queue_prompt_payload @server.PromptServer.instance.routes.get("/distributed/worker_ws") -async def worker_ws_endpoint(request: web.Request) -> web.WebSocketResponse: +async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse: """WebSocket endpoint for worker prompt dispatch.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error + ws = web.WebSocketResponse(heartbeat=30) await ws.prepare(request) @@ -118,6 +118,9 @@ async def worker_ws_endpoint(request: web.Request) -> web.WebSocketResponse: @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching") async def clear_launching_state(request: web.Request) -> web.StreamResponse: """Clear the launching flag when worker is confirmed running.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: wm = get_worker_manager() data = await request.json() @@ -127,8 +130,10 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse: worker_id = str(data.get("worker_id")).strip() config = load_config() - if not validate_worker_id(worker_id, config): - return await handle_api_error(request, f"Worker {worker_id} not found", 404) + try: + worker_id = require_worker_id(worker_id, config) + except ValueError as exc: + return await handle_api_error(request, exc, 404) # Clear launching flag in managed processes if worker_id in wm.processes: @@ -372,6 +377,9 @@ def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any] @server.PromptServer.instance.routes.get("/distributed/local_log") async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse: """Return this instance's in-memory ComfyUI log buffer.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: from app.logger import get_logs except Exception as e: @@ -417,6 +425,9 @@ async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.get("/distributed/network_info") async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse: """Get network interfaces and recommend best IP for master.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: loop = asyncio.get_running_loop() info = await loop.run_in_executor(None, _collect_network_info_sync) @@ -432,6 +443,9 @@ async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.get("/distributed/system_info") async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse: """Get system information including machine ID for local worker detection.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: import socket @@ -456,6 +470,9 @@ async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/launch_worker") async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse: """Launch a worker process from the UI.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: wm = get_worker_manager() data = await request.json() @@ -467,8 +484,10 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse: # Find worker config config = load_config() - if not validate_worker_id(worker_id, config): - return await handle_api_error(request, f"Worker {worker_id} not found", 404) + try: + worker_id = require_worker_id(worker_id, config) + except ValueError as exc: + return await handle_api_error(request, exc, 404) worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None) if not worker: return await handle_api_error(request, f"Worker {worker_id} not found", 404) @@ -487,7 +506,7 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse: is_running = process.poll() is None else: # Restored process without subprocess object - is_running = wm._is_process_running(proc_info['pid']) + is_running = wm.is_process_running(proc_info['pid']) if is_running: return await handle_api_error(request, "Worker already running (managed by UI)", 409) @@ -517,6 +536,9 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.post("/distributed/stop_worker") async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse: """Stop a worker process that was launched from the UI.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: wm = get_worker_manager() data = await request.json() @@ -526,8 +548,10 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse: worker_id = str(data.get("worker_id")).strip() config = load_config() - if not validate_worker_id(worker_id, config): - return await handle_api_error(request, f"Worker {worker_id} not found", 404) + try: + worker_id = require_worker_id(worker_id, config) + except ValueError as exc: + return await handle_api_error(request, exc, 404) success, message = wm.stop_worker(worker_id) @@ -547,6 +571,9 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.get("/distributed/managed_workers") async def get_managed_workers_endpoint(request: web.Request) -> web.StreamResponse: """Get list of workers managed by this UI instance.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: managed = get_worker_manager().get_managed_workers() return web.json_response({ @@ -560,6 +587,9 @@ async def get_managed_workers_endpoint(request: web.Request) -> web.StreamRespon @server.PromptServer.instance.routes.get("/distributed/local-worker-status") async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamResponse: """Check status of all local workers (localhost/no host specified).""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: config = load_config() worker_statuses = {} @@ -630,6 +660,9 @@ async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamRe @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}") async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse: """Get log content for a specific worker.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: wm = get_worker_manager() worker_id = request.match_info['worker_id'] @@ -673,6 +706,9 @@ async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse: @server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}") async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResponse: """Proxy a remote worker log request to the worker's local in-memory log endpoint.""" + auth_error = await authorization_error_or_none(request) + if auth_error is not None: + return auth_error try: worker_id = str(request.match_info["worker_id"]).strip() config = load_config() @@ -695,6 +731,7 @@ async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResp async with session.get( worker_url, params={"lines": str(lines_to_read)}, + headers=distributed_auth_headers(config), timeout=aiohttp.ClientTimeout(total=5), ) as resp: if resp.status >= 400: diff --git a/bootstrap/__init__.py b/bootstrap/__init__.py new file mode 100644 index 0000000..2905566 --- /dev/null +++ b/bootstrap/__init__.py @@ -0,0 +1,5 @@ +"""Bootstrap entrypoints for ComfyUI-Distributed.""" + +from .entrypoint import build_node_mappings, initialize_runtime + +__all__ = ["build_node_mappings", "initialize_runtime"] diff --git a/bootstrap/entrypoint.py b/bootstrap/entrypoint.py new file mode 100644 index 0000000..fb2ca61 --- /dev/null +++ b/bootstrap/entrypoint.py @@ -0,0 +1,66 @@ +"""Single startup owner for node surface and runtime initialization.""" +from __future__ import annotations + +import atexit +import os +from functools import lru_cache +from typing import Any + + +@lru_cache(maxsize=1) +def build_node_mappings() -> tuple[dict[str, Any], dict[str, str]]: + """Build and cache the full exported node mapping surface.""" + from ..nodes import ( + NODE_CLASS_MAPPINGS as distributed_class_mappings, + NODE_DISPLAY_NAME_MAPPINGS as distributed_display_name_mappings, + ) + from ..nodes.distributed_upscale import ( + NODE_CLASS_MAPPINGS as upscale_class_mappings, + NODE_DISPLAY_NAME_MAPPINGS as upscale_display_name_mappings, + ) + + return ( + {**distributed_class_mappings, **upscale_class_mappings}, + {**distributed_display_name_mappings, **upscale_display_name_mappings}, + ) + + +@lru_cache(maxsize=1) +def _runtime_initialized() -> dict[str, bool]: + return {"done": False} + + +def initialize_runtime(prompt_server: Any | None = None) -> None: + """Initialize distributed runtime side effects once.""" + state = _runtime_initialized() + if state["done"]: + return + + import server + + from ..api import bootstrap_routes + from ..api.queue_orchestration import ensure_distributed_state + from ..upscale.job_store import ensure_tile_jobs_initialized + from ..utils.config import CONFIG_FILE, ensure_config_exists + from ..utils.logging import debug_log + from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup + + ensure_config_exists() + bootstrap_routes() + + server_instance = prompt_server if prompt_server is not None else server.PromptServer.instance + ensure_distributed_state(server_instance) + ensure_tile_jobs_initialized() + + atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup + if not os.environ.get("COMFYUI_IS_WORKER"): + atexit.register(sync_cleanup) + delayed_auto_launch() + register_async_signals() + + state["done"] = True + + node_mappings, _ = build_node_mappings() + debug_log("Loaded Distributed nodes.") + debug_log(f"Config file: {CONFIG_FILE}") + debug_log(f"Available nodes: {list(node_mappings.keys())}") diff --git a/conftest.py b/conftest.py index db98903..fe87efb 100644 --- a/conftest.py +++ b/conftest.py @@ -1,28 +1,4 @@ -# conftest.py — project-level pytest configuration. -# -# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports -# (from .distributed import ...) that fail when pytest tries to import it as a -# standalone module during Package.setup() for the root package node. -# -# Fix: patch Package.setup() to skip the root-package's __init__.py import. -# All actual package context is provided by each test module via -# importlib.util.spec_from_file_location with synthetic stub packages. - -from _pytest.python import Package - -_orig_pkg_setup = Package.setup - - -def _patched_pkg_setup(self) -> None: - # Skip the root package setup — its __init__.py uses relative imports - # that require a parent package (ComfyUI's plugin loader) which is not - # available in the test environment. - if self.path == self.config.rootpath: - return - _orig_pkg_setup(self) - - -Package.setup = _patched_pkg_setup +"""Project-level pytest collection rules for plugin-style package layout.""" collect_ignore = [ "__init__.py", diff --git a/distributed.py b/distributed.py index 0a3ad22..a1f7d9d 100644 --- a/distributed.py +++ b/distributed.py @@ -1,55 +1,15 @@ -""" -ComfyUI-Distributed: thin entry point. -All implementation lives in workers/, nodes/, api/. -""" -import atexit -import os +"""Compatibility facade for legacy imports from `distributed`.""" +from __future__ import annotations -import server - -from .utils.config import ensure_config_exists -from .utils.logging import debug_log -from .utils.network import cleanup_client_session -from .workers import get_worker_manager -from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup -from .upscale.job_store import ensure_tile_jobs_initialized -from .nodes import ( - NODE_CLASS_MAPPINGS, - NODE_DISPLAY_NAME_MAPPINGS, - ImageBatchDivider, - DistributedCollectorNode, - DistributedSeed, - DistributedModelName, - DistributedValue, - AudioBatchDivider, - DistributedEmptyImage, - DistributedListSplitter, - DistributedListCollector, - DistributedBranch, - DistributedBranchCollector, - AnyType, - ByPassTypeTuple, - any_type, +from .bootstrap.entrypoint import ( + build_node_mappings, + initialize_runtime, ) -from . import api # noqa: F401 - triggers all @routes.* registrations -from .api.queue_orchestration import ensure_distributed_state -ensure_config_exists() +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings() -# Aiohttp session cleanup -async def _cleanup_session(): - await cleanup_client_session() - - -atexit.register(lambda: None) # placeholder; real cleanup in sync_cleanup - -# Initialize distributed job state on prompt_server -prompt_server = server.PromptServer.instance -ensure_distributed_state(prompt_server) -ensure_tile_jobs_initialized() - -# Worker startup -if not os.environ.get('COMFYUI_IS_WORKER'): - atexit.register(sync_cleanup) - delayed_auto_launch() - register_async_signals() +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "initialize_runtime", +] diff --git a/docs/comfyui-distributed-api.md b/docs/comfyui-distributed-api.md index 21383be..bf39575 100644 --- a/docs/comfyui-distributed-api.md +++ b/docs/comfyui-distributed-api.md @@ -53,7 +53,6 @@ Queue a workflow for distributed execution. "delegate_master": false, "enabled_worker_ids": ["1", "2"], "workers": ["1", "2"], - "auto_prepare": true, "trace_execution_id": "exec_1700000000_ab12cd" } ``` @@ -74,10 +73,6 @@ Queue a workflow for distributed execution. - The explicit worker IDs to consider for this run. - `workers` (optional, array of strings or objects with `id`) - Transitional alias for `enabled_worker_ids` used by older clients. -- `auto_prepare` (optional, boolean) - - Kept for wire compatibility. - - Backend orchestration always runs with auto-prepare semantics. - - If top-level `prompt` is omitted, backend will attempt `workflow.prompt`. - `trace_execution_id` (optional, string) - Passed through to orchestration logs. - Server log lines include the marker as `[exec:]`. @@ -105,8 +100,7 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto ```json { "prompt_id": "", - "worker_count": 2, - "auto_prepare_supported": true + "worker_count": 2 } ``` diff --git a/nodes/__init__.py b/nodes/__init__.py index d83057b..d611a4a 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -10,15 +10,11 @@ from .utilities import ( any_type, ) from .collector import DistributedCollectorNode -from .list_splitter import DistributedListSplitter -from .list_collector import DistributedListCollector from .branch import DistributedBranch from .branch_collector import DistributedBranchCollector NODE_CLASS_MAPPINGS = { "DistributedCollector": DistributedCollectorNode, - "DistributedListSplitter": DistributedListSplitter, - "DistributedListCollector": DistributedListCollector, "DistributedBranch": DistributedBranch, "DistributedBranchCollector": DistributedBranchCollector, "DistributedSeed": DistributedSeed, @@ -30,8 +26,6 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { "DistributedCollector": "Distributed Collector", - "DistributedListSplitter": "Distributed List Splitter", - "DistributedListCollector": "Distributed List Collector", "DistributedBranch": "Distributed Branch", "DistributedBranchCollector": "Distributed Branch Collector", "DistributedSeed": "Distributed Seed", diff --git a/nodes/branch_collector.py b/nodes/branch_collector.py index 9b0f5f9..174c2dc 100644 --- a/nodes/branch_collector.py +++ b/nodes/branch_collector.py @@ -1,16 +1,17 @@ import asyncio -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import aiohttp import torch +from ..utils.auth import distributed_auth_headers from ..utils.async_helpers import run_async_in_server_loop -from ..utils.config import is_master_delegate_only +from ..utils.config import is_master_delegate_only, load_config from ..utils.image import encode_tensor_png_data_url, ensure_contiguous from ..utils.logging import log from ..utils.network import get_client_session -from ..utils.worker_ids import parse_enabled_worker_ids +from ..utils.worker_ids import coerce_enabled_worker_ids from .context_kwargs import parse_distributed_hidden_context from .hidden_inputs import build_distributed_hidden_inputs from .queue_wait import collect_worker_queue_results @@ -29,7 +30,7 @@ class BranchRunContext: multi_job_id: str = "" is_worker: bool = False master_url: str = "" - enabled_worker_ids: str = "[]" + enabled_worker_ids: list[str] = field(default_factory=list) worker_id: str = "" assigned_branch: int = -1 delegate_only: bool = False @@ -190,6 +191,7 @@ class DistributedBranchCollector: async with session.post( url, json=payload, + headers=distributed_auth_headers(load_config()), timeout=aiohttp.ClientTimeout(total=60), ) as response: response.raise_for_status() @@ -236,7 +238,7 @@ class DistributedBranchCollector: return tuple(outputs) delegate_mode = bool(run_context.delegate_only or is_master_delegate_only()) - enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids) + enabled_workers = coerce_enabled_worker_ids(run_context.enabled_worker_ids) slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count) expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count) expected_workers = set(expected_worker_ids) @@ -282,6 +284,8 @@ class DistributedBranchCollector: finally: async with prompt_server.distributed_jobs_lock: prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None) + if hasattr(prompt_server, "distributed_job_allowed_workers"): + prompt_server.distributed_job_allowed_workers.pop(run_context.multi_job_id, None) missing_workers = sorted(expected_workers - workers_done) if missing_workers: diff --git a/nodes/collector.py b/nodes/collector.py index bbf43c1..bb7331b 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -3,7 +3,7 @@ import base64 import io import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import aiohttp @@ -12,12 +12,13 @@ import comfy.model_management import torch from comfy.utils import ProgressBar +from ..utils.auth import distributed_auth_headers from ..utils.logging import debug_log, log from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only from ..utils.constants import HEARTBEAT_INTERVAL from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous from ..utils.network import build_worker_url, get_client_session, probe_worker -from ..utils.worker_ids import parse_enabled_worker_ids +from ..utils.worker_ids import coerce_enabled_worker_ids from ..utils.audio_payload import encode_audio_payload from ..utils.async_helpers import run_async_in_server_loop from .context_kwargs import parse_distributed_hidden_context @@ -31,7 +32,7 @@ class CollectorRunContext: multi_job_id: str = "" is_worker: bool = False master_url: str = "" - enabled_worker_ids: str = "[]" + enabled_worker_ids: list[str] = field(default_factory=list) worker_batch_size: int = 1 worker_id: str = "" pass_through: bool = False @@ -142,6 +143,7 @@ class DistributedCollectorNode: async with session.post( url, json=payload, + headers=distributed_auth_headers(load_config()), timeout=aiohttp.ClientTimeout(total=60), ) as response: response.raise_for_status() @@ -280,6 +282,8 @@ class DistributedCollectorNode: async with prompt_server.distributed_jobs_lock: if multi_job_id in prompt_server.distributed_pending_jobs: del prompt_server.distributed_pending_jobs[multi_job_id] + if hasattr(prompt_server, "distributed_job_allowed_workers"): + prompt_server.distributed_job_allowed_workers.pop(multi_job_id, None) def _build_master_inputs(self, images, audio, delegate_mode, multi_job_id, num_workers): if delegate_mode: @@ -452,7 +456,7 @@ class DistributedCollectorNode: delegate_only, ): delegate_mode = delegate_only or is_master_delegate_only() - enabled_workers = parse_enabled_worker_ids(enabled_worker_ids) + enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids) expected_workers = set(enabled_workers) num_workers = len(expected_workers) if num_workers == 0: diff --git a/nodes/context_kwargs.py b/nodes/context_kwargs.py index 5eb3804..acc3173 100644 --- a/nodes/context_kwargs.py +++ b/nodes/context_kwargs.py @@ -1,13 +1,26 @@ from typing import Any +from ..utils.parsing import coerce_bool +from ..utils.worker_ids import coerce_enabled_worker_ids + + +def _parse_bool(value: Any, default: bool = False) -> bool: + return coerce_bool(value, default=default) + + +def _parse_enabled_workers(value: Any) -> list[str]: + if value is None: + return [] + return coerce_enabled_worker_ids(value) + def parse_distributed_hidden_context(kwargs: dict[str, Any]) -> dict[str, Any]: """Parse common distributed hidden inputs from node kwargs.""" return { - "multi_job_id": str(kwargs.get("multi_job_id", "")), - "is_worker": bool(kwargs.get("is_worker", False)), - "master_url": str(kwargs.get("master_url", "")), - "enabled_worker_ids": str(kwargs.get("enabled_worker_ids", "[]")), - "worker_id": str(kwargs.get("worker_id", "")), - "delegate_only": bool(kwargs.get("delegate_only", False)), + "multi_job_id": str(kwargs.get("multi_job_id", "") or ""), + "is_worker": _parse_bool(kwargs.get("is_worker", False)), + "master_url": str(kwargs.get("master_url", "") or ""), + "enabled_worker_ids": _parse_enabled_workers(kwargs.get("enabled_worker_ids", "[]")), + "worker_id": str(kwargs.get("worker_id", "") or ""), + "delegate_only": _parse_bool(kwargs.get("delegate_only", False)), } diff --git a/nodes/distributed_upscale.py b/nodes/distributed_upscale.py index c71503b..0e757c8 100644 --- a/nodes/distributed_upscale.py +++ b/nodes/distributed_upscale.py @@ -1,4 +1,3 @@ -import json import math from functools import wraps from typing import Any, Callable @@ -18,6 +17,7 @@ from ..upscale.processing_args import UpscaleCoreArgs from ..upscale.modes.single_gpu import SingleGpuModeMixin from ..upscale.modes.static import StaticModeMixin from ..upscale.modes.dynamic import DynamicModeMixin +from ..utils.worker_ids import coerce_enabled_worker_ids def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]: """Decorator to wrap async methods for synchronous execution.""" @@ -30,21 +30,10 @@ def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]: ) return sync_func -def _parse_enabled_worker_ids(enabled_worker_ids): - """Parse enabled worker IDs from either JSON or list input.""" - if isinstance(enabled_worker_ids, list): - return [str(worker_id) for worker_id in enabled_worker_ids] - if not enabled_worker_ids: - return [] - if isinstance(enabled_worker_ids, str): - try: - parsed = json.loads(enabled_worker_ids) - except json.JSONDecodeError: - log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.") - return [] - if isinstance(parsed, list): - return [str(wid) for wid in parsed] - return [] + +def _parse_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]: + """Backward-compatible alias for enabled-worker normalization.""" + return coerce_enabled_worker_ids(enabled_worker_ids) class UltimateSDUpscaleDistributed( DynamicModeMixin, @@ -255,7 +244,7 @@ class UltimateSDUpscaleDistributed( # Ensure mode consistency across master/workers via shared threshold # Determine mode (must match master's logic) - enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids) + enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids) num_workers = len(enabled_workers) # Compute number of tiles for this image to decide if tile distribution makes sense _, height, width, _ = upscaled_image.shape @@ -264,8 +253,8 @@ class UltimateSDUpscaleDistributed( mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) # For USDU-style processing, we want tile distribution whenever workers are available - # and there is more than one tile to process, even if batch == 1. - if num_workers > 0 and num_tiles_per_image > 1: + # and there is more than one tile to process for single-image runs. + if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1: mode = "static" debug_log(f"USDU Dist Worker - Batch size {batch_size}") @@ -324,14 +313,14 @@ class UltimateSDUpscaleDistributed( ) # Parse enabled workers - enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids) + enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids) num_workers = len(enabled_workers) # Determine processing mode mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold) # Prefer tile-based static distribution when workers are available and there are multiple tiles, - # even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue. - if num_workers > 0 and num_tiles_per_image > 1: + # for single-image jobs to spread tiles across GPUs like the legacy tile queue. + if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1: mode = "static" log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}") @@ -370,18 +359,14 @@ class UltimateSDUpscaleDistributed( all_tiles, num_tiles_per_image) def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str: - """Determines processing mode per requested policy: - - any workers => prefer static (tile-based) for USDU - - no workers => single_gpu - """ + """Determine mode from worker availability and configured threshold.""" if num_workers == 0: return "single_gpu" - # Default to static when distributed; master/worker may still override if special cases arise + threshold = max(1, int(dynamic_threshold)) + if int(batch_size) >= threshold: + return "dynamic" return "static" -# Ensure initialization before registering routes -ensure_tile_jobs_initialized() - # Node registration NODE_CLASS_MAPPINGS = { "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed, diff --git a/nodes/hidden_inputs.py b/nodes/hidden_inputs.py index 8128805..a327986 100644 --- a/nodes/hidden_inputs.py +++ b/nodes/hidden_inputs.py @@ -1,6 +1,15 @@ from typing import Any +def build_worker_identity_hidden_inputs() -> dict[str, tuple[str, dict[str, Any]]]: + """Hidden-input contract for worker identity context.""" + return { + "is_worker": ("BOOLEAN", {"default": False}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + "worker_id": ("STRING", {"default": ""}), + } + + def build_distributed_hidden_inputs( *, include_assigned_branch: bool = False, @@ -10,12 +19,10 @@ def build_distributed_hidden_inputs( ) -> dict[str, tuple[str, dict[str, Any]]]: hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = { "multi_job_id": ("STRING", {"default": ""}), - "is_worker": ("BOOLEAN", {"default": False}), "master_url": ("STRING", {"default": ""}), - "enabled_worker_ids": ("STRING", {"default": "[]"}), - "worker_id": ("STRING", {"default": ""}), "delegate_only": ("BOOLEAN", {"default": False}), } + hidden_inputs.update(build_worker_identity_hidden_inputs()) if include_assigned_branch: hidden_inputs["assigned_branch"] = ( "INT", diff --git a/nodes/list_collector.py b/nodes/list_collector.py deleted file mode 100644 index 4a275e8..0000000 --- a/nodes/list_collector.py +++ /dev/null @@ -1,238 +0,0 @@ -import asyncio -from dataclasses import dataclass -from typing import Any - -import aiohttp -import torch - -from ..utils.async_helpers import run_async_in_server_loop -from ..utils.config import is_master_delegate_only -from ..utils.image import encode_tensor_png_data_url, ensure_contiguous -from ..utils.logging import debug_log, log -from ..utils.network import get_client_session -from ..utils.worker_ids import parse_enabled_worker_ids -from .context_kwargs import parse_distributed_hidden_context -from .hidden_inputs import build_distributed_hidden_inputs -from .queue_wait import collect_worker_queue_results -from .runtime_helpers import ( - get_prompt_server_instance as _get_prompt_server_instance, - throw_if_processing_interrupted as _throw_if_processing_interrupted, -) - - -@dataclass(frozen=True) -class ListCollectorRunContext: - multi_job_id: str = "" - is_worker: bool = False - master_url: str = "" - enabled_worker_ids: str = "[]" - worker_id: str = "" - delegate_only: bool = False - - -class DistributedListCollector: - INPUT_IS_LIST = True - OUTPUT_IS_LIST = (True,) - - @classmethod - def INPUT_TYPES(cls: type["DistributedListCollector"]) -> dict[str, Any]: - return { - "required": { - "images": ("IMAGE",), - }, - "hidden": build_distributed_hidden_inputs(), - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("images",) - FUNCTION = "run" - CATEGORY = "image" - - def _build_run_context(self, **kwargs: Any) -> ListCollectorRunContext: - return ListCollectorRunContext(**parse_distributed_hidden_context(kwargs)) - - def run( - self, - images: list[Any] | torch.Tensor | None, - **kwargs: Any, - ) -> tuple[list[Any]]: - context = self._build_run_context(**kwargs) - image_list = self._normalize_image_list(images) - if not context.multi_job_id: - return (image_list,) - - return run_async_in_server_loop( - self.execute( - image_list, - context=context, - ) - ) - - def _normalize_image_list(self, images: list[Any] | torch.Tensor | None) -> list[Any]: - if images is None: - return [] - if isinstance(images, list): - return list(images) - return [images] - - async def send_list_to_master( - self, - image_list: list[Any], - multi_job_id: str, - master_url: str, - worker_id: str, - ) -> None: - if not image_list: - return - - sendable = [] - for idx, image in enumerate(image_list): - if not isinstance(image, torch.Tensor): - continue - batch = image - if batch.ndim == 3: - batch = batch.unsqueeze(0) - if batch.ndim != 4 or batch.shape[0] <= 0: - continue - sendable.append((idx, batch)) - - if not sendable: - return - - session = await get_client_session() - url = f"{master_url}/distributed/job_complete" - - for send_pos, (image_index, image_batch) in enumerate(sendable): - payload = { - "job_id": str(multi_job_id), - "worker_id": str(worker_id), - "batch_idx": int(image_index), - "image": encode_tensor_png_data_url(image_batch, 0), - "is_last": bool(send_pos == len(sendable) - 1), - } - try: - async with session.post( - url, - json=payload, - timeout=aiohttp.ClientTimeout(total=60), - ) as response: - response.raise_for_status() - except Exception as exc: - log(f"Worker - Failed to send list item to master: {exc}") - debug_log(f"Worker - Full error details: URL={url}") - raise - - def _store_worker_result(self, worker_images, item): - worker_id = str(item.get("worker_id", "")) - tensor = item.get("tensor") - image_index = item.get("image_index") - if not worker_id or tensor is None or image_index is None: - return 0 - - worker_images.setdefault(worker_id, {}) - worker_images[worker_id][int(image_index)] = tensor - return 1 - - def _reorder_and_combine_list(self, worker_images, worker_order, master_images, delegate_mode): - ordered = [] - if not delegate_mode: - ordered.extend(master_images) - - ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])] - seen = set() - - for worker_id in ordered_worker_ids: - seen.add(worker_id) - worker_entries = worker_images.get(worker_id, {}) - for image_index in sorted(worker_entries.keys()): - ordered.append(worker_entries[image_index]) - - for worker_id in sorted(worker_images.keys()): - if worker_id in seen: - continue - worker_entries = worker_images.get(worker_id, {}) - for image_index in sorted(worker_entries.keys()): - ordered.append(worker_entries[image_index]) - - combined = [] - for image in ordered: - if isinstance(image, torch.Tensor): - if image.is_cuda: - image = image.cpu() - image = ensure_contiguous(image) - combined.append(image) - - return combined - - async def execute(self, *args: Any, **kwargs: Any) -> tuple[list[Any]]: - """Compatibility wrapper with a uniform execute() signature across collectors.""" - return await self._execute_list(*args, **kwargs) - - async def _execute_list( - self, - images: list[Any] | torch.Tensor | None, - context: ListCollectorRunContext | None = None, - ) -> tuple[list[Any]]: - run_context = context or ListCollectorRunContext() - image_list = self._normalize_image_list(images) - - if run_context.is_worker: - debug_log( - "Worker - Job " - f"{run_context.multi_job_id} complete. Sending {len(image_list)} image list item(s) to master" - ) - await self.send_list_to_master( - image_list, - run_context.multi_job_id, - run_context.master_url, - run_context.worker_id, - ) - return (image_list,) - - delegate_mode = bool(run_context.delegate_only or is_master_delegate_only()) - enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids) - expected_workers = set(enabled_workers) - if not expected_workers: - return (image_list,) - - prompt_server = _get_prompt_server_instance() - - async with prompt_server.distributed_jobs_lock: - if run_context.multi_job_id not in prompt_server.distributed_pending_jobs: - prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue() - debug_log( - "Master - Initialized queue early for list collector job " - f"{run_context.multi_job_id}" - ) - - master_images = [] if delegate_mode else image_list - - worker_images = {} - workers_done: set[str] = set() - - def _handle_queue_result(result: dict[str, Any]) -> None: - self._store_worker_result(worker_images, result) - - try: - workers_done = await collect_worker_queue_results( - prompt_server=prompt_server, - multi_job_id=run_context.multi_job_id, - expected_workers=expected_workers, - on_result=_handle_queue_result, - timeout_log_prefix=( - "Master - List collector heartbeat timeout. " - "Still waiting for workers: " - ), - throw_if_interrupted=_throw_if_processing_interrupted, - ) - - finally: - async with prompt_server.distributed_jobs_lock: - prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None) - - combined = self._reorder_and_combine_list(worker_images, enabled_workers, master_images, delegate_mode) - debug_log( - "Master - List collector job " - f"{run_context.multi_job_id} complete. Combined {len(combined)} image list item(s)." - ) - return (combined,) diff --git a/nodes/list_splitter.py b/nodes/list_splitter.py deleted file mode 100644 index 7b3cf95..0000000 --- a/nodes/list_splitter.py +++ /dev/null @@ -1,220 +0,0 @@ -import asyncio -from dataclasses import dataclass -from typing import Any - -import aiohttp - -from ..upscale.job_models import BaseJobState, ImageJobState -from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job -from ..utils.async_helpers import run_async_in_server_loop -from ..utils.logging import debug_log -from ..utils.network import get_client_session -from .utilities import _chunk_bounds - - -@dataclass(frozen=True) -class ListSplitterRunContext: - participant_index: int = 0 - total_participants: int = 1 - multi_job_id: str = "" - is_worker: bool = False - master_url: str = "" - worker_id: str = "" - - -class DistributedListSplitter: - INPUT_IS_LIST = True - OUTPUT_IS_LIST = (True,) - - @classmethod - def INPUT_TYPES(cls: type["DistributedListSplitter"]) -> dict[str, Any]: - return { - "required": { - "images": ("IMAGE",), - "mode": (["static", "dynamic"], {"default": "static"}), - }, - "hidden": { - "participant_index": ("INT", {"default": 0, "min": 0, "max": 1024}), - "total_participants": ("INT", {"default": 1, "min": 1, "max": 1024}), - "multi_job_id": ("STRING", {"default": ""}), - "is_worker": ("BOOLEAN", {"default": False}), - "master_url": ("STRING", {"default": ""}), - "worker_id": ("STRING", {"default": ""}), - }, - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("images",) - FUNCTION = "split" - CATEGORY = "image" - - def _build_run_context(self, **kwargs: Any) -> ListSplitterRunContext: - try: - participant_index = int(kwargs.get("participant_index", 0)) - except (TypeError, ValueError): - participant_index = 0 - try: - total_participants = int(kwargs.get("total_participants", 1)) - except (TypeError, ValueError): - total_participants = 1 - - return ListSplitterRunContext( - participant_index=participant_index, - total_participants=max(total_participants, 1), - multi_job_id=str(kwargs.get("multi_job_id", "")), - is_worker=bool(kwargs.get("is_worker", False)), - master_url=str(kwargs.get("master_url", "")), - worker_id=str(kwargs.get("worker_id", "")), - ) - - def split( - self, - images: list[Any] | Any | None, - mode: str = "static", - **kwargs: Any, - ) -> tuple[list[Any]]: - context = self._build_run_context(**kwargs) - if images is None: - return ([],) - - image_list = images if isinstance(images, list) else [images] - if mode != "dynamic" or not context.multi_job_id: - return self._split_static( - image_list, - context.participant_index, - context.total_participants, - ) - - return run_async_in_server_loop( - self._split_dynamic( - image_list, - context.multi_job_id, - context.is_worker, - context.master_url, - context.worker_id, - ) - ) - - def _split_static(self, image_list, participant_index, total_participants): - split_count = max(1, int(total_participants or 1)) - if split_count <= 1: - return (list(image_list),) - - bounds = _chunk_bounds(len(image_list), split_count) - index = int(participant_index or 0) - if index < 0: - index = 0 - if index >= split_count: - index = split_count - 1 - - start, end = bounds[index] - return (list(image_list[start:end]),) - - async def _ensure_local_dynamic_queue(self, multi_job_id, total_items): - prompt_server = ensure_tile_jobs_initialized() - async with prompt_server.distributed_tile_jobs_lock: - existing = prompt_server.distributed_pending_tile_jobs.get(multi_job_id) - if isinstance(existing, ImageJobState): - return - await init_dynamic_job( - multi_job_id, - batch_size=total_items, - enabled_workers=[], - all_indices=list(range(total_items)), - ) - - async def _pull_items_from_local_queue(self, multi_job_id, worker_id): - prompt_server = ensure_tile_jobs_initialized() - pulled = [] - worker_key = str(worker_id or "master") - - while True: - 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): - break - pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None) - if pending_queue is None: - break - - try: - item_idx = pending_queue.get_nowait() - except asyncio.QueueEmpty: - break - - job_data.assigned_to_workers.setdefault(worker_key, []).append(item_idx) - pulled.append(int(item_idx)) - - return pulled - - async def _request_item_from_master(self, multi_job_id, master_url, worker_id): - session = await get_client_session() - url = f"{master_url}/distributed/request_list_item" - try: - async with session.post( - url, - json={"multi_job_id": str(multi_job_id), "worker_id": str(worker_id or "")}, - timeout=aiohttp.ClientTimeout(total=10), - ) as response: - if response.status == 404: - return "not_ready", None - if response.status != 200: - return "error", None - payload = await response.json() - return "ok", payload.get("item_idx") - except Exception: - return "error", None - - async def _pull_items_from_master(self, multi_job_id, master_url, worker_id): - pulled = [] - consecutive_not_ready = 0 - max_not_ready_polls = 40 - - while True: - status, item_idx = await self._request_item_from_master(multi_job_id, master_url, worker_id) - - if status == "not_ready": - if consecutive_not_ready >= max_not_ready_polls: - break - consecutive_not_ready += 1 - await asyncio.sleep(0.25) - continue - - if status == "error": - if pulled: - break - await asyncio.sleep(0.25) - continue - - consecutive_not_ready = 0 - if item_idx is None: - break - - try: - pulled.append(int(item_idx)) - except (TypeError, ValueError) as exc: - debug_log(f"DistributedListSplitter dynamic mode: invalid item index '{item_idx}' ignored: {exc}") - continue - - return pulled - - async def _split_dynamic(self, image_list, multi_job_id, is_worker, master_url, worker_id): - if not image_list: - return ([],) - - if is_worker: - pulled_indices = await self._pull_items_from_master(multi_job_id, master_url, worker_id) - else: - await self._ensure_local_dynamic_queue(multi_job_id, len(image_list)) - pulled_indices = await self._pull_items_from_local_queue(multi_job_id, worker_id) - - selected_images = [ - image_list[index] - for index in pulled_indices - if isinstance(index, int) and 0 <= index < len(image_list) - ] - debug_log( - f"DistributedListSplitter dynamic mode: role={'worker' if is_worker else 'master'} " - f"pulled={len(selected_images)} items for job {multi_job_id}" - ) - return (selected_images,) diff --git a/nodes/utilities.py b/nodes/utilities.py index bc43bfc..143b8bf 100644 --- a/nodes/utilities.py +++ b/nodes/utilities.py @@ -3,6 +3,8 @@ import json from typing import Any from ..utils.logging import debug_log, log +from ..utils.worker_ids import coerce_enabled_worker_ids, parse_worker_index, worker_value_key +from .hidden_inputs import build_worker_identity_hidden_inputs def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]: @@ -40,8 +42,7 @@ class DistributedSeed: }), }, "hidden": { - "is_worker": ("BOOLEAN", {"default": False}), - "worker_id": ("STRING", {"default": ""}), + **build_worker_identity_hidden_inputs(), }, } @@ -55,30 +56,24 @@ class DistributedSeed: seed: int, is_worker: bool = False, worker_id: str = "", + enabled_worker_ids: str = "[]", ) -> tuple[int]: if not is_worker: # Master node: pass through original values debug_log(f"Distributor - Master: seed={seed}") return (seed,) else: - # Worker node: apply offset based on worker index - # Find worker index from enabled_worker_ids - try: - # Worker IDs are passed as "worker_0", "worker_1", etc. - if worker_id.startswith("worker_"): - worker_index = int(worker_id.split("_")[1]) - else: - # Fallback: try to parse as direct index - worker_index = int(worker_id) - + enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids) + worker_index = parse_worker_index(worker_id, enabled_workers) + if worker_index is not None: offset = worker_index + 1 new_seed = seed + offset debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}") return (new_seed,) - except (ValueError, IndexError) as e: - debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}") - # Fallback: return original seed - return (seed,) + + debug_log(f"Distributor - Error parsing worker_id '{worker_id}': no worker index resolved") + # Fallback: return original seed + return (seed,) # Define ByPassTypeTuple for flexible return types @@ -105,8 +100,7 @@ class DistributedValue: "worker_values": ("STRING", {"default": "{}"}), }, "hidden": { - "is_worker": ("BOOLEAN", {"default": False}), - "worker_id": ("STRING", {"default": ""}), + **build_worker_identity_hidden_inputs(), }, } @@ -149,12 +143,14 @@ class DistributedValue: worker_values: str | dict[str, Any] = "{}", is_worker: bool = False, worker_id: str = "", + enabled_worker_ids: str = "[]", ) -> tuple[Any]: values = {} value_type = "STRING" try: - values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values + raw_values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values + values = dict(raw_values) if isinstance(raw_values, dict) else {} if not isinstance(values, dict): values = {} except json.JSONDecodeError as e: @@ -165,26 +161,23 @@ class DistributedValue: value_type = values.get("_type", inferred_type) if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}: value_type = inferred_type - values["_type"] = value_type coerced_default = self._coerce_safe(default_value, value_type) if not is_worker: debug_log(f"DistributedValue - Master: returning default '{coerced_default}'") return (coerced_default,) - try: - if worker_id.startswith("worker_"): - idx = int(worker_id.split("_")[1]) - else: - idx = int(worker_id) - key = str(idx + 1) # worker_0 → key "1" (1-indexed) - raw = values.get(key, "") - if raw: + enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids) + direct_key = str(worker_id).strip() + lookup_key = direct_key if direct_key in values else worker_value_key(worker_id, enabled_workers) + raw = values.get(lookup_key) + if raw is not None and raw != "": + try: coerced = self._coerce(raw, value_type) - debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'") + debug_log(f"DistributedValue - Worker key {lookup_key}: returning '{coerced}'") return (coerced,) - except (ValueError, IndexError) as e: - debug_log(f"DistributedValue - Error: {e}") + except (TypeError, ValueError) as e: + debug_log(f"DistributedValue - Error coercing worker value for key {lookup_key}: {e}") debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'") return (coerced_default,) diff --git a/pyproject.toml b/pyproject.toml index dde9500..3f0c660 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,10 @@ name = "ComfyUI-Distributed" description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud" version = "1.4.1" license = {file = "LICENSE"} -dependencies = [] +dependencies = [ + "aiohttp>=3.9,<4", + "Pillow>=10,<12", +] [project.urls] Repository = "https://github.com/robertvoy/ComfyUI-Distributed" diff --git a/tests/api/harness.py b/tests/api/harness.py new file mode 100644 index 0000000..042dbaa --- /dev/null +++ b/tests/api/harness.py @@ -0,0 +1,163 @@ +"""Shared stubs/utilities for API route module unit tests.""" +from __future__ import annotations + +import sys +import types +from typing import Any, Callable + + +class RoutesStub: + """Simple decorator-style route registrar used in unit tests.""" + + def get(self, _path: str): + return lambda fn: fn + + def post(self, _path: str): + return lambda fn: fn + + +class ClientTimeoutStub: + def __init__(self, total: float | None = None): + self.total = total + + +class WSMsgTypeStub: + TEXT = "TEXT" + ERROR = "ERROR" + CLOSED = "CLOSED" + + +class WebSocketResponseStub: + def __init__(self, *args: Any, **kwargs: Any): + self.args = args + self.kwargs = kwargs + + async def prepare(self, _request: Any) -> None: + return None + + async def send_json(self, _payload: Any) -> None: + return None + + def __aiter__(self): + async def _empty(): + if False: + yield None + + return _empty() + + +class FormDataStub: + def add_field(self, *_args: Any, **_kwargs: Any) -> None: + return None + + +def reset_package_namespace(package_name: str) -> None: + """Clear test package modules from sys.modules.""" + prefix = f"{package_name}." + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(prefix): + del sys.modules[mod_name] + + +def ensure_namespace_package(module_name: str) -> types.ModuleType: + """Create a namespace-like package module in sys.modules.""" + module = types.ModuleType(module_name) + module.__path__ = [] + sys.modules[module_name] = module + return module + + +def bootstrap_test_package( + package_name: str, + *, + with_api: bool = True, + with_utils: bool = True, + with_workers: bool = False, + with_upscale: bool = False, + with_orchestration: bool = False, +) -> None: + """Create a clean package scaffold for module-level import tests.""" + reset_package_namespace(package_name) + ensure_namespace_package(package_name) + + if with_api: + ensure_namespace_package(f"{package_name}.api") + if with_utils: + ensure_namespace_package(f"{package_name}.utils") + if with_workers: + ensure_namespace_package(f"{package_name}.workers") + if with_upscale: + ensure_namespace_package(f"{package_name}.upscale") + if with_orchestration: + ensure_namespace_package(f"{package_name}.api.orchestration") + + +def install_server_stub(prompt_server_instance: Any | None = None) -> Any: + """Install a minimal `server` module exposing PromptServer.instance.routes.""" + if prompt_server_instance is None: + prompt_server_instance = types.SimpleNamespace() + if not hasattr(prompt_server_instance, "routes"): + prompt_server_instance.routes = RoutesStub() + + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance) + sys.modules["server"] = server_module + return prompt_server_instance + + +def install_aiohttp_stub( + json_response_factory: Callable[..., Any], +) -> bool: + """Install/augment `aiohttp` with common members used by API tests.""" + created = False + if "aiohttp" not in sys.modules: + created = True + aiohttp_module = types.ModuleType("aiohttp") + sys.modules["aiohttp"] = aiohttp_module + + aiohttp_module = sys.modules["aiohttp"] + if not hasattr(aiohttp_module, "ClientTimeout"): + aiohttp_module.ClientTimeout = ClientTimeoutStub + if not hasattr(aiohttp_module, "WSMsgType"): + aiohttp_module.WSMsgType = WSMsgTypeStub + if not hasattr(aiohttp_module, "FormData"): + aiohttp_module.FormData = FormDataStub + + web_obj = getattr(aiohttp_module, "web", None) + if web_obj is None: + web_obj = types.SimpleNamespace() + aiohttp_module.web = web_obj + web_obj.json_response = json_response_factory + if not hasattr(web_obj, "WebSocketResponse"): + web_obj.WebSocketResponse = WebSocketResponseStub + + return created + + +def cleanup_optional_module(module_name: str, created: bool) -> None: + """Remove stubbed module only when this test created it.""" + if created: + sys.modules.pop(module_name, None) + + +def install_request_guards_stub(package_name: str) -> None: + """Install permissive request guard for isolated route tests.""" + request_guards_module = types.ModuleType(f"{package_name}.api.request_guards") + + async def _authorization_error_or_none(_request): + return None + + request_guards_module.authorization_error_or_none = _authorization_error_or_none + sys.modules[f"{package_name}.api.request_guards"] = request_guards_module + + +def install_endpoint_policy_passthrough_stub(package_name: str) -> None: + """Install endpoint policy stub that executes operation directly.""" + endpoint_policy_module = types.ModuleType(f"{package_name}.api.endpoint_policy") + + async def _run_authorized_endpoint(_request, operation, unexpected_status=500): + _ = unexpected_status + return await operation() + + endpoint_policy_module.run_authorized_endpoint = _run_authorized_endpoint + sys.modules[f"{package_name}.api.endpoint_policy"] = endpoint_policy_module diff --git a/tests/api/test_auth_guard_integration.py b/tests/api/test_auth_guard_integration.py new file mode 100644 index 0000000..7d63dcb --- /dev/null +++ b/tests/api/test_auth_guard_integration.py @@ -0,0 +1,423 @@ +import asyncio +import copy +import importlib.util +import sys +import types +import unittest +from contextlib import asynccontextmanager +from pathlib import Path + +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, + install_server_stub, +) + + +AUTH_TOKEN = "integration-secret" + + +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + + +class _FakeRequest: + def __init__(self, *, headers=None, json_payload=None, post_payload=None, query=None): + self.headers = headers or {} + self._json_payload = json_payload or {} + self._post_payload = post_payload or {} + self.query = query or {} + + async def json(self): + return self._json_payload + + async def post(self): + return self._post_payload + + +def _load_module(package_name: str, rel_path: str, module_name: str): + module_path = Path(__file__).resolve().parents[2] / rel_path + spec = importlib.util.spec_from_file_location(f"{package_name}.{module_name}", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _install_common_utils(package_name: str, config_data: dict): + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: copy.deepcopy(config_data) + config_module.save_config = lambda _cfg: True + + @asynccontextmanager + async def _config_transaction(): + yield config_module.load_config() + + config_module.config_transaction = _config_transaction + sys.modules[f"{package_name}.utils.config"] = config_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def _handle_api_error(_request, error, status=500): + return _FakeResponse({"status": "error", "message": str(error)}, status=status) + + network_module.handle_api_error = _handle_api_error + network_module.normalize_host = lambda value: value + network_module.build_worker_url = lambda worker, path="": f"http://{worker.get('host', '127.0.0.1')}{path}" + network_module.get_client_session = lambda: None + network_module.probe_worker = lambda *_args, **_kwargs: {"ok": True} + sys.modules[f"{package_name}.utils.network"] = network_module + + +def _load_real_guard_stack(package_name: str): + _load_module(package_name, "utils/auth.py", "utils.auth") + _load_module(package_name, "utils/parsing.py", "utils.parsing") + _load_module(package_name, "api/schemas.py", "api.schemas") + _load_module(package_name, "api/request_guards.py", "api.request_guards") + _load_module(package_name, "api/endpoint_policy.py", "api.endpoint_policy") + + +def _load_request_guards_only_module(): + package_name = "dist_auth_guard_testpkg" + config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}} + bootstrap_test_package(package_name, with_api=True, with_utils=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + request_guards = sys.modules[f"{package_name}.api.request_guards"] + schemas = sys.modules[f"{package_name}.api.schemas"] + cleanup_optional_module("aiohttp", created_aiohttp_stub) + return request_guards, schemas + + +def _load_config_routes_real_guard_module(): + package_name = "dist_config_auth_integration_testpkg" + config_data = { + "workers": [], + "master": {"host": "127.0.0.1"}, + "settings": {"distributed_api_token": AUTH_TOKEN}, + "tunnel": {}, + } + bootstrap_test_package(package_name, with_api=True, with_utils=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + module = _load_module(package_name, "api/config_routes.py", "api.config_routes") + cleanup_optional_module("aiohttp", created_aiohttp_stub) + return module + + +def _load_tunnel_routes_real_guard_module(): + package_name = "dist_tunnel_auth_integration_testpkg" + config_data = { + "workers": [], + "master": {"host": "127.0.0.1"}, + "settings": {"distributed_api_token": AUTH_TOKEN}, + "tunnel": {}, + } + bootstrap_test_package(package_name, with_api=True, with_utils=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + + cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare") + cloudflare_module.cloudflare_tunnel_manager = types.SimpleNamespace( + get_status=lambda: {"active": False}, + start_tunnel=lambda: {"active": True}, + stop_tunnel=lambda: {"active": False}, + ) + sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module + + module = _load_module(package_name, "api/tunnel_routes.py", "api.tunnel_routes") + cleanup_optional_module("aiohttp", created_aiohttp_stub) + return module + + +def _load_usdu_routes_real_guard_module(): + package_name = "dist_usdu_auth_integration_testpkg" + config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}} + bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + + usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management") + usdu_management_module.MAX_PAYLOAD_SIZE = 1024 * 1024 + sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module + + prompt_server_holder = { + "value": types.SimpleNamespace( + distributed_tile_jobs_lock=asyncio.Lock(), + distributed_pending_tile_jobs={}, + ) + } + + job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") + job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"] + job_store_module.init_dynamic_job = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.upscale.job_store"] = job_store_module + + job_models_module = types.ModuleType(f"{package_name}.upscale.job_models") + + class _BaseJobState: + pass + + class _ImageJobState(_BaseJobState): + def __init__(self): + self.mode = "dynamic" + self.queue = asyncio.Queue() + self.pending_images = asyncio.Queue() + self.completed_images = {} + self.worker_status = {} + self.assigned_to_workers = {} + + class _TileJobState(_BaseJobState): + def __init__(self): + self.mode = "static" + self.queue = asyncio.Queue() + self.pending_tasks = asyncio.Queue() + self.completed_tasks = {} + self.worker_status = {} + self.assigned_to_workers = {} + self.batched_static = False + + job_models_module.BaseJobState = _BaseJobState + job_models_module.ImageJobState = _ImageJobState + job_models_module.TileJobState = _TileJobState + sys.modules[f"{package_name}.upscale.job_models"] = job_models_module + + parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers") + parsers_module.parse_tiles_from_form = lambda _data: [] + sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module + + module = _load_module(package_name, "api/usdu_routes.py", "api.usdu_routes") + module._prompt_server_holder = prompt_server_holder + module._ImageJobState = _ImageJobState + cleanup_optional_module("aiohttp", created_aiohttp_stub) + return module + + +def _install_torch_stub_if_missing(): + if "torch" in sys.modules: + return False + + torch_module = types.ModuleType("torch") + torch_module.cuda = types.SimpleNamespace( + is_available=lambda: False, + empty_cache=lambda: None, + ipc_collect=lambda: None, + ) + sys.modules["torch"] = torch_module + return True + + +def _load_worker_routes_real_guard_module(): + package_name = "dist_worker_auth_integration_testpkg" + config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}, "workers": []} + bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + created_torch_stub = _install_torch_stub_if_missing() + + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.CHUNK_SIZE = 65536 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + workers_module = types.ModuleType(f"{package_name}.workers") + workers_module.get_worker_manager = lambda: types.SimpleNamespace( + processes={}, + save_processes=lambda: None, + is_process_running=lambda _pid: False, + ) + sys.modules[f"{package_name}.workers"] = workers_module + + detection_module = types.ModuleType(f"{package_name}.workers.detection") + detection_module.get_machine_id = lambda: "machine-id" + detection_module.is_docker_environment = lambda: False + detection_module.is_runpod_environment = lambda: False + sys.modules[f"{package_name}.workers.detection"] = detection_module + + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + + class _PromptValidationError(Exception): + def __init__(self, message="validation error"): + super().__init__(message) + self.validation_error = {} + self.node_errors = {} + + async def _queue_prompt_payload(*_args, **_kwargs): + return "prompt-id" + + async_helpers_module.PromptValidationError = _PromptValidationError + async_helpers_module.queue_prompt_payload = _queue_prompt_payload + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + module = _load_module(package_name, "api/worker_routes.py", "api.worker_routes") + cleanup_optional_module("aiohttp", created_aiohttp_stub) + if created_torch_stub: + sys.modules.pop("torch", None) + return module + + +def _load_job_routes_real_guard_module(): + package_name = "dist_job_auth_integration_testpkg" + config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}} + bootstrap_test_package(package_name, with_api=True, with_utils=True) + install_server_stub() + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + created_torch_stub = _install_torch_stub_if_missing() + + _install_common_utils(package_name, config_data) + _load_real_guard_stack(package_name) + + image_module = types.ModuleType(f"{package_name}.utils.image") + image_module.pil_to_tensor = lambda _img: None + image_module.ensure_contiguous = lambda tensor: tensor + sys.modules[f"{package_name}.utils.image"] = image_module + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.JOB_INIT_GRACE_PERIOD = 0.1 + constants_module.MEMORY_CLEAR_DELAY = 0.0 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state") + runtime_state = types.SimpleNamespace( + distributed_jobs_lock=asyncio.Lock(), + distributed_pending_jobs={}, + ) + runtime_state_module.ensure_distributed_runtime_state = lambda: runtime_state + sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module + + orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration") + orchestration_module.orchestrate_distributed_execution = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.api.queue_orchestration"] = orchestration_module + + queue_request_module = types.ModuleType(f"{package_name}.api.queue_request") + queue_request_module.parse_queue_request_payload = lambda data: data + sys.modules[f"{package_name}.api.queue_request"] = queue_request_module + + module = _load_module(package_name, "api/job_routes.py", "api.job_routes") + cleanup_optional_module("aiohttp", created_aiohttp_stub) + if created_torch_stub: + sys.modules.pop("torch", None) + return module + + +request_guards_module, schemas_module = _load_request_guards_only_module() +config_routes_module = _load_config_routes_real_guard_module() +tunnel_routes_module = _load_tunnel_routes_real_guard_module() +usdu_routes_module = _load_usdu_routes_real_guard_module() +worker_routes_module = _load_worker_routes_real_guard_module() +job_routes_module = _load_job_routes_real_guard_module() + + +class AuthGuardContractTests(unittest.IsolatedAsyncioTestCase): + async def test_is_authorized_request_missing_token_fails(self): + request = _FakeRequest(headers={}) + config = {"settings": {"distributed_api_token": AUTH_TOKEN}} + self.assertFalse(schemas_module.is_authorized_request(request, config)) + + async def test_is_authorized_request_matching_token_succeeds(self): + request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN}) + config = {"settings": {"distributed_api_token": AUTH_TOKEN}} + self.assertTrue(schemas_module.is_authorized_request(request, config)) + + async def test_authorization_error_or_none_returns_403_for_invalid_token(self): + request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: "wrong-token"}) + response = await request_guards_module.authorization_error_or_none(request) + self.assertIsNotNone(response) + self.assertEqual(response.status, 403) + + async def test_authorization_error_or_none_returns_none_for_valid_token(self): + request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN}) + response = await request_guards_module.authorization_error_or_none(request) + self.assertIsNone(response) + + +class RouteAuthIntegrationTests(unittest.IsolatedAsyncioTestCase): + def _auth_headers(self): + return {schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN} + + async def test_config_endpoint_enforces_real_guard(self): + unauthorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers={})) + self.assertEqual(unauthorized.status, 403) + + authorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers=self._auth_headers())) + self.assertEqual(authorized.status, 200) + self.assertIn("settings", authorized.payload) + + async def test_tunnel_endpoint_enforces_real_guard(self): + unauthorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers={})) + self.assertEqual(unauthorized.status, 403) + + authorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers=self._auth_headers())) + self.assertEqual(authorized.status, 200) + self.assertEqual(authorized.payload.get("status"), "success") + + async def test_usdu_endpoint_enforces_real_guard(self): + unauthorized = await usdu_routes_module.job_status_endpoint(_FakeRequest(headers={})) + self.assertEqual(unauthorized.status, 403) + + authorized = await usdu_routes_module.job_status_endpoint( + _FakeRequest(headers=self._auth_headers(), query={"multi_job_id": "job-1"}) + ) + self.assertEqual(authorized.status, 200) + self.assertIn("ready", authorized.payload) + + async def test_worker_endpoint_enforces_real_guard(self): + worker_routes_module._collect_network_info_sync = lambda: { + "interfaces": [], + "recommended_host": "127.0.0.1", + } + unauthorized = await worker_routes_module.get_network_info_endpoint(_FakeRequest(headers={})) + self.assertEqual(unauthorized.status, 403) + + authorized = await worker_routes_module.get_network_info_endpoint( + _FakeRequest(headers=self._auth_headers()) + ) + self.assertEqual(authorized.status, 200) + self.assertEqual(authorized.payload.get("status"), "success") + + async def test_job_endpoint_enforces_real_guard(self): + unauthorized = await job_routes_module.prepare_job_endpoint( + _FakeRequest(headers={}, json_payload={"multi_job_id": "job-1"}) + ) + self.assertEqual(unauthorized.status, 403) + + authorized = await job_routes_module.prepare_job_endpoint( + _FakeRequest(headers=self._auth_headers(), json_payload={"multi_job_id": "job-1"}) + ) + self.assertEqual(authorized.status, 200) + self.assertEqual(authorized.payload.get("status"), "success") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/api/test_config_routes.py b/tests/api/test_config_routes.py index a435547..63f703a 100644 --- a/tests/api/test_config_routes.py +++ b/tests/api/test_config_routes.py @@ -3,9 +3,19 @@ import importlib.util import sys import types import unittest +from contextlib import asynccontextmanager from pathlib import Path from unittest.mock import patch +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, + install_endpoint_policy_passthrough_stub, + install_request_guards_stub, + install_server_stub, +) + class _FakeResponse: def __init__(self, payload, status=200): @@ -25,45 +35,18 @@ def _load_config_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py" package_name = "dist_api_config_testpkg" - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] + bootstrap_test_package(package_name, with_api=True, with_utils=True) - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg + schemas_module = types.ModuleType(f"{package_name}.api.schemas") + schemas_module.is_authorized_request = lambda _request, _config: True + sys.modules[f"{package_name}.api.schemas"] = schemas_module - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg - - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status) - ) - sys.modules["aiohttp"] = aiohttp_module - - class _Routes: - def get(self, _path): - def _decorator(fn): - return fn - return _decorator - - def post(self, _path): - def _decorator(fn): - return fn - return _decorator - - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) - sys.modules["server"] = server_module + install_request_guards_stub(package_name) + install_endpoint_policy_passthrough_stub(package_name) + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + install_server_stub() logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_log = lambda *_args, **_kwargs: None @@ -73,7 +56,16 @@ def _load_config_routes_module(): network_module = types.ModuleType(f"{package_name}.utils.network") async def _handle_api_error(_request, error, status=500): - return _FakeResponse({"status": "error", "message": str(error)}, status=status) + if isinstance(error, list): + message = "; ".join(str(item) for item in error) + error_payload = [str(item) for item in error] + else: + message = str(error) + error_payload = str(error) + return _FakeResponse( + {"status": "error", "error": error_payload, "message": message}, + status=status, + ) network_module.handle_api_error = _handle_api_error network_module.normalize_host = lambda value: value @@ -89,6 +81,12 @@ def _load_config_routes_module(): config_module = types.ModuleType(f"{package_name}.utils.config") config_module.load_config = lambda: copy.deepcopy(default_config) config_module.save_config = lambda _cfg: True + + @asynccontextmanager + async def _config_transaction(): + yield config_module.load_config() + + config_module.config_transaction = _config_transaction sys.modules[f"{package_name}.utils.config"] = config_module spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path) @@ -96,8 +94,7 @@ def _load_config_routes_module(): assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) + cleanup_optional_module("aiohttp", created_aiohttp_stub) return module @@ -118,9 +115,7 @@ class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase): async def test_update_config_valid_field_persists(self): cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}} - with patch.object(config_routes, "load_config", return_value=cfg), patch.object( - config_routes, "save_config", return_value=True - ): + with patch.object(config_routes, "load_config", return_value=cfg): response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True})) self.assertEqual(response.status, 200) diff --git a/tests/api/test_job_routes.py b/tests/api/test_job_routes.py index b553624..85790c5 100644 --- a/tests/api/test_job_routes.py +++ b/tests/api/test_job_routes.py @@ -11,6 +11,14 @@ from unittest.mock import AsyncMock, patch import numpy as np import torch +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, + install_request_guards_stub, + install_server_stub, +) + class _FakeResponse: def __init__(self, payload, status=200): @@ -19,8 +27,9 @@ class _FakeResponse: class _FakeRequest: - def __init__(self, payload): + def __init__(self, payload, headers=None): self._payload = payload + self.headers = headers or {} async def json(self): return self._payload @@ -30,53 +39,19 @@ def _load_job_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py" package_name = "dist_api_queue_testpkg" - # Reset package namespace to avoid stale module state across test runs. - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + bootstrap_test_package(package_name, with_api=True, with_utils=True) # aiohttp.web stub - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status) - ) - sys.modules["aiohttp"] = aiohttp_module + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) # server module stub with route decorators - class _Routes: - def get(self, _path): - def _decorator(fn): - return fn - return _decorator - - def post(self, _path): - def _decorator(fn): - return fn - return _decorator - prompt_server_instance = types.SimpleNamespace( - routes=_Routes(), distributed_jobs_lock=None, distributed_pending_jobs={}, ) - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance) - sys.modules["server"] = server_module + install_server_stub(prompt_server_instance) # torch stub (only needed to satisfy import) created_torch_stub = False @@ -157,6 +132,47 @@ def _load_job_routes_module(): constants_module.JOB_INIT_GRACE_PERIOD = 10.0 sys.modules[f"{package_name}.utils.constants"] = constants_module + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: {"settings": {}} + sys.modules[f"{package_name}.utils.config"] = config_module + + runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state") + + def _ensure_distributed_runtime_state(server_instance=None): + ps = server_instance or prompt_server_instance + if not hasattr(ps, "distributed_pending_jobs"): + ps.distributed_pending_jobs = {} + if not hasattr(ps, "distributed_jobs_lock") or ps.distributed_jobs_lock is None: + ps.distributed_jobs_lock = asyncio.Lock() + if not hasattr(ps, "distributed_job_allowed_workers"): + ps.distributed_job_allowed_workers = {} + if not hasattr(ps, "distributed_pending_tile_jobs"): + ps.distributed_pending_tile_jobs = {} + if not hasattr(ps, "distributed_tile_jobs_lock"): + ps.distributed_tile_jobs_lock = asyncio.Lock() + return types.SimpleNamespace( + distributed_pending_jobs=ps.distributed_pending_jobs, + distributed_jobs_lock=ps.distributed_jobs_lock, + distributed_job_allowed_workers=ps.distributed_job_allowed_workers, + distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs, + distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock, + ) + + runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state + runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance + sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module + + schemas_module = types.ModuleType(f"{package_name}.api.schemas") + schemas_module.is_authorized_request = lambda _request, _config: True + schemas_module.require_bool_literal = ( + lambda value, field_name="value": value + if isinstance(value, bool) + else (_ for _ in ()).throw(ValueError(f"Field '{field_name}' must be a boolean literal.")) + ) + sys.modules[f"{package_name}.api.schemas"] = schemas_module + + install_request_guards_stub(package_name) + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local") sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module @@ -172,7 +188,6 @@ def _load_job_routes_module(): client_id: str delegate_master: object enabled_worker_ids: list - auto_prepare: bool trace_execution_id: object def _parse_queue_request_payload(data): @@ -193,7 +208,6 @@ def _load_job_routes_module(): client_id=client_id, delegate_master=data.get("delegate_master"), enabled_worker_ids=enabled, - auto_prepare=bool(data.get("auto_prepare", True)), trace_execution_id=data.get("trace_execution_id"), ) @@ -206,13 +220,11 @@ def _load_job_routes_module(): assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) - if created_torch_stub: - sys.modules.pop("torch", None) + cleanup_optional_module("aiohttp", created_aiohttp_stub) + cleanup_optional_module("torch", created_torch_stub) if created_pil_stub: - sys.modules.pop("PIL.Image", None) - sys.modules.pop("PIL", None) + cleanup_optional_module("PIL.Image", True) + cleanup_optional_module("PIL", True) return module @@ -239,7 +251,6 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(response.status, 200) self.assertEqual(response.payload.get("prompt_id"), "prompt_123") - self.assertTrue(response.payload.get("auto_prepare_supported")) async def test_distributed_queue_missing_prompt_returns_400(self): request = _FakeRequest( @@ -276,8 +287,10 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase): async def test_job_complete_accepts_audio_payload(self): queue = asyncio.Queue() - job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock() - job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue} + prompt_server = job_routes.server.PromptServer.instance + prompt_server.distributed_jobs_lock = asyncio.Lock() + prompt_server.distributed_pending_jobs = {"job-1": queue} + prompt_server.distributed_job_allowed_workers = {"job-1": {"worker-1"}} request = _FakeRequest( { "job_id": "job-1", @@ -300,6 +313,28 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(queued["audio"]["sample_rate"], 44100) self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4)) + async def test_job_complete_rejects_worker_not_in_job_allowlist(self): + queue = asyncio.Queue() + prompt_server = job_routes.server.PromptServer.instance + prompt_server.distributed_jobs_lock = asyncio.Lock() + prompt_server.distributed_pending_jobs = {"job-allow": queue} + prompt_server.distributed_job_allowed_workers = {"job-allow": {"worker-expected"}} + request = _FakeRequest( + { + "job_id": "job-allow", + "worker_id": "worker-unexpected", + "batch_idx": 0, + "image": "data:image/png;base64,AAAA", + "is_last": True, + } + ) + + with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"): + response = await job_routes.job_complete_endpoint(request) + + self.assertEqual(response.status, 403) + self.assertIn("unauthorized", response.payload.get("message", "").lower()) + def test_decode_audio_payload_rejects_bad_shape(self): bad = { "sample_rate": 44100, diff --git a/tests/api/test_media_sync.py b/tests/api/test_media_sync.py index 840e52a..284fb3a 100644 --- a/tests/api/test_media_sync.py +++ b/tests/api/test_media_sync.py @@ -4,30 +4,29 @@ import types import unittest from pathlib import Path +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, +) + + +class _AiohttpResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + def _load_media_sync_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py" package_name = "dist_ms_testpkg" - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - orch_pkg = types.ModuleType(f"{package_name}.api.orchestration") - orch_pkg.__path__ = [] - sys.modules[f"{package_name}.api.orchestration"] = orch_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + bootstrap_test_package( + package_name, + with_api=True, + with_utils=True, + with_orchestration=True, + ) logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_log = lambda *_args, **_kwargs: None @@ -48,22 +47,17 @@ def _load_media_sync_module(): trace_module.trace_info = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.trace_logger"] = trace_module - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") + auth_module = types.ModuleType(f"{package_name}.utils.auth") + auth_module.distributed_auth_headers = lambda _config=None: {} + sys.modules[f"{package_name}.utils.auth"] = auth_module - class _ClientTimeout: - def __init__(self, total=None): - pass + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: {"settings": {}} + sys.modules[f"{package_name}.utils.config"] = config_module - class _FormData: - def add_field(self, *args, **kwargs): - pass - - aiohttp_module.ClientTimeout = _ClientTimeout - aiohttp_module.FormData = _FormData - sys.modules["aiohttp"] = aiohttp_module + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _AiohttpResponse(payload, status=status) + ) spec = importlib.util.spec_from_file_location( f"{package_name}.api.orchestration.media_sync", @@ -73,8 +67,7 @@ def _load_media_sync_module(): assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) + cleanup_optional_module("aiohttp", created_aiohttp_stub) return module @@ -260,3 +253,7 @@ class RewritePromptMediaInputsTests(unittest.TestCase): if __name__ == "__main__": unittest.main() + class _AiohttpResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status diff --git a/tests/api/test_queue_orchestration.py b/tests/api/test_queue_orchestration.py index d6977ef..6ee5fd8 100644 --- a/tests/api/test_queue_orchestration.py +++ b/tests/api/test_queue_orchestration.py @@ -5,32 +5,22 @@ import types import unittest from pathlib import Path +from tests.api.harness import ( + bootstrap_test_package, + install_server_stub, +) + def _bootstrap_package(package_name): - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration") - orchestration_pkg.__path__ = [] - sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + bootstrap_test_package( + package_name, + with_api=True, + with_utils=True, + with_orchestration=True, + ) prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188) - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance) - sys.modules["server"] = server_module + install_server_stub(prompt_server_instance) async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") @@ -65,9 +55,37 @@ def _bootstrap_package(package_name): trace_logger_module.trace_debug = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module + runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state") + + def _ensure_distributed_runtime_state(server_instance=None): + ps = server_instance or prompt_server_instance + if not hasattr(ps, "distributed_pending_jobs"): + ps.distributed_pending_jobs = {} + if not hasattr(ps, "distributed_jobs_lock"): + ps.distributed_jobs_lock = asyncio.Lock() + if not hasattr(ps, "distributed_job_allowed_workers"): + ps.distributed_job_allowed_workers = {} + if not hasattr(ps, "distributed_pending_tile_jobs"): + ps.distributed_pending_tile_jobs = {} + if not hasattr(ps, "distributed_tile_jobs_lock"): + ps.distributed_tile_jobs_lock = asyncio.Lock() + return types.SimpleNamespace( + distributed_pending_jobs=ps.distributed_pending_jobs, + distributed_jobs_lock=ps.distributed_jobs_lock, + distributed_job_allowed_workers=ps.distributed_job_allowed_workers, + distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs, + distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock, + ) + + runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state + runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance + sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module + schemas_module = types.ModuleType(f"{package_name}.api.schemas") - schemas_module.parse_positive_int = lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default - schemas_module.parse_positive_float = ( + schemas_module.coerce_positive_int = ( + lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default + ) + schemas_module.coerce_positive_float = ( lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 else default ) sys.modules[f"{package_name}.api.schemas"] = schemas_module diff --git a/tests/api/test_tunnel_routes.py b/tests/api/test_tunnel_routes.py index 1767909..1e017cc 100644 --- a/tests/api/test_tunnel_routes.py +++ b/tests/api/test_tunnel_routes.py @@ -4,6 +4,14 @@ import types import unittest from pathlib import Path +from tests.api.harness import ( + bootstrap_test_package, + install_aiohttp_stub, + install_endpoint_policy_passthrough_stub, + install_request_guards_stub, + install_server_stub, +) + class _FakeResponse: def __init__(self, payload, status=200): @@ -15,38 +23,9 @@ def _load_tunnel_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py" package_name = "dist_tunnel_routes_testpkg" - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg - - class _Routes: - def get(self, _path): - return lambda fn: fn - - def post(self, _path): - return lambda fn: fn - - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) - sys.modules["server"] = server_module - - aiohttp_module = types.ModuleType("aiohttp") - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status) - ) - sys.modules["aiohttp"] = aiohttp_module + bootstrap_test_package(package_name, with_api=True, with_utils=True) + install_server_stub() + install_aiohttp_stub(lambda payload, status=200: _FakeResponse(payload, status=status)) cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare") @@ -80,6 +59,9 @@ def _load_tunnel_routes_module(): network_module.handle_api_error = _handle_api_error sys.modules[f"{package_name}.utils.network"] = network_module + install_request_guards_stub(package_name) + install_endpoint_policy_passthrough_stub(package_name) + spec = importlib.util.spec_from_file_location(f"{package_name}.api.tunnel_routes", module_path) module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None diff --git a/tests/api/test_usdu_routes.py b/tests/api/test_usdu_routes.py index e74a2a6..1a31d7e 100644 --- a/tests/api/test_usdu_routes.py +++ b/tests/api/test_usdu_routes.py @@ -9,6 +9,14 @@ from pathlib import Path from PIL import Image +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, + install_request_guards_stub, + install_server_stub, +) + class _FakeResponse: def __init__(self, payload, status=200): @@ -30,43 +38,30 @@ class _FakeRequest: return self._post_payload -class _Routes: - def post(self, _path): - def _decorator(fn): - return fn - - return _decorator - - def get(self, _path): - def _decorator(fn): - return fn - - return _decorator - - def _load_usdu_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py" package_name = "dist_api_usdu_testpkg" - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] + bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True) - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg + schemas_module = types.ModuleType(f"{package_name}.api.schemas") - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg + def _require_bool_literal(value, field_name: str = "value"): + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized == "true": + return True + if normalized == "false": + return False + raise ValueError(f"Field '{field_name}' must be a boolean literal.") - upscale_pkg = types.ModuleType(f"{package_name}.upscale") - upscale_pkg.__path__ = [] - sys.modules[f"{package_name}.upscale"] = upscale_pkg + schemas_module.require_bool_literal = _require_bool_literal + schemas_module.is_authorized_request = lambda _request, _config: True + sys.modules[f"{package_name}.api.schemas"] = schemas_module - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + install_request_guards_stub(package_name) prompt_server_holder = { "value": types.SimpleNamespace( @@ -75,23 +70,23 @@ def _load_usdu_routes_module(): ) } - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status) - ) - sys.modules["aiohttp"] = aiohttp_module - - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) - sys.modules["server"] = server_module + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + install_server_stub() logging_module = types.ModuleType(f"{package_name}.utils.logging") logging_module.debug_log = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.logging"] = logging_module + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: {"settings": {}} + sys.modules[f"{package_name}.utils.config"] = config_module + + usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management") + usdu_management_module.MAX_PAYLOAD_SIZE = 1024 + sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module + network_module = types.ModuleType(f"{package_name}.utils.network") async def _handle_api_error(_request, error, status=500): @@ -174,8 +169,7 @@ def _load_usdu_routes_module(): assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) + cleanup_optional_module("aiohttp", created_aiohttp_stub) module.web = types.SimpleNamespace( json_response=lambda payload, status=200: _FakeResponse(payload, status=status) @@ -238,7 +232,8 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase): response = await usdu_routes.request_image_endpoint(request) self.assertEqual(response.status, 200) - self.assertEqual(response.payload.get("image_idx"), 7) + self.assertEqual(response.payload.get("kind"), "image") + self.assertEqual(response.payload.get("task_idx"), 7) self.assertEqual(response.payload.get("estimated_remaining"), 0) self.assertEqual(job_data.assigned_to_workers["worker-a"], [7]) self.assertIn("worker-a", job_data.worker_status) @@ -254,38 +249,11 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase): response = await usdu_routes.request_image_endpoint(request) self.assertEqual(response.status, 200) - self.assertEqual(response.payload.get("tile_idx"), 4) + self.assertEqual(response.payload.get("kind"), "tile") + self.assertEqual(response.payload.get("task_idx"), 4) self.assertTrue(response.payload.get("batched_static")) self.assertEqual(job_data.assigned_to_workers["worker-a"], [4]) - async def test_init_list_queue_initializes_pending_items(self): - request = _FakeRequest( - json_payload={ - "multi_job_id": "list-job-1", - "list_size": 3, - "enabled_workers": ["worker-a", "worker-b"], - } - ) - response = await usdu_routes.init_list_queue_endpoint(request) - - self.assertEqual(response.status, 200) - self.assertEqual(response.payload.get("status"), "success") - self.assertEqual(response.payload.get("remaining"), 3) - - async def test_request_list_item_assigns_next_index(self): - prompt_server = usdu_routes._prompt_server_holder["value"] - job_data = usdu_routes._ImageJobState("list-job-2") - await job_data.pending_images.put(11) - prompt_server.distributed_pending_tile_jobs["list-job-2"] = job_data - - request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "list-job-2"}) - response = await usdu_routes.request_list_item_endpoint(request) - - self.assertEqual(response.status, 200) - self.assertEqual(response.payload.get("item_idx"), 11) - self.assertEqual(response.payload.get("estimated_remaining"), 0) - self.assertEqual(job_data.assigned_to_workers["worker-a"], [11]) - async def test_submit_tiles_completion_signal_enqueues_last_marker(self): prompt_server = usdu_routes._prompt_server_holder["value"] job_data = usdu_routes._TileJobState("job-4") diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index c016d1d..6efde16 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -8,6 +8,14 @@ from collections import deque from pathlib import Path from unittest.mock import patch +from tests.api.harness import ( + bootstrap_test_package, + cleanup_optional_module, + install_aiohttp_stub, + install_request_guards_stub, + install_server_stub, +) + class _FakeResponse: def __init__(self, payload, status=200): @@ -16,10 +24,11 @@ class _FakeResponse: class _FakeRequest: - def __init__(self, payload=None, match_info=None, query=None): + def __init__(self, payload=None, match_info=None, query=None, headers=None): self._payload = payload self.match_info = match_info or {} self.query = query or {} + self.headers = headers or {} async def json(self): return self._payload @@ -49,8 +58,8 @@ class _FakeHTTPClientSession: self._status = status self.calls = [] - def get(self, url, params=None, timeout=None): - self.calls.append({"url": url, "params": params, "timeout": timeout}) + def get(self, url, params=None, headers=None, timeout=None): + self.calls.append({"url": url, "params": params, "headers": headers, "timeout": timeout}) return _FakeHTTPClientResponse(self._payload, status=self._status) @@ -67,7 +76,7 @@ class _DummyWorkerManager: } return 12345 - def _is_process_running(self, _pid): + def is_process_running(self, _pid): return False def save_processes(self): @@ -89,21 +98,7 @@ def _load_worker_routes_module(): module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py" package_name = "dist_api_worker_testpkg" - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - api_pkg = types.ModuleType(f"{package_name}.api") - api_pkg.__path__ = [] - sys.modules[f"{package_name}.api"] = api_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True) workers_pkg = types.ModuleType(f"{package_name}.workers") workers_pkg.__path__ = [] @@ -119,59 +114,10 @@ def _load_worker_routes_module(): detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan" sys.modules[f"{package_name}.workers.detection"] = detection_module - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") - - class _ClientTimeout: - def __init__(self, total=None): - self.total = total - - class _WSMsgType: - TEXT = "TEXT" - ERROR = "ERROR" - CLOSED = "CLOSED" - - class _WebSocketResponse: - def __init__(self, *args, **kwargs): - self.args = args - self.kwargs = kwargs - - async def prepare(self, _request): - return None - - async def send_json(self, _payload): - return None - - def __aiter__(self): - async def _empty(): - if False: - yield None - return _empty() - - aiohttp_module.ClientTimeout = _ClientTimeout - aiohttp_module.WSMsgType = _WSMsgType - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status), - WebSocketResponse=_WebSocketResponse, - ) - sys.modules["aiohttp"] = aiohttp_module - - class _Routes: - def get(self, _path): - def _decorator(fn): - return fn - return _decorator - - def post(self, _path): - def _decorator(fn): - return fn - return _decorator - - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) - sys.modules["server"] = server_module + created_aiohttp_stub = install_aiohttp_stub( + lambda payload, status=200: _FakeResponse(payload, status=status) + ) + install_server_stub() created_torch_stub = False if "torch" not in sys.modules: @@ -248,19 +194,29 @@ def _load_worker_routes_module(): def _validate_worker_id(worker_id, config): return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", [])) + def _require_worker_id(worker_id, config, field_name="worker_id"): + _ = field_name + worker_id_str = str(worker_id).strip() + if any(str(worker.get("id")) == worker_id_str for worker in config.get("workers", [])): + return worker_id_str + raise ValueError(f"Worker {worker_id_str} not found") + schemas_module.require_fields = _require_fields schemas_module.validate_worker_id = _validate_worker_id + schemas_module.require_worker_id = _require_worker_id + schemas_module.is_authorized_request = lambda _request, _config: True + schemas_module.distributed_auth_headers = lambda _config: {} sys.modules[f"{package_name}.api.schemas"] = schemas_module + install_request_guards_stub(package_name) + spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path) module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) - if created_torch_stub: - sys.modules.pop("torch", None) + cleanup_optional_module("aiohttp", created_aiohttp_stub) + cleanup_optional_module("torch", created_torch_stub) return module diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index af5d908..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,2 +0,0 @@ -# This conftest.py marks the tests/ directory as the pytest collection root, -# preventing pytest from traversing into the parent package's __init__.py. diff --git a/tests/unit/nodes/test_distributed_upscale.py b/tests/unit/nodes/test_distributed_upscale.py index 5c055b5..cb14d17 100644 --- a/tests/unit/nodes/test_distributed_upscale.py +++ b/tests/unit/nodes/test_distributed_upscale.py @@ -37,6 +37,14 @@ def _bootstrap_package(package_name): logging_module.log = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.logging"] = logging_module + worker_ids_module = types.ModuleType(f"{package_name}.utils.worker_ids") + worker_ids_module.coerce_enabled_worker_ids = lambda value: ( + [str(v) for v in value] + if isinstance(value, list) + else [] + ) + sys.modules[f"{package_name}.utils.worker_ids"] = worker_ids_module + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") async_helpers_module.run_async_in_server_loop = lambda coro, timeout=None: (_ for _ in ()).throw( RuntimeError(f"run_async_in_server_loop should not be used in these tests (timeout={timeout})") @@ -139,6 +147,7 @@ def _load_distributed_upscale_module(): package_name = "dist_distributed_upscale_testpkg" _bootstrap_package(package_name) _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs") + _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts") _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args") return _load_module(package_name, "nodes/distributed_upscale.py", "nodes.distributed_upscale") diff --git a/tests/unit/nodes/test_distributed_value.py b/tests/unit/nodes/test_distributed_value.py index 53809ff..e66a128 100644 --- a/tests/unit/nodes/test_distributed_value.py +++ b/tests/unit/nodes/test_distributed_value.py @@ -28,6 +28,18 @@ class DistributedValueTests(unittest.TestCase): root_pkg.__path__ = [] sys.modules[pkg_name] = root_pkg + nodes_pkg = types.ModuleType(f"{pkg_name}.nodes") + nodes_pkg.__path__ = [] + sys.modules[f"{pkg_name}.nodes"] = nodes_pkg + + hidden_inputs_mod = types.ModuleType(f"{pkg_name}.nodes.hidden_inputs") + hidden_inputs_mod.build_worker_identity_hidden_inputs = lambda: { + "is_worker": ("BOOLEAN", {"default": False}), + "enabled_worker_ids": ("STRING", {"default": "[]"}), + "worker_id": ("STRING", {"default": ""}), + } + sys.modules[f"{pkg_name}.nodes.hidden_inputs"] = hidden_inputs_mod + utils_pkg = types.ModuleType(f"{pkg_name}.utils") utils_pkg.__path__ = [] sys.modules[f"{pkg_name}.utils"] = utils_pkg @@ -37,6 +49,16 @@ class DistributedValueTests(unittest.TestCase): logging_mod.log = lambda *_a, **_k: None sys.modules[f"{pkg_name}.utils.logging"] = logging_mod + worker_ids_path = Path(__file__).resolve().parents[3] / "utils" / "worker_ids.py" + worker_ids_spec = importlib.util.spec_from_file_location( + f"{pkg_name}.utils.worker_ids", + worker_ids_path, + ) + worker_ids_mod = importlib.util.module_from_spec(worker_ids_spec) + assert worker_ids_spec is not None and worker_ids_spec.loader is not None + worker_ids_spec.loader.exec_module(worker_ids_mod) + sys.modules[f"{pkg_name}.utils.worker_ids"] = worker_ids_mod + spec = importlib.util.spec_from_file_location( f"{pkg_name}.nodes.utilities", module_path ) diff --git a/tests/unit/nodes/test_list_collector.py b/tests/unit/nodes/test_list_collector.py deleted file mode 100644 index 19843ef..0000000 --- a/tests/unit/nodes/test_list_collector.py +++ /dev/null @@ -1,237 +0,0 @@ -import asyncio -import base64 -import importlib.util -import io -import json -import sys -import types -import unittest -from pathlib import Path - -import torch -from PIL import Image - - -def _bootstrap_package(package_name): - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - nodes_pkg = types.ModuleType(f"{package_name}.nodes") - nodes_pkg.__path__ = [] - sys.modules[f"{package_name}.nodes"] = nodes_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg - - logging_module = types.ModuleType(f"{package_name}.utils.logging") - logging_module.debug_log = lambda *_args, **_kwargs: None - logging_module.log = lambda *_args, **_kwargs: None - sys.modules[f"{package_name}.utils.logging"] = logging_module - - config_module = types.ModuleType(f"{package_name}.utils.config") - config_module.get_worker_timeout_seconds = lambda: 0.2 - config_module.is_master_delegate_only = lambda: False - sys.modules[f"{package_name}.utils.config"] = config_module - - constants_module = types.ModuleType(f"{package_name}.utils.constants") - constants_module.HEARTBEAT_INTERVAL = 1.0 - sys.modules[f"{package_name}.utils.constants"] = constants_module - - image_module = types.ModuleType(f"{package_name}.utils.image") - - def ensure_contiguous(tensor): - return tensor.contiguous() - - def tensor_to_pil(image_batch, index): - tensor = image_batch[index].detach().cpu().clamp(0, 1) - array = (tensor.numpy() * 255.0).astype("uint8") - return Image.fromarray(array) - - def encode_tensor_png_data_url(image_batch, batch_index=0): - image = tensor_to_pil(image_batch, batch_index) - byte_io = io.BytesIO() - image.save(byte_io, format="PNG", compress_level=0) - encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8") - return f"data:image/png;base64,{encoded}" - - image_module.ensure_contiguous = ensure_contiguous - image_module.tensor_to_pil = tensor_to_pil - image_module.encode_tensor_png_data_url = encode_tensor_png_data_url - sys.modules[f"{package_name}.utils.image"] = image_module - - network_module = types.ModuleType(f"{package_name}.utils.network") - - async def get_client_session(): - raise RuntimeError("test should patch get_client_session") - - network_module.get_client_session = get_client_session - sys.modules[f"{package_name}.utils.network"] = network_module - - async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") - async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro) - sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module - - -def _load_module(package_name, module_rel_path, module_name): - module_path = Path(__file__).resolve().parents[3] / module_rel_path - spec = importlib.util.spec_from_file_location( - f"{package_name}.{module_name}", - module_path, - ) - module = importlib.util.module_from_spec(spec) - assert spec is not None and spec.loader is not None - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def _load_list_collector_module(): - package_name = "dist_list_collector_testpkg" - _bootstrap_package(package_name) - _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs") - _load_module(package_name, "nodes/runtime_helpers.py", "nodes.runtime_helpers") - _load_module(package_name, "nodes/queue_wait.py", "nodes.queue_wait") - _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids") - return _load_module(package_name, "nodes/list_collector.py", "nodes.list_collector") - - -collector_module = _load_list_collector_module() - - -def _context(**overrides): - defaults = { - "multi_job_id": "", - "is_worker": False, - "master_url": "", - "enabled_worker_ids": "[]", - "worker_id": "", - "delegate_only": False, - } - defaults.update(overrides) - return collector_module.ListCollectorRunContext(**defaults) - - -class _FakeResponseCtx: - def __init__(self, recorder, payload, url): - self._recorder = recorder - self._payload = payload - self._url = url - - async def __aenter__(self): - self._recorder.append({"url": self._url, "payload": self._payload}) - return self - - async def __aexit__(self, exc_type, exc, tb): - return False - - def raise_for_status(self): - return None - - -class _FakeSession: - def __init__(self): - self.calls = [] - - def post(self, url, json=None, timeout=None): - return _FakeResponseCtx(self.calls, json, url) - - -class _FakePromptServer: - def __init__(self): - self.distributed_pending_jobs = {} - self.distributed_jobs_lock = asyncio.Lock() - - -class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase): - def setUp(self): - self.node = collector_module.DistributedListCollector() - - async def test_worker_mode_sends_images_to_master(self): - fake_session = _FakeSession() - - async def _get_session(): - return fake_session - - collector_module.get_client_session = _get_session - - images = [ - torch.zeros((1, 4, 4, 3), dtype=torch.float32), - torch.ones((1, 4, 4, 3), dtype=torch.float32), - ] - - result = await self.node.execute( - images, - context=_context( - multi_job_id="job-list", - is_worker=True, - master_url="http://master.local:8188", - worker_id="worker-a", - ), - ) - - self.assertEqual(len(result[0]), 2) - self.assertEqual(len(fake_session.calls), 2) - self.assertEqual(fake_session.calls[0]["url"], "http://master.local:8188/distributed/job_complete") - self.assertEqual(fake_session.calls[0]["payload"]["batch_idx"], 0) - self.assertFalse(fake_session.calls[0]["payload"]["is_last"]) - self.assertEqual(fake_session.calls[1]["payload"]["batch_idx"], 1) - self.assertTrue(fake_session.calls[1]["payload"]["is_last"]) - - async def test_master_mode_collects_results_as_list(self): - fake_server = _FakePromptServer() - collector_module._get_prompt_server_instance = lambda: fake_server - - queue = asyncio.Queue() - worker_a = torch.full((1, 4, 4, 3), 0.25, dtype=torch.float32) - worker_b = torch.full((1, 4, 4, 3), 0.75, dtype=torch.float32) - await queue.put({"tensor": worker_a, "worker_id": "worker-a", "image_index": 0, "is_last": True}) - await queue.put({"tensor": worker_b, "worker_id": "worker-b", "image_index": 0, "is_last": True}) - fake_server.distributed_pending_jobs["job-list"] = queue - - master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)] - - result = await self.node.execute( - master_images, - context=_context( - multi_job_id="job-list", - enabled_worker_ids=json.dumps(["worker-a", "worker-b"]), - ), - ) - - self.assertEqual(len(result[0]), 3) - self.assertTrue(torch.equal(result[0][0], master_images[0])) - self.assertTrue(torch.equal(result[0][1], worker_a)) - self.assertTrue(torch.equal(result[0][2], worker_b)) - - async def test_delegate_only_mode_returns_worker_images_only(self): - fake_server = _FakePromptServer() - collector_module._get_prompt_server_instance = lambda: fake_server - - queue = asyncio.Queue() - worker_img = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32) - await queue.put({"tensor": worker_img, "worker_id": "worker-a", "image_index": 0, "is_last": True}) - fake_server.distributed_pending_jobs["job-list"] = queue - - master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)] - - result = await self.node.execute( - master_images, - context=_context( - multi_job_id="job-list", - enabled_worker_ids=json.dumps(["worker-a"]), - delegate_only=True, - ), - ) - - self.assertEqual(len(result[0]), 1) - self.assertTrue(torch.equal(result[0][0], worker_img)) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/nodes/test_list_splitter.py b/tests/unit/nodes/test_list_splitter.py deleted file mode 100644 index 849447f..0000000 --- a/tests/unit/nodes/test_list_splitter.py +++ /dev/null @@ -1,261 +0,0 @@ -import asyncio -import importlib.util -import sys -import types -import unittest -from dataclasses import dataclass, field -from pathlib import Path - -import torch - - -def _bootstrap_package(package_name): - for mod_name in list(sys.modules): - if mod_name == package_name or mod_name.startswith(f"{package_name}."): - del sys.modules[mod_name] - - root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] - sys.modules[package_name] = root_pkg - - nodes_pkg = types.ModuleType(f"{package_name}.nodes") - nodes_pkg.__path__ = [] - sys.modules[f"{package_name}.nodes"] = nodes_pkg - - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg - - upscale_pkg = types.ModuleType(f"{package_name}.upscale") - upscale_pkg.__path__ = [] - sys.modules[f"{package_name}.upscale"] = upscale_pkg - - logging_module = types.ModuleType(f"{package_name}.utils.logging") - logging_module.debug_log = lambda *_args, **_kwargs: None - logging_module.log = lambda *_args, **_kwargs: None - sys.modules[f"{package_name}.utils.logging"] = logging_module - - async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") - async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro) - sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module - - network_module = types.ModuleType(f"{package_name}.utils.network") - - async def _get_client_session(): - raise RuntimeError("test should patch get_client_session") - - network_module.get_client_session = _get_client_session - sys.modules[f"{package_name}.utils.network"] = network_module - - job_models_module = types.ModuleType(f"{package_name}.upscale.job_models") - - class BaseJobState: - pass - - @dataclass - class ImageJobState(BaseJobState): - multi_job_id: str - mode: str = field(default="dynamic", init=False) - queue: asyncio.Queue = field(default_factory=asyncio.Queue) - pending_images: asyncio.Queue = field(default_factory=asyncio.Queue) - completed_images: dict = field(default_factory=dict) - worker_status: dict = field(default_factory=dict) - assigned_to_workers: dict = field(default_factory=dict) - - @property - def pending_tasks(self): - return self.pending_images - - job_models_module.BaseJobState = BaseJobState - job_models_module.ImageJobState = ImageJobState - sys.modules[f"{package_name}.upscale.job_models"] = job_models_module - - prompt_server_holder = { - "value": types.SimpleNamespace( - distributed_tile_jobs_lock=asyncio.Lock(), - distributed_pending_tile_jobs={}, - ) - } - - job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") - - def ensure_tile_jobs_initialized(): - return prompt_server_holder["value"] - - async def init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None): - prompt_server = prompt_server_holder["value"] - async with prompt_server.distributed_tile_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_tile_jobs: - return - job_data = ImageJobState(multi_job_id=multi_job_id) - for idx in (all_indices if all_indices is not None else list(range(batch_size))): - await job_data.pending_images.put(idx) - prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data - - job_store_module.ensure_tile_jobs_initialized = ensure_tile_jobs_initialized - job_store_module.init_dynamic_job = init_dynamic_job - job_store_module._prompt_server_holder = prompt_server_holder - sys.modules[f"{package_name}.upscale.job_store"] = job_store_module - - -def _load_module(package_name, module_rel_path, module_name): - module_path = Path(__file__).resolve().parents[3] / module_rel_path - spec = importlib.util.spec_from_file_location( - f"{package_name}.{module_name}", - module_path, - ) - module = importlib.util.module_from_spec(spec) - assert spec is not None and spec.loader is not None - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def _load_list_splitter_module(): - package_name = "dist_list_splitter_testpkg" - _bootstrap_package(package_name) - _load_module(package_name, "nodes/utilities.py", "nodes.utilities") - module = _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter") - return module - - -splitter_module = _load_list_splitter_module() - - -class _FakeResponse: - def __init__(self, payload, status=200): - self.payload = payload - self.status = status - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc, tb): - return False - - async def json(self): - return self.payload - - -class _FakeSession: - def __init__(self, responses): - self.responses = list(responses) - self.calls = [] - - def post(self, url, json=None, timeout=None): - self.calls.append({"url": url, "json": json}) - if not self.responses: - return _FakeResponse({"item_idx": None}, status=200) - payload = self.responses.pop(0) - return _FakeResponse(payload, status=200) - - -class DistributedListSplitterTests(unittest.TestCase): - def setUp(self): - self.node = splitter_module.DistributedListSplitter() - prompt_server = splitter_module.ensure_tile_jobs_initialized() - prompt_server.distributed_pending_tile_jobs = {} - - def _images(self, count): - return [torch.full((1, 4, 4, 3), float(idx), dtype=torch.float32) for idx in range(count)] - - def test_single_participant_returns_full_list(self): - images = self._images(5) - out = self.node.split(images, mode="static", participant_index=0, total_participants=1)[0] - self.assertEqual(len(out), 5) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in out], [0, 1, 2, 3, 4]) - - def test_two_participants_split_six_images_evenly(self): - images = self._images(6) - p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0] - p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0] - self.assertEqual(len(p0), 3) - self.assertEqual(len(p1), 3) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2]) - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4, 5]) - - def test_three_participants_split_seven_images(self): - images = self._images(7) - p0 = self.node.split(images, mode="static", participant_index=0, total_participants=3)[0] - p1 = self.node.split(images, mode="static", participant_index=1, total_participants=3)[0] - p2 = self.node.split(images, mode="static", participant_index=2, total_participants=3)[0] - self.assertEqual([len(p0), len(p1), len(p2)], [3, 2, 2]) - - def test_more_participants_than_images_yields_empty_slices(self): - images = self._images(2) - outputs = [self.node.split(images, mode="static", participant_index=i, total_participants=5)[0] for i in range(5)] - self.assertEqual([len(chunk) for chunk in outputs], [1, 1, 0, 0, 0]) - - def test_mixed_dimensions_are_preserved(self): - images = [ - torch.zeros((1, 4, 4, 3), dtype=torch.float32), - torch.zeros((1, 8, 6, 3), dtype=torch.float32), - torch.zeros((1, 3, 9, 3), dtype=torch.float32), - ] - p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0] - p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0] - self.assertEqual([tuple(t.shape) for t in p0], [(1, 4, 4, 3), (1, 8, 6, 3)]) - self.assertEqual([tuple(t.shape) for t in p1], [(1, 3, 9, 3)]) - - def test_dynamic_mode_master_pulls_from_local_queue(self): - images = self._images(4) - - pulled = self.node.split( - images, - mode="dynamic", - multi_job_id="job-dynamic", - is_worker=False, - worker_id="master", - )[0] - - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [0, 1, 2, 3]) - - def test_dynamic_mode_worker_pulls_via_http(self): - images = self._images(5) - fake_session = _FakeSession([ - {"item_idx": 1, "estimated_remaining": 3}, - {"item_idx": 3, "estimated_remaining": 2}, - {"item_idx": None}, - ]) - - async def _get_session(): - return fake_session - - splitter_module.get_client_session = _get_session - - pulled = self.node.split( - images, - mode="dynamic", - multi_job_id="job-dynamic", - is_worker=True, - master_url="http://master.local:8188", - worker_id="worker-a", - )[0] - - self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [1, 3]) - self.assertGreaterEqual(len(fake_session.calls), 3) - - def test_dynamic_mode_exhausted_queue_returns_empty(self): - images = self._images(3) - - first_pull = self.node.split( - images, - mode="dynamic", - multi_job_id="job-dynamic", - is_worker=False, - worker_id="master", - )[0] - second_pull = self.node.split( - images, - mode="dynamic", - multi_job_id="job-dynamic", - is_worker=False, - worker_id="master", - )[0] - - self.assertEqual(len(first_pull), 3) - self.assertEqual(second_pull, []) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/system/test_config.py b/tests/unit/system/test_config.py index fa05043..4ac0729 100644 --- a/tests/unit/system/test_config.py +++ b/tests/unit/system/test_config.py @@ -21,14 +21,13 @@ def _load_config_module(): root_pkg.__path__ = [] sys.modules[package_name] = root_pkg - logging_module = types.ModuleType(f"{package_name}.logging") - logging_module.log = lambda *_args, **_kwargs: None - logging_module.debug_log = lambda *_args, **_kwargs: None - sys.modules[f"{package_name}.logging"] = logging_module - constants_module = types.ModuleType(f"{package_name}.constants") constants_module.GPU_CONFIG_FILE = "gpu_config.json" constants_module.HEARTBEAT_TIMEOUT = 30 + constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8 + constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4 + constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2 + constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120 sys.modules[f"{package_name}.constants"] = constants_module spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path) diff --git a/tests/unit/system/test_distributed.py b/tests/unit/system/test_distributed.py index 49f2ce5..3cf868b 100644 --- a/tests/unit/system/test_distributed.py +++ b/tests/unit/system/test_distributed.py @@ -1,21 +1,17 @@ import importlib.util -import os import sys import types import unittest from pathlib import Path -def _load_distributed_module(is_worker): +def _load_distributed_module(): module_path = Path(__file__).resolve().parents[3] / "distributed.py" package_name = "dist_distributed_entry_testpkg" calls = { - "ensure_config_exists": 0, - "ensure_distributed_state": 0, - "ensure_tile_jobs_initialized": 0, - "delayed_auto_launch": 0, - "register_async_signals": 0, - "sync_cleanup_registered": 0, + "build_node_mappings": 0, + "initialize_runtime": 0, + "initialize_runtime_args": [], } for mod_name in list(sys.modules): @@ -26,126 +22,49 @@ def _load_distributed_module(is_worker): root_pkg.__path__ = [] sys.modules[package_name] = root_pkg - # Absolute dependency used by module. - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace()) - sys.modules["server"] = server_module + bootstrap_pkg = types.ModuleType(f"{package_name}.bootstrap") + bootstrap_pkg.__path__ = [] + sys.modules[f"{package_name}.bootstrap"] = bootstrap_pkg - utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] - sys.modules[f"{package_name}.utils"] = utils_pkg + entrypoint_module = types.ModuleType(f"{package_name}.bootstrap.entrypoint") - config_module = types.ModuleType(f"{package_name}.utils.config") - config_module.ensure_config_exists = lambda: calls.__setitem__("ensure_config_exists", calls["ensure_config_exists"] + 1) - sys.modules[f"{package_name}.utils.config"] = config_module + def _build_node_mappings(): + calls["build_node_mappings"] += 1 + return ({"MockNode": object}, {"MockNode": "Mock Node"}) - logging_module = types.ModuleType(f"{package_name}.utils.logging") - logging_module.debug_log = lambda *_args, **_kwargs: None - sys.modules[f"{package_name}.utils.logging"] = logging_module + def _initialize_runtime(prompt_server=None): + calls["initialize_runtime"] += 1 + calls["initialize_runtime_args"].append(prompt_server) - network_module = types.ModuleType(f"{package_name}.utils.network") + entrypoint_module.build_node_mappings = _build_node_mappings + entrypoint_module.initialize_runtime = _initialize_runtime + sys.modules[f"{package_name}.bootstrap.entrypoint"] = entrypoint_module - async def _cleanup_client_session(): - return None - - network_module.cleanup_client_session = _cleanup_client_session - sys.modules[f"{package_name}.utils.network"] = network_module - - workers_pkg = types.ModuleType(f"{package_name}.workers") - workers_pkg.__path__ = [] - workers_pkg.get_worker_manager = lambda: None - sys.modules[f"{package_name}.workers"] = workers_pkg - - startup_module = types.ModuleType(f"{package_name}.workers.startup") - startup_module.delayed_auto_launch = lambda: calls.__setitem__("delayed_auto_launch", calls["delayed_auto_launch"] + 1) - startup_module.register_async_signals = lambda: calls.__setitem__("register_async_signals", calls["register_async_signals"] + 1) - startup_module.sync_cleanup = lambda: None - sys.modules[f"{package_name}.workers.startup"] = startup_module - - upscale_pkg = types.ModuleType(f"{package_name}.upscale") - upscale_pkg.__path__ = [] - sys.modules[f"{package_name}.upscale"] = upscale_pkg - - job_store_module = types.ModuleType(f"{package_name}.upscale.job_store") - job_store_module.ensure_tile_jobs_initialized = lambda: calls.__setitem__( - "ensure_tile_jobs_initialized", calls["ensure_tile_jobs_initialized"] + 1 - ) - sys.modules[f"{package_name}.upscale.job_store"] = job_store_module - - nodes_module = types.ModuleType(f"{package_name}.nodes") - nodes_module.NODE_CLASS_MAPPINGS = {} - nodes_module.NODE_DISPLAY_NAME_MAPPINGS = {} - nodes_module.ImageBatchDivider = object - nodes_module.DistributedCollectorNode = object - nodes_module.DistributedSeed = object - nodes_module.DistributedModelName = object - nodes_module.DistributedValue = object - nodes_module.AudioBatchDivider = object - nodes_module.DistributedEmptyImage = object - nodes_module.DistributedListSplitter = object - nodes_module.DistributedListCollector = object - nodes_module.DistributedBranch = object - nodes_module.DistributedBranchCollector = object - nodes_module.AnyType = object - nodes_module.ByPassTypeTuple = tuple - nodes_module.any_type = "*" - sys.modules[f"{package_name}.nodes"] = nodes_module - - api_module = types.ModuleType(f"{package_name}.api") - sys.modules[f"{package_name}.api"] = api_module - - queue_orch_module = types.ModuleType(f"{package_name}.api.queue_orchestration") - queue_orch_module.ensure_distributed_state = lambda _ps: calls.__setitem__( - "ensure_distributed_state", calls["ensure_distributed_state"] + 1 - ) - sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orch_module - - import atexit - - original_register = atexit.register - atexit.register = lambda fn: calls.__setitem__("sync_cleanup_registered", calls["sync_cleanup_registered"] + 1) or fn - - original_env = os.environ.get("COMFYUI_IS_WORKER") - try: - if is_worker: - os.environ["COMFYUI_IS_WORKER"] = "1" - else: - os.environ.pop("COMFYUI_IS_WORKER", None) - - spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path) - module = importlib.util.module_from_spec(spec) - assert spec is not None and spec.loader is not None - sys.modules[spec.name] = module - spec.loader.exec_module(module) - finally: - atexit.register = original_register - if original_env is None: - os.environ.pop("COMFYUI_IS_WORKER", None) - else: - os.environ["COMFYUI_IS_WORKER"] = original_env - - return calls + spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module, calls class DistributedEntryTests(unittest.TestCase): - def test_import_initializes_core_state_for_worker(self): - calls = _load_distributed_module(is_worker=True) + def test_import_builds_mappings_without_runtime_bootstrap(self): + module, calls = _load_distributed_module() - self.assertEqual(calls["ensure_config_exists"], 1) - self.assertEqual(calls["ensure_distributed_state"], 1) - self.assertEqual(calls["ensure_tile_jobs_initialized"], 1) - self.assertEqual(calls["delayed_auto_launch"], 0) - self.assertEqual(calls["register_async_signals"], 0) + self.assertIn("MockNode", module.NODE_CLASS_MAPPINGS) + self.assertEqual(module.NODE_DISPLAY_NAME_MAPPINGS["MockNode"], "Mock Node") + self.assertEqual(calls["build_node_mappings"], 1) + self.assertEqual(calls["initialize_runtime"], 0) - def test_import_triggers_master_startup_hooks(self): - calls = _load_distributed_module(is_worker=False) + def test_initialize_runtime_delegates_to_bootstrap_entrypoint(self): + module, calls = _load_distributed_module() + sentinel = object() - self.assertEqual(calls["ensure_config_exists"], 1) - self.assertEqual(calls["ensure_distributed_state"], 1) - self.assertEqual(calls["ensure_tile_jobs_initialized"], 1) - self.assertEqual(calls["delayed_auto_launch"], 1) - self.assertEqual(calls["register_async_signals"], 1) - self.assertGreaterEqual(calls["sync_cleanup_registered"], 2) + module.initialize_runtime(sentinel) + + self.assertEqual(calls["initialize_runtime"], 1) + self.assertEqual(calls["initialize_runtime_args"], [sentinel]) if __name__ == "__main__": diff --git a/tests/unit/system/test_logging.py b/tests/unit/system/test_logging.py index 587fb03..cdf172e 100644 --- a/tests/unit/system/test_logging.py +++ b/tests/unit/system/test_logging.py @@ -1,63 +1,29 @@ import io -import json -import os -import tempfile import unittest from contextlib import redirect_stdout +from unittest.mock import patch from utils import logging as logging_utils class LoggingUtilsTests(unittest.TestCase): - def setUp(self): - self.original_config_file = logging_utils.GPU_CONFIG_FILE - self.state = logging_utils._debug_state() - self.original_cache = self.state.enabled - self.original_cache_time = self.state.updated_at + def test_is_debug_enabled_delegates_to_config(self): + with patch.object(logging_utils, "get_debug_enabled", return_value=True) as mocked: + self.assertTrue(logging_utils.is_debug_enabled()) + mocked.assert_called_once_with(default=False) - def tearDown(self): - logging_utils.GPU_CONFIG_FILE = self.original_config_file - self.state.enabled = self.original_cache - self.state.updated_at = self.original_cache_time - - def _reset_cache(self): - self.state.enabled = None - self.state.updated_at = 0.0 - - def test_is_debug_enabled_reads_config_flag(self): - with tempfile.NamedTemporaryFile("w", delete=False) as temp_file: - json.dump({"settings": {"debug": True}}, temp_file) - temp_path = temp_file.name - self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path)) - - logging_utils.GPU_CONFIG_FILE = temp_path - self._reset_cache() - - self.assertTrue(logging_utils.is_debug_enabled()) - - def test_is_debug_enabled_handles_invalid_json(self): - with tempfile.NamedTemporaryFile("w", delete=False) as temp_file: - temp_file.write("{invalid-json") - temp_path = temp_file.name - self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path)) - - logging_utils.GPU_CONFIG_FILE = temp_path - self._reset_cache() - - self.assertFalse(logging_utils.is_debug_enabled()) + def test_is_debug_enabled_handles_config_errors(self): + with patch.object(logging_utils, "get_debug_enabled", side_effect=RuntimeError("boom")): + self.assertFalse(logging_utils.is_debug_enabled()) def test_debug_log_emits_only_when_enabled(self): - self.state.enabled = True - self.state.updated_at = 1e18 output = io.StringIO() - with redirect_stdout(output): + with patch.object(logging_utils, "is_debug_enabled", return_value=True), redirect_stdout(output): logging_utils.debug_log("hello") self.assertIn("[Distributed] hello", output.getvalue()) - self.state.enabled = False - self.state.updated_at = 1e18 output = io.StringIO() - with redirect_stdout(output): + with patch.object(logging_utils, "is_debug_enabled", return_value=False), redirect_stdout(output): logging_utils.debug_log("hidden") self.assertEqual(output.getvalue(), "") diff --git a/tests/unit/system/test_prompt_transform.py b/tests/unit/system/test_prompt_transform.py index ee92728..d6158bd 100644 --- a/tests/unit/system/test_prompt_transform.py +++ b/tests/unit/system/test_prompt_transform.py @@ -81,16 +81,6 @@ def _delegate_prompt(): } -def _list_splitter_prompt(): - """1(List source) → 2(DistributedListSplitter) → 3(DistributedListCollector) → 4(SaveImage).""" - return { - "1": {"class_type": "ImageListSource", "inputs": {}}, - "2": {"class_type": "DistributedListSplitter", "inputs": {"images": ["1", 0]}}, - "3": {"class_type": "DistributedListCollector", "inputs": {"images": ["2", 0]}}, - "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}}, - } - - def _branch_prompt(): """1 → 2(DistributedBranch) -> slot0:3->4, slot1:5->6.""" return { @@ -285,14 +275,6 @@ class PrunePromptForWorkerTests(unittest.TestCase): self.assertIn("2", result) self.assertNotIn("3", result) - def test_list_collector_is_treated_as_distributed_anchor(self): - prompt = _list_splitter_prompt() - result = pt.prune_prompt_for_worker(prompt) - self.assertIn("1", result) - self.assertIn("2", result) - self.assertIn("3", result) - self.assertNotIn("4", result) - def test_branch_anchor_keeps_downstream_for_later_branch_pruning(self): prompt = _branch_prompt() result = pt.prune_prompt_for_worker(prompt) @@ -385,22 +367,6 @@ class GenerateJobIdMapTests(unittest.TestCase): job_map = pt.generate_job_id_map(idx, "run") self.assertEqual(job_map["5"], "run_5") - def test_maps_list_collector_nodes(self): - prompt = { - "8": {"class_type": "DistributedListCollector", "inputs": {}}, - } - idx = pt.PromptIndex(prompt) - job_map = pt.generate_job_id_map(idx, "run") - self.assertEqual(job_map["8"], "run_8") - - def test_maps_list_splitter_nodes(self): - prompt = { - "7": {"class_type": "DistributedListSplitter", "inputs": {}}, - } - idx = pt.PromptIndex(prompt) - job_map = pt.generate_job_id_map(idx, "run") - self.assertEqual(job_map["7"], "run_7") - def test_maps_branch_collector_nodes(self): prompt = { "10": {"class_type": "DistributedBranchCollector", "inputs": {}}, @@ -508,9 +474,9 @@ class ApplyOverridesSeedTests(unittest.TestCase): result = _apply(self._seed_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) - def test_worker_id_reflects_index_in_enabled_list(self): + def test_worker_id_uses_canonical_worker_id(self): result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"]) - self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1") + self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b") def test_master_sets_is_worker_false(self): result = _apply(self._seed_prompt(), "master") @@ -566,9 +532,9 @@ class ApplyOverridesValueTests(unittest.TestCase): result = _apply(self._value_prompt(), "worker-a") self.assertTrue(result["1"]["inputs"]["is_worker"]) - def test_worker_id_reflects_index_in_enabled_list(self): + def test_worker_id_uses_canonical_worker_id(self): result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"]) - self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1") + self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b") def test_master_sets_is_worker_false(self): result = _apply(self._value_prompt(), "master") @@ -579,62 +545,6 @@ class ApplyOverridesValueTests(unittest.TestCase): self.assertEqual(result["1"]["inputs"]["worker_id"], "") -class ApplyOverridesListSplitterTests(unittest.TestCase): - def _splitter_prompt(self): - return {"1": {"class_type": "DistributedListSplitter", "inputs": {}}} - - def test_master_gets_participant_index_zero(self): - result = _apply(self._splitter_prompt(), "master", enabled_worker_ids=["worker-a", "worker-b"]) - self.assertEqual(result["1"]["inputs"]["participant_index"], 0) - self.assertEqual(result["1"]["inputs"]["total_participants"], 3) - self.assertFalse(result["1"]["inputs"]["is_worker"]) - self.assertNotIn("master_url", result["1"]["inputs"]) - self.assertEqual(result["1"]["inputs"]["worker_id"], "master") - self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1") - - def test_worker_gets_incremented_participant_index(self): - result = _apply(self._splitter_prompt(), "worker-a", enabled_worker_ids=["worker-a", "worker-b"]) - self.assertEqual(result["1"]["inputs"]["participant_index"], 1) - self.assertEqual(result["1"]["inputs"]["total_participants"], 3) - self.assertTrue(result["1"]["inputs"]["is_worker"]) - self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a") - self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com") - - def test_delegate_mode_shifts_worker_to_index_zero(self): - result = _apply( - self._splitter_prompt(), - "worker-a", - enabled_worker_ids=["worker-a", "worker-b"], - delegate_master=True, - ) - self.assertEqual(result["1"]["inputs"]["participant_index"], 0) - self.assertEqual(result["1"]["inputs"]["total_participants"], 2) - - -class ApplyOverridesListCollectorTests(unittest.TestCase): - def _collector_prompt(self): - return {"1": {"class_type": "DistributedListCollector", "inputs": {}}} - - def test_worker_sets_master_url_and_worker_id(self): - result = _apply(self._collector_prompt(), "worker-a", enabled_worker_ids=["worker-a"]) - self.assertTrue(result["1"]["inputs"]["is_worker"]) - self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com") - self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a") - - def test_master_sets_delegate_only_and_clears_worker_fields(self): - prompt = { - "1": { - "class_type": "DistributedListCollector", - "inputs": {"master_url": "stale", "worker_id": "stale"}, - } - } - result = _apply(prompt, "master", enabled_worker_ids=["worker-a"], delegate_master=True) - self.assertFalse(result["1"]["inputs"]["is_worker"]) - self.assertTrue(result["1"]["inputs"]["delegate_only"]) - self.assertNotIn("master_url", result["1"]["inputs"]) - self.assertNotIn("worker_id", result["1"]["inputs"]) - - class ApplyOverridesBranchTests(unittest.TestCase): def test_master_gets_branch_zero_worker_gets_branch_one(self): prompt = _branch_prompt() diff --git a/tests/unit/system/test_queue_request.py b/tests/unit/system/test_queue_request.py index 156388a..1517af8 100644 --- a/tests/unit/system/test_queue_request.py +++ b/tests/unit/system/test_queue_request.py @@ -26,27 +26,26 @@ class QueueRequestPayloadTests(unittest.TestCase): def test_normalizes_enabled_worker_ids(self): payload_data = self._base_payload() - payload_data["enabled_worker_ids"] = ["a", 2, 3] + payload_data["enabled_worker_ids"] = [" worker-a ", "worker-b", "worker-a"] payload_data["delegate_master"] = True payload = parse_queue_request_payload( payload_data ) - self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"]) + self.assertEqual(payload.enabled_worker_ids, ["worker-a", "worker-b"]) self.assertTrue(payload.delegate_master) def test_supports_legacy_workers_field(self): payload_data = self._base_payload() payload_data.pop("enabled_worker_ids", None) - payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}] + payload_data["workers"] = [{"id": "w1"}, "w2", {"id": "w3"}, {"name": "no-id"}] payload = parse_queue_request_payload( payload_data ) - self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"]) + self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "w3"]) - def test_supports_auto_prepare_prompt_fallback(self): + def test_supports_workflow_prompt_fallback(self): payload_data = self._base_payload() payload_data.pop("prompt", None) - payload_data["auto_prepare"] = True payload_data["workflow"] = { "prompt": { "10": {"class_type": "DistributedCollector"}, @@ -56,7 +55,6 @@ class QueueRequestPayloadTests(unittest.TestCase): payload_data ) self.assertIn("10", payload.prompt) - self.assertTrue(payload.auto_prepare) def test_normalizes_trace_execution_id(self): payload_data = self._base_payload() @@ -74,10 +72,6 @@ class QueueRequestPayloadTests(unittest.TestCase): ) self.assertIsNone(payload.trace_execution_id) - def test_auto_prepare_defaults_true(self): - payload = parse_queue_request_payload(self._base_payload()) - self.assertTrue(payload.auto_prepare) - def test_workers_field_must_be_list(self): payload_data = self._base_payload() payload_data.pop("enabled_worker_ids", None) @@ -91,7 +85,7 @@ class QueueRequestPayloadTests(unittest.TestCase): with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"): parse_queue_request_payload(payload_data) - def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self): + def test_auto_prepare_is_ignored_for_backward_compat(self): payload_data = self._base_payload() payload_data.pop("prompt", None) payload_data["auto_prepare"] = False @@ -100,13 +94,6 @@ class QueueRequestPayloadTests(unittest.TestCase): } payload = parse_queue_request_payload(payload_data) self.assertIn("10", payload.prompt) - self.assertFalse(payload.auto_prepare) - - def test_auto_prepare_must_be_boolean(self): - payload_data = self._base_payload() - payload_data["auto_prepare"] = "true" - with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"): - parse_queue_request_payload(payload_data) def test_invalid_delegate_master_type_raises(self): payload_data = self._base_payload() @@ -120,6 +107,12 @@ class QueueRequestPayloadTests(unittest.TestCase): with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"): parse_queue_request_payload(payload_data) + def test_legacy_index_worker_tokens_are_rejected(self): + payload_data = self._base_payload() + payload_data["enabled_worker_ids"] = ["worker-a", "0", "worker_1"] + with self.assertRaisesRegex(ValueError, "legacy index token"): + parse_queue_request_payload(payload_data) + def test_invalid_top_level_payload_raises(self): with self.assertRaisesRegex(ValueError, "Expected a JSON object body"): parse_queue_request_payload(["not", "an", "object"]) diff --git a/tests/unit/upscale/test_dynamic.py b/tests/unit/upscale/test_dynamic.py index e4686f7..9e0a150 100644 --- a/tests/unit/upscale/test_dynamic.py +++ b/tests/unit/upscale/test_dynamic.py @@ -4,6 +4,7 @@ import sys import types import unittest from pathlib import Path +from types import SimpleNamespace import numpy as np from PIL import Image @@ -67,6 +68,8 @@ def _bootstrap_package(package_name): sys.modules[f"{package_name}.utils.config"] = config_module constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.JOB_POLL_INTERVAL = 0.0 + constants_module.JOB_POLL_MAX_ATTEMPTS = 3 constants_module.TILE_WAIT_TIMEOUT = 5.0 constants_module.TILE_SEND_TIMEOUT = 5.0 sys.modules[f"{package_name}.utils.constants"] = constants_module @@ -110,6 +113,7 @@ def _load_dynamic_mode_module(): package_name = "dist_dynamic_mode_testpkg" _bootstrap_package(package_name) _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args") + _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts") return _load_module(package_name, "upscale/modes/dynamic.py", "upscale.modes.dynamic") @@ -143,24 +147,33 @@ class _DummyDynamicNode(dynamic_module.DynamicModeMixin): def _poll_job_ready(self, *_args, **_kwargs): return self._poll_ready - async def _request_image_from_master(self, *_args, **_kwargs): - if self._assignments: - return self._assignments.pop(0) - return (None, 0) + async def check_job_status(self, *_args, **_kwargs): + return self._poll_ready - def _slice_conditioning(self, positive, negative, _image_idx): + async def request_assignment(self, *_args, **_kwargs): + if self._assignments: + image_idx, estimated_remaining = self._assignments.pop(0) + return SimpleNamespace( + kind="image" if image_idx is not None else "none", + task_idx=image_idx, + estimated_remaining=estimated_remaining, + batched_static=False, + ) + return SimpleNamespace(kind="none", task_idx=None, estimated_remaining=0, batched_static=False) + + def slice_conditioning(self, positive, negative, _image_idx): return positive, negative - def _process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs): + def process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs): return local_image - async def _send_heartbeat_to_master(self, *_args, **_kwargs): + async def send_heartbeat(self, *_args, **_kwargs): self.heartbeat_count += 1 - async def _send_full_image_to_master(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last): + async def send_full_image(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last): self.sent_images.append((image_idx, is_last, local_image.size)) - async def _send_worker_complete_signal(self, *_args, **_kwargs): + async def send_worker_complete_signal(self, *_args, **_kwargs): self.completion_sent = True diff --git a/tests/unit/upscale/test_job_store.py b/tests/unit/upscale/test_job_store.py index 338cb65..91f69a6 100644 --- a/tests/unit/upscale/test_job_store.py +++ b/tests/unit/upscale/test_job_store.py @@ -27,9 +27,33 @@ def _bootstrap_package(package_name): logging_module.debug_log = lambda *_args, **_kwargs: None sys.modules[f"{package_name}.utils.logging"] = logging_module - server_module = types.ModuleType("server") - server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace()) - sys.modules["server"] = server_module + prompt_server = types.SimpleNamespace() + + runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state") + + def _ensure_distributed_runtime_state(server_instance=None): + 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() + if not hasattr(ps, "distributed_job_allowed_workers"): + ps.distributed_job_allowed_workers = {} + if not hasattr(ps, "distributed_pending_tile_jobs"): + ps.distributed_pending_tile_jobs = {} + if not hasattr(ps, "distributed_tile_jobs_lock"): + ps.distributed_tile_jobs_lock = asyncio.Lock() + return types.SimpleNamespace( + distributed_pending_jobs=ps.distributed_pending_jobs, + distributed_jobs_lock=ps.distributed_jobs_lock, + distributed_job_allowed_workers=ps.distributed_job_allowed_workers, + distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs, + distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock, + ) + + runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state + runtime_state_module.get_prompt_server_instance = lambda: prompt_server + sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module def _load_module(package_name, module_rel_path, module_name): @@ -58,7 +82,7 @@ job_models_module, job_store_module = _load_job_store_modules() class JobStoreTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): - self.prompt_server = job_store_module.server.PromptServer.instance + self.prompt_server = job_store_module.get_prompt_server_instance() self.prompt_server.distributed_pending_tile_jobs = {} self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock() diff --git a/tests/unit/upscale/test_single_gpu.py b/tests/unit/upscale/test_single_gpu.py index 6a2985b..3abad75 100644 --- a/tests/unit/upscale/test_single_gpu.py +++ b/tests/unit/upscale/test_single_gpu.py @@ -105,6 +105,7 @@ def _load_single_gpu_mode_module(): _bootstrap_package(package_name) _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args") _load_module(package_name, "upscale/tile_processing.py", "upscale.tile_processing") + _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts") return _load_module(package_name, "upscale/modes/single_gpu.py", "upscale.modes.single_gpu") diff --git a/tests/unit/upscale/test_static.py b/tests/unit/upscale/test_static.py index 13d274a..f1b6e2c 100644 --- a/tests/unit/upscale/test_static.py +++ b/tests/unit/upscale/test_static.py @@ -4,6 +4,7 @@ import sys import types import unittest from pathlib import Path +from types import SimpleNamespace import torch @@ -166,6 +167,16 @@ def _load_static_mode_module(): sys.modules[tile_processing_spec.name] = tile_processing_module tile_processing_spec.loader.exec_module(tile_processing_module) + mode_contexts_path = Path(__file__).resolve().parents[3] / "upscale" / "mode_contexts.py" + mode_contexts_spec = importlib.util.spec_from_file_location( + f"{package_name}.upscale.mode_contexts", + mode_contexts_path, + ) + mode_contexts_module = importlib.util.module_from_spec(mode_contexts_spec) + assert mode_contexts_spec is not None and mode_contexts_spec.loader is not None + sys.modules[mode_contexts_spec.name] = mode_contexts_module + mode_contexts_spec.loader.exec_module(mode_contexts_module) + spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", module_path) module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None @@ -198,14 +209,20 @@ class _FakeStaticWorker(static_mode.StaticModeMixin): def _poll_job_ready(self, *_args, **_kwargs): return self.job_ready - async def _request_tile_from_master(self, *_args, **_kwargs): + async def request_assignment(self, *_args, **_kwargs): self.request_calls += 1 - return self.tile_sequence.pop(0) + tile_idx, estimated_remaining, batched_static = self.tile_sequence.pop(0) + return SimpleNamespace( + kind="tile" if tile_idx is not None else "none", + task_idx=tile_idx, + estimated_remaining=estimated_remaining, + batched_static=batched_static, + ) - async def _send_heartbeat_to_master(self, *_args, **_kwargs): + async def send_heartbeat(self, *_args, **_kwargs): self.heartbeat_calls += 1 - async def send_tiles_batch_to_master( + async def send_tiles_batch( self, processed_tiles, _multi_job_id, diff --git a/upscale/job_state.py b/upscale/job_state.py index 451ccb8..8e4b7ce 100644 --- a/upscale/job_state.py +++ b/upscale/job_state.py @@ -22,6 +22,10 @@ class JobStateMixin: return dict(job_data.completed_images) return {} + async def all_completed_tasks(self, multi_job_id): + """Public completed-task accessor.""" + return await self._get_all_completed_tasks(multi_job_id) + async def _get_next_image_index(self, multi_job_id): """Get next image index from pending queue for master.""" prompt_server = ensure_tile_jobs_initialized() @@ -39,6 +43,10 @@ class JobStateMixin: except asyncio.TimeoutError: return None + async def next_image_index(self, multi_job_id): + """Public next-image accessor.""" + return await self._get_next_image_index(multi_job_id) + async def _get_next_tile_index(self, multi_job_id): """Get next tile index from pending queue for master in static mode.""" prompt_server = ensure_tile_jobs_initialized() @@ -56,6 +64,10 @@ class JobStateMixin: except asyncio.TimeoutError: return None + async def next_tile_index(self, multi_job_id): + """Public next-tile accessor.""" + return await self._get_next_tile_index(multi_job_id) + async def _get_total_completed_count(self, multi_job_id): """Get total count of all completed images (master + workers).""" prompt_server = ensure_tile_jobs_initialized() @@ -67,6 +79,10 @@ class JobStateMixin: return len(job_data.completed_tasks) return 0 + async def total_completed_count(self, multi_job_id): + """Public completed-count accessor.""" + return await self._get_total_completed_count(multi_job_id) + async def _get_all_completed_images(self, multi_job_id): """Get all completed images.""" prompt_server = ensure_tile_jobs_initialized() @@ -76,6 +92,10 @@ class JobStateMixin: return job_data.completed_images.copy() return {} + async def all_completed_images(self, multi_job_id): + """Public completed-image accessor.""" + return await self._get_all_completed_images(multi_job_id) + async def _get_pending_count(self, multi_job_id): """Get count of pending images in the queue.""" prompt_server = ensure_tile_jobs_initialized() @@ -87,6 +107,10 @@ class JobStateMixin: return job_data.pending_tasks.qsize() return 0 + async def pending_count(self, multi_job_id): + """Public pending-count accessor.""" + return await self._get_pending_count(multi_job_id) + async def _drain_worker_results_queue(self, multi_job_id): """Drain pending worker results from queue and update completed images.""" prompt_server = ensure_tile_jobs_initialized() @@ -130,6 +154,14 @@ class JobStateMixin: return collected + async def drain_worker_results_queue(self, multi_job_id): + """Public worker-result drain API.""" + return await self._drain_worker_results_queue(multi_job_id) + async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): """Check for timed out workers and requeue their assigned images.""" return await check_and_requeue_timed_out_workers(multi_job_id, batch_size) + + async def check_and_requeue_timed_out_workers(self, multi_job_id, batch_size): + """Public timeout-requeue API.""" + return await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size) diff --git a/upscale/job_store.py b/upscale/job_store.py index 87d2ff9..f049116 100644 --- a/upscale/job_store.py +++ b/upscale/job_store.py @@ -4,9 +4,8 @@ import time from dataclasses import dataclass from typing import Any -import server - from ..utils.logging import debug_log +from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance from .job_models import BaseJobState, ImageJobState, TileJobState # Configure maximum payload size (50MB default, configurable via environment variable) @@ -25,20 +24,21 @@ class JobQueueInitConfig: def ensure_tile_jobs_initialized() -> Any: """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] + prompt_server = get_prompt_server_instance() + state = ensure_distributed_runtime_state(prompt_server) + if not isinstance(state.distributed_pending_tile_jobs, dict): + debug_log("Resetting invalid distributed_pending_tile_jobs state to empty dict.") + state.distributed_pending_tile_jobs = {} + ensure_distributed_runtime_state(prompt_server) + + invalid_job_ids = [ + job_id + for job_id, job_data in state.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 state.distributed_pending_tile_jobs[job_id] return prompt_server diff --git a/upscale/mode_contexts.py b/upscale/mode_contexts.py new file mode 100644 index 0000000..07092d9 --- /dev/null +++ b/upscale/mode_contexts.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +class TileOpsCollaborator: + """Typed facade over tile-related operations used by USDU modes.""" + + def __init__(self, delegate: Any) -> None: + self._delegate = delegate + + def round_to_multiple(self, value: int) -> int: + return self._delegate.round_to_multiple(value) + + def calculate_tiles( + self, + width: int, + height: int, + tile_width: int, + tile_height: int, + force_uniform_tiles: bool, + ) -> list[tuple[int, int]]: + return self._delegate.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + + def create_tile_mask( + self, + width: int, + height: int, + tx: int, + ty: int, + tile_width: int, + tile_height: int, + mask_blur: int, + ) -> Any: + return self._delegate.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur) + + def blend_tile( + self, + base_image: Any, + tile_pil: Any, + x1: int, + y1: int, + size: tuple[int, int], + tile_mask: Any, + padding: int, + ) -> Any: + return self._delegate.blend_tile(base_image, tile_pil, x1, y1, size, tile_mask, padding) + + def slice_conditioning(self, positive: Any, negative: Any, image_idx: int) -> tuple[Any, Any]: + return self._delegate.slice_conditioning(positive, negative, image_idx) + + def process_and_blend_tile( + self, + tile_idx: int, + tile_pos: tuple[int, int], + upscaled_image: Any, + result_image: Any, + model: Any, + positive: Any, + negative: Any, + vae: Any, + seed: int, + steps: int, + cfg: float, + sampler_name: str, + scheduler: str, + denoise: float, + tile_width: int, + tile_height: int, + padding: int, + mask_blur: int, + image_width: int, + image_height: int, + force_uniform_tiles: bool, + tiled_decode: bool, + batch_idx: int, + ) -> Any: + return self._delegate.process_and_blend_tile( + tile_idx, + tile_pos, + upscaled_image, + result_image, + model, + positive, + negative, + vae, + seed, + steps, + cfg, + sampler_name, + scheduler, + denoise, + tile_width, + tile_height, + padding, + mask_blur, + image_width, + image_height, + force_uniform_tiles, + tiled_decode, + batch_idx=batch_idx, + ) + + +class JobStateCollaborator: + """Typed facade over job-state operations used by USDU modes.""" + + def __init__(self, delegate: Any) -> None: + self._delegate = delegate + + async def next_tile_index(self, multi_job_id: str) -> int | None: + return await self._delegate.next_tile_index(multi_job_id) + + async def all_completed_tasks(self, multi_job_id: str) -> dict[int, dict[str, Any]]: + return await self._delegate.all_completed_tasks(multi_job_id) + + async def check_and_requeue_timed_out_workers(self, multi_job_id: str, expected_total: int) -> int: + return await self._delegate.check_and_requeue_timed_out_workers(multi_job_id, expected_total) + + async def next_image_index(self, multi_job_id: str) -> int | None: + return await self._delegate.next_image_index(multi_job_id) + + async def drain_worker_results_queue(self, multi_job_id: str) -> int: + return await self._delegate.drain_worker_results_queue(multi_job_id) + + async def total_completed_count(self, multi_job_id: str) -> int: + return await self._delegate.total_completed_count(multi_job_id) + + async def pending_count(self, multi_job_id: str) -> int: + return await self._delegate.pending_count(multi_job_id) + + async def all_completed_images(self, multi_job_id: str) -> dict[int, Any]: + return await self._delegate.all_completed_images(multi_job_id) + + +class WorkerCommsCollaborator: + """Typed facade over worker/master communication operations used by USDU modes.""" + + def __init__(self, delegate: Any) -> None: + self._delegate = delegate + + async def request_assignment(self, multi_job_id: str, master_url: str, worker_id: str) -> Any: + return await self._delegate.request_assignment(multi_job_id, master_url, worker_id) + + async def send_tiles_batch( + self, + processed_tiles: list[dict[str, Any]], + multi_job_id: str, + master_url: str, + padding: int, + worker_id: str, + is_final_flush: bool = False, + ) -> None: + await self._delegate.send_tiles_batch( + processed_tiles, + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=is_final_flush, + ) + + async def send_heartbeat(self, multi_job_id: str, master_url: str, worker_id: str) -> None: + await self._delegate.send_heartbeat(multi_job_id, master_url, worker_id) + + async def check_job_status(self, multi_job_id: str, master_url: str) -> bool: + return await self._delegate.check_job_status(multi_job_id, master_url) + + async def async_yield(self) -> None: + await self._delegate.async_yield() + + async def send_full_image( + self, + image_pil: Any, + image_idx: int, + multi_job_id: str, + master_url: str, + worker_id: str, + is_last: bool, + ) -> None: + await self._delegate.send_full_image( + image_pil, + image_idx, + multi_job_id, + master_url, + worker_id, + is_last, + ) + + async def send_worker_complete_signal(self, multi_job_id: str, master_url: str, worker_id: str) -> None: + await self._delegate.send_worker_complete_signal(multi_job_id, master_url, worker_id) + + +class ResultCollectorCollaborator: + """Typed facade over result-collection operations used by dynamic mode.""" + + def __init__(self, delegate: Any) -> None: + self._delegate = delegate + + async def mark_image_completed(self, multi_job_id: str, image_idx: int, image_pil: Any) -> None: + await self._delegate.mark_image_completed(multi_job_id, image_idx, image_pil) + + async def collect_dynamic_images( + self, + multi_job_id: str, + remaining_to_collect: int, + num_workers: int, + batch_size: int, + master_processed_count: int, + ) -> dict[int, Any]: + return await self._delegate.collect_dynamic_images( + multi_job_id, + remaining_to_collect, + num_workers, + batch_size, + master_processed_count, + ) + + +@dataclass(frozen=True) +class SingleGpuModeContext: + tile_ops: TileOpsCollaborator + + +@dataclass(frozen=True) +class StaticModeContext: + tile_ops: TileOpsCollaborator + job_state: JobStateCollaborator + worker_comms: WorkerCommsCollaborator + + +@dataclass(frozen=True) +class DynamicModeContext: + tile_ops: TileOpsCollaborator + job_state: JobStateCollaborator + worker_comms: WorkerCommsCollaborator + result_collector: ResultCollectorCollaborator diff --git a/upscale/modes/dynamic.py b/upscale/modes/dynamic.py index dc2998a..29c3e5b 100644 --- a/upscale/modes/dynamic.py +++ b/upscale/modes/dynamic.py @@ -1,14 +1,26 @@ from __future__ import annotations -import asyncio, torch +import asyncio, time, torch from typing import Any from PIL import Image from ...utils.logging import debug_log, log from ...utils.image import tensor_to_pil, pil_to_tensor from ...utils.async_helpers import run_async_in_server_loop from ...utils.config import get_worker_timeout_seconds -from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT +from ...utils.constants import ( + JOB_POLL_INTERVAL, + JOB_POLL_MAX_ATTEMPTS, + TILE_WAIT_TIMEOUT, + TILE_SEND_TIMEOUT, +) from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job +from ..mode_contexts import ( + DynamicModeContext, + JobStateCollaborator, + ResultCollectorCollaborator, + TileOpsCollaborator, + WorkerCommsCollaborator, +) from ..processing_args import UpscaleCoreArgs @@ -17,47 +29,82 @@ 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`). + - 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`). + - WorkerCommsMixin (`request_assignment`, `send_full_image`, `send_heartbeat`). """ + def _build_dynamic_mode_context(self) -> DynamicModeContext: + """Build explicit collaborators for dynamic mode execution.""" + return DynamicModeContext( + tile_ops=TileOpsCollaborator(self), + job_state=JobStateCollaborator(self), + worker_comms=WorkerCommsCollaborator(self), + result_collector=ResultCollectorCollaborator(self), + ) + + def _poll_job_ready_dynamic( + self, + multi_job_id: str, + master_url: str, + worker_id: str | None = None, + max_attempts: int = JOB_POLL_MAX_ATTEMPTS, + mode_context: DynamicModeContext | None = None, + ) -> bool: + """Poll master for job readiness using explicit worker-comms collaborator.""" + context = mode_context or self._build_dynamic_mode_context() + for attempt in range(max_attempts): + ready = run_async_in_server_loop( + context.worker_comms.check_job_status(multi_job_id, master_url), + timeout=5.0, + ) + if ready: + if worker_id: + debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts") + else: + debug_log(f"Job {multi_job_id} ready after {attempt} attempts") + return True + time.sleep(JOB_POLL_INTERVAL) + return False + def _handle_master_dynamic_idle_state( self, multi_job_id, batch_size, consecutive_retries, max_consecutive_retries, + mode_context: DynamicModeContext | None = None, ): """Handle queue-empty state while master waits for worker progress.""" + context = mode_context or self._build_dynamic_mode_context() drained_count = run_async_in_server_loop( - self._drain_worker_results_queue(multi_job_id), + context.job_state.drain_worker_results_queue(multi_job_id), timeout=5.0, ) - run_async_in_server_loop(self._async_yield(), timeout=0.1) + run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) requeued_count = run_async_in_server_loop( - self._check_and_requeue_timed_out_workers(multi_job_id, batch_size), + context.job_state.check_and_requeue_timed_out_workers(multi_job_id, batch_size), timeout=5.0, ) - run_async_in_server_loop(self._async_yield(), timeout=0.1) + run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) if requeued_count > 0: log(f"Requeued {requeued_count} images from timed out workers") return False, True, 0 completed_now = run_async_in_server_loop( - self._get_total_completed_count(multi_job_id), + context.job_state.total_completed_count(multi_job_id), timeout=1.0, ) log(f"USDU Dist: Images progress {completed_now}/{batch_size}") if completed_now >= batch_size: return True, False, consecutive_retries - run_async_in_server_loop(self._async_yield(), timeout=0.1) + run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) pending_count = run_async_in_server_loop( - self._get_pending_count(multi_job_id), + context.job_state.pending_count(multi_job_id), timeout=1.0, ) if pending_count > 0: @@ -82,10 +129,12 @@ class DynamicModeMixin: processed_count, result_images, upscaled_image, + mode_context: DynamicModeContext | None = None, ): """Collect remaining worker outputs and convert final images back to tensor.""" + context = mode_context or self._build_dynamic_mode_context() all_completed = run_async_in_server_loop( - self._get_all_completed_images(multi_job_id), + context.job_state.all_completed_images(multi_job_id), timeout=5.0, ) remaining_to_collect = batch_size - len(all_completed) @@ -93,7 +142,7 @@ class DynamicModeMixin: debug_log(f"Waiting for {remaining_to_collect} more images from workers") collection_timeout = float(get_worker_timeout_seconds()) collected_images = run_async_in_server_loop( - self._async_collect_dynamic_images( + context.result_collector.collect_dynamic_images( multi_job_id, remaining_to_collect, num_workers, @@ -128,8 +177,10 @@ class DynamicModeMixin: force_uniform_tiles: bool, multi_job_id: str, enabled_workers: list[str], + mode_context: DynamicModeContext | None = None, ) -> tuple[torch.Tensor]: """Dynamic mode for large batches - assigns whole images to workers dynamically, including master.""" + context = mode_context or self._build_dynamic_mode_context() # Get batch size and dimensions batch_size, height, width, _ = upscaled_image.shape num_workers = len(enabled_workers) @@ -142,7 +193,7 @@ class DynamicModeMixin: debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.") # Calculate tiles for processing - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) # Initialize job queue for communication try: @@ -167,7 +218,7 @@ class DynamicModeMixin: while processed_count < batch_size: # Try to get an image to process image_idx = run_async_in_server_loop( - self._get_next_image_index(multi_job_id), + context.job_state.next_image_index(multi_job_id), timeout=5.0 # Short timeout to allow frequent checks ) @@ -183,7 +234,7 @@ class DynamicModeMixin: image_seed = core_args.seed # Pre-slice conditioning once per image (not per tile) - positive_sliced, negative_sliced = self._slice_conditioning( + positive_sliced, negative_sliced = context.tile_ops.slice_conditioning( core_args.positive, core_args.negative, image_idx, @@ -193,7 +244,7 @@ class DynamicModeMixin: source_tensor = pil_to_tensor(local_image) if single_tensor.is_cuda: source_tensor = source_tensor.cuda() - local_image = self._process_and_blend_tile( + local_image = context.tile_ops.process_and_blend_tile( tile_idx, pos, source_tensor, local_image, core_args.model, positive_sliced, @@ -213,20 +264,20 @@ class DynamicModeMixin: ) # Yield after each tile to minimize worker downtime - run_async_in_server_loop(self._async_yield(), timeout=0.1) + run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) # Note: No per-tile drain here – that's what makes this "per-image" result_images[image_idx] = local_image # Mark as completed run_async_in_server_loop( - self._mark_image_completed(multi_job_id, image_idx, local_image), + context.result_collector.mark_image_completed(multi_job_id, image_idx, local_image), timeout=5.0 ) # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing) drained_count = run_async_in_server_loop( - self._drain_worker_results_queue(multi_job_id), + context.job_state.drain_worker_results_queue(multi_job_id), timeout=5.0 ) @@ -235,19 +286,20 @@ class DynamicModeMixin: # NEW: Log overall progress (includes master's image + any drained workers) completed_now = run_async_in_server_loop( - self._get_total_completed_count(multi_job_id), + context.job_state.total_completed_count(multi_job_id), timeout=1.0 ) log(f"USDU Dist: Images progress {completed_now}/{batch_size}") # Yield to allow workers to get new images after completing one - run_async_in_server_loop(self._async_yield(), timeout=0.1) + run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1) else: should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state( multi_job_id=multi_job_id, batch_size=batch_size, consecutive_retries=consecutive_retries, max_consecutive_retries=max_consecutive_retries, + mode_context=context, ) if should_break: break @@ -262,6 +314,7 @@ class DynamicModeMixin: processed_count=processed_count, result_images=result_images, upscaled_image=upscaled_image, + mode_context=context, ) debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete") @@ -282,33 +335,43 @@ class DynamicModeMixin: worker_id: str, enabled_worker_ids: list[str] | str, dynamic_threshold: int, + mode_context: DynamicModeContext | None = None, ) -> tuple[torch.Tensor]: """Worker processing in dynamic mode - processes whole images.""" + context = mode_context or self._build_dynamic_mode_context() # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) + tile_width = context.tile_ops.round_to_multiple(tile_width) + tile_height = context.tile_ops.round_to_multiple(tile_height) # Get dimensions and tile grid batch_size, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}") # Keep track of processed images for is_last detection processed_count = 0 # Poll for job readiness to avoid races during master init - max_poll_attempts = 20 # ~20s at 1s sleep - if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts): + max_poll_attempts = JOB_POLL_MAX_ATTEMPTS + if not self._poll_job_ready_dynamic( + multi_job_id, + master_url, + worker_id=worker_id, + max_attempts=max_poll_attempts, + mode_context=context, + ): log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") return (upscaled_image,) # Loop to request and process images while True: # Request an image to process - image_idx, estimated_remaining = run_async_in_server_loop( - self._request_image_from_master(multi_job_id, master_url, worker_id), - timeout=TILE_WAIT_TIMEOUT + assignment = run_async_in_server_loop( + context.worker_comms.request_assignment(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT, ) + image_idx = assignment.task_idx if assignment.kind == "image" else None + estimated_remaining = assignment.estimated_remaining if image_idx is None: debug_log(f"USDU Dist Worker - No more images to process") @@ -330,7 +393,7 @@ class DynamicModeMixin: image_seed = core_args.seed # Pre-slice conditioning once per image (not per tile) - positive_sliced, negative_sliced = self._slice_conditioning( + positive_sliced, negative_sliced = context.tile_ops.slice_conditioning( core_args.positive, core_args.negative, image_idx, @@ -340,7 +403,7 @@ class DynamicModeMixin: source_tensor = pil_to_tensor(local_image) if single_tensor.is_cuda: source_tensor = source_tensor.cuda() - local_image = self._process_and_blend_tile( + local_image = context.tile_ops.process_and_blend_tile( tile_idx, pos, source_tensor, local_image, core_args.model, positive_sliced, @@ -359,7 +422,7 @@ class DynamicModeMixin: batch_idx=image_idx, ) run_async_in_server_loop( - self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id), timeout=5.0 ) @@ -368,13 +431,19 @@ class DynamicModeMixin: # 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), + context.worker_comms.send_full_image( + 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), + context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id), timeout=5.0 ) if is_last: @@ -387,7 +456,7 @@ class DynamicModeMixin: debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal") try: run_async_in_server_loop( - self._send_worker_complete_signal(multi_job_id, master_url, worker_id), + context.worker_comms.send_worker_complete_signal(multi_job_id, master_url, worker_id), timeout=TILE_SEND_TIMEOUT ) except Exception as e: diff --git a/upscale/modes/single_gpu.py b/upscale/modes/single_gpu.py index b6b10b6..24a64a2 100644 --- a/upscale/modes/single_gpu.py +++ b/upscale/modes/single_gpu.py @@ -1,23 +1,19 @@ from __future__ import annotations import math, torch -from typing import Any, Protocol, cast +from typing import Any from ...utils.logging import log from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil +from ..mode_contexts import SingleGpuModeContext, TileOpsCollaborator from ..processing_args import UpscaleCoreArgs from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch -class _SingleGpuOps(Protocol): - def round_to_multiple(self, value: int) -> int: ... - def calculate_tiles(self, width: int, height: int, tile_width: int, tile_height: int, force_uniform_tiles: bool): ... - def create_tile_mask(self, width: int, height: int, tx: int, ty: int, tile_width: int, tile_height: int, mask_blur: int): ... - def blend_tile(self, base_image: Any, tile_pil: Any, x1: int, y1: int, size: tuple[int, int], tile_mask: Any, padding: int): ... - def extract_batch_tile_with_padding(self, source_batch: Any, tx: int, ty: int, tile_width: int, tile_height: int, padding: int, force_uniform_tiles: bool): ... - def process_tiles_batch(self, *args: Any, **kwargs: Any): ... - - class SingleGpuModeMixin: + def _build_single_gpu_mode_context(self) -> SingleGpuModeContext: + """Build explicit collaborators for single-GPU mode execution.""" + return SingleGpuModeContext(tile_ops=TileOpsCollaborator(self)) + def process_single_gpu( self, upscaled_image: torch.Tensor, @@ -27,9 +23,11 @@ class SingleGpuModeMixin: padding: int, mask_blur: int, force_uniform_tiles: bool, + mode_context: SingleGpuModeContext | None = None, ) -> tuple[torch.Tensor]: """Process all tiles on a single GPU (no distribution), batching per tile like USDU.""" - ops = cast(_SingleGpuOps, self) + context = mode_context or self._build_single_gpu_mode_context() + ops = context.tile_ops # Round tile dimensions tile_width = ops.round_to_multiple(tile_width) tile_height = ops.round_to_multiple(tile_height) diff --git a/upscale/modes/static.py b/upscale/modes/static.py index 354d5d1..4c6c69e 100644 --- a/upscale/modes/static.py +++ b/upscale/modes/static.py @@ -18,6 +18,12 @@ from ..job_store import ( cleanup_job, drain_results_queue, get_completed_count, mark_task_completed, ) from ..job_models import TileJobState +from ..mode_contexts import ( + JobStateCollaborator, + StaticModeContext, + TileOpsCollaborator, + WorkerCommsCollaborator, +) from ..processing_args import UpscaleCoreArgs from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch @@ -29,14 +35,30 @@ class StaticModeMixin: 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`). + - WorkerCommsMixin (`send_tiles_batch`, `request_assignment`, `send_heartbeat`). """ - def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS): + def _build_static_mode_context(self) -> StaticModeContext: + """Build explicit collaborators for static mode execution.""" + return StaticModeContext( + tile_ops=TileOpsCollaborator(self), + job_state=JobStateCollaborator(self), + worker_comms=WorkerCommsCollaborator(self), + ) + + def _poll_job_ready( + self, + multi_job_id, + master_url, + worker_id=None, + max_attempts=JOB_POLL_MAX_ATTEMPTS, + mode_context: StaticModeContext | None = None, + ): """Poll master for job readiness to avoid worker/master initialization race.""" + context = mode_context or self._build_static_mode_context() for attempt in range(max_attempts): ready = run_async_in_server_loop( - self._check_job_status(multi_job_id, master_url), + context.worker_comms.check_job_status(multi_job_id, master_url), timeout=5.0 ) if ready: @@ -111,12 +133,14 @@ class StaticModeMixin: padding, worker_id, is_final_flush=False, + mode_context: StaticModeContext | None = None, ): """Send accumulated tile payloads to master and return a fresh accumulator.""" + context = mode_context or self._build_static_mode_context() if not processed_tiles: if is_final_flush: run_async_in_server_loop( - self.send_tiles_batch_to_master( + context.worker_comms.send_tiles_batch( [], multi_job_id, master_url, @@ -128,7 +152,7 @@ class StaticModeMixin: ) return processed_tiles run_async_in_server_loop( - self.send_tiles_batch_to_master( + context.worker_comms.send_tiles_batch( processed_tiles, multi_job_id, master_url, @@ -167,8 +191,10 @@ class StaticModeMixin: tiled_decode, width, height, + mode_context: StaticModeContext | None = None, ): """Process one tile_id across the batch and blend into result_images.""" + context = mode_context or self._build_static_mode_context() source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0) if upscaled_image.is_cuda: source_batch = source_batch.cuda() @@ -202,7 +228,7 @@ class StaticModeMixin: result_images, processed_batch, b, - self.blend_tile, + context.tile_ops.blend_tile, x1, y1, ew, @@ -221,15 +247,17 @@ class StaticModeMixin: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, master_url, - worker_id, enabled_workers): + worker_id, enabled_workers, + mode_context: StaticModeContext | None = None): """Worker static mode processing with optional dynamic queue pulling.""" + context = mode_context or self._build_static_mode_context() # Round tile dimensions - tile_width = self.round_to_multiple(tile_width) - tile_height = self.round_to_multiple(tile_height) + tile_width = context.tile_ops.round_to_multiple(tile_width) + tile_height = context.tile_ops.round_to_multiple(tile_height) # Get dimensions and calculate tiles _, height, width, _ = upscaled_image.shape - all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) + all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles) num_tiles_per_image = len(all_tiles) batch_size = upscaled_image.shape[0] total_tiles = batch_size * num_tiles_per_image @@ -241,24 +269,31 @@ class StaticModeMixin: working_images.append(image_pil.copy()) tile_masks = [] for tx, ty in all_tiles: - tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) + tile_masks.append(context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)) # Dynamic queue mode (static processing): process batched-per-tile log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}") processed_count = 0 max_poll_attempts = JOB_POLL_MAX_ATTEMPTS - if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts): + if not self._poll_job_ready( + multi_job_id, + master_url, + worker_id=worker_id, + max_attempts=max_poll_attempts, + mode_context=context, + ): log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting") return (upscaled_image,) # Main processing loop - pull tile ids from queue while True: # Request a tile to process - tile_idx, estimated_remaining, batched_static = run_async_in_server_loop( - self._request_tile_from_master(multi_job_id, master_url, worker_id), - timeout=TILE_WAIT_TIMEOUT + assignment = run_async_in_server_loop( + context.worker_comms.request_assignment(multi_job_id, master_url, worker_id), + timeout=TILE_WAIT_TIMEOUT, ) + tile_idx = assignment.task_idx if assignment.kind == "tile" else None if tile_idx is None: debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process") @@ -298,7 +333,7 @@ class StaticModeMixin: tile_pil = tensor_to_pil(processed_batch, b) if tile_pil.size != (ew, eh): tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS) - working_images[b] = self.blend_tile( + working_images[b] = context.tile_ops.blend_tile( working_images[b], tile_pil, x1, @@ -322,7 +357,7 @@ class StaticModeMixin: # Send heartbeat try: run_async_in_server_loop( - self._send_heartbeat_to_master(multi_job_id, master_url, worker_id), + context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id), timeout=5.0 ) except Exception as e: @@ -331,20 +366,39 @@ class StaticModeMixin: # Send tiles in batches within loop if len(processed_tiles) >= MAX_BATCH: processed_tiles = self._flush_tiles_to_master( - processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False + processed_tiles, + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=False, + mode_context=context, ) # Send any remaining tiles processed_tiles = self._flush_tiles_to_master( - processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=True + processed_tiles, + multi_job_id, + master_url, + padding, + worker_id, + is_final_flush=True, + mode_context=context, ) debug_log(f"Worker {worker_id} completed all assigned and requeued tiles") return (upscaled_image,) - async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total): + async def _async_collect_and_monitor_static( + self, + multi_job_id, + total_tiles, + expected_total, + mode_context: StaticModeContext | None = None, + ): """Async helper for collection and monitoring in static mode. Returns collected tasks dict. Caller should check if all tasks are complete.""" + context = mode_context or self._build_static_mode_context() last_progress_log = time.time() progress_interval = 5.0 last_heartbeat_check = time.time() @@ -362,7 +416,7 @@ class StaticModeMixin: # Check and requeue timed-out workers periodically current_time = time.time() if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL: - requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total) + requeued_count = await context.job_state.check_and_requeue_timed_out_workers(multi_job_id, expected_total) if requeued_count > 0: log(f"Requeued {requeued_count} tasks from timed-out workers") last_heartbeat_check = current_time @@ -395,7 +449,7 @@ class StaticModeMixin: await asyncio.sleep(0.1) # Get all completed tasks for return - return await self._get_all_completed_tasks(multi_job_id) + return await context.job_state.all_completed_tasks(multi_job_id) def _init_master_static_state( self, @@ -408,7 +462,9 @@ class StaticModeMixin: tile_width, tile_height, mask_blur, + mode_context: StaticModeContext | None = None, ): + context = mode_context or self._build_static_mode_context() _, height, width, _ = upscaled_image.shape result_images = [tensor_to_pil(upscaled_image[b:b + 1], 0).copy() for b in range(batch_size)] @@ -420,7 +476,7 @@ class StaticModeMixin: debug_log(f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}") tile_masks = [ - self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur) + context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur) for tx, ty in all_tiles ] return result_images, tile_masks, width, height @@ -452,14 +508,16 @@ class StaticModeMixin: tiled_decode, width, height, + mode_context: StaticModeContext | None = None, ): + context = mode_context or self._build_static_mode_context() processed_count = 0 consecutive_no_tile = 0 max_consecutive_no_tile = 2 while processed_count < total_tiles: comfy.model_management.throw_exception_if_processing_interrupted() - tile_id = run_async_in_server_loop(self._get_next_tile_index(multi_job_id), timeout=5.0) + tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0) if tile_id is None: consecutive_no_tile += 1 if consecutive_no_tile >= max_consecutive_no_tile: @@ -495,6 +553,7 @@ class StaticModeMixin: tiled_decode, width, height, + mode_context=context, ) log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})") return processed_count @@ -527,14 +586,21 @@ class StaticModeMixin: tiled_decode, width, height, + mode_context: StaticModeContext | None = None, ): + context = mode_context or self._build_static_mode_context() remaining_tiles = total_tiles - master_processed_count if remaining_tiles <= 0: - return run_async_in_server_loop(self._get_all_completed_tasks(multi_job_id), timeout=5.0) + return run_async_in_server_loop(context.job_state.all_completed_tasks(multi_job_id), timeout=5.0) debug_log(f"Master waiting for {remaining_tiles} tiles from workers") collected_tasks = run_async_in_server_loop( - self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles), + self._async_collect_and_monitor_static( + multi_job_id, + total_tiles, + expected_total=total_tiles, + mode_context=context, + ), timeout=None, ) @@ -545,7 +611,7 @@ class StaticModeMixin: log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures") while True: comfy.model_management.throw_exception_if_processing_interrupted() - tile_id = run_async_in_server_loop(self._get_next_tile_index(multi_job_id), timeout=5.0) + tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0) if tile_id is None: break self._master_process_one_tile( @@ -574,6 +640,7 @@ class StaticModeMixin: tiled_decode, width, height, + mode_context=context, ) return collected_tasks @@ -588,7 +655,9 @@ class StaticModeMixin: tile_width, tile_height, padding, + mode_context: StaticModeContext | None = None, ): + context = mode_context or self._build_static_mode_context() def _sort_key(item): global_idx, tile_data = item batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image) @@ -609,7 +678,7 @@ class StaticModeMixin: tile_mask = tile_masks[tile_idx] extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding) extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding) - result_images[batch_idx] = self.blend_tile( + result_images[batch_idx] = context.tile_ops.blend_tile( result_images[batch_idx], tile_pil, x, @@ -632,8 +701,10 @@ class StaticModeMixin: seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers, - all_tiles, num_tiles_per_image): + all_tiles, num_tiles_per_image, + mode_context: StaticModeContext | None = None): """Static mode master processing with optional dynamic queue pulling.""" + context = mode_context or self._build_static_mode_context() batch_size = upscaled_image.shape[0] total_tiles = batch_size * num_tiles_per_image result_images, tile_masks, width, height = self._init_master_static_state( @@ -646,6 +717,7 @@ class StaticModeMixin: tile_width=tile_width, tile_height=tile_height, mask_blur=mask_blur, + mode_context=context, ) try: @@ -675,6 +747,7 @@ class StaticModeMixin: tiled_decode=tiled_decode, width=width, height=height, + mode_context=context, ) collected_tasks = self._collect_remaining_static_tiles( @@ -704,6 +777,7 @@ class StaticModeMixin: tiled_decode=tiled_decode, width=width, height=height, + mode_context=context, ) self._blend_static_collected_tiles( @@ -716,6 +790,7 @@ class StaticModeMixin: tile_width=tile_width, tile_height=tile_height, padding=padding, + mode_context=context, ) result_tensor = self._result_images_to_tensor(result_images, batch_size, upscaled_image) diff --git a/upscale/result_collector.py b/upscale/result_collector.py index a0517bd..f5fdd66 100644 --- a/upscale/result_collector.py +++ b/upscale/result_collector.py @@ -16,7 +16,7 @@ class ResultCollectorMixin: 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. + - `self.async_yield(...)` optional helper from WorkerCommsMixin. """ def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]: @@ -196,6 +196,10 @@ class ResultCollectorMixin: if isinstance(job_data, ImageJobState): job_data.completed_images[image_idx] = image_pil + async def mark_image_completed(self, multi_job_id, image_idx, image_pil): + """Public image-completion API.""" + await self._mark_image_completed(multi_job_id, image_idx, image_pil) + async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count): """Collect remaining processed images from workers.""" _ = master_processed_count @@ -207,3 +211,20 @@ class ResultCollectorMixin: remaining_to_collect=remaining_to_collect, batch_size=batch_size, ) + + async def collect_dynamic_images( + self, + multi_job_id, + remaining_to_collect, + num_workers, + batch_size, + master_processed_count, + ): + """Public dynamic image-collection API.""" + return await self._async_collect_dynamic_images( + multi_job_id, + remaining_to_collect, + num_workers, + batch_size, + master_processed_count, + ) diff --git a/upscale/tile_ops.py b/upscale/tile_ops.py index f8b235f..1f6565e 100644 --- a/upscale/tile_ops.py +++ b/upscale/tile_ops.py @@ -371,6 +371,10 @@ class TileOpsMixin: return positive_sliced, negative_sliced + def slice_conditioning(self, positive, negative, batch_idx): + """Public conditioning-slice API.""" + return self._slice_conditioning(positive, negative, batch_idx) + def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tile_width, tile_height, @@ -399,6 +403,59 @@ class TileOpsMixin: return result_image + def process_and_blend_tile( + self, + tile_idx, + tile_pos, + upscaled_image, + result_image, + model, + positive, + negative, + vae, + seed, + steps, + cfg, + sampler_name, + scheduler, + denoise, + tile_width, + tile_height, + padding, + mask_blur, + image_width, + image_height, + force_uniform_tiles, + tiled_decode, + batch_idx: int = 0, + ): + """Public tile-processing API.""" + return self._process_and_blend_tile( + tile_idx, + tile_pos, + upscaled_image, + result_image, + model, + positive, + negative, + vae, + seed, + steps, + cfg, + sampler_name, + scheduler, + denoise, + tile_width, + tile_height, + padding, + mask_blur, + image_width, + image_height, + force_uniform_tiles, + tiled_decode, + batch_idx=batch_idx, + ) + def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles, model, positive, negative, vae, seed, steps, cfg, sampler_name, scheduler, denoise, tiled_decode, tile_width, tile_height, padding, diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py index 06b11c2..2727a3b 100644 --- a/upscale/worker_comms.py +++ b/upscale/worker_comms.py @@ -1,25 +1,45 @@ import asyncio, io, json, time -from typing import Any +from dataclasses import dataclass +from typing import Any, Literal import aiohttp from PIL import Image from ..utils.logging import debug_log, log +from ..utils.auth import distributed_auth_headers +from ..utils.config import load_config from ..utils.network import get_client_session from ..utils.constants import TILE_SEND_TIMEOUT -from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, send_heartbeat_to_master +from ..utils.usdu_management import MAX_PAYLOAD_SIZE, send_heartbeat_to_master from ..utils.image import tensor_to_pil +WorkAssignmentKind = Literal["image", "tile", "none"] + + +@dataclass(frozen=True, slots=True) +class WorkAssignment: + """Canonical typed representation of a single work-item assignment.""" + + kind: WorkAssignmentKind + task_idx: int | None + estimated_remaining: int = 0 + batched_static: bool = False + + class WorkerCommsMixin: - async def _send_heartbeat_to_master( + @staticmethod + def _master_auth_headers() -> dict[str, str]: + return distributed_auth_headers(load_config()) + + async def send_heartbeat( self, multi_job_id: str, master_url: str, worker_id: str, ) -> None: - """Proxy heartbeat helper used by worker processing mixins.""" + """Send worker heartbeat to the master.""" await send_heartbeat_to_master(multi_job_id, master_url, worker_id) - async def send_tiles_batch_to_master( + async def send_tiles_batch( self, processed_tiles: list[dict[str, Any]], multi_job_id: str, @@ -63,12 +83,8 @@ class WorkerCommsMixin: i = 0 chunk_index = 0 while i < total_tiles: - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('padding', str(padding)) - metadata = [] + chunk_images: list[tuple[int, int, bytes]] = [] used = 0 j = i while j < total_tiles: @@ -80,7 +96,7 @@ class WorkerCommsMixin: break # Accept this tile in this chunk metadata.append(meta) - data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png') + chunk_images.append((j - i, j, img_bytes)) used += len(img_bytes) + overhead j += 1 @@ -89,14 +105,28 @@ class WorkerCommsMixin: # Single oversized tile, send anyway meta = encoded[j]['meta'] metadata.append(meta) - data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png') + chunk_images.append((0, j, encoded[j]['bytes'])) j += 1 chunk_size = j - i is_chunk_last = (j >= total_tiles) - data.add_field('is_last', str(bool(is_final_flush and is_chunk_last))) - data.add_field('batch_size', str(chunk_size)) - data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') + + def _build_chunk_form() -> aiohttp.FormData: + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('padding', str(padding)) + data.add_field('is_last', str(bool(is_final_flush and is_chunk_last))) + data.add_field('batch_size', str(chunk_size)) + data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') + for relative_idx, source_idx, img_bytes in chunk_images: + data.add_field( + f'tile_{relative_idx}', + io.BytesIO(img_bytes), + filename=f'tile_{source_idx}.png', + content_type='image/png', + ) + return data # Retry logic with exponential backoff max_retries = 5 @@ -105,7 +135,11 @@ class WorkerCommsMixin: try: session = await get_client_session() url = f"{master_url}/distributed/submit_tiles" - async with session.post(url, data=data) as response: + async with session.post( + url, + data=_build_chunk_form(), + headers=self._master_auth_headers(), + ) as response: response.raise_for_status() break except Exception as e: @@ -130,7 +164,7 @@ class WorkerCommsMixin: session = await get_client_session() url = f"{master_url}/distributed/submit_tiles" - async with session.post(url, data=data) as response: + async with session.post(url, data=data, headers=self._master_auth_headers()) as response: response.raise_for_status() debug_log(f"Worker {worker_id} sent static completion signal") @@ -157,7 +191,7 @@ class WorkerCommsMixin: async with session.post(url, json={ 'worker_id': str(worker_id), 'multi_job_id': multi_job_id - }) as response: + }, headers=self._master_auth_headers()) as response: if response.status == 200: return await response.json() if response.status == 404: @@ -181,41 +215,83 @@ class WorkerCommsMixin: return None - async def _request_image_from_master(self, multi_job_id, master_url, worker_id): - """Request an image index to process from master in dynamic mode.""" - data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id) + @staticmethod + def _parse_work_assignment(data: dict[str, Any] | None) -> WorkAssignment: + """Normalize assignment payloads into a single discriminated contract.""" if not data: - return None, 0 - image_idx = data.get('image_idx') - estimated_remaining = data.get('estimated_remaining', 0) - return image_idx, estimated_remaining + return WorkAssignment(kind="none", task_idx=None) - async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): - """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" + kind_value = data.get("kind") + if kind_value is None: + # Backward-compatible parsing for older masters. + if "image_idx" in data: + kind_value = "image" + data = { + "kind": "image", + "task_idx": data.get("image_idx"), + "estimated_remaining": data.get("estimated_remaining", 0), + } + elif "tile_idx" in data: + kind_value = "tile" + data = { + "kind": "tile", + "task_idx": data.get("tile_idx"), + "estimated_remaining": data.get("estimated_remaining", 0), + "batched_static": data.get("batched_static", False), + } + else: + kind_value = "none" + + if kind_value not in {"image", "tile", "none"}: + return WorkAssignment(kind="none", task_idx=None) + + task_idx_raw = data.get("task_idx") + if task_idx_raw is None: + task_idx = None + else: + try: + task_idx = int(task_idx_raw) + except (TypeError, ValueError): + task_idx = None + + try: + estimated_remaining = int(data.get("estimated_remaining", 0) or 0) + except (TypeError, ValueError): + estimated_remaining = 0 + + return WorkAssignment( + kind=kind_value, + task_idx=task_idx, + estimated_remaining=estimated_remaining, + batched_static=bool(data.get("batched_static", False)), + ) + + async def request_assignment(self, multi_job_id, master_url, worker_id) -> WorkAssignment: + """Request one assignment and parse into the canonical discriminated contract.""" data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id) - if not data: - return None, 0, False - tile_idx = data.get('tile_idx') - estimated_remaining = data.get('estimated_remaining', 0) - batched_static = data.get('batched_static', False) - return tile_idx, estimated_remaining, batched_static + return self._parse_work_assignment(data) - async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, - master_url, worker_id, is_last): + async def send_full_image(self, image_pil, image_idx, multi_job_id, + master_url, worker_id, is_last): """Send a processed full image back to master in dynamic mode.""" # Serialize image to PNG byte_io = io.BytesIO() image_pil.save(byte_io, format='PNG', compress_level=0) - byte_io.seek(0) - - # Prepare form data - data = aiohttp.FormData() - data.add_field('multi_job_id', multi_job_id) - data.add_field('worker_id', str(worker_id)) - data.add_field('image_idx', str(image_idx)) - data.add_field('is_last', str(is_last)) - data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', - content_type='image/png') + image_bytes = byte_io.getvalue() + + def _build_image_form() -> aiohttp.FormData: + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('image_idx', str(image_idx)) + data.add_field('is_last', str(is_last)) + data.add_field( + 'full_image', + io.BytesIO(image_bytes), + filename=f'image_{image_idx}.png', + content_type='image/png', + ) + return data # Retry logic max_retries = 5 @@ -226,7 +302,11 @@ class WorkerCommsMixin: session = await get_client_session() url = f"{master_url}/distributed/submit_image" - async with session.post(url, data=data) as response: + async with session.post( + url, + data=_build_image_form(), + headers=self._master_auth_headers(), + ) as response: response.raise_for_status() debug_log(f"Successfully sent image {image_idx} to master") return @@ -240,7 +320,7 @@ class WorkerCommsMixin: log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}") raise - async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id): + async def send_worker_complete_signal(self, multi_job_id, master_url, worker_id): """Send completion signal to master in dynamic mode.""" # Send a dummy request with is_last=True data = aiohttp.FormData() @@ -252,16 +332,16 @@ class WorkerCommsMixin: session = await get_client_session() url = f"{master_url}/distributed/submit_image" - async with session.post(url, data=data) as response: + async with session.post(url, data=data, headers=self._master_auth_headers()) as response: response.raise_for_status() debug_log(f"Worker {worker_id} sent completion signal") - async def _check_job_status(self, multi_job_id, master_url): + async def check_job_status(self, multi_job_id, master_url): """Check if job is ready on the master.""" try: session = await get_client_session() url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}" - async with session.get(url) as response: + async with session.get(url, headers=self._master_auth_headers()) as response: if response.status == 200: data = await response.json() return data.get('ready', False) @@ -270,6 +350,6 @@ class WorkerCommsMixin: debug_log(f"Job status check failed: {e}") return False - async def _async_yield(self): + async def async_yield(self): """Simple async yield to allow event loop processing.""" await asyncio.sleep(0) diff --git a/utils/async_helpers.py b/utils/async_helpers.py index 3a36879..a9b4e9b 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -6,7 +6,6 @@ import threading import time import uuid import execution -import server from typing import Optional, Any, Coroutine from .network import get_server_loop @@ -43,10 +42,11 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) - # Schedule on server's event loop loop = get_server_loop() - asyncio.run_coroutine_threadsafe(wrapper(), loop) + task_future = asyncio.run_coroutine_threadsafe(wrapper(), loop) # Wait for completion if not event.wait(timeout): + task_future.cancel() raise TimeoutError(f"Async operation timed out after {timeout} seconds") if error: @@ -54,7 +54,10 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) - return result -prompt_server = server.PromptServer.instance +def _prompt_server_instance(): + import server + + return server.PromptServer.instance def _summarize_node_errors(node_errors: dict) -> str: @@ -111,6 +114,7 @@ async def queue_prompt_payload( client_id: str | None = None, ) -> str: """Validate and queue a prompt via ComfyUI's prompt queue.""" + prompt_server = _prompt_server_instance() payload = {"prompt": prompt_obj} payload = prompt_server.trigger_on_prompt(payload) prompt = payload["prompt"] diff --git a/utils/auth.py b/utils/auth.py new file mode 100644 index 0000000..85d6ea8 --- /dev/null +++ b/utils/auth.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import hmac +from typing import Any + + +AUTH_HEADER_NAME = "X-Distributed-Token" + + +def get_distributed_api_token(config: dict[str, Any] | None) -> str: + """Return configured distributed API token or empty string when unset.""" + settings = (config or {}).get("settings", {}) if isinstance(config, dict) else {} + token = settings.get("distributed_api_token") or settings.get("api_token") + return str(token).strip() if token is not None else "" + + +def distributed_auth_headers(config: dict[str, Any] | None) -> dict[str, str]: + """Build outbound auth headers for distributed internal API calls.""" + token = get_distributed_api_token(config) + if not token: + return {} + return {AUTH_HEADER_NAME: token} + + +def is_authorized_request(request: Any, config: dict[str, Any] | None) -> bool: + """Validate a request against the configured shared distributed token.""" + expected_token = get_distributed_api_token(config) + if not expected_token: + return True + + provided_token = "" + headers = getattr(request, "headers", None) + if headers is not None: + provided_token = str(headers.get(AUTH_HEADER_NAME, "")).strip() + + return hmac.compare_digest(provided_token, expected_token) diff --git a/utils/cloudflare/process_reader.py b/utils/cloudflare/process_reader.py index e476a36..c8c5107 100644 --- a/utils/cloudflare/process_reader.py +++ b/utils/cloudflare/process_reader.py @@ -62,7 +62,10 @@ class ProcessReader: return loop = self._loop - for raw_line in iter(process.stdout.readline, ""): + while True: + raw_line = process.stdout.readline() + if not raw_line: + break line = self._normalize_line(raw_line) if not line: continue diff --git a/utils/config.py b/utils/config.py index b6b4ecb..6b483bc 100644 --- a/utils/config.py +++ b/utils/config.py @@ -9,10 +9,16 @@ from collections.abc import AsyncIterator from contextlib import asynccontextmanager from functools import lru_cache from typing import Any -from .logging import log # Import defaults for timeout fallbacks -from .constants import GPU_CONFIG_FILE, HEARTBEAT_TIMEOUT +from .constants import ( + GPU_CONFIG_FILE, + HEARTBEAT_TIMEOUT, + ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, + ORCHESTRATION_MEDIA_SYNC_TIMEOUT, + ORCHESTRATION_WORKER_PREP_CONCURRENCY, + ORCHESTRATION_WORKER_PROBE_CONCURRENCY, +) CONFIG_FILE = GPU_CONFIG_FILE @@ -36,6 +42,11 @@ def _config_lock() -> asyncio.Lock: def _config_path() -> str: return CONFIG_FILE + +def _log(message: str) -> None: + print(f"[Distributed] {message}") + + def get_default_config() -> dict[str, Any]: """Returns the default configuration dictionary. Single source of truth.""" return { @@ -47,10 +58,10 @@ def get_default_config() -> dict[str, Any]: "stop_workers_on_master_exit": True, "master_delegate_only": False, "websocket_orchestration": True, - "worker_probe_concurrency": 8, - "worker_prep_concurrency": 4, - "media_sync_concurrency": 2, - "media_sync_timeout_seconds": 120 + "worker_probe_concurrency": ORCHESTRATION_WORKER_PROBE_CONCURRENCY, + "worker_prep_concurrency": ORCHESTRATION_WORKER_PREP_CONCURRENCY, + "media_sync_concurrency": ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, + "media_sync_timeout_seconds": ORCHESTRATION_MEDIA_SYNC_TIMEOUT, }, "tunnel": { "status": "stopped", @@ -107,7 +118,7 @@ def load_config() -> dict[str, Any]: loaded = json.load(f) state.cache = _merge_with_defaults(loaded, get_default_config()) except Exception as e: - log(f"Error loading config, using defaults: {e}") + _log(f"Error loading config, using defaults: {e}") state.cache = get_default_config() state.mtime = mtime @@ -128,8 +139,8 @@ def save_config(config: dict[str, Any]) -> bool: try: os.unlink(tmp_path) except OSError as cleanup_error: - log(f"Error removing temporary config file '{tmp_path}': {cleanup_error}") - log(f"Error saving config: {e}") + _log(f"Error removing temporary config file '{tmp_path}': {cleanup_error}") + _log(f"Error saving config: {e}") return False @@ -149,11 +160,8 @@ def ensure_config_exists() -> None: """Creates default config file if it doesn't exist. Used by __init__.py""" if not os.path.exists(_config_path()): default_config = get_default_config() - if save_config(default_config): - from .logging import debug_log - debug_log("Created default config file") - else: - log("Could not create default config file") + if not save_config(default_config): + _log("Could not create default config file") def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int: """Return the unified worker timeout (seconds). @@ -181,3 +189,12 @@ def is_master_delegate_only() -> bool: return bool(cfg.get('settings', {}).get('master_delegate_only', False)) except Exception: return False + + +def get_debug_enabled(default: bool = False) -> bool: + """Return whether distributed debug logging is enabled in config settings.""" + try: + cfg = load_config() + return bool(cfg.get("settings", {}).get("debug", default)) + except Exception: + return bool(default) diff --git a/utils/logging.py b/utils/logging.py index 905470e..0a5fef8 100644 --- a/utils/logging.py +++ b/utils/logging.py @@ -1,50 +1,15 @@ """ Shared logging utilities for ComfyUI-Distributed. """ -import os -import json -import time -from dataclasses import dataclass -from functools import lru_cache -from .constants import GPU_CONFIG_FILE +from .config import get_debug_enabled -_DEBUG_TTL: float = 5.0 - - -@dataclass -class _DebugState: - enabled: bool | None = None - updated_at: float = 0.0 - - -@lru_cache(maxsize=1) -def _debug_state() -> _DebugState: - return _DebugState() def is_debug_enabled() -> bool: - """Check if debug is enabled.""" - state = _debug_state() - - now = time.monotonic() - if state.enabled is not None and (now - state.updated_at) < _DEBUG_TTL: - return state.enabled - - enabled = state.enabled if state.enabled is not None else False - if os.path.exists(GPU_CONFIG_FILE): - try: - with open(GPU_CONFIG_FILE, "r", encoding="utf-8") as f: - config = json.load(f) - enabled = bool(config.get("settings", {}).get("debug", False)) - except (OSError, json.JSONDecodeError, ValueError) as exc: - # Fall back to the previous cached value to avoid flapping behavior. - print( - "[Distributed] Failed to read debug flag from config; " - f"using cached value {enabled}: {exc}" - ) - - state.enabled = enabled - state.updated_at = now - return enabled + """Check whether debug logging is currently enabled.""" + try: + return get_debug_enabled(default=False) + except Exception: + return False def debug_log(message: str) -> None: """Log debug messages only if debug is enabled in config.""" diff --git a/utils/network.py b/utils/network.py index 9317f48..e9c5892 100644 --- a/utils/network.py +++ b/utils/network.py @@ -5,7 +5,6 @@ Network and API utilities for ComfyUI-Distributed. """ import asyncio import aiohttp -import server from dataclasses import dataclass from functools import lru_cache from typing import Any @@ -23,6 +22,12 @@ class _NetworkState: def _network_state() -> _NetworkState: return _NetworkState() + +def _prompt_server_instance() -> Any: + import server + + return server.PromptServer.instance + async def get_client_session() -> aiohttp.ClientSession: """Get or create a shared aiohttp client session.""" state = _network_state() @@ -50,28 +55,45 @@ async def handle_api_error(request: web.Request, error: Any, status: int = 500) 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) + return web.json_response( + { + "status": "error", + "error": messages, + "message": "; ".join(messages), + }, + status=status, + ) if isinstance(error, DistributedError): debug_log(f"API Error [{status}] ({error.__class__.__name__}): {error}") return web.json_response( - {"error": str(error), "error_type": error.__class__.__name__}, + { + "status": "error", + "error": str(error), + "message": str(error), + "error_type": error.__class__.__name__, + }, status=status, ) message = str(error) debug_log(f"API Error [{status}]: {message}") - return web.json_response({"error": message}, status=status) + return web.json_response( + { + "status": "error", + "error": message, + "message": message, + }, + status=status, + ) def get_server_port() -> int: """Get the ComfyUI server port.""" - import server - return server.PromptServer.instance.port + return _prompt_server_instance().port def get_server_loop() -> asyncio.AbstractEventLoop: """Get the ComfyUI server event loop.""" - import server - return server.PromptServer.instance.loop + return _prompt_server_instance().loop def normalize_host(value: Any) -> Any: @@ -94,7 +116,7 @@ def build_worker_url(worker: dict[str, Any], endpoint: str = "") -> str: port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188) if not host: - host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1" + host = getattr(_prompt_server_instance(), "address", "127.0.0.1") or "127.0.0.1" if host.startswith(("http://", "https://")): base = host.rstrip("/") @@ -148,7 +170,7 @@ def build_master_url( from .config import load_config config = load_config() - prompt_server_instance = prompt_server_instance or server.PromptServer.instance + prompt_server_instance = prompt_server_instance or _prompt_server_instance() master_cfg = (config or {}).get("master", {}) or {} configured_host = (master_cfg.get("host") or "").strip() configured_port = master_cfg.get("port") diff --git a/utils/parsing.py b/utils/parsing.py new file mode 100644 index 0000000..44bba0d --- /dev/null +++ b/utils/parsing.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from typing import Any + + +def require_bool_literal(value: Any, *, field_name: str = "value") -> bool: + """Strict boolean parser for API boundaries.""" + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized == "true": + return True + if normalized == "false": + return False + raise ValueError(f"Field '{field_name}' must be a boolean literal ('true' or 'false').") + + +def coerce_bool(value: Any, default: bool = False) -> bool: + """Permissive boolean parser for internal/runtime values.""" + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off", ""}: + return False + return default + if value is None: + return default + if isinstance(value, (int, float)): + return bool(value) + return default diff --git a/utils/process.py b/utils/process.py index 42bf24d..29a629d 100644 --- a/utils/process.py +++ b/utils/process.py @@ -126,13 +126,23 @@ def is_process_alive(pid: int) -> bool: def terminate_process(process, timeout=5): """Gracefully terminate a process with timeout.""" - if process.poll() is None: # Still running - process.terminate() + wait_timeout = max(float(timeout), 0.1) + if process.poll() is not None: + return True + + process.terminate() + try: + process.wait(timeout=wait_timeout) + return True + except (subprocess.TimeoutExpired, TimeoutError): + process.kill() try: - process.wait(timeout=timeout) - except (subprocess.TimeoutExpired, TimeoutError): - process.kill() - process.wait() + process.wait(timeout=wait_timeout) + return True + except (subprocess.TimeoutExpired, TimeoutError) as exc: + raise TimeoutError( + f"Process did not terminate within {wait_timeout} seconds after kill." + ) from exc def get_python_executable(): """Get the Python executable path.""" diff --git a/utils/runtime_state.py b/utils/runtime_state.py new file mode 100644 index 0000000..add366c --- /dev/null +++ b/utils/runtime_state.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class DistributedRuntimeState: + """Typed distributed runtime state bound to PromptServer.instance.""" + + distributed_pending_jobs: dict[str, asyncio.Queue] = field(default_factory=dict) + distributed_jobs_lock: asyncio.Lock = field(default_factory=asyncio.Lock) + distributed_job_allowed_workers: dict[str, set[str]] = field(default_factory=dict) + distributed_pending_tile_jobs: dict[str, Any] = field(default_factory=dict) + distributed_tile_jobs_lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +def get_prompt_server_instance() -> Any: + import server + + return server.PromptServer.instance + + +def ensure_distributed_runtime_state(server_instance: Any | None = None) -> DistributedRuntimeState: + """Return canonical distributed runtime state for the server instance.""" + prompt_server = server_instance or get_prompt_server_instance() + state = getattr(prompt_server, "distributed_runtime_state", None) + if not isinstance(state, DistributedRuntimeState): + state = DistributedRuntimeState() + setattr(prompt_server, "distributed_runtime_state", state) + + # Backward-compatible aliases for legacy call sites. + prompt_server.distributed_pending_jobs = state.distributed_pending_jobs + prompt_server.distributed_jobs_lock = state.distributed_jobs_lock + prompt_server.distributed_job_allowed_workers = state.distributed_job_allowed_workers + prompt_server.distributed_pending_tile_jobs = state.distributed_pending_tile_jobs + prompt_server.distributed_tile_jobs_lock = state.distributed_tile_jobs_lock + return state diff --git a/utils/usdu_managment.py b/utils/usdu_management.py similarity index 58% rename from utils/usdu_managment.py rename to utils/usdu_management.py index 757efb5..7dd6214 100644 --- a/utils/usdu_managment.py +++ b/utils/usdu_management.py @@ -1,10 +1,6 @@ -"""Backward-compatibility shim for USDU helpers. +"""Canonical USDU helper exports. -Route handlers and job logic now live in: -- upscale.job_store -- upscale.job_timeout -- upscale.payload_parsers -- api.usdu_routes +This module is the stable import target for USDU queue/heartbeat helpers. """ from ..upscale.conditioning import clone_conditioning, clone_control_chain @@ -12,15 +8,17 @@ 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, + get_completed_count, init_dynamic_job, + init_job_queue, init_static_job_batched, + mark_task_completed, ) from ..upscale.job_timeout import check_and_requeue_timed_out_workers from ..upscale.payload_parsers import parse_tiles_from_form +from .auth import distributed_auth_headers +from .config import load_config from .logging import debug_log from .network import get_client_session @@ -28,26 +26,18 @@ from .network import get_client_session async def send_heartbeat_to_master(multi_job_id: str, master_url: str, worker_id: str) -> None: """Send heartbeat to master from worker-side processing loops.""" try: - data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)} + data = {"multi_job_id": multi_job_id, "worker_id": str(worker_id)} session = await get_client_session() url = f"{master_url}/distributed/heartbeat" - async with session.post(url, json=data) as response: + async with session.post( + url, + json=data, + headers=distributed_auth_headers(load_config()), + ) as response: response.raise_for_status() except Exception as e: debug_log(f"Heartbeat failed: {e}") - -# Backward-compatible aliases for legacy imports. -_check_and_requeue_timed_out_workers = check_and_requeue_timed_out_workers -_cleanup_job = cleanup_job -_drain_results_queue = drain_results_queue -_get_completed_count = get_completed_count -_init_job_queue = init_job_queue -_mark_task_completed = mark_task_completed -_parse_tiles_from_form = parse_tiles_from_form -_send_heartbeat_to_master = send_heartbeat_to_master - - __all__ = [ "MAX_PAYLOAD_SIZE", "check_and_requeue_timed_out_workers", @@ -58,14 +48,6 @@ __all__ = [ "mark_task_completed", "parse_tiles_from_form", "send_heartbeat_to_master", - "_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", diff --git a/utils/usdu_utils.py b/utils/usdu_utils.py index 4fd594f..4370bbf 100644 --- a/utils/usdu_utils.py +++ b/utils/usdu_utils.py @@ -2,31 +2,17 @@ import numpy as np from PIL import Image, ImageFilter import torch import torch.nn.functional as F -from torchvision.transforms import GaussianBlur import math from typing import Any +from .image import pil_to_tensor, tensor_to_pil + if (not hasattr(Image, 'Resampling')): # For older versions of Pillow Image.Resampling = Image BLUR_KERNEL_SIZE = 15 -def tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image: - # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels] - # and returns an RGB PIL Image. Assumes channels=3 - return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8)) - - -def pil_to_tensor(image: Image.Image) -> torch.Tensor: - # Takes a PIL image and returns a tensor of shape [1, height, width, channels] - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image).unsqueeze(0) - if len(image.shape) == 3: # If the image is grayscale, add a channel dimension - image = image.unsqueeze(-1) - return image - - def controlnet_hint_to_pil(tensor: torch.Tensor, batch_index: int = 0) -> Image.Image: return tensor_to_pil(tensor.movedim(1, -1), batch_index) diff --git a/utils/worker_ids.py b/utils/worker_ids.py index f5549f7..1d44bc2 100644 --- a/utils/worker_ids.py +++ b/utils/worker_ids.py @@ -1,20 +1,98 @@ import json -from typing import Any +import re + +_LEGACY_INDEX_TOKEN_RE = re.compile(r"^(?:worker_)?\d+$") -def parse_enabled_worker_ids(enabled_worker_ids: str | list[Any] | None) -> list[str]: - """Parse and de-duplicate enabled worker IDs from a JSON list payload.""" - try: - raw = json.loads(enabled_worker_ids) if isinstance(enabled_worker_ids, str) else enabled_worker_ids - except Exception: - raw = [] +def _require_canonical_worker_id(value: object) -> str: + worker_id = str(value).strip() + if not worker_id: + raise ValueError("enabled_worker_ids cannot contain blank IDs") + if _LEGACY_INDEX_TOKEN_RE.match(worker_id): + raise ValueError( + f"enabled_worker_ids entry '{worker_id}' uses a legacy index token; " + "provide stable explicit worker IDs" + ) + return worker_id + + +def require_enabled_worker_ids(enabled_worker_ids: list[str] | None) -> list[str]: + """Strictly validate canonical enabled-worker list.""" + if enabled_worker_ids is None: + return [] + if not isinstance(enabled_worker_ids, list): + raise ValueError("enabled_worker_ids must decode to a list") workers = [] seen = set() - for worker_id in raw if isinstance(raw, list) else []: - worker_id_str = str(worker_id) + for worker_id in enabled_worker_ids: + worker_id_str = _require_canonical_worker_id(worker_id) if worker_id_str in seen: continue seen.add(worker_id_str) workers.append(worker_id_str) return workers + + +def decode_enabled_worker_ids(value: str | list[str] | None) -> list[str]: + """Decode boundary payload into canonical enabled-worker list.""" + raw = json.loads(value) if isinstance(value, str) else value + return require_enabled_worker_ids(raw) + + +def coerce_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]: + """Best-effort boundary decoding for permissive internal call sites.""" + try: + return decode_enabled_worker_ids(enabled_worker_ids) + except (TypeError, ValueError, json.JSONDecodeError): + return [] + + +def parse_worker_index( + worker_id: str, + enabled_worker_ids: list[str], +) -> int | None: + """Resolve worker index from canonical ID or legacy index token.""" + workers = list(enabled_worker_ids or []) + worker_token = str(worker_id).strip() + if not worker_token: + return None + + if workers and worker_token in workers: + return workers.index(worker_token) + + normalized = worker_token[len("worker_") :] if worker_token.startswith("worker_") else worker_token + + try: + parsed_index = int(normalized) + except (TypeError, ValueError): + return None + if parsed_index < 0: + return None + return parsed_index + + +def resolve_worker_identity( + worker_id: str, + enabled_worker_ids: list[str], +) -> tuple[str, int | None]: + """Return canonical worker id and resolved index when possible.""" + worker_token = str(worker_id).strip() + worker_index = parse_worker_index(worker_token, enabled_worker_ids) + workers = list(enabled_worker_ids or []) + if worker_index is None: + return worker_token, None + if 0 <= worker_index < len(workers): + return workers[worker_index], worker_index + return worker_token, worker_index + + +def worker_value_key( + worker_id: str, + enabled_worker_ids: list[str], +) -> str: + """Return a stable key used for worker-value maps.""" + worker_token, worker_index = resolve_worker_identity(worker_id, enabled_worker_ids) + if worker_index is not None: + return str(worker_index + 1) + return worker_token diff --git a/web/constants.js b/web/constants.js index 21073ff..2ed8e67 100644 --- a/web/constants.js +++ b/web/constants.js @@ -225,8 +225,6 @@ export const ENDPOINTS = { export const NODE_CLASSES = { DISTRIBUTED_COLLECTOR: 'DistributedCollector', - DISTRIBUTED_LIST_SPLITTER: 'DistributedListSplitter', - DISTRIBUTED_LIST_COLLECTOR: 'DistributedListCollector', DISTRIBUTED_BRANCH: 'DistributedBranch', DISTRIBUTED_BRANCH_COLLECTOR: 'DistributedBranchCollector', DISTRIBUTED_SEED: 'DistributedSeed', diff --git a/web/executionUtils.js b/web/executionUtils.js index fb662ff..dc01dd7 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -7,13 +7,11 @@ export function setupInterceptor(extension) { api.queuePrompt = async (number, prompt) => { if (extension.isEnabled) { const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0; - const hasListSplitter = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_SPLITTER).length > 0; - const hasListCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_COLLECTOR).length > 0; const hasBranch = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH).length > 0; const hasBranchCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH_COLLECTOR).length > 0; const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0; - if (hasCollector || hasListSplitter || hasListCollector || hasBranch || hasBranchCollector || hasDistUpscale) { + if (hasCollector || hasBranch || hasBranchCollector || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); // Immediate status check for instant feedback checkAllWorkerStatuses(extension); @@ -87,7 +85,6 @@ export async function executeParallelDistributed(extension, promptWrapper) { workers: activeWorkers.map((worker) => ({ id: worker.id })), client_id: api.clientId, delegate_master: Boolean(extension.config?.settings?.master_delegate_only), - auto_prepare: true, trace_execution_id: traceExecutionId, }); if (queueResponse?.prompt_id) { diff --git a/workers/__init__.py b/workers/__init__.py index 79f5b3e..41b766e 100644 --- a/workers/__init__.py +++ b/workers/__init__.py @@ -1,15 +1,9 @@ from .process_manager import WorkerProcessManager from functools import lru_cache -# Keep monitor module importable via workers package for subprocess integration. -from . import worker_monitor as worker_monitor # noqa: F401 - - @lru_cache(maxsize=1) def _worker_manager() -> WorkerProcessManager: - manager = WorkerProcessManager() - manager.queues = {} - return manager + return WorkerProcessManager() def get_worker_manager() -> WorkerProcessManager: return _worker_manager() diff --git a/workers/process/lifecycle.py b/workers/process/lifecycle.py index 2f20ca6..dd02bc6 100644 --- a/workers/process/lifecycle.py +++ b/workers/process/lifecycle.py @@ -74,6 +74,13 @@ class ProcessLifecycle: os.path.dirname(os.path.dirname(__file__)), "worker_monitor.py", ) + project_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) + existing_pythonpath = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"{project_root}{os.pathsep}{existing_pythonpath}" + if existing_pythonpath + else project_root + ) monitored_cmd = [get_python_executable(), monitor_script] + cmd log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n") log_handle.flush() diff --git a/workers/process/persistence.py b/workers/process/persistence.py index a0e800a..3224a8a 100644 --- a/workers/process/persistence.py +++ b/workers/process/persistence.py @@ -15,7 +15,7 @@ class ProcessPersistence: for worker_id, proc_info in managed_processes.items(): pid = proc_info.get("pid") - if pid and self._manager._is_process_running(pid): + if pid and self._manager.is_process_running(pid): self._manager.processes[worker_id] = { "pid": pid, "process": None, @@ -33,7 +33,7 @@ class ProcessPersistence: managed_processes = {} for worker_id, proc_info in self._manager.processes.items(): - is_running, _ = self._manager._check_worker_process(worker_id, proc_info) + is_running, _ = self._manager.check_worker_process(worker_id, proc_info) if not is_running: continue managed_processes[worker_id] = { diff --git a/workers/process_manager.py b/workers/process_manager.py index 0f170f4..9c8209e 100644 --- a/workers/process_manager.py +++ b/workers/process_manager.py @@ -5,10 +5,11 @@ from .process.root_discovery import ComfyRootDiscovery class WorkerProcessManager: - """Thin composition wrapper around worker process subsystems.""" + """Worker process service composed from launch, lifecycle, and persistence subsystems.""" def __init__(self): self.processes = {} + self.queues = {} self._root_discovery = ComfyRootDiscovery() self._launch_builder = LaunchCommandBuilder() self._lifecycle = ProcessLifecycle(self) @@ -18,9 +19,6 @@ class WorkerProcessManager: def find_comfy_root(self): return self._root_discovery.find_comfy_root() - def _find_windows_terminal(self): - return self._launch_builder._find_windows_terminal() - def build_launch_command(self, worker_config, comfy_root): return self._launch_builder.build_launch_command(worker_config, comfy_root) @@ -42,11 +40,11 @@ class WorkerProcessManager: def save_processes(self): return self._persistence.save_processes() - def _is_process_running(self, pid): + def is_process_running(self, pid): return self._lifecycle._is_process_running(pid) - def _check_worker_process(self, worker_id, proc_info): + def check_worker_process(self, worker_id, proc_info): return self._lifecycle._check_worker_process(worker_id, proc_info) - def _kill_process_tree(self, pid): + def kill_process_tree(self, pid): return self._lifecycle._kill_process_tree(pid) diff --git a/workers/startup.py b/workers/startup.py index 9b43e5d..6b6523c 100644 --- a/workers/startup.py +++ b/workers/startup.py @@ -47,7 +47,7 @@ def auto_launch_workers() -> None: # Check if already running if str(worker_id) in wm.processes: proc_info = wm.processes[str(worker_id)] - if wm._is_process_running(proc_info['pid']): + if wm.is_process_running(proc_info['pid']): debug_log(f"Worker {worker_name} already running, skipping") continue diff --git a/workers/worker_monitor.py b/workers/worker_monitor.py index 3caa93e..c5e9809 100644 --- a/workers/worker_monitor.py +++ b/workers/worker_monitor.py @@ -6,53 +6,17 @@ and terminates the worker if the master dies. import os import sys import time -import subprocess import platform import signal from typing import Any try: + from ..utils.process import is_process_alive, launch_process_with_timeout, terminate_process + from ..utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT +except ImportError: + # Fallback when executed as a standalone script. from utils.process import is_process_alive, launch_process_with_timeout, terminate_process from utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT -except ImportError: - # Fallback if running from different context - def is_process_alive(pid: int | str) -> bool: - """Best-effort fallback process liveness probe.""" - try: - pid = int(pid) - except (TypeError, ValueError): - return False - - if os.name == "nt": - try: - output = subprocess.check_output( - ["tasklist", "/FI", f"PID eq {pid}"], - stderr=subprocess.DEVNULL, - text=True, - timeout=5.0, - ) - except Exception: - return False - return str(pid) in output - - try: - os.kill(pid, 0) - return True - except OSError: - return False - - def launch_process_with_timeout( - command: list[str], - timeout_seconds: float = 10.0, - **kwargs: Any, - ): - _ = kwargs - raise RuntimeError( - "launch_process_with_timeout unavailable: utils.process import failed" - ) - - WORKER_CHECK_INTERVAL = 2.0 - PROCESS_TERMINATION_TIMEOUT = 5.0 def monitor_and_run(master_pid: int, command: list[str]) -> int: """Run command and monitor master process; return process exit code.""" @@ -98,15 +62,8 @@ def monitor_and_run(master_pid: int, command: list[str]) -> int: if worker_process.poll() is None: # Still running try: terminate_process(worker_process, timeout=PROCESS_TERMINATION_TIMEOUT) - except NameError: - # Fallback if terminate_process wasn't imported - worker_process.terminate() - try: - worker_process.wait(timeout=PROCESS_TERMINATION_TIMEOUT) - except subprocess.TimeoutExpired: - print("[Distributed] Worker didn't terminate gracefully, forcing kill...") - worker_process.kill() - worker_process.wait() + except TimeoutError: + print("[Distributed] Worker did not terminate within timeout.") print("[Distributed] Worker terminated.")