feat(collector): optional compressed worker → master result transfer (#97)
* feat(collector): compressed worker -> master result transfer Workers sent every frame as an uncompressed PNG (compress_level=0), base64-encoded inside a JSON request, one request per frame. For cloud workers this makes result download (egress) the dominant transfer cost: a 121-frame 720p batch is ~446 MB per worker. Add `transfer_format` / `transfer_quality` optional widgets to Distributed Collector: - legacy_png (default): unchanged behaviour - png, ffv1: lossless, bit-identical master tensors to legacy_png - jpeg, webp: lossy stills - h264, h265, av1: lossy video of the whole batch Compressed results are posted as multipart/form-data chunks (binary, no base64, ~32 MB soft cap per request, configurable with COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES). The master decodes each chunk off the event loop and enqueues one item per frame, so collection, ordering and audio handling are unchanged. Encoding also runs in an executor on the worker. Compatibility: - /distributed/job_complete still accepts the legacy JSON envelope. - A worker talking to an older master gets a 400 on the first chunk and resends everything as legacy_png. - Missing video encoders fall back to h264 (lossy) or png (ffv1). Measured on 121 frames at 1280x720: png 163 MB (2.7x), ffv1 128 MB (3.5x), h264 q90 2.4 MB (186x, 40.2 dB), av1 q90 0.4 MB. Also clips IMAGE values before 8-bit quantization on the new path so out-of-range values no longer wrap around. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6 * fix(collector): wire transfer options through V3 node; bound finalized payload size Addresses review on #97. V3 integration - Add optional `transfer_format` / `transfer_quality` inputs to the registered nodes/v3.py DistributedCollector schema (legacy defaults) and accept/forward them in execute() to the runtime helper. - tests/test_v3_collector_transfer.py checks the registered node's public schema and argument forwarding without a ComfyUI checkout; the opt-in real-runtime acceptance (tests/helpers/v3_runtime_check.py) now allows exactly these appended optional inputs, verifies their contract, and checks forwarding through the real executor. Payload size bound - encode_chunks() enforces limits on finalized bytes (after encoder flush and muxer finalization). Encoder lookahead kept the pre-flush buffer near zero, so a chunk could hold the whole batch and finalize far above the limit. Oversized video chunks are re-encoded with proportionally fewer frames; the next chunk is sized from observed bytes/frame. - Multi-frame chunks stay within TRANSFER_MAX_CHUNK_BYTES (32 MiB); a single frame may use up to TRANSFER_MAX_REQUEST_BYTES (96 MiB, new, COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES) minus envelope reserve. A frame that cannot fit raises FrameTooLargeError. - The worker encodes all chunks before sending, so an unfittable frame fails before anything is delivered (no partial batch on the master). - Audio rides on the last chunk only when the request stays in bounds; otherwise it is sent as a final audio-only completion. - Legacy fallback triggers only on an older master's rejection ("Invalid JSON payload"); other errors, e.g. size limits, raise with a pointer to the limit settings instead of resending as legacy_png. - The master returns 413 when a request exceeds --max-upload-size. Tests: codec flush-overshoot and single-frame cases, worker fallback / nothing-sent / audio placement, and real aiohttp server tests with client_max_size covering bounded high-entropy video, 413 reporting, and unfittable frames. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6 --------- Co-authored-by: ymiyad <ymiyad@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
ymiyad
co-authored by
Claude Opus 5.5
ymiyad
Claude Opus 5.5
ymiyad
parent
6f85129e45
commit
6edd3fda38
@@ -4,6 +4,8 @@ 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
|
||||
@@ -26,6 +28,17 @@ class _FakeRequest:
|
||||
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"
|
||||
@@ -46,6 +59,7 @@ def _load_job_routes_module():
|
||||
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
|
||||
@@ -394,5 +408,144 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user