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:
+63
-90
@@ -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")
|
||||
|
||||
@@ -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"}
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)]]
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()))
|
||||
|
||||
@@ -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))
|
||||
@@ -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,
|
||||
}
|
||||
]
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
@@ -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):
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user