diff --git a/README.md b/README.md index f5725c5..8397368 100644 --- a/README.md +++ b/README.md @@ -38,11 +38,12 @@ ComfyUI_temp_khvcc_00034_@0 25x -ComfyUI Distributed supports three types of workers: +ComfyUI Distributed supports four types of workers: - **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch) - **Remote Workers** - GPUs on other computers in your network - **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels +- **Serverless (Runpod)** - Queue-based endpoint pools for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md). > For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md) diff --git a/api/__init__.py b/api/__init__.py index d20f665..a2c9ad8 100644 --- a/api/__init__.py +++ b/api/__init__.py @@ -1,4 +1,5 @@ from . import config_routes # noqa: F401 +from . import serverless_routes # noqa: F401 from . import tunnel_routes # noqa: F401 from . import worker_routes # noqa: F401 from . import job_routes # noqa: F401 diff --git a/api/config_routes.py b/api/config_routes.py index 02d1afb..f5cecd0 100644 --- a/api/config_routes.py +++ b/api/config_routes.py @@ -1,5 +1,6 @@ import json from contextlib import asynccontextmanager +from typing import Any from aiohttp import web import server @@ -24,6 +25,7 @@ except ImportError: save_config(config) from ..utils.logging import debug_log, log from ..utils.network import handle_api_error, normalize_host +from ..utils.serverless_config import public_config, validate_serverless_worker def _positive_int(value): @@ -86,22 +88,188 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None: target[key] = normalizer(value) if (normalizer and value is not None) else value +_SECRET_FIELDS = {"api_key", "clear_api_key"} +_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"} +_SERVERLESS_ONLY_FIELDS = { + "provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured", +} +_SERVERLESS_CONFIG_FIELDS = { + "worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency", + "timeout_seconds", "enabled", "name", "api_key", "clear_api_key", +} + + +class _ActivePoolCheckError(RuntimeError): + pass + + +def _error_response(message, status): + return web.json_response({"status": "error", "message": message}, status=status) + + +def _is_secret_field_name(key): + normalized = "".join(character for character in key.casefold() if character.isalnum()) + return normalized in {"apikey", "clearapikey"} or "credential" in normalized + + +def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0): + if isinstance(value, dict): + for key, item in value.items(): + field_name = key.casefold() if isinstance(key, str) else "" + if field_name and _is_secret_field_name(field_name): + if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS): + return True + if _contains_secret_fields( + item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1 + ): + return True + elif isinstance(value, (list, tuple)): + return any( + _contains_secret_fields( + item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1 + ) + for item in value + ) + return False + + +def _serverless_secret_store(): + from .orchestration.serverless_runtime import get_secret_store + return get_secret_store() + + +def _serverless_pool_is_active(worker_id): + from .orchestration.serverless_runtime import get_manager + manager = get_manager() + tasks = getattr(manager, "tasks", {}) + runs = getattr(manager, "runs", {}) + for run_id, task in tasks.items(): + if task.done(): + continue + run = runs.get(run_id) + worker = run.get("worker") if isinstance(run, dict) else None + if not isinstance(worker, dict): + continue + pool_id = worker.get("pool_id", worker.get("id")) + if str(pool_id) == str(worker_id): + return True + return False + + +def _ensure_pool_idle(worker_id): + try: + active = _serverless_pool_is_active(worker_id) + except Exception: + raise _ActivePoolCheckError from None + if active: + raise _ActivePoolCheckError + + +def _worker_by_id(workers, worker_id): + for worker in workers: + if isinstance(worker, dict) and worker.get("id") == worker_id: + return worker + return None + + +def _worker_type(worker): + return worker.get("type", "local") if isinstance(worker, dict) else None + + +def _strip_sensitive_fields(value: Any) -> Any: + if isinstance(value, dict): + return { + key: _strip_sensitive_fields(item) + for key, item in value.items() + if not (isinstance(key, str) and _is_secret_field_name(key)) + } + if isinstance(value, list): + return [_strip_sensitive_fields(item) for item in value] + return value + + +def _clean_serverless_worker(worker): + clean = _strip_sensitive_fields(dict(worker)) + for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}: + clean.pop(field, None) + return clean + + +def _validate_bulk_workers(workers): + validated = [] + for worker in workers: + if isinstance(worker, dict) and worker.get("type") == "serverless": + clean = _clean_serverless_worker(worker) + validated.append(validate_serverless_worker(clean)) + else: + validated.append(worker) + return validated + + +def _clean_bulk_type_conversions(old_workers, new_workers): + old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None} + for index, worker in enumerate(new_workers): + if not isinstance(worker, dict) or worker.get("id") is None: + continue + old = old_by_id.get(str(worker["id"])) + if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless": + continue + converted = dict(worker) + for field in _SERVERLESS_ONLY_FIELDS: + converted.pop(field, None) + new_workers[index] = converted + return new_workers + + +def _bulk_serverless_changes(old_workers, new_workers): + old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None} + new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None} + secret_ids = set() + for worker_id, old in old_by_id.items(): + new = new_by_id.get(worker_id) + if old.get("type") == "serverless": + if new is None or new.get("type") != "serverless": + _ensure_pool_idle(worker_id) + secret_ids.add(worker_id) + elif (old.get("endpoint_id") != new.get("endpoint_id") + or old.get("concurrency", 1) != new.get("concurrency", 1)): + _ensure_pool_idle(worker_id) + elif new is not None and new.get("type") == "serverless": + secret_ids.add(worker_id) + return secret_ids + + +def _restore_secret(store, worker_id, prior_secret): + try: + if prior_secret is None: + store.delete(worker_id) + else: + store.set(worker_id, prior_secret) + except Exception: + pass + + @server.PromptServer.instance.routes.get("/distributed/config") async def get_config_endpoint(request): - config = load_config() - return web.json_response(config) + try: + config = load_config() + return web.json_response(public_config(config, _serverless_secret_store())) + except Exception: + return _error_response("Unable to read config", 500) @server.PromptServer.instance.routes.post("/distributed/config") async def update_config_endpoint(request): - """Bulk config update with schema validation.""" + """Bulk config update with schema validation and credential redaction.""" try: data = await request.json() - except Exception as e: - return await handle_api_error(request, f"Invalid JSON payload: {e}", 400) + except Exception: + return _error_response("Invalid JSON payload", 400) if not isinstance(data, dict): - return await handle_api_error(request, "Config payload must be an object", 400) + return _error_response("Config payload must be an object", 400) + if _contains_secret_fields(data): + return _error_response("API keys cannot be set through bulk config", 400) validated_settings = {} validated_root = {} @@ -118,9 +286,15 @@ async def update_config_endpoint(request): continue if validator and not validator(value): - errors.append(f"{key}: value {value!r} failed validation") + errors.append(f"{key}: failed validation") continue + if key == "workers": + try: + value = _validate_bulk_workers(value) + except (TypeError, ValueError): + return _error_response("Invalid serverless worker configuration", 400) + if key in _SETTINGS_FIELDS: validated_settings[key] = value else: @@ -133,15 +307,45 @@ async def update_config_endpoint(request): "message": "; ".join(errors), }, status=400) + store = None + changed_secrets = [] + secret_changes_applied = False try: async with config_transaction() as config: + if "workers" in validated_root: + old_workers = config.get("workers", []) + validated_root["workers"] = _clean_bulk_type_conversions( + old_workers, validated_root["workers"] + ) + changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"]) + for worker in validated_root["workers"]: + if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False): + store = store or _serverless_secret_store() + if worker["id"] in changed_ids or not store.has(worker["id"]): + return _error_response("Save an API key before enabling this pool", 400) + if changed_ids: + store = store or _serverless_secret_store() + for worker_id in sorted(changed_ids): + prior_secret = store.get(worker_id) + changed_secrets.append((worker_id, prior_secret)) + if prior_secret is not None: + store.delete(worker_id) + secret_changes_applied = True + 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) + response_config = public_config(config, store or _serverless_secret_store()) + + return web.json_response({"status": "success", "config": response_config}) + except _ActivePoolCheckError: + return _error_response("Wait for active serverless jobs before changing this pool", 409) + except Exception: + if secret_changes_applied and store is not None: + for worker_id, prior_secret in changed_secrets: + _restore_secret(store, worker_id, prior_secret) + return _error_response("Unable to update config", 500) @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}") @@ -164,79 +368,205 @@ async def queue_status_endpoint(request): @server.PromptServer.instance.routes.post("/distributed/config/update_worker") async def update_worker_endpoint(request): + store = None + prior_secret = None + secret_changed = False + worker_id = None try: - data = await request.json() + try: + data = await request.json() + except Exception: + return _error_response("Invalid JSON payload", 400) + if not isinstance(data, dict): + return _error_response("Worker payload must be an object", 400) + worker_id = data.get("worker_id") - if worker_id is None: - return await handle_api_error(request, "Missing worker_id", 400) + return _error_response("Missing worker_id", 400) + if _contains_secret_fields(data, allow_root_serverless_secrets=True): + return _error_response("Invalid worker credentials", 400) + + result = None async with config_transaction() as config: - worker_found = False workers = config.setdefault("workers", []) + existing = _worker_by_id(workers, worker_id) + current_type = _worker_type(existing) + target_type = data.get("type", current_type or "local") - for worker in workers: - if worker["id"] == worker_id: - _apply_field_patch(worker, data, _WORKER_FIELDS) - worker_found = True - break + if target_type == "serverless": + unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS + if unknown: + return _error_response("Unknown serverless worker field", 400) + if not isinstance(worker_id, str) or not worker_id.strip(): + return _error_response("worker_id must be a non-empty string", 400) + if "api_key" in data and not isinstance(data["api_key"], str): + return _error_response("Invalid worker credentials", 400) + if "clear_api_key" in data and type(data["clear_api_key"]) is not bool: + return _error_response("Invalid worker credentials", 400) + supplied_key = data.get("api_key", "") + clear_key = data.get("clear_api_key", False) + if supplied_key and clear_key: + return _error_response("Choose either api_key or clear_api_key", 400) + if "enabled" in data and type(data["enabled"]) is not bool: + return _error_response("enabled must be a boolean", 400) + if "name" in data and not isinstance(data["name"], str): + return _error_response("name must be a string", 400) - if not worker_found: - # If worker not found and all required fields are provided, create new worker - if all(key in data for key in ["name", "port", "cuda_device"]): - new_worker = { - "id": worker_id, - "name": data["name"], - "host": normalize_host(data.get("host", "localhost")), - "port": data["port"], - "cuda_device": data["cuda_device"], - "enabled": data.get("enabled", False), - "extra_args": data.get("extra_args", ""), - "type": data.get("type", "local") - } - workers.append(new_worker) + candidate = dict(existing or {}) + candidate["id"] = worker_id + candidate["type"] = "serverless" + candidate.setdefault("provider", "runpod") + candidate.setdefault("enabled", False) + for field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "enabled", "name"): + if field in data: + candidate[field] = data[field] + candidate = _clean_serverless_worker(candidate) + validated = validate_serverless_worker(candidate) + if "enabled" in validated and type(validated["enabled"]) is not bool: + return _error_response("enabled must be a boolean", 400) + if "name" in validated and not isinstance(validated["name"], str): + return _error_response("name must be a string", 400) + + store = _serverless_secret_store() + prior_secret = store.get(worker_id) + converting_to_serverless = current_type != "serverless" + next_secret = None if converting_to_serverless else prior_secret + if supplied_key: + next_secret = supplied_key + elif clear_key: + next_secret = None + + endpoint_changed = ( + current_type == "serverless" + and existing.get("endpoint_id") != validated.get("endpoint_id") + ) + concurrency_changed = ( + current_type == "serverless" + and existing.get("concurrency", 1) != validated.get("concurrency", 1) + ) + if current_type == "serverless" and ( + endpoint_changed or concurrency_changed or next_secret != prior_secret + ): + _ensure_pool_idle(worker_id) + + if validated.get("enabled", False) and next_secret is None: + return _error_response("Save an API key before enabling this pool; disable it before clearing the key", 400) + + if next_secret != prior_secret: + if next_secret is None: + store.delete(worker_id) + else: + store.set(worker_id, next_secret) + secret_changed = True + + if existing is None: + workers.append(validated) else: - return await handle_api_error( - request, - f"Worker {worker_id} not found and missing required fields for creation", - 404, - ) + existing.clear() + existing.update(validated) + result = {"status": "success", "api_key_configured": next_secret is not None} + + elif existing is not None: + cleaned_existing = _strip_sensitive_fields(existing) + existing.clear() + existing.update(cleaned_existing) + if current_type == "serverless": + if _contains_secret_fields(data): + return _error_response("Invalid worker credentials", 400) + _ensure_pool_idle(worker_id) + store = _serverless_secret_store() + prior_secret = store.get(worker_id) + if prior_secret is not None: + store.delete(worker_id) + secret_changed = True + for field in _SERVERLESS_ONLY_FIELDS: + existing.pop(field, None) + elif _contains_secret_fields(data): + return _error_response("API keys are only accepted for serverless workers", 400) + _apply_field_patch(existing, data, _WORKER_FIELDS) + result = {"status": "success"} + elif all(key in data for key in ["name", "port", "cuda_device"]): + if _contains_secret_fields(data): + return _error_response("API keys are only accepted for serverless workers", 400) + new_worker = { + "id": worker_id, + "name": data["name"], + "host": normalize_host(data.get("host", "localhost")), + "port": data["port"], + "cuda_device": data["cuda_device"], + "enabled": data.get("enabled", False), + "extra_args": data.get("extra_args", ""), + "type": data.get("type", "local") + } + workers.append(new_worker) + result = {"status": "success"} + else: + return _error_response( + "Worker not found and missing required fields for creation", 404 + ) + + return web.json_response(result) + except _ActivePoolCheckError: + if secret_changed and store is not None: + _restore_secret(store, worker_id, prior_secret) + return _error_response("Wait for active serverless jobs before changing this pool", 409) + except ValueError: + if secret_changed and store is not None: + _restore_secret(store, worker_id, prior_secret) + return _error_response("Invalid worker configuration", 400) + except Exception: + if secret_changed and store is not None: + _restore_secret(store, worker_id, prior_secret) + return _error_response("Unable to update worker configuration", 500) - return web.json_response({"status": "success"}) - except Exception as e: - return await handle_api_error(request, e, 400) @server.PromptServer.instance.routes.post("/distributed/config/delete_worker") async def delete_worker_endpoint(request): + store = None + prior_secret = None + secret_changed = False + worker_id = None try: - data = await request.json() + try: + data = await request.json() + except Exception: + return _error_response("Invalid JSON payload", 400) + if not isinstance(data, dict): + return _error_response("Worker payload must be an object", 400) worker_id = data.get("worker_id") - if worker_id is None: - return await handle_api_error(request, "Missing worker_id", 400) - + return _error_response("Missing worker_id", 400) + + result = None async with config_transaction() as config: workers = config.get("workers", []) - - # Find and remove the worker - worker_index = -1 - for i, worker in enumerate(workers): - if worker["id"] == worker_id: - worker_index = i - break - - if worker_index == -1: - return await handle_api_error(request, f"Worker {worker_id} not found", 404) - - # Remove the worker - removed_worker = workers.pop(worker_index) - - return web.json_response({ + removed = _worker_by_id(workers, worker_id) + if removed is None: + return _error_response("Worker not found", 404) + if _worker_type(removed) == "serverless": + _ensure_pool_idle(worker_id) + store = _serverless_secret_store() + prior_secret = store.get(str(worker_id)) + if prior_secret is not None: + store.delete(str(worker_id)) + secret_changed = True + workers.remove(removed) + result = { "status": "success", - "message": f"Worker {removed_worker.get('name', worker_id)} deleted" - }) - except Exception as e: - return await handle_api_error(request, e, 400) + "message": "Worker deleted" if _worker_type(removed) == "serverless" + else f"Worker {removed.get('name', worker_id)} deleted", + } + + return web.json_response(result) + except _ActivePoolCheckError: + if secret_changed and store is not None: + _restore_secret(store, str(worker_id), prior_secret) + return _error_response("Wait for active serverless jobs before changing this pool", 409) + except Exception: + if secret_changed and store is not None: + _restore_secret(store, str(worker_id), prior_secret) + return _error_response("Unable to delete worker", 500) @server.PromptServer.instance.routes.post("/distributed/config/update_setting") async def update_setting_endpoint(request): diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 79fe59f..e2c02a4 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -55,6 +55,11 @@ async def worker_ws_is_active(worker): async def _probe_worker_active(worker, use_websocket, semaphore): async with semaphore: + if worker.get("type") == "serverless": + from .serverless_runtime import get_manager + # Authentication/configuration failures must surface, not silently skip paid work. + status = await get_manager().health(worker) + return worker, status["ok"] is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker)) return worker, is_active @@ -203,6 +208,10 @@ def _extract_queue_remaining(payload): async def _probe_worker_queue(worker, semaphore, probe_timeout): async with semaphore: + if worker.get("type") == "serverless": + from .serverless_runtime import get_manager + status = await get_manager().health(worker) + return {"worker": worker, "queue_remaining": status["queueRemaining"]} worker_url = build_worker_url(worker) payload = await probe_worker(worker_url, timeout=probe_timeout) if payload is None: diff --git a/api/orchestration/runpod_client.py b/api/orchestration/runpod_client.py new file mode 100644 index 0000000..64ac958 --- /dev/null +++ b/api/orchestration/runpod_client.py @@ -0,0 +1,97 @@ +"""Small Runpod queue API client. Never retries a potentially billable POST.""" +import asyncio +import json +import re + +import aiohttp + +API_BASE = "https://api.runpod.ai/v2" +MAX_REQUEST_BYTES = 10 * 1024 * 1024 +_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z") + + +class RunpodError(RuntimeError): + """Sanitized provider failure safe to display in the panel.""" + + +class AmbiguousSubmission(RunpodError): + """Submission may have been accepted: never retry it automatically.""" + + +def checked_id(value): + if not isinstance(value, str) or not _ID.fullmatch(value): + raise ValueError("Invalid Runpod endpoint or job ID") + return value + + +class RunpodClient: + def __init__(self, session, endpoint_id, api_key): + self.session = session + self.endpoint_id = checked_id(endpoint_id) + if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"): + raise ValueError("A valid Runpod API key is required") + self._api_key = api_key + + async def _request(self, method, path, payload=None): + # Base origin is code-owned, never a user-supplied URL. Do not follow redirects + # with credentials, or include response bodies/transport errors in exceptions. + url = f"{API_BASE}/{self.endpoint_id}/{path}" + attempts = 3 if method == "GET" else 1 + for attempt in range(attempts): + try: + async with self.session.request( + method, url, + headers={"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}, + json=payload, + timeout=aiohttp.ClientTimeout(total=30), + allow_redirects=False, + ) as response: + status = response.status + if status != 200: + if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts: + await asyncio.sleep(0.5 * (2 ** attempt)) + continue + if method == "POST" and path == "run" and status >= 500: + raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.") + raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})") + try: + result = await response.json() + except (ValueError, aiohttp.ContentTypeError): + if method == "POST" and path == "run": + raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None + raise RunpodError("Runpod returned an invalid JSON response") from None + if not isinstance(result, dict): + if method == "POST" and path == "run": + raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.") + raise RunpodError("Runpod returned an invalid response") + return result + except (asyncio.TimeoutError, aiohttp.ClientError, OSError): + if method == "POST" and path == "run": + raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None + if method == "GET" and attempt + 1 < attempts: + await asyncio.sleep(0.5 * (2 ** attempt)) + continue + raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None + raise RunpodError("Runpod request failed") + + async def health(self): + return await self._request("GET", "health") + + async def submit(self, workflow, images, timeout_seconds): + payload = { + "input": {"workflow": workflow, "images": images}, + "policy": {"executionTimeout": timeout_seconds * 1000, "ttl": timeout_seconds * 1000}, + } + if len(json.dumps(payload, allow_nan=False).encode("utf-8")) > MAX_REQUEST_BYTES: + raise ValueError("Serverless request exceeds the 10 MiB safety limit; reduce input media size") + result = await self._request("POST", "run", payload) + try: + return checked_id(result.get("id")) + except ValueError: + raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None + + async def status(self, job_id): + return await self._request("GET", f"status/{checked_id(job_id)}") + + async def cancel(self, job_id): + return await self._request("POST", f"cancel/{checked_id(job_id)}") diff --git a/api/orchestration/serverless_ledger.py b/api/orchestration/serverless_ledger.py new file mode 100644 index 0000000..5f96020 --- /dev/null +++ b/api/orchestration/serverless_ledger.py @@ -0,0 +1,57 @@ +"""Private, atomic job-ID journal. Contains identifiers/status only, never payloads.""" +import json +import os +from pathlib import Path +import tempfile +import threading + +TERMINAL = {"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"} +FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"} + + +class JobLedger: + def __init__(self, path): + self.path = Path(path) + self._lock = threading.RLock() + + def _read(self): + if self.path.is_symlink(): + raise RuntimeError("Unsafe serverless job journal path") + if not self.path.exists(): + return {} + try: + data = json.loads(self.path.read_text(encoding="utf-8")) + if not isinstance(data, dict) or any(not isinstance(v, dict) or set(v) - FIELDS for v in data.values()): + raise ValueError() + return data + except (ValueError, OSError): + raise RuntimeError("Serverless job journal is unreadable; preserve it for recovery before submitting more work") from None + + def all(self): + with self._lock: + return list(self._read().values()) + + def update(self, record_id, **fields): + if set(fields) - FIELDS: + raise ValueError("Unexpected serverless journal field") + with self._lock: + data = self._read() + data.setdefault(record_id, {"record_id": record_id}).update(fields) + # Keep every unresolved job, plus a bounded terminal history. + completed = sorted((k for k, v in data.items() if v.get("status") in TERMINAL), key=lambda k: data[k].get("updated_at", 0)) + for key in completed[:-200]: + data.pop(key) + if self.path.parent.is_symlink(): + raise RuntimeError("Unsafe serverless journal directory") + self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=".jobs-", dir=self.path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + json.dump(data, stream, allow_nan=False) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, self.path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + return dict(data[record_id]) diff --git a/api/orchestration/serverless_media.py b/api/orchestration/serverless_media.py new file mode 100644 index 0000000..d474eb6 --- /dev/null +++ b/api/orchestration/serverless_media.py @@ -0,0 +1,571 @@ +"""Bounded image transport for generic Runpod-backed ComfyUI jobs.""" + +import asyncio +import base64 +import binascii +import hashlib +import io +import ipaddress +import os +import re +import socket +import stat +import warnings +from pathlib import Path, PurePosixPath +from urllib.parse import urljoin, urlsplit + +import aiohttp +import numpy as np +from PIL import Image, UnidentifiedImageError + + +MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024 +MAX_INPUT_TOTAL_BYTES = 7 * 1024 * 1024 +MAX_OUTPUT_FILES = 1024 +MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024 +MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024 +MAX_IMAGE_PIXELS = 40_000_000 +MAX_OUTPUT_ARRAY_BYTES = 256 * 1024 * 1024 +MAX_REDIRECTS = 3 +_DOWNLOAD_TIMEOUT_SECONDS = 60 +_CHUNK_SIZE = 64 * 1024 + +_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} +_UNSUPPORTED_IMAGE_EXTENSIONS = {".gif", ".tif", ".tiff", ".avif", ".heic", ".ico", ".svg"} +_AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus", ".aiff", ".aif", ".wma"} +_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"} +_MEDIA_EXTENSIONS = _IMAGE_EXTENSIONS | _UNSUPPORTED_IMAGE_EXTENSIONS | _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS +_MEDIA_ANNOTATION = re.compile(r"\s+\[([A-Za-z0-9_-]{1,64})\]$") +_PREFIX_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\Z", re.ASCII) +_SAFE_FILENAME_RE = re.compile(r"[A-Za-z0-9_.-]{1,255}\Z", re.ASCII) + + +class ServerlessMediaError(ValueError): + """Sanitized input, artifact, or download validation failure.""" + + +def _folder_roots(folder_paths): + roots = {} + for kind, getter_name in ( + ("input", "get_input_directory"), + ("output", "get_output_directory"), + ("temp", "get_temp_directory"), + ): + getter = getattr(folder_paths, getter_name, None) + if not callable(getter): + raise ServerlessMediaError("ComfyUI media directories are unavailable") + try: + roots[kind] = Path(getter()).resolve(strict=True) + except Exception: + raise ServerlessMediaError("ComfyUI media directories are unavailable") from None + return roots + + +def _inside_roots(path, roots): + for root in roots.values(): + try: + path.relative_to(root) + return True + except ValueError: + continue + return False + + +def _normalize_media_reference(value): + if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value: + raise ServerlessMediaError("Image input reference is invalid") + value = value.replace("\\", "/") + explicit_root = None + annotation = _MEDIA_ANNOTATION.search(value) + if annotation: + marker = annotation.group(1).lower() + if marker not in {"input", "output", "temp"}: + raise ServerlessMediaError("Image input annotation is unsupported") + explicit_root = marker + value = value[: annotation.start()] + if "://" in value: + raise ServerlessMediaError("Remote media references are not supported as image inputs") + if re.match(r"^[A-Za-z]:", value) or value.startswith("//"): + raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories") + + path = PurePosixPath(value) + if any(part == ".." for part in path.parts): + raise ServerlessMediaError("Image input path traversal is not allowed") + if not path.parts or value in {".", ""}: + raise ServerlessMediaError("Image input reference is invalid") + + if explicit_root: + value = f"{value} [{explicit_root}]" + suffix = Path(path.name).suffix.lower() + return value, suffix + + +def _resolve_input_path(value, suffix, folder_paths, roots): + if suffix not in _IMAGE_EXTENSIONS: + if suffix in _MEDIA_EXTENSIONS: + raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport") + raise ServerlessMediaError("Unsupported serverless image file type") + + try: + source = Path(value) + if source.is_absolute(): + annotation = _MEDIA_ANNOTATION.search(value) + unannotated = value[:annotation.start()] if annotation else value + resolved = Path(unannotated).resolve(strict=True) + else: + resolver = getattr(folder_paths, "get_annotated_filepath", None) + if not callable(resolver): + raise ServerlessMediaError("ComfyUI media path resolver is unavailable") + resolved = Path(resolver(value)).resolve(strict=True) + except ServerlessMediaError: + raise + except Exception: + raise ServerlessMediaError("Image input file is missing or invalid") from None + + if not _inside_roots(resolved, roots): + raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories") + try: + mode = resolved.stat().st_mode + except OSError: + raise ServerlessMediaError("Image input file is missing or invalid") from None + if not stat.S_ISREG(mode): + raise ServerlessMediaError("Image input is not a regular file") + return resolved + + +def _iter_image_references(workflow): + if not isinstance(workflow, dict): + raise ServerlessMediaError("Serverless workflow must be an API graph") + references = [] + for node_id, node in workflow.items(): + if not isinstance(node, dict): + continue + class_type = node.get("class_type") + inputs = node.get("inputs", {}) + if class_type == "LoadImage" and not isinstance(inputs, dict): + raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported") + if not isinstance(inputs, dict): + continue + if class_type == "LoadImage": + value = inputs.get("image") + if not isinstance(value, str): + raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported") + references.append((node_id, inputs, "image", value, True)) + + for key in ("image", "file"): + if key == "image" and class_type == "LoadImage": + continue + value = inputs.get(key) + if not isinstance(value, str): + continue + normalized, suffix = _normalize_media_reference(value) + if suffix in _IMAGE_EXTENSIONS: + references.append((node_id, inputs, key, normalized, False)) + elif suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS: + raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport") + elif suffix in _UNSUPPORTED_IMAGE_EXTENSIONS or (key == "image" and suffix): + raise ServerlessMediaError("Unsupported serverless image file type") + + for key in ("audio", "video"): + value = inputs.get(key) + if not isinstance(value, str): + continue + _normalized, suffix = _normalize_media_reference(value) + if suffix in _MEDIA_EXTENSIONS: + raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport") + return references + + +def _read_limited(path, expected_size): + try: + descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)) + with os.fdopen(descriptor, "rb") as handle: + info = os.fstat(handle.fileno()) + if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES: + raise ServerlessMediaError("Image input exceeds the permitted size") + data = handle.read(MAX_INPUT_IMAGE_BYTES + 1) + except ServerlessMediaError: + raise + except OSError: + raise ServerlessMediaError("Image input file could not be read safely") from None + if len(data) > MAX_INPUT_IMAGE_BYTES or len(data) != expected_size: + raise ServerlessMediaError("Image input exceeds the permitted size or changed while being read") + return data + + +async def prepare_images(workflow: dict) -> list[dict]: + """Pack referenced local images and rewrite only their graph input references. + + ``workflow`` is expected to be a private API-graph copy owned by the caller. + Returned ``name`` values are flat generated names suitable for ``input.images``. + """ + references = _iter_image_references(workflow) + if not references: + return [] + + # Import only in ComfyUI, keeping this module importable in the serverless client. + try: + import folder_paths + except Exception: + raise ServerlessMediaError("ComfyUI media directories are unavailable") from None + + roots = _folder_roots(folder_paths) + resolved_refs = [] + unique_paths = {} + total_source_bytes = 0 + for node_id, inputs, key, value, _is_load_image in references: + normalized, suffix = _normalize_media_reference(value) + if key == "image" and suffix not in _IMAGE_EXTENSIONS: + if suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS: + raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport") + raise ServerlessMediaError("Unsupported serverless image file type") + path = _resolve_input_path(normalized, suffix, folder_paths, roots) + if path not in unique_paths: + try: + size = path.stat().st_size + except OSError: + raise ServerlessMediaError("Image input file is missing or invalid") from None + if size < 0 or size > MAX_INPUT_IMAGE_BYTES: + raise ServerlessMediaError("Image input exceeds the permitted size") + total_source_bytes += size + if total_source_bytes > MAX_INPUT_TOTAL_BYTES: + raise ServerlessMediaError("Total serverless image inputs exceed the permitted size") + unique_paths[path] = {"size": size, "suffix": suffix} + resolved_refs.append((inputs, key, path)) + + packed_by_name = {} + path_to_name = {} + for path, metadata in unique_paths.items(): + data = _read_limited(path, metadata["size"]) + digest = hashlib.sha256(data).hexdigest() + name = f"serverless_{digest}{metadata['suffix']}" + path_to_name[path] = name + packed_by_name.setdefault(name, {"name": name, "image": base64.b64encode(data).decode("ascii")}) + + for inputs, key, path in resolved_refs: + inputs[key] = path_to_name[path] + return [packed_by_name[name] for name in sorted(packed_by_name)] + + +def _validate_artifact_url(url): + if not isinstance(url, str) or not url or any(ord(char) < 0x20 or char.isspace() for char in url): + raise ServerlessMediaError("Serverless image artifact URL is invalid") + try: + parsed = urlsplit(url) + port = parsed.port + except ValueError: + raise ServerlessMediaError("Serverless image artifact URL is invalid") from None + if ( + parsed.scheme.lower() != "https" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or port not in (None, 443) + or parsed.fragment + ): + raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL") + host = parsed.hostname.rstrip(".").lower() + if host in {"localhost", "metadata", "metadata.google.internal"} or host.endswith((".localhost", ".local", ".internal")): + raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL") + try: + address = ipaddress.ip_address(host) + except ValueError: + return url + if not address.is_global: + raise ServerlessMediaError("Serverless image artifact URL must resolve to a public address") + return url + + +class _PublicResolver(aiohttp.abc.AbstractResolver): + """Resolve once, reject non-public answers, and pin the checked address.""" + + async def _getaddrinfo(self, host, port, family): + return await asyncio.get_running_loop().getaddrinfo( + host, port, family=family, type=socket.SOCK_STREAM + ) + + async def resolve(self, host, port=0, family=socket.AF_INET): + try: + results = await self._getaddrinfo(host, port, family) + except (OSError, asyncio.TimeoutError): + raise OSError("Public artifact host could not be resolved") from None + if not results: + raise OSError("Public artifact host could not be resolved") + + resolved = [] + seen = set() + for result_family, _socktype, proto, _canonname, sockaddr in results: + try: + address = ipaddress.ip_address(sockaddr[0]) + except (ValueError, TypeError, IndexError): + raise OSError("Public artifact host resolved to an invalid address") from None + if not address.is_global: + raise OSError("Public artifact host resolved to a non-public address") + key = (str(address), result_family) + if key in seen: + continue + seen.add(key) + resolved.append({ + "hostname": host, + "host": str(address), + "port": port, + "family": result_family, + "proto": proto, + "flags": 0, + }) + if not resolved: + raise OSError("Public artifact host could not be resolved") + return resolved + + async def close(self): + return None + + +async def download_public_artifact(url: str, max_bytes: int = MAX_OUTPUT_IMAGE_BYTES) -> bytes: + """Download a signed public HTTPS artifact without provider credentials.""" + if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or not 0 < max_bytes <= MAX_OUTPUT_IMAGE_BYTES: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + current_url = _validate_artifact_url(url) + resolver = _PublicResolver() + connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False, limit=4) + timeout = aiohttp.ClientTimeout( + total=_DOWNLOAD_TIMEOUT_SECONDS, + connect=10, + sock_read=15, + ) + try: + async with aiohttp.ClientSession( + connector=connector, + timeout=timeout, + trust_env=False, + cookie_jar=aiohttp.DummyCookieJar(), + ) as session: + for redirect_count in range(MAX_REDIRECTS + 1): + _validate_artifact_url(current_url) + async with session.get(current_url, allow_redirects=False) as response: + if response.status in {301, 302, 303, 307, 308}: + location = response.headers.get("Location") + if not location or redirect_count >= MAX_REDIRECTS: + raise ServerlessMediaError("Serverless image artifact redirect is invalid") + current_url = _validate_artifact_url(urljoin(current_url, location)) + continue + if response.status != 200: + raise ServerlessMediaError("Serverless image artifact download failed") + payload = bytearray() + async for chunk in response.content.iter_chunked(_CHUNK_SIZE): + if not isinstance(chunk, bytes): + raise ServerlessMediaError("Serverless image artifact response is invalid") + if len(payload) + len(chunk) > max_bytes: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + payload.extend(chunk) + if not payload: + raise ServerlessMediaError("Serverless image artifact is empty") + return bytes(payload) + except ServerlessMediaError: + raise + except Exception: + raise ServerlessMediaError("Serverless image artifact download failed") from None + raise ServerlessMediaError("Serverless image artifact download failed") + + +def _validated_manifest(manifest): + if not isinstance(manifest, list) or len(manifest) > MAX_OUTPUT_FILES: + raise ServerlessMediaError("Serverless image manifest is invalid") + normalized = [] + seen_prefixes = set() + prefixes = [] + for item in manifest: + if not isinstance(item, dict): + raise ServerlessMediaError("Serverless image manifest is invalid") + if item.get("kind") != "image": + raise ServerlessMediaError("Only image collectors are supported by serverless media transport") + for key in ("collector_id", "job_id", "node_id"): + value = item.get(key) + if not isinstance(value, (str, int)) or isinstance(value, bool) or not str(value): + raise ServerlessMediaError("Serverless image manifest is invalid") + prefix = item.get("prefix") + if not isinstance(prefix, str) or not _PREFIX_RE.fullmatch(prefix) or prefix in {".", ".."}: + raise ServerlessMediaError("Serverless image manifest contains an unsafe filename prefix") + folded = prefix.casefold() + if folded in seen_prefixes: + raise ServerlessMediaError("Serverless image manifest prefixes must be unique") + seen_prefixes.add(folded) + prefixes.append(prefix) + normalized.append({ + "collector_id": str(item["collector_id"]), + "job_id": str(item["job_id"]), + "node_id": str(item["node_id"]), + "prefix": prefix, + }) + for first in prefixes: + for second in prefixes: + if first != second and second.casefold().startswith(first.casefold() + "_"): + raise ServerlessMediaError("Serverless image manifest prefixes are ambiguous") + return normalized + + +def _artifact_filename_matches(filename, prefixes): + if not isinstance(filename, str): + return None + basename = filename.replace("\\", "/").rsplit("/", 1)[-1] + matches = [item for item in prefixes if basename.startswith(item["prefix"] + "_")] + if len(matches) > 1: + raise ServerlessMediaError("Serverless image artifact matches multiple collectors") + return matches[0] if matches else None + + +def _canonical_artifact_identity(entry, filename): + artifact_path = entry.get("artifact_path") + artifact_type = entry.get("artifact_type", "output") + if artifact_path is not None: + if not isinstance(artifact_path, str) or not artifact_path or "\x00" in artifact_path: + raise ServerlessMediaError("Serverless image artifact identity is invalid") + normalized_path = os.path.normcase(os.path.normpath(artifact_path.replace("\\", "/"))) + else: + subfolder = entry.get("subfolder", "") + if not isinstance(subfolder, str) or "\x00" in subfolder: + raise ServerlessMediaError("Serverless image artifact identity is invalid") + normalized_path = os.path.normcase(os.path.normpath(f"{subfolder}/{filename}")) + return str(artifact_type), normalized_path + + +def _decode_base64_data(value, max_bytes=MAX_OUTPUT_IMAGE_BYTES): + if not isinstance(value, str) or not value: + raise ServerlessMediaError("Serverless image artifact data is invalid") + max_bytes = min(max_bytes, MAX_OUTPUT_IMAGE_BYTES) + max_encoded = ((max_bytes + 2) // 3) * 4 + if len(value) > max_encoded: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + padding = len(value) - len(value.rstrip("=")) + estimated_size = (len(value) // 4) * 3 - padding + if estimated_size > max_bytes: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + try: + payload = base64.b64decode(value, validate=True) + except (binascii.Error, ValueError): + raise ServerlessMediaError("Serverless image artifact data is invalid") from None + if not payload: + raise ServerlessMediaError("Serverless image artifact is empty") + if len(payload) > max_bytes: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + return payload + + +def _decode_image_array(payload, total_array_bytes): + try: + with warnings.catch_warnings(): + warnings.simplefilter("error", Image.DecompressionBombWarning) + image = Image.open(io.BytesIO(payload)) + width, height = image.size + if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS: + raise ServerlessMediaError("Serverless image dimensions exceed the permitted limit") + has_alpha = image.mode in {"RGBA", "LA"} or "transparency" in image.info + mode = "RGBA" if has_alpha else "RGB" + output_bytes = width * height * (4 if has_alpha else 3) * 4 + if total_array_bytes + output_bytes > MAX_OUTPUT_ARRAY_BYTES: + raise ServerlessMediaError("Decoded serverless images exceed the permitted memory limit") + converted = image.convert(mode) + converted.load() + array = np.array(converted, dtype=np.float32, copy=True) + np.multiply(array, 1.0 / 255.0, out=array) + return array, output_bytes + except ServerlessMediaError: + raise + except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning): + raise ServerlessMediaError("Serverless image artifact could not be decoded") from None + except Exception: + raise ServerlessMediaError("Serverless image artifact could not be decoded") from None + + +async def collect_images(output: dict, manifest: list, download=None) -> dict[str, list[np.ndarray]]: + """Validate and decode expected image artifacts, grouped by serverless job ID.""" + if not isinstance(output, dict): + raise ServerlessMediaError("Serverless job output is invalid") + nested = output.get("output") + data = nested if isinstance(nested, dict) else output + for candidate in (output, data): + if candidate.get("error") or candidate.get("errors"): + raise ServerlessMediaError("Serverless job reported an execution error") + + expected = _validated_manifest(manifest) + if not expected: + return {} + artifacts = data.get("images", []) + if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES: + raise ServerlessMediaError("Serverless job returned too many image files") + + selected = {item["prefix"]: [] for item in expected} + identities = set() + total_payload_bytes = 0 + for entry in artifacts: + if not isinstance(entry, dict): + continue + filename = entry.get("filename") + manifest_entry = _artifact_filename_matches(filename, expected) + if manifest_entry is None: + continue + normalized_filename = filename.replace("\\", "/") + if "/" in normalized_filename or not _SAFE_FILENAME_RE.fullmatch(filename): + raise ServerlessMediaError("Serverless image artifact filename is invalid") + if Path(filename).suffix.lower() not in _IMAGE_EXTENSIONS: + raise ServerlessMediaError("Serverless image artifact has an unsupported file type") + if "artifact_type" in entry and entry["artifact_type"] != "output": + raise ServerlessMediaError("Expected serverless image artifact was not an output image") + identity = _canonical_artifact_identity(entry, filename) + if identity in identities: + raise ServerlessMediaError("Duplicate serverless image artifact identity") + identities.add(identity) + match = re.match(rf"^{re.escape(manifest_entry['prefix'])}_(\d+)(?:_|\.)", filename) + if match is None: + raise ServerlessMediaError("Expected serverless image filename counter is invalid") + artifact_type = entry.get("type") + remaining_bytes = MAX_OUTPUT_TOTAL_BYTES - total_payload_bytes + if remaining_bytes <= 0: + raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size") + if artifact_type == "base64": + payload = _decode_base64_data(entry.get("data"), remaining_bytes) + elif artifact_type in {"gcs_url", "s3_url", "url"}: + url = entry.get("data") or entry.get("url") + url = _validate_artifact_url(url) + try: + if download is None: + payload = await download_public_artifact( + url, max_bytes=min(MAX_OUTPUT_IMAGE_BYTES, remaining_bytes) + ) + else: + payload = await download(url) + except ServerlessMediaError: + raise + except Exception: + raise ServerlessMediaError("Serverless image artifact download failed") from None + if not isinstance(payload, bytes) or not payload: + raise ServerlessMediaError("Serverless image artifact response is invalid") + if len(payload) > MAX_OUTPUT_IMAGE_BYTES: + raise ServerlessMediaError("Serverless image artifact exceeds the permitted size") + else: + raise ServerlessMediaError("Serverless image artifact transport type is unsupported") + total_payload_bytes += len(payload) + if total_payload_bytes > MAX_OUTPUT_TOTAL_BYTES: + raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size") + selected[manifest_entry["prefix"]].append((int(match.group(1)), filename, identity, payload)) + + if any(not selected[item["prefix"]] for item in expected): + raise ServerlessMediaError("Serverless job is missing an expected collector image") + + result = {} + total_array_bytes = 0 + for manifest_entry in expected: + prefix = manifest_entry["prefix"] + rows = sorted(selected[prefix], key=lambda row: (row[0], row[1], row[2])) + arrays = [] + shape = None + for _counter, _filename, _identity, payload in rows: + array, added_bytes = _decode_image_array(payload, total_array_bytes) + total_array_bytes += added_bytes + if shape is None: + shape = array.shape + elif array.shape != shape: + raise ServerlessMediaError("Images from one serverless collector have inconsistent dimensions") + arrays.append(array) + result.setdefault(manifest_entry["job_id"], []).extend(arrays) + return result diff --git a/api/orchestration/serverless_prompt.py b/api/orchestration/serverless_prompt.py new file mode 100644 index 0000000..e45ec6c --- /dev/null +++ b/api/orchestration/serverless_prompt.py @@ -0,0 +1,391 @@ +"""Compile prepared distributed prompts into worker-side ComfyUI image jobs. + +This module is deliberately pure: it only copies and rewrites an API prompt +mapping and never imports or executes ComfyUI nodes. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +import re + + +_SUPPORTED_DISTRIBUTED = { + "DistributedCollector", + "DistributedSeed", + "DistributedValue", + "DistributedModelName", +} +_REMOVED_SCALAR_NODES = { + "DistributedSeed", + "DistributedValue", + "DistributedModelName", +} +_PRIMITIVE_NODES = { + "PrimitiveBoolean": False, + "PrimitiveFloat": 0.0, + "PrimitiveInt": 0, + "PrimitiveString": "", +} +_REMOVED_NODE_OUTPUTS = { + "DistributedCollector": {0, 1}, + "DistributedSeed": {0}, + "DistributedValue": {0}, + "DistributedModelName": {0}, +} +_SEED_DEFAULT = 1125899906842 + + +def _is_link(value): + return isinstance(value, list) and len(value) == 2 + + +def _is_unsupported_distributed(class_type): + if not isinstance(class_type, str) or class_type in _SUPPORTED_DISTRIBUTED: + return False + lowered = class_type.lower() + return "distributed" in lowered or "batchdivider" in lowered + + +def _safe_component(value): + cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_-") + return cleaned[:24] or "empty" + + +def _filename_prefix(job_id, worker_id, collector_id, kind): + identity = json.dumps( + [job_id, worker_id, collector_id, kind], + ensure_ascii=False, + separators=(",", ":"), + ) + digest = hashlib.sha256(identity.encode("utf-8")).hexdigest() + parts = ( + "serverless", + _safe_component(job_id), + _safe_component(worker_id), + _safe_component(collector_id), + kind, + digest, + ) + return "_".join(parts) + + +def _worker_index(worker_id): + """Mirror nodes.utilities worker-index parsing; None means fallback.""" + try: + if worker_id.startswith("worker_"): + return int(worker_id.split("_")[1]) + return int(worker_id) + except (ValueError, IndexError): + return None + + +def _coerce(value, value_type): + if value_type == "INT": + return int(float(value)) + if value_type == "FLOAT": + return float(value) + return value + + +def _coerce_safe(value, value_type): + try: + return _coerce(value, value_type) + except (TypeError, ValueError): + return value + + +def _stringify_model_name(value): + if isinstance(value, str): + return value + if isinstance(value, (int, float, bool)): + return str(value) + try: + return json.dumps(value, indent=4) + except (TypeError, ValueError): + return str(value) + + +def _model_name_result(text): + if isinstance(text, list): + values = [_stringify_model_name(value) for value in text] + return values[0] if len(values) == 1 else values + return _stringify_model_name(text) + + +def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]: + """Compile a prepared worker prompt to an image-only serverless graph. + + Every DistributedCollector's original ``images`` input is connected to a + new SaveImage node. Only those savers and their upstream dependencies are + retained. DistributedSeed/Value/ModelName links are resolved as scalar + literals; unsupported node execution is never attempted. + """ + if not isinstance(prepared_prompt, dict): + raise ValueError("prepared prompt must be an API graph mapping") + + source_nodes = {} + for raw_id, node in prepared_prompt.items(): + node_id = str(raw_id) + if node_id in source_nodes: + raise ValueError("prompt has duplicate node IDs after string normalization") + if not isinstance(node, dict): + raise ValueError(f"prompt node {node_id!r} must be an object") + class_type = node.get("class_type") + if not isinstance(class_type, str): + raise ValueError(f"prompt node {node_id!r} is missing a class_type") + if _is_unsupported_distributed(class_type): + raise ValueError(f"unsupported distributed or batch node: {class_type}") + inputs = node.get("inputs", {}) + if not isinstance(inputs, dict): + raise ValueError(f"prompt node {node_id!r} inputs must be an object") + source_nodes[node_id] = node + + if not source_nodes: + raise ValueError("prepared prompt contains no nodes") + + # Validate every API link before any graph rewriting. Slots on the nodes + # removed by this compiler have a known schema and are checked exactly. + for node_id, node in source_nodes.items(): + for input_name, value in node.get("inputs", {}).items(): + if not _is_link(value): + continue + upstream_id, output_slot = value + upstream_id = str(upstream_id) + if upstream_id not in source_nodes: + raise ValueError( + f"invalid link at node {node_id!r} input {input_name!r}: " + f"source {upstream_id!r} does not exist" + ) + if isinstance(output_slot, bool) or not isinstance(output_slot, int) or output_slot < 0: + raise ValueError( + f"invalid link at node {node_id!r} input {input_name!r}: " + "output slot must be a nonnegative integer" + ) + upstream_class = source_nodes[upstream_id]["class_type"] + valid_slots = _REMOVED_NODE_OUTPUTS.get(upstream_class) + if valid_slots is not None and output_slot not in valid_slots: + raise ValueError( + f"invalid output slot {output_slot} for removed node " + f"{upstream_class} {upstream_id!r}" + ) + + collector_ids = sorted( + (node_id for node_id, node in source_nodes.items() if node["class_type"] == "DistributedCollector"), + key=str, + ) + if not collector_ids: + raise ValueError("prepared prompt contains no DistributedCollector nodes") + + # Only image collectors are supported by the current generic handler. + for collector_id in collector_ids: + inputs = source_nodes[collector_id].get("inputs", {}) + if "audio" in inputs: + raise ValueError("collector audio input is unsupported by this image-only compiler") + if not _is_link(inputs.get("images")): + raise ValueError(f"collector {collector_id!r} is missing a linked image input") + source_id = str(inputs["images"][0]) + source_class = source_nodes[source_id]["class_type"] + if source_class in _REMOVED_SCALAR_NODES: + raise ValueError(f"collector {collector_id!r} image input must link to an image node") + + # A collector cannot be used, directly or through processing, as another + # collector's input: this first compiler handles disconnected jobs only. + for collector_id in collector_ids: + pending = [str(source_nodes[collector_id]["inputs"]["images"][0])] + seen = set() + while pending: + node_id = pending.pop() + if node_id in seen: + continue + seen.add(node_id) + if source_nodes[node_id]["class_type"] == "DistributedCollector": + raise ValueError("collector-to-collector dependence is unsupported") + for value in source_nodes[node_id].get("inputs", {}).values(): + if _is_link(value): + pending.append(str(value[0])) + + # Resolve only stock primitive nodes. A visited set makes malformed + # primitive-to-primitive cycles fail clearly instead of recursing forever. + def resolve_primitive(node_id, visiting=None): + node_id = str(node_id) + node = source_nodes[node_id] + class_type = node["class_type"] + if class_type not in _PRIMITIVE_NODES: + raise ValueError("linked value must be literal or supported PrimitiveInt/Float/String/Boolean") + if visiting is None: + visiting = set() + if node_id in visiting: + raise ValueError(f"cycle while resolving primitive value at node {node_id!r}") + visiting.add(node_id) + value = node.get("inputs", {}).get("value", _PRIMITIVE_NODES[class_type]) + if _is_link(value): + parent_id, output_slot = str(value[0]), value[1] + if output_slot != 0: + raise ValueError(f"invalid output slot {output_slot} for primitive node {parent_id!r}") + if source_nodes[parent_id]["class_type"] not in _PRIMITIVE_NODES: + raise ValueError("linked primitive value must come from a stock Primitive node") + result = resolve_primitive(parent_id, visiting) + else: + result = copy.deepcopy(value) + visiting.remove(node_id) + return result + + def resolve_literal_or_primitive(value, context): + if _is_link(value): + source_id, output_slot = str(value[0]), value[1] + if output_slot != 0 or source_nodes[source_id]["class_type"] not in _PRIMITIVE_NODES: + raise ValueError(f"{context} must be a literal or supported Primitive link") + try: + return resolve_primitive(source_id) + except ValueError as exc: + if "cycle" in str(exc): + raise + raise ValueError(f"{context} must be a literal or supported Primitive link") from None + return copy.deepcopy(value) + + def resolve_hidden(node_id, name, default): + value = source_nodes[node_id].get("inputs", {}).get(name, default) + if _is_link(value): + raise ValueError(f"{name} on distributed node {node_id!r} must be a literal") + return copy.deepcopy(value) + + resolved_distributed = {} + + def resolve_distributed(node_id, output_slot): + node_id = str(node_id) + class_type = source_nodes[node_id]["class_type"] + if output_slot != 0: + raise ValueError(f"invalid output slot {output_slot} for {class_type} {node_id!r}") + cache_key = (node_id, output_slot) + if cache_key in resolved_distributed: + return copy.deepcopy(resolved_distributed[cache_key]) + inputs = source_nodes[node_id].get("inputs", {}) + if class_type == "DistributedSeed": + seed = resolve_literal_or_primitive(inputs.get("seed", _SEED_DEFAULT), f"seed on {node_id!r}") + is_worker = resolve_hidden(node_id, "is_worker", False) + worker_id = resolve_hidden(node_id, "worker_id", "") + if is_worker: + worker_index = _worker_index(worker_id) + if worker_index is not None: + seed = seed + worker_index + 1 + result = seed + elif class_type == "DistributedValue": + default_value = resolve_literal_or_primitive( + inputs.get("default_value", ""), f"default_value on {node_id!r}" + ) + worker_values = resolve_literal_or_primitive( + inputs.get("worker_values", "{}"), f"worker_values on {node_id!r}" + ) + is_worker = resolve_hidden(node_id, "is_worker", False) + worker_id = resolve_hidden(node_id, "worker_id", "") + try: + values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values + except json.JSONDecodeError: + values = {} + if not isinstance(values, dict): + values = {} + value_type = values.get("_type", "STRING") + fallback = _coerce_safe(default_value, value_type) + result = fallback + if is_worker: + worker_index = _worker_index(worker_id) + if worker_index is not None: + raw_value = values.get(str(worker_index + 1), "") + if raw_value: + try: + result = _coerce(raw_value, value_type) + except (ValueError, IndexError): + result = fallback + elif class_type == "DistributedModelName": + text = resolve_literal_or_primitive( + inputs.get("text", ""), f"DistributedModelName text on {node_id!r}" + ) + result = _model_name_result(text) + else: + raise ValueError(f"unsupported distributed node: {class_type}") + resolved_distributed[cache_key] = copy.deepcopy(result) + return result + + # The prepared collector values define the artifact identity. The fallback + # for absent hidden values is stable and retains the collector's job scope. + collector_specs = [] + max_numeric_id = 0 + for node_id in source_nodes: + try: + max_numeric_id = max(max_numeric_id, int(node_id)) + except ValueError: + pass + next_numeric_id = max_numeric_id + 1 + for collector_id in collector_ids: + inputs = source_nodes[collector_id].get("inputs", {}) + raw_job_id = inputs.get("multi_job_id", collector_id) + if _is_link(raw_job_id): + raise ValueError(f"collector {collector_id!r} multi_job_id must be a literal") + job_id = str(collector_id if raw_job_id in (None, "") else raw_job_id) + raw_worker_id = inputs.get("worker_id", "") + if _is_link(raw_worker_id): + raise ValueError(f"collector {collector_id!r} worker_id must be a literal") + worker_id = "" if raw_worker_id is None else str(raw_worker_id) + saver_id = str(next_numeric_id) + next_numeric_id += 1 + prefix = _filename_prefix(job_id, worker_id, collector_id, "image") + collector_specs.append( + { + "collector_id": collector_id, + "job_id": job_id, + "node_id": saver_id, + "kind": "image", + "prefix": prefix, + "source_link": copy.deepcopy(inputs["images"]), + } + ) + + saver_nodes = {} + for spec in collector_specs: + saver_nodes[spec["node_id"]] = { + "class_type": "SaveImage", + "inputs": { + "images": spec["source_link"], + "filename_prefix": spec["prefix"], + }, + "_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"}, + } + + all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()} + all_nodes.update(copy.deepcopy(saver_nodes)) + needed = set() + pending = [spec["node_id"] for spec in collector_specs] + while pending: + node_id = pending.pop() + if node_id in needed: + continue + needed.add(node_id) + node = all_nodes[node_id] + if node["class_type"] == "DistributedCollector": + raise ValueError("collector-to-collector dependence is unsupported") + for input_name, value in list(node.get("inputs", {}).items()): + if not _is_link(value): + continue + upstream_id, output_slot = str(value[0]), value[1] + upstream_class = source_nodes[upstream_id]["class_type"] + if upstream_class in _REMOVED_SCALAR_NODES: + node["inputs"][input_name] = resolve_distributed(upstream_id, output_slot) + else: + pending.append(upstream_id) + + compiled = { + node_id: copy.deepcopy(node) + for node_id, node in all_nodes.items() + if node_id in needed + } + if any("distributed" in node["class_type"].lower() for node in compiled.values()): + raise ValueError("compiled prompt still contains a distributed dependency") + + manifest = [ + {key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")} + for spec in collector_specs + ] + return compiled, manifest diff --git a/api/orchestration/serverless_runtime.py b/api/orchestration/serverless_runtime.py new file mode 100644 index 0000000..eeebfff --- /dev/null +++ b/api/orchestration/serverless_runtime.py @@ -0,0 +1,275 @@ +"""Master-owned asynchronous Runpod jobs, independent of worker callbacks.""" +import asyncio +import time +import uuid +from pathlib import Path + +from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission +from .serverless_ledger import JobLedger, TERMINAL + + +async def decode_results(output, manifest): + from .serverless_media import collect_images + import torch + images = await collect_images(output, manifest) + return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()} + + +def state_directory(): + import folder_paths + return Path(folder_paths.get_user_directory()) / "distributed" / "serverless" + + +def get_secret_store(): + from ...utils.serverless_config import SecretStore + return SecretStore(state_directory() / "credentials.json") + + +def get_manager(): + import server + from ...utils.network import get_client_session + ps = server.PromptServer.instance + manager = getattr(ps, "distributed_serverless", None) + if manager is None: + manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session) + ps.distributed_serverless = manager + return manager + + +class ServerlessManager: + def __init__(self, prompt_server, secrets, ledger_path, session_factory): + self.server = prompt_server + self.secrets = secrets + self.ledger = JobLedger(ledger_path) + self.session_factory = session_factory + self.tasks = {} + self.runs = {} + self.pool_limits = {} + self.health_cache = {} + self.health_locks = {} + self.poll_seconds = 2.0 + + async def _client(self, worker): + pool_id = worker.get("pool_id", worker["id"]) + key = self.secrets.get(pool_id) + return RunpodClient(await self.session_factory(), worker["endpoint_id"], key) + + async def health(self, worker, *, fresh=False): + pool_id = worker.get("pool_id", worker["id"]) + cache_key = (pool_id, worker["endpoint_id"]) + lock = self.health_locks.setdefault(cache_key, asyncio.Lock()) + async with lock: + cached = self.health_cache.get(cache_key) + if cached and not fresh and time.monotonic() - cached[0] < 10: + return cached[1] + client = await self._client(worker) + data = await client.health() + jobs = data.get("jobs", {}) + if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict): + raise RunpodError("Runpod returned an invalid endpoint health record") + try: + pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0))) + except (ValueError, TypeError): + raise RunpodError("Runpod returned invalid queue counts") from None + result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"} + self.health_cache[cache_key] = (time.monotonic(), result) + return result + + def _semaphore(self, worker): + pool_id = worker.get("pool_id", worker["id"]) + limit = worker["concurrency"] + current = self.pool_limits.get(pool_id) + if current is None: + current = (limit, asyncio.Semaphore(limit)) + self.pool_limits[pool_id] = current + elif current[0] != limit: + if any(not task.done() and self.runs[rid]["worker"].get("pool_id", self.runs[rid]["worker"]["id"]) == pool_id for rid, task in self.tasks.items()): + raise ValueError("Wait for this serverless pool's current jobs before changing concurrency") + current = (limit, asyncio.Semaphore(limit)) + self.pool_limits[pool_id] = current + return current[1] + + def ensure_capacity(self, count): + if sum(not task.done() for task in self.tasks.values()) + count > 64: + raise ValueError("Distributed already has too many pending serverless jobs (limit 64)") + + async def start(self, execution_id, prompt_id, prepared_jobs): + self.ensure_capacity(len(prepared_jobs)) + new_runs = [] + # Journal every local slot before starting any billed submission. + for prepared in prepared_jobs: + worker = prepared["worker"] + semaphore = self._semaphore(worker) + record_id = uuid.uuid4().hex + now = time.time() + self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id, + worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"], + endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now) + run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id, + provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"]) + self.runs[record_id] = run + new_runs.append((record_id, run, semaphore)) + for record_id, run, semaphore in new_runs: + self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}") + # Bound in-memory completed tracking as well as the on-disk journal. + done = [rid for rid, task in self.tasks.items() if task.done()] + for rid in done[:-200]: + self.tasks.pop(rid, None) + self.runs.pop(rid, None) + + def is_pending(self, collector_job_id, participant_id): + return any(not task.done() and run["worker"]["id"] == participant_id and + any(m["job_id"] == collector_job_id for m in run["manifest"]) + for rid, task in self.tasks.items() for run in [self.runs[rid]]) + + def owns(self, collector_job_id, participant_id): + return any(run["worker"]["id"] == participant_id and + any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values()) + + async def _deliver(self, run, images=None, error=None): + async with self.server.distributed_jobs_lock: + for binding in run["manifest"]: + queue = self.server.distributed_pending_jobs.get(binding["job_id"]) + if queue is None: + continue # A cancelled/finished collector must never be resurrected. + if error: + queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True}) + else: + batch = images[binding["job_id"]] + for index, tensor in enumerate(batch): + queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index, + "tensor": tensor, "is_last": index == len(batch) - 1}) + + async def _cancel_provider(self, run): + if run["provider_job_id"] and run["client"]: + try: + result = await run["client"].cancel(run["provider_job_id"]) + return isinstance(result, dict) and result.get("status") == "CANCELLED" + except Exception: + return False + return False + + def _check_master(self, run): + queue = getattr(self.server, "prompt_queue", None) + if queue is None: + return + history = queue.get_history(run["prompt_id"]) + record = history.get(run["prompt_id"], {}) + if record.get("status", {}).get("status_str") == "error": + raise RunpodError("Master workflow failed; serverless work stopped") + if not record and hasattr(queue, "get_current_queue"): + running, pending = queue.get_current_queue() + if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]): + raise RunpodError("Master workflow was removed from the queue; serverless work stopped") + + async def _run(self, run, semaphore): + rid = run["record_id"] + submitted = False + provider_terminal = False + try: + async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())): + async with semaphore: + self._check_master(run) + client = await self._client(run["worker"]) + run["client"] = client + self.ledger.update(rid, status="SUBMITTING", updated_at=time.time()) + submitted = True + job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"]) + run["provider_job_id"] = job_id + # Persist immediately: all subsequent actions use this exact ID. + self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time()) + while True: + self._check_master(run) + response = await client.status(job_id) + status = response.get("status") + if status not in {"IN_QUEUE", "IN_PROGRESS", "COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"}: + raise RunpodError("Runpod returned an unknown job state") + provider_terminal = status in TERMINAL + if status == "COMPLETED": + output = response.get("output") + if not isinstance(output, dict) or output.get("error") or output.get("errors"): + raise RunpodError("Serverless workflow reported an output error; inspect the provider job") + self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time()) + images = await decode_results(output, run["manifest"]) + await self._deliver(run, images=images) + self.ledger.update(rid, status="COMPLETED", updated_at=time.time()) + return + self.ledger.update(rid, status=status, updated_at=time.time()) + if status in TERMINAL: + raise RunpodError(f"Serverless job ended with {status}") + await asyncio.sleep(self.poll_seconds) + except AmbiguousSubmission as exc: + self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time()) + await self._deliver(run, error=str(exc)) + except (asyncio.TimeoutError, asyncio.CancelledError) as exc: + confirmed = await self._cancel_provider(run) + if submitted and not run["provider_job_id"]: + status = "UNKNOWN" + error = "Submission interrupted before its job ID was received; check Runpod before retrying" + elif isinstance(exc, asyncio.TimeoutError): + status = "TIMED_OUT" if not run["provider_job_id"] or confirmed else "CANCEL_REQUESTED" + error = "Serverless job exceeded its queue/execution deadline" + else: + status = "CANCELLED" if not run["provider_job_id"] or confirmed else "CANCEL_REQUESTED" + error = "Serverless job cancelled" if confirmed or not submitted else "Serverless cancellation requested; provider confirmation is pending" + self.ledger.update(rid, status=status, error=error, updated_at=time.time()) + await self._deliver(run, error=error) + except Exception as exc: + # If polling/decoding failed, stop only our own provider job, never the pool. + confirmed = await self._cancel_provider(run) if not provider_terminal else True + error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract" + status = "CANCEL_REQUESTED" if run["provider_job_id"] and not confirmed else "FAILED" + self.ledger.update(rid, status=status, error=error, updated_at=time.time()) + await self._deliver(run, error=error) + finally: + # Completed tracking retains IDs/bindings only, not potentially large + # base64 inputs, workflow metadata, or a client holding credentials. + run.pop("workflow", None) + run.pop("images", None) + run["client"] = None + + async def wait_for_execution(self, execution_id): + tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id] + await asyncio.gather(*tasks, return_exceptions=False) + + async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None): + targets = [] + for rid, task in self.tasks.items(): + run = self.runs[rid] + if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id: + continue + if execution_id and run["execution_id"] != execution_id: + continue + if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]): + continue + if not task.done(): + task.cancel() + targets.append(task) + if targets: + await asyncio.gather(*targets, return_exceptions=True) + # A restarted master must still be able to cancel IDs it journalled. + # No workflow payloads are persisted, so never resume or resubmit them. + recovered = 0 + for rec in self.ledger.all(): + task = self.tasks.get(rec["record_id"]) + if (task is not None and not task.done()) or rec.get("status") in TERMINAL: + continue + if worker_id and rec.get("worker_id") != worker_id: + continue + if execution_id and rec.get("execution_id") != execution_id: + continue + if collector_job_id or not rec.get("job_id"): + continue + try: + client = await self._client({"id": rec["worker_id"], "endpoint_id": rec["endpoint_id"]}) + result = await client.cancel(rec["job_id"]) + status = "CANCELLED" if result.get("status") == "CANCELLED" else "CANCEL_REQUESTED" + except Exception: + status = "CANCEL_REQUESTED" + self.ledger.update(rec["record_id"], status=status, updated_at=time.time()) + recovered += 1 + return len(targets) + recovered + + def public_jobs(self, worker_id): + return [{key: rec.get(key) for key in ("job_id", "status", "worker_id", "execution_id", "error")} + for rec in reversed(self.ledger.all()) if rec.get("worker_id") == worker_id][:20] diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index ab18e7f..b86e7cb 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -74,6 +74,14 @@ def _resolve_enabled_workers(config, requested_ids=None): elif not worker.get("enabled", False): continue + if worker.get("type") == "serverless": + from ..utils.serverless_config import validate_serverless_worker + pool = validate_serverless_worker(worker) + for slot in range(pool["jobs_per_run"]): + workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}", + name=f"{pool.get('name', worker_id)} / job {slot + 1}")) + continue + raw_port = worker.get("port", worker.get("listen_port", 8188)) try: port = int(raw_port or 8188) @@ -154,6 +162,17 @@ async def _prepare_worker_payload( """Prepare one worker prompt payload with bounded concurrency and media-sync timeout.""" async with worker_prep_semaphore: worker_prompt = prompt_index.copy_prompt() + if worker.get("type") == "serverless": + from .orchestration.serverless_prompt import compile_serverless_prompt + from .orchestration.serverless_media import prepare_images + worker_prompt = apply_participant_overrides( + prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids, + job_id_map, "", delegate_master, prompt_index, + ) + workflow, manifest = compile_serverless_prompt(worker_prompt) + images = await prepare_images(workflow) + return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images} + worker_master_url = build_master_callback_url( worker, config=config, @@ -330,9 +349,6 @@ async def orchestrate_distributed_execution( queue_result.get("node_errors", {}), ) - for job_id in job_id_map.values(): - await _ensure_distributed_queue(job_id) - master_prompt = prompt_index.copy_prompt() master_prompt = apply_participant_overrides( master_prompt, @@ -387,7 +403,19 @@ async def orchestrate_distributed_execution( ] ) - if worker_payloads: + serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"] + regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"] + manager = None + if serverless_jobs: + from .orchestration.serverless_runtime import get_manager + manager = get_manager() + manager.ensure_capacity(len(serverless_jobs)) + + # Do not leave collector queues behind when offline compilation/media validation fails. + for job_id in job_id_map.values(): + await _ensure_distributed_queue(job_id) + + if regular_payloads: await asyncio.gather( *[ dispatch_worker_prompt( @@ -398,17 +426,38 @@ async def orchestrate_distributed_execution( use_websocket=use_websocket, trace_execution_id=execution_trace_id, ) - for worker, wprompt in worker_payloads + for worker, wprompt in regular_payloads ] ) - queue_result = await queue_prompt_payload( - master_prompt, - workflow_meta, - client_id, - include_queue_metadata=True, - ) + try: + queue_result = await queue_prompt_payload( + master_prompt, + workflow_meta, + client_id, + include_queue_metadata=True, + ) + except Exception: + if serverless_jobs: + async with prompt_server.distributed_jobs_lock: + for job_id in job_id_map.values(): + prompt_server.distributed_pending_jobs.pop(job_id, None) + raise prompt_id = queue_result["prompt_id"] + if manager is not None: + try: + await manager.start(execution_trace_id, prompt_id, serverless_jobs) + except Exception: + # Wake queued collectors with an explicit error if journal/capacity setup failed. + await manager.cancel(execution_id=execution_trace_id) + async with prompt_server.distributed_jobs_lock: + for prepared in serverless_jobs: + for binding in prepared["manifest"]: + q = prompt_server.distributed_pending_jobs.get(binding["job_id"]) + if q is not None: + q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True, + "serverless_error": "Serverless job setup failed; no automatic retry was attempted"}) + raise prompt_number = queue_result["number"] node_errors = queue_result.get("node_errors", {}) trace_debug( diff --git a/api/serverless_routes.py b/api/serverless_routes.py new file mode 100644 index 0000000..eb4ba68 --- /dev/null +++ b/api/serverless_routes.py @@ -0,0 +1,58 @@ +"""Master-only Runpod pool monitoring. Health checks never submit inference.""" +from aiohttp import web +import server + +from ..utils.config import load_config +from .orchestration.serverless_runtime import get_manager +from .orchestration.runpod_client import RunpodError + + +def _pool(worker_id): + return next((w for w in load_config().get("workers", []) + if w.get("id") == worker_id and w.get("type") == "serverless"), None) + + +async def _status(worker_id, fresh=False): + worker = _pool(worker_id) + if worker is None: + return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404) + try: + manager = get_manager() + result = await manager.health(worker, fresh=fresh) + return web.json_response(dict(result, jobs=manager.public_jobs(worker_id))) + except Exception as exc: + message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key" + return web.json_response({"ok": False, "queueRemaining": 0, "message": message}) + + +@server.PromptServer.instance.routes.post("/distributed/serverless/test") +async def test_serverless_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + except Exception: + return web.json_response({"ok": False, "message": "Invalid request"}, status=400) + return await _status(worker_id, fresh=True) + + +@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}") +async def serverless_status_endpoint(request): + return await _status(request.match_info["worker_id"]) + + +@server.PromptServer.instance.routes.post("/distributed/serverless/cancel") +async def cancel_serverless_endpoint(request): + try: + data = await request.json() + worker_id = data.get("worker_id") + except Exception: + return web.json_response({"ok": False, "message": "Invalid request"}, status=400) + if _pool(worker_id) is None: + return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404) + try: + manager = get_manager() + count = await manager.cancel(worker_id=worker_id) + return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id), + "message": "Cancellation requested for this master's jobs only"}) + except Exception: + return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502) diff --git a/api/worker_routes.py b/api/worker_routes.py index dfdf767..709fb77 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -460,6 +460,9 @@ async def launch_worker_endpoint(request): if not worker: return await handle_api_error(request, f"Worker {worker_id} not found", 404) + if worker.get("type") == "serverless": + return await handle_api_error(request, "Serverless pools are managed by Runpod, not local processes", 400) + # Ensure consistent string ID worker_id_str = worker_id @@ -554,6 +557,9 @@ async def get_local_worker_status_endpoint(request): worker_statuses = {} for worker in config.get("workers", []): + # Serverless pools have no host, port, or local process. + if worker.get("type") == "serverless": + continue # Only check local workers host = normalize_host(worker.get("host")) or "" if not host or host in ["localhost", "127.0.0.1"]: diff --git a/docs/serverless-workers.md b/docs/serverless-workers.md new file mode 100644 index 0000000..dff76d8 --- /dev/null +++ b/docs/serverless-workers.md @@ -0,0 +1,119 @@ +# Runpod serverless pools (preview) + +A **Serverless (Runpod)** worker is a queue-based endpoint/pool, not a permanent +ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one +normal ComfyUI API workflow per logical slot and collects the returned images. +The endpoint workers **do not need ComfyUI-Distributed installed**. + +## Setup + +1. Add a worker in the Distributed panel and select **Serverless (Runpod)**. +2. Enter its name, Runpod **endpoint ID** (not a URL), and API key. +3. Start with **Jobs per run = 1**, **Concurrent jobs = 1**, and **Timeout = 1800**. +4. Save, then choose **Test Saved Endpoint**. This checks authentication and + endpoint health only; it does not start GPU inference or prove model/node compatibility. +5. Enable the pool when you want queued Distributed workflows to use it. +6. Use an image **Distributed Collector** and queue through Distributed as usual. + Enable **Delegate Only** if the master should collect without generating locally. + +Keys are write-only. Leave the key field blank to retain a saved key, or select +**Clear saved API key** to remove it (disable the pool first). A saved key is +required before enabling a pool. Use HTTPS or a trusted local connection when +entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted +users: anyone able to queue arbitrary workflows can run code and spend endpoint credits. + +## Endpoint contract + +The endpoint must accept this generic input shape; named application modes are +not used: + +```json +{ + "input": { + "workflow": { + "1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}}, + "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}} + }, + "images": [] + } +} +``` + +For image inputs, `images` contains `{ "name": "input.png", "image": "" }` +entries and the handler must rewrite `LoadImage` references to the staged files. +The adapter transfers local input images; it does not copy models or install nodes. +All referenced models and custom nodes must already exist in the endpoint image. + +On completion, Runpod's `output.images` must contain artifact records such as: + +```json +{ + "filename": "example_00001_.png", + "subfolder": "", + "artifact_type": "output", + "artifact_path": "example_00001_.png", + "type": "base64", + "data": "" +} +``` + +Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"` +with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without +Runpod credentials; private-network destinations and unsafe redirects are refused. +The handler must preserve injected `SaveImage.filename_prefix` values so results +can be matched to the right collector and slot. Duplicate basenames in different +subfolders remain different artifact identities. + +A Runpod job marked `COMPLETED` is still rejected if `output.error` or +`output.errors` is nonempty, required collector outputs are missing, or image +artifacts are invalid. Failures are surfaced rather than silently returning an +incomplete serverless batch. + +## Scheduling and cancellation + +- **Jobs per run** adds logical collector participants. Each receives its own + Distributed Seed/Value overrides. It is not a physical GPU count. +- **Concurrent jobs** bounds this master's outstanding submissions per pool. + Other applications using the same endpoint are unaffected. +- **Load balance** selects one participant, so it does not fan out every pool slot. +- **Timeout** includes waiting for a local slot, Runpod queue/cold-start time, + execution, and output retrieval. Provider policy limits are submitted as well. +- A zero-worker endpoint can be healthy: Runpod scales from zero on submission. +- Interrupting a collector or using the panel cancellation action requests + cancellation only for jobs submitted by this master. It never purges the endpoint. +- A network failure during submission is **not retried automatically**: the job + may already be billable. Check the Runpod job list before queuing again. + +## Persistence and recovery + +Credentials live outside `gpu_config.json` and Git in the master ComfyUI user +folder under `distributed/serverless/credentials.json`, with restrictive POSIX +permissions. They are omitted from public configuration, workflow exports, and +browser configuration state. The first version requires a POSIX (Linux/macOS) +master for serverless secret storage; existing non-serverless workers remain +available on Windows. + +`distributed/serverless/jobs.json` records provider job IDs and sanitized states, +not image data, workflow payloads, or API keys. A restarted master cannot resume +an interrupted collector. It does **not** automatically resubmit jobs. Retained +IDs are available through `/distributed/serverless/status/`; the pool's +cancel action can request cancellation of known pre-restart jobs. If no job ID +was received, check Runpod manually. Cancellation is best-effort and may race +with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation. + +## First-version limits + +- Image collectors only. Audio/video transfer and audio collectors are rejected. +- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed + node types are rejected before submission. +- Downstream post-collector processing remains local; the remote graph contains + only the upstream compute required for image saves. +- Inputs must be supported images under ComfyUI's input/output/temp roots. +- Image input bytes are bounded to 7 MiB total; the complete request is bounded + to 10 MiB. Large-media/object-storage upload support is not included. +- No remote GPU discovery, endpoint provisioning, model synchronization, automatic + scaling changes, or inbound master callback/tunnel is needed. + +This feature is verified with offline provider/transport tests. A successful +health test is not a paid end-to-end generation test; test your deployed handler +and model inventory with one slot before increasing concurrency. diff --git a/nodes/collector.py b/nodes/collector.py index d577b0f..90a5da0 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -151,6 +151,15 @@ class DistributedCollectorNode: ) return result + @staticmethod + def _get_serverless_manager(): + """Return the serverless manager only for collector jobs with serverless slots.""" + manager = getattr(prompt_server, "distributed_serverless", None) + if manager is not None: + return manager + from ..api.orchestration.serverless_runtime import get_manager + return get_manager() + async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id): """Send an image batch, optionally with audio, or an audio-only completion.""" encoded_audio = encode_audio_payload(audio) @@ -339,6 +348,17 @@ class DistributedCollectorNode: seen_enabled.add(worker_id_str) enabled_workers.append(worker_id_str) expected_workers = set(enabled_workers) + serverless_workers = { + worker_id for worker_id in expected_workers if worker_id.startswith("serverless:") + } + serverless_manager = None + + def get_serverless_manager(): + nonlocal serverless_manager + if serverless_workers and serverless_manager is None: + serverless_manager = self._get_serverless_manager() + return serverless_manager + num_workers = len(expected_workers) if num_workers == 0: return (images, audio if audio is not None else self.EMPTY_AUDIO) @@ -405,152 +425,177 @@ class DistributedCollectorNode: workers_done.add(done_worker_id) p.update(1) # +1 per completed expected worker + def consume_result(result): + nonlocal collected_count + if "serverless_error" in result: + error = result.get("serverless_error") + message = str(error).strip() if error is not None else "unspecified serverless failure" + participant = result.get("worker_id", "unknown") + raise RuntimeError(f"Serverless participant {participant} failed: {message}") + + result_worker_id = str(result["worker_id"]) + is_last = result.get("is_last", False) + collected_count += self._store_worker_result(worker_images, result) + debug_log( + f"Master - Got canonical result from worker {result_worker_id}, " + f"image {result.get('image_index', 0)}, is_last={is_last}" + ) + result_audio = result.get("audio") + if result_audio is not None: + worker_audio[result_worker_id] = result_audio + debug_log(f"Master - Got audio from worker {result_worker_id}") + if is_last: + mark_worker_done(result_worker_id) + try: while len(workers_done) < num_workers: - # Check for user interruption to abort collection promptly + # Check for user interruption to abort collection promptly. comfy.model_management.throw_exception_if_processing_interrupted() try: - # Get the queue again each time to ensure we have the right reference async with prompt_server.distributed_jobs_lock: q = prompt_server.distributed_pending_jobs[multi_job_id] - current_size = q.qsize() - result = await asyncio.wait_for(q.get(), timeout=slice_timeout) - worker_id = result['worker_id'] - is_last = result.get('is_last', False) - count = self._store_worker_result(worker_images, result) - collected_count += count - debug_log( - f"Master - Got canonical result from worker {worker_id}, " - f"image {result.get('image_index', 0)}, is_last={is_last}" - ) + consume_result(result) - # Collect audio data if present - result_audio = result.get('audio') - if result_audio is not None: - worker_audio[worker_id] = result_audio - debug_log(f"Master - Got audio from worker {worker_id}") - - # Record activity and refresh timeout baseline + # Real results refresh the regular-worker timeout baseline. last_activity = time.time() base_timeout = float(get_worker_timeout_seconds()) - if is_last: - mark_worker_done(worker_id) - except asyncio.TimeoutError: - # If we still have time, continue polling; otherwise handle timeout if (time.time() - last_activity) < base_timeout: comfy.model_management.throw_exception_if_processing_interrupted() continue - # Re-check for user interruption after timeout expiry comfy.model_management.throw_exception_if_processing_interrupted() - missing_workers = set(str(w) for w in enabled_workers) - workers_done + missing_workers = expected_workers - workers_done + + # A queue delivery can race the timed wait. Process it before + # declaring an owned provider slot stopped without a result. + async with prompt_server.distributed_jobs_lock: + current_q = prompt_server.distributed_pending_jobs.get(multi_job_id) + queue_has_results = current_q is not None and not current_q.empty() + if queue_has_results: + continue + + missing_serverless = missing_workers & serverless_workers + if missing_serverless: + manager = get_serverless_manager() + if manager is None: + raise RuntimeError("Serverless manager is unavailable for pending collector slots") + pending_serverless = [] + for participant_id in sorted(missing_serverless): + if manager.is_pending(multi_job_id, participant_id): + pending_serverless.append(participant_id) + elif manager.owns(multi_job_id, participant_id): + raise RuntimeError( + f"Serverless participant {participant_id} stopped without delivering a result" + ) + # Provider deadlines are managed by the serverless runtime; + # do not apply the regular worker heartbeat timeout to it. + if pending_serverless: + continue + + missing_regular = missing_workers - serverless_workers elapsed = time.time() - last_activity - for missing_worker_id in sorted(missing_workers): + for missing_worker_id in sorted(missing_regular): log( "Master - Heartbeat timeout: " f"worker={missing_worker_id}, elapsed={elapsed:.1f}s" ) - log( - f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} " - f"(elapsed={elapsed:.1f}s)" - ) + if missing_regular: + log( + f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} " + f"(elapsed={elapsed:.1f}s)" + ) - # Probe missing workers' /prompt endpoints to check if they are actively processing + # Probe only regular worker /prompt endpoints. any_busy = False - try: - cfg = load_config() - cfg_workers = cfg.get('workers', []) - for wid in list(missing_workers): - wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None) - if not wrec: - debug_log(f"Collector probe: worker {wid} not found in config") - continue - worker_url = build_worker_url(wrec) - try: - payload = await probe_worker(worker_url, timeout=2.0) - queue_remaining = None - if payload is not None: - queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0)) - debug_log( - "Collector probe: worker " - f"{wid} online={payload is not None} queue_remaining={queue_remaining}" - ) - if payload is not None and queue_remaining and queue_remaining > 0: - any_busy = True - log( - f"Master - Probe grace: worker {wid} appears busy " - f"(queue_remaining={queue_remaining}). Continuing to wait." + if missing_regular: + try: + cfg = load_config() + cfg_workers = cfg.get('workers', []) + for wid in sorted(missing_regular): + wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None) + if not wrec: + debug_log(f"Collector probe: worker {wid} not found in config") + continue + worker_url = build_worker_url(wrec) + try: + payload = await probe_worker(worker_url, timeout=2.0) + queue_remaining = None + if payload is not None: + queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0)) + debug_log( + "Collector probe: worker " + f"{wid} online={payload is not None} queue_remaining={queue_remaining}" ) - break - except Exception as e: - debug_log(f"Collector probe failed for worker {wid}: {e}") - except Exception as e: - debug_log(f"Collector probe setup error: {e}") + if payload is not None and queue_remaining and queue_remaining > 0: + any_busy = True + log( + f"Master - Probe grace: worker {wid} appears busy " + f"(queue_remaining={queue_remaining}). Continuing to wait." + ) + break + except Exception as e: + debug_log(f"Collector probe failed for worker {wid}: {e}") + except Exception as e: + debug_log(f"Collector probe setup error: {e}") if any_busy: - # Refresh last_activity and continue waiting last_activity = time.time() - # Refresh base timeout in case the user changed it in UI base_timeout = float(get_worker_timeout_seconds()) continue - - # Check queue size again with lock + async with prompt_server.distributed_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_jobs: - final_q = prompt_server.distributed_pending_jobs[multi_job_id] - final_size = final_q.qsize() - - # Try to drain any remaining items - remaining_items = [] + final_q = prompt_server.distributed_pending_jobs.get(multi_job_id) + remaining_items = [] + if final_q is not None: while not final_q.empty(): try: - item = final_q.get_nowait() - remaining_items.append(item) + remaining_items.append(final_q.get_nowait()) except asyncio.QueueEmpty: break - - if remaining_items: - # Process them - for item in remaining_items: - worker_id = item['worker_id'] - is_last = item.get('is_last', False) - - collected_count += self._store_worker_result(worker_images, item) - - if is_last: - mark_worker_done(worker_id) - else: - log(f"Master - Queue {multi_job_id} no longer exists!") + for item in remaining_items: + consume_result(item) break - except comfy.model_management.InterruptProcessingException: - # Cleanup queue on interruption and re-raise to abort prompt cleanly + + total_collected = sum(len(imgs) for imgs in worker_images.values()) + + # Remove the queue before combining so a late provider callback cannot + # resurrect results for a collector that has completed. async with prompt_server.distributed_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_jobs: - del prompt_server.distributed_pending_jobs[multi_job_id] + prompt_server.distributed_pending_jobs.pop(multi_job_id, None) + + combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers) + serverless_results = set(worker_images) & serverless_workers + try: + combined = self._reorder_and_combine_tensors( + worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images + ) + if combined is None: + debug_log(f"Master - Job {multi_job_id} complete with audio only") + else: + debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " + f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") + return (combined, combined_audio) + except Exception as e: + if serverless_results: + raise RuntimeError( + "Serverless result batch shape mismatch while combining participants " + f"{sorted(serverless_results)}: {e}" + ) from e + log(f"Master - Error combining images: {e}") + # Preserve collected audio for regular workers when image assembly fails. + return (images, combined_audio) + except BaseException: + # Remove the queue first, preventing cancelled work from delivering into + # a collector that is already unwinding. + async with prompt_server.distributed_jobs_lock: + prompt_server.distributed_pending_jobs.pop(multi_job_id, None) + if serverless_workers: + try: + manager = get_serverless_manager() + if manager is not None: + await manager.cancel(collector_job_id=multi_job_id) + except Exception as cancel_error: + debug_log(f"Collector - Serverless cancellation failed: {cancel_error}") raise - - total_collected = sum(len(imgs) for imgs in worker_images.values()) - - # Clean up job queue - async with prompt_server.distributed_jobs_lock: - if multi_job_id in prompt_server.distributed_pending_jobs: - del prompt_server.distributed_pending_jobs[multi_job_id] - - combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers) - try: - combined = self._reorder_and_combine_tensors( - worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images - ) - if combined is None: - debug_log(f"Master - Job {multi_job_id} complete with audio only") - else: - debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total " - f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})") - - return (combined, combined_audio) - except Exception as e: - log(f"Master - Error combining images: {e}") - # Preserve collected audio even when image assembly fails. - return (images, combined_audio) diff --git a/tests/api/test_config_routes.py b/tests/api/test_config_routes.py index a435547..fc33018 100644 --- a/tests/api/test_config_routes.py +++ b/tests/api/test_config_routes.py @@ -38,17 +38,10 @@ def _load_config_routes_module(): sys.modules[f"{package_name}.api"] = api_pkg utils_pkg = types.ModuleType(f"{package_name}.utils") - utils_pkg.__path__ = [] + utils_pkg.__path__ = [str(module_path.parents[1] / "utils")] sys.modules[f"{package_name}.utils"] = utils_pkg - created_aiohttp_stub = False - if "aiohttp" not in sys.modules: - created_aiohttp_stub = True - aiohttp_module = types.ModuleType("aiohttp") - aiohttp_module.web = types.SimpleNamespace( - json_response=lambda payload, status=200: _FakeResponse(payload, status=status) - ) - sys.modules["aiohttp"] = aiohttp_module + import aiohttp # Load the real dependency; isolate the fake response at the tested module. class _Routes: def get(self, _path): @@ -96,8 +89,14 @@ def _load_config_routes_module(): assert spec is not None and spec.loader is not None spec.loader.exec_module(module) - if created_aiohttp_stub: - sys.modules.pop("aiohttp", None) + class _EmptySecretStore: + def has(self, _worker_id): + return False + + module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status=status) + ) + module._serverless_secret_store = lambda: _EmptySecretStore() return module diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_distributed_queue.py index ac6a333..c9b85eb 100644 --- a/tests/api/test_distributed_queue.py +++ b/tests/api/test_distributed_queue.py @@ -205,6 +205,10 @@ def _load_job_routes_module(): module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None spec.loader.exec_module(module) + # Keep response assertions independent of real aiohttp import ordering. + module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status=status) + ) if created_aiohttp_stub: sys.modules.pop("aiohttp", None) diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py index 56aabb7..f6d0b2e 100644 --- a/tests/api/test_worker_routes.py +++ b/tests/api/test_worker_routes.py @@ -257,6 +257,8 @@ def _load_worker_routes_module(): module = importlib.util.module_from_spec(spec) assert spec is not None and spec.loader is not None spec.loader.exec_module(module) + # Keep this module's response double independent of aiohttp import order. + module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status)) if created_aiohttp_stub: sys.modules.pop("aiohttp", None) diff --git a/tests/test_runpod_client.py b/tests/test_runpod_client.py new file mode 100644 index 0000000..5b6bc10 --- /dev/null +++ b/tests/test_runpod_client.py @@ -0,0 +1,98 @@ +import importlib.util +import sys +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] + +def load_module(): + spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + +client = load_module() + +class Response: + def __init__(self, status=200, payload=None): + self.status, self.payload = status, payload + async def __aenter__(self): + return self + async def __aexit__(self, *args): + return False + async def json(self): + return self.payload + +class Session: + def __init__(self, responses): + self.responses, self.calls = list(responses), [] + def request(self, method, url, **kwargs): + self.calls.append((method, url, kwargs)) + response = self.responses.pop(0) + if isinstance(response, Exception): + raise response + return response + +class ClientTests(unittest.IsolatedAsyncioTestCase): + async def test_submit_exact_generic_contract_and_once(self): + session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})]) + api = client.RunpodClient(session, "endpoint123", "secret-value") + result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800) + self.assertEqual(result, "job-1") + method, url, kwargs = session.calls[0] + self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run")) + self.assertEqual(set(kwargs["json"]["input"]), {"workflow", "images"}) + self.assertEqual(kwargs["json"]["policy"]["ttl"], 1800000) + self.assertFalse(kwargs["allow_redirects"]) + self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value") + + async def test_submit_timeout_is_ambiguous_and_never_retried(self): + session = Session([TimeoutError("secret-value in unsafe error")]) + api = client.RunpodClient(session, "endpoint123", "secret-value") + with self.assertRaises(client.AmbiguousSubmission) as ctx: + await api.submit({}, [], 30) + self.assertNotIn("secret-value", str(ctx.exception)) + self.assertEqual(len(session.calls), 1) + + async def test_http_submit_failure_not_retried_or_leaked(self): + for status in [401, 429, 500, 302]: + session = Session([Response(status, {"error": "secret-value"})]) + api = client.RunpodClient(session, "endpoint123", "secret-value") + with self.assertRaises(client.RunpodError) as ctx: + await api.submit({}, [], 30) + self.assertEqual(len(session.calls), 1) + self.assertNotIn("secret-value", str(ctx.exception)) + + async def test_health_zero_workers_is_valid(self): + session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})]) + api = client.RunpodClient(session, "endpoint123", "key") + self.assertEqual((await api.health())["workers"]["idle"], 0) + + async def test_invalid_provider_id_no_network(self): + session = Session([]) + with self.assertRaises(ValueError): + client.RunpodClient(session, "https://attacker.invalid", "key") + api = client.RunpodClient(session, "endpoint123", "key") + with self.assertRaises(ValueError): + await api.status("../other") + self.assertFalse(session.calls) + + async def test_status_and_cancel_only_target_exact_job(self): + session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})]) + api = client.RunpodClient(session, "endpoint123", "key") + await api.status("job-1") + await api.cancel("job-1") + self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")]) + + async def test_missing_id_is_ambiguous(self): + api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key") + with self.assertRaises(client.AmbiguousSubmission): + await api.submit({}, [], 30) + + async def test_large_payload_rejected_before_submit(self): + session = Session([]) + api = client.RunpodClient(session, "endpoint123", "key") + with self.assertRaises(ValueError): + await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30) + self.assertFalse(session.calls) diff --git a/tests/test_serverless_collector.py b/tests/test_serverless_collector.py new file mode 100644 index 0000000..b565823 --- /dev/null +++ b/tests/test_serverless_collector.py @@ -0,0 +1,202 @@ +import asyncio + +import pytest +import torch + +from test_collector_list_inputs import _load_collector_module + + +class _FakeManager: + def __init__(self, *, pending=False, owned=True): + self.pending = pending + self.owned = owned + self.pending_calls = [] + self.owned_calls = [] + self.cancel_calls = [] + + def is_pending(self, collector_job_id, participant_id): + self.pending_calls.append((collector_job_id, participant_id)) + return self.pending + + def owns(self, collector_job_id, participant_id): + self.owned_calls.append((collector_job_id, participant_id)) + return self.owned + + async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None): + self.cancel_calls.append({ + "worker_id": worker_id, + "execution_id": execution_id, + "collector_job_id": collector_job_id, + }) + return 1 + + +def _set_up_collector(module, job_id, queue, manager=None): + module.prompt_server.distributed_jobs_lock = asyncio.Lock() + module.prompt_server.distributed_pending_jobs = {job_id: queue} + if manager is not None: + module.prompt_server.distributed_serverless = manager + return module.DistributedCollectorNode() + + +def test_serverless_pending_waits_past_normal_worker_timeout(): + module = _load_collector_module() + job_id = "collector-pending" + participant = "serverless:pool:1" + queue = asyncio.Queue() + manager = _FakeManager(pending=True) + collector = _set_up_collector(module, job_id, queue, manager) + module.get_worker_timeout_seconds = lambda: 0.01 + module.HEARTBEAT_INTERVAL = 0.1 + probe_calls = [] + module.probe_worker = lambda *args, **kwargs: probe_calls.append(args) + + async def deliver_after_provider_work(): + await asyncio.sleep(0.14) + queue.put_nowait({ + "worker_id": participant, + "image_index": 0, + "tensor": torch.ones(1, 2, 2, 2), + "is_last": True, + }) + + async def run_collection(): + producer = asyncio.create_task(deliver_after_provider_work()) + result = await collector.execute( + images=None, + audio=None, + multi_job_id=job_id, + enabled_worker_ids=f'["{participant}"]', + delegate_only=True, + ) + await producer + return result + + images, _audio = asyncio.run(run_collection()) + + assert tuple(images.shape) == (1, 2, 2, 2) + assert manager.pending_calls + assert probe_calls == [] + assert job_id not in module.prompt_server.distributed_pending_jobs + + +def test_serverless_queue_error_raises_and_cancels_only_collector_binding(): + module = _load_collector_module() + job_id = "collector-error" + queue = asyncio.Queue() + queue.put_nowait({ + "worker_id": "serverless:pool:1", + "serverless_error": "Provider execution failed safely", + "is_last": True, + }) + manager = _FakeManager() + collector = _set_up_collector(module, job_id, queue, manager) + + with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"): + asyncio.run(collector.execute( + images=None, + audio=None, + multi_job_id=job_id, + enabled_worker_ids='["serverless:pool:1"]', + delegate_only=True, + )) + + assert manager.cancel_calls == [{ + "worker_id": None, + "execution_id": None, + "collector_job_id": job_id, + }] + assert job_id not in module.prompt_server.distributed_pending_jobs + + +def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success(): + module = _load_collector_module() + job_id = "collector-stopped" + queue = asyncio.Queue() + manager = _FakeManager(pending=False, owned=True) + collector = _set_up_collector(module, job_id, queue, manager) + module.get_worker_timeout_seconds = lambda: 0.01 + module.HEARTBEAT_INTERVAL = 0.1 + module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs") + + with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"): + asyncio.run(collector.execute( + images=None, + audio=None, + multi_job_id=job_id, + enabled_worker_ids='["serverless:pool:1"]', + delegate_only=True, + )) + + assert manager.owned_calls + assert manager.cancel_calls[0]["collector_job_id"] == job_id + assert job_id not in module.prompt_server.distributed_pending_jobs + + +def test_interruption_cancels_serverless_job_and_removes_collector_queue(): + module = _load_collector_module() + job_id = "collector-interrupted" + queue = asyncio.Queue() + manager = _FakeManager(pending=True) + collector = _set_up_collector(module, job_id, queue, manager) + interrupted = module.comfy.model_management.InterruptProcessingException("stop") + module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted) + + with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"): + asyncio.run(collector.execute( + images=None, + audio=None, + multi_job_id=job_id, + enabled_worker_ids='["serverless:pool:1"]', + delegate_only=True, + )) + + assert manager.cancel_calls[0]["collector_job_id"] == job_id + assert job_id not in module.prompt_server.distributed_pending_jobs + + +def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback(): + module = _load_collector_module() + job_id = "collector-shape-mismatch" + queue = asyncio.Queue() + queue.put_nowait({ + "worker_id": "serverless:pool:1", + "image_index": 0, + "tensor": torch.ones(1, 3, 2, 2), + "is_last": True, + }) + manager = _FakeManager() + collector = _set_up_collector(module, job_id, queue, manager) + local_images = torch.zeros(1, 2, 2, 2) + + with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"): + asyncio.run(collector.execute( + images=local_images, + audio=None, + multi_job_id=job_id, + enabled_worker_ids='["serverless:pool:1"]', + )) + + assert manager.cancel_calls[0]["collector_job_id"] == job_id + assert job_id not in module.prompt_server.distributed_pending_jobs + + +def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager(): + module = _load_collector_module() + job_id = "collector-regular-timeout" + queue = asyncio.Queue() + collector = _set_up_collector(module, job_id, queue) + module.get_worker_timeout_seconds = lambda: 0.01 + module.HEARTBEAT_INTERVAL = 0.1 + local_images = torch.zeros(1, 2, 2, 2) + collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager") + + images, _audio = asyncio.run(collector.execute( + images=local_images, + audio=None, + multi_job_id=job_id, + enabled_worker_ids='["worker-a"]', + )) + + assert torch.equal(images, local_images) + assert job_id not in module.prompt_server.distributed_pending_jobs diff --git a/tests/test_serverless_config.py b/tests/test_serverless_config.py new file mode 100644 index 0000000..4048f94 --- /dev/null +++ b/tests/test_serverless_config.py @@ -0,0 +1,364 @@ +import importlib.util +import json +import os +import stat +import sys +import tempfile +import threading +import unittest +from pathlib import Path +from unittest import mock + + +def _load_serverless_config_module(*, without_fcntl=False): + module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py" + module_name = ( + "dist_serverless_config_no_fcntl_test_module" + if without_fcntl + else "dist_serverless_config_test_module" + ) + sys.modules.pop(module_name, None) + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + if without_fcntl: + with mock.patch.dict(sys.modules, {"fcntl": None}): + spec.loader.exec_module(module) + else: + spec.loader.exec_module(module) + return module + + +serverless_config = _load_serverless_config_module() + + +class ValidateServerlessWorkerTests(unittest.TestCase): + def test_returns_copy_with_runpod_defaults(self): + worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"} + + result = serverless_config.validate_serverless_worker(worker) + + self.assertEqual( + result, + { + "id": "worker-1", + "endpoint_id": "ep_123-abc", + "name": "GPU", + "type": "serverless", + "provider": "runpod", + "jobs_per_run": 1, + "concurrency": 1, + "timeout_seconds": 1800, + }, + ) + self.assertIsNot(result, worker) + self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}) + + def test_accepts_inclusive_numeric_bounds(self): + for jobs, concurrency, timeout in ((1, 1, 30), (16, 16, 86400)): + with self.subTest(jobs=jobs, concurrency=concurrency, timeout=timeout): + result = serverless_config.validate_serverless_worker( + { + "id": "worker-1", + "endpoint_id": "ep-123", + "jobs_per_run": jobs, + "concurrency": concurrency, + "timeout_seconds": timeout, + } + ) + self.assertEqual(result["jobs_per_run"], jobs) + self.assertEqual(result["concurrency"], concurrency) + self.assertEqual(result["timeout_seconds"], timeout) + + def test_rejects_missing_or_empty_id_and_endpoint(self): + base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"} + for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")): + with self.subTest(field=field, value=value): + worker = dict(base) + worker[field] = value + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker(worker) + + worker = dict(base) + del worker["id"] + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker(worker) + + def test_endpoint_id_must_be_ascii_opaque_identifier(self): + for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"): + with self.subTest(endpoint_id=endpoint_id): + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker( + {"id": "worker-1", "endpoint_id": endpoint_id} + ) + + def test_endpoint_id_matches_runpod_client_length_limit(self): + valid_id = "e" * 200 + result = serverless_config.validate_serverless_worker( + {"id": "worker-1", "endpoint_id": valid_id} + ) + self.assertEqual(result["endpoint_id"], valid_id) + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker( + {"id": "worker-1", "endpoint_id": "e" * 201} + ) + + def test_enabled_must_be_boolean_when_supplied(self): + base = {"id": "worker-1", "endpoint_id": "ep-123"} + for enabled in (True, False): + with self.subTest(enabled=enabled): + self.assertIs( + serverless_config.validate_serverless_worker( + {**base, "enabled": enabled} + )["enabled"], + enabled, + ) + for enabled in ("false", 0, 1, None): + with self.subTest(enabled=enabled): + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker( + {**base, "enabled": enabled} + ) + + def test_rejects_unknown_provider_and_non_serverless_type(self): + base = {"id": "worker-1", "endpoint_id": "ep-123"} + for extra in ({"provider": "other"}, {"type": "local"}): + with self.subTest(extra=extra): + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker({**base, **extra}) + + def test_numeric_fields_reject_bool_float_and_out_of_range(self): + base = {"id": "worker-1", "endpoint_id": "ep-123"} + invalid_values = { + "jobs_per_run": (True, 1.0, 0, 17), + "concurrency": (False, 2.5, 0, 17), + "timeout_seconds": (True, 30.0, 29, 86401), + } + for field, values in invalid_values.items(): + for value in values: + with self.subTest(field=field, value=value): + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker({**base, field: value}) + + def test_requires_worker_object(self): + for worker in (None, [], "worker"): + with self.subTest(worker=worker): + with self.assertRaises(ValueError): + serverless_config.validate_serverless_worker(worker) + + +class SecretStoreTests(unittest.TestCase): + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.root = Path(self.tempdir.name) + self.path = self.root / "private" / "runpod-secrets.json" + + def tearDown(self): + self.tempdir.cleanup() + + def test_set_get_has_delete_and_persists_as_json_mapping(self): + store = serverless_config.SecretStore(self.path) + store.set("worker-1", "secret-value") + + self.assertEqual(store.get("worker-1"), "secret-value") + self.assertTrue(store.has("worker-1")) + self.assertIsNone(store.get("missing")) + self.assertFalse(store.has("missing")) + self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"}) + self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value") + + store.delete("worker-1") + self.assertIsNone(store.get("worker-1")) + self.assertFalse(store.has("worker-1")) + + def test_rejects_api_keys_with_newlines_or_over_4096_characters(self): + store = serverless_config.SecretStore(self.path) + for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097): + with self.subTest(key_length=len(key)): + with self.assertRaises(serverless_config.SecretStoreError): + store.set("worker-1", key) + store.set("worker-1", "k" * 4096) + self.assertEqual(store.get("worker-1"), "k" * 4096) + + def test_new_secret_file_and_directory_have_restrictive_modes(self): + store = serverless_config.SecretStore(self.path) + store.set("worker-1", "secret-value") + + self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600) + self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700) + + def test_missing_store_reads_as_empty_without_creating_files(self): + store = serverless_config.SecretStore(self.path) + + self.assertIsNone(store.get("worker-1")) + self.assertFalse(store.has("worker-1")) + self.assertFalse(self.path.parent.exists()) + + def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self): + self.path.parent.mkdir() + corrupt = b'{"partial": "NEVER-LEAK-THIS-secret' + self.path.write_bytes(corrupt) + store = serverless_config.SecretStore(self.path) + + for operation in ( + lambda: store.get("worker-1"), + lambda: store.set("worker-1", "replacement-secret"), + lambda: store.delete("worker-1"), + ): + with self.subTest(operation=operation): + with self.assertRaises(Exception) as raised: + operation() + self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception)) + self.assertNotIn("replacement-secret", str(raised.exception)) + self.assertEqual(self.path.read_bytes(), corrupt) + + def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self): + for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'): + with self.subTest(corrupt=corrupt): + self.path.parent.mkdir(exist_ok=True) + self.path.write_bytes(corrupt) + with self.assertRaises(Exception) as raised: + serverless_config.SecretStore(self.path).get("worker") + self.assertNotIn("opaque-stored-value", str(raised.exception)) + + def test_rejects_symlink_file_and_preserves_link_target(self): + target = self.root / "target.json" + target.write_text('{"worker-1": "target-secret"}', encoding="utf-8") + self.path.parent.mkdir() + self.path.symlink_to(target) + store = serverless_config.SecretStore(self.path) + + with self.assertRaises(Exception): + store.get("worker-1") + with self.assertRaises(Exception): + store.set("worker-1", "replacement-secret") + self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}') + + def test_rejects_symlinked_parent_directory(self): + real_dir = self.root / "real" + real_dir.mkdir() + link_dir = self.root / "linked" + link_dir.symlink_to(real_dir, target_is_directory=True) + store = serverless_config.SecretStore(link_dir / "secrets.json") + + with self.assertRaises(Exception): + store.set("worker-1", "secret-value") + self.assertFalse((real_dir / "secrets.json").exists()) + + def test_rejects_parent_traversal_in_store_path(self): + store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json") + + with self.assertRaises(Exception): + store.set("worker-1", "secret-value") + self.assertFalse((self.root / "secrets.json").exists()) + + def test_threaded_updates_from_independent_instances_are_not_lost(self): + workers = 12 + barrier = threading.Barrier(workers) + errors = [] + + def write_one(index): + try: + barrier.wait(timeout=5) + serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}") + except Exception as exc: # collected to fail in the test thread + errors.append(exc) + + threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish") + self.assertEqual(errors, []) + saved = json.loads(self.path.read_text(encoding="utf-8")) + self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)}) + + +class PublicConfigTests(unittest.TestCase): + def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self): + with tempfile.TemporaryDirectory() as tempdir: + store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json") + store.set("serverless-1", "stored-api-key") + config = { + "workers": [ + { + "id": "serverless-1", + "type": "serverless", + "provider": "runpod", + "endpoint_id": "ep-123", + "api_key": "inline-secret", + "clear_api_key": "another-secret", + "credential": "credential-secret", + "nested": {"runpod_credential": "nested-secret", "safe": "kept"}, + "name": "remote GPU", + }, + {"id": "serverless-2", "type": "serverless", "name": "not configured"}, + {"id": "local-1", "type": "local", "api_key": "local-secret"}, + ], + "settings": {"mode": "preserved", "credentials_note": "remove-me"}, + } + original = json.loads(json.dumps(config)) + + result = serverless_config.public_config(config, store) + + first, second, local = result["workers"] + self.assertEqual(first["api_key_configured"], True) + self.assertEqual(second["api_key_configured"], False) + self.assertNotIn("api_key_configured", local) + self.assertNotIn("api_key", first) + self.assertNotIn("clear_api_key", first) + self.assertNotIn("credential", first) + self.assertNotIn("runpod_credential", first["nested"]) + self.assertEqual(first["nested"]["safe"], "kept") + self.assertNotIn("api_key", local) + self.assertNotIn("credentials_note", result["settings"]) + self.assertEqual(first["name"], "remote GPU") + self.assertEqual(config, original) + self.assertIsNot(result, config) + self.assertIsNot(result["workers"], config["workers"]) + + def test_handles_non_dict_config_without_mutating_it(self): + config = ["safe", {"api_key": "secret"}] + result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path())) + self.assertEqual(result, ["safe", {}]) + self.assertEqual(config, ["safe", {"api_key": "secret"}]) + + def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self): + module = _load_serverless_config_module(without_fcntl=True) + store = module.SecretStore(self.root_path()) + config = { + "workers": [ + {"id": "local-1", "type": "local", "api_key": "inline-secret"} + ] + } + + result = module.public_config(config, store) + + self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]}) + self.assertEqual(config["workers"][0]["api_key"], "inline-secret") + + def test_secret_operations_fail_clearly_without_fcntl(self): + module = _load_serverless_config_module(without_fcntl=True) + store = module.SecretStore(self.root_path()) + unsupported_message = "Serverless credential storage is currently supported on POSIX masters only" + + for operation in ( + lambda: store.get("worker-1"), + lambda: store.set("worker-1", "api-key"), + lambda: store.delete("worker-1"), + lambda: store.has("worker-1"), + ): + with self.subTest(operation=operation): + with self.assertRaisesRegex(module.SecretStoreError, unsupported_message): + operation() + + @staticmethod + def root_path(): + return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json" + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serverless_integration.py b/tests/test_serverless_integration.py new file mode 100644 index 0000000..79ac7de --- /dev/null +++ b/tests/test_serverless_integration.py @@ -0,0 +1,166 @@ +"""Offline integration: real graph preparation + orchestration, fake provider.""" +import asyncio +import aiohttp as real_aiohttp +import importlib.util +import json +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + +ROOT = Path(__file__).resolve().parents[1] +PKG = "dist_serverless_integration" + + +def module(name, **attrs): + value = types.ModuleType(name) + value.__dict__.update(attrs) + return value + + +def load_modules(): + modules = {} + for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]: + modules[PKG + part] = module(PKG + part, __path__=[str(directory)]) + ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn)) + modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps)) + modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None) + modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None) + modules[PKG + ".utils.config"] = module("config", load_config=lambda: {}) + modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock()) + modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master") + modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock()) + modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), sync_worker_media=AsyncMock()) + with patch.dict(sys.modules, modules): + spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py") + q = importlib.util.module_from_spec(spec) + spec.loader.exec_module(q) + # Keep imported package namespace only; restore shared Comfy runtime stub. + modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)}) + return q, modules, ps + + +class IntegrationTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.q, modules, self.ps = load_modules() + self.patch = patch.dict(sys.modules, modules) + self.patch.start() + self.addCleanup(self.patch.stop) + self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30} + self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}} + self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate)) + self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1}) + self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None) + sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager) + sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[])) + self.prompt = { + "1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}}, + "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}}, + "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}}, + } + + async def test_real_compiler_slots_to_master_collector_without_callbacks(self): + result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client") + self.assertEqual(result[2], 2) + self.q.dispatch_worker_prompt.assert_not_awaited() + self.manager.start.assert_awaited_once() + jobs = self.manager.start.call_args.args[2] + self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"]) + self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2) + for job in jobs: + graph = job["workflow"] + self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values())) + self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values())) + self.assertNotIn("http://master", json.dumps(graph)) + master_graph = self.q.queue_prompt_payload.call_args.args[0] + self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"]) + + async def test_compile_pack_provider_decode_and_collect_real_components(self): + """Exercise the whole image path, mocking only the provider and Comfy queue.""" + import base64 + import io + import tempfile + import torch + from PIL import Image + from test_collector_list_inputs import _load_collector_module + + sys.modules["aiohttp"] = real_aiohttp + for leaf in ("serverless_media", "serverless_runtime"): + name = PKG + ".api.orchestration." + leaf + spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py")) + real_module = importlib.util.module_from_spec(spec) + sys.modules[name] = real_module + spec.loader.exec_module(real_module) + runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"] + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + image_buffer = io.BytesIO() + Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG") + image_bytes = image_buffer.getvalue() + (root / "source.png").write_bytes(image_bytes) + paths = module("folder_paths", get_input_directory=lambda: directory, + get_output_directory=lambda: directory, get_temp_directory=lambda: directory, + get_annotated_filepath=lambda name: str(root / name)) + submissions = [] + + class FakeProvider: + async def submit(self, graph, images, timeout_seconds): + self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage") + submissions.append((graph, images)) + self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage") + assert self_image == images[0]["name"] + assert base64.b64decode(images[0]["image"]) == image_bytes + return "provider-job-" + str(len(submissions)) + + async def status(self, _job_id): + filename = self.prefix + "_00001_.png" + return {"status": "COMPLETED", "output": {"images": [{ + "filename": filename, "artifact_path": filename, "artifact_type": "output", + "type": "base64", "data": base64.b64encode(image_bytes).decode(), + }]}} + + manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None) + manager._client = AsyncMock(side_effect=lambda _: FakeProvider()) + self.ps.distributed_serverless = manager + self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}} + collector_module = _load_collector_module() + collector_module.prompt_server = self.ps + sys.modules["aiohttp"] = real_aiohttp + with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}): + result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client") + master_graph = self.q.queue_prompt_payload.call_args.args[0] + inputs = master_graph["2"]["inputs"] + images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute( + images=None, audio=None, multi_job_id=inputs["multi_job_id"], + enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True, + ), timeout=3) + await asyncio.gather(*manager.tasks.values()) + self.assertEqual(result[2], 2) + self.assertEqual(len(submissions), 2) + self.assertEqual(tuple(images.shape), (2, 3, 2, 3)) + self.assertEqual(images.dtype, torch.float32) + self.assertTrue(torch.equal(images[0], images[1])) + self.assertEqual(float(images[0, 0, 0, 0]), 1.0) + self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"]) + self.assertEqual(self.ps.distributed_pending_jobs, {}) + self.q.dispatch_worker_prompt.assert_not_awaited() + + async def test_audio_rejected_before_any_dispatch_or_master_queue(self): + self.prompt["2"]["inputs"]["audio"] = ["9", 0] + self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}} + with self.assertRaises(ValueError): + await self.q.orchestrate_distributed_execution(self.prompt, None, "client") + self.manager.start.assert_not_awaited() + self.q.dispatch_worker_prompt.assert_not_awaited() + self.q.queue_prompt_payload.assert_not_awaited() + + async def test_disabled_pool_excluded(self): + self.pool["enabled"] = False + self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), []) + + async def test_master_validation_failure_submits_no_paid_job(self): + self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph") + with self.assertRaisesRegex(ValueError, "invalid local graph"): + await self.q.orchestrate_distributed_execution(self.prompt, None, "client") + self.manager.start.assert_not_awaited() diff --git a/tests/test_serverless_media.py b/tests/test_serverless_media.py new file mode 100644 index 0000000..054258c --- /dev/null +++ b/tests/test_serverless_media.py @@ -0,0 +1,419 @@ +import base64 +import importlib.util +import io +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + +from PIL import Image + +ROOT = Path(__file__).resolve().parents[1] +MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py" + + +def load_module(): + spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +media = load_module() + + +def png_bytes(mode="RGB", size=(2, 1), color=None): + if color is None: + color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30) + buffer = io.BytesIO() + Image.new(mode, size, color).save(buffer, format="PNG") + return buffer.getvalue() + + +class FolderPathsFixture: + def __init__(self, root): + self.input = root / "input" + self.output = root / "output" + self.temp = root / "temp" + for directory in (self.input, self.output, self.temp): + directory.mkdir(parents=True, exist_ok=True) + + def get_input_directory(self): + return str(self.input) + + def get_output_directory(self): + return str(self.output) + + def get_temp_directory(self): + return str(self.temp) + + def get_annotated_filepath(self, filename): + value = filename.replace("\\", "/") + # Match real ComfyUI folder_paths: root annotations are suffixes, + # not directory prefixes (output/foo.png is an input subdirectory). + for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)): + suffix = f" [{marker}]" + if value.endswith(suffix): + return str(root / value[:-len(suffix)]) + return str(self.input / value) + + +class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.paths = FolderPathsFixture(Path(self.tempdir.name)) + self.folder_paths = types.ModuleType("folder_paths") + self.folder_paths.get_input_directory = self.paths.get_input_directory + self.folder_paths.get_output_directory = self.paths.get_output_directory + self.folder_paths.get_temp_directory = self.paths.get_temp_directory + self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath + self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths}) + self.folder_patch.start() + self.addCleanup(self.folder_patch.stop) + self.addCleanup(self.tempdir.cleanup) + + async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self): + payload = png_bytes() + (self.paths.input / "cat.png").write_bytes(payload) + workflow = { + "1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}}, + "2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}}, + "3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}}, + } + + packed = await media.prepare_images(workflow) + + self.assertEqual(len(packed), 1) + self.assertEqual(set(packed[0]), {"name", "image"}) + self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$") + self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload) + self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"]) + self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"]) + self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged") + self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable") + + async def test_resolves_annotated_input_output_and_temp_sources(self): + sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp")) + workflow = {} + for index, (directory, filename) in enumerate(sources, 1): + (directory / filename).write_bytes(png_bytes()) + root = directory.name + workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}} + + packed = await media.prepare_images(workflow) + + self.assertEqual(len(packed), 3) + self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()}) + self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed}) + + async def test_directory_prefix_is_not_a_root_annotation(self): + (self.paths.input / "output").mkdir() + payload = png_bytes(color=(100, 0, 0)) + (self.paths.input / "output" / "same.png").write_bytes(payload) + (self.paths.output / "same.png").write_bytes(png_bytes()) + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}} + packed = await media.prepare_images(workflow) + self.assertEqual(base64.b64decode(packed[0]["image"]), payload) + + async def test_rejects_unknown_root_annotation(self): + (self.paths.input / "cat.png").write_bytes(png_bytes()) + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}} + with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"): + await media.prepare_images(workflow) + + async def test_recognizes_custom_file_image_reference(self): + (self.paths.input / "source.bmp").write_bytes(png_bytes()) + workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}} + + packed = await media.prepare_images(workflow) + + self.assertEqual(len(packed), 1) + self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"]) + + async def test_rejects_missing_and_dynamic_loadimage_values(self): + cases = ( + {"1": {"class_type": "LoadImage", "inputs": {}}}, + {"1": {"class_type": "LoadImage"}}, + {"1": {"class_type": "LoadImage", "inputs": None}}, + {"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}}, + ) + for workflow in cases: + with self.subTest(workflow=workflow): + with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"): + await media.prepare_images(workflow) + + async def test_rejects_missing_source_and_unapproved_extension(self): + for value in ("absent.png", "animation.gif", "clip.mp4"): + with self.subTest(value=value): + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}} + with self.assertRaises(media.ServerlessMediaError): + await media.prepare_images(workflow) + + async def test_rejects_unsupported_audio_or_video_references_without_reading(self): + for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")): + with self.subTest(key=key): + workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}} + with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"): + await media.prepare_images(workflow) + + async def test_rejects_parent_traversal_and_absolute_outside_roots(self): + outside = Path(self.tempdir.name).parent / "outside.png" + outside.write_bytes(png_bytes()) + self.addCleanup(lambda: outside.unlink(missing_ok=True)) + for value in ("../outside.png", str(outside)): + with self.subTest(value=value): + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}} + with self.assertRaises(media.ServerlessMediaError): + await media.prepare_images(workflow) + + async def test_rejects_symlink_escape_from_allowed_root(self): + outside = Path(self.tempdir.name).parent / "linked-outside.png" + outside.write_bytes(png_bytes()) + self.addCleanup(lambda: outside.unlink(missing_ok=True)) + (self.paths.input / "link.png").symlink_to(outside) + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}} + + with self.assertRaises(media.ServerlessMediaError): + await media.prepare_images(workflow) + + async def test_rejects_individual_and_aggregate_input_size_overages(self): + huge = self.paths.input / "large.jpg" + with huge.open("wb") as handle: + handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1) + workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}} + with self.assertRaisesRegex(media.ServerlessMediaError, "size"): + await media.prepare_images(workflow) + + huge.unlink() + first = self.paths.input / "first.png" + second = self.paths.input / "second.png" + first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1)) + second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1)) + workflow = { + "1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}}, + "2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}}, + } + with self.assertRaisesRegex(media.ServerlessMediaError, "size"): + await media.prepare_images(workflow) + + +class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase): + def manifest(self, prefix="render"): + return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}] + + def artifact(self, filename, payload=None, **extra): + entry = { + "filename": filename, + "artifact_type": "output", + "artifact_path": f"/safe/output/{filename}", + "subfolder": "", + "type": "base64", + "data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"), + } + entry.update(extra) + return entry + + async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self): + output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]} + + result = await media.collect_images(output, self.manifest()) + + self.assertEqual(list(result), ["job-1"]) + self.assertEqual(len(result["job-1"]), 2) + self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"])) + self.assertEqual(result["job-1"][0].shape, (1, 2, 3)) + self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255) + + async def test_groups_images_under_manifest_job_ids(self): + manifest = [ + self.manifest("first")[0], + {"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"}, + ] + output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]} + + result = await media.collect_images(output, manifest) + + self.assertEqual(set(result), {"job-1", "job-2"}) + self.assertEqual(len(result["job-1"]), 1) + self.assertEqual(len(result["job-2"]), 1) + + async def test_preserves_rgba_alpha(self): + output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]} + + result = await media.collect_images(output, self.manifest()) + + self.assertEqual(result["job-1"][0].shape, (1, 2, 4)) + self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255) + + async def test_signed_url_is_downloaded_only_through_injected_downloader(self): + output = {"images": [{ + "filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url", + "url": "https://storage.example.test/signed", + }]} + fetch = AsyncMock(return_value=png_bytes()) + + result = await media.collect_images(output, self.manifest(), download=fetch) + + self.assertEqual(result["job-1"][0].shape, (1, 2, 3)) + fetch.assert_awaited_once_with("https://storage.example.test/signed") + + async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self): + output = {"images": [ + self.artifact("preview_00001_.png"), + self.artifact("render_00001_.png"), + {"filename": "mesh.obj", "type": "url", "url": "not a fetch target"}, + ]} + + result = await media.collect_images(output, self.manifest()) + + self.assertEqual(len(result["job-1"]), 1) + with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"): + await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest()) + + async def test_accepts_compiler_length_flat_safe_prefix(self): + prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64 + manifest = self.manifest(prefix) + output = {"images": [self.artifact(f"{prefix}_00001_.png")]} + + result = await media.collect_images(output, manifest) + + self.assertEqual(len(result["job-1"]), 1) + + async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self): + duplicate = self.manifest() * 2 + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images({"images": []}, duplicate) + unsafe = self.manifest("../render") + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images({"images": []}, unsafe) + + async def test_rejects_duplicate_canonical_artifact_identity(self): + item = self.artifact("render_00001_.png") + with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"): + await media.collect_images({"images": [item, dict(item)]}, self.manifest()) + + async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self): + wrong_type = self.artifact("render_00001_.png", artifact_type="temp") + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images({"images": [wrong_type]}, self.manifest()) + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest()) + + async def test_rejects_provider_errors_even_when_job_completed(self): + for output in ( + {"error": "raw provider secret", "images": []}, + {"errors": ["raw provider secret"], "images": []}, + {"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}}, + ): + with self.subTest(output=output): + with self.assertRaises(media.ServerlessMediaError) as ctx: + await media.collect_images(output, self.manifest()) + self.assertNotIn("raw provider secret", str(ctx.exception)) + + async def test_rejects_invalid_base64_and_malformed_image(self): + for item in ( + self.artifact("render_00001_.png", data="!not base64!"), + self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()), + ): + with self.subTest(item=item): + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images({"images": [item]}, self.manifest()) + + async def test_validates_shape_consistency_within_collector(self): + output = {"images": [ + self.artifact("render_00001_.png", png_bytes(size=(2, 1))), + self.artifact("render_00002_.png", png_bytes(size=(3, 1))), + ]} + with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"): + await media.collect_images(output, self.manifest()) + + async def test_bounds_output_file_count_and_bytes(self): + too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)} + with self.assertRaisesRegex(media.ServerlessMediaError, "files"): + await media.collect_images(too_many, self.manifest()) + + with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4): + output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]} + with self.assertRaisesRegex(media.ServerlessMediaError, "size"): + await media.collect_images(output, self.manifest()) + + async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self): + output = {"images": [ + self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()), + self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()), + ]} + with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3): + with self.assertRaisesRegex(media.ServerlessMediaError, "size"): + await media.collect_images(output, self.manifest()) + + async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self): + for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"): + with self.subTest(url=url): + fetch = AsyncMock(return_value=png_bytes()) + output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]} + with self.assertRaises(media.ServerlessMediaError): + await media.collect_images(output, self.manifest(), download=fetch) + fetch.assert_not_awaited() + + async def test_public_resolver_rejects_non_global_dns_answers(self): + resolver = media._PublicResolver() + resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))]) + self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1) + resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))]) + with self.assertRaises(OSError): + await resolver.resolve("cdn.example.test", 443) + + async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self): + class Content: + def __init__(self, chunks): + self.chunks = iter(chunks) + async def iter_chunked(self, _size): + for chunk in self.chunks: + yield chunk + + class Response: + def __init__(self, status, headers=None, chunks=()): + self.status, self.headers = status, headers or {} + self.content = Content(chunks) + async def __aenter__(self): + return self + async def __aexit__(self, *_args): + return False + + class Session: + def __init__(self, responses): + self.responses, self.calls = list(responses), [] + async def __aenter__(self): + return self + async def __aexit__(self, *_args): + return False + def get(self, url, **kwargs): + self.calls.append((url, kwargs)) + return self.responses.pop(0) + + session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])]) + with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session: + payload = await media.download_public_artifact("https://storage.example.test/signed") + + self.assertEqual(payload, b"abcdef") + self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"]) + self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls)) + self.assertFalse(create_session.call_args.kwargs["trust_env"]) + self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver) + self.assertNotIn("Authorization", session.calls[0][1].get("headers", {})) + await create_session.call_args.kwargs["connector"].close() + + blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})]) + with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session: + with self.assertRaises(media.ServerlessMediaError): + await media.download_public_artifact("https://storage.example.test/signed") + self.assertEqual(len(blocked.calls), 1) + await blocked_session.call_args.kwargs["connector"].close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serverless_prompt.py b/tests/test_serverless_prompt.py new file mode 100644 index 0000000..a463dfb --- /dev/null +++ b/tests/test_serverless_prompt.py @@ -0,0 +1,335 @@ +import importlib.util +import json +import sys +import types +from pathlib import Path + +import pytest + + +MODULE_PATH = ( + Path(__file__).resolve().parents[1] + / "api" + / "orchestration" + / "serverless_prompt.py" +) +_spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH) +assert _spec is not None and _spec.loader is not None +_serverless_prompt = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_serverless_prompt) +compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt + + +def _prepared_prompt(): + return { + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}}, + "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}}, + "3": { + "class_type": "DistributedSeed", + "inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"}, + }, + "4": { + "class_type": "DistributedCollector", + "inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True}, + }, + "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}}, + "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}}, + "7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}}, + "8": {"class_type": "UnrelatedNode", "inputs": {}}, + } + + +def _run(prompt): + return compile_serverless_prompt(prompt) + + +def _load_prompt_transform(): + module_path = MODULE_PATH.with_name("prompt_transform.py") + package_name = "serverless_prompt_transform_testpkg" + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + for suffix in ("api", "api.orchestration", "utils"): + pkg = types.ModuleType(f"{package_name}.{suffix}") + pkg.__path__ = [] + sys.modules[f"{package_name}.{suffix}"] = pkg + logging_module = types.ModuleType(f"{package_name}.utils.logging") + setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None) + setattr(logging_module, "log", lambda *_args, **_kwargs: None) + sys.modules[f"{package_name}.utils.logging"] = logging_module + spec = importlib.util.spec_from_file_location( + f"{package_name}.api.orchestration.prompt_transform", module_path + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_compiles_prompt_prepared_by_existing_worker_helpers(): + prompt_transform = _load_prompt_transform() + original = { + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}}, + "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}}, + "3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, + "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}}, + } + pruned = prompt_transform.prune_prompt_for_worker(original) + prompt_index = prompt_transform.PromptIndex(pruned) + prepared = prompt_transform.apply_participant_overrides( + pruned, + participant_id="worker-a", + enabled_worker_ids=["worker-a"], + job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"), + master_url="http://master.invalid", + delegate_master=False, + prompt_index=prompt_index, + ) + + graph, manifest = _run(prepared) + + assert graph["2"]["inputs"]["seed"] == 101 + assert set(graph) == {"1", "2", "7"} + assert graph["7"]["class_type"] == "SaveImage" + assert manifest[0]["job_id"] == "run_4" + assert manifest[0]["collector_id"] == "4" + + +def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream(): + original = _prepared_prompt() + + graph, manifest = _run(original) + + assert manifest == [ + { + "collector_id": "4", + "job_id": "job/42", + "node_id": "9", + "kind": "image", + "prefix": graph["9"]["inputs"]["filename_prefix"], + } + ] + prefix = manifest[0]["prefix"] + assert graph["9"] == { + "class_type": "SaveImage", + "inputs": {"images": ["2", 0], "filename_prefix": prefix}, + "_meta": {"title": "Serverless output: collector 4 (image)"}, + } + assert "/" not in prefix and "\\" not in prefix + assert graph["2"]["inputs"]["seed"] == 41 + assert set(graph) == {"1", "2", "9"} + assert all("Distributed" not in node["class_type"] for node in graph.values()) + assert original["3"]["inputs"]["seed"] == 40 + assert original["4"]["inputs"]["images"] == ["2", 0] + + +def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts(): + prompt = _prepared_prompt() + prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe") + + first_graph, first_manifest = _run(prompt) + second_graph, second_manifest = _run(prompt) + + assert first_manifest == second_manifest + assert first_graph == second_graph + prefix = first_manifest[0]["prefix"] + assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"] + assert not any(char in prefix for char in ("/", "\\", ":")) + assert first_manifest[0]["job_id"] == "../run/a" + + +def test_compiles_multiple_disconnected_collectors_in_stable_order(): + prompt = { + "10": {"class_type": "LoadImage", "inputs": {}}, + "20": {"class_type": "LoadImage", "inputs": {}}, + "2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}}, + "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}}, + "30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}}, + } + + graph, manifest = _run(prompt) + + assert [item["collector_id"] for item in manifest] == ["1", "2"] + assert [item["node_id"] for item in manifest] == ["31", "32"] + assert graph["31"]["inputs"]["images"] == ["10", 0] + assert graph["32"]["inputs"]["images"] == ["20", 0] + assert set(graph) == {"10", "20", "31", "32"} + + +def test_rejects_audio_input_instead_of_emitting_audio_saver(): + prompt = { + "1": {"class_type": "LoadAudio", "inputs": {}}, + "2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}}, + } + + with pytest.raises(ValueError, match="audio.*unsupported"): + _run(prompt) + + +def test_rejects_collector_without_image_input(): + with pytest.raises(ValueError, match="collector.*image"): + _run({"1": {"class_type": "DistributedCollector", "inputs": {}}}) + + +@pytest.mark.parametrize( + "class_type", + ["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"], +) +def test_rejects_unsupported_distributed_and_batch_nodes(class_type): + prompt = { + "1": {"class_type": class_type, "inputs": {}}, + "2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + "3": {"class_type": "LoadImage", "inputs": {}}, + } + + with pytest.raises(ValueError, match="unsupported"): + _run(prompt) + + +def test_rejects_collector_to_collector_dependency_through_processing_node(): + prompt = { + "1": {"class_type": "LoadImage", "inputs": {}}, + "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}}, + "3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + } + + with pytest.raises(ValueError, match="collector-to-collector"): + _run(prompt) + + +@pytest.mark.parametrize( + "link", + [["missing", 0], ["1", -1], ["1", "0"], ["1", True]], +) +def test_rejects_malformed_api_links(link): + prompt = { + "1": {"class_type": "LoadImage", "inputs": {}}, + "2": {"class_type": "DistributedCollector", "inputs": {"images": link}}, + } + + with pytest.raises(ValueError, match="link"): + _run(prompt) + + +def test_rejects_unexpected_output_slot_from_removed_distributed_node(): + prompt = { + "1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}}, + "2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}}, + "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, + } + + with pytest.raises(ValueError, match="output slot"): + _run(prompt) + + +def test_distributed_seed_matches_worker_offset_and_master_passthrough(): + prompt = { + "1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}}, + "2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}}, + "3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + } + + graph, _manifest = _run(prompt) + + assert graph["3"]["inputs"] == {"seed": 12, "steps": 7} + assert "1" not in graph and "2" not in graph + + +def test_distributed_value_uses_one_based_worker_values_and_type_coercion(): + prompt = { + "1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}}, + "2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}}, + "3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + } + + graph, _manifest = _run(prompt) + + assert graph["3"]["inputs"] == {"steps": 8, "seed": 3} + assert "1" not in graph and "2" not in graph + + +def test_resolves_stock_primitive_links_for_distributed_seed_and_value(): + prompt = { + "1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}}, + "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}}, + "3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}}, + "4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}}, + "5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}}, + "6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}}, + "7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}}, + } + + graph, _manifest = _run(prompt) + + assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75} + assert not {"1", "2", "3", "4", "5"}.intersection(graph) + + +def test_rejects_dynamic_linked_seed_or_value_without_executing_it(): + prompt = { + "1": {"class_type": "ArbitraryNode", "inputs": {}}, + "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}}, + "3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + } + + with pytest.raises(ValueError, match="literal or supported Primitive"): + _run(prompt) + + +def test_primitive_resolution_detects_cycles(): + prompt = { + "1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}}, + "2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}}, + "3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}}, + "4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}}, + "5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}}, + } + + with pytest.raises(ValueError, match="cycle"): + _run(prompt) + + +def test_model_name_trivial_text_passthrough_is_inlined(): + prompt = { + "1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}}, + "2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}}, + "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}}, + } + + graph, _manifest = _run(prompt) + + assert graph["2"]["inputs"]["model_name"] == "model.safetensors" + assert "1" not in graph + + +def test_rejects_dynamic_linked_model_name_text(): + prompt = { + "1": {"class_type": "ArbitraryNode", "inputs": {}}, + "2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}}, + "3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}}, + "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}}, + } + + with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"): + _run(prompt) + + +def test_resolves_boolean_and_float_primitive_links(): + prompt = { + "1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}}, + "2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}, + "3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}}, + "4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}}, + "5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}}, + "6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}}, + } + + graph, _manifest = _run(prompt) + + assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5} + assert not {"1", "2", "3", "4"}.intersection(graph) diff --git a/tests/test_serverless_routes.py b/tests/test_serverless_routes.py new file mode 100644 index 0000000..1264c3f --- /dev/null +++ b/tests/test_serverless_routes.py @@ -0,0 +1,463 @@ +import copy +import importlib.util +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = copy.deepcopy(payload) + self.status = status + + +class _FakeRequest: + def __init__(self, payload=None, error=None): + self.payload = payload + self.error = error + + async def json(self): + if self.error is not None: + raise self.error + return self.payload + + +class _Routes: + def get(self, _path): + return lambda handler: handler + + def post(self, _path): + return lambda handler: handler + + +def _load_modules(): + repo = Path(__file__).resolve().parents[1] + package_name = "dist_serverless_routes_testpkg" + for name in list(sys.modules): + if name == package_name or name.startswith(f"{package_name}."): + del sys.modules[name] + + root = types.ModuleType(package_name) + root.__path__ = [str(repo)] + sys.modules[package_name] = root + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [str(repo / "api")] + sys.modules[f"{package_name}.api"] = api_pkg + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [str(repo / "utils")] + sys.modules[f"{package_name}.utils"] = utils_pkg + orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration") + orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")] + sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.log = lambda *_args, **_kwargs: None + logging_module.debug_log = lambda *_args, **_kwargs: None + sys.modules[logging_module.__name__] = logging_module + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.HEARTBEAT_TIMEOUT = 30 + sys.modules[constants_module.__name__] = constants_module + + def load_file(module_name, path): + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py") + serverless_config = load_file( + f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py" + ) + + runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime") + runtime.secret_store = None + runtime.manager = types.SimpleNamespace(tasks={}, runs={}) + runtime.get_secret_store = lambda: runtime.secret_store + runtime.get_manager = lambda: runtime.manager + sys.modules[runtime.__name__] = runtime + + import aiohttp # Ensure the real dependency is loaded; never replace global aiohttp. + previous_server = sys.modules.get("server") + server = types.ModuleType("server") + server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) + sys.modules["server"] = server + + network = types.ModuleType(f"{package_name}.utils.network") + + async def handle_api_error(_request, error, status=500): + return _FakeResponse({"status": "error", "message": str(error)}, status) + + network.handle_api_error = handle_api_error + network.normalize_host = lambda value: value + sys.modules[network.__name__] = network + + routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py") + routes.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status) + ) + + if previous_server is None: + sys.modules.pop("server", None) + else: + sys.modules["server"] = previous_server + return config_module, serverless_config, runtime, routes + + +config_module, serverless_config, runtime_module, config_routes = _load_modules() + + +class _RunningTask: + def done(self): + return False + + +class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.tempdir = tempfile.TemporaryDirectory() + self.root = Path(self.tempdir.name) + config_module.CONFIG_FILE = str(self.root / "gpu_config.json") + config_module.invalidate_config_cache() + self.config = config_module.get_default_config() + self.config["workers"] = [] + self.assertTrue(config_module.save_config(self.config)) + self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json") + runtime_module.secret_store = self.secrets + runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}) + + def tearDown(self): + config_module.invalidate_config_cache() + self.tempdir.cleanup() + + def read_config_file(self): + return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8")) + + async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self): + self.secrets.set("pool-1", "private-api-key") + self.config["workers"] = [ + {"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1", + "api_key": "inline-api-key", "name": "GPU"} + ] + Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8") + config_module.invalidate_config_cache() + + response = await config_routes.get_config_endpoint(_FakeRequest()) + + self.assertEqual(response.status, 200) + worker = response.payload["workers"][0] + self.assertTrue(worker["api_key_configured"]) + self.assertNotIn("api_key", worker) + self.assertNotIn("private-api-key", json.dumps(response.payload)) + self.assertNotIn("inline-api-key", json.dumps(response.payload)) + + async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self): + secret = "runpod-secret-never-returned" + payload = { + "worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod", + "endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2, + "timeout_seconds": 600, "enabled": True, "api_key": secret, + } + + response = await config_routes.update_worker_endpoint(_FakeRequest(payload)) + + self.assertEqual(response.status, 200) + self.assertTrue(response.payload["api_key_configured"]) + self.assertNotIn(secret, json.dumps(response.payload)) + saved = self.read_config_file()["workers"][0] + self.assertEqual(saved["endpoint_id"], "ep-123") + self.assertEqual(saved["jobs_per_run"], 3) + self.assertEqual(saved["concurrency"], 2) + self.assertEqual(saved["timeout_seconds"], 600) + for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"): + self.assertNotIn(obsolete, saved) + self.assertEqual(self.secrets.get("pool-1"), secret) + + async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self): + self.secrets.set("pool-1", "saved-secret") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool", + "endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3, + "timeout_seconds": 900, "enabled": False, + }] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.update_worker_endpoint( + _FakeRequest({"worker_id": "pool-1", "enabled": True}) + ) + + worker = self.read_config_file()["workers"][0] + self.assertTrue(worker["enabled"]) + self.assertEqual(worker["endpoint_id"], "ep-123") + self.assertEqual(worker["jobs_per_run"], 4) + self.assertEqual(worker["concurrency"], 3) + self.assertEqual(worker["timeout_seconds"], 900) + self.assertEqual(self.secrets.get("pool-1"), "saved-secret") + self.assertTrue(response.payload["api_key_configured"]) + + async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self): + secret = "must-not-leak" + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "type": "serverless", "provider": "runpod", + "endpoint_id": "ep-123", "concurrency": True, "api_key": secret, + })) + + self.assertEqual(response.status, 400) + self.assertNotIn(secret, json.dumps(response.payload)) + self.assertEqual(self.read_config_file()["workers"], []) + self.assertIsNone(self.secrets.get("pool-1")) + + async def test_json_decode_error_is_generic_and_does_not_echo_body(self): + secret = "body-secret-value" + response = await config_routes.update_worker_endpoint( + _FakeRequest(error=ValueError(f"bad json near {secret}")) + ) + self.assertEqual(response.status, 400) + self.assertNotIn(secret, json.dumps(response.payload)) + + async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self): + secret = "bulk-secret-never-write" + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", + "endpoint_id": "ep-123", "nested": {"api_key": secret}, + }] + })) + + self.assertEqual(response.status, 400) + self.assertNotIn(secret, json.dumps(response.payload)) + self.assertEqual(self.read_config_file()["workers"], []) + self.assertFalse(self.secrets.path.exists()) + + async def test_bulk_update_validates_and_returns_redacted_serverless_config(self): + self.secrets.set("pool-1", "saved-secret") + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", + "endpoint_id": "ep-123", "name": "Pool", "enabled": True, + "host": "obsolete-host", "port": 8188, "cuda_device": 0, + }] + })) + + self.assertEqual(response.status, 200) + self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"]) + saved = self.read_config_file()["workers"][0] + for obsolete in ("host", "port", "cuda_device"): + self.assertNotIn(obsolete, saved) + self.assertNotIn("api_key", response.payload["config"]["workers"][0]) + self.assertEqual(self.secrets.get("pool-1"), "saved-secret") + + async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self): + self.secrets.set("pool-1", "bulk-clear-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{ + "id": "pool-1", "type": "local", "port": 8188, + "provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True, + }] + })) + + self.assertEqual(response.status, 200) + self.assertIsNone(self.secrets.get("pool-1")) + saved = self.read_config_file()["workers"][0] + for field in ("provider", "endpoint_id", "api_key_configured"): + self.assertNotIn(field, saved) + + async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self): + existing = {"id": "local-1", "type": "local", "port": 8188} + self.config["workers"] = [existing] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}] + })) + + self.assertEqual(response.status, 400) + self.assertEqual(self.read_config_file()["workers"], [existing]) + + async def test_delete_serverless_worker_clears_saved_key(self): + self.secrets.set("pool-1", "delete-me") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123" + }] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) + + self.assertEqual(response.status, 200) + self.assertEqual(self.read_config_file()["workers"], []) + self.assertIsNone(self.secrets.get("pool-1")) + + async def test_type_conversion_clears_secret_and_removes_serverless_fields(self): + self.secrets.set("pool-1", "old-key") + self.config["workers"] = [{ + "id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod", + "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600, + }] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188, + "cuda_device": 0, "enabled": False, + })) + + self.assertEqual(response.status, 200) + self.assertIsNone(self.secrets.get("pool-1")) + worker = self.read_config_file()["workers"][0] + self.assertEqual(worker["type"], "local") + for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"): + self.assertNotIn(serverless_field, worker) + + async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self): + self.secrets.set("pool-1", "saved-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", + "enabled": False, + }] + self.assertTrue(config_module.save_config(self.config)) + runtime_module.manager = types.SimpleNamespace( + tasks={"run-1": _RunningTask()}, + runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + ) + + enabled = await config_routes.update_worker_endpoint( + _FakeRequest({"worker_id": "pool-1", "enabled": True}) + ) + blocked = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "endpoint_id": "ep-456", + })) + + self.assertEqual(enabled.status, 200) + self.assertEqual(blocked.status, 409) + self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123") + + async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self): + self.secrets.set("pool-1", "saved-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + + keep = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "api_key": "", + })) + clear = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "clear_api_key": True, + })) + + self.assertTrue(keep.payload["api_key_configured"]) + self.assertFalse(clear.payload["api_key_configured"]) + self.assertIsNone(self.secrets.get("pool-1")) + + async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self): + with self.subTest(case="inline key"): + config = config_module.get_default_config() + config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}] + self.assertFalse(config_module.save_config(config)) + with self.subTest(case="invalid serverless record"): + config = config_module.get_default_config() + config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}] + self.assertFalse(config_module.save_config(config)) + + async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self): + self.secrets.set("pool-1", "saved-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + runtime_module.manager = types.SimpleNamespace( + tasks={"run-1": _RunningTask()}, + runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + ) + + credential_change = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "api_key": "new-key", + })) + type_change = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0, + })) + deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) + + self.assertEqual(credential_change.status, 409) + self.assertEqual(type_change.status, 409) + self.assertEqual(deletion.status, 409) + self.assertEqual(self.secrets.get("pool-1"), "saved-key") + self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") + + async def test_secret_write_is_rolled_back_when_config_persistence_fails(self): + secret = "write-then-rollback" + with patch.object(config_module, "save_config", return_value=False): + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + "api_key": secret, + })) + + self.assertEqual(response.status, 500) + self.assertIsNone(self.secrets.get("pool-1")) + self.assertEqual(self.read_config_file()["workers"], []) + self.assertEqual(config_module.load_config()["workers"], []) + + async def test_regular_worker_delete_does_not_require_posix_secret_storage(self): + self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}] + self.assertTrue(config_module.save_config(self.config)) + with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")): + response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"})) + self.assertEqual(response.status, 200) + self.assertEqual(self.read_config_file()["workers"], []) + + async def test_enabling_serverless_without_key_is_rejected(self): + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True, + })) + self.assertEqual(response.status, 400) + self.assertEqual(self.read_config_file()["workers"], []) + + async def test_bulk_concurrency_change_is_blocked_during_active_run(self): + self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}] + self.assertTrue(config_module.save_config(self.config)) + runtime_module.manager = types.SimpleNamespace( + tasks={"run-1": _RunningTask()}, + runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + ) + updated = copy.deepcopy(self.config["workers"]) + updated[0]["concurrency"] = 2 + response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated})) + self.assertEqual(response.status, 409) + self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1) + + async def test_non_serverless_update_behavior_is_preserved(self): + self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}] + self.assertTrue(config_module.save_config(self.config)) + + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "local-1", "name": "Renamed", "enabled": True, + })) + + self.assertEqual(response.status, 200) + worker = self.read_config_file()["workers"][0] + self.assertEqual(worker["name"], "Renamed") + self.assertTrue(worker["enabled"]) + + async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self): + raw = { + "workers": [{ + "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + "host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old", + }] + } + normalized = config_module._merge_with_defaults(raw, config_module.get_default_config()) + worker = normalized["workers"][0] + for obsolete in ("host", "port", "cuda_device", "extra_args"): + self.assertNotIn(obsolete, worker) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serverless_runtime.py b/tests/test_serverless_runtime.py new file mode 100644 index 0000000..8b9bb35 --- /dev/null +++ b/tests/test_serverless_runtime.py @@ -0,0 +1,143 @@ +import asyncio +import importlib.util +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + +ROOT = Path(__file__).resolve().parents[1] +PACKAGE = "dist_serverless_runtime_tests" +for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]: + module = types.ModuleType(name) + module.__path__ = [str(path)] + sys.modules[name] = module + +def load(name): + spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py")) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + +runtime = load("serverless_runtime") + +class Secrets: + def get(self, worker_id): + return "private-test-key" + +class RuntimeTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.temp = tempfile.TemporaryDirectory() + self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock()) + self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object())) + self.manager.poll_seconds = 0 + self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"} + self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]} + + async def asyncTearDown(self): + await self.manager.cancel() + self.temp.cleanup() + + async def test_job_id_persisted_before_status_and_results_injected(self): + api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock()) + async def status(job_id): + self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job") + return {"status": "COMPLETED", "output": {"images": []}} + api.status = status + with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})): + await self.manager.start("execution", "prompt", [self.prepared]) + await self.manager.wait_for_execution("execution") + self.assertEqual(api.submit.await_count, 1) + item = self.server.distributed_pending_jobs["collector-job"].get_nowait() + self.assertEqual(item["worker_id"], "pool:1") + self.assertEqual(item["tensor"], "tensor") + self.assertTrue(item["is_last"]) + self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED") + + async def test_output_error_marks_failed_and_delivers_explicit_error(self): + api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock()) + with patch.object(runtime, "RunpodClient", return_value=api): + await self.manager.start("execution", "prompt", [self.prepared]) + await self.manager.wait_for_execution("execution") + record = self.manager.ledger.all()[0] + self.assertEqual(record["status"], "FAILED") + item = self.server.distributed_pending_jobs["collector-job"].get_nowait() + self.assertIn("serverless_error", item) + self.assertNotIn("private-test-key", str(record) + str(item)) + + async def test_ambiguous_submit_never_retries(self): + api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock()) + with patch.object(runtime, "RunpodClient", return_value=api): + await self.manager.start("execution", "prompt", [self.prepared]) + await self.manager.wait_for_execution("execution") + self.assertEqual(api.submit.await_count, 1) + self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN") + + async def test_cancel_targets_owned_job_only(self): + polled = asyncio.Event() + async def status(_job): + polled.set() + await asyncio.Event().wait() + api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"})) + with patch.object(runtime, "RunpodClient", return_value=api): + await self.manager.start("execution", "prompt", [self.prepared]) + await polled.wait() + await self.manager.cancel(worker_id="other-pool") + api.cancel.assert_not_awaited() + await self.manager.cancel(worker_id="pool") + api.cancel.assert_awaited_once_with("owned-job") + + async def test_per_pool_concurrency_bounds_submission(self): + gate = asyncio.Event() + reached = asyncio.Event() + count = 0 + async def submit(*args): + nonlocal count + count += 1 + reached.set() + return f"provider-{count}" + async def status(_job): + await gate.wait() + return {"status": "COMPLETED", "output": {"images": []}} + api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock()) + second = dict(self.prepared, worker=dict(self.worker, id="pool:2")) + with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})): + await self.manager.start("execution", "prompt", [self.prepared, second]) + await reached.wait() + await asyncio.sleep(0) + self.assertEqual(count, 1) + gate.set() + await self.manager.wait_for_execution("execution") + self.assertEqual(count, 2) + + async def test_cancel_persisted_job_after_master_restart(self): + self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1) + api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"})) + with patch.object(runtime, "RunpodClient", return_value=api) as client: + await self.manager.cancel(worker_id="pool") + api.cancel.assert_awaited_once_with("old-job") + self.assertEqual(client.call_args.args[1], "old-endpoint") + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED") + + async def test_master_failure_cancels_job_without_partial_success(self): + self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}}) + api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"})) + with patch.object(runtime, "RunpodClient", return_value=api): + await self.manager.start("execution", "prompt", [self.prepared]) + await self.manager.wait_for_execution("execution") + api.submit.assert_not_awaited() + self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED") + + async def test_deadline_includes_local_slot_wait(self): + self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01) + api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"})) + # A blocked slot expires without a billable submission. + self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0)) + with patch.object(runtime, "RunpodClient", return_value=api): + await self.manager.start("execution", "prompt", [self.prepared]) + await self.manager.wait_for_execution("execution") + api.submit.assert_not_awaited() + self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT") diff --git a/tests/test_serverless_status.py b/tests/test_serverless_status.py new file mode 100644 index 0000000..544e091 --- /dev/null +++ b/tests/test_serverless_status.py @@ -0,0 +1,58 @@ +import importlib.util +import json +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + +ROOT = Path(__file__).resolve().parents[1] +PKG = "dist_serverless_status_tests" + + +class StatusTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + def mod(name, **kwargs): + result = types.ModuleType(name) + result.__dict__.update(kwargs) + return result + self.manager = types.SimpleNamespace(health=AsyncMock(return_value={"ok": True, "queueRemaining": 0, "message": "Ready (scales from zero)"}), public_jobs=lambda _: [], cancel=AsyncMock(return_value=1)) + self.pool = {"id": "pool", "type": "serverless", "endpoint_id": "ep"} + modules = {PKG: mod(PKG, __path__=[str(ROOT)]), PKG + ".api": mod("api", __path__=[str(ROOT/'api')]), + PKG + ".utils": mod("utils", __path__=[]), + PKG + ".utils.config": mod("config", load_config=lambda: {"workers": [self.pool]}), + PKG + ".api.orchestration": mod("orchestration", __path__=[str(ROOT/'api/orchestration')]), + PKG + ".api.orchestration.serverless_runtime": mod("runtime", get_manager=lambda: self.manager), + "server": mod("server", PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda f: f, post=lambda _: lambda f: f))))} + with patch.dict(sys.modules, modules): + spec = importlib.util.spec_from_file_location(PKG+'.api.serverless_routes', ROOT/'api/serverless_routes.py') + self.routes = importlib.util.module_from_spec(spec) + spec.loader.exec_module(self.routes) + + async def test_health_only_never_submits_and_supports_scale_to_zero(self): + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"})) + result = await self.routes.test_serverless_endpoint(request) + self.assertEqual(result.status, 200) + self.assertTrue(json.loads(result.text)["ok"]) + self.manager.health.assert_awaited_once_with(self.pool, fresh=True) + self.manager.cancel.assert_not_awaited() + + async def test_unknown_pool_cannot_cancel_other_jobs(self): + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "other"})) + result = await self.routes.cancel_serverless_endpoint(request) + self.assertEqual(result.status, 404) + self.manager.cancel.assert_not_awaited() + + async def test_transport_error_does_not_leak_credentials(self): + self.manager.health.side_effect = RuntimeError("Bearer secret-token https://signed.invalid/?token=value") + request = types.SimpleNamespace(match_info={"worker_id": "pool"}) + result = await self.routes.serverless_status_endpoint(request) + self.assertFalse(json.loads(result.text)["ok"]) + self.assertNotIn("secret-token", result.text) + self.assertNotIn("signed.invalid", result.text) + + async def test_cancel_uses_only_requested_pool(self): + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"})) + result = await self.routes.cancel_serverless_endpoint(request) + self.assertEqual(result.status, 200) + self.manager.cancel.assert_awaited_once_with(worker_id="pool") diff --git a/utils/config.py b/utils/config.py index c580415..cff821d 100644 --- a/utils/config.py +++ b/utils/config.py @@ -2,6 +2,7 @@ Configuration management for ComfyUI-Distributed. """ import asyncio +import copy import os import json from contextlib import asynccontextmanager @@ -15,6 +16,57 @@ _config_cache = None _config_mtime = 0.0 _config_lock = asyncio.Lock() +_SERVERLESS_OBSOLETE_WORKER_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"} + + +def _contains_serverless_secret(value): + if isinstance(value, dict): + return any( + (isinstance(key, str) and _is_secret_field_name(key)) + or _contains_serverless_secret(item) + for key, item in value.items() + ) + if isinstance(value, (list, tuple)): + return any(_contains_serverless_secret(item) for item in value) + return False + + +def _is_secret_field_name(key): + normalized = "".join(character for character in key.casefold() if character.isalnum()) + return normalized in {"apikey", "clearapikey"} or "credential" in normalized + + +def _normalize_serverless_workers(config): + """Remove obsolete local-worker fields without adding local defaults.""" + if not isinstance(config, dict): + return config + workers = config.get("workers") + if isinstance(workers, list): + for worker in workers: + if isinstance(worker, dict) and worker.get("type") == "serverless": + for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS: + worker.pop(field, None) + return config + + +def _validated_config_for_save(config): + """Return a safe normalized copy ready for gpu_config.json.""" + if _contains_serverless_secret(config): + raise ValueError("API keys must be stored outside the config file") + saved = _normalize_serverless_workers(copy.deepcopy(config)) + workers = saved.get("workers") if isinstance(saved, dict) else None + if isinstance(workers, list) and any( + isinstance(worker, dict) and worker.get("type") == "serverless" for worker in workers + ): + from .serverless_config import validate_serverless_worker + for index, worker in enumerate(workers): + if isinstance(worker, dict) and worker.get("type") == "serverless": + saved_worker = validate_serverless_worker(worker) + for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS: + saved_worker.pop(field, None) + workers[index] = saved_worker + return saved + def _config_path(): return CONFIG_FILE @@ -62,7 +114,7 @@ def _merge_with_defaults(data, defaults): if key not in merged: merged[key] = value - return merged + return _normalize_serverless_workers(merged) def invalidate_config_cache(): @@ -97,11 +149,12 @@ def load_config(): return _config_cache def save_config(config): - """Saves the configuration to file.""" + """Saves normalized configuration without allowing API keys in the config file.""" tmp_path = f"{_config_path()}.tmp" try: + saved_config = _validated_config_for_save(config) with open(tmp_path, 'w', encoding='utf-8') as f: - json.dump(config, f, indent=2) + json.dump(saved_config, f, indent=2) f.flush() os.fsync(f.fileno()) os.replace(tmp_path, _config_path()) @@ -121,12 +174,21 @@ async def config_transaction(): """Acquire config lock, yield loaded config, and save if changed.""" async with _config_lock: config = load_config() + original_config = copy.deepcopy(config) original_snapshot = json.dumps(config, sort_keys=True) - yield config + try: + yield config + except BaseException: + if isinstance(config, dict): + config.clear() + config.update(original_config) + raise updated_snapshot = json.dumps(config, sort_keys=True) - if updated_snapshot != original_snapshot: - if not save_config(config): - raise RuntimeError("Failed to save config") + if updated_snapshot != original_snapshot and not save_config(config): + if isinstance(config, dict): + config.clear() + config.update(original_config) + raise RuntimeError("Failed to save config") def ensure_config_exists(): """Creates default config file if it doesn't exist. Used by __init__.py""" diff --git a/utils/network.py b/utils/network.py index 6d7f33f..d24cc92 100644 --- a/utils/network.py +++ b/utils/network.py @@ -87,6 +87,8 @@ def _split_host_and_port(host): def build_worker_url(worker, endpoint=""): """Construct the worker base URL with optional endpoint.""" + if worker.get("type") == "serverless": + raise ValueError("Serverless pools have no ComfyUI worker URL") host = (worker.get("host") or "").strip() port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188) diff --git a/utils/serverless_config.py b/utils/serverless_config.py new file mode 100644 index 0000000..e828305 --- /dev/null +++ b/utils/serverless_config.py @@ -0,0 +1,380 @@ +"""Validation and private storage for Runpod serverless worker settings. + +Secret storage is supported only on POSIX masters. The module remains +importable elsewhere so ordinary worker configuration and redaction continue +to work, but secret-store operations fail explicitly without POSIX locking. +""" + +from __future__ import annotations + +import copy +import json +import os +import re +import secrets +import stat +import threading +from contextlib import contextmanager +from pathlib import Path +from typing import Iterator + +try: + import fcntl +except ImportError: # Windows and other platforms without POSIX file locking. + fcntl = None + + +class ServerlessConfigError(ValueError): + """A serverless worker configuration is invalid.""" + + +class SecretStoreError(RuntimeError): + """The secret store could not be accessed safely.""" + + +_ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII) +_UNSUPPORTED_SECRET_STORE_MESSAGE = ( + "Serverless credential storage is currently supported on POSIX masters only" +) +_NUMERIC_FIELDS = { + "jobs_per_run": (1, 16, 1), + "concurrency": (1, 16, 1), + "timeout_seconds": (30, 86400, 1800), +} +_SENSITIVE_FIELD_NAMES = {"api_key", "clear_api_key"} +_THREAD_LOCKS: dict[str, threading.Lock] = {} +_THREAD_LOCKS_GUARD = threading.Lock() + + +def validate_serverless_worker(worker: dict) -> dict: + """Return a validated copy of a Runpod serverless worker record. + + Provider credentials are intentionally not required here; they are stored + separately by :class:`SecretStore`. + """ + if not isinstance(worker, dict): + raise ServerlessConfigError("worker must be an object") + + try: + result = copy.deepcopy(worker) + except Exception: + raise ServerlessConfigError("worker must be copyable") from None + + worker_id = result.get("id") + if not isinstance(worker_id, str) or not worker_id.strip(): + raise ServerlessConfigError("id must be a non-empty string") + + endpoint_id = result.get("endpoint_id") + if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id): + raise ServerlessConfigError("endpoint_id must be an opaque Runpod identifier") + + if result.get("type", "serverless") != "serverless": + raise ServerlessConfigError("type must be serverless") + if result.get("provider", "runpod") != "runpod": + raise ServerlessConfigError("provider is not supported") + if "enabled" in result and type(result["enabled"]) is not bool: + raise ServerlessConfigError("enabled must be a boolean") + + for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items(): + value = result.get(field, default) + if type(value) is not int or not minimum <= value <= maximum: + raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}") + result[field] = value + + result["type"] = "serverless" + result["provider"] = "runpod" + return result + + +def _thread_lock_for(path: str) -> threading.Lock: + with _THREAD_LOCKS_GUARD: + lock = _THREAD_LOCKS.get(path) + if lock is None: + lock = threading.Lock() + _THREAD_LOCKS[path] = lock + return lock + + +class SecretStore: + """Atomically persist worker-id to API-key mappings in a private JSON file.""" + + def __init__(self, path: Path): + self.path = Path(path) + + def get(self, worker_id: str) -> str | None: + self._require_supported() + worker_id = self._validate_worker_id(worker_id) + try: + with self._locked(create_parent=False) as (parent_fd, filename): + data = self._read_mapping(parent_fd, filename) + return data.get(worker_id) + except FileNotFoundError: + return None + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + + def set(self, worker_id: str, key: str) -> None: + self._require_supported() + worker_id = self._validate_worker_id(worker_id) + if ( + not isinstance(key, str) + or not key + or len(key) > 4096 + or "\r" in key + or "\n" in key + ): + raise SecretStoreError("secret must be 1 to 4096 characters and contain no CR or LF") + try: + with self._locked(create_parent=True) as (parent_fd, filename): + data = self._read_mapping(parent_fd, filename) + data[worker_id] = key + self._write_mapping(parent_fd, filename, data) + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + + def delete(self, worker_id: str) -> None: + self._require_supported() + worker_id = self._validate_worker_id(worker_id) + try: + with self._locked(create_parent=False) as (parent_fd, filename): + data = self._read_mapping(parent_fd, filename) + if worker_id in data: + del data[worker_id] + self._write_mapping(parent_fd, filename, data) + except FileNotFoundError: + return + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + + def has(self, worker_id: str) -> bool: + return self.get(worker_id) is not None + + @staticmethod + def _require_supported() -> None: + # Never substitute a weaker Windows store: all operations retain the + # dir_fd/O_NOFOLLOW and flock protections implemented below. + if os.name != "posix" or fcntl is None: + raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE) + + @staticmethod + def _validate_worker_id(worker_id: str) -> str: + if not isinstance(worker_id, str) or not worker_id.strip(): + raise SecretStoreError("worker_id must be a non-empty string") + return worker_id + + def _absolute_path(self) -> Path: + path = self.path + if ".." in path.parts: + raise SecretStoreError("unsafe secret store path") + if not path.is_absolute(): + path = Path.cwd() / path + if len(path.parts) < 2 or path.name in ("", ".", ".."): + raise SecretStoreError("unsafe secret store path") + return path + + @contextmanager + def _locked(self, create_parent: bool) -> Iterator[tuple[int, str]]: + lock_api = fcntl + if os.name != "posix" or lock_api is None: + raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE) + path = self._absolute_path() + filename = path.name + lock_name = f".{filename}.lock" + thread_lock = _thread_lock_for(str(path)) + + with thread_lock: + parent_fd = self._open_parent(path, create=create_parent) + lock_fd = None + try: + flags = os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW + lock_fd = os.open(lock_name, flags, 0o600, dir_fd=parent_fd) + lock_stat = os.fstat(lock_fd) + if not stat.S_ISREG(lock_stat.st_mode): + raise SecretStoreError("unsafe secret store lock") + os.fchmod(lock_fd, 0o600) + lock_api.flock(lock_fd, lock_api.LOCK_EX) + self._check_target(parent_fd, filename) + yield parent_fd, filename + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + finally: + if lock_fd is not None: + try: + lock_api.flock(lock_fd, lock_api.LOCK_UN) + except OSError: + pass + os.close(lock_fd) + os.close(parent_fd) + + @staticmethod + def _open_parent(path: Path, create: bool) -> int: + parts = path.parts + current_fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC) + try: + for part in parts[1:-1]: + created = False + try: + next_fd = os.open( + part, + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=current_fd, + ) + except FileNotFoundError: + if not create: + raise + try: + os.mkdir(part, 0o700, dir_fd=current_fd) + created = True + except FileExistsError: + pass + next_fd = os.open( + part, + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=current_fd, + ) + if created: + os.fchmod(next_fd, 0o700) + os.close(current_fd) + current_fd = next_fd + return current_fd + except Exception: + os.close(current_fd) + raise + + @staticmethod + def _check_target(parent_fd: int, filename: str) -> None: + try: + target_stat = os.stat(filename, dir_fd=parent_fd, follow_symlinks=False) + except FileNotFoundError: + return + if not stat.S_ISREG(target_stat.st_mode): + raise SecretStoreError("unsafe secret store file") + + @classmethod + def _read_mapping(cls, parent_fd: int, filename: str) -> dict[str, str]: + cls._check_target(parent_fd, filename) + try: + fd = os.open( + filename, + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=parent_fd, + ) + except FileNotFoundError: + return {} + except Exception: + raise SecretStoreError("Unable to access secret store") from None + + try: + file_stat = os.fstat(fd) + if not stat.S_ISREG(file_stat.st_mode): + raise SecretStoreError("unsafe secret store file") + os.fchmod(fd, 0o600) + with os.fdopen(fd, "rb") as stream: + raw = stream.read() + try: + mapping = json.loads(raw.decode("utf-8"), object_pairs_hook=cls._unique_object) + except Exception: + raise SecretStoreError("secret store is corrupt") from None + if not isinstance(mapping, dict) or any( + not isinstance(worker_id, str) + or not worker_id.strip() + or not isinstance(key, str) + or not key + for worker_id, key in mapping.items() + ): + raise SecretStoreError("secret store is corrupt") + return mapping + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + finally: + try: + os.close(fd) + except OSError: + pass + + @staticmethod + def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, value in pairs: + if key in result: + raise ValueError("duplicate key") + result[key] = value + return result + + @staticmethod + def _write_mapping(parent_fd: int, filename: str, mapping: dict[str, str]) -> None: + payload = json.dumps(mapping, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + temp_name = f".{filename}.{secrets.token_hex(12)}.tmp" + temp_fd = None + try: + flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW + temp_fd = os.open(temp_name, flags, 0o600, dir_fd=parent_fd) + os.fchmod(temp_fd, 0o600) + with os.fdopen(temp_fd, "wb") as stream: + temp_fd = None + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + SecretStore._check_target(parent_fd, filename) + os.replace(temp_name, filename, src_dir_fd=parent_fd, dst_dir_fd=parent_fd) + os.fsync(parent_fd) + except SecretStoreError: + raise + except Exception: + raise SecretStoreError("Unable to access secret store") from None + finally: + if temp_fd is not None: + try: + os.close(temp_fd) + except OSError: + pass + try: + os.unlink(temp_name, dir_fd=parent_fd) + except OSError: + pass + + +def _redact_credentials(value): + if isinstance(value, dict): + redacted = {} + for key, item in value.items(): + field_name = key.casefold() if isinstance(key, str) else "" + if field_name in _SENSITIVE_FIELD_NAMES or "credential" in field_name: + continue + redacted[key] = _redact_credentials(item) + return redacted + if isinstance(value, list): + return [_redact_credentials(item) for item in value] + if isinstance(value, tuple): + return tuple(_redact_credentials(item) for item in value) + return value + + +def public_config(config, store: SecretStore): + """Return a deep-copied public configuration with credentials removed.""" + result = _redact_credentials(copy.deepcopy(config)) + if not isinstance(result, dict): + return result + + workers = result.get("workers") + if isinstance(workers, list): + for worker in workers: + if isinstance(worker, dict) and worker.get("type") == "serverless": + worker_id = worker.get("id") + configured = ( + store.has(worker_id) + if isinstance(worker_id, str) and worker_id.strip() + else False + ) + worker["api_key_configured"] = configured + return result diff --git a/web/apiClient.js b/web/apiClient.js index bda5895..7028c1f 100644 --- a/web/apiClient.js +++ b/web/apiClient.js @@ -7,7 +7,7 @@ export function createApiClient(baseUrl) { const request = async ( endpoint, options = {}, - { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {}, + { retries = TIMEOUTS.MAX_RETRIES, retry = true, redactErrors = false } = {}, ) => { const maxAttempts = retry ? retries : 1; let lastError; @@ -36,7 +36,7 @@ export function createApiClient(baseUrl) { return await response.json(); } catch (error) { lastError = error; - console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`); + console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${redactErrors ? "[redacted]" : error.message}`); if (attempt < maxAttempts - 1) { await new Promise(resolve => setTimeout(resolve, delay)); delay *= 2; // Exponential backoff @@ -86,9 +86,35 @@ export function createApiClient(baseUrl) { }, async updateWorker(workerId, data) { - return request('/distributed/config/update_worker', { + const apiKey = data?.api_key; + try { + return await request('/distributed/config/update_worker', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId, ...data }) + }, { retry: false, redactErrors: Boolean(apiKey) }); + } catch (error) { + if (apiKey && typeof error?.message === "string") { + error.message = error.message.split(apiKey).join("[redacted]"); + } + throw error; + } + }, + + async testServerlessWorker(workerId) { + return request('/distributed/serverless/test', { method: 'POST', - body: JSON.stringify({ worker_id: workerId, ...data }) + body: JSON.stringify({ worker_id: workerId }), + }, { retry: false }); + }, + + async getServerlessStatus(workerId) { + return request(`/distributed/serverless/status/${encodeURIComponent(workerId)}`); + }, + + async cancelServerlessWorker(workerId) { + return request('/distributed/serverless/cancel', { + method: 'POST', + body: JSON.stringify({ worker_id: workerId }), }, { retry: false }); }, diff --git a/web/executionUtils.js b/web/executionUtils.js index ae601d4..8102bc6 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -112,6 +112,21 @@ export async function performPreflightCheck(extension, workers) { const startTime = Date.now(); const checkPromises = workers.map(async (worker) => { + if (worker.type === "serverless") { + try { + const providerStatus = await extension.api.getServerlessStatus(worker.id); + if (providerStatus?.ok) { + extension.log(`Serverless worker ${worker.name} is ready to scale`, "debug"); + return { worker, active: true }; + } + extension.log(`Serverless worker ${worker.name} is not ready: ${providerStatus?.message || "unavailable"}`, "debug"); + return { worker, active: false }; + } catch (error) { + extension.log(`Serverless worker ${worker.name} status check failed: ${error.message}`, "debug"); + return { worker, active: false }; + } + } + const workerUrl = getWorkerUrl(extension, worker); extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug"); diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index d59cb5d..4a00049 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -18,6 +18,19 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) { const workerState = extension.state.getWorker(workerId); const isLaunching = Boolean(workerState?.launching); + if (worker.type === "serverless") { + const providerStatus = workerState?.serverlessStatus || {}; + if (newStatus.online && newStatus.processing) { + extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, `Runpod active (${newStatus.queueCount || 0} queued)`, false); + } else if (newStatus.online) { + extension.ui.updateStatusDot(workerId, STATUS_COLORS.ONLINE_GREEN, providerStatus.message || "Runpod ready to scale", false); + } else if (worker.enabled) { + extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, providerStatus.message || "Runpod endpoint unavailable", false); + } + updateWorkerControls(extension, workerId); + return true; + } + if (isLaunching && !newStatus.online) { extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true); } else if (newStatus.online && newStatus.processing) { diff --git a/web/tests/serverlessWorkers.test.js b/web/tests/serverlessWorkers.test.js new file mode 100644 index 0000000..d7d481e --- /dev/null +++ b/web/tests/serverlessWorkers.test.js @@ -0,0 +1,412 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client" } }), { virtual: true }); + +import { createApiClient } from "../apiClient.js"; +import { createWorkerSettingsForm } from "../ui/settingsForm.js"; +import { DistributedUI } from "../ui.js"; +import { checkWorkerStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js"; +import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js"; +import { isRemoteWorker, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js"; +import { performPreflightCheck } from "../executionUtils.js"; + +function createDocumentMock() { + const elements = new Map(); + const makeElement = (tagName) => { + const element = { + tagName, + children: [], + style: {}, + classList: { + values: new Set(), + add(...names) { names.forEach((name) => this.values.add(name)); }, + remove(...names) { names.forEach((name) => this.values.delete(name)); }, + toggle(name, force) { + const enabled = force ?? !this.values.has(name); + if (enabled) this.values.add(name); + else this.values.delete(name); + return enabled; + }, + }, + appendChild(child) { this.children.push(child); child.parentNode = this; return child; }, + dispatchEvent() { this.onchange?.({ target: this }); }, + }; + Object.defineProperty(element, "id", { + get() { return this._id || ""; }, + set(value) { this._id = value; elements.set(value, this); }, + }); + return element; + }; + return { + elements, + document: { + createElement: makeElement, + getElementById: (id) => elements.get(id) || null, + }, + }; +} + +function createFormUi() { + return { + createFormGroup(label, value, id, type = "text", placeholder = "") { + const group = document.createElement("div"); + const input = document.createElement("input"); + input.id = id; + input.value = String(value ?? ""); + input.type = type; + input.placeholder = placeholder; + group.appendChild(input); + return { group, input, label }; + }, + createButton(text, onclick) { + const button = document.createElement("button"); + button.textContent = text; + button.onclick = onclick; + return button; + }, + createButtonGroup(buttons) { + const group = document.createElement("div"); + buttons.forEach((button) => group.appendChild(button)); + return group; + }, + }; +} + +function makeSaveExtension(worker, updateWorker = vi.fn().mockResolvedValue({})) { + return { + config: { workers: [worker] }, + api: { updateWorker }, + state: { updateWorker: vi.fn(), setWorkerExpanded: vi.fn() }, + app: { extensionManager: { toast: { add: vi.fn() } } }, + panelElement: null, + }; +} + +describe("Runpod serverless worker panel", () => { + let originalDocument; + let originalWindow; + let elements; + + beforeEach(() => { + originalDocument = globalThis.document; + originalWindow = globalThis.window; + const mock = createDocumentMock(); + elements = mock.elements; + globalThis.document = mock.document; + globalThis.window = { + location: { hostname: "localhost", protocol: "http:", port: "8188", origin: "http://localhost:8188" }, + }; + }); + + afterEach(() => { + globalThis.document = originalDocument; + globalThis.window = originalWindow; + vi.restoreAllMocks(); + }); + + it("classifies Runpod serverless as remote for worker allocation", () => { + expect(isRemoteWorker({}, { type: "serverless" })).toBe(true); + }); + + it("shows only serverless settings and keeps credentials write-only", () => { + const worker = { + id: "w1", name: "Runpod", type: "remote", host: "worker.example", port: 8188, + cuda_device: 0, extra_args: "--test", api_key_configured: true, + }; + const form = createWorkerSettingsForm(createFormUi(), {}, worker); + const type = elements.get("worker-type-w1"); + type.value = "serverless"; + type.dispatchEvent(); + + expect(type.children.map((option) => [option.value, option.textContent])).toContainEqual([ + "serverless", "Serverless (Runpod)", + ]); + for (const id of ["host-group-w1", "port-group-w1", "cuda-group-w1", "args-group-w1"]) { + expect(elements.get(id).style.display).toBe("none"); + } + for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "concurrency-w1", "timeout-seconds-w1"]) { + expect(elements.has(id)).toBe(true); + } + for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "concurrency-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) { + expect(elements.get(id).style.display).toBe("flex"); + } + expect(elements.get("serverless-help-w1").textContent).toContain("image Collectors only"); + expect(elements.get("api-key-w1").type).toBe("password"); + expect(elements.get("api-key-w1").value).toBe(""); + expect(elements.get("api-key-w1").placeholder).toMatch(/leave blank/i); + expect(elements.get("jobs-per-run-w1").value).toBe("1"); + expect(elements.get("jobs-per-run-w1").min).toBe("1"); + expect(elements.get("jobs-per-run-w1").max).toBe("16"); + expect(elements.get("concurrency-w1").value).toBe("1"); + expect(elements.get("concurrency-w1").min).toBe("1"); + expect(elements.get("concurrency-w1").max).toBe("16"); + expect(elements.get("timeout-seconds-w1").value).toBe("1800"); + expect(elements.get("timeout-seconds-w1").min).toBe("30"); + expect(elements.get("timeout-seconds-w1").max).toBe("86400"); + expect(form.children.length).toBeGreaterThan(0); + }); + + it("saves fixed Runpod settings without enabling or retaining the API key", async () => { + const worker = { id: "w2", name: "Old", type: "remote", enabled: false, api_key_configured: false }; + const extension = makeSaveExtension(worker); + elements.set("name-w2", { value: "Runpod worker" }); + elements.set("worker-type-w2", { value: "serverless" }); + elements.set("endpoint-id-w2", { value: "ep_123" }); + elements.set("api-key-w2", { value: "secret-token" }); + elements.set("clear-api-key-w2", { checked: false }); + elements.set("jobs-per-run-w2", { value: "4" }); + elements.set("concurrency-w2", { value: "3" }); + elements.set("timeout-seconds-w2", { value: "2400" }); + + await saveWorkerSettings(extension, "w2"); + + expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", { + name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123", + api_key: "secret-token", jobs_per_run: 4, concurrency: 3, timeout_seconds: 2400, enabled: false, + }); + expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true }); + expect(worker).not.toHaveProperty("api_key"); + expect(JSON.stringify(extension.config)).not.toContain("secret-token"); + expect(elements.get("api-key-w2").value).toBe(""); + expect(extension.state.updateWorker).toHaveBeenCalledWith("w2", { enabled: false }); + expect(JSON.stringify(extension.state.updateWorker.mock.calls)).not.toContain("secret-token"); + }); + + it("omits a blank API key so the saved key is preserved", async () => { + const worker = { id: "w-blank", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true }; + const extension = makeSaveExtension(worker); + elements.set("name-w-blank", { value: "Runpod" }); + elements.set("worker-type-w-blank", { value: "serverless" }); + elements.set("endpoint-id-w-blank", { value: "ep_123" }); + elements.set("api-key-w-blank", { value: "" }); + elements.set("clear-api-key-w-blank", { checked: false }); + elements.set("jobs-per-run-w-blank", { value: "1" }); + elements.set("concurrency-w-blank", { value: "1" }); + elements.set("timeout-seconds-w-blank", { value: "1800" }); + + await saveWorkerSettings(extension, "w-blank"); + + const payload = extension.api.updateWorker.mock.calls[0][1]; + expect(payload).not.toHaveProperty("api_key"); + expect(payload).not.toHaveProperty("clear_api_key"); + expect(worker.api_key_configured).toBe(true); + expect(worker).not.toHaveProperty("api_key"); + }); + + it("omits a blank key to preserve it and sends an explicit clear flag when requested", async () => { + const worker = { id: "w3", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true }; + const extension = makeSaveExtension(worker); + elements.set("name-w3", { value: "Runpod" }); + elements.set("worker-type-w3", { value: "serverless" }); + elements.set("endpoint-id-w3", { value: "ep_123" }); + elements.set("api-key-w3", { value: "" }); + elements.set("clear-api-key-w3", { checked: true }); + elements.set("jobs-per-run-w3", { value: "1" }); + elements.set("concurrency-w3", { value: "1" }); + elements.set("timeout-seconds-w3", { value: "1800" }); + + await saveWorkerSettings(extension, "w3"); + + expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", { + name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123", + clear_api_key: true, jobs_per_run: 1, concurrency: 1, timeout_seconds: 1800, enabled: true, + }); + expect(worker.api_key_configured).toBe(false); + expect(worker).not.toHaveProperty("api_key"); + }); + + it("uses provider status instead of a direct ComfyUI URL probe", async () => { + const statusElement = { textContent: "" }; + elements.set("serverless-status-w4", statusElement); + const status = { + ok: true, queueRemaining: 0, message: "Ready to scale", + jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w4", execution_id: "exec-1" }], + }; + const extension = { + config: { workers: [{ id: "w4", name: "Runpod", type: "serverless", enabled: true }] }, + api: { getServerlessStatus: vi.fn().mockResolvedValue(status), probeWorker: vi.fn() }, + statusCheckAbortController: null, + state: { + setWorkerStatus: vi.fn(), + updateWorker: vi.fn(), + getWorkerStatus: vi.fn().mockReturnValue({ online: true, processing: false, queueCount: 0 }), + getWorker: vi.fn().mockReturnValue({}), + isWorkerLaunching: vi.fn().mockReturnValue(false), + }, + updateWorkerCard: vi.fn().mockReturnValue(true), + log: vi.fn(), + }; + + await checkWorkerStatus(extension, extension.config.workers[0]); + + expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w4"); + expect(extension.api.probeWorker).not.toHaveBeenCalled(); + expect(extension.state.setWorkerStatus).toHaveBeenCalledWith("w4", { + online: true, processing: true, queueCount: 0, + }); + expect(statusElement.textContent).toContain("Ready to scale"); + expect(statusElement.textContent).toContain("RUNNING"); + }); + + it("considers a healthy zero-worker Runpod pool ready and never builds a worker URL", async () => { + const worker = { id: "w-ready", name: "Runpod", type: "serverless", enabled: true }; + const extension = { + api: { + getServerlessStatus: vi.fn().mockResolvedValue({ ok: true, queueRemaining: 0, message: "Ready to scale" }), + probeWorker: vi.fn(), + }, + log: vi.fn(), + }; + + await expect(performPreflightCheck(extension, [worker])).resolves.toEqual([worker]); + expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w-ready"); + expect(extension.api.probeWorker).not.toHaveBeenCalled(); + expect(getWorkerUrl(extension, worker)).toBeNull(); + }); + + it("does not render or invoke direct ComfyUI launch, stop, or log controls", async () => { + const ui = new DistributedUI(); + const worker = { id: "w-controls", name: "Runpod", type: "serverless", enabled: true }; + const controls = ui.createControlsSection({ dynamic: true }, worker, {}, true); + const collect = (element) => [element, ...(element.children || []).flatMap(collect)]; + const elementsInControls = collect(controls); + const extension = { + config: { workers: [worker] }, + api: { + launchWorker: vi.fn(), stopWorker: vi.fn(), getRemoteWorkerLog: vi.fn(), + }, + }; + + expect(elementsInControls.filter((element) => element.tagName === "button")).toHaveLength(0); + await launchWorker(extension, worker.id); + await stopWorker(extension, worker.id); + await viewWorkerLog(extension, worker.id, true); + expect(extension.api.launchWorker).not.toHaveBeenCalled(); + expect(extension.api.stopWorker).not.toHaveBeenCalled(); + expect(extension.api.getRemoteWorkerLog).not.toHaveBeenCalled(); + }); + + it("uses the serverless status, test, and cancel routes with the frozen request shapes", async () => { + const originalFetch = globalThis.fetch; + const statusPayload = { + ok: true, + queueRemaining: 0, + message: "Ready to scale", + jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w5", execution_id: "exec-1" }], + }; + globalThis.fetch = vi.fn() + .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint test passed" }) }) + .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue(statusPayload) }) + .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true }) }); + try { + const client = createApiClient("http://localhost:8188"); + await client.testServerlessWorker("w5"); + await expect(client.getServerlessStatus("w5")).resolves.toEqual(statusPayload); + await client.cancelServerlessWorker("w5"); + + expect(globalThis.fetch).toHaveBeenNthCalledWith(1, "http://localhost:8188/distributed/serverless/test", expect.objectContaining({ + method: "POST", body: JSON.stringify({ worker_id: "w5" }), + })); + expect(globalThis.fetch).toHaveBeenNthCalledWith(2, "http://localhost:8188/distributed/serverless/status/w5", expect.any(Object)); + expect(globalThis.fetch).toHaveBeenNthCalledWith(3, "http://localhost:8188/distributed/serverless/cancel", expect.objectContaining({ + method: "POST", body: JSON.stringify({ worker_id: "w5" }), + })); + } finally { + globalThis.fetch = originalFetch; + } + }); + + it("tests only the saved Runpod endpoint and reports the backend message", async () => { + const button = document.createElement("button"); + button.id = "test-serverless-w6"; + button.textContent = "Test Saved Endpoint"; + const extension = { + api: { + testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }), + queueDistributed: vi.fn(), + }, + app: { extensionManager: { toast: { add: vi.fn() } } }, + }; + + await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" }); + + expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6"); + expect(extension.api.queueDistributed).not.toHaveBeenCalled(); + expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({ + severity: "success", detail: "Endpoint reachable", + })); + expect(button.textContent).toBe("Test Saved Endpoint"); + expect(button.disabled).toBe(false); + }); + + it("redacts API-key-bearing worker update errors from logs and user-visible exceptions", async () => { + const originalFetch = globalThis.fetch; + const logSpy = vi.spyOn(console, "log").mockImplementation(() => {}); + const secret = "never-log-this-runpod-key"; + globalThis.fetch = vi.fn().mockResolvedValue({ + ok: false, + status: 400, + json: vi.fn().mockResolvedValue({ message: `Rejected key ${secret}` }), + }); + try { + const client = createApiClient("http://localhost:8188"); + await expect(client.updateWorker("w-secret", { api_key: secret })).rejects.toThrow("[redacted]"); + expect(logSpy.mock.calls.flat().join(" ")).not.toContain(secret); + } finally { + globalThis.fetch = originalFetch; + } + }); + + it("never sends VRAM operations to serverless workers", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn(); + vi.useFakeTimers(); + const button = { + textContent: "Clear Worker VRAM", style: { cssText: "" }, disabled: false, + classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } }, + }; + const extension = { enabledWorkers: [{ id: "s-vram", name: "Runpod", type: "serverless", enabled: true }] }; + try { + await handleClearMemory(extension, button); + expect(globalThis.fetch).not.toHaveBeenCalled(); + } finally { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + globalThis.fetch = originalFetch; + } + }); + + it("cancels each enabled serverless worker once without changing normal worker interrupts", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = vi.fn().mockResolvedValue({ ok: true }); + vi.useFakeTimers(); + const button = { + textContent: "Interrupt Workers", style: { cssText: "" }, disabled: false, + classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } }, + }; + const workers = [ + { id: "s1", name: "Serverless 1", type: "serverless", enabled: true }, + { id: "s2", name: "Serverless 2", type: "serverless", enabled: true }, + { id: "s3", name: "Disabled", type: "serverless", enabled: false }, + { id: "local", name: "Local", type: "local", enabled: true, port: 8189 }, + ]; + const extension = { + enabledWorkers: workers.filter((worker) => worker.enabled), + api: { cancelServerlessWorker: vi.fn().mockResolvedValue({ ok: true }) }, + log: vi.fn(), + }; + try { + await handleInterruptWorkers(extension, button); + + expect(extension.api.cancelServerlessWorker).toHaveBeenCalledTimes(2); + expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(1, "s1"); + expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(2, "s2"); + expect(globalThis.fetch).toHaveBeenCalledTimes(1); + expect(globalThis.fetch.mock.calls[0][0]).toBe("http://localhost:8189/interrupt"); + } finally { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + globalThis.fetch = originalFetch; + } + }); +}); diff --git a/web/ui.js b/web/ui.js index 2f992a1..adc4b6a 100644 --- a/web/ui.js +++ b/web/ui.js @@ -7,6 +7,14 @@ import { createLogModal } from './ui/logModal.js'; import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js'; import { isRemoteWorker } from './workerSettings.js'; +const escapeWorkerInfo = (value) => String(value ?? "").replace(/[&<>"']/g, (character) => ({ + "&": "&", + "<": "<", + ">": ">", + '"': """, + "'": "'", +}[character])); + const cardConfigs = { master: { checkbox: { @@ -58,7 +66,15 @@ const cardConfigs = { infoText: (data, extension) => { const isRemote = isRemoteWorker(extension, data); const isCloud = data.type === 'cloud'; - + const isServerless = data.type === 'serverless'; + + if (isServerless) { + const endpoint = data.endpoint_id || "Endpoint not configured"; + const jobsPerRun = data.jobs_per_run ?? 1; + const concurrency = data.concurrency ?? 1; + return `${escapeWorkerInfo(data.name)}
Runpod Serverless • ${escapeWorkerInfo(endpoint)} • ${escapeWorkerInfo(jobsPerRun)} jobs/run • ${escapeWorkerInfo(concurrency)} slots`; + } + if (isCloud) { // For cloud workers, don't show port (it's always 443) return `${data.name}
${data.host}`; @@ -591,18 +607,24 @@ export class DistributedUI { controlsWrapper.appendChild(badge); } else if (config.dynamic && data) { if (isRemote) { - const isCloud = data.type === 'cloud'; - const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; - const workerTypeBadge = this.createInfoBox(workerTypeText); - workerTypeBadge.title = "Worker is externally hosted"; - controlsWrapper.appendChild(workerTypeBadge); + if (data.type === "serverless") { + const workerTypeBadge = this.createInfoBox("Runpod Serverless worker"); + workerTypeBadge.title = "Provider-managed workers scale on demand"; + controlsWrapper.appendChild(workerTypeBadge); + } else { + const isCloud = data.type === 'cloud'; + const workerTypeText = isCloud ? "Cloud worker" : "Remote worker"; + const workerTypeBadge = this.createInfoBox(workerTypeText); + workerTypeBadge.title = "Worker is externally hosted"; + controlsWrapper.appendChild(workerTypeBadge); - const logBtn = this.createButton('View Log', () => viewWorkerLog(extension, data.id, true)); - logBtn.id = `log-${data.id}`; - logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; - logBtn.classList.add("btn--log"); - logBtn.title = "View remote worker log"; - controlsWrapper.appendChild(logBtn); + const logBtn = this.createButton('View Log', () => viewWorkerLog(extension, data.id, true)); + logBtn.id = `log-${data.id}`; + logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; + logBtn.classList.add("btn--log"); + logBtn.title = "View remote worker log"; + controlsWrapper.appendChild(logBtn); + } } else { const controls = this.createWorkerControls(data.id, { launch: () => launchWorker(extension, data.id), diff --git a/web/ui/entityCard.js b/web/ui/entityCard.js index 69bcd14..4ac988c 100644 --- a/web/ui/entityCard.js +++ b/web/ui/entityCard.js @@ -1,6 +1,20 @@ import { updateWorkerControls, toggleWorkerExpanded } from "../workerLifecycle.js"; import { isRemoteWorker } from "../workerSettings.js"; +function formatServerlessStatus(status) { + if (!status) return "Runpod status pending"; + const message = status.message || (status.ok ? "Ready to scale" : "Runpod status unavailable"); + const queue = Number(status.queueRemaining); + const parts = [message]; + if (Number.isFinite(queue)) parts.push(`${Math.max(0, queue)} queued`); + const jobs = (Array.isArray(status.jobs) ? status.jobs : []) + .slice(0, 2) + .map((job) => String(job?.status || "").trim()) + .filter(Boolean); + if (jobs.length) parts.push(`Jobs: ${jobs.join(", ")}`); + return parts.join(" • "); +} + export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { const config = cardConfigs[entityType] || {}; const isPlaceholder = entityType === 'blueprint' || entityType === 'add'; @@ -74,6 +88,16 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { infoSpan.innerHTML = config.infoText(data, extension); workerContent.appendChild(infoSpan); + if (isWorker && data.type === "serverless") { + const statusMessage = document.createElement("small"); + statusMessage.id = `serverless-status-${data.id}`; + statusMessage.className = "dist-worker-info__meta dist-serverless-status"; + statusMessage.style.display = "block"; + statusMessage.setAttribute("aria-live", "polite"); + statusMessage.textContent = formatServerlessStatus(extension.state.getWorker(data.id)?.serverlessStatus); + infoSpan.appendChild(statusMessage); + } + infoRow.appendChild(workerContent); let settingsArrow; diff --git a/web/ui/settingsForm.js b/web/ui/settingsForm.js index 41ac7c2..e3b95d0 100644 --- a/web/ui/settingsForm.js +++ b/web/ui/settingsForm.js @@ -1,5 +1,5 @@ import { BUTTON_STYLES } from '../constants.js'; -import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings } from '../workerSettings.js'; +import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings, testServerlessWorker } from '../workerSettings.js'; export function createWorkerSettingsForm(ui, extension, worker) { const form = document.createElement("div"); @@ -33,9 +33,14 @@ export function createWorkerSettingsForm(ui, extension, worker) { cloudOption.value = "cloud"; cloudOption.textContent = "Cloud"; + const serverlessOption = document.createElement("option"); + serverlessOption.value = "serverless"; + serverlessOption.textContent = "Serverless (Runpod)"; + typeSelect.appendChild(localOption); typeSelect.appendChild(remoteOption); typeSelect.appendChild(cloudOption); + typeSelect.appendChild(serverlessOption); const runpodText = document.createElement("a"); runpodText.id = `runpod-text-${worker.id}`; @@ -54,18 +59,32 @@ export function createWorkerSettingsForm(ui, extension, worker) { const cudaGroup = document.getElementById(`cuda-group-${worker.id}`); const argsGroup = document.getElementById(`args-group-${worker.id}`); const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`); + const serverlessGroups = [ + `endpoint-id-group-${worker.id}`, + `api-key-group-${worker.id}`, + `clear-api-key-group-${worker.id}`, + `jobs-per-run-group-${worker.id}`, + `concurrency-group-${worker.id}`, + `timeout-seconds-group-${worker.id}`, + `serverless-test-group-${worker.id}`, + `serverless-help-${worker.id}`, + ].map((id) => document.getElementById(id)); - if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) { + if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput || serverlessGroups.some((group) => !group)) { return; } + const hideServerless = () => serverlessGroups.forEach((group) => { group.style.display = "none"; }); + if (workerType === "local") { + hideServerless(); hostGroup.style.display = "none"; portGroup.style.display = "flex"; cudaGroup.style.display = "flex"; argsGroup.style.display = "flex"; runpodTextElem.style.display = "none"; } else if (workerType === "remote") { + hideServerless(); hostGroup.style.display = "flex"; portGroup.style.display = "flex"; cudaGroup.style.display = "none"; @@ -76,6 +95,7 @@ export function createWorkerSettingsForm(ui, extension, worker) { hostInput.value = ""; } } else if (workerType === "cloud") { + hideServerless(); hostGroup.style.display = "flex"; portGroup.style.display = "flex"; cudaGroup.style.display = "none"; @@ -86,6 +106,13 @@ export function createWorkerSettingsForm(ui, extension, worker) { if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") { hostInput.value = ""; } + } else if (workerType === "serverless") { + hostGroup.style.display = "none"; + portGroup.style.display = "none"; + cudaGroup.style.display = "none"; + argsGroup.style.display = "none"; + runpodTextElem.style.display = "none"; + serverlessGroups.forEach((group) => { group.style.display = "flex"; }); } }; }; @@ -114,6 +141,77 @@ export function createWorkerSettingsForm(ui, extension, worker) { argsGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex"; form.appendChild(argsGroup.group); + const serverlessHelp = document.createElement("small"); + serverlessHelp.id = `serverless-help-${worker.id}`; + serverlessHelp.style.display = "none"; + serverlessHelp.textContent = "Preview: image Collectors only. The endpoint must provide the generic workflow + images handler and required models/nodes. Enabling it permits paid Runpod jobs."; + form.appendChild(serverlessHelp); + + const endpointGroup = ui.createFormGroup("Endpoint ID:", worker.endpoint_id || "", `endpoint-id-${worker.id}`, "text", "Runpod endpoint ID"); + endpointGroup.group.id = `endpoint-id-group-${worker.id}`; + endpointGroup.group.style.display = "none"; + form.appendChild(endpointGroup.group); + + const apiKeyGroup = ui.createFormGroup("Runpod API Key (write-only):", "", `api-key-${worker.id}`, "password", worker.api_key_configured ? "Leave blank to keep saved API key" : "Runpod API key"); + apiKeyGroup.group.id = `api-key-group-${worker.id}`; + apiKeyGroup.input.autocomplete = "new-password"; + apiKeyGroup.input.spellcheck = false; + apiKeyGroup.group.style.display = "none"; + form.appendChild(apiKeyGroup.group); + + const clearApiKeyGroup = document.createElement("div"); + clearApiKeyGroup.id = `clear-api-key-group-${worker.id}`; + clearApiKeyGroup.style.cssText = "display: none; align-items: center; gap: 8px;"; + const clearApiKeyInput = document.createElement("input"); + clearApiKeyInput.type = "checkbox"; + clearApiKeyInput.id = `clear-api-key-${worker.id}`; + const clearApiKeyLabel = document.createElement("label"); + clearApiKeyLabel.htmlFor = clearApiKeyInput.id; + clearApiKeyLabel.textContent = "Clear saved API key (disable pool first)"; + clearApiKeyGroup.appendChild(clearApiKeyInput); + clearApiKeyGroup.appendChild(clearApiKeyLabel); + clearApiKeyInput.onchange = () => { + if (clearApiKeyInput.checked) apiKeyGroup.input.value = ""; + apiKeyGroup.input.disabled = clearApiKeyInput.checked; + }; + form.appendChild(clearApiKeyGroup); + + const jobsPerRunGroup = ui.createFormGroup("Jobs per run:", worker.jobs_per_run ?? 1, `jobs-per-run-${worker.id}`, "number"); + jobsPerRunGroup.group.id = `jobs-per-run-group-${worker.id}`; + jobsPerRunGroup.input.min = "1"; + jobsPerRunGroup.input.max = "16"; + jobsPerRunGroup.input.step = "1"; + jobsPerRunGroup.input.title = "Maximum jobs submitted in one run"; + jobsPerRunGroup.group.style.display = "none"; + form.appendChild(jobsPerRunGroup.group); + + const concurrencyGroup = ui.createFormGroup("Concurrent jobs (slots):", worker.concurrency ?? 1, `concurrency-${worker.id}`, "number"); + concurrencyGroup.group.id = `concurrency-group-${worker.id}`; + concurrencyGroup.input.min = "1"; + concurrencyGroup.input.max = "16"; + concurrencyGroup.input.step = "1"; + concurrencyGroup.input.title = "Maximum concurrent jobs per serverless worker"; + concurrencyGroup.group.style.display = "none"; + form.appendChild(concurrencyGroup.group); + + const timeoutGroup = ui.createFormGroup("Timeout (seconds):", worker.timeout_seconds ?? 1800, `timeout-seconds-${worker.id}`, "number"); + timeoutGroup.group.id = `timeout-seconds-group-${worker.id}`; + timeoutGroup.input.min = "30"; + timeoutGroup.input.max = "86400"; + timeoutGroup.input.step = "1"; + timeoutGroup.group.style.display = "none"; + form.appendChild(timeoutGroup.group); + + const testServerlessButton = ui.createButton("Test Saved Endpoint", () => testServerlessWorker(extension, worker.id)); + testServerlessButton.id = `test-serverless-${worker.id}`; + testServerlessButton.title = "Test the saved Runpod endpoint without submitting inference"; + testServerlessButton.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; + const testServerlessGroup = document.createElement("div"); + testServerlessGroup.id = `serverless-test-group-${worker.id}`; + testServerlessGroup.style.cssText = "display: none;"; + testServerlessGroup.appendChild(testServerlessButton); + form.appendChild(testServerlessGroup); + const saveBtn = ui.createButton("Save", () => saveWorkerSettings(extension, worker.id), "background-color: #4a7c4a;"); saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; @@ -128,7 +226,9 @@ export function createWorkerSettingsForm(ui, extension, worker) { typeSelect.onchange = createOnChangeHandler(); - if (worker.type === "cloud") { + if (worker.type === "serverless") { + typeSelect.value = "serverless"; + } else if (worker.type === "cloud") { typeSelect.value = "cloud"; runpodText.style.display = "block"; } else if (isRemoteWorker(extension, worker)) { diff --git a/web/workerLifecycle.js b/web/workerLifecycle.js index 8acdc2e..ffeaeda 100644 --- a/web/workerLifecycle.js +++ b/web/workerLifecycle.js @@ -41,6 +41,28 @@ function setButtonVisibility(button, visible) { button.style.display = visible ? "" : "none"; } +function updateServerlessStatusMessage(workerId, status) { + const messageElement = document.getElementById(`serverless-status-${workerId}`); + if (!messageElement) return; + + const message = typeof status?.message === "string" && status.message + ? status.message + : status?.ok ? "Ready to scale" : "Runpod status unavailable"; + const queueRemaining = Number(status?.queueRemaining); + const parts = [message]; + if (Number.isFinite(queueRemaining)) { + parts.push(`${Math.max(0, queueRemaining)} queued`); + } + const jobStatuses = (Array.isArray(status?.jobs) ? status.jobs : []) + .slice(0, 2) + .map((job) => String(job?.status || "").trim()) + .filter(Boolean); + if (jobStatuses.length) { + parts.push(`Jobs: ${jobStatuses.join(", ")}`); + } + messageElement.textContent = parts.join(" • "); +} + export async function checkAllWorkerStatuses(extension) { if (_statusCheckRunning || !extension.panelElement) { return; @@ -136,11 +158,60 @@ export async function checkMasterStatus(extension) { // Helper to build worker URL export function getWorkerUrl(extension, worker, endpoint = "") { + if (worker?.type === "serverless") { + return null; + } return buildWorkerUrl(worker, endpoint, window.location); } export async function checkWorkerStatus(extension, worker) { // Assume caller ensured enabled; proceed with check + if (worker.type === "serverless") { + try { + const result = await extension.api.getServerlessStatus(worker.id); + const queueRemaining = Number(result?.queueRemaining); + const normalizedQueue = Number.isFinite(queueRemaining) ? Math.max(0, queueRemaining) : 0; + const normalizedStatus = { + ok: Boolean(result?.ok), + queueRemaining: normalizedQueue, + message: typeof result?.message === "string" ? result.message : "", + jobs: (Array.isArray(result?.jobs) ? result.jobs : []).map((job) => ({ + job_id: String(job?.job_id || ""), + status: String(job?.status || ""), + worker_id: String(job?.worker_id || ""), + execution_id: String(job?.execution_id || ""), + })), + }; + const hasActiveJobs = normalizedStatus.jobs.some((job) => + /^(waiting|submitting|in_queue|queued|pending|running|processing|starting|in_progress|downloading|cancel_requested)$/i.test(job.status.trim()) + ); + const status = { + online: normalizedStatus.ok, + processing: normalizedQueue > 0 || hasActiveJobs, + queueCount: normalizedQueue, + }; + extension.state.updateWorker(worker.id, { serverlessStatus: normalizedStatus }); + extension.state.setWorkerStatus(worker.id, status); + updateServerlessStatusMessage(worker.id, normalizedStatus); + const updatedInPlace = extension.updateWorkerCard?.(worker.id, status); + if (!updatedInPlace) { + updateWorkerControls(extension, worker.id); + } + } catch (error) { + if (error?.name === "AbortError") return; + const failedStatus = { ok: false, queueRemaining: 0, message: error?.message || "Runpod status unavailable", jobs: [] }; + extension.state.updateWorker(worker.id, { serverlessStatus: failedStatus }); + extension.state.setWorkerStatus(worker.id, { online: false, processing: false, queueCount: 0 }); + updateServerlessStatusMessage(worker.id, failedStatus); + const updatedInPlace = extension.updateWorkerCard?.(worker.id, extension.state.getWorkerStatus(worker.id)); + if (!updatedInPlace) { + updateWorkerControls(extension, worker.id); + } + extension.log(`Serverless worker ${worker.id} status check failed: ${error.message}`, "debug"); + } + return; + } + const workerUrl = getWorkerUrl(extension, worker); try { @@ -206,6 +277,9 @@ export async function checkWorkerStatus(extension, worker) { export async function launchWorker(extension, workerId) { const worker = extension.config.workers.find((w) => w.id === workerId); + if (!worker || worker.type === "serverless") { + return; + } // If worker is disabled, enable it first if (!worker.enabled) { @@ -273,6 +347,9 @@ export async function launchWorker(extension, workerId) { export async function stopWorker(extension, workerId) { const worker = extension.config.workers.find((w) => w.id === workerId); + if (!worker || worker.type === "serverless") { + return; + } const stopBtn = document.querySelectorAll(`#controls-${workerId} button`)[1]; // Provide immediate feedback @@ -349,8 +426,12 @@ export async function loadManagedWorkers(extension) { try { const result = await extension.api.getManagedWorkers(); - // Check for launching workers + const serverlessWorkerIds = new Set( + (extension.config?.workers || []).filter((worker) => worker.type === "serverless").map((worker) => String(worker.id)) + ); + // Managed-worker lifecycle state is ComfyUI-process-only; serverless is provider-managed. for (const [workerId, info] of Object.entries(result.managed_workers)) { + if (serverlessWorkerIds.has(String(workerId))) continue; extension.state.setWorkerManaged(workerId, info); // If worker is marked as launching, add to launchingWorkers set @@ -386,6 +467,13 @@ export function updateWorkerControls(extension, workerId) { const stopBtn = document.getElementById(`stop-${workerId}`); const logBtn = document.getElementById(`log-${workerId}`); + if (worker.type === "serverless") { + setButtonVisibility(launchBtn, false); + setButtonVisibility(stopBtn, false); + setButtonVisibility(logBtn, false); + return; + } + if (isRemoteWorker(extension, worker)) { setButtonVisibility(launchBtn, false); setButtonVisibility(stopBtn, false); @@ -438,6 +526,9 @@ export function updateWorkerControls(extension, workerId) { export async function viewWorkerLog(extension, workerId, isRemote = false) { const worker = extension.config.workers.find((w) => w.id === workerId); + if (worker?.type === "serverless") { + return; + } const isRemoteLog = isRemote || (worker ? isRemoteWorker(extension, worker) : false); const managedInfo = extension.state.getWorker(workerId).managed; if (!isRemoteLog && !managedInfo?.log_file) { @@ -498,6 +589,7 @@ export async function refreshLog(extension, workerId, silent = false) { try { const worker = extension.config.workers.find((w) => w.id === workerId); + if (worker?.type === "serverless") return; const isRemoteLog = worker ? isRemoteWorker(extension, worker) : false; const data = isRemoteLog ? await extension.api.getRemoteWorkerLog(workerId, 300) diff --git a/web/workerSettings.js b/web/workerSettings.js index 9956c54..9d03d03 100644 --- a/web/workerSettings.js +++ b/web/workerSettings.js @@ -18,7 +18,7 @@ export function isRemoteWorker(extension, worker) { const workerType = String(worker?.type || "").toLowerCase(); // Explicit type always wins over host heuristics. - if (workerType === "cloud" || workerType === "remote") { + if (workerType === "cloud" || workerType === "remote" || workerType === "serverless") { return true; } if (workerType === "local") { @@ -37,6 +37,134 @@ export function isCloudWorker(extension, worker) { return worker.type === "cloud"; } +function showServerlessValidationError(extension, detail) { + extension.app.extensionManager.toast.add({ + severity: "error", + summary: "Validation Error", + detail, + life: 3000, + }); +} + +async function saveServerlessWorkerSettings(extension, worker, workerId, name) { + const endpointId = String(document.getElementById(`endpoint-id-${workerId}`).value || "").trim(); + const apiKey = document.getElementById(`api-key-${workerId}`).value || ""; + const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked); + const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value); + const concurrency = Number(document.getElementById(`concurrency-${workerId}`).value); + const timeoutSeconds = Number(document.getElementById(`timeout-seconds-${workerId}`).value); + + if (!name.trim()) { + showServerlessValidationError(extension, "Worker name is required"); + return; + } + if (!endpointId) { + showServerlessValidationError(extension, "Endpoint ID is required for Runpod serverless workers"); + return; + } + if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) { + showServerlessValidationError(extension, "Jobs per run must be an integer between 1 and 16"); + return; + } + if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 16) { + showServerlessValidationError(extension, "Concurrency must be an integer between 1 and 16"); + return; + } + if (!Number.isInteger(timeoutSeconds) || timeoutSeconds < 30 || timeoutSeconds > 86400) { + showServerlessValidationError(extension, "Timeout must be an integer between 30 and 86400 seconds"); + return; + } + + const updateData = { + name: name.trim(), + type: "serverless", + provider: "runpod", + endpoint_id: endpointId, + ...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}), + ...(clearApiKey ? { clear_api_key: true } : {}), + jobs_per_run: jobsPerRun, + concurrency, + timeout_seconds: timeoutSeconds, + // Never enable a serverless worker as a side effect of saving credentials. + enabled: worker.enabled, + }; + + try { + await extension.api.updateWorker(workerId, updateData); + document.getElementById(`api-key-${workerId}`).value = ""; + + worker.name = name.trim(); + worker.type = "serverless"; + worker.provider = "runpod"; + worker.endpoint_id = endpointId; + worker.jobs_per_run = jobsPerRun; + worker.concurrency = concurrency; + worker.timeout_seconds = timeoutSeconds; + worker.api_key_configured = clearApiKey + ? false + : apiKey.length > 0 || Boolean(worker.api_key_configured); + // The credential is write-only: never copy it into extension config/state. + delete worker.api_key; + delete worker.host; + delete worker.port; + delete worker.cuda_device; + delete worker.extra_args; + + extension.state.updateWorker(workerId, { enabled: worker.enabled }); + emitWorkersChanged(extension); + extension.app.extensionManager.toast.add({ + severity: "success", + summary: "Settings Saved", + detail: `Worker ${worker.name} settings updated`, + life: 3000, + }); + extension.state.setWorkerExpanded(workerId, false); + if (extension.panelElement) { + renderSidebarContent(extension, extension.panelElement); + } + } catch (error) { + extension.app.extensionManager.toast.add({ + severity: "error", + summary: "Save Failed", + detail: error.message, + life: 5000, + }); + } +} + +export async function testServerlessWorker(extension, workerId) { + const button = document.getElementById(`test-serverless-${workerId}`); + const originalText = button?.textContent; + if (button) { + button.disabled = true; + button.textContent = "Testing..."; + } + try { + // This endpoint only checks the saved endpoint configuration; it does not enqueue inference. + const result = await extension.api.testServerlessWorker(workerId); + extension.app.extensionManager.toast.add({ + severity: result?.ok ? "success" : "error", + summary: result?.ok ? "Runpod Endpoint Ready" : "Runpod Endpoint Test Failed", + detail: result?.message || (result?.ok ? "Saved endpoint is reachable" : "Endpoint test failed"), + life: 4000, + }); + return result; + } catch (error) { + extension.app.extensionManager.toast.add({ + severity: "error", + summary: "Runpod Endpoint Test Failed", + detail: error.message, + life: 5000, + }); + return null; + } finally { + if (button) { + button.disabled = false; + button.textContent = originalText; + } + } +} + export async function saveWorkerSettings(extension, workerId) { const worker = extension.config.workers.find((w) => w.id === workerId); if (!worker) { @@ -44,8 +172,12 @@ export async function saveWorkerSettings(extension, workerId) { } // Get form values - const name = document.getElementById(`name-${workerId}`).value; const workerType = document.getElementById(`worker-type-${workerId}`).value; + if (workerType === "serverless") { + await saveServerlessWorkerSettings(extension, worker, workerId, document.getElementById(`name-${workerId}`).value); + return; + } + const name = document.getElementById(`name-${workerId}`).value; const isRemote = workerType === "remote" || workerType === "cloud"; const isCloud = workerType === "cloud"; const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; @@ -209,6 +341,21 @@ export function cancelWorkerSettings(extension, workerId) { document.getElementById(`port-${workerId}`).value = worker.port; document.getElementById(`cuda-${workerId}`).value = worker.cuda_device || 0; document.getElementById(`args-${workerId}`).value = worker.extra_args || ""; + const endpointInput = document.getElementById(`endpoint-id-${workerId}`); + if (endpointInput) endpointInput.value = worker.endpoint_id || ""; + const apiKeyInput = document.getElementById(`api-key-${workerId}`); + if (apiKeyInput) { + apiKeyInput.value = ""; + apiKeyInput.disabled = false; + } + const clearApiKeyInput = document.getElementById(`clear-api-key-${workerId}`); + if (clearApiKeyInput) clearApiKeyInput.checked = false; + const jobsPerRunInput = document.getElementById(`jobs-per-run-${workerId}`); + if (jobsPerRunInput) jobsPerRunInput.value = worker.jobs_per_run ?? 1; + const concurrencyInput = document.getElementById(`concurrency-${workerId}`); + if (concurrencyInput) concurrencyInput.value = worker.concurrency ?? 1; + const timeoutInput = document.getElementById(`timeout-seconds-${workerId}`); + if (timeoutInput) timeoutInput.value = worker.timeout_seconds ?? 1800; // Reset remote checkbox const remoteCheckbox = document.getElementById(`remote-${workerId}`); diff --git a/web/workerUtils.js b/web/workerUtils.js index a7d9e90..47782b1 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -19,12 +19,17 @@ export async function handleWorkerOperation(extension, button, operation, succes setButtonStateClass("btn--working"); try { - const urlsToProcess = extension.enabledWorkers.map(w => ({ - name: w.name, - url: getWorkerUrl(extension, w) + const workersToProcess = extension.enabledWorkers.filter((worker) => + operation.endpoint !== ENDPOINTS.CLEAR_MEMORY || worker.type !== "serverless" + ); + const targets = workersToProcess.map((worker) => ({ + name: worker.name, + worker, + serverless: worker.type === "serverless", + url: worker.type === "serverless" ? null : getWorkerUrl(extension, worker), })); - - if (urlsToProcess.length === 0) { + + if (targets.length === 0) { button.textContent = "No Workers"; setButtonStateClass("btn--error"); setTimeout(() => { @@ -37,14 +42,19 @@ export async function handleWorkerOperation(extension, button, operation, succes return; } - const promises = urlsToProcess.map(target => - fetch(`${target.url}${operation.endpoint}`, { - method: 'POST', + const promises = targets.map((target) => { + if (target.serverless && operation.name === "Interrupt") { + return extension.api.cancelServerlessWorker(target.worker.id) + .then((result) => ({ ok: result?.ok !== false, name: target.name })) + .catch(() => ({ ok: false, name: target.name })); + } + return fetch(`${target.url}${operation.endpoint}`, { + method: 'POST', mode: 'cors' }) .then(response => ({ ok: response.ok, name: target.name })) - .catch(() => ({ ok: false, name: target.name })) - ); + .catch(() => ({ ok: false, name: target.name })); + }); const results = await Promise.all(promises); const failures = results.filter(r => !r.ok); diff --git a/workers/detection.py b/workers/detection.py index ec98731..dc14a7b 100644 --- a/workers/detection.py +++ b/workers/detection.py @@ -10,6 +10,8 @@ from ..utils.logging import debug_log async def is_local_worker(worker_config): """Check if a worker is running on the same machine as the master.""" + if worker_config.get('type') == 'serverless': + return False host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local': return True diff --git a/workers/process_manager.py b/workers/process_manager.py index 823cb5b..e575b36 100644 --- a/workers/process_manager.py +++ b/workers/process_manager.py @@ -22,6 +22,8 @@ class WorkerProcessManager: return self._launch_builder.build_launch_command(worker_config, comfy_root) def launch_worker(self, worker_config, show_window=False): + if worker_config.get("type") == "serverless": + raise ValueError("Serverless pools cannot be launched as local processes") return self._lifecycle.launch_worker(worker_config, show_window=show_window) def stop_worker(self, worker_id): diff --git a/workers/startup.py b/workers/startup.py index d42d226..1f6e6ac 100644 --- a/workers/startup.py +++ b/workers/startup.py @@ -35,6 +35,8 @@ def auto_launch_workers(): launched_count = 0 for worker in workers: + if worker.get('type') == 'serverless': + continue if worker.get('enabled', False): worker_id = worker.get('id') worker_name = worker.get('name', f'Worker {worker_id}')