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.
This commit is contained in:
Robert Wojciechowski
2026-10-13 21:33:02 +00:00
committed by r
parent 7060e31edb
commit 06be0e5dee
13 changed files with 358 additions and 95 deletions
+4 -8
View File
@@ -218,8 +218,7 @@ def _bulk_serverless_changes(old_workers, new_workers):
if new is None or new.get("type") != "serverless":
_ensure_pool_idle(worker_id)
secret_ids.add(worker_id)
elif (old.get("endpoint_id") != new.get("endpoint_id")
or old.get("concurrency", 1) != new.get("concurrency", 1)):
elif old.get("endpoint_id") != new.get("endpoint_id"):
_ensure_pool_idle(worker_id)
elif new is not None and new.get("type") == "serverless":
secret_ids.add(worker_id)
@@ -401,7 +400,8 @@ async def _update_worker_endpoint(request, secret_changes):
candidate["type"] = "serverless"
candidate.setdefault("provider", "runpod")
candidate.setdefault("enabled", False)
for field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "enabled", "name"):
# Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name"):
if field in data:
candidate[field] = data[field]
candidate = _clean_serverless_worker(candidate)
@@ -424,12 +424,8 @@ async def _update_worker_endpoint(request, secret_changes):
current_type == "serverless"
and existing.get("endpoint_id") != validated.get("endpoint_id")
)
concurrency_changed = (
current_type == "serverless"
and existing.get("concurrency", 1) != validated.get("concurrency", 1)
)
if current_type == "serverless" and (
endpoint_changed or concurrency_changed or next_secret != prior_secret
endpoint_changed or next_secret != prior_secret
):
_ensure_pool_idle(worker_id)
+17 -12
View File
@@ -12,6 +12,15 @@ from .serverless_state import (
)
# 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
@@ -85,7 +94,7 @@ class ServerlessManager:
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"]:
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
@@ -122,18 +131,14 @@ class ServerlessManager:
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"])
limit = worker["concurrency"]
current = self.pool_limits.get(pool_id)
if current is None:
current = (limit, asyncio.Semaphore(limit))
self.pool_limits[pool_id] = current
elif current[0] != limit:
if self.pool_is_active(pool_id):
raise ValueError("Wait for this serverless worker's current jobs before changing concurrency")
current = (limit, asyncio.Semaphore(limit))
self.pool_limits[pool_id] = current
return current[1]
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:
+23 -14
View File
@@ -2,14 +2,15 @@
A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
normal ComfyUI API workflow per logical slot and collects the returned images.
normal ComfyUI API workflow per requested image and collects the returned images.
The endpoint workers **do not need ComfyUI-Distributed installed**.
## Setup
1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
2. Enter its name, Runpod **endpoint ID** (not a URL), and API key.
3. Start with **Jobs per run = 1**, **Concurrent jobs = 1**, and **Timeout = 1800**.
2. Enter its name, Runpod **endpoint ID** (not a URL; a pasted `https://api.runpod.ai/v2/<id>/…`
URL is trimmed to the ID), and API key.
3. Start with **Images per run = 1** and **Timeout = 1800**.
4. Save, then choose **Test Saved Endpoint**. This checks authentication and
endpoint health only; it does not start GPU inference or prove model/node compatibility.
5. Enable the worker when you want queued Distributed workflows to use it.
@@ -72,15 +73,23 @@ incomplete serverless batch.
## Scheduling and cancellation
- **Jobs per run** adds logical collector participants. Each receives its own
Distributed Seed/Value overrides. It is not a physical GPU count.
- **Concurrent jobs** bounds this master's outstanding submissions per worker,
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 job slot.
- **Timeout** includes waiting for a local slot, Runpod queue/cold-start time,
execution, and output retrieval. Provider policy limits are submitted as well.
- **Images per run** (1–16) is how many separate Runpod jobs this worker submits
for each queued run. Each returns one image and receives its own Distributed
Seed/Value overrides. All of them are submitted at once.
- How many of those jobs actually run in parallel is decided by the Runpod
endpoint's **max workers**, not by Distributed. Jobs beyond that wait in
Runpod's queue. To run every image at the same time, keep Images per run at or
below the endpoint's max workers.
- Distributed also keeps a fixed limit of 16 outstanding jobs per serverless
worker (and 64 across all of them). Retained jobs after a restart, ambiguous
submissions, and unconfirmed cancellations keep counting against it until
reconciled. Other applications using the same endpoint are unaffected.
- Older configs may contain a `concurrency` field from the previous separate
setting. It is still accepted when loading or saving, but ignored and dropped.
- **Load balance** selects one participant, so it sends one job instead of
Images per run.
- **Timeout** includes Runpod queue/cold-start time, execution, and output
retrieval. Provider policy limits are submitted as well.
- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
- Interrupting a collector or using the panel cancellation action requests
cancellation only for jobs submitted by this master. It never purges the endpoint.
@@ -104,7 +113,7 @@ 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 worker's endpoint, concurrency, API key, worker type,
Unresolved records protect the worker's endpoint, API key, worker type,
and deletion even when no local task remains. Disabling the worker 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
@@ -161,4 +170,4 @@ are finished. No credentials or ledger files need to be edited by hand.
This feature is verified with offline provider/transport tests. A successful
health test is not a paid end-to-end generation test; test your deployed handler
and model inventory with one slot before increasing concurrency.
and model inventory with one image per run before increasing it.
+18 -6
View File
@@ -54,7 +54,6 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
"type": "serverless",
"provider": "runpod",
"jobs_per_run": 1,
"concurrency": 1,
"timeout_seconds": 1800,
},
)
@@ -62,21 +61,35 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
def test_accepts_inclusive_numeric_bounds(self):
for jobs, concurrency, timeout in ((1, 1, 30), (16, 16, 86400)):
with self.subTest(jobs=jobs, concurrency=concurrency, timeout=timeout):
for jobs, timeout in ((1, 30), (16, 86400)):
with self.subTest(jobs=jobs, timeout=timeout):
result = serverless_config.validate_serverless_worker(
{
"id": "worker-1",
"endpoint_id": "ep-123",
"jobs_per_run": jobs,
"concurrency": concurrency,
"timeout_seconds": timeout,
}
)
self.assertEqual(result["jobs_per_run"], jobs)
self.assertEqual(result["concurrency"], concurrency)
self.assertEqual(result["timeout_seconds"], timeout)
def test_legacy_concurrency_is_accepted_and_dropped(self):
# Configs saved before Images per run replaced the separate setting.
for legacy in (1, 2, 16, 17, True, "3"):
with self.subTest(legacy=legacy):
result = serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": legacy}
)
self.assertNotIn("concurrency", result)
self.assertEqual(result["jobs_per_run"], 2)
def test_endpoint_url_is_rejected_with_a_plain_message(self):
with self.assertRaisesRegex(ValueError, "not a URL"):
serverless_config.validate_serverless_worker(
{"id": "worker-1", "endpoint_id": "https://api.runpod.ai/v2/abc123/run"}
)
def test_rejects_missing_or_empty_id_and_endpoint(self):
base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
for field, value in (("id", None), ("id", " "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
@@ -138,7 +151,6 @@ class ValidateServerlessWorkerTests(unittest.TestCase):
base = {"id": "worker-1", "endpoint_id": "ep-123"}
invalid_values = {
"jobs_per_run": (True, 1.0, 0, 17),
"concurrency": (False, 2.5, 0, 17),
"timeout_seconds": (True, 30.0, 29, 86401),
}
for field, values in invalid_values.items():
+9 -3
View File
@@ -16,7 +16,7 @@ class LifecycleTests(unittest.IsolatedAsyncioTestCase):
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.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "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)
@@ -25,6 +25,10 @@ class LifecycleTests(unittest.IsolatedAsyncioTestCase):
self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
self.decode_patch.start()
self.addCleanup(self.decode_patch.stop)
# These regressions exercise capacity holding, so use a limit of one.
self.limit_patch = patch.object(runtime, "PER_WORKER_JOB_LIMIT", 1)
self.limit_patch.start()
self.addCleanup(self.limit_patch.stop)
async def asyncTearDown(self):
await self.manager.cancel()
@@ -176,8 +180,10 @@ class LifecycleTests(unittest.IsolatedAsyncioTestCase):
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)])
broken = self.second()
del broken["worker"]["endpoint_id"]
with self.assertRaises(KeyError):
await self.manager.start("exec", "prompt", [self.prepared, broken])
await self.manager.cancel(execution_id="exec")
self.assertFalse(self.manager.pool_is_active("pool"))
self.api.submit.assert_not_awaited()
+25 -9
View File
@@ -210,7 +210,8 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
saved = self.read_config_file()["workers"][0]
self.assertEqual(saved["endpoint_id"], "ep-123")
self.assertEqual(saved["jobs_per_run"], 3)
self.assertEqual(saved["concurrency"], 2)
# Legacy field from older browser tabs is accepted but not persisted.
self.assertNotIn("concurrency", saved)
self.assertEqual(saved["timeout_seconds"], 600)
for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"):
self.assertNotIn(obsolete, saved)
@@ -233,7 +234,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
self.assertTrue(worker["enabled"])
self.assertEqual(worker["endpoint_id"], "ep-123")
self.assertEqual(worker["jobs_per_run"], 4)
self.assertEqual(worker["concurrency"], 3)
self.assertNotIn("concurrency", worker)
self.assertEqual(worker["timeout_seconds"], 900)
self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
self.assertTrue(response.payload["api_key_configured"])
@@ -242,7 +243,7 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
secret = "must-not-leak"
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "type": "serverless", "provider": "runpod",
"endpoint_id": "ep-123", "concurrency": True, "api_key": secret,
"endpoint_id": "ep-123", "jobs_per_run": True, "api_key": secret,
}))
self.assertEqual(response.status, 400)
@@ -507,8 +508,8 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
self.assertEqual(response.status, 400)
self.assertEqual(self.read_config_file()["workers"], [])
async def test_bulk_concurrency_change_is_blocked_during_active_run(self):
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
async def test_bulk_endpoint_change_is_blocked_during_active_run(self):
self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}]
self.assertTrue(config_module.save_config(self.config))
runtime_module.manager = types.SimpleNamespace(
tasks={"run-1": _RunningTask()},
@@ -516,10 +517,10 @@ class ServerlessConfigRouteTests(_ConfigRouteFixture):
pool_is_active=lambda worker_id: worker_id == "pool-1",
)
updated = copy.deepcopy(self.config["workers"])
updated[0]["concurrency"] = 2
updated[0]["endpoint_id"] = "ep-other"
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
self.assertEqual(response.status, 409)
self.assertEqual(self.read_config_file()["workers"][0]["concurrency"], 1)
self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
async def test_non_serverless_update_behavior_is_preserved(self):
self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}]
@@ -574,19 +575,34 @@ class PersistedPoolGuardTests(_ConfigRouteFixture):
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}):
for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"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)]):
for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")]):
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_jobs_allow_images_per_run_and_legacy_concurrency_saves(self):
# Upgrading from a config with concurrency=1 must not trip the busy guard,
# and the image count stays editable while old jobs are unresolved.
self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN")
response = await config_routes.update_worker_endpoint(_FakeRequest({
"worker_id": "pool-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2,
}))
self.assertEqual(response.status, 200)
saved = self.read_config_file()["workers"][0]
self.assertEqual(saved["jobs_per_run"], 2)
self.assertNotIn("concurrency", saved)
response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [dict(saved, jobs_per_run=3, concurrency=5)]}))
self.assertEqual(response.status, 200)
self.assertEqual(self.read_config_file()["workers"][0]["jobs_per_run"], 3)
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"}))
+30 -3
View File
@@ -34,7 +34,7 @@ class RuntimeTests(unittest.IsolatedAsyncioTestCase):
self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
self.manager.poll_seconds = 0
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "concurrency": 1, "timeout_seconds": 30, "type": "serverless"}
self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "timeout_seconds": 30, "type": "serverless"}
self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
async def asyncTearDown(self):
@@ -90,7 +90,34 @@ class RuntimeTests(unittest.IsolatedAsyncioTestCase):
await self.manager.cancel(worker_id="pool")
api.cancel.assert_awaited_once_with("owned-job")
async def test_per_pool_concurrency_bounds_submission(self):
async def test_all_images_of_a_run_submit_at_once_ignoring_legacy_concurrency(self):
self.assertEqual(runtime.PER_WORKER_JOB_LIMIT, 16)
gate = asyncio.Event()
submitted = []
async def submit(*args):
submitted.append(args)
return f"provider-{len(submitted)}"
async def status(_job):
await gate.wait()
return {"status": "COMPLETED", "output": {"images": []}}
api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
# A config saved before the setting was removed still carries concurrency=1.
legacy = dict(self.worker, concurrency=1)
jobs = [dict(self.prepared, worker=dict(legacy, id=f"pool:{i}")) for i in range(1, 4)]
with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
await self.manager.start("execution", "prompt", jobs)
for _ in range(50):
if len(submitted) == 3:
break
await asyncio.sleep(0)
self.assertEqual(len(submitted), 3)
gate.set()
await self.manager.wait_for_execution("execution")
async def test_per_pool_job_limit_bounds_submission(self):
patcher = patch.object(runtime, "PER_WORKER_JOB_LIMIT", 1)
patcher.start()
self.addCleanup(patcher.stop)
gate = asyncio.Event()
reached = asyncio.Event()
count = 0
@@ -135,7 +162,7 @@ class RuntimeTests(unittest.IsolatedAsyncioTestCase):
self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
# A blocked slot expires without a billable submission.
self.manager.pool_limits["pool"] = (1, asyncio.Semaphore(0))
self.manager.pool_limits["pool"] = asyncio.Semaphore(0)
with patch.object(runtime, "RunpodClient", return_value=api):
await self.manager.start("execution", "prompt", [self.prepared])
await self.manager.wait_for_execution("execution")
+13 -3
View File
@@ -38,11 +38,17 @@ _ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII)
_UNSUPPORTED_SECRET_STORE_MESSAGE = (
"Serverless credential storage is currently supported on POSIX masters only"
)
# Images per run is also this worker's effective parallelism: every image job
# is submitted at once (the runtime's fixed PER_WORKER_JOB_LIMIT equals this),
# and the Runpod endpoint's max workers decides how many actually run together.
MAX_JOBS_PER_RUN = 16
_NUMERIC_FIELDS = {
"jobs_per_run": (1, 16, 1),
"concurrency": (1, 16, 1),
"jobs_per_run": (1, MAX_JOBS_PER_RUN, 1),
"timeout_seconds": (30, 86400, 1800),
}
# Older configs saved a separate concurrency setting. It is accepted so they
# still load, then dropped: PER_WORKER_JOB_LIMIT replaces it.
_LEGACY_IGNORED_FIELDS = ("concurrency",)
_THREAD_LOCKS: dict[str, threading.Lock] = {}
_THREAD_LOCKS_GUARD = threading.Lock()
@@ -67,7 +73,9 @@ def validate_serverless_worker(worker: dict) -> dict:
endpoint_id = result.get("endpoint_id")
if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id):
raise ServerlessConfigError("endpoint_id must be an opaque Runpod identifier")
raise ServerlessConfigError(
"Endpoint ID must be just the Runpod endpoint ID (letters, numbers, - or _), not a URL"
)
if result.get("type", "serverless") != "serverless":
raise ServerlessConfigError("type must be serverless")
@@ -81,6 +89,8 @@ def validate_serverless_worker(worker: dict) -> dict:
if type(value) is not int or not minimum <= value <= maximum:
raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}")
result[field] = value
for field in _LEGACY_IGNORED_FIELDS:
result.pop(field, None)
result["type"] = "serverless"
result["provider"] = "runpod"
+16 -1
View File
@@ -15,7 +15,8 @@ import { parseHostInput, getMasterUrl as buildMasterUrl } from './urlUtils.js';
const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
class DistributedExtension {
// Exported for tests; ComfyUI only uses the registerExtension side effect below.
export class DistributedExtension {
constructor() {
this.config = null;
this.originalQueuePrompt = api.queuePrompt.bind(api);
@@ -196,6 +197,7 @@ class DistributedExtension {
async updateWorkerEnabled(workerId, enabled) {
const worker = this.config.workers.find(w => w.id === workerId);
const previousEnabled = worker ? Boolean(worker.enabled) : !enabled;
if (worker) {
worker.enabled = enabled;
this.state.updateWorker(workerId, { enabled });
@@ -219,6 +221,19 @@ class DistributedExtension {
await this.api.updateWorker(workerId, { enabled });
} catch (error) {
this.log("Error updating worker: " + error.message, "error");
// The backend kept the old state (e.g. a serverless worker without a
// saved API key); show that instead of a toggle that did not stick.
if (worker) {
worker.enabled = previousEnabled;
this.state.updateWorker(workerId, { enabled: previousEnabled });
this._emitWorkersChanged();
}
this.app?.extensionManager?.toast?.add({
severity: "error",
summary: enabled ? "Could not enable worker" : "Could not disable worker",
detail: error.message,
life: 5000,
});
}
if (this.panelElement) {
+112 -14
View File
@@ -7,7 +7,7 @@ import { createWorkerSettingsForm } from "../ui/settingsForm.js";
import { DistributedUI } from "../ui.js";
import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
import { isRemoteWorker, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
import { isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
import { performPreflightCheck } from "../executionUtils.js";
function createDocumentMock() {
@@ -125,12 +125,14 @@ describe("Runpod serverless worker panel", () => {
expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
}
expect(elements.get("runpod-text-saved").style.display).toBe("none");
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "concurrency", "timeout-seconds", "serverless-test"]) {
for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
}
// The separate concurrency setting was removed; a legacy value is ignored.
expect(elements.has("concurrency-saved")).toBe(false);
expect(elements.has("concurrency-group-saved")).toBe(false);
expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
expect(elements.get("jobs-per-run-saved").value).toBe("2");
expect(elements.get("concurrency-saved").value).toBe("1");
expect(elements.get("timeout-seconds-saved").value).toBe("900");
expect(elements.get("api-key-saved").value).toBe("");
expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
@@ -170,10 +172,11 @@ describe("Runpod serverless worker panel", () => {
for (const id of ["host-group-w1", "port-group-w1", "cuda-group-w1", "args-group-w1"]) {
expect(elements.get(id).style.display).toBe("none");
}
for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "concurrency-w1", "timeout-seconds-w1"]) {
for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "timeout-seconds-w1"]) {
expect(elements.has(id)).toBe(true);
}
for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "concurrency-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) {
expect(elements.has("concurrency-w1")).toBe(false);
for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) {
expect(elements.get(id).style.display).toBe("flex");
}
expect(elements.has("serverless-help-w1")).toBe(false);
@@ -183,9 +186,9 @@ describe("Runpod serverless worker panel", () => {
expect(elements.get("jobs-per-run-w1").value).toBe("1");
expect(elements.get("jobs-per-run-w1").min).toBe("1");
expect(elements.get("jobs-per-run-w1").max).toBe("16");
expect(elements.get("concurrency-w1").value).toBe("1");
expect(elements.get("concurrency-w1").min).toBe("1");
expect(elements.get("concurrency-w1").max).toBe("16");
expect(elements.get("jobs-per-run-w1").title).toBe(
"Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.",
);
expect(elements.get("timeout-seconds-w1").value).toBe("1800");
expect(elements.get("timeout-seconds-w1").min).toBe("30");
expect(elements.get("timeout-seconds-w1").max).toBe("86400");
@@ -218,14 +221,13 @@ describe("Runpod serverless worker panel", () => {
elements.set("api-key-w2", { value: "secret-token" });
elements.set("clear-api-key-w2", { checked: false });
elements.set("jobs-per-run-w2", { value: "4" });
elements.set("concurrency-w2", { value: "3" });
elements.set("timeout-seconds-w2", { value: "2400" });
await saveWorkerSettings(extension, "w2");
expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
api_key: "secret-token", jobs_per_run: 4, concurrency: 3, timeout_seconds: 2400, enabled: false,
api_key: "secret-token", jobs_per_run: 4, timeout_seconds: 2400, enabled: false,
});
expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
expect(worker).not.toHaveProperty("api_key");
@@ -235,6 +237,104 @@ describe("Runpod serverless worker panel", () => {
expect(JSON.stringify(extension.state.updateWorker.mock.calls)).not.toContain("secret-token");
});
it("labels the single parallelism setting Images per run and drops a legacy concurrency value", async () => {
const worker = { id: "legacy", name: "Runpod", type: "serverless", provider: "runpod", enabled: false,
api_key_configured: true, jobs_per_run: 1, concurrency: 1 };
const extension = makeSaveExtension(worker);
elements.set("name-legacy", { value: "Runpod" });
elements.set("worker-type-legacy", { value: "serverless" });
elements.set("endpoint-id-legacy", { value: "ep_123" });
elements.set("api-key-legacy", { value: "" });
elements.set("clear-api-key-legacy", { checked: false });
elements.set("jobs-per-run-legacy", { value: "2" });
elements.set("timeout-seconds-legacy", { value: "1800" });
await saveWorkerSettings(extension, "legacy");
const payload = extension.api.updateWorker.mock.calls[0][1];
expect(payload.jobs_per_run).toBe(2);
expect(payload).not.toHaveProperty("concurrency");
expect(worker.jobs_per_run).toBe(2);
expect(worker).not.toHaveProperty("concurrency");
elements.set("jobs-per-run-legacy", { value: "17" });
await saveWorkerSettings(extension, "legacy");
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
expect.objectContaining({ detail: "Images per run must be a whole number from 1 to 16" }),
);
});
it("trims a pasted Runpod URL to the endpoint ID and explains other invalid IDs", async () => {
expect(normalizeRunpodEndpointId(" abc123xyz ")).toBe("abc123xyz");
expect(normalizeRunpodEndpointId("https://api.runpod.ai/v2/abc123xyz/run")).toBe("abc123xyz");
expect(normalizeRunpodEndpointId("api.runpod.ai/v2/abc_12-3")).toBe("abc_12-3");
expect(normalizeRunpodEndpointId("https://www.runpod.io/console/serverless/user/endpoint/abc")).toBe(
"https://www.runpod.io/console/serverless/user/endpoint/abc",
);
const worker = { id: "url", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
const extension = makeSaveExtension(worker);
elements.set("name-url", { value: "Runpod" });
elements.set("worker-type-url", { value: "serverless" });
elements.set("endpoint-id-url", { value: "https://api.runpod.ai/v2/abc123xyz/runsync" });
elements.set("api-key-url", { value: "" });
elements.set("clear-api-key-url", { checked: false });
elements.set("jobs-per-run-url", { value: "1" });
elements.set("timeout-seconds-url", { value: "1800" });
await saveWorkerSettings(extension, "url");
expect(extension.api.updateWorker.mock.calls[0][1].endpoint_id).toBe("abc123xyz");
elements.set("endpoint-id-url", { value: "https://www.runpod.io/console/serverless/endpoint/abc" });
await saveWorkerSettings(extension, "url");
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
expect.objectContaining({ detail: expect.stringMatching(/just the Runpod endpoint ID/) }),
);
});
it("asks before a type change deletes the saved API key, and cancels cleanly", async () => {
const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
const extension = makeSaveExtension(worker);
elements.set("name-conv", { value: "Runpod" });
elements.set("worker-type-conv", { value: "remote" });
elements.set("host-conv", { value: "192.168.1.20" });
elements.set("port-conv", { value: "8188" });
const confirmSpy = vi.fn().mockReturnValue(false);
globalThis.confirm = confirmSpy;
try {
await saveWorkerSettings(extension, "conv");
expect(confirmSpy).toHaveBeenCalledWith(expect.stringMatching(/saved Runpod API key will be deleted/));
expect(extension.api.updateWorker).not.toHaveBeenCalled();
expect(worker.type).toBe("serverless");
confirmSpy.mockReturnValue(true);
await saveWorkerSettings(extension, "conv");
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
expect(extension.api.updateWorker.mock.calls[0][1].type).toBe("remote");
} finally {
delete globalThis.confirm;
}
});
it("does not ask for confirmation when a serverless worker without a key changes type", async () => {
const worker = { id: "nokey2", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: false };
const extension = makeSaveExtension(worker);
elements.set("name-nokey2", { value: "Runpod" });
elements.set("worker-type-nokey2", { value: "remote" });
elements.set("host-nokey2", { value: "192.168.1.20" });
elements.set("port-nokey2", { value: "8188" });
const confirmSpy = vi.fn().mockReturnValue(false);
globalThis.confirm = confirmSpy;
try {
await saveWorkerSettings(extension, "nokey2");
expect(confirmSpy).not.toHaveBeenCalled();
expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
} finally {
delete globalThis.confirm;
}
});
it("omits a blank API key so the saved key is preserved", async () => {
const worker = { id: "w-blank", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
const extension = makeSaveExtension(worker);
@@ -244,7 +344,6 @@ describe("Runpod serverless worker panel", () => {
elements.set("api-key-w-blank", { value: "" });
elements.set("clear-api-key-w-blank", { checked: false });
elements.set("jobs-per-run-w-blank", { value: "1" });
elements.set("concurrency-w-blank", { value: "1" });
elements.set("timeout-seconds-w-blank", { value: "1800" });
await saveWorkerSettings(extension, "w-blank");
@@ -265,14 +364,13 @@ describe("Runpod serverless worker panel", () => {
elements.set("api-key-w3", { value: "" });
elements.set("clear-api-key-w3", { checked: true });
elements.set("jobs-per-run-w3", { value: "1" });
elements.set("concurrency-w3", { value: "1" });
elements.set("timeout-seconds-w3", { value: "1800" });
await saveWorkerSettings(extension, "w3");
expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
clear_api_key: true, jobs_per_run: 1, concurrency: 1, timeout_seconds: 1800, enabled: true,
clear_api_key: true, jobs_per_run: 1, timeout_seconds: 1800, enabled: true,
});
expect(worker.api_key_configured).toBe(false);
expect(worker).not.toHaveProperty("api_key");
@@ -317,7 +415,7 @@ describe("Runpod serverless worker panel", () => {
expect(elements.get("serverless-status-compact").textContent).toBe("Queued");
expect(elements.get("endpoint-id-compact").value).toBe("ep_private_id");
expect(Number(elements.get("jobs-per-run-compact").value)).toBe(1);
expect(Number(elements.get("concurrency-compact").value)).toBe(1);
expect(elements.has("concurrency-compact")).toBe(false);
const box = elements.get("controls-compact").children[0].children[0];
expect(box.textContent).toBe("Runpod Serverless endpoint");
expect(box.classList.values.has("dist-info-box")).toBe(true);
+54
View File
@@ -0,0 +1,54 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
vi.mock("../../scripts/app.js", () => ({ app: { registerExtension: vi.fn() } }), { virtual: true });
vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client", queuePrompt: vi.fn() } }), { virtual: true });
import { DistributedExtension } from "../main.js";
function makeExtension(worker, updateWorker) {
return {
config: { workers: [worker] },
state: { updateWorker: vi.fn() },
api: { updateWorker },
app: { extensionManager: { toast: { add: vi.fn() } } },
ui: { updateStatusDot: vi.fn() },
log: vi.fn(),
_emitWorkersChanged: vi.fn(),
panelElement: null,
};
}
describe("worker enable toggle", () => {
let originalDocument;
beforeEach(() => {
originalDocument = globalThis.document;
globalThis.document = { getElementById: () => null };
});
afterEach(() => {
globalThis.document = originalDocument;
});
it("reverts the toggle and explains why when the backend rejects the change", async () => {
const worker = { id: "w1", type: "serverless", enabled: false };
const extension = makeExtension(worker, vi.fn().mockRejectedValue(new Error("Save an API key before enabling this worker")));
await DistributedExtension.prototype.updateWorkerEnabled.call(extension, "w1", true);
expect(worker.enabled).toBe(false);
expect(extension.state.updateWorker).toHaveBeenLastCalledWith("w1", { enabled: false });
expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
severity: "error", summary: "Could not enable worker", detail: "Save an API key before enabling this worker",
}));
});
it("keeps the new state when the backend accepts it", async () => {
const worker = { id: "w2", type: "serverless", enabled: false };
const extension = makeExtension(worker, vi.fn().mockResolvedValue({ status: "success" }));
await DistributedExtension.prototype.updateWorkerEnabled.call(extension, "w2", true);
expect(worker.enabled).toBe(true);
expect(extension.state.updateWorker).toHaveBeenLastCalledWith("w2", { enabled: true });
expect(extension.app.extensionManager.toast.add).not.toHaveBeenCalled();
});
});
+6 -12
View File
@@ -1,6 +1,9 @@
import { BUTTON_STYLES } from '../constants.js';
import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings, testServerlessWorker } from '../workerSettings.js';
export const IMAGES_PER_RUN_TOOLTIP =
"Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.";
export function createWorkerSettingsForm(ui, extension, worker) {
const form = document.createElement("div");
form.style.cssText = "display: flex; flex-direction: column; gap: 8px;";
@@ -59,7 +62,6 @@ export function createWorkerSettingsForm(ui, extension, worker) {
apiKeyGroup.group,
clearApiKeyGroup,
jobsPerRunGroup.group,
concurrencyGroup.group,
timeoutGroup.group,
testServerlessGroup,
];
@@ -169,24 +171,16 @@ export function createWorkerSettingsForm(ui, extension, worker) {
};
form.appendChild(clearApiKeyGroup);
const jobsPerRunGroup = ui.createFormGroup("Jobs per run:", worker.jobs_per_run ?? 1, `jobs-per-run-${worker.id}`, "number");
const jobsPerRunGroup = ui.createFormGroup("Images per run (1–16):", worker.jobs_per_run ?? 1, `jobs-per-run-${worker.id}`, "number");
jobsPerRunGroup.group.id = `jobs-per-run-group-${worker.id}`;
jobsPerRunGroup.input.min = "1";
jobsPerRunGroup.input.max = "16";
jobsPerRunGroup.input.step = "1";
jobsPerRunGroup.input.title = "Maximum jobs submitted in one run";
jobsPerRunGroup.input.title = IMAGES_PER_RUN_TOOLTIP;
jobsPerRunGroup.group.title = IMAGES_PER_RUN_TOOLTIP;
jobsPerRunGroup.group.style.display = "none";
form.appendChild(jobsPerRunGroup.group);
const concurrencyGroup = ui.createFormGroup("Concurrent jobs (slots):", worker.concurrency ?? 1, `concurrency-${worker.id}`, "number");
concurrencyGroup.group.id = `concurrency-group-${worker.id}`;
concurrencyGroup.input.min = "1";
concurrencyGroup.input.max = "16";
concurrencyGroup.input.step = "1";
concurrencyGroup.input.title = "Maximum concurrent jobs per serverless worker";
concurrencyGroup.group.style.display = "none";
form.appendChild(concurrencyGroup.group);
const timeoutGroup = ui.createFormGroup("Timeout (seconds):", worker.timeout_seconds ?? 1800, `timeout-seconds-${worker.id}`, "number");
timeoutGroup.group.id = `timeout-seconds-group-${worker.id}`;
timeoutGroup.input.min = "30";
+31 -10
View File
@@ -37,6 +37,19 @@ export function isCloudWorker(extension, worker) {
return worker.type === "cloud";
}
// Mirrors the backend rule in utils/serverless_config.py.
const RUNPOD_ENDPOINT_ID_RE = /^[A-Za-z0-9_-]{1,200}$/;
/**
* Accept a pasted Runpod API URL (https://api.runpod.ai/v2/<id>/...) and keep
* only the endpoint ID; anything else is returned trimmed for validation.
*/
export function normalizeRunpodEndpointId(value) {
const text = String(value || "").trim();
const match = text.match(/^(?:https?:\/\/)?api\.runpod\.ai\/v2\/([A-Za-z0-9_-]{1,200})(?:[/?#].*)?$/i);
return match ? match[1] : text;
}
function showServerlessValidationError(extension, detail) {
extension.app.extensionManager.toast.add({
severity: "error",
@@ -47,11 +60,10 @@ function showServerlessValidationError(extension, detail) {
}
async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
const endpointId = String(document.getElementById(`endpoint-id-${workerId}`).value || "").trim();
const endpointId = normalizeRunpodEndpointId(document.getElementById(`endpoint-id-${workerId}`).value);
const apiKey = document.getElementById(`api-key-${workerId}`).value || "";
const clearApiKey = Boolean(document.getElementById(`clear-api-key-${workerId}`).checked);
const jobsPerRun = Number(document.getElementById(`jobs-per-run-${workerId}`).value);
const concurrency = Number(document.getElementById(`concurrency-${workerId}`).value);
const timeoutSeconds = Number(document.getElementById(`timeout-seconds-${workerId}`).value);
if (!name.trim()) {
@@ -62,12 +74,15 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
showServerlessValidationError(extension, "Endpoint ID is required for Runpod serverless workers");
return;
}
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
showServerlessValidationError(extension, "Jobs per run must be an integer between 1 and 16");
if (!RUNPOD_ENDPOINT_ID_RE.test(endpointId)) {
showServerlessValidationError(
extension,
"Enter just the Runpod endpoint ID (for example abc123xyz), not a URL. It is shown on the endpoint's page in the Runpod console.",
);
return;
}
if (!Number.isInteger(concurrency) || concurrency < 1 || concurrency > 16) {
showServerlessValidationError(extension, "Concurrency must be an integer between 1 and 16");
if (!Number.isInteger(jobsPerRun) || jobsPerRun < 1 || jobsPerRun > 16) {
showServerlessValidationError(extension, "Images per run must be a whole number from 1 to 16");
return;
}
if (!Number.isInteger(timeoutSeconds) || timeoutSeconds < 30 || timeoutSeconds > 86400) {
@@ -83,7 +98,6 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
...(apiKey.length > 0 && !clearApiKey ? { api_key: apiKey } : {}),
...(clearApiKey ? { clear_api_key: true } : {}),
jobs_per_run: jobsPerRun,
concurrency,
timeout_seconds: timeoutSeconds,
// Never enable a serverless worker as a side effect of saving credentials.
enabled: worker.enabled,
@@ -98,7 +112,8 @@ async function saveServerlessWorkerSettings(extension, worker, workerId, name) {
worker.provider = "runpod";
worker.endpoint_id = endpointId;
worker.jobs_per_run = jobsPerRun;
worker.concurrency = concurrency;
// Replaced by a fixed backend limit; drop any value loaded from an older config.
delete worker.concurrency;
worker.timeout_seconds = timeoutSeconds;
worker.api_key_configured = clearApiKey
? false
@@ -178,6 +193,14 @@ export async function saveWorkerSettings(extension, workerId) {
return;
}
const name = document.getElementById(`name-${workerId}`).value;
if (worker.type === "serverless" && worker.api_key_configured) {
// The backend deletes the stored key when a worker stops being serverless.
const proceed = confirm(
`Change "${worker.name}" from Serverless (Runpod) to another type?\n\n` +
"Its saved Runpod API key will be deleted. You would need to enter it again to switch back."
);
if (!proceed) return;
}
const isRemote = workerType === "remote" || workerType === "cloud";
const isCloud = workerType === "cloud";
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
@@ -352,8 +375,6 @@ export function cancelWorkerSettings(extension, workerId) {
if (clearApiKeyInput) clearApiKeyInput.checked = false;
const jobsPerRunInput = document.getElementById(`jobs-per-run-${workerId}`);
if (jobsPerRunInput) jobsPerRunInput.value = worker.jobs_per_run ?? 1;
const concurrencyInput = document.getElementById(`concurrency-${workerId}`);
if (concurrencyInput) concurrencyInput.value = worker.concurrency ?? 1;
const timeoutInput = document.getElementById(`timeout-seconds-${workerId}`);
if (timeoutInput) timeoutInput.value = worker.timeout_seconds ?? 1800;