Robert Wojciechowski
6b94242110
* feat: add Runpod serverless image collector pools * fix: prevent worker settings clipping and remove preview copy * fix: initialize worker settings before mounting the form * fix: validate complete serverless requests against 10 MB limit * fix: simplify serverless cards and keep status dots circular * fix: match serverless endpoint box to master card * fix: preserve serverless capacity and recovery across job failures Retain durable reservations until provider termination or explicit recovery, protect unresolved pools and credentials, and add scoped reconciliation. Consolidate lifecycle and secret policies and split prompt/media stages with regression coverage. * fix(serverless): disable Clear saved API key while the worker is enabled The checkbox read "Clear saved API key (disable pool first)", but the panel never defines "pool" and saving with it ticked on an enabled worker was rejected by the backend. - Disable the checkbox while the worker is enabled, with a tooltip saying to turn the worker off first. Also disable it when no key is saved. - Use "worker" instead of "pool" in user-facing serverless errors, the serverless guide and the README entry. * feat(serverless): replace Concurrent jobs with a single Images per run setting - Remove the separate Concurrent jobs field; rename Jobs per run to Images per run (1-16) with a tooltip explaining that Runpod's max workers decides how many GPUs run them. - Backend uses a fixed per-worker limit of 16 outstanding jobs, so all images of a run are submitted at once. Unresolved jobs still count against it. The limit is constant, so it never has to change while jobs are unresolved. - Legacy 'concurrency' in saved configs and older browser tabs is accepted and ignored/dropped; it no longer triggers the busy guard. - Revert the enable toggle and show the reason when the backend rejects it. - Trim a pasted api.runpod.ai/v2/<id> URL to the endpoint ID; plainer endpoint-ID error otherwise. - Confirm before a type change deletes a saved Runpod API key. - Update docs and tests.
184 lines
11 KiB
Python
184 lines
11 KiB
Python
"""Offline integration: real graph preparation + orchestration, fake provider."""
|
|
import asyncio
|
|
import aiohttp as real_aiohttp
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
PKG = "dist_serverless_integration"
|
|
|
|
|
|
def module(name, **attrs):
|
|
value = types.ModuleType(name)
|
|
value.__dict__.update(attrs)
|
|
return value
|
|
|
|
|
|
def load_modules():
|
|
modules = {}
|
|
for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]:
|
|
modules[PKG + part] = module(PKG + part, __path__=[str(directory)])
|
|
ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn))
|
|
modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps))
|
|
modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None)
|
|
modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None)
|
|
modules[PKG + ".utils.config"] = module("config", load_config=lambda: {})
|
|
modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock())
|
|
modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master")
|
|
modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock())
|
|
modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), fetch_worker_has_node=AsyncMock(return_value=False), sync_worker_media=AsyncMock())
|
|
with patch.dict(sys.modules, modules):
|
|
spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py")
|
|
q = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(q)
|
|
# Keep imported package namespace only; restore shared Comfy runtime stub.
|
|
modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)})
|
|
return q, modules, ps
|
|
|
|
|
|
class IntegrationTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
self.q, modules, self.ps = load_modules()
|
|
self.patch = patch.dict(sys.modules, modules)
|
|
self.patch.start()
|
|
self.addCleanup(self.patch.stop)
|
|
self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30}
|
|
self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}}
|
|
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
|
|
self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
|
|
self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None)
|
|
sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager)
|
|
sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[]))
|
|
self.prompt = {
|
|
"1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
|
|
"2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}},
|
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
|
|
}
|
|
|
|
async def test_real_compiler_slots_to_master_collector_without_callbacks(self):
|
|
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
|
self.assertEqual(result[2], 2)
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
self.manager.start.assert_awaited_once()
|
|
jobs = self.manager.start.call_args.args[2]
|
|
self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"])
|
|
self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2)
|
|
for job in jobs:
|
|
graph = job["workflow"]
|
|
self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values()))
|
|
self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values()))
|
|
self.assertNotIn("http://master", json.dumps(graph))
|
|
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
|
self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"])
|
|
|
|
async def test_compile_pack_provider_decode_and_collect_real_components(self):
|
|
"""Exercise the whole image path, mocking only the provider and Comfy queue."""
|
|
import base64
|
|
import io
|
|
import tempfile
|
|
import torch
|
|
from PIL import Image
|
|
from test_collector_list_inputs import _load_collector_module
|
|
|
|
sys.modules["aiohttp"] = real_aiohttp
|
|
for leaf in ("serverless_media", "serverless_runtime"):
|
|
name = PKG + ".api.orchestration." + leaf
|
|
spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py"))
|
|
real_module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = real_module
|
|
spec.loader.exec_module(real_module)
|
|
runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"]
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
image_buffer = io.BytesIO()
|
|
Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG")
|
|
image_bytes = image_buffer.getvalue()
|
|
(root / "source.png").write_bytes(image_bytes)
|
|
paths = module("folder_paths", get_input_directory=lambda: directory,
|
|
get_output_directory=lambda: directory, get_temp_directory=lambda: directory,
|
|
get_annotated_filepath=lambda name: str(root / name))
|
|
submissions = []
|
|
|
|
class FakeProvider:
|
|
async def submit(self, graph, images, timeout_seconds):
|
|
self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage")
|
|
submissions.append((graph, images))
|
|
self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage")
|
|
assert self_image == images[0]["name"]
|
|
assert base64.b64decode(images[0]["image"]) == image_bytes
|
|
return "provider-job-" + str(len(submissions))
|
|
|
|
async def status(self, _job_id):
|
|
filename = self.prefix + "_00001_.png"
|
|
return {"status": "COMPLETED", "output": {"images": [{
|
|
"filename": filename, "artifact_path": filename, "artifact_type": "output",
|
|
"type": "base64", "data": base64.b64encode(image_bytes).decode(),
|
|
}]}}
|
|
|
|
manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None)
|
|
manager._client = AsyncMock(side_effect=lambda _: FakeProvider())
|
|
self.ps.distributed_serverless = manager
|
|
self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}}
|
|
collector_module = _load_collector_module()
|
|
collector_module.prompt_server = self.ps
|
|
sys.modules["aiohttp"] = real_aiohttp
|
|
with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}):
|
|
result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
|
master_graph = self.q.queue_prompt_payload.call_args.args[0]
|
|
inputs = master_graph["2"]["inputs"]
|
|
images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute(
|
|
images=None, audio=None, multi_job_id=inputs["multi_job_id"],
|
|
enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True,
|
|
), timeout=3)
|
|
await asyncio.gather(*manager.tasks.values())
|
|
self.assertEqual(result[2], 2)
|
|
self.assertEqual(len(submissions), 2)
|
|
self.assertEqual(tuple(images.shape), (2, 3, 2, 3))
|
|
self.assertEqual(images.dtype, torch.float32)
|
|
self.assertTrue(torch.equal(images[0], images[1]))
|
|
self.assertEqual(float(images[0, 0, 0, 0]), 1.0)
|
|
self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"])
|
|
self.assertEqual(self.ps.distributed_pending_jobs, {})
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
|
|
async def test_audio_rejected_before_any_dispatch_or_master_queue(self):
|
|
self.prompt["2"]["inputs"]["audio"] = ["9", 0]
|
|
self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}}
|
|
with self.assertRaises(ValueError):
|
|
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
|
self.manager.start.assert_not_awaited()
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|
|
|
|
async def test_disabled_pool_excluded(self):
|
|
self.pool["enabled"] = False
|
|
self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), [])
|
|
|
|
async def test_master_validation_failure_submits_no_paid_job(self):
|
|
self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph")
|
|
with self.assertRaisesRegex(ValueError, "invalid local graph"):
|
|
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
|
self.manager.start.assert_not_awaited()
|
|
|
|
async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self):
|
|
media_module = sys.modules[PKG + ".api.orchestration.serverless_media"]
|
|
media_module.prepare_images.return_value = [{
|
|
"name": "private-image-name.png",
|
|
"image": "x" * 10_000_001,
|
|
}]
|
|
|
|
with self.assertRaisesRegex(ValueError, "10,000,") as ctx:
|
|
await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
|
|
|
|
self.assertNotIn("private-image-name.png", str(ctx.exception))
|
|
self.assertNotIn("x" * 100, str(ctx.exception))
|
|
self.q.dispatch_worker_prompt.assert_not_awaited()
|
|
self.q.queue_prompt_payload.assert_not_awaited()
|
|
self.manager.start.assert_not_awaited()
|
|
self.assertEqual(self.ps.distributed_pending_jobs, {})
|