robertvoy-ComfyUI-Distributed/tests/api/test_distributed_queue.py

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)

    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()