nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge feat/runpod-serverless-workers into feat/dynamic-image-list-queue

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-10-13 21:33:28 +00:00
parent dcddf23db2 06be0e5dee
commit 00a9ba05a3
17 changed files with 407 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_serverless_config.py
+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():
tests/test_serverless_lifecycle.py
+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()
tests/test_serverless_routes.py
+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"}))
tests/test_serverless_runtime.py
+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")