diff --git a/api/image_list_routes.py b/api/image_list_routes.py index 92a303d..2a3569a 100644 --- a/api/image_list_routes.py +++ b/api/image_list_routes.py @@ -5,6 +5,7 @@ from aiohttp import web import server from ..image_list.runtime import ( + IMAGE_TRANSFER_FORMATS, MAX_REQUEST_BYTES, PROTOCOL_VERSION, decode_result, @@ -39,6 +40,7 @@ async def image_list_capabilities_endpoint(_request): "protocol_version": PROTOCOL_VERSION, "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"], "transport": "typed-float-tensor-json-v1", + "image_transfer_formats": list(IMAGE_TRANSFER_FORMATS), }) diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 4213f31..c91b9a2 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -188,7 +188,8 @@ async def dispatch_worker_prompt( return prompt_id if return_prompt_id else None -async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): +async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *, + required_transfer_formats=()): """Require protocol/node support and report which direct workers are idle. This performs read-only GETs only; it never submits a prompt or contacts a @@ -196,6 +197,14 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): """ if not workers: raise RuntimeError("Distributed Image List requires at least one selected direct worker") + if (not isinstance(required_transfer_formats, (list, tuple, set, frozenset)) + or any(type(fmt) is not str for fmt in required_transfer_formats)): + raise ValueError("required_transfer_formats must contain image transfer format names") + required_formats = set(required_transfer_formats) - {"exact"} + if required_formats: + from ...serverless_bridge.image_transfer import validate_options + for fmt in required_formats: + validate_options(fmt) session = await get_client_session() required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"} semaphore = asyncio.Semaphore(8) @@ -242,6 +251,17 @@ async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0): "upgrade that worker before running this feature; no work was submitted" ) + if required_formats: + advertised = capability.get("image_transfer_formats") + if (not isinstance(advertised, list) + or any(type(fmt) is not str for fmt in advertised) + or not required_formats.issubset(advertised)): + raise RuntimeError( + f"Image-list worker {worker_id} lacks required image_transfer_formats " + f"{sorted(required_formats)}; upgrade that worker before compressed IMAGE returns; " + "no work was submitted" + ) + check = "queue" queue_url = build_worker_url(worker, "/queue") async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response: diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 002ef68..99c711c 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -361,6 +361,14 @@ async def orchestrate_distributed_execution( "DistributedMediaCollector or Ultimate SD Upscale nodes in one prompt" ) validate_image_list_regions(prompt_obj, explicit_targets=True) + from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options + required_transfer_formats = set() + for end_id in prompt_index.nodes_for_class("DistributedImageListEnd"): + end_inputs = prompt_obj[end_id].get("inputs", {}) + fmt, _ = validate_options(end_inputs.get("transfer_format", DEFAULT_FORMAT), + end_inputs.get("transfer_quality", DEFAULT_QUALITY)) + if fmt != DEFAULT_FORMAT: + required_transfer_formats.add(fmt) workers = _resolve_direct_workers_for_image_list(config, requested_ids) if _has_selected_serverless_workers(config, requested_ids): from ..image_list.serverless import resolve_serverless_workers @@ -421,7 +429,9 @@ async def orchestrate_distributed_execution( "No selected worker is reachable for Distributed Image List in orchestrator-only mode" ) direct_workers = [worker for worker in active_workers if worker.get("type") != "serverless"] - idle_ids = (await preflight_image_list_workers(direct_workers, require_idle=delegate_master) + transfer_kwargs = ({"required_transfer_formats": sorted(required_transfer_formats)} + if required_transfer_formats else {}) + idle_ids = (await preflight_image_list_workers(direct_workers, require_idle=delegate_master, **transfer_kwargs) if direct_workers else {}) active_workers = [worker for worker in active_workers if worker.get("type") == "serverless" or worker["id"] in idle_ids] diff --git a/docs/image-list-queue.md b/docs/image-list-queue.md index e20acfc..d22dc74 100644 --- a/docs/image-list-queue.md +++ b/docs/image-list-queue.md @@ -62,6 +62,32 @@ End's optional `texts` is a connection-only socket. For workflows saved before t File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Direct-worker copies use separate content-addressed paths under `input/distributed-media/`; each direct worker's rewritten prompt is reused across items without reusing item tokens. Cloud requests use the PR2 bounded file-input adapters instead. +## Image return format and quality + +End adds `transfer_format` (default **exact**) and `transfer_quality` (default **90**). +These apply to returned images from **both direct workers and serverless pools**. +Master-local results bypass encoding. Input images, masks, generic values and +returned masks/text retain their exact existing transport. + +- `exact` preserves floating-point dtype and values. +- `legacy_png`, `png`, `ffv1` preserve quantized 8-bit image pixels, not the + original floating-point colors. `legacy_png` is PNG without compression. +- `jpeg`, `webp`, `h264`, `h265`, `av1` are lossy choices; quality is 1–100. + Video codecs encode each assigned image independently, not the whole source list. +- Dimensions, floating dtype and alpha survive; alpha is an exact sidecar. +- Direct workers must advertise the selected codec in `image_transfer_formats`; + older workers fail before submission for non-exact choices. Exact remains + compatible with the existing protocol. Serverless endpoints need the updated + complete bridge; health probing does not prove codec availability. +- The selected software encoder must be installed. Encoding failures stop the + item; there is no silent format fallback. No hardware encoder is selected. + +Compression can reduce transfer size, but does **not** bypass decoded tensor, +conservative codec scratch, or complete request limits. For details and codec +limits, see [image return options](serverless-media.md#image-return-format-and-quality). +Mixed local/remote results may have different color fidelity when compression is +selected, because local images are deliberately not quantized. + ## Serverless pools and bounded capacity Select `media_profile: "comfy-media-v1"` for each cloud pool. The endpoint must @@ -87,7 +113,8 @@ Provider health/authentication probing does not verify bridge/model installation - Each cloud request uses the standard media-v1 envelope and existing durable serverless job lifecycle. The private item graph replaces Start/End with typed inputs and explicit exporter roots, not independent all-output execution. - IMAGE/MASK tensor data and text are preserved by the typed codecs. There is no + IMAGE returns are exact by default, or explicitly encoded as selected at End. + MASK tensor data, input images and text retain exact typed codecs. There is no image-list-specific endpoint protocol or new visible node socket. - Cloud `value_N` tensors must have native IMAGE `[B,H,W,C]` or MASK `[B,H,W]` shapes, not arbitrary rank-one/rank-two tensor payloads. Typed tensors must be @@ -118,7 +145,7 @@ Provider health/authentication probing does not verify bridge/model installation - Queue preparation works on an execution copy, leaving the saved canvas unchanged. Unused per-item preview/save branches are pruned by dependency reachability, not node names. Required computational utilities marked `OUTPUT_NODE` stay in the graph; explicit `partial_execution_targets` prevent them from also becoming independent outer-master roots. End and legitimate post-End or unrelated outputs are retained. Bypass/rejoin branches still fail boundary validation rather than being silently removed. - Every selected direct worker must advertise `execution_targets: true` as well as the existing image-list protocol. Older workers fail preflight before dispatch. Both HTTP and WebSocket item submissions target End only; the participating master uses its existing End-only local executor. Cloud jobs use explicit media exporter roots. Submit through `/distributed/queue`: an ordinary all-output `/prompt` request does not receive this relaxation. - **Target selection is not a sandbox.** A required utility still runs its own save, preview or other side effects when it executes for an item. Keep intentional final saving after End. -- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Direct-worker transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each direct-worker tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. Serverless inline limits are smaller, as described above. In-process master items use tensors directly and are not subject to those HTTP limits. +- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Direct-worker input transport and default exact returns preserve floating-point tensor dtype and values. Opt-in image return encoding does not relax decoded-size limits. Each direct-worker tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. Serverless inline limits are smaller, as described above. In-process master items use tensors directly and are not subject to those HTTP limits. - Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically. - Worker errors fail the region. For direct workers, the worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue. Serverless jobs instead use their configured serverless timeout and durable manager lifecycle; they are excluded from direct-history/liveness polling. - User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are ambiguous, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network. diff --git a/image_list/runtime.py b/image_list/runtime.py index be78017..ae15cf0 100644 --- a/image_list/runtime.py +++ b/image_list/runtime.py @@ -15,6 +15,15 @@ import time from dataclasses import dataclass, field from typing import Any +try: + from ..serverless_bridge.image_transfer import ( + FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options, + ) +except ImportError: # Standalone image_list.runtime (without extension bootstrap). + from serverless_bridge.image_transfer import ( + FORMATS as IMAGE_TRANSFER_FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options, + ) + # v2: Start carries generic value_N slots instead of fixed reference/text/strength/seed fields. PROTOCOL_VERSION = 2 START_CLASS = "DistributedImageListStart" @@ -257,6 +266,8 @@ def decode_value(data: Any, expected_kind: str) -> Any: return None if not isinstance(data, dict) or data.get("type") != expected_kind: raise ValueError(f"Expected {expected_kind} transport value") + if "encoding" in data or "descriptor" in data: + raise ValueError("Compressed transport is supported only for the returned IMAGE") if expected_kind in ("IMAGE", "MASK", "TENSOR"): torch = _torch() dtype_name = data.get("dtype") @@ -734,10 +745,48 @@ def register_assignment(context_id: str, item_index: int, worker_id: str, token: state.expected_workers[item_index] = worker_id +def _image_codec(): + # Exact returns and inputs must not import optional media codec dependencies. + if __package__ == "image_list": + from serverless_bridge import image_codec + else: + from ..serverless_bridge import image_codec + return image_codec + + +def _decode_return_image(data: Any) -> Any: + if (not isinstance(data, dict) or set(data) != {"type", "encoding", "descriptor", "data"} + or data["type"] != "IMAGE" or data["encoding"] != "image-transfer-v1"): + raise ValueError("Invalid compressed IMAGE envelope") + descriptor = data["descriptor"] + if not isinstance(descriptor, dict) or descriptor.get("encoding") != "image-transfer-v1": + raise ValueError("Invalid compressed IMAGE descriptor") + shape = descriptor.get("shape") + if (not isinstance(shape, list) or len(shape) != 4 + or any(type(dim) is not int or dim < 1 for dim in shape) + or shape[0] != 1 or shape[-1] not in (3, 4)): + raise ValueError("Compressed IMAGE must have shape [1,H,W,3|4]") + encoded = data["data"] + if not isinstance(encoded, str) or len(encoded) > ((MAX_PAYLOAD_BYTES + 2) // 3) * 4: + raise ValueError("Compressed IMAGE data is missing or exceeds the transport limit") + try: + raw = base64.b64decode(encoded, validate=True) + except (ValueError, TypeError) as exc: + raise ValueError("Compressed IMAGE data must be valid base64") from exc + if len(raw) > MAX_PAYLOAD_BYTES: + raise ValueError("Compressed IMAGE exceeds the transport limit") + return _image_codec().decode_image(descriptor, raw, max_bytes=MAX_PAYLOAD_BYTES) + + def decode_result(result: Any) -> dict[str, Any]: if not isinstance(result, dict): raise ValueError("result must be an object") - image = decode_value(result.get("image"), "IMAGE") + image_data = result.get("image") + if isinstance(image_data, dict) and ("encoding" in image_data or "descriptor" in image_data): + encode_request_payload(result) # Bound combined JSON/base64 before codec allocation. + image = _decode_return_image(image_data) + else: + image = decode_value(image_data, "IMAGE") if image is None or image.ndim != 4 or image.shape[0] != 1 or image.shape[-1] not in (3, 4): raise ValueError("result.image must contain exactly one RGB/RGBA image") mask = decode_value(result.get("mask"), "MASK") @@ -773,11 +822,28 @@ def normalize_worker_result(images: Any, masks: Any = None, texts: Any = None) - return {"image": image, "mask": mask, "text": text} -def encode_result(result: dict[str, Any]) -> dict[str, Any]: +def encode_result(result: dict[str, Any], *, transfer_format=DEFAULT_FORMAT, + transfer_quality=DEFAULT_QUALITY) -> dict[str, Any]: + fmt, quality = validate_options(transfer_format, transfer_quality) normalized = normalize_worker_result(result.get("image"), result.get("mask"), result.get("text")) - return {"image": encode_value(normalized["image"], "IMAGE"), - "mask": encode_value(normalized["mask"], "MASK"), - "text": encode_value(normalized["text"], "STRING")} + if fmt == "exact": + image = encode_value(normalized["image"], "IMAGE") + else: + value = normalized["image"] + # Reject before moving/copying a potentially oversized device tensor. + if value.numel() * value.element_size() > MAX_PAYLOAD_BYTES: + raise ValueError("IMAGE item exceeds the image-list transport limit") + descriptor, raw = _image_codec().encode_image( + value.detach().to("cpu").contiguous(), transfer_format=fmt, + transfer_quality=quality, max_bytes=MAX_PAYLOAD_BYTES, + ) + image = {"type": "IMAGE", "encoding": "image-transfer-v1", "descriptor": descriptor, + "data": base64.b64encode(raw).decode("ascii")} + encoded = {"image": image, "mask": encode_value(normalized["mask"], "MASK"), + "text": encode_value(normalized["text"], "STRING")} + if fmt != "exact": + encode_request_payload(encoded) + return encoded def _collect_output_lists(results: list[dict[str, Any]]) -> tuple[list[Any], list[Any], list[Any]]: @@ -872,8 +938,17 @@ async def _run_image_list_slots(state, config, workers, region, controller): import aiohttp serverless_slots = list(controller.slots) if controller is not None else [] + end_inputs = state.original_prompt[state.end_id]["inputs"] + transfer_format, _ = validate_options( + end_inputs.get("transfer_format", DEFAULT_FORMAT), + end_inputs.get("transfer_quality", DEFAULT_QUALITY), + ) + transfer_requirements = ({"required_transfer_formats": (transfer_format,)} + if transfer_format != "exact" else {}) # Direct capability/idle probes must never be sent to a serverless pool. - worker_statuses = await preflight_image_list_workers(workers, require_idle=not state.master_participates) if workers else {} + worker_statuses = await preflight_image_list_workers( + workers, require_idle=not state.master_participates, **transfer_requirements, + ) if workers else {} workers = [worker for worker in workers if worker["id"] in worker_statuses] if not workers and not serverless_slots and not state.master_participates: raise RuntimeError("No reachable, feature-compatible image-list workers are available; master execution is disabled") @@ -913,6 +988,13 @@ async def _run_image_list_slots(state, config, workers, region, controller): if worker is None: raise RuntimeError(f"Worker {worker_id} was removed before dispatch") assert session is not None + if transfer_requirements: + # Recheck on every compressed assignment, before staging or POSTing. + statuses = await preflight_image_list_workers( + [worker], require_idle=True, **transfer_requirements, + ) + if worker_id not in statuses: + raise RuntimeError(f"Image-list worker {worker_id} unavailable before compressed return dispatch") token = create_worker_task_token() register_assignment(context_id, index, worker_id, token) item_payload = {"context_id": context_id, "item": encode_item(state.items[index])} diff --git a/image_list/serverless.py b/image_list/serverless.py index 360cccf..867c13f 100644 --- a/image_list/serverless.py +++ b/image_list/serverless.py @@ -150,11 +150,17 @@ async def prepare_serverless_item(prompt, region, item, worker, *, request_id, j raise ValueError("Image-list region is missing its boundary nodes") graph = {nid: copy.deepcopy(prompt[nid]) for nid in sorted(ids)} end_inputs = graph[end].get("inputs", {}) + from ..serverless_bridge.image_transfer import DEFAULT_FORMAT, DEFAULT_QUALITY, validate_options + transfer_format, transfer_quality = validate_options( + end_inputs.get("transfer_format", DEFAULT_FORMAT), + end_inputs.get("transfer_quality", DEFAULT_QUALITY)) if "images" not in end_inputs: raise ValueError("Image-list End requires images") graph[end] = {"class_type": "DistributedMediaCollector", "inputs": { **{key: end_inputs[key] for key in ("images", "masks", "texts") if key in end_inputs}, "multi_job_id": job_id, "worker_id": worker["id"]}} + if transfer_format != DEFAULT_FORMAT: + graph[end]["inputs"].update(transfer_format=transfer_format, transfer_quality=transfer_quality) del graph[start] if type(item.get("index")) is not int or item["index"] < 0: raise ValueError("Image-list item index must be a nonnegative integer") diff --git a/nodes/image_list.py b/nodes/image_list.py index 1057b3c..3c39ae4 100644 --- a/nodes/image_list.py +++ b/nodes/image_list.py @@ -6,6 +6,8 @@ import json from comfy_api.v0_0_2 import io from ..utils.async_helpers import run_async_in_server_loop +from ..serverless_bridge.image_transfer import (FORMATS, DEFAULT_FORMAT, DEFAULT_QUALITY, + FORMAT_TOOLTIP, QUALITY_TOOLTIP, validate_options) from ..image_list.runtime import ( MAX_VALUE_SLOTS, VALUE_NAMES, @@ -155,6 +157,10 @@ class DistributedImageListEnd(io.ComfyNode): io.Mask.Input("masks", optional=True, lazy=True), io.String.Input("texts", optional=True, lazy=True, force_input=True), io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"), + io.Combo.Input("transfer_format", options=list(FORMATS), default=DEFAULT_FORMAT, + optional=True, tooltip=FORMAT_TOOLTIP), + io.Int.Input("transfer_quality", default=DEFAULT_QUALITY, min=1, max=100, + optional=True, tooltip=QUALITY_TOOLTIP), ], outputs=[ io.Image.Output(display_name="image_list", is_output_list=True), @@ -195,7 +201,10 @@ class DistributedImageListEnd(io.ComfyNode): @classmethod def execute(cls, images=None, masks=None, texts=None, context=None, worker_result_token="", worker_item_index=-1, master_callback_url="", - image_list_local_execution=False): + image_list_local_execution=False, transfer_format=DEFAULT_FORMAT, + transfer_quality=DEFAULT_QUALITY): + transfer_format, transfer_quality = validate_options(unwrap_single(transfer_format), + unwrap_single(transfer_quality)) worker_result_token = unwrap_single(worker_result_token) master_callback_url = unwrap_single(master_callback_url) if worker_result_token: @@ -209,7 +218,8 @@ class DistributedImageListEnd(io.ComfyNode): set_master_result(context_values[0], int(unwrap_single(worker_item_index)), str(worker_result_token), result) else: - payload = encode_result(result) + payload = encode_result(result, transfer_format=transfer_format, + transfer_quality=transfer_quality) run_async_in_server_loop( submit_worker_result(str(master_callback_url), str(worker_result_token), context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60, diff --git a/tests/helpers/image_list_network_check.py b/tests/helpers/image_list_network_check.py index 1f58ecf..7b9ed41 100644 --- a/tests/helpers/image_list_network_check.py +++ b/tests/helpers/image_list_network_check.py @@ -458,7 +458,14 @@ async def coordinator(): ALLOW_MASTER = True config['settings']['master_delegate_only'] = False participation_counts = {} - for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]: + for mode, selected, fmt in [('mixed', [w['id'] for w in workers], 'exact'), + ('master-only', [], 'exact'), + ('mixed-png', [w['id'] for w in workers], 'png'), + ('master-only-png', [], 'png')]: + case_prompt = copy.deepcopy(prompt) + if fmt != 'exact': + case_prompt['4']['inputs'].update(transfer_format=fmt, transfer_quality=73) + saved_case = copy.deepcopy(case_prompt) executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets) before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS) before_utilities, before_events = len(UTILITY_CALLS), len(ITEM_EXECUTIONS) @@ -472,8 +479,8 @@ async def coordinator(): patch.object(orchestration, 'queue_prompt_payload', enqueue), \ patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15): queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution( - prompt, None, 'test', selected), 45) - assert prompt == saved_prompt, (mode, 'saved/source graph changed') + case_prompt, None, 'test', selected), 45) + assert case_prompt == saved_case and prompt == saved_prompt, (mode, 'saved/source graph changed') assert executor.success, (mode, executor.status_messages) assert queued[2] == len(selected) assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item' @@ -499,7 +506,7 @@ async def coordinator(): assert '5' not in submission['node_ids'] remote_calls.extend(calls) assert not current['running'] and not current['pending'], current - if mode == 'mixed': + if mode.startswith('mixed'): assert remote_calls, 'mixed execution failed to use direct workers' else: assert not remote_calls and len(local_calls) == COUNT @@ -507,13 +514,18 @@ async def coordinator(): images, masks, texts = CAPTURED[-1] assert len(images) == len(masks) == len(texts) == COUNT for i, row in enumerate(source_rows()): - assert torch.equal(images[i], row['image'] + row['strength']), (mode, i) + expected = row['image'] + row['strength'] + if fmt == 'png' and i in remote_calls: + expected = expected.clone() + expected[..., :3] = (expected[..., :3].clamp(0, 1) * 255).to(torch.uint8).float() / 255 + assert torch.equal(images[i], expected), (mode, i) assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i) assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts) assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)} print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts)) - assert len(outer_executions) == 4 and all(row['captures'] == 1 for row in outer_executions) + assert len(outer_executions) == 2 + len(participation_counts) and all(row['captures'] == 1 for row in outer_executions) + print('IMAGE_LIST_TRANSFER_HTTP_OK mixed-png and master-only-png; masks/text/alpha/local images exact') print('IMAGE_LIST_TARGETS_OK ' + json.dumps({ 'modes': ['delegate-only', 'mixed', 'master-only'], 'outer_targets': outer_executions[0]['targets'], 'worker_targets': ['4'], diff --git a/tests/test_image_list_transfer_options.py b/tests/test_image_list_transfer_options.py new file mode 100644 index 0000000..94972b6 --- /dev/null +++ b/tests/test_image_list_transfer_options.py @@ -0,0 +1,70 @@ +"""End policy reaches cloud exports and initial worker compatibility preflight.""" +import asyncio +import copy +import unittest + +import pytest +import torch + +import test_image_list_serverless as cloud +import test_image_list_master_selection as selection + +FORMATS = ('exact', 'legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1') + + +@pytest.mark.parametrize('fmt', FORMATS) +def test_cloud_returns_use_selected_image_format_but_inputs_remain_exact(fmt): + async def check(): + prompt = cloud.source() + prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=73) + original = copy.deepcopy(prompt) + value = cloud.item(dtype=torch.float64) + job = await cloud.prepare(prompt, value) + assert prompt == original + for binding in job['manifest']: + controls = job['workflow'][binding['node_id']]['inputs'] + if binding['kind'] == 'IMAGE' and fmt != 'exact': + assert controls['transfer_format'] == fmt + assert controls['transfer_quality'] == 73 + else: + assert 'transfer_format' not in controls and 'transfer_quality' not in controls + for artifact in job['media']['artifacts']: + assert artifact['descriptor']['encoding'] == 'tensor-le-v1' + inputs = await cloud.decoded_inputs(job) + for ref in job['media']['inputs']: + assert torch.equal(inputs[ref['artifact_id']], value[ref['input_name']]) + cloud.handler._validate_request({k: job[k] for k in ('workflow', 'images', 'media')}) + asyncio.run(check()) + + +@pytest.mark.parametrize('fmt,quality', [('JPEG',90),('jpeg',True),('jpeg',101),('jpeg','90')]) +def test_cloud_rejects_invalid_policy(fmt, quality): + prompt = cloud.source() + prompt['ns/end']['inputs'].update(transfer_format=fmt, transfer_quality=quality) + with pytest.raises(ValueError): + asyncio.run(cloud.prepare(prompt)) + + +class TransferPreflightTests(unittest.IsolatedAsyncioTestCase): + asyncSetUp = selection.MasterSelectionTests.asyncSetUp + queue = selection.MasterSelectionTests.queue + + async def test_queue_requests_image_codec_capability(self): + self.prompt['3']['inputs'].update(transfer_format='jpeg', transfer_quality=73) + await self.queue() + self.assertEqual(self.preflight.call_args.kwargs['required_transfer_formats'], ['jpeg']) + + async def test_invalid_policy_fails_before_worker_probe_or_queue(self): + self.prompt['3']['inputs'].update(transfer_format='bad-codec') + with self.assertRaises(ValueError): + await self.queue() + self.q.select_active_workers.assert_not_awaited() + self.preflight.assert_not_awaited() + self.q.queue_prompt_payload.assert_not_awaited() + + async def test_missing_capability_cannot_silently_fall_back_to_master(self): + self.prompt['3']['inputs'].update(transfer_format='jpeg') + self.preflight.side_effect = RuntimeError('upgrade worker for image return encoding') + with self.assertRaisesRegex(RuntimeError, 'upgrade'): + await self.queue() + self.q.queue_prompt_payload.assert_not_awaited() diff --git a/tests/test_image_list_transfer_runtime.py b/tests/test_image_list_transfer_runtime.py new file mode 100644 index 0000000..1abc3ff --- /dev/null +++ b/tests/test_image_list_transfer_runtime.py @@ -0,0 +1,265 @@ +"""CPU-only direct-worker IMAGE return codecs and capability negotiation.""" +import asyncio +import base64 +import importlib.util +import json +from pathlib import Path +import sys +import types +from unittest.mock import AsyncMock, Mock, patch + +import pytest +import torch + +from image_list import runtime as rt +from serverless_bridge import image_codec +from serverless_bridge.image_transfer import FORMATS + +ROOT = Path(__file__).resolve().parents[1] + + +def result(channels=4, dtype=torch.float32): + image = torch.linspace(-0.1, 1.1, 17 * 19 * channels, dtype=dtype).reshape(1, 17, 19, channels) + return {"image": image, "mask": torch.full((1, 17, 19), 0.123456789, dtype=torch.float64), + "text": "exact text: \u00e9\n"} + + +def legacy_value(value, kind): + if isinstance(value, str): + return {"type": kind, "value": value} + raw = value.contiguous().view(torch.uint8).numpy().tobytes() + return {"type": kind, "dtype": str(value.dtype).removeprefix("torch."), + "shape": list(value.shape), "data": base64.b64encode(raw).decode("ascii")} + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_exact_result_and_inputs_are_byte_for_byte_legacy(dtype): + source = result(dtype=dtype) + expected = {name: legacy_value(source[name], kind) + for name, kind in (("image", "IMAGE"), ("mask", "MASK"), ("text", "STRING"))} + body = rt.encode_request_payload(expected) + with patch.object(image_codec, "encode_image", side_effect=AssertionError("exact used codec")): + assert rt.encode_request_payload(rt.encode_result(source)) == body + assert rt.encode_request_payload(rt.encode_result(source, transfer_format="exact", transfer_quality=90)) == body + decoded = rt.decode_result(expected) + assert torch.equal(decoded["image"].view(torch.uint8), source["image"].view(torch.uint8)) + item = {"image": source["image"], "mask": source["mask"], "index": 0, "values": {"value_1": source["image"]}} + encoded = rt.encode_item(item) + assert encoded["image"] == expected["image"] + assert encoded["mask"] == expected["mask"] + assert encoded["values"]["value_1"] == legacy_value(source["image"], "TENSOR") + + +@pytest.mark.parametrize("fmt", FORMATS[1:]) +@pytest.mark.parametrize("channels", [3, 4]) +def test_all_compressed_formats_roundtrip_only_main_image(fmt, channels): + source = result(channels) + encoded = rt.encode_result(source, transfer_format=fmt, transfer_quality=90) + image = encoded["image"] + assert set(image) == {"type", "encoding", "descriptor", "data"} + assert image["type"] == "IMAGE" + assert image["encoding"] == "image-transfer-v1" + assert image["descriptor"]["format"] == fmt + assert image["descriptor"]["shape"] == [1, 17, 19, channels] + assert image["descriptor"]["quality"] == 90 + assert base64.b64decode(image["data"], validate=True) + decoded = rt.decode_result(json.loads(rt.encode_request_payload(encoded))) + assert decoded["image"].shape == source["image"].shape + assert decoded["image"].dtype == source["image"].dtype + assert decoded["image"].device.type == "cpu" + assert torch.isfinite(decoded["image"]).all() + pixels = source["image"][..., :3].float().mul(255).clamp(0, 255).to(torch.uint8).float().div(255) + if fmt in ("legacy_png", "png", "ffv1"): + assert torch.equal(decoded["image"][..., :3], pixels) + else: + assert (decoded["image"][..., :3] - pixels).abs().mean() < 0.08 + if channels == 4: + assert torch.equal(decoded["image"][..., 3], source["image"][..., 3]) + assert encoded["mask"] == legacy_value(source["mask"], "MASK") + assert torch.equal(decoded["mask"], source["mask"]) + assert decoded["text"] == source["text"] + + +@pytest.mark.parametrize("fmt,quality", [("nope", 90), (["node", 0], 90), ("png", True), ("png", 0), ("exact", 101)]) +def test_invalid_options_fail_closed(fmt, quality): + with pytest.raises(ValueError): + rt.encode_result(result(), transfer_format=fmt, transfer_quality=quality) + + +def test_compressed_inputs_masks_and_generic_tensors_are_rejected(): + source = result() + encoded = rt.encode_result(source, transfer_format="png") + for kind in ("IMAGE", "MASK", "TENSOR", "STRING"): + with pytest.raises(ValueError): + rt.decode_value({**encoded["image"], "type": kind}, kind) + item = rt.encode_item({"image": source["image"], "index": 0}) + with pytest.raises(ValueError): + rt.decode_item({**item, "image": encoded["image"]}) + for key, kind in (("mask", "MASK"), ("text", "STRING")): + with pytest.raises(ValueError): + rt.decode_result({**encoded, key: {**encoded["image"], "type": kind}}) + + +@pytest.mark.parametrize("shape", [[2, 17, 19, 3], [1, 17, 19, 1], [17, 19, 3], [True, 17, 19, 3], [1, 0, 19, 3]]) +def test_invalid_image_shapes_rejected_before_codec_decode(shape): + encoded = rt.encode_result(result(), transfer_format="png") + encoded["image"]["descriptor"]["shape"] = shape + with patch.object(image_codec, "decode_image", side_effect=AssertionError("native codec reached")): + with pytest.raises(ValueError): + rt.decode_result(encoded) + + +@pytest.mark.parametrize("change", [ + {"encoding": "image-transfer-v2"}, {"type": "MASK"}, {"extra": 1}, {"data": "%%%"}, + {"descriptor": None}, {"descriptor": {"encoding": "tensor-le-v1", "shape": [1, 17, 19, 4]}}, +]) +def test_malformed_compressed_envelopes_rejected(change): + encoded = rt.encode_result(result(), transfer_format="png") + encoded["image"].update(change) + with pytest.raises(ValueError): + rt.decode_result(encoded) + + +def test_single_frame_and_image_list_not_bridge_budgets(monkeypatch): + with pytest.raises(ValueError): + rt.encode_result({"image": torch.zeros(2, 17, 19, 3)}, transfer_format="png") + with patch.object(image_codec, "encode_image", wraps=image_codec.encode_image) as encode: + encoded = rt.encode_result(result(), transfer_format="png") + assert encode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES + with patch.object(image_codec, "decode_image", wraps=image_codec.decode_image) as decode: + rt.decode_result(encoded) + assert decode.call_args.kwargs["max_bytes"] == rt.MAX_PAYLOAD_BYTES + monkeypatch.setattr(rt, "MAX_PAYLOAD_BYTES", 64) + with pytest.raises(ValueError, match="limit"): + rt.encode_result({"image": torch.empty(1, 4, 4, 3, device="meta")}, transfer_format="png") + with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")): + with pytest.raises(ValueError, match="limit"): + rt.decode_result(encoded) + + +def test_compressed_request_combined_budget(monkeypatch): + encoded = rt.encode_result(result(), transfer_format="png") + monkeypatch.setattr(rt, "MAX_REQUEST_BYTES", len(rt.encode_request_payload(encoded)) - 1) + with pytest.raises(ValueError, match="request"): + rt.encode_result(result(), transfer_format="png") + with patch.object(image_codec, "decode_image", side_effect=AssertionError("oversize reached codec")): + with pytest.raises(ValueError, match="request"): + rt.decode_result(encoded) + + +@pytest.fixture +def dispatch_module(monkeypatch): + name = "_image_list_transfer_tests" + for suffix in ("", ".api", ".api.orchestration", ".utils"): + module = types.ModuleType(name + suffix) + module.__path__ = [str(ROOT / suffix.lstrip(".").replace(".", "/"))] + monkeypatch.setitem(sys.modules, name + suffix, module) + monkeypatch.setitem(sys.modules, name + ".utils.logging", types.SimpleNamespace(log=lambda *a: None, debug_log=lambda *a: None)) + monkeypatch.setitem(sys.modules, name + ".utils.network", types.SimpleNamespace( + build_worker_url=lambda worker, path="": "http://" + worker["id"] + path, + get_client_session=AsyncMock(), probe_worker=AsyncMock())) + spec = importlib.util.spec_from_file_location(name + ".api.orchestration.dispatch", ROOT / "api/orchestration/dispatch.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def capability(): + return {"supported": True, "protocol_version": 2, "execution_targets": True, + "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"]} + + +class Response: + status = 200 + + def __init__(self, value): + self.value = value + + async def __aenter__(self): + return self + + async def __aexit__(self, *args): + pass + + async def json(self): + return self.value + + +@pytest.mark.parametrize("advertised", [None, [], ["exact"], "png", ["png", 3]]) +def test_nonexact_requires_advertised_capability_before_queue(dispatch_module, advertised): + cap = capability() + if advertised is not None: + cap["image_transfer_formats"] = advertised + session = types.SimpleNamespace(get=Mock(return_value=Response(cap)), post=Mock()) + dispatch_module.get_client_session = AsyncMock(return_value=session) + with pytest.raises(RuntimeError, match="upgrade.*no work was submitted"): + asyncio.run(dispatch_module.preflight_image_list_workers( + [{"id": "worker"}], required_transfer_formats=["png"])) + assert session.get.call_count == 1 + session.post.assert_not_called() + + +@pytest.mark.parametrize("requested", [(), ("exact",), ("png", "jpeg")]) +def test_exact_legacy_or_advertised_nonexact_capability_succeeds(dispatch_module, requested): + cap = capability() + if "png" in requested: + cap["image_transfer_formats"] = list(FORMATS) + session = types.SimpleNamespace(get=Mock(side_effect=[Response(cap), Response({"queue_running": [], "queue_pending": []})])) + dispatch_module.get_client_session = AsyncMock(return_value=session) + assert asyncio.run(dispatch_module.preflight_image_list_workers( + [{"id": "worker"}], required_transfer_formats=requested)) == {"worker": True} + + +def test_capability_endpoint_advertises_shared_formats(monkeypatch): + from aiohttp import web + name = "_transfer_routes" + monkeypatch.setitem(sys.modules, name, types.ModuleType(name)) + monkeypatch.setitem(sys.modules, name + ".api", types.ModuleType(name + ".api")) + monkeypatch.setitem(sys.modules, name + ".image_list", types.ModuleType(name + ".image_list")) + monkeypatch.setitem(sys.modules, name + ".image_list.runtime", rt) + monkeypatch.setitem(sys.modules, "server", types.SimpleNamespace(PromptServer=types.SimpleNamespace( + instance=types.SimpleNamespace(routes=web.RouteTableDef())))) + spec = importlib.util.spec_from_file_location(name + ".api.image_list_routes", ROOT / "api/image_list_routes.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + response = asyncio.run(module.image_list_capabilities_endpoint(None)) + assert json.loads(response.text)["image_transfer_formats"] == list(FORMATS) + + +@pytest.mark.parametrize("reject_call", [None, 1, 2, 3]) +def test_runtime_initial_and_each_nonexact_assignment_preflight(reject_call): + from test_image_list_runtime_boundaries import RuntimeBoundaryTests + + async def exercise(): + fixture = RuntimeBoundaryTests() + await fixture.asyncSetUp() + try: + context = fixture.create(2) + state = fixture.rt.get_master_job(context) + state.original_prompt[state.end_id]["inputs"].update(transfer_format="png", transfer_quality=87) + calls = [] + + async def preflight(workers, **kwargs): + calls.append(kwargs) + assert kwargs["required_transfer_formats"] == ("png",) + if len(calls) == reject_call: + raise RuntimeError("upgrade worker; no work was submitted") + return {"worker": True} + + sys.modules["_image_list_boundaries.api.orchestration.dispatch"].preflight_image_list_workers = preflight + if reject_call: + with pytest.raises(RuntimeError, match="upgrade"): + await fixture.rt.run_image_list_job(context) + assert len(fixture.uploads) == max(0, reject_call - 2) + assert len(fixture.prompts) == max(0, reject_call - 2) + else: + await fixture.rt.run_image_list_job(context) + assert len(calls) == 3 + assert len(fixture.uploads) == 2 + assert not fixture.rt._MASTER_JOBS + finally: + # Do not call doCleanups(): this borrowed IsolatedAsyncioTestCase has + # no runner, but its sys.modules patch still must be restored. + fixture.patch_modules.stop() + + asyncio.run(exercise())