Robert Wojciechowski
59ccbabe79
Resolve nodes/collector.py: keep the serverless manager accessor and main's compressed result transfer. Update test stubs for the merged code: fetch_worker_has_node (#99) in the serverless integration test, and the real aiohttp.web in the transfer integration test.
556 lines
21 KiB
Python
556 lines
21 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
class _FakeResponse:
|
|
def __init__(self, payload, status=200):
|
|
self.payload = payload
|
|
self.status = status
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
async def json(self):
|
|
return self._payload
|
|
|
|
|
|
def _register_real_transfer_codec(package_name, utils_pkg):
|
|
"""Load the real utils/transfer_codec.py into the synthetic test package."""
|
|
codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
|
|
codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
|
|
codec_module = importlib.util.module_from_spec(codec_spec)
|
|
sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
|
|
codec_spec.loader.exec_module(codec_module)
|
|
utils_pkg.transfer_codec = codec_module
|
|
return codec_module
|
|
|
|
|
|
def _load_job_routes_module():
|
|
module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
|
|
package_name = "dist_api_queue_testpkg"
|
|
|
|
# Reset package namespace to avoid stale module state across test runs.
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
api_pkg = types.ModuleType(f"{package_name}.api")
|
|
api_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.api"] = api_pkg
|
|
|
|
utils_pkg = types.ModuleType(f"{package_name}.utils")
|
|
utils_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.utils"] = utils_pkg
|
|
_register_real_transfer_codec(package_name, utils_pkg)
|
|
|
|
# aiohttp.web stub
|
|
created_aiohttp_stub = False
|
|
if "aiohttp" not in sys.modules:
|
|
created_aiohttp_stub = True
|
|
aiohttp_module = types.ModuleType("aiohttp")
|
|
aiohttp_module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
sys.modules["aiohttp"] = aiohttp_module
|
|
|
|
# server module stub with route decorators
|
|
class _Routes:
|
|
def get(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
def post(self, _path):
|
|
def _decorator(fn):
|
|
return fn
|
|
return _decorator
|
|
|
|
prompt_server_instance = types.SimpleNamespace(
|
|
routes=_Routes(),
|
|
distributed_jobs_lock=None,
|
|
distributed_pending_jobs={},
|
|
)
|
|
server_module = types.ModuleType("server")
|
|
server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
|
|
sys.modules["server"] = server_module
|
|
|
|
# torch stub (only needed to satisfy import)
|
|
created_torch_stub = False
|
|
if "torch" not in sys.modules:
|
|
created_torch_stub = True
|
|
torch_module = types.ModuleType("torch")
|
|
torch_module.cuda = types.SimpleNamespace(
|
|
is_available=lambda: False,
|
|
empty_cache=lambda: None,
|
|
ipc_collect=lambda: None,
|
|
)
|
|
sys.modules["torch"] = torch_module
|
|
|
|
# PIL stub (only needed to satisfy import)
|
|
created_pil_stub = False
|
|
if "PIL" not in sys.modules:
|
|
created_pil_stub = True
|
|
pil_module = types.ModuleType("PIL")
|
|
image_module = types.ModuleType("PIL.Image")
|
|
pil_module.Image = image_module
|
|
sys.modules["PIL"] = pil_module
|
|
sys.modules["PIL.Image"] = image_module
|
|
|
|
logging_module = types.ModuleType(f"{package_name}.utils.logging")
|
|
logging_module.debug_log = lambda *_args, **_kwargs: None
|
|
logging_module.log = lambda *_args, **_kwargs: None
|
|
sys.modules[f"{package_name}.utils.logging"] = logging_module
|
|
|
|
image_module = types.ModuleType(f"{package_name}.utils.image")
|
|
image_module.pil_to_tensor = lambda *_args, **_kwargs: None
|
|
image_module.ensure_contiguous = lambda tensor: tensor
|
|
sys.modules[f"{package_name}.utils.image"] = image_module
|
|
|
|
audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
|
|
|
|
def _decode_audio_payload(payload):
|
|
if payload is None:
|
|
return None
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Field 'audio' must be an object when provided.")
|
|
|
|
encoded = payload.get("data")
|
|
shape = payload.get("shape")
|
|
dtype = payload.get("dtype", "float32")
|
|
sample_rate = payload.get("sample_rate", 44100)
|
|
if not isinstance(encoded, str) or not encoded.strip():
|
|
raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
|
|
if not isinstance(shape, list) or len(shape) != 3:
|
|
raise ValueError("Field 'audio.shape' must be a 3-item list.")
|
|
if dtype != "float32":
|
|
raise ValueError("Field 'audio.dtype' must be 'float32'.")
|
|
try:
|
|
shape_tuple = tuple(int(dim) for dim in shape)
|
|
except Exception as exc:
|
|
raise ValueError("Field 'audio.shape' must contain integers.") from exc
|
|
|
|
raw = base64.b64decode(encoded, validate=True)
|
|
expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
|
|
if len(raw) != expected_bytes:
|
|
raise ValueError("Field 'audio.data' byte size mismatch.")
|
|
|
|
waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
|
|
return {"waveform": waveform, "sample_rate": int(sample_rate)}
|
|
|
|
audio_payload_module.decode_audio_payload = _decode_audio_payload
|
|
sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
|
|
|
|
network_module = types.ModuleType(f"{package_name}.utils.network")
|
|
|
|
async def _handle_api_error(_request, error, status=500):
|
|
return _FakeResponse({"status": "error", "message": str(error)}, status=status)
|
|
|
|
network_module.handle_api_error = _handle_api_error
|
|
sys.modules[f"{package_name}.utils.network"] = network_module
|
|
|
|
constants_module = types.ModuleType(f"{package_name}.utils.constants")
|
|
constants_module.MEMORY_CLEAR_DELAY = 0.0
|
|
constants_module.JOB_INIT_GRACE_PERIOD = 10.0
|
|
sys.modules[f"{package_name}.utils.constants"] = constants_module
|
|
|
|
async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
|
|
async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
|
|
sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
|
|
|
|
queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
|
|
queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
|
|
sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
|
|
|
|
@dataclass(frozen=True)
|
|
class _QueuePayload:
|
|
prompt: dict
|
|
workflow_meta: object
|
|
client_id: str
|
|
delegate_master: object
|
|
enabled_worker_ids: list
|
|
auto_prepare: bool
|
|
trace_execution_id: object
|
|
|
|
def _parse_queue_request_payload(data):
|
|
if not isinstance(data, dict):
|
|
raise ValueError("Expected a JSON object body")
|
|
prompt = data.get("prompt")
|
|
if not isinstance(prompt, dict):
|
|
raise ValueError("Field 'prompt' must be an object")
|
|
enabled = data.get("enabled_worker_ids")
|
|
if not isinstance(enabled, list):
|
|
raise ValueError("enabled_worker_ids required")
|
|
client_id = data.get("client_id")
|
|
if not isinstance(client_id, str) or not client_id.strip():
|
|
raise ValueError("client_id required")
|
|
return _QueuePayload(
|
|
prompt=prompt,
|
|
workflow_meta=data.get("workflow"),
|
|
client_id=client_id,
|
|
delegate_master=data.get("delegate_master"),
|
|
enabled_worker_ids=enabled,
|
|
auto_prepare=bool(data.get("auto_prepare", True)),
|
|
trace_execution_id=data.get("trace_execution_id"),
|
|
)
|
|
|
|
queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
|
|
queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
|
|
sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
# Keep response assertions independent of real aiohttp import ordering.
|
|
module.web = types.SimpleNamespace(
|
|
json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
|
|
)
|
|
|
|
if created_aiohttp_stub:
|
|
sys.modules.pop("aiohttp", None)
|
|
if created_torch_stub:
|
|
sys.modules.pop("torch", None)
|
|
if created_pil_stub:
|
|
sys.modules.pop("PIL.Image", None)
|
|
sys.modules.pop("PIL", None)
|
|
|
|
return module
|
|
|
|
|
|
job_routes = _load_job_routes_module()
|
|
|
|
|
|
class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
"auto_prepare": True,
|
|
}
|
|
)
|
|
with patch.object(
|
|
job_routes,
|
|
"orchestrate_distributed_execution",
|
|
new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
|
|
):
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
|
|
self.assertEqual(response.payload.get("number"), 42)
|
|
self.assertEqual(response.payload.get("node_errors"), {})
|
|
self.assertTrue(response.payload.get("auto_prepare_supported"))
|
|
|
|
async def test_distributed_queue_missing_prompt_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"enabled_worker_ids": ["w1"],
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("prompt", response.payload.get("message", "").lower())
|
|
|
|
async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"prompt": {"1": {"class_type": "Node"}},
|
|
"client_id": "client-1",
|
|
}
|
|
)
|
|
response = await job_routes.distributed_queue_endpoint(request)
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
|
|
|
|
|
|
class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
|
def _encoded_audio_payload(self):
|
|
waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
|
|
return {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2, 4],
|
|
"dtype": "float32",
|
|
"data": base64.b64encode(waveform.tobytes()).decode("ascii"),
|
|
}
|
|
|
|
async def test_job_complete_accepts_audio_payload(self):
|
|
queue = asyncio.Queue()
|
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"image": "data:image/png;base64,AAAA",
|
|
"audio": self._encoded_audio_payload(),
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
queued = await queue.get()
|
|
self.assertEqual(queued["worker_id"], "worker-1")
|
|
self.assertTrue(queued["is_last"])
|
|
self.assertIsNotNone(queued["audio"])
|
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
|
self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
|
|
|
|
async def test_job_complete_accepts_audio_without_image(self):
|
|
queue = asyncio.Queue()
|
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "audio-only-job",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"audio": self._encoded_audio_payload(),
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
decode_image.assert_not_called()
|
|
queued = await queue.get()
|
|
self.assertIsNone(queued["tensor"])
|
|
self.assertEqual(queued["audio"]["sample_rate"], 44100)
|
|
|
|
async def test_job_complete_rejects_payload_without_image_or_audio(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("image or audio", response.payload.get("message", "").lower())
|
|
|
|
async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"image": 123,
|
|
"audio": self._encoded_audio_payload(),
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("image", response.payload.get("message", "").lower())
|
|
|
|
async def test_job_complete_returns_400_for_malformed_audio(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-1",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 0,
|
|
"audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
|
|
"is_last": True,
|
|
}
|
|
)
|
|
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("audio.shape", response.payload.get("message", "").lower())
|
|
|
|
def test_decode_audio_payload_rejects_bad_shape(self):
|
|
bad = {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2],
|
|
"dtype": "float32",
|
|
"data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
|
|
}
|
|
with self.assertRaises(ValueError):
|
|
job_routes._decode_audio_payload(bad)
|
|
|
|
def test_decode_audio_payload_rejects_bad_dtype(self):
|
|
payload = {
|
|
"sample_rate": 44100,
|
|
"shape": [1, 2, 4],
|
|
"dtype": "float16",
|
|
"data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
|
|
}
|
|
with self.assertRaises(ValueError):
|
|
job_routes._decode_audio_payload(payload)
|
|
|
|
|
|
class _FakeFileField:
|
|
def __init__(self, data):
|
|
self.file = io.BytesIO(data)
|
|
|
|
|
|
class _FakeMultipartRequest:
|
|
content_type = "multipart/form-data"
|
|
|
|
def __init__(self, envelope, parts):
|
|
from multidict import MultiDict
|
|
|
|
fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
|
|
fields += [("media", _FakeFileField(part)) for part in parts]
|
|
self._form = MultiDict(fields)
|
|
|
|
async def post(self):
|
|
return self._form
|
|
|
|
async def json(self):
|
|
raise ValueError("not JSON")
|
|
|
|
|
|
class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
|
|
codec = job_routes.transfer_codec
|
|
|
|
def setUp(self):
|
|
self.queue = asyncio.Queue()
|
|
job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
|
|
job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
|
|
|
|
def _frames(self, count=4, height=20, width=30):
|
|
rng = np.random.default_rng(1)
|
|
return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
|
|
|
|
def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
|
|
chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
|
|
parts = chunk.pop("parts")
|
|
envelope = {
|
|
"job_id": "job-z",
|
|
"worker_id": "worker-1",
|
|
"is_last": is_last,
|
|
**chunk,
|
|
"start_index": chunk["start_index"] + start_offset,
|
|
}
|
|
if audio is not None:
|
|
envelope["audio"] = audio
|
|
envelope.update(overrides)
|
|
return _FakeMultipartRequest(envelope, parts)
|
|
|
|
async def _drain(self):
|
|
items = []
|
|
while not self.queue.empty():
|
|
items.append(await self.queue.get())
|
|
return items
|
|
|
|
async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
|
|
frames = self._frames()
|
|
response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
|
|
|
|
self.assertEqual(response.status, 200)
|
|
items = await self._drain()
|
|
self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
|
|
self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
|
|
self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
|
|
for item, frame in zip(items, frames):
|
|
self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
|
|
self.assertEqual(item["tensor"].dtype, torch.float32)
|
|
expected = frame.astype(np.float32) / 255.0
|
|
self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
|
|
|
|
async def test_non_final_chunk_never_marks_worker_done(self):
|
|
response = await job_routes.job_complete_endpoint(
|
|
self._request_for(self._frames(count=2), "png", is_last=False)
|
|
)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
self.assertFalse(any(item["is_last"] for item in await self._drain()))
|
|
|
|
async def test_audio_is_attached_only_to_final_frame(self):
|
|
audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
|
|
response = await job_routes.job_complete_endpoint(
|
|
self._request_for(self._frames(count=3), "jpeg", audio=audio)
|
|
)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
items = await self._drain()
|
|
self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
|
|
self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
|
|
|
|
async def test_rejects_frame_count_mismatch(self):
|
|
response = await job_routes.job_complete_endpoint(
|
|
self._request_for(self._frames(count=2), "ffv1", count=5)
|
|
)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("frames", response.payload["message"])
|
|
self.assertTrue(self.queue.empty())
|
|
|
|
async def test_rejects_bad_envelope(self):
|
|
response = await job_routes.job_complete_endpoint(
|
|
self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
|
|
)
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("format", response.payload["message"])
|
|
self.assertIn("is_last", response.payload["message"])
|
|
|
|
async def test_rejects_invalid_envelope_json(self):
|
|
response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
|
|
|
|
self.assertEqual(response.status, 400)
|
|
self.assertIn("envelope", response.payload["message"])
|
|
|
|
async def test_unknown_job_returns_404(self):
|
|
job_routes.prompt_server.distributed_pending_jobs = {}
|
|
with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
|
|
response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
|
|
|
|
self.assertEqual(response.status, 404)
|
|
|
|
async def test_legacy_json_payload_still_accepted(self):
|
|
request = _FakeRequest(
|
|
{
|
|
"job_id": "job-z",
|
|
"worker_id": "worker-1",
|
|
"batch_idx": 3,
|
|
"image": "data:image/png;base64,AAAA",
|
|
"is_last": True,
|
|
}
|
|
)
|
|
with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
|
|
response = await job_routes.job_complete_endpoint(request)
|
|
|
|
self.assertEqual(response.status, 200)
|
|
items = await self._drain()
|
|
self.assertEqual(len(items), 1)
|
|
self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|