Merge feat/runpod-serverless-workers (now including main) into feat/dynamic-image-list-queue

Resolve api/queue_orchestration.py, nodes/v3.py and
tests/helpers/v3_runtime_check.py to keep both the worker sink (#99)
and the image-list nodes. Stub nodes.image_list in the V3 collector
transfer test.

api/job_routes.py

@@ -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}, "

api/orchestration/media_sync.py

@@ -148,6 +148,30 @@ async def fetch_worker_path_separator(worker, trace_execution_id=None):
         return None
 
 
+async def fetch_worker_has_node(worker, node_class, trace_execution_id=None):
+    """Best-effort check that a worker has ``node_class`` registered.
+
+    Uses ComfyUI's ``GET /object_info/{node_class}``, which returns ``{}`` for
+    unknown classes. Any failure counts as "not available" so callers fall
+    back to nodes every ComfyUI install has.
+    """
+    url = build_worker_url(worker, f"/object_info/{node_class}")
+    session = await get_client_session()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            if resp.status != 200:
+                return False
+            payload = await resp.json()
+            return isinstance(payload, dict) and node_class in payload
+    except Exception as exc:
+        message = f"Failed to probe node '{node_class}' on worker {worker.get('id')}: {exc}"
+        if trace_execution_id:
+            trace_debug(trace_execution_id, message)
+        else:
+            debug_log(f"[Distributed] {message}")
+        return False
+
+
 async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
     """Upload one media file to worker iff missing or hash-mismatched."""
     session = await get_client_session()

api/orchestration/prompt_transform.py

@@ -328,8 +328,19 @@ def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-def prune_prompt_for_worker(prompt_obj):
-    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+WORKER_SINK_CLASS = "DistributedWorkerSink"
+
+
+def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
+    """Prune worker prompt to distributed nodes and their upstream dependencies.
+
+    Removing downstream nodes can leave the worker prompt without an output
+    node. One is re-added after each such distributed node: the hidden
+    ``DistributedWorkerSink`` when ``use_worker_sink`` is set (the worker must
+    have it registered), otherwise PreviewImage/PreviewAudio for workers
+    running older versions. The preview nodes write every frame to the
+    worker's temp directory, which the sink avoids.
+    """
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
     upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
     distributed_ids = collector_ids + upscale_ids
@@ -364,7 +375,13 @@ def prune_prompt_for_worker(prompt_obj):
                 and len(inputs["audio"]) == 2
             )
 
-            if image_connected:
+            if use_worker_sink and (image_connected or audio_connected):
+                pruned_prompt[next_id()] = {
+                    "inputs": {"value": [dist_id, 0 if image_connected else 1]},
+                    "class_type": WORKER_SINK_CLASS,
+                    "_meta": {"title": "Distributed Worker Sink (auto-added)"},
+                }
+            elif image_connected:
                 preview_id = next_id()
                 pruned_prompt[preview_id] = {
                     "inputs": {"images": [dist_id, 0]},

api/queue_orchestration.py

@@ -21,7 +21,12 @@ from .orchestration.dispatch import (
     select_active_workers,
     select_least_busy_worker,
 )
-from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
+from .orchestration.media_sync import (
+    convert_paths_for_platform,
+    fetch_worker_has_node,
+    fetch_worker_path_separator,
+    sync_worker_media,
+)
 from .orchestration.prompt_transform import (
     PromptIndex,
     apply_participant_overrides,
@@ -29,6 +34,7 @@ from .orchestration.prompt_transform import (
     generate_job_id_map,
     prepare_delegate_master_prompt,
     prune_prompt_for_worker,
+    WORKER_SINK_CLASS,
     _resolve_direct_workers_for_image_list,
 )
 from ..image_list.runtime import validate_image_list_regions
@@ -187,12 +193,21 @@ async def _prepare_worker_payload(
 
         worker_type = str(worker.get("type") or "local").strip().lower()
         is_remote_like = bool(worker.get("host")) and worker_type != "local"
-        if is_remote_like:
-            path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
-            if path_separator:
-                worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+        # Older workers lack the sink node; they keep the PreviewImage terminator.
+        has_sink_task = asyncio.ensure_future(
+            fetch_worker_has_node(worker, WORKER_SINK_CLASS, trace_execution_id=trace_execution_id)
+        )
+        try:
+            if is_remote_like:
+                path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
+                if path_separator:
+                    worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+            use_worker_sink = await has_sink_task
+        finally:
+            if not has_sink_task.done():
+                has_sink_task.cancel()
 
-        worker_prompt = prune_prompt_for_worker(worker_prompt)
+        worker_prompt = prune_prompt_for_worker(worker_prompt, use_worker_sink=use_worker_sink)
         worker_prompt = apply_participant_overrides(
             worker_prompt,
             worker["id"],