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(), 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, {})
|