Robert Wojciechowski
6b94242110
* feat: add Runpod serverless image collector pools * fix: prevent worker settings clipping and remove preview copy * fix: initialize worker settings before mounting the form * fix: validate complete serverless requests against 10 MB limit * fix: simplify serverless cards and keep status dots circular * fix: match serverless endpoint box to master card * 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. * fix(serverless): disable Clear saved API key while the worker is enabled The checkbox read "Clear saved API key (disable pool first)", but the panel never defines "pool" and saving with it ticked on an enabled worker was rejected by the backend. - Disable the checkbox while the worker is enabled, with a tooltip saying to turn the worker off first. Also disable it when no key is saved. - Use "worker" instead of "pool" in user-facing serverless errors, the serverless guide and the README entry. * feat(serverless): replace Concurrent jobs with a single Images per run setting - Remove the separate Concurrent jobs field; rename Jobs per run to Images per run (1-16) with a tooltip explaining that Runpod's max workers decides how many GPUs run them. - Backend uses a fixed per-worker limit of 16 outstanding jobs, so all images of a run are submitted at once. Unresolved jobs still count against it. The limit is constant, so it never has to change while jobs are unresolved. - Legacy 'concurrency' in saved configs and older browser tabs is accepted and ignored/dropped; it no longer triggers the busy guard. - Revert the enable toggle and show the reason when the backend rejects it. - Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer endpoint-ID error otherwise. - Confirm before a type change deletes a saved Runpod API key. - Update docs and tests.
381 lines
20 KiB
Python
381 lines
20 KiB
Python
"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
|
|
from .serverless_ledger import JobLedger
|
|
from .serverless_state import (
|
|
PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
|
|
occupies_capacity, terminal_provider_status,
|
|
)
|
|
|
|
|
|
# Fixed per-worker cap on this master's outstanding Runpod jobs; equals the
|
|
# Images per run maximum (MAX_JOBS_PER_RUN in utils/serverless_config.py), so
|
|
# every image of a run is submitted at once. Deliberately not configurable or
|
|
# derived from jobs_per_run: unresolved jobs from earlier runs (ambiguous
|
|
# submissions, unconfirmed cancels) still count against it, and a constant
|
|
# never has to change while such jobs exist.
|
|
PER_WORKER_JOB_LIMIT = 16
|
|
|
|
|
|
async def decode_results(output, manifest):
|
|
from .serverless_media import collect_images
|
|
import torch
|
|
images = await collect_images(output, manifest)
|
|
return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
|
|
|
|
|
|
def state_directory():
|
|
import folder_paths
|
|
return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
|
|
|
|
|
|
def get_secret_store():
|
|
from ...utils.serverless_config import SecretStore
|
|
return SecretStore(state_directory() / "credentials.json")
|
|
|
|
|
|
def get_manager():
|
|
import server
|
|
from ...utils.network import get_client_session
|
|
ps = server.PromptServer.instance
|
|
manager = getattr(ps, "distributed_serverless", None)
|
|
if manager is None:
|
|
manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
|
|
ps.distributed_serverless = manager
|
|
return manager
|
|
|
|
|
|
class ServerlessManager:
|
|
def __init__(self, prompt_server, secrets, ledger_path, session_factory):
|
|
self.server = prompt_server
|
|
self.secrets = secrets
|
|
self.ledger = JobLedger(ledger_path)
|
|
self.session_factory = session_factory
|
|
self.tasks = {}
|
|
self.runs = {}
|
|
self.pool_limits = {}
|
|
self.health_cache = {}
|
|
self.health_locks = {}
|
|
self.poll_seconds = 2.0
|
|
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 < PER_WORKER_JOB_LIMIT:
|
|
# 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"])
|
|
key = self.secrets.get(pool_id)
|
|
return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
|
|
|
|
async def health(self, worker, *, fresh=False):
|
|
pool_id = worker.get("pool_id", worker["id"])
|
|
cache_key = (pool_id, worker["endpoint_id"])
|
|
lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
|
|
async with lock:
|
|
cached = self.health_cache.get(cache_key)
|
|
if cached and not fresh and time.monotonic() - cached[0] < 10:
|
|
return cached[1]
|
|
client = await self._client(worker)
|
|
data = await client.health()
|
|
jobs = data.get("jobs", {})
|
|
if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
|
|
raise RunpodError("Runpod returned an invalid endpoint health record")
|
|
try:
|
|
pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
|
|
except (ValueError, TypeError):
|
|
raise RunpodError("Runpod returned invalid queue counts") from None
|
|
result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
|
|
self.health_cache[cache_key] = (time.monotonic(), result)
|
|
return result
|
|
|
|
def _semaphore(self, worker):
|
|
# The limit is a constant, so it never has to be swapped while this
|
|
# worker still has unresolved jobs.
|
|
pool_id = worker.get("pool_id", worker["id"])
|
|
semaphore = self.pool_limits.get(pool_id)
|
|
if semaphore is None:
|
|
semaphore = asyncio.Semaphore(PER_WORKER_JOB_LIMIT)
|
|
self.pool_limits[pool_id] = semaphore
|
|
return semaphore
|
|
|
|
def ensure_capacity(self, count):
|
|
if sum(not task.done() for task in self.tasks.values()) + count > 64:
|
|
raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
|
|
|
|
async def start(self, execution_id, prompt_id, prepared_jobs):
|
|
self.ensure_capacity(len(prepared_jobs))
|
|
new_runs = []
|
|
# Journal every local slot before starting any billed submission.
|
|
for prepared in prepared_jobs:
|
|
worker = prepared["worker"]
|
|
semaphore = self._semaphore(worker)
|
|
record_id = uuid.uuid4().hex
|
|
now = time.time()
|
|
self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
|
|
worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
|
|
endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
|
|
run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
|
|
provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
|
|
self.runs[record_id] = run
|
|
new_runs.append((record_id, run, semaphore))
|
|
for record_id, run, semaphore in new_runs:
|
|
self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
|
|
# Bound in-memory completed tracking as well as the on-disk journal.
|
|
done = [rid for rid, task in self.tasks.items() if task.done()]
|
|
for rid in done[:-200]:
|
|
self.tasks.pop(rid, None)
|
|
self.runs.pop(rid, None)
|
|
|
|
def is_pending(self, collector_job_id, participant_id):
|
|
return any(not task.done() and run["worker"]["id"] == participant_id and
|
|
any(m["job_id"] == collector_job_id for m in run["manifest"])
|
|
for rid, task in self.tasks.items() for run in [self.runs[rid]])
|
|
|
|
def owns(self, collector_job_id, participant_id):
|
|
return any(run["worker"]["id"] == participant_id and
|
|
any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
|
|
|
|
async def _deliver(self, run, images=None, error=None):
|
|
async with self.server.distributed_jobs_lock:
|
|
for binding in run["manifest"]:
|
|
queue = self.server.distributed_pending_jobs.get(binding["job_id"])
|
|
if queue is None:
|
|
continue # A cancelled/finished collector must never be resurrected.
|
|
if error:
|
|
queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
|
|
else:
|
|
batch = images[binding["job_id"]]
|
|
for index, tensor in enumerate(batch):
|
|
queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
|
|
"tensor": tensor, "is_last": index == len(batch) - 1})
|
|
|
|
async def _cancel_provider(self, run):
|
|
if run["provider_job_id"] and run["client"]:
|
|
try:
|
|
result = await run["client"].cancel(run["provider_job_id"])
|
|
return terminal_provider_status(result)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
def _check_master(self, run):
|
|
queue = getattr(self.server, "prompt_queue", None)
|
|
if queue is None:
|
|
return
|
|
history = queue.get_history(run["prompt_id"])
|
|
record = history.get(run["prompt_id"], {})
|
|
if record.get("status", {}).get("status_str") == "error":
|
|
raise RunpodError("Master workflow failed; serverless work stopped")
|
|
if not record and hasattr(queue, "get_current_queue"):
|
|
running, pending = queue.get_current_queue()
|
|
if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
|
|
raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
|
|
|
|
async def _run(self, run, semaphore):
|
|
rid = run["record_id"]
|
|
submitted = False
|
|
provider_terminal = False
|
|
try:
|
|
async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
|
|
async with semaphore:
|
|
await self._reserve_provider_capacity(run)
|
|
client = await self._client(run["worker"])
|
|
run["client"] = client
|
|
submitted = True
|
|
job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
|
|
run["provider_job_id"] = job_id
|
|
# Persist immediately: all subsequent actions use this exact ID.
|
|
self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
|
|
while True:
|
|
self._check_master(run)
|
|
response = await client.status(job_id)
|
|
status = response.get("status")
|
|
if not isinstance(status, str) or status not in PROVIDER_STATES:
|
|
raise RunpodError("Runpod returned an unknown job state")
|
|
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"):
|
|
raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
|
|
self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
|
|
images = await decode_results(output, run["manifest"])
|
|
await self._deliver(run, images=images)
|
|
self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
|
|
return
|
|
self.ledger.update(rid, status=status, updated_at=time.time())
|
|
if 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:
|
|
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 = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
|
|
error = "Serverless job exceeded its queue/execution deadline"
|
|
else:
|
|
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.
|
|
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 = 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:
|
|
# Completed tracking retains IDs/bindings only, not potentially large
|
|
# base64 inputs, workflow metadata, or a client holding credentials.
|
|
run.pop("workflow", None)
|
|
run.pop("images", None)
|
|
run["client"] = None
|
|
# 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 worker")
|
|
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 = {}
|
|
for rid, task in self.tasks.items():
|
|
run = self.runs[rid]
|
|
if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
|
|
continue
|
|
if execution_id and run["execution_id"] != execution_id:
|
|
continue
|
|
if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
|
|
continue
|
|
if not task.done():
|
|
# 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.values(), return_exceptions=True)
|
|
recovered = 0
|
|
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):
|
|
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)]]
|