nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(collector): optional compressed worker → master result transfer (#97)

Browse Source 
* 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
2026-10-11 20:50:23 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 ymiyad
parent 6f85129e45
commit 6edd3fda38
12 changed files with 1839 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/job_routes.py
+153 -24
View File
@@ -15,6 +15,7 @@ from ..utils.logging import debug_log
from ..utils.image import pil_to_tensor, ensure_contiguous
from ..utils.network import handle_api_error
from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
from ..utils import transfer_codec
try:
from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
except ImportError:
@@ -26,8 +27,16 @@ from .queue_request import parse_queue_request_payload
prompt_server = server.PromptServer.instance
# Canonical worker result envelope accepted by POST /distributed/job_complete:
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
# Worker result envelopes accepted by POST /distributed/job_complete:
#
# application/json (legacy, one frame per request):
# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool,
# "audio"?: {...} }
#
# multipart/form-data (compressed, many frames per request):
# envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
# "is_last", "audio"?: {...} }
# media: one part per frame for still formats, or one video file for video formats
def _decode_image_sync(image_path):
@@ -270,8 +279,136 @@ async def check_file_endpoint(request):
return await handle_api_error(request, e, 500)
async def _enqueue_worker_items(multi_job_id, items):
"""Put decoded items on the job queue, waiting briefly for the master to create it.
Returns the queue size after enqueueing, or None if the job never appeared.
"""
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
while True:
async with prompt_server.distributed_jobs_lock:
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
if pending is not None:
for item in items:
await pending.put(item)
return pending.qsize()
if time.monotonic() > deadline:
return None
await asyncio.sleep(0.05)
def _decode_compressed_chunk_sync(envelope, parts):
"""Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
frames = transfer_codec.decode_chunk(envelope, parts)
return [
ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
for frame in frames
]
def _validate_compressed_envelope(envelope):
errors = []
if not isinstance(envelope, dict):
return ["envelope: expected JSON object"]
for key in ("job_id", "worker_id"):
value = envelope.get(key)
if not isinstance(value, str) or not value.strip():
errors.append(f"{key}: expected non-empty string")
start_index = envelope.get("start_index")
if not isinstance(start_index, int) or isinstance(start_index, bool) or start_index < 0:
errors.append("start_index: expected non-negative integer")
if not isinstance(envelope.get("is_last"), bool):
errors.append("is_last: expected boolean")
try:
if transfer_codec.normalize_format(envelope.get("format")) == transfer_codec.LEGACY_FORMAT:
errors.append("format: legacy_png must be sent as JSON")
except transfer_codec.TransferCodecError as exc:
errors.append(f"format: {exc}")
audio_payload = envelope.get("audio")
if audio_payload is not None and not isinstance(audio_payload, dict):
errors.append("audio: expected object when provided")
return errors
async def _job_complete_compressed(request):
try:
form = await request.post()
except Exception as exc:
too_large = getattr(web, "HTTPRequestEntityTooLarge", None)
if too_large is not None and isinstance(exc, too_large):
return await handle_api_error(
request, f"Request body exceeds the master's upload limit (--max-upload-size): {exc}", 413)
return await handle_api_error(request, f"Invalid multipart payload: {exc}", 400)
raw_envelope = form.get("envelope")
if isinstance(raw_envelope, (bytes, bytearray)):
raw_envelope = raw_envelope.decode("utf-8")
elif hasattr(raw_envelope, "file"):
raw_envelope = raw_envelope.file.read().decode("utf-8")
try:
envelope = json.loads(raw_envelope) if isinstance(raw_envelope, str) else None
except ValueError as exc:
return await handle_api_error(request, f"envelope: invalid JSON ({exc})", 400)
if envelope is None:
return await handle_api_error(request, "envelope: missing field", 400)
errors = _validate_compressed_envelope(envelope)
if errors:
return await handle_api_error(request, errors, 400)
parts = []
for field in form.getall("media", []):
if hasattr(field, "file"):
parts.append(field.file.read())
elif isinstance(field, (bytes, bytearray)):
parts.append(bytes(field))
else:
return await handle_api_error(request, "media: expected binary file parts", 400)
loop = asyncio.get_running_loop()
try:
tensors = await loop.run_in_executor(None, _decode_compressed_chunk_sync, envelope, parts)
decoded_audio = _decode_audio_payload(envelope.get("audio")) if envelope.get("audio") is not None else None
except ValueError as exc:
return await handle_api_error(request, exc, 400)
multi_job_id = envelope["job_id"].strip()
worker_id = envelope["worker_id"].strip()
start_index = envelope["start_index"]
is_last = envelope["is_last"]
last_pos = len(tensors) - 1
items = [
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": start_index + pos,
"is_last": bool(is_last and pos == last_pos),
"audio": decoded_audio if pos == last_pos else None,
}
for pos, tensor in enumerate(tensors)
]
queue_size = await _enqueue_worker_items(multi_job_id, items)
if queue_size is None:
return await handle_api_error(request, "job not initialized", 404)
debug_log(
f"job_complete received {envelope['format']} chunk - job_id: {multi_job_id}, worker: {worker_id}, "
f"frames: {start_index}..{start_index + len(items) - 1}, "
f"bytes: {sum(len(p) for p in parts)}, is_last: {is_last}, queue_size: {queue_size}"
)
return web.json_response({"status": "success", "frames": len(items)})
@server.PromptServer.instance.routes.post("/distributed/job_complete")
async def job_complete_endpoint(request):
content_type = str(getattr(request, "content_type", "") or "").lower()
if content_type.startswith("multipart/"):
try:
return await _job_complete_compressed(request)
except Exception as e:
return await handle_api_error(request, e)
try:
data = await request.json()
except Exception as exc:
@@ -316,28 +453,20 @@ async def job_complete_endpoint(request):
multi_job_id = job_id.strip()
worker_id = worker_id.strip()
pending = None
queue_size = 0
deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
while pending is None:
async with prompt_server.distributed_jobs_lock:
pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
if pending is not None:
await pending.put(
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": int(batch_idx),
"is_last": is_last,
"audio": decoded_audio,
}
)
queue_size = pending.qsize()
break
if time.monotonic() > deadline:
return await handle_api_error(request, "job not initialized", 404)
await asyncio.sleep(0.05)
queue_size = await _enqueue_worker_items(
multi_job_id,
[
{
"tensor": tensor,
"worker_id": worker_id,
"image_index": int(batch_idx),
"is_last": is_last,
"audio": decoded_audio,
}
],
)
if queue_size is None:
return await handle_api_error(request, "job not initialized", 404)
debug_log(
f"job_complete received canonical envelope - job_id: {multi_job_id}, "