Merge feat/runpod-serverless-workers into feat/dynamic-image-list-queue
This commit is contained in:
Robert Wojciechowski
@@ -48,7 +48,7 @@ ComfyUI Distributed supports four types of workers:
|
||||
- **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
|
||||
- **Remote Workers** - GPUs on other computers in your network
|
||||
- **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
|
||||
- **Serverless (Runpod)** - Queue-based endpoint pools for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
- **Serverless (Runpod)** - Queue-based Runpod endpoints for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
|
||||
|
||||
> For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
|
||||
|
||||
|
||||
+9
-13
@@ -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)
|
||||
@@ -299,7 +298,7 @@ async def update_config_endpoint(request):
|
||||
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)
|
||||
return _error_response("Save an API key before enabling this worker", 400)
|
||||
if changed_ids:
|
||||
store = store or _serverless_secret_store()
|
||||
for worker_id in sorted(changed_ids):
|
||||
@@ -316,7 +315,7 @@ async def update_config_endpoint(request):
|
||||
secret_changes.commit()
|
||||
return response
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to update config", 500)
|
||||
|
||||
@@ -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,17 +424,13 @@ 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)
|
||||
|
||||
if validated.get("enabled", False) and next_secret is None:
|
||||
return _error_response("Save an API key before enabling this pool; disable it before clearing the key", 400)
|
||||
return _error_response("Save an API key before enabling this worker; turn the worker off before clearing its key", 400)
|
||||
|
||||
if next_secret != prior_secret:
|
||||
if next_secret is None:
|
||||
@@ -489,7 +485,7 @@ async def _update_worker_endpoint(request, secret_changes):
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
|
||||
except ValueError:
|
||||
return _error_response("Invalid worker configuration", 400)
|
||||
except Exception:
|
||||
@@ -537,7 +533,7 @@ async def _delete_worker_endpoint(request, secret_changes):
|
||||
|
||||
return web.json_response(result)
|
||||
except _ActivePoolCheckError:
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this pool", 409)
|
||||
return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
|
||||
except Exception:
|
||||
return _error_response("Unable to delete worker", 500)
|
||||
|
||||
|
||||
@@ -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 pool'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:
|
||||
@@ -320,7 +325,7 @@ class ServerlessManager:
|
||||
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")
|
||||
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()
|
||||
|
||||
@@ -83,7 +83,7 @@ async def _recover(request, *, acknowledge=False):
|
||||
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)
|
||||
return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this worker"}, status=409)
|
||||
except Exception:
|
||||
return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
|
||||
|
||||
|
||||
@@ -461,7 +461,7 @@ async def launch_worker_endpoint(request):
|
||||
return await handle_api_error(request, f"Worker {worker_id} not found", 404)
|
||||
|
||||
if worker.get("type") == "serverless":
|
||||
return await handle_api_error(request, "Serverless pools are managed by Runpod, not local processes", 400)
|
||||
return await handle_api_error(request, "Serverless workers are managed by Runpod, not local processes", 400)
|
||||
|
||||
# Ensure consistent string ID
|
||||
worker_id_str = worker_id
|
||||
|
||||
+34
-24
@@ -1,24 +1,26 @@
|
||||
# Runpod serverless pools (preview)
|
||||
# Runpod serverless workers (preview)
|
||||
|
||||
A **Serverless (Runpod)** worker is a queue-based endpoint/pool, not a permanent
|
||||
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 pool when you want queued Distributed workflows to use it.
|
||||
5. Enable the worker when you want queued Distributed workflows to use it.
|
||||
6. Use an image **Distributed Collector** and queue through Distributed as usual.
|
||||
Enable **Delegate Only** if the master should collect without generating locally.
|
||||
|
||||
Keys are write-only. Leave the key field blank to retain a saved key, or select
|
||||
**Clear saved API key** to remove it (disable the pool first). A saved key is
|
||||
required before enabling a pool. Use HTTPS or a trusted local connection when
|
||||
**Clear saved API key** to remove it. The checkbox is unavailable while the
|
||||
worker is enabled, so turn the worker off first. A saved key is required before
|
||||
enabling a worker. Use HTTPS or a trusted local connection when
|
||||
entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
|
||||
users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
|
||||
|
||||
@@ -71,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 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.
|
||||
- **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.
|
||||
@@ -98,13 +108,13 @@ available on Windows.
|
||||
`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
|
||||
not image data, workflow payloads, or API keys. A restarted master cannot resume
|
||||
an interrupted collector. It does **not** automatically resubmit jobs. Retained
|
||||
IDs are available through `/distributed/serverless/status/<worker_id>`; the pool's
|
||||
IDs are available through `/distributed/serverless/status/<worker_id>`; the worker's
|
||||
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;
|
||||
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
|
||||
`UNKNOWN` and remain protected.
|
||||
@@ -122,7 +132,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
Network errors, expired IDs, or ambiguous responses do not.
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID"}
|
||||
{"worker_id": "YOUR_WORKER_ID"}
|
||||
```
|
||||
|
||||
3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
|
||||
@@ -130,7 +140,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
queued**, send `POST /distributed/serverless/resolve` for that one record:
|
||||
|
||||
```json
|
||||
{"worker_id": "YOUR_POOL_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
|
||||
{"worker_id": "YOUR_WORKER_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
|
||||
```
|
||||
|
||||
This records `RESOLVED`: an explicit operator acknowledgment, **not** a
|
||||
@@ -139,7 +149,7 @@ Recovery uses the same trusted-master API as the panel; it never resubmits a job
|
||||
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
|
||||
Worker setting 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
|
||||
@@ -160,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.
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"}))
|
||||
|
||||
@@ -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")
|
||||
|
||||
+1
-1
@@ -88,7 +88,7 @@ def _split_host_and_port(host):
|
||||
def build_worker_url(worker, endpoint=""):
|
||||
"""Construct the worker base URL with optional endpoint."""
|
||||
if worker.get("type") == "serverless":
|
||||
raise ValueError("Serverless pools have no ComfyUI worker URL")
|
||||
raise ValueError("Serverless workers have no ComfyUI worker URL")
|
||||
host = (worker.get("host") or "").strip()
|
||||
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
|
||||
|
||||
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
@@ -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,15 +186,32 @@ 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");
|
||||
expect(form.children.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("disables Clear saved API key while the worker is enabled and explains why", () => {
|
||||
vi.spyOn(document, "getElementById").mockReturnValue(null);
|
||||
const base = { name: "Runpod", type: "serverless", endpoint_id: "ep_123" };
|
||||
createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "on", enabled: true, api_key_configured: true });
|
||||
createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "off", enabled: false, api_key_configured: true });
|
||||
createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "nokey", enabled: false, api_key_configured: false });
|
||||
|
||||
expect(elements.get("clear-api-key-on").disabled).toBe(true);
|
||||
expect(elements.get("clear-api-key-group-on").title).toMatch(/turn this worker off/i);
|
||||
expect(elements.get("clear-api-key-off").disabled).toBe(false);
|
||||
expect(elements.get("clear-api-key-nokey").disabled).toBe(true);
|
||||
expect(elements.get("clear-api-key-group-nokey").title).toMatch(/no api key is saved/i);
|
||||
for (const id of ["on", "off", "nokey"]) {
|
||||
expect(elements.get(`clear-api-key-group-${id}`).title).not.toMatch(/pool/i);
|
||||
}
|
||||
});
|
||||
|
||||
it("saves fixed Runpod settings without enabling or retaining the API key", async () => {
|
||||
const worker = { id: "w2", name: "Old", type: "remote", enabled: false, api_key_configured: false };
|
||||
const extension = makeSaveExtension(worker);
|
||||
@@ -201,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");
|
||||
@@ -218,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);
|
||||
@@ -227,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");
|
||||
@@ -248,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");
|
||||
@@ -300,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);
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
+17
-13
@@ -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,
|
||||
];
|
||||
@@ -150,7 +152,17 @@ export function createWorkerSettingsForm(ui, extension, worker) {
|
||||
clearApiKeyInput.id = `clear-api-key-${worker.id}`;
|
||||
const clearApiKeyLabel = document.createElement("label");
|
||||
clearApiKeyLabel.htmlFor = clearApiKeyInput.id;
|
||||
clearApiKeyLabel.textContent = "Clear saved API key (disable pool first)";
|
||||
clearApiKeyLabel.textContent = "Clear saved API key";
|
||||
// The backend rejects clearing the key of an enabled worker; say so up front.
|
||||
// The form is rebuilt when the enable toggle changes, so this stays current.
|
||||
const clearApiKeyBlockedReason = !worker.api_key_configured
|
||||
? "No API key is saved for this worker."
|
||||
: worker.enabled
|
||||
? "Turn this worker off before clearing its API key. An enabled serverless worker needs a saved key."
|
||||
: "";
|
||||
clearApiKeyInput.disabled = Boolean(clearApiKeyBlockedReason);
|
||||
clearApiKeyGroup.title = clearApiKeyBlockedReason || "Remove the saved Runpod API key when you save.";
|
||||
clearApiKeyLabel.style.cssText = clearApiKeyBlockedReason ? "opacity: 0.5; cursor: not-allowed;" : "cursor: pointer;";
|
||||
clearApiKeyGroup.appendChild(clearApiKeyInput);
|
||||
clearApiKeyGroup.appendChild(clearApiKeyLabel);
|
||||
clearApiKeyInput.onchange = () => {
|
||||
@@ -159,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
@@ -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;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user