fix: preserve serverless capacity and recovery across job failures

Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage.
This commit is contained in:
Robert Wojciechowski
2026-10-11 06:56:01 +00:00
parent 371517ae14
commit 87c0cda9f6
17 changed files with 924 additions and 162 deletions
+63 -90
View File
@@ -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")
+1 -1
View File
@@ -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"}
+21 -5
View File
@@ -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)
+37 -8
View File
@@ -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
+142 -42
View File
@@ -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)]]
+29
View File
@@ -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)
+42 -2
View File
@@ -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)
+43 -2
View File
@@ -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/<worker_id>`. 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.
+1 -1
View File
@@ -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")
+34
View File
@@ -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()))
+192
View File
@@ -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))
+57
View File
@@ -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,
}
]
+148 -3
View File
@@ -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()
+40
View File
@@ -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")
+1 -5
View File
@@ -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):
+9
View File
@@ -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
+64 -3
View File
@@ -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