diff --git a/api/config_routes.py b/api/config_routes.py index f5cecd0..882d380 100644 --- a/api/config_routes.py +++ b/api/config_routes.py @@ -25,7 +25,12 @@ 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 +from ..utils.secret_fields import is_secret_field_name +from ..utils.serverless_config import ( + ServerlessSecretChanges, + public_config, + validate_serverless_worker, +) def _positive_int(value): @@ -107,16 +112,11 @@ 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 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( @@ -140,20 +140,7 @@ def _serverless_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 + return get_manager().pool_is_active(worker_id) def _ensure_pool_idle(worker_id): @@ -181,7 +168,7 @@ def _strip_sensitive_fields(value: Any) -> Any: return { key: _strip_sensitive_fields(item) for key, item in value.items() - if not (isinstance(key, str) and _is_secret_field_name(key)) + if not is_secret_field_name(key) } if isinstance(value, list): return [_strip_sensitive_fields(item) for item in value] @@ -239,16 +226,6 @@ def _bulk_serverless_changes(old_workers, new_workers): 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): try: @@ -308,43 +285,39 @@ async def update_config_endpoint(request): }, status=400) store = None - changed_secrets = [] - secret_changes_applied = False + response_config = None 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): + with ServerlessSecretChanges() as secret_changes: + async with config_transaction() as config: + if "workers" in validated_root: + old_workers = config.get("workers", []) + validated_root["workers"] = _clean_bulk_type_conversions( + old_workers, validated_root["workers"] + ) + changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"]) + for worker in validated_root["workers"]: + if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False): + store = store or _serverless_secret_store() + if worker["id"] in changed_ids or not store.has(worker["id"]): + return _error_response("Save an API key before enabling this pool", 400) + if changed_ids: 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 + for worker_id in sorted(changed_ids): + prior_secret = secret_changes.capture(store, worker_id) + if prior_secret is not None: + secret_changes.delete(store, worker_id) - settings = config.setdefault("settings", {}) - settings.update(validated_settings) - for key, value in validated_root.items(): - config[key] = value - response_config = public_config(config, store or _serverless_secret_store()) - - return web.json_response({"status": "success", "config": response_config}) + settings = config.setdefault("settings", {}) + settings.update(validated_settings) + for key, value in validated_root.items(): + config[key] = value + response_config = public_config(config, store or _serverless_secret_store()) + response = web.json_response({"status": "success", "config": response_config}) + secret_changes.commit() + return response except _ActivePoolCheckError: - return _error_response("Wait for active serverless jobs before changing this pool", 409) + return _error_response("Finish or reconcile unresolved 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) @@ -368,9 +341,19 @@ async def queue_status_endpoint(request): @server.PromptServer.instance.routes.post("/distributed/config/update_worker") async def update_worker_endpoint(request): + return await _run_with_serverless_secret_changes(_update_worker_endpoint, request) + + +async def _run_with_serverless_secret_changes(endpoint, request): + with ServerlessSecretChanges() as secret_changes: + response = await endpoint(request, secret_changes) + if response.status < 400: + secret_changes.commit() + return response + + +async def _update_worker_endpoint(request, secret_changes): store = None - prior_secret = None - secret_changed = False worker_id = None try: try: @@ -429,7 +412,7 @@ async def update_worker_endpoint(request): return _error_response("name must be a string", 400) store = _serverless_secret_store() - prior_secret = store.get(worker_id) + prior_secret = secret_changes.capture(store, worker_id) converting_to_serverless = current_type != "serverless" next_secret = None if converting_to_serverless else prior_secret if supplied_key: @@ -455,10 +438,9 @@ async def update_worker_endpoint(request): if next_secret != prior_secret: if next_secret is None: - store.delete(worker_id) + secret_changes.delete(store, worker_id) else: - store.set(worker_id, next_secret) - secret_changed = True + secret_changes.set(store, worker_id, next_secret) if existing is None: workers.append(validated) @@ -476,10 +458,9 @@ async def update_worker_endpoint(request): return _error_response("Invalid worker credentials", 400) _ensure_pool_idle(worker_id) store = _serverless_secret_store() - prior_secret = store.get(worker_id) + prior_secret = secret_changes.capture(store, worker_id) if prior_secret is not None: - store.delete(worker_id) - secret_changed = True + secret_changes.delete(store, worker_id) for field in _SERVERLESS_ONLY_FIELDS: existing.pop(field, None) elif _contains_secret_fields(data): @@ -508,24 +489,20 @@ async def update_worker_endpoint(request): 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) + return _error_response("Finish or reconcile unresolved 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) @server.PromptServer.instance.routes.post("/distributed/config/delete_worker") async def delete_worker_endpoint(request): + return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request) + + +async def _delete_worker_endpoint(request, secret_changes): store = None - prior_secret = None - secret_changed = False worker_id = None try: try: @@ -547,10 +524,10 @@ async def delete_worker_endpoint(request): if _worker_type(removed) == "serverless": _ensure_pool_idle(worker_id) store = _serverless_secret_store() - prior_secret = store.get(str(worker_id)) + secret_worker_id = str(worker_id) + prior_secret = secret_changes.capture(store, secret_worker_id) if prior_secret is not None: - store.delete(str(worker_id)) - secret_changed = True + secret_changes.delete(store, secret_worker_id) workers.remove(removed) result = { "status": "success", @@ -560,12 +537,8 @@ async def delete_worker_endpoint(request): 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) + return _error_response("Finish or reconcile unresolved 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") diff --git a/api/orchestration/serverless_ledger.py b/api/orchestration/serverless_ledger.py index 5f96020..932815c 100644 --- a/api/orchestration/serverless_ledger.py +++ b/api/orchestration/serverless_ledger.py @@ -5,7 +5,7 @@ from pathlib import Path import tempfile import threading -TERMINAL = {"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"} +from .serverless_state import TERMINAL FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"} diff --git a/api/orchestration/serverless_media.py b/api/orchestration/serverless_media.py index 4bf79c1..a8b8091 100644 --- a/api/orchestration/serverless_media.py +++ b/api/orchestration/serverless_media.py @@ -491,8 +491,8 @@ def _decode_image_array(payload, total_array_bytes): 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.""" +def _image_output(output): + """Unwrap handler output and reject explicit execution errors.""" if not isinstance(output, dict): raise ServerlessMediaError("Serverless job output is invalid") nested = output.get("output") @@ -501,9 +501,11 @@ async def collect_images(output: dict, manifest: list, download=None) -> dict[st if candidate.get("error") or candidate.get("errors"): raise ServerlessMediaError("Serverless job reported an execution error") - expected = _validated_manifest(manifest) - if not expected: - return {} + return data + + +async def _select_image_payloads(data, expected, download): + """Match artifact identities and fetch bounded bytes, without image decoding.""" artifacts = data.get("images", []) if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES: raise ServerlessMediaError("Serverless job returned too many image files") @@ -565,7 +567,11 @@ async def collect_images(output: dict, manifest: list, download=None) -> dict[st if any(not selected[item["prefix"]] for item in expected): raise ServerlessMediaError("Serverless job is missing an expected collector image") + return selected + +def _decode_image_groups(selected, expected): + """Decode in stable counter order, bounding array memory across collectors.""" result = {} total_array_bytes = 0 for manifest_entry in expected: @@ -583,3 +589,13 @@ async def collect_images(output: dict, manifest: list, download=None) -> dict[st arrays.append(array) result.setdefault(manifest_entry["job_id"], []).extend(arrays) return result + + +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.""" + data = _image_output(output) + expected = _validated_manifest(manifest) + if not expected: + return {} + selected = await _select_image_payloads(data, expected, download) + return _decode_image_groups(selected, expected) diff --git a/api/orchestration/serverless_prompt.py b/api/orchestration/serverless_prompt.py index e45ec6c..1168a56 100644 --- a/api/orchestration/serverless_prompt.py +++ b/api/orchestration/serverless_prompt.py @@ -115,14 +115,8 @@ def _model_name_result(text): 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. - """ +def _validate_prepared_prompt(prepared_prompt): + """Validate graph shape, links, collectors, and supported topology.""" if not isinstance(prepared_prompt, dict): raise ValueError("prepared prompt must be an API graph mapping") @@ -207,6 +201,11 @@ def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]: if _is_link(value): pending.append(str(value[0])) + return source_nodes, collector_ids + + +def _make_scalar_resolver(source_nodes): + """Build pure resolvers for distributed scalar nodes and stock primitives.""" # 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): @@ -309,6 +308,11 @@ def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]: resolved_distributed[cache_key] = copy.deepcopy(result) return result + return resolve_distributed + + +def _build_savers(source_nodes, collector_ids): + """Construct stable SaveImage nodes and their collector metadata.""" # 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 = [] @@ -354,6 +358,11 @@ def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]: "_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"}, } + return collector_specs, saver_nodes + + +def _prune_graph(source_nodes, collector_specs, saver_nodes, resolve_distributed): + """Keep saver ancestors, replacing distributed scalar links with literals.""" all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()} all_nodes.update(copy.deepcopy(saver_nodes)) needed = set() @@ -384,6 +393,26 @@ def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]: if any("distributed" in node["class_type"].lower() for node in compiled.values()): raise ValueError("compiled prompt still contains a distributed dependency") + return compiled + + +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. + """ + source_nodes, collector_ids = _validate_prepared_prompt(prepared_prompt) + resolve_distributed = _make_scalar_resolver(source_nodes) + collector_specs, saver_nodes = _build_savers(source_nodes, collector_ids) + compiled = _prune_graph( + source_nodes, + collector_specs, + saver_nodes, + resolve_distributed, + ) manifest = [ {key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")} for spec in collector_specs diff --git a/api/orchestration/serverless_runtime.py b/api/orchestration/serverless_runtime.py index eeebfff..128ef4a 100644 --- a/api/orchestration/serverless_runtime.py +++ b/api/orchestration/serverless_runtime.py @@ -5,7 +5,11 @@ import uuid from pathlib import Path from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission -from .serverless_ledger import JobLedger, TERMINAL +from .serverless_ledger import JobLedger +from .serverless_state import ( + PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved, + occupies_capacity, terminal_provider_status, +) async def decode_results(output, manifest): @@ -48,6 +52,48 @@ class ServerlessManager: self.health_cache = {} self.health_locks = {} self.poll_seconds = 2.0 + self.capacity_changed = asyncio.Condition() + self.recovery_lock = asyncio.Lock() + self._recover_local_records() + + def _recover_local_records(self): + # WAITING is committed before admission; it cannot have reached Runpod. + # SUBMITTING is an intent record, so a lost ID must remain reserved. + for record in self.ledger.all(): + if record.get("status") == "WAITING" and not record.get("job_id"): + self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time()) + elif record.get("status") == "SUBMITTING" and not record.get("job_id"): + self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time()) + + def _task_is_active(self, record_id): + task = self.tasks.get(record_id) + return task is not None and not task.done() + + def pool_is_active(self, worker_id): + """Protect both live local work and unresolved pre-restart provider work.""" + for rid, run in self.runs.items(): + pool_id = run["worker"].get("pool_id", run["worker"]["id"]) + if str(pool_id) == str(worker_id) and self._task_is_active(rid): + return True + return any(belongs_to_pool(record, worker_id) and is_unresolved(record) + for record in self.ledger.all()) + + async def _reserve_provider_capacity(self, run): + pool_id = run["worker"].get("pool_id", run["worker"]["id"]) + async with self.capacity_changed: + while True: + self._check_master(run) + outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record) + for record in self.ledger.all()) + if outstanding < run["worker"]["concurrency"]: + # Claim atomically before any await that could submit a billed job. + self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time()) + return + await self.capacity_changed.wait() + + async def _notify_capacity(self): + async with self.capacity_changed: + self.capacity_changed.notify_all() async def _client(self, worker): pool_id = worker.get("pool_id", worker["id"]) @@ -83,7 +129,7 @@ class ServerlessManager: 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()): + if self.pool_is_active(pool_id): raise ValueError("Wait for this serverless pool's current jobs before changing concurrency") current = (limit, asyncio.Semaphore(limit)) self.pool_limits[pool_id] = current @@ -144,10 +190,10 @@ class ServerlessManager: 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" + return terminal_provider_status(result) except Exception: - return False - return False + return None + return None def _check_master(self, run): queue = getattr(self.server, "prompt_queue", None) @@ -169,10 +215,9 @@ class ServerlessManager: try: async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())): async with semaphore: - self._check_master(run) + await self._reserve_provider_capacity(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 @@ -182,9 +227,9 @@ class ServerlessManager: 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"}: + if not isinstance(status, str) or status not in PROVIDER_STATES: raise RunpodError("Runpod returned an unknown job state") - provider_terminal = status in TERMINAL + provider_terminal = status in PROVIDER_TERMINAL if status == "COMPLETED": output = response.get("output") if not isinstance(output, dict) or output.get("error") or output.get("errors"): @@ -195,30 +240,30 @@ class ServerlessManager: 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: + if provider_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) + terminal = "FAILED" if provider_terminal else 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" + status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT") 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" + status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED") + error = "Serverless work stopped" if terminal 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 + terminal = "FAILED" if provider_terminal else await self._cancel_provider(run) 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" + status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED") self.ledger.update(rid, status=status, error=error, updated_at=time.time()) await self._deliver(run, error=error) finally: @@ -227,13 +272,61 @@ class ServerlessManager: run.pop("workflow", None) run.pop("images", None) run["client"] = None + # A released local semaphore alone cannot release a durable reservation. + await self._notify_capacity() 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) + def _detached_records(self, worker_id=None, execution_id=None): + return [record for record in self.ledger.all() + if is_unresolved(record) and not self._task_is_active(record["record_id"]) + and (worker_id is None or belongs_to_pool(record, worker_id)) + and (execution_id is None or record.get("execution_id") == execution_id)] + + async def _recover_provider_record(self, record, *, cancel=False): + """Reconcile an exact saved ID, never submit or resume a workflow.""" + if not record.get("job_id"): + return False + status = None + try: + client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]}) + response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"])) + status = terminal_provider_status(response) + except Exception: + pass # Unreachable/expired is not proof of provider termination. + if status is not None or cancel: + self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time()) + return status is not None + + async def reconcile(self, worker_id=None): + """Read terminal states for detached jobs and release only proven capacity.""" + resolved = 0 + async with self.recovery_lock: + for record in self._detached_records(worker_id): + resolved += await self._recover_provider_record(record) + await self._notify_capacity() + return resolved + + async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False): + """Explicit operator escape hatch for lost IDs/expired provider records. + + This records a human's verification, not a cancellation or provider + status. It is never called automatically or allowed for live tasks. + """ + if provider_checked is not True: + raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it") + async with self.recovery_lock: + record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None) + if record is None: + raise ValueError("Unresolved inactive job not found for this pool") + self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(), + error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed") + await self._notify_capacity() + async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None): - targets = [] + 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: @@ -243,33 +336,40 @@ class ServerlessManager: 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) + # Repeated requests must not interrupt the first cancellation's cleanup. + if not task.cancelling(): + task.cancel() + targets[rid] = 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. + await asyncio.gather(*targets.values(), return_exceptions=True) 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 + async with self.recovery_lock: + for record in self._detached_records(worker_id, execution_id): + rid = record["record_id"] + # A task cancelled before its first step never enters _run's handlers. + # A failed start() may also leave a journalled but unstarted local slot. + if record.get("status") == "WAITING" and not record.get("job_id"): + run = self.runs.get(rid) + if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])): + self.ledger.update(rid, status="CANCELLED", updated_at=time.time()) + run.pop("workflow", None) + run.pop("images", None) + await self._deliver(run, error="Serverless job cancelled before submission") + continue + # Do not send two cancellation requests for the same live job in + # one action. Its unresolved reservation survives for later recovery. + if rid in targets or collector_job_id or not record.get("job_id"): + continue + await self._recover_provider_record(record, cancel=True) + recovered += 1 + await self._notify_capacity() 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] + records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)] + # Never hide unresolved IDs behind the bounded completed history. + visible = [record for record in records if is_unresolved(record)] + visible.extend(record for record in records if not is_unresolved(record)) + unresolved_count = sum(is_unresolved(record) for record in records) + return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")} + for record in visible[:max(20, unresolved_count)]] diff --git a/api/orchestration/serverless_state.py b/api/orchestration/serverless_state.py new file mode 100644 index 0000000..2778a74 --- /dev/null +++ b/api/orchestration/serverless_state.py @@ -0,0 +1,29 @@ +"""Shared rules for provider capacity and durable recovery protection. + +A local task ending is not evidence that provider work ended. RESOLVED is an +explicit operator acknowledgment, never a provider-reported state. +""" + +PROVIDER_TERMINAL = frozenset({"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"}) +TERMINAL = PROVIDER_TERMINAL | {"RESOLVED"} +PROVIDER_STATES = PROVIDER_TERMINAL | {"IN_QUEUE", "IN_PROGRESS"} + + +def terminal_provider_status(response): + status = response.get("status") if isinstance(response, dict) else None + return status if isinstance(status, str) and status in PROVIDER_TERMINAL else None + + +def is_unresolved(record): + return record.get("status") not in TERMINAL + + +def occupies_capacity(record): + # WAITING is journalled before admission and cannot have been submitted. + # Every other unresolved state is conservatively reserved, including a + # submission whose provider ID was lost and an unconfirmed cancellation. + return is_unresolved(record) and (record.get("status") != "WAITING" or bool(record.get("job_id"))) + + +def belongs_to_pool(record, worker_id): + return str(record.get("worker_id")) == str(worker_id) diff --git a/api/serverless_routes.py b/api/serverless_routes.py index eb4ba68..27592df 100644 --- a/api/serverless_routes.py +++ b/api/serverless_routes.py @@ -16,13 +16,15 @@ 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) + jobs = [] try: manager = get_manager() + jobs = manager.public_jobs(worker_id) result = await manager.health(worker, fresh=fresh) - return web.json_response(dict(result, jobs=manager.public_jobs(worker_id))) + return web.json_response(dict(result, jobs=jobs)) 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}) + return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs}) @server.PromptServer.instance.routes.post("/distributed/serverless/test") @@ -56,3 +58,41 @@ async def cancel_serverless_endpoint(request): "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) + + +async def _recover(request, *, acknowledge=False): + try: + data = await request.json() + worker_id = data.get("worker_id") + if not isinstance(worker_id, str) or not worker_id.strip(): + raise ValueError() + record_id = data.get("record_id") + if acknowledge: + if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True: + raise ValueError() + except Exception: + return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, 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() + if acknowledge: + await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True) + result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"} + else: + result = {"resolved": await manager.reconcile(worker_id=worker_id)} + return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id))) + except ValueError: + return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this pool"}, status=409) + except Exception: + return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502) + + +@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile") +async def reconcile_serverless_endpoint(request): + return await _recover(request) + + +@server.PromptServer.instance.routes.post("/distributed/serverless/resolve") +async def resolve_serverless_endpoint(request): + return await _recover(request, acknowledge=True) diff --git a/docs/serverless-workers.md b/docs/serverless-workers.md index e0b87be..b3abd2f 100644 --- a/docs/serverless-workers.md +++ b/docs/serverless-workers.md @@ -73,8 +73,10 @@ incomplete serverless batch. - **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. +- **Concurrent jobs** bounds this master's outstanding submissions per pool, + including retained jobs after a restart. Ambiguous submissions and unconfirmed + cancellations keep their capacity reservation until reconciled. 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. @@ -101,6 +103,45 @@ 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. +Unresolved records protect the pool's endpoint, concurrency, API key, worker type, +and deletion even when no local task remains. Disabling a pool is still allowed; +it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are +closed locally, while interrupted `SUBMITTING` records without IDs become +`UNKNOWN` and remain protected. + +### Reconcile or acknowledge recovery + +Recovery uses the same trusted-master API as the panel; it never resubmits a job. + +1. Read `GET /distributed/serverless/status/`. The `jobs` list includes + every unresolved record with its local `record_id`, provider `job_id` (if + received), and state. These records remain visible if endpoint health fails. +2. Request cancellation from the panel, or send + `POST /distributed/serverless/reconcile` with the body below to **read** the + retained provider IDs. Only a confirmed terminal state releases protection. + Network errors, expired IDs, or ambiguous responses do not. + + ```json + {"worker_id": "YOUR_POOL_ID"} + ``` + +3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod + manually first. **Only after verifying that the job is no longer running or + queued**, send `POST /distributed/serverless/resolve` for that one record: + + ```json + {"worker_id": "YOUR_POOL_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true} + ``` + + This records `RESOLVED`: an explicit operator acknowledgment, **not** a + provider-confirmed cancellation. It cannot resolve a live local task and does + not stop remote compute. It releases the record's reservation and may let + already-waiting jobs proceed. Never automate this acknowledgment as a retry + mechanism; an unverified job can still be billable. + +Pool changes become available once all its local tasks and unresolved records +are finished. No credentials or ledger files need to be edited by hand. + ## First-version limits - Image collectors only. Audio/video transfer and audio collectors are rejected. diff --git a/tests/test_config.py b/tests/test_config.py index f5eaa0e..0af47a3 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -18,7 +18,7 @@ def _load_config_module(): del sys.modules[mod_name] root_pkg = types.ModuleType(package_name) - root_pkg.__path__ = [] + root_pkg.__path__ = [str(module_path.parent)] sys.modules[package_name] = root_pkg logging_module = types.ModuleType(f"{package_name}.logging") diff --git a/tests/test_serverless_config.py b/tests/test_serverless_config.py index 4048f94..4bdad8c 100644 --- a/tests/test_serverless_config.py +++ b/tests/test_serverless_config.py @@ -5,6 +5,7 @@ import stat import sys import tempfile import threading +import types import unittest from pathlib import Path from unittest import mock @@ -17,6 +18,11 @@ def _load_serverless_config_module(*, without_fcntl=False): if without_fcntl else "dist_serverless_config_test_module" ) + package_name = module_name + "_pkg" + package = types.ModuleType(package_name) + package.__path__ = [str(module_path.parent)] + sys.modules[package_name] = package + module_name = package_name + ".serverless_config" 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 @@ -320,6 +326,34 @@ class PublicConfigTests(unittest.TestCase): self.assertIsNot(result, config) self.assertIsNot(result["workers"], config["workers"]) + def test_redacts_normalized_secret_field_aliases(self): + config = { + "workers": [{ + "id": "serverless-1", + "type": "serverless", + "apiKey": "camel-api-key", + "APIKEY": "upper-api-key", + "clearApiKey": "camel-clear-key", + "CLEAR_API_KEY": "upper-clear-key", + "runpod_CREDENTIAL": "provider-credential", + "nested": {"credential_note": "nested-credential"}, + "safe": "preserved", + }] + } + + result = serverless_config.public_config( + config, serverless_config.SecretStore(self.root_path()) + ) + + worker = result["workers"][0] + self.assertEqual(worker, { + "id": "serverless-1", + "type": "serverless", + "nested": {}, + "safe": "preserved", + "api_key_configured": False, + }) + 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())) diff --git a/tests/test_serverless_lifecycle.py b/tests/test_serverless_lifecycle.py new file mode 100644 index 0000000..2bc80fc --- /dev/null +++ b/tests/test_serverless_lifecycle.py @@ -0,0 +1,192 @@ +"""Billing-safe lifecycle regressions: local completion is not remote completion.""" +import asyncio +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + +from test_serverless_runtime import runtime, Secrets + + +class LifecycleTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock()) + self.path = Path(self.temp.name) / "jobs.json" + self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object())) + self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"} + self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]} + self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"})) + self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api) + self.client_patch.start() + self.addCleanup(self.client_patch.stop) + self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]})) + self.decode_patch.start() + self.addCleanup(self.decode_patch.stop) + + async def asyncTearDown(self): + await self.manager.cancel() + + def second(self, **changes): + return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes)) + + def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"): + self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id, + status=status, created_at=1, updated_at=1) + + async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self): + cancel_started, release_cancel = asyncio.Event(), asyncio.Event() + self.api.submit.side_effect = ["job-1", "job-2"] + self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}] + async def cancel(_job_id): + cancel_started.set() + await release_cancel.wait() + return {"status": "CANCELLED"} + self.api.cancel.side_effect = cancel + await self.manager.start("exec", "prompt", [self.prepared, self.second()]) + try: + await asyncio.wait_for(cancel_started.wait(), 1) + await asyncio.sleep(0.01) + self.assertEqual(self.api.submit.await_count, 1) + finally: + release_cancel.set() + await self.manager.wait_for_execution("exec") + self.assertEqual(self.api.submit.await_count, 2) + + async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self): + self.api.status.side_effect = runtime.RunpodError("poll failed") + self.api.cancel.return_value = {"status": "IN_PROGRESS"} + await self.manager.start("first", "prompt", [self.prepared]) + await self.manager.wait_for_execution("first") + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED") + self.assertTrue(self.manager.pool_is_active("pool")) + await self.manager.start("second", "prompt", [self.second()]) + try: + await asyncio.sleep(0.01) + self.assertEqual(self.api.submit.await_count, 1) + self.api.status.side_effect = None + self.api.status.return_value = {"status": "COMPLETED", "output": {}} + await self.manager.reconcile(worker_id="pool") + await asyncio.wait_for(self.manager.wait_for_execution("second"), 1) + self.assertEqual(self.api.submit.await_count, 2) + finally: + self.api.cancel.return_value = {"status": "CANCELLED"} + await self.manager.cancel() + + async def test_restarted_manager_reserves_known_unresolved_jobs(self): + self.retained("IN_PROGRESS") + self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object())) + self.assertTrue(self.manager.pool_is_active("pool")) + self.assertFalse(self.manager.pool_is_active("other")) + await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)]) + await self.manager.wait_for_execution("exec") + self.api.submit.assert_not_awaited() + self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT") + + async def test_recovery_closes_never_submitted_waiting_records(self): + self.retained("WAITING", job_id=None) + self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object())) + self.assertFalse(self.manager.pool_is_active("pool")) + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED") + self.api.submit.assert_not_awaited() + + async def test_recovery_keeps_interrupted_submission_unresolved(self): + self.retained("SUBMITTING", job_id=None) + self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object())) + self.assertTrue(self.manager.pool_is_active("pool")) + self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN") + await self.manager.reconcile(worker_id="pool") + self.api.status.assert_not_awaited() + self.api.submit.assert_not_awaited() + + async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self): + self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown") + await self.manager.start("exec", "prompt", [self.prepared]) + await self.manager.wait_for_execution("exec") + record = self.manager.ledger.all()[0] + self.assertTrue(self.manager.pool_is_active("pool")) + with self.assertRaises(ValueError): + await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False) + with self.assertRaises(ValueError): + await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True) + await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True) + self.assertFalse(self.manager.pool_is_active("pool")) + self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED") + self.api.submit.assert_awaited_once() + + async def test_cancel_already_completed_provider_job_is_terminal(self): + self.retained() + self.api.cancel.return_value = {"status": "COMPLETED"} + await self.manager.cancel(worker_id="pool") + self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED") + self.assertFalse(self.manager.pool_is_active("pool")) + + async def test_reconciliation_failure_never_releases_capacity(self): + self.retained() + self.api.status.side_effect = runtime.RunpodError("not reachable") + await self.manager.reconcile(worker_id="pool") + self.assertTrue(self.manager.pool_is_active("pool")) + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED") + + async def test_cancel_before_first_task_step_finalizes_local_record(self): + await self.manager.start("exec", "prompt", [self.prepared]) + await self.manager.cancel(worker_id="pool") + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED") + self.assertFalse(self.server.distributed_pending_jobs["collector"].empty()) + self.assertFalse(self.manager.pool_is_active("pool")) + self.api.submit.assert_not_awaited() + + async def test_live_record_cannot_be_acknowledged_or_reconciled(self): + polling = asyncio.Event() + async def status(_job_id): + polling.set() + await asyncio.Event().wait() + self.api.status.side_effect = status + await self.manager.start("exec", "prompt", [self.prepared]) + await asyncio.wait_for(polling.wait(), 1) + record = self.manager.ledger.all()[0] + with self.assertRaises(ValueError): + await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True) + await self.manager.reconcile(worker_id="pool") + self.assertEqual(self.api.status.await_count, 1) + self.assertTrue(self.manager.pool_is_active("pool")) + + async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self): + polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event() + async def status(_job_id): + polling.set() + await asyncio.Event().wait() + async def cancel(_job_id): + cancelling.set() + await release.wait() + return {"status": "CANCELLED"} + self.api.status.side_effect = status + self.api.cancel.side_effect = cancel + await self.manager.start("exec", "prompt", [self.prepared]) + await polling.wait() + first = asyncio.create_task(self.manager.cancel(worker_id="pool")) + await cancelling.wait() + second = asyncio.create_task(self.manager.cancel(worker_id="pool")) + await asyncio.sleep(0) + release.set() + await asyncio.gather(first, second) + self.api.cancel.assert_awaited_once_with("job") + self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED") + + async def test_partial_start_cleanup_does_not_leave_waiting_protection(self): + with self.assertRaises(ValueError): + await self.manager.start("exec", "prompt", [self.prepared, self.second(concurrency=2)]) + await self.manager.cancel(execution_id="exec") + self.assertFalse(self.manager.pool_is_active("pool")) + self.api.submit.assert_not_awaited() + + async def test_public_jobs_never_hides_unresolved_recovery_ids(self): + for index in range(22): + self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index) + for index in range(5): + self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index) + records = self.manager.public_jobs("pool") + self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)}) + self.assertTrue(all(record["status"] == "UNKNOWN" for record in records)) diff --git a/tests/test_serverless_prompt.py b/tests/test_serverless_prompt.py index a463dfb..d559529 100644 --- a/tests/test_serverless_prompt.py +++ b/tests/test_serverless_prompt.py @@ -333,3 +333,60 @@ def test_resolves_boolean_and_float_primitive_links(): assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5} assert not {"1", "2", "3", "4"}.intersection(graph) + + +def test_compiled_graph_matches_golden_prompt_and_manifest(): + prompt = _prepared_prompt() + prompt["3"]["inputs"]["seed"] = ["10", 0] + prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}} + prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}} + prompt["12"] = { + "class_type": "DistributedValue", + "inputs": { + "default_value": ["11", 0], + "worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}), + "is_worker": True, + "worker_id": "worker_0", + }, + } + prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}} + prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}} + prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0]) + + graph, manifest = _run(prompt) + + expected_prefix = ( + "serverless_job_42_worker_0_4_image_" + "1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244" + ) + assert list(graph) == ["1", "2", "15"] + assert graph == { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {}, + "_meta": {"title": "source"}, + }, + "2": { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "seed": 41, + "cfg": 2.5, + "model_name": "model.safetensors", + }, + }, + "15": { + "class_type": "SaveImage", + "inputs": {"images": ["2", 0], "filename_prefix": expected_prefix}, + "_meta": {"title": "Serverless output: collector 4 (image)"}, + }, + } + assert manifest == [ + { + "collector_id": "4", + "job_id": "job/42", + "node_id": "15", + "kind": "image", + "prefix": expected_prefix, + } + ] diff --git a/tests/test_serverless_routes.py b/tests/test_serverless_routes.py index 1264c3f..6185c26 100644 --- a/tests/test_serverless_routes.py +++ b/tests/test_serverless_routes.py @@ -77,7 +77,7 @@ def _load_modules(): runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime") runtime.secret_store = None - runtime.manager = types.SimpleNamespace(tasks={}, runs={}) + runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False) runtime.get_secret_store = lambda: runtime.secret_store runtime.get_manager = lambda: runtime.manager sys.modules[runtime.__name__] = runtime @@ -117,7 +117,7 @@ class _RunningTask: return False -class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): +class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase): def setUp(self): self.tempdir = tempfile.TemporaryDirectory() self.root = Path(self.tempdir.name) @@ -128,7 +128,7 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): 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={}) + runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False) def tearDown(self): config_module.invalidate_config_cache() @@ -137,6 +137,8 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): def read_config_file(self): return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8")) + +class ServerlessConfigRouteTests(_ConfigRouteFixture): async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self): self.secrets.set("pool-1", "private-api-key") self.config["workers"] = [ @@ -155,6 +157,43 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): self.assertNotIn("private-api-key", json.dumps(response.payload)) self.assertNotIn("inline-api-key", json.dumps(response.payload)) + async def test_get_config_redacts_normalized_secret_field_aliases(self): + aliases = { + "apiKey": "camel-api-key", + "APIKEY": "upper-api-key", + "clearApiKey": "camel-clear-key", + "CLEAR_API_KEY": "upper-clear-key", + "runpod_CREDENTIAL": "provider-credential", + } + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + **aliases, + "nested": {"credential_note": "nested-credential"}, + }] + 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] + for alias, value in aliases.items(): + self.assertNotIn(alias, worker) + self.assertNotIn(value, json.dumps(response.payload)) + self.assertNotIn("credential_note", worker["nested"]) + self.assertNotIn("nested-credential", json.dumps(response.payload)) + + async def test_bulk_config_rejects_normalized_secret_field_aliases(self): + for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"): + with self.subTest(secret_field=secret_field): + secret = f"private-{secret_field}" + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{"id": "pool-1", "type": "local", secret_field: secret}] + })) + self.assertEqual(response.status, 400) + self.assertNotIn(secret, json.dumps(response.payload)) + self.assertEqual(self.read_config_file()["workers"], []) + async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self): secret = "runpod-secret-never-returned" payload = { @@ -326,6 +365,7 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + pool_is_active=lambda worker_id: worker_id == "pool-1", ) enabled = await config_routes.update_worker_endpoint( @@ -376,6 +416,7 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + pool_is_active=lambda worker_id: worker_id == "pool-1", ) credential_change = await config_routes.update_worker_endpoint(_FakeRequest({ @@ -405,6 +446,52 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(self.read_config_file()["workers"], []) self.assertEqual(config_module.load_config()["workers"], []) + async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self): + self.secrets.set("pool-1", "bulk-restore-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + + with patch.object(config_module, "save_config", return_value=False): + response = await config_routes.update_config_endpoint(_FakeRequest({ + "workers": [{"id": "pool-1", "type": "local", "port": 8188}] + })) + + self.assertEqual(response.status, 500) + self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key") + self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") + + async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self): + self.secrets.set("pool-1", "prior-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + + with patch.object(config_module, "save_config", return_value=False): + response = await config_routes.update_worker_endpoint(_FakeRequest({ + "worker_id": "pool-1", "api_key": "replacement-key", + })) + + self.assertEqual(response.status, 500) + self.assertEqual(self.secrets.get("pool-1"), "prior-key") + self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless") + + async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self): + self.secrets.set("pool-1", "delete-restore-key") + self.config["workers"] = [{ + "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", + }] + self.assertTrue(config_module.save_config(self.config)) + + with patch.object(config_module, "save_config", return_value=False): + response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) + + self.assertEqual(response.status, 500) + self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key") + self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1") + 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)) @@ -426,6 +513,7 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): runtime_module.manager = types.SimpleNamespace( tasks={"run-1": _RunningTask()}, runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}}, + pool_is_active=lambda worker_id: worker_id == "pool-1", ) updated = copy.deepcopy(self.config["workers"]) updated[0]["concurrency"] = 2 @@ -459,5 +547,62 @@ class ServerlessConfigRouteTests(unittest.IsolatedAsyncioTestCase): self.assertNotIn(obsolete, worker) +class PersistedPoolGuardTests(_ConfigRouteFixture): + """Use the actual manager/ledger across restart, not a mocked idle result.""" + + def setUp(self): + super().setUp() + from test_serverless_runtime import runtime + from unittest.mock import AsyncMock + self.runtime = runtime + self.session_factory = AsyncMock() + self.ledger_path = self.root / "jobs.json" + self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory) + runtime_module.manager = self.manager + self.secrets.set("pool-1", "saved-key") + self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}] + self.assertTrue(config_module.save_config(self.config)) + + async def test_retained_jobs_block_destructive_changes_after_restart(self): + original = self.read_config_file() + for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"): + for restart in (False, True): + with self.subTest(status=status, restart=restart): + self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id") + if restart: + self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory) + runtime_module.manager = self.manager + else: + self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True) + for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"concurrency": 2}, {"type": "local", "port": 8188}): + response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1"))) + self.assertEqual(response.status, 409, fields) + response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) + self.assertEqual(response.status, 409) + worker = original["workers"][0] + for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")], [dict(worker, concurrency=2)]): + response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers})) + self.assertEqual(response.status, 409, workers) + self.assertEqual(self.read_config_file(), original) + self.assertEqual(self.secrets.get("pool-1"), "saved-key") + self.session_factory.assert_not_awaited() + + async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self): + self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN") + response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"})) + self.assertEqual(response.status, 200) + self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused") + response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True})) + self.assertEqual(response.status, 409) + self.assertEqual(self.secrets.get("pool-1"), "saved-key") + + async def test_terminal_or_other_pool_records_do_not_block_deletion(self): + self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED") + self.manager.ledger.update("other", worker_id="other", status="UNKNOWN") + response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"})) + self.assertEqual(response.status, 200) + self.assertIsNone(self.secrets.get("pool-1")) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_serverless_status.py b/tests/test_serverless_status.py index 544e091..5412a49 100644 --- a/tests/test_serverless_status.py +++ b/tests/test_serverless_status.py @@ -56,3 +56,43 @@ class StatusTests(unittest.IsolatedAsyncioTestCase): result = await self.routes.cancel_serverless_endpoint(request) self.assertEqual(result.status, 200) self.manager.cancel.assert_awaited_once_with(worker_id="pool") + + async def test_reconcile_only_reads_the_requested_pools_jobs(self): + self.manager.reconcile = AsyncMock(return_value=2) + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"})) + result = await self.routes.reconcile_serverless_endpoint(request) + self.assertEqual(result.status, 200) + self.assertEqual(json.loads(result.text)["resolved"], 2) + self.manager.reconcile.assert_awaited_once_with(worker_id="pool") + self.manager.cancel.assert_not_awaited() + + async def test_resolve_requires_explicit_boolean_confirmation_and_record_id(self): + self.manager.acknowledge_recovery = AsyncMock() + for payload in ({"worker_id": "pool", "record_id": "record"}, + {"worker_id": "pool", "record_id": "record", "provider_checked": "true"}, + {"worker_id": "pool", "provider_checked": True}): + with self.subTest(payload=payload): + result = await self.routes.resolve_serverless_endpoint(types.SimpleNamespace(json=AsyncMock(return_value=payload))) + self.assertEqual(result.status, 400) + self.manager.acknowledge_recovery.assert_not_awaited() + + async def test_resolve_is_scoped_and_does_not_claim_provider_cancellation(self): + self.manager.acknowledge_recovery = AsyncMock() + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True})) + result = await self.routes.resolve_serverless_endpoint(request) + self.assertEqual(result.status, 200) + self.manager.acknowledge_recovery.assert_awaited_once_with("pool", "record", provider_checked=True) + self.assertIn("not provider-confirmed", json.loads(result.text)["message"]) + self.manager.cancel.assert_not_awaited() + + async def test_live_or_mismatched_record_refuses_manual_resolution(self): + self.manager.acknowledge_recovery = AsyncMock(side_effect=ValueError("Unresolved inactive job not found for this pool")) + request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True})) + result = await self.routes.resolve_serverless_endpoint(request) + self.assertEqual(result.status, 409) + + async def test_failed_health_still_exposes_recovery_ids(self): + self.manager.health.side_effect = RuntimeError("private transport error") + self.manager.public_jobs = lambda _: [{"record_id": "record", "status": "UNKNOWN"}] + result = await self.routes.serverless_status_endpoint(types.SimpleNamespace(match_info={"worker_id": "pool"})) + self.assertEqual(json.loads(result.text)["jobs"][0]["record_id"], "record") diff --git a/utils/config.py b/utils/config.py index cff821d..324762e 100644 --- a/utils/config.py +++ b/utils/config.py @@ -7,6 +7,7 @@ import os import json from contextlib import asynccontextmanager from .logging import log +from .secret_fields import is_secret_field_name as _is_secret_field_name # Import defaults for timeout fallbacks from .constants import HEARTBEAT_TIMEOUT @@ -31,11 +32,6 @@ def _contains_serverless_secret(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): diff --git a/utils/secret_fields.py b/utils/secret_fields.py new file mode 100644 index 0000000..862c0e7 --- /dev/null +++ b/utils/secret_fields.py @@ -0,0 +1,9 @@ +"""Shared classification for credential-bearing configuration fields.""" + + +def is_secret_field_name(key: object) -> bool: + """Return whether a field name identifies an API key or credential.""" + if not isinstance(key, str): + return False + normalized = "".join(character for character in key.casefold() if character.isalnum()) + return normalized in {"apikey", "clearapikey"} or "credential" in normalized \ No newline at end of file diff --git a/utils/serverless_config.py b/utils/serverless_config.py index e828305..29a8f94 100644 --- a/utils/serverless_config.py +++ b/utils/serverless_config.py @@ -18,6 +18,8 @@ from contextlib import contextmanager from pathlib import Path from typing import Iterator +from .secret_fields import is_secret_field_name + try: import fcntl except ImportError: # Windows and other platforms without POSIX file locking. @@ -41,7 +43,6 @@ _NUMERIC_FIELDS = { "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() @@ -344,12 +345,72 @@ class SecretStore: pass +class ServerlessSecretChanges: + """Apply serverless secret changes and restore them unless committed. + + Use this around the associated config transaction so a failed config save + compensates any secret-store writes or deletions made along the way. + """ + + def __init__(self): + self._originals: dict[tuple[int, str], tuple[SecretStore, str, str | None]] = {} + self._mutations: list[tuple[int, str]] = [] + self._committed = False + + def __enter__(self) -> "ServerlessSecretChanges": + return self + + def __exit__(self, _exc_type, _exc_value, _traceback) -> bool: + if not self._committed: + self.rollback() + return False + + def capture(self, store: SecretStore, worker_id: str) -> str | None: + """Capture and return the original key once per store/worker pair.""" + mutation_id = (id(store), worker_id) + if mutation_id not in self._originals: + self._originals[mutation_id] = (store, worker_id, store.get(worker_id)) + return self._originals[mutation_id][2] + + def set(self, store: SecretStore, worker_id: str, key: str) -> None: + self._prepare_mutation(store, worker_id) + store.set(worker_id, key) + + def delete(self, store: SecretStore, worker_id: str) -> None: + self._prepare_mutation(store, worker_id) + store.delete(worker_id) + + def commit(self) -> None: + """Keep the secret changes after the corresponding config save.""" + self._committed = True + + def rollback(self) -> None: + """Best-effort restoration without masking the original route error.""" + for mutation_id in reversed(self._mutations): + store, worker_id, prior_secret = self._originals[mutation_id] + try: + if prior_secret is None: + store.delete(worker_id) + else: + store.set(worker_id, prior_secret) + except Exception: + pass + self._mutations.clear() + + def _prepare_mutation(self, store: SecretStore, worker_id: str) -> None: + self.capture(store, worker_id) + mutation_id = (id(store), worker_id) + if mutation_id not in self._mutations: + # Record before applying: a store operation may fail after replacing + # its file, so the exception path must still attempt compensation. + self._mutations.append(mutation_id) + + 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: + if is_secret_field_name(key): continue redacted[key] = _redact_credentials(item) return redacted